FileSystemDataSourceDiscovery [skip ci]

cpp/src/arrow/dataset/discovery.h

@@ -87,6 +87,7 @@ class ARROW_DS_EXPORT DataSourceDiscovery {
   ExpressionPtr root_partition_;
 };
 
+// TODO(kszucs): a static Defaults method would be nice
 struct FileSystemDiscoveryOptions {
   // For the purposes of applying the partition scheme, paths will be stripped
   // of the partition_base_dir. Files not matching the partition_base_dir

python/pyarrow/_dataset.pyx

@@ -29,7 +29,7 @@ from pyarrow.lib cimport *
 # from pyarrow.lib import ArrowException
 # from pyarrow.lib import as_buffer
 # from pyarrow.lib cimport RecordBatch, MemoryPool, Schema, check_status
-from pyarrow._fs cimport FileSystem
+from pyarrow._fs cimport FileSystem, Selector
 from pyarrow.includes.libarrow_dataset cimport *
 
 
@@ -100,26 +100,122 @@ cdef class FileFormat:
         self.wrapped = sp
         self.format = sp.get()
 
+    cdef inline shared_ptr[CFileFormat] unwrap(self):
+        return self.wrapped
+
     # @property
     # def name(self):
 
 
-cdef class CsvFileFormat(FileFormat):
-    pass
-
-
-cdef class FeatherFileFormat(FileFormat):
-    pass
-
-
-cdef class JsonFileFormat(FileFormat):
-    pass
-
-
 cdef class ParquetFileFormat(FileFormat):
     pass
 
 
+cdef class FileSystemDiscoveryOptions:
+
+    cdef:
+        CFileSystemDiscoveryOptions options
+
+     # Avoid mistakingly creating attributes
+    __slots__ = ()
+
+    def __init__(self, partition_base_dir=None, exclude_invalid_files=None,
+                 list ignore_prefixes=None):
+        if partition_base_dir is not None:
+            self.partition_base_dir = partition_base_dir
+        if exclude_invalid_files is not None:
+            self.exclude_invalid_files = exclude_invalid_files
+        if ignore_prefixes is not None:
+            self.ignore_prefixes = ignore_prefixes
+
+    cdef inline CFileSystemDiscoveryOptions unwrap(self):
+        return self.options
+
+    @property
+    def partition_base_dir(self):
+        return frombytes(self.options.partition_base_dir)
+
+    @partition_base_dir.setter
+    def partition_base_dir(self, value):
+        self.options.partition_base_dir = tobytes(value)
+
+    @property
+    def exclude_invalid_files(self):
+        return self.options.exclude_invalid_files
+
+    @exclude_invalid_files.setter
+    def exclude_invalid_files(self, bint value):
+        self.options.exclude_invalid_files = value
+
+    @property
+    def ignore_prefixes(self):
+        return [frombytes(p) for p in self.options.ignore_prefixes]
+
+    @ignore_prefixes.setter
+    def ignore_prefixes(self, values):
+        self.options.ignore_prefixes = [tobytes(v) for v in values]
+
+
+cdef class DataSourceDiscovery:
+
+    cdef:
+        shared_ptr[CDataSourceDiscovery] wrapped
+        CDataSourceDiscovery* discovery
+
+    cdef init(self, shared_ptr[CDataSourceDiscovery]& sp):
+        self.wrapped = sp
+        self.discovery = sp.get()
+
+    @staticmethod
+    cdef wrap(shared_ptr[CDataSourceDiscovery]& sp):
+        cdef DataSourceDiscovery self = \
+            DataSourceDiscovery.__new__(DataSourceDiscovery)
+        self.init(sp)
+        return self
+
+    cdef inline shared_ptr[CDataSourceDiscovery] unwrap(self):
+        return self.wrapped
+
+    def inspect(self):
+        cdef CResult[shared_ptr[CSchema]] result
+        with nogil:
+            result = self.discovery.Inspect()
+        return pyarrow_wrap_schema(GetResultValue(result))
+
+    def schema(self):
+        return pyarrow_wrap_schema(self.discovery.schema())
+
+    def finish(self):
+        cdef CResult[shared_ptr[CDataSource]] result
+        with nogil:
+            result = self.discovery.Finish()
+        return DataSource.wrap(GetResultValue(result))
+
+
+cdef class FileSystemDataSourceDiscovery(DataSourceDiscovery):
+
+    cdef:
+        CFileSystemDataSourceDiscovery* filesystem_discovery
+
+    def __init__(self, FileSystem filesystem, Selector selector,
+                 FileFormat format, FileSystemDiscoveryOptions options=None):
+        # TODO(kszucs): support instantiating from explicit paths
+        cdef CResult[shared_ptr[CDataSourceDiscovery]] result
+
+        options = options or FileSystemDiscoveryOptions()
+        result = CFileSystemDataSourceDiscovery.Make(
+            filesystem.unwrap(),
+            selector.unwrap(),
+            format.unwrap(),
+            options.unwrap()
+        )
+        self.init(GetResultValue(result))
+
+    cdef init(self, shared_ptr[CDataSourceDiscovery]& sp):
+        DataSourceDiscovery.init(self, sp)
+        self.filesystem_discovery = <CFileSystemDataSourceDiscovery*> sp.get()
+
+
 cdef class FileSource:
 
     cdef:
@@ -401,14 +497,15 @@ cdef class Dataset:
         self.wrapped = sp
         self.dataset = sp.get()
 
+    cdef inline shared_ptr[CDataset] unwrap(self):
+        return self.wrapped
+
     # TODO(kszucs): pass ScanContext
     def new_scan(self):
-        cdef:
-            shared_ptr[CScannerBuilder] builder
-            shared_ptr[CScanner] scanner
+        cdef shared_ptr[CScannerBuilder] builder
         builder = GetResultValue(self.dataset.NewScan())
-        scanner = GetResultValue(builder.get().Finish())
-        return Scanner.wrap(scanner)
+        # return ScannerBuilder(self, context)
+        return ScannerBuilder.wrap(builder)
 
     @property
     def sources(self):
@@ -426,8 +523,10 @@ cdef class ScanOptions:
         shared_ptr[CScanOptions] wrapped
         CScanOptions* options
 
-    def __init__(self):
+    def __init__(self, Schema schema=None):
         cdef shared_ptr[CScanOptions] options = CScanOptions.Defaults()
+        if schema is not None:
+            options.get().schema = pyarrow_unwrap_schema(schema)
         self.init(options)
 
     cdef init(self, const shared_ptr[CScanOptions]& sp):
@@ -522,6 +621,37 @@ cdef class ScannerBuilder:
         shared_ptr[CScannerBuilder] wrapped
         CScannerBuilder* builder
 
+    def __init__(self, Dataset dataset, ScanContext context):
+        cdef shared_ptr[CScannerBuilder] builder
+        builder = make_shared[CScannerBuilder](
+            dataset.unwrap(), context.unwrap()
+        )
+        self.init(builder)
+
+    cdef void init(self, shared_ptr[CScannerBuilder]& sp):
+        self.wrapped = sp
+        self.builder = sp.get()
+
+    @staticmethod
+    cdef wrap(shared_ptr[CScannerBuilder]& sp):
+        cdef ScannerBuilder self = ScannerBuilder.__new__(ScannerBuilder)
+        self.init(sp)
+        return self
+
+    cdef inline shared_ptr[CScannerBuilder] unwrap(self):
+        return self.wrapped
+
+    def project(self, columns):
+        cdef vector[c_string] cols = [tobytes(c) for c in columns]
+        print(cols)
+        check_status(self.builder.Project(cols))
+
+    def finish(self):
+        return Scanner.wrap(GetResultValue(self.builder.Finish()))
+
+    # def filter_(self):
+    #     pass
+
 
 cdef class Scanner:
 

python/pyarrow/_fs.pxd

@@ -46,6 +46,8 @@ cdef class Selector:
     cdef:
         CSelector selector
 
+    cdef inline CSelector unwrap(self)
+
 
 cdef class FileSystem:
     cdef:

python/pyarrow/_fs.pyx

@@ -147,6 +147,9 @@ cdef class Selector:
         self.recursive = recursive
         self.allow_non_existent = allow_non_existent
 
+    cdef inline CSelector unwrap(self):
+        return self.selector
+
     @property
     def base_dir(self):
         return frombytes(self.selector.base_dir)

python/pyarrow/dataset.py

@@ -26,13 +26,17 @@ from pyarrow._dataset import (
     Dataset,
     DataSource,
     DataFragment,
+    ParquetFileFormat,
     FileSource,
     SimpleDataFragment,
     ScanOptions,
     ScanTask,
     Scanner,
+    ScannerBuilder,
     ScanOptions,
     ScanContext,
     SimpleDataSource,
-    TreeDataSource
+    TreeDataSource,
+    FileSystemDiscoveryOptions,
+    FileSystemDataSourceDiscovery
 )

python/pyarrow/includes/libarrow_dataset.pxd

@@ -30,6 +30,12 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
         pass
 
 
+cdef extern from "arrow/dataset/api.h" namespace "arrow::fs" nogil:
+
+    ctypedef shared_ptr[CFileSystem] CFileSystemPtr "arrow::fs::FileSystemPtr"
+    ctypedef vector[CFileStats] CFileStatsVector "arrow::fs::FileStatsVector"
+
+
 cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
 
     # cdef enum CFilter" arrow::dataset::Filter::type":
@@ -50,9 +56,13 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         pass
 
     cdef cppclass CScanOptions "arrow::dataset::ScanOptions":
+        # CExpressionPtr filter
+        # shared_ptr[CExpressionEvaluator] evaluator
+        shared_ptr[CSchema] schema
+        # shared_ptr[CRecordBatchProjector] projector
+
         @staticmethod
         shared_ptr[CScanOptions] Defaults()
-        # const shared_ptr[DataSelector]& selector()
 
     cdef cppclass CScanContext "arrow::dataset::ScanContext":
         CMemoryPool* pool
@@ -158,6 +168,8 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
                              CScanOptionsPtr opts,
                              CDataFragmentPtr* out)
 
+    ctypedef shared_ptr[CFileFormat] CFileFormatPtr "arrow::dataset::FileFormatPtr"
+
     cdef cppclass CFileBasedDataFragment \
             "arrow::dataset::FileBasedDataFragment"(CDataFragment):
         CFileBasedDataFragment(const CFileSource& source,
@@ -263,3 +275,34 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CParquetFragment(const CFileSource& source,
                          shared_ptr[CScanOptions] options)
         c_bool splittable()
+
+    ############################### Discovery #################################
+
+    cdef cppclass CFileSystemDiscoveryOptions "arrow::dataset::FileSystemDiscoveryOptions":
+        c_string partition_base_dir
+        c_bool exclude_invalid_files
+        vector[c_string] ignore_prefixes
+
+    cdef cppclass CDataSourceDiscovery "arrow::dataset::DataSourceDiscovery":
+        CResult[shared_ptr[CSchema]] Inspect()
+        CResult[CDataSourcePtr] Finish()
+        shared_ptr[CSchema] schema()
+        CStatus SetSchema(shared_ptr[CSchema])
+        # CPartitionSchemePtr partition_scheme()
+        # CStatus SetPartitionScheme(CPartitionSchemePtr partition_scheme)
+        # CExpressionPtr root_partition()
+        # CStatus SetRootPartition(CExpressionPtr partition)
+
+    ctypedef shared_ptr[CDataSourceDiscovery] CDataSourceDiscoveryPtr "arrow::dataset::DataSourceDiscovery"
+
+    cdef cppclass CFileSystemDataSourceDiscovery "arrow::dataset::FileSystemDataSourceDiscovery"(CDataSourceDiscovery):
+        @staticmethod
+        CResult[CDataSourceDiscoveryPtr] Make(CFileSystemPtr filesytem,
+                                              CFileStatsVector paths,
+                                              CFileFormatPtr format,
+                                              CFileSystemDiscoveryOptions options)
+        @staticmethod
+        CResult[CDataSourceDiscoveryPtr] Make(CFileSystemPtr filesytem,
+                                              CSelector,
+                                              CFileFormatPtr format,
+                                              CFileSystemDiscoveryOptions options)

python/pyarrow/tests/test_dataset.py

@@ -20,22 +20,26 @@ from collections.abc import Generator
 import pytest
 
 import pyarrow as pa
-from pyarrow.fs import _MockFileSystem
+from pyarrow.fs import _MockFileSystem, Selector as FileSelector
 
 try:
     from pyarrow.dataset import (
         Dataset,
         DataSource,
         DataFragment,
+        ParquetFileFormat,
         SimpleDataFragment,
         ScanOptions,
         FileSource,
         ScanTask,
         Scanner,
+        ScannerBuilder,
         ScanOptions,
         ScanContext,
         SimpleDataSource,
-        TreeDataSource
+        TreeDataSource,
+        FileSystemDiscoveryOptions,
+        FileSystemDataSourceDiscovery
     )
 except ImportError as e:
     ds = None
@@ -83,6 +87,11 @@ def tree_data_source(simple_data_source):
     return TreeDataSource([simple_data_source] * 2)
 
 
+@pytest.fixture
+def dataset(simple_data_source, tree_data_source):
+    return Dataset([simple_data_source, tree_data_source])
+
+
 def test_scan_options():
     pass
 
@@ -143,7 +152,10 @@ def test_dataset(simple_data_source, tree_data_source):
     for source in dataset.sources:
         assert isinstance(source, DataSource)
 
-    scanner = dataset.new_scan()
+    builder = dataset.new_scan()
+    assert isinstance(builder, ScannerBuilder)
+
+    scanner = builder.finish()
     assert isinstance(scanner, Scanner)
 
     for task in scanner.scan():
@@ -152,9 +164,10 @@ def test_dataset(simple_data_source, tree_data_source):
             assert isinstance(record_batch, pa.RecordBatch)
 
 
-def test_scanner(simple_data_source):
+def test_scanner(schema, simple_data_source):
     sources = [simple_data_source]
-    options = ScanOptions()
+    # FIXME(kszucs): if schema is not set to options it segfaults
+    options = ScanOptions(schema=schema)
     context = ScanContext()
 
     scanner = Scanner(sources, options, context)
@@ -163,8 +176,17 @@ def test_scanner(simple_data_source):
         for record_batch in task.scan():  # call it execute?
             assert isinstance(record_batch, pa.RecordBatch)
 
-    # table = scanner.to_table()
-    # assert isinstance(table, pa.Table)
+    table = scanner.to_table()
+    assert isinstance(table, pa.Table)
+
+
+def test_scanner_builder(dataset):
+    context = ScanContext()
+    builder = ScannerBuilder(dataset, context)
+
+    builder.project([])
+    scanner = builder.finish()
+    assert isinstance(scanner, Scanner)
 
 
 def test_projector():
@@ -185,3 +207,11 @@ def test_file_source(fs):
 
 def test_file_system_data_source():
     pass
+
+
+def test_file_system_discovery(fs):
+    selector = FileSelector('/base', recursive=True)
+    fileformat = ParquetFileFormat()
+    discovery = FileSystemDataSourceDiscovery(fs, selector, fileformat)
+
+