Working scanner [skip ci]

python/pyarrow/_dataset.pyx

@@ -124,6 +124,37 @@ cdef class FileSource:
 
     cdef:
         shared_ptr[CFileSource] wrapped
+        CFileSource* source
+
+    cdef readonly:
+        FileSystem fs
+
+    def __init__(self, path, FileSystem filesystem, compression=None):
+        cdef shared_ptr[CFileSource] source
+
+        # need to hold a reference for the filesystem
+        self.fs = filesystem
+
+        source.reset(new CFileSource(
+            tobytes(path),
+            self.fs.unwrap().get(),
+            _get_compression_type(compression)
+        ))
+
+        self.init(source)
+
+    cdef init(self, shared_ptr[CFileSource]& sp):
+        self.wrapped = sp
+        self.source = sp.get()
+
+    @staticmethod
+    cdef wrap(shared_ptr[CFileSource]& sp):
+        cdef FileSource self = FileSource.__new__(FileSource)
+        self.init(sp)
+        return self
+
+    cdef inline shared_ptr[CFileSource] unwrap(self):
+        return self.wrapped
 
     def equals(self, FileSource other):
         return self.wrapped.get() == other.wrapped.get()
@@ -142,10 +173,6 @@ cdef class FileSource:
     def path(self):
         return frombytes(self.wrapped.get().path())
 
-    # @property
-    # def fs(self):
-    #     FileSystem.wrap(self.wrapped.get().filesystem())
-
     # def open(self):
     #     CStatus Open(shared_ptr[CRandomAccessFile]* out)
 
@@ -160,11 +187,14 @@ cdef class DataFragment:
         self.fragment = sp.get()
 
     @staticmethod
-    cdef DataFragment wrap(shared_ptr[CDataFragment]& sp):
+    cdef wrap(shared_ptr[CDataFragment]& sp):
         cdef DataFragment self = DataFragment.__new__(DataFragment)
         self.init(sp)
         return self
 
+    cdef inline shared_ptr[CDataFragment] unwrap(self):
+        return self.wrapped
+
     def scan(self, ScanContext context=None):
         cdef:
             CResult[CScanTaskIterator] iterator_result
@@ -283,11 +313,11 @@ cdef class SimpleDataSource(DataSource):
     def __init__(self, data_fragments):
         cdef:
             DataFragment fragment
-            vector[shared_ptr[CDataFragment]] fragments
+            CDataFragmentVector fragments
             shared_ptr[CSimpleDataSource] simple_source
 
         for fragment in data_fragments:
-            fragments.push_back(fragment.wrapped)
+            fragments.push_back(fragment.unwrap())
 
         simple_source = make_shared[CSimpleDataSource](fragments)
         self.init(<shared_ptr[CDataSource]> simple_source)
@@ -305,7 +335,7 @@ cdef class TreeDataSource(DataSource):
     def __init__(self, data_sources):
         cdef:
             DataSource child
-            vector[shared_ptr[CDataSource]] children
+            CDataSourceVector children
             shared_ptr[CTreeDataSource] tree_source
 
         for child in data_sources:
@@ -328,7 +358,7 @@ cdef class FileSystemDataSource(DataSource):
                  path_partitions, file_format):
         cdef:
             DataSource child
-            vector[shared_ptr[CDataSource]] children
+            CDataSourceVector children
             shared_ptr[CFileSystemDataSource] filesystem_source
 
         # TODO(kszucs)
@@ -371,19 +401,13 @@ cdef class Dataset:
         self.wrapped = sp
         self.dataset = sp.get()
 
+    # TODO(kszucs): pass ScanContext
     def new_scan(self):
         cdef:
-            CResult[CScannerBuilderPtr] builder_result
-            CResult[CScannerPtr] scanner_result
-            CScannerBuilderPtr builder
-            CScannerPtr scanner
-
-        builder_result = self.dataset.NewScan()
-        builder = GetResultValue(builder_result)
-
-        scanner_result = builder.get().Finish()
-        scanner = GetResultValue(scanner_result)
-
+            shared_ptr[CScannerBuilder] builder
+            shared_ptr[CScanner] scanner
+        builder = GetResultValue(self.dataset.NewScan())
+        scanner = GetResultValue(builder.get().Finish())
         return Scanner.wrap(scanner)
 
     @property
@@ -403,7 +427,7 @@ cdef class ScanOptions:
         CScanOptions* options
 
     def __init__(self):
-        cdef shared_ptr[CScanOptions] options
+        cdef shared_ptr[CScanOptions] options = CScanOptions.Defaults()
         self.init(options)
 
     cdef init(self, const shared_ptr[CScanOptions]& sp):
@@ -411,7 +435,7 @@ cdef class ScanOptions:
         self.options = sp.get()
 
     @staticmethod
-    cdef ScanOptions wrap(const shared_ptr[CScanOptions]& sp):
+    cdef wrap(const shared_ptr[CScanOptions]& sp):
         cdef ScanOptions self = ScanOptions.__new__(ScanOptions)
         self.init(sp)
         return self
@@ -423,48 +447,48 @@ cdef class ScanOptions:
 cdef class ScanContext:
 
     cdef:
-        CScanContextPtr wrapped
-        CScanContext* context
+        shared_ptr[CScanContext] wrapped
+        CScanContext *context
 
     def __init__(self, MemoryPool memory_pool=None):
-        cdef CScanContextPtr context
+        cdef shared_ptr[CScanContext] context = make_shared[CScanContext]()
 
         if memory_pool is not None:
             context.get().pool = memory_pool.pool
 
         self.init(context)
 
-    cdef init(self, CScanContextPtr& sp):
+    cdef init(self, shared_ptr[CScanContext]& sp):
         self.wrapped = sp
         self.context = sp.get()
 
     @staticmethod
-    cdef ScanContext wrap(CScanContextPtr& sp):
+    cdef ScanContext wrap(shared_ptr[CScanContext]& sp):
         cdef ScanContext self = ScanContext.__new__(ScanContext)
         self.init(sp)
         return self
 
-    cdef inline CScanContextPtr unwrap(self):
+    cdef inline shared_ptr[CScanContext] unwrap(self):
         return self.wrapped
 
 
 cdef class ScanTask:
 
     cdef:
-        CScanTaskPtr wrapped
+        shared_ptr[CScanTask] wrapped
         CScanTask* task
 
-    cdef init(self, CScanTaskPtr& sp):
+    cdef init(self, shared_ptr[CScanTask]& sp):
         self.wrapped = sp
         self.task = self.wrapped.get()
 
     @staticmethod
-    cdef wrap(CScanTaskPtr& sp):
+    cdef wrap(shared_ptr[CScanTask]& sp):
         cdef SimpleScanTask self = SimpleScanTask.__new__(SimpleScanTask)
         self.init(sp)
         return self
 
-    cdef inline CScanTaskPtr unwrap(self):
+    cdef inline shared_ptr[CScanTask] unwrap(self):
         return self.wrapped
 
     def scan(self):
@@ -487,55 +511,52 @@ cdef class SimpleScanTask(ScanTask):
     cdef:
         CSimpleScanTask* simple_task
 
-    cdef init(self, CScanTaskPtr& sp):
+    cdef init(self, shared_ptr[CScanTask]& sp):
         ScanTask.init(self, sp)
         self.simple_task = <CSimpleScanTask*> sp.get()
 
 
-# cdef class ScannerBuilder:
+cdef class ScannerBuilder:
 
-#     cdef:
-#         CScannerBuilderPtr wrapped
-#         CScannerBuilder* builder
+    cdef:
+        shared_ptr[CScannerBuilder] wrapped
+        CScannerBuilder* builder
 
 
 cdef class Scanner:
 
     cdef:
-        CScannerPtr wrapped
+        shared_ptr[CScanner] wrapped
         CScanner* scanner
 
-    # def __init__(self, data_sources, ScanOptions options, ScanContext context):
-    #     cdef:
-    #         DataSource source
-    #         CDataSourceVector sources
-    #         CSimpleScanner* simple_scanner
-    #         unique_ptr[CScanner] scanner
+    def __init__(self, data_sources, ScanOptions options, ScanContext context):
+        cdef:
+            DataSource source
+            CDataSourceVector sources
+            shared_ptr[CScanner] scanner
 
-    #     for source in data_sources:
-    #         sources.push_back(source.unwrap())
+        for source in data_sources:
+            sources.push_back(source.unwrap())
 
-    #     simple_scanner = new CSimpleScanner(
-    #         sources, options.unwrap(), context.unwrap()
-    #     )
-    #     scanner.reset(simple_scanner)
-    #     self.init(scanner)
+        scanner = make_shared[CScanner](
+            sources, options.unwrap(), context.unwrap()
+        )
+        self.init(scanner)
 
-    cdef void init(self, CScannerPtr& sp):
+    cdef void init(self, shared_ptr[CScanner]& sp):
         self.wrapped = sp
-        self.scanner = self.wrapped.get()
+        self.scanner = sp.get()
 
     @staticmethod
-    cdef wrap(CScannerPtr& sp):
+    cdef wrap(shared_ptr[CScanner]& sp):
         cdef Scanner self = Scanner.__new__(Scanner)
         self.init(sp)
         return self
 
     def scan(self):
         cdef:
-            CResult[CScanTaskIterator] iterator_result
             CScanTaskIterator iterator
-            CScanTaskPtr task
+            shared_ptr[CScanTask] task
 
         iterator = move(GetResultValue(move(self.scanner.Scan())))
 
@@ -547,16 +568,9 @@ cdef class Scanner:
                 yield ScanTask.wrap(task)
 
     def to_table(self):
-        cdef:
-            shared_ptr[CTable] table
-            CResult[shared_ptr[CTable]] result
+        cdef CResult[shared_ptr[CTable]] result
 
         with nogil:
             result = self.scanner.ToTable()
-        table = GetResultValue(result)
 
-        return pyarrow_wrap_table(table)
-
-
-# cdef class ScannerBuilder:
-#     pass
+        return pyarrow_wrap_table(GetResultValue(result))

python/pyarrow/_fs.pxd

@@ -57,6 +57,8 @@ cdef class FileSystem:
     @staticmethod
     cdef wrap(shared_ptr[CFileSystem]& sp)
 
+    cdef inline shared_ptr[CFileSystem] unwrap(self)
+
 
 cdef class LocalFileSystem(FileSystem):
     cdef:

python/pyarrow/_fs.pyx

@@ -188,18 +188,23 @@ cdef class FileSystem:
     cdef wrap(shared_ptr[CFileSystem]& sp):
         cdef FileSystem self
 
-        typ = tobytes(sp.get().type())
-        print(FileSystem.__subclasses__)
-        for klass in FileSystem.__subclasses__:
-            if klass.type_id == typ:
-                self = klass.__new__(klass)
-                break
+        # TODO(kszucs): add S3
+        typ = frombytes(sp.get().type())
+        if typ == 'local':
+            self = LocalFileSystem.__new__(LocalFileSystem)
+        elif typ == 'mock':
+            self = _MockFileSystem.__new__(_MockFileSystem)
+        elif typ == 'subtree':
+            self = SubTreeFileSystem.__new__(SubTreeFileSystem)
         else:
             raise TypeError('Cannot wrap FileSystem pointer')
 
         self.init(sp)
         return self
 
+    cdef inline shared_ptr[CFileSystem] unwrap(self):
+        return self.wrapped
+
     def get_target_stats(self, paths_or_selector):
         """Get statistics for the given target.
 

python/pyarrow/dataset.py

@@ -26,6 +26,7 @@ from pyarrow._dataset import (
     Dataset,
     DataSource,
     DataFragment,
+    FileSource,
     SimpleDataFragment,
     ScanOptions,
     ScanTask,

python/pyarrow/includes/libarrow_dataset.pxd

@@ -36,11 +36,6 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     #     CFilter_EXPRESSION" arrow::dataset::Filter::type::EXPRESSION"
     #     CFilter_GENERIC" arrow::dataset::Filter::GENERIC"
 
-
-    cdef cppclass CDataSourceVector "arrow::dataset::DataSourceVector"(
-            vector[shared_ptr[CDataSource]]):
-        pass
-
     cdef cppclass CWriteOptions "arrow::dataset::WriteOptions":
         pass
 
@@ -55,110 +50,80 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         pass
 
     cdef cppclass CScanOptions "arrow::dataset::ScanOptions":
+        @staticmethod
+        shared_ptr[CScanOptions] Defaults()
         # const shared_ptr[DataSelector]& selector()
-        pass
-
-    cdef cppclass CScanOptionsPtr "arrow::dataset::ScanOptionsPtr"(
-            shared_ptr[CScanOptions]):
-        pass
 
     cdef cppclass CScanContext "arrow::dataset::ScanContext":
         CMemoryPool* pool
 
-    cdef cppclass CScanContextPtr "arrow::dataset::ScanContextPtr"(
-            shared_ptr[CScanContext]):
-        pass
+    ctypedef shared_ptr[CScanOptions] CScanOptionsPtr "arrow::dataset::ScanOptionsPtr"
+    ctypedef shared_ptr[CScanContext] CScanContextPtr "arrow::dataset::ScanContextPtr"
 
     cdef cppclass CScanTask" arrow::dataset::ScanTask":
         CResult[CRecordBatchIterator] Scan()
 
-    cdef cppclass CScanTaskPtr "arrow::dataset::ScanTaskPtr"(
-            shared_ptr[CScanTask]):
-        pass
-
-    cdef cppclass CScanTaskIterator "arrow::dataset::ScanTaskIterator"(
-            CIterator[CScanTaskPtr]):
-        pass
-
     cdef cppclass CSimpleScanTask "arrow::dataset::SimpleScanTask"(CScanTask):
         pass
 
+    ctypedef shared_ptr[CScanTask] CScanTaskPtr "arrow::dataset::ScanTaskPtr"
+    ctypedef CIterator[CScanTaskPtr] CScanTaskIterator "arrow::dataset::ScanTaskIterator"
+
+    cdef cppclass CScanner "arrow::dataset::Scanner":
+        CResult[CScanTaskIterator] Scan()
+        CResult[shared_ptr[CTable]] ToTable()
+
+    ctypedef shared_ptr[CScanner] CScannerPtr "arrow::dataset::ScannerPtr"
+
     cdef cppclass CScannerBuilder "arrow::dataset::ScannerBuilder":
-        CScannerBuilder(shared_ptr[CDataset] dataset,
-                        shared_ptr[CScanContext] scan_context)
-        # CStatus Project(const vector[c_string]& columns)
+        CScannerBuilder(shared_ptr[CDataset], CScanContextPtr scan_context)
+        CStatus Project(const vector[c_string]& columns)
         # CStatus Filter(const CExpression& filter)
         # CStatus Filter(shared_ptr[CExpression] filter)
         # CStatus UseThreads(c_bool use_threads)
         CResult[CScannerPtr] Finish()
         shared_ptr[CSchema] schema() const
 
-    cdef cppclass CScannerBuilderPtr "arrow::dataset::ScannerBuilderPtr"(
-            shared_ptr[CScannerBuilder]):
-        pass
-
-    cdef cppclass CScanner "arrow::dataset::Scanner":
-        CResult[CScanTaskIterator] Scan()
-        CResult[shared_ptr[CTable]] ToTable()
-
-    cdef cppclass CScannerPtr "arrow::dataset::ScannerPtr"(
-            shared_ptr[CScanner]):
-        pass
+    ctypedef shared_ptr[CScannerBuilder] CScannerBuilderPtr "arrow::dataset::ScannerBuilderPtr"
 
     cdef cppclass CDataFragment "arrow::dataset::DataFragment":
         CResult[CScanTaskIterator] Scan(CScanContextPtr context)
         c_bool splittable()
         CScanOptionsPtr scan_options()
 
-    cdef cppclass CDataFragmentPtr "arrow::dataset::DataFragmentPtr"(
-            shared_ptr[CDataFragment]):
-        pass
-
-    cdef cppclass CDataFragmentIterator "arrow::dataset::DataFragmentIterator"(
-            CIterator[CDataFragmentPtr]):
-        pass
-
-    cdef cppclass CDataFragmentVector "arrow::dataset::DataFragmentVector"(
-            vector[CDataFragmentPtr]):
-        pass
+    ctypedef shared_ptr[CDataFragment] CDataFragmentPtr "arrow::dataset::DataFragmentPtr"
+    ctypedef vector[CDataFragmentPtr] CDataFragmentVector "arrow::dataset::DataFragmentVector"
+    ctypedef CIterator[CDataFragmentPtr] CDataFragmentIterator "arrow::dataset::DataFragmentIterator"
 
     cdef cppclass CSimpleDataFragment "arrow::dataset::SimpleDataFragment"(
             CDataFragment):
         CSimpleDataFragment(vector[shared_ptr[CRecordBatch]] record_batches,
-                            shared_ptr[CScanOptions] scan_options)
+                            CScanOptionsPtr scan_options)
 
     cdef cppclass CDataSource "arrow::dataset::DataSource":
-        CDataFragmentIterator GetFragments(shared_ptr[CScanOptions] options)
+        CDataFragmentIterator GetFragments(CScanOptionsPtr options)
         # const shared_ptr[CExpression]& partition_expression()
         c_string type()
 
-    cdef cppclass CDataSourcePtr "arrow::dataset::DataSourcePtr"(
-            shared_ptr[CDataSource]):
+    cdef cppclass CSimpleDataSource "arrow::dataset::SimpleDataSource"(CDataSource):
         pass
 
-    cdef cppclass CDataSourceVector "arrow::dataset::DataSourceVector"(
-            vector[shared_ptr[CDataSource]]):
+    cdef cppclass CTreeDataSource "arrow::dataset::TreeDataSource"(CDataSource):
         pass
 
-    cdef cppclass CSimpleDataSource "arrow::dataset::SimpleDataSource"(
-            CDataSource):
-        pass
-
-    cdef cppclass CTreeDataSource "arrow::dataset::TreeDataSource"(
-            CDataSource):
-        pass
+    ctypedef shared_ptr[CDataSource] CDataSourcePtr "arrow::dataset::DataSourcePtr"
+    ctypedef vector[CDataSourcePtr] CDataSourceVector "arrow::dataset::DataSourceVector"
 
     cdef cppclass CDataset "arrow::dataset::Dataset":
         @staticmethod
-        CResult[CDatasetPtr] Make(CDataSourceVector sources,
-                                  shared_ptr[CSchema] schema)
+        CResult[shared_ptr[CDataset]] Make(CDataSourceVector sources,
+                                           shared_ptr[CSchema] schema)
+        CResult[CScannerBuilderPtr] NewScan(CScanContextPtr context)
         CResult[CScannerBuilderPtr] NewScan()
         const CDataSourceVector& sources()
         shared_ptr[CSchema] schema()
 
-    cdef cppclass CDatasetPtr "arrow::dataset::DatasetPtr"(
-            shared_ptr[CDataset]):
-        pass
+    ctypedef shared_ptr[CDataset] CDatasetPtr "arrow::dataset::DatasetPtr"
 
     ############################### File ######################################
 
@@ -167,7 +132,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         c_string file_type()
 
     cdef cppclass CFileSource "arrow::dataset::FileSource":
-        CFileSource(c_string path, CFileSystem filesystem,
+        CFileSource(c_string path, CFileSystem* filesystem,
                     CompressionType compression)
         c_bool operator==(const CFileSource other)
         CompressionType compression()
@@ -186,23 +151,23 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CStatus Inspect(const CFileSource& source,
                         shared_ptr[CSchema]* out) const
         CStatus ScanFile(const CFileSource& source,
-                         shared_ptr[CScanOptions] scan_options,
-                         shared_ptr[CScanContext] scan_context,
+                         CScanOptionsPtr scan_options,
+                         CScanContextPtr scan_context,
                          CScanTaskIterator* out) const
         CStatus MakeFragment(const CFileSource& location,
-                             shared_ptr[CScanOptions] opts,
-                             shared_ptr[CDataFragment]* out)
+                             CScanOptionsPtr opts,
+                             CDataFragmentPtr* out)
 
     cdef cppclass CFileBasedDataFragment \
             "arrow::dataset::FileBasedDataFragment"(CDataFragment):
         CFileBasedDataFragment(const CFileSource& source,
                                shared_ptr[CFileFormat] format,
-                               shared_ptr[CScanOptions] scan_options)
-        CStatus Scan(shared_ptr[CScanContext] scan_context,
+                               CScanOptionsPtr scan_options)
+        CStatus Scan(CScanContextPtr scan_context,
                      shared_ptr[CScanTaskIterator]* out)
         const CFileSource& source()
         shared_ptr[CFileFormat] format()
-        shared_ptr[CScanOptions] scan_options()
+        CScanOptionsPtr scan_options()
 
     cdef cppclass CFileSystemDataSource \
             "arrow::dataset::FileSystemDataSource"(CDataSource):
@@ -210,11 +175,10 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CStatus Make(CFileSystem* filesystem,
                      const CSelector& selector,
                      shared_ptr[CFileFormat] format,
-                     shared_ptr[CScanOptions] scan_options,
+                     CScanOptionsPtr scan_options,
                      shared_ptr[CFileSystemDataSource]* out)
         c_string type()
-        shared_ptr[CDataFragmentIterator] GetFragments(
-            shared_ptr[CScanOptions] options)
+        shared_ptr[CDataFragmentIterator] GetFragments(CScanOptionsPtr options)
 
     ############################### File CSV ##################################
 

python/pyarrow/lib.pxd

@@ -533,6 +533,8 @@ cdef class _CRecordBatchReader:
         shared_ptr[CRecordBatchReader] reader
 
 
+cdef CompressionType _get_compression_type(object name) except *
+
 cdef get_input_stream(object source, c_bool use_memory_map,
                       shared_ptr[CInputStream]* reader)
 cdef get_reader(object source, c_bool use_memory_map,

python/pyarrow/tests/test_dataset.py

@@ -29,6 +29,7 @@ try:
         DataFragment,
         SimpleDataFragment,
         ScanOptions,
+        FileSource,
         ScanTask,
         Scanner,
         ScanOptions,
@@ -151,30 +152,35 @@ def test_dataset(simple_data_source, tree_data_source):
             assert isinstance(record_batch, pa.RecordBatch)
 
 
-# def test_scanner(simple_data_source):
-#     sources = [simple_data_source]
-#     options = ScanOptions()
-#     context = ScanContext()
-#     scanner = SimpleScanner(sources, options, context)
+def test_scanner(simple_data_source):
+    sources = [simple_data_source]
+    options = ScanOptions()
+    context = ScanContext()
 
-#     for task in scanner.scan():
-#         assert isinstance(task, ScanTask)
+    scanner = Scanner(sources, options, context)
+    for task in scanner.scan():
+        assert isinstance(task, ScanTask)
+        for record_batch in task.scan():  # call it execute?
+            assert isinstance(record_batch, pa.RecordBatch)
+
+    # table = scanner.to_table()
+    # assert isinstance(table, pa.Table)
 
 
 def test_projector():
     pass
 
 
-# def test_file_source(fs):
-#     source1 = FileSource('/path/to/file.ext', fs, compression=None)
-#     source2 = FileSource('/path/to/file.ext.gz', fs, compression='gzip')
+def test_file_source(fs):
+    source1 = FileSource('/path/to/file.ext', fs, compression=None)
+    source2 = FileSource('/path/to/file.ext.gz', fs, compression='gzip')
 
-#     assert source1.path == '/path/to/file.ext'
-#     assert source1.fs == fs
-#     assert source1.compression is None
-#     assert source2.path == '/path/to/file.ext.gz'
-#     assert source2.fs == fs
-#     assert source2.compression == 'gzip'
+    assert source1.path == '/path/to/file.ext'
+    assert source1.fs == fs
+    assert source1.compression == 0  # None
+    assert source2.path == '/path/to/file.ext.gz'
+    assert source2.fs == fs
+    assert source2.compression == 2  # 'gzip'
 
 
 def test_file_system_data_source():