Test parquet data discovery [skip ci]

python/pyarrow/_dataset.pyx

@@ -93,9 +93,14 @@ cdef class FileFormat:
     cdef inline shared_ptr[CFileFormat] unwrap(self):
         return self.wrapped
 
+    def name(self):
+        return frombytes(self.format.name())
+
 
 cdef class ParquetFileFormat(FileFormat):
-    pass
+
+    def __init__(self):
+        self.init(shared_ptr[CFileFormat](new CParquetFileFormat()))
 
 
 cdef class PartitionScheme:
@@ -109,7 +114,7 @@ cdef class PartitionScheme:
         self.scheme = sp.get()
 
     @staticmethod
-    cdef wrap(shared_ptr[CPartitionScheme]& sp):
+    cdef wrap(const shared_ptr[CPartitionScheme]& sp):
         cdef PartitionScheme self
 
         # TODO(kszucs): either choose type() or name() but consistently
@@ -212,6 +217,14 @@ cdef class DataSourceDiscovery:
     cdef inline shared_ptr[CDataSourceDiscovery] unwrap(self):
         return self.wrapped
 
+    @property
+    def partition_scheme(self):
+        return PartitionScheme.wrap(self.discovery.partition_scheme())
+
+    @partition_scheme.setter
+    def partition_scheme(self, PartitionScheme scheme not None):
+        check_status(self.discovery.SetPartitionScheme(scheme.unwrap()))
+
     def inspect(self):
         cdef CResult[shared_ptr[CSchema]] result
         with nogil:
@@ -233,8 +246,9 @@ cdef class FileSystemDataSourceDiscovery(DataSourceDiscovery):
     cdef:
         CFileSystemDataSourceDiscovery* filesystem_discovery
 
-    def __init__(self, FileSystem filesystem, Selector selector,
-                 FileFormat format, FileSystemDiscoveryOptions options=None):
+    def __init__(self, FileSystem filesystem not None,
+                 Selector selector not None, FileFormat format not None,
+                 FileSystemDiscoveryOptions options=None):
         # TODO(kszucs): support instantiating from explicit paths
         cdef CResult[shared_ptr[CDataSourceDiscovery]] result
 
@@ -261,7 +275,7 @@ cdef class FileSource:
     cdef readonly:
         FileSystem fs
 
-    def __init__(self, path, FileSystem filesystem, compression=None):
+    def __init__(self, path, FileSystem filesystem not None, compression=None):
         cdef shared_ptr[CFileSource] source
 
         # need to hold a reference for the filesystem
@@ -413,10 +427,6 @@ cdef class DataSource:
     cdef shared_ptr[CDataSource] unwrap(self):
         return self.wrapped
 
-    @property
-    def type(self):
-        return frombytes(self.source.type())
-
     # def __iter__(self):
     #     for fragment in self.fragments():
     #         yield fragment
@@ -486,8 +496,8 @@ cdef class FileSystemDataSource(DataSource):
     cdef:
         CFileSystemDataSource* filesystem_source
 
-    def __init__(self, filesystem, file_stats, source_partition,
-                 path_partitions, file_format):
+    def __init__(self, FileSystem filesystem not None, file_stats,
+                 source_partition, path_partitions, file_format):
         cdef:
             DataSource child
             CDataSourceVector children
@@ -679,7 +689,6 @@ cdef class ScannerBuilder:
 
     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):

python/pyarrow/dataset.py

@@ -37,6 +37,8 @@ from pyarrow._dataset import (
     ScanContext,
     SimpleDataSource,
     TreeDataSource,
+    FileSystemDataSource,
     FileSystemDiscoveryOptions,
-    FileSystemDataSourceDiscovery
+    FileSystemDataSourceDiscovery,
+    SchemaPartitionScheme
 )

python/pyarrow/includes/libarrow_dataset.pxd

@@ -267,15 +267,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
 
     cdef cppclass CParquetFileFormat "arrow::dataset::ParquetFileFormat"(
             CFileFormat):
-        c_string name()
-        c_bool IsKnownExtension(const c_string& ext)
-        CStatus ScanFile(const CFileSource& source,
-                         shared_ptr[CScanOptions] scan_options,
-                         shared_ptr[CScanContext] scan_context,
-                         shared_ptr[CScanTaskIterator]* out)
-        CStatus MakeFragment(const CFileSource& source,
-                             shared_ptr[CScanOptions] opts,
-                             shared_ptr[CDataFragment]* out)
+        pass
 
     cdef cppclass CParquetFragment "arrow::dataset::ParquetFragment"(
             CFileBasedDataFragment):
@@ -283,6 +275,22 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
                          shared_ptr[CScanOptions] options)
         c_bool splittable()
 
+    ############################### Partitioning ##############################
+
+    cdef cppclass CPartitionScheme "arrow::dataset::PartitionScheme":
+        c_string name() const
+        CResult[CExpressionPtr] Parse(const c_string& path) const
+
+    ctypedef shared_ptr[CPartitionScheme] CPartitionSchemePtr "arrow::dataset::PartitionSchemePtr"
+
+    cdef cppclass CSchemaPartitionScheme "arrow::dataset::SchemaPartitionScheme"(CPartitionScheme):
+        CSchemaPartitionScheme(shared_ptr[CSchema] schema)
+        const shared_ptr[CSchema]& schema()
+
+    cdef cppclass CHivePartitionScheme "arrow::dataset::HivePartitionScheme"(CPartitionScheme):
+        CHivePartitionScheme(shared_ptr[CSchema] schema)
+        # vector[CUnconvertedKey] GetUnconvertedKeys(const c_string& path) const;
+
     ############################### Discovery #################################
 
     cdef cppclass CFileSystemDiscoveryOptions "arrow::dataset::FileSystemDiscoveryOptions":
@@ -295,8 +303,8 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CResult[CDataSourcePtr] Finish()
         shared_ptr[CSchema] schema()
         CStatus SetSchema(shared_ptr[CSchema])
-        # CPartitionSchemePtr partition_scheme()
-        # CStatus SetPartitionScheme(CPartitionSchemePtr partition_scheme)
+        CPartitionSchemePtr partition_scheme()
+        CStatus SetPartitionScheme(CPartitionSchemePtr partition_scheme)
         CExpressionPtr root_partition()
         CStatus SetRootPartition(CExpressionPtr partition)
 
@@ -313,17 +321,3 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
                                               CSelector,
                                               CFileFormatPtr format,
                                               CFileSystemDiscoveryOptions options)
-
-    ############################### Partitioning ##############################
-
-    cdef cppclass CPartitionScheme "arrow::dataset::PartitionScheme":
-        c_string name() const
-        CResult[CExpressionPtr] Parse(const c_string& path) const
-
-    cdef cppclass CSchemaPartitionScheme "arrow::dataset::SchemaPartitionScheme"(CPartitionScheme):
-        CSchemaPartitionScheme(shared_ptr[CSchema] schema)
-        const shared_ptr[CSchema]& schema()
-
-    cdef cppclass CHivePartitionScheme "arrow::dataset::HivePartitionScheme"(CPartitionScheme):
-        CHivePartitionScheme(shared_ptr[CSchema] schema)
-        # vector[CUnconvertedKey] GetUnconvertedKeys(const c_string& path) const;

python/pyarrow/tests/test_dataset.py

@@ -16,6 +16,10 @@
 # under the License.
 
 from collections.abc import Generator
+try:
+    import pathlib
+except ImportError:
+    import pathlib2 as pathlib  # py2 compat
 
 import pytest
 
@@ -38,8 +42,10 @@ try:
         ScanContext,
         SimpleDataSource,
         TreeDataSource,
+        FileSystemDataSource,
         FileSystemDiscoveryOptions,
-        FileSystemDataSourceDiscovery
+        FileSystemDataSourceDiscovery,
+        SchemaPartitionScheme
     )
 except ImportError as e:
     ds = None
@@ -51,9 +57,23 @@ pytestmark = pytest.mark.dataset
 
 
 @pytest.fixture
-def fs():
-    return _MockFileSystem()
+@pytest.mark.parquet
+def fs(table):
+    import pyarrow.parquet as pq
 
+    fs = _MockFileSystem()
+
+    directories = [
+        'subdir/1/xxx',
+        'subdir/2/yyy',
+    ]
+    for i, directory in enumerate(directories):
+        path = '{}/file{}.parquet'.format(directory, i)
+        fs.create_dir(directory)
+        with fs.open_output_stream(path) as out:
+            pq.write_table(table, out)
+
+    return fs
 
 @pytest.fixture
 def schema():
@@ -72,6 +92,11 @@ def record_batch(schema):
     return pa.record_batch(data, schema=schema)
 
 
+@pytest.fixture
+def table(record_batch):
+    return pa.Table.from_batches([record_batch] * 10)
+
+
 @pytest.fixture
 def simple_data_fragment(record_batch):
     return SimpleDataFragment([record_batch] * 5)
@@ -116,9 +141,7 @@ def test_simple_data_fragment(record_batch):
 
 def test_simple_data_source(simple_data_fragment):
     source = SimpleDataSource([simple_data_fragment] * 4)
-
     assert isinstance(source, SimpleDataSource)
-    assert source.type == 'simple_data_source'
 
     result = source.fragments()
     assert isinstance(result, Generator)
@@ -131,9 +154,7 @@ def test_simple_data_source(simple_data_fragment):
 
 def test_tree_data_source(simple_data_source):
     source = TreeDataSource([simple_data_source] * 2)
-
     assert isinstance(source, TreeDataSource)
-    assert source.type == 'tree_data_source'
 
     result = source.fragments()
     assert isinstance(result, Generator)
@@ -183,7 +204,6 @@ def test_scanner(schema, simple_data_source):
 def test_scanner_builder(dataset):
     context = ScanContext()
     builder = ScannerBuilder(dataset, context)
-
     builder.project([])
     scanner = builder.finish()
     assert isinstance(scanner, Scanner)
@@ -196,7 +216,6 @@ def test_projector():
 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 == 0  # None
@@ -210,6 +229,40 @@ def test_file_system_data_source():
 
 
 def test_file_system_discovery(fs):
-    selector = FileSelector('/base', recursive=True)
-    fileformat = ParquetFileFormat()
-    discovery = FileSystemDataSourceDiscovery(fs, selector, fileformat)
+    selector = FileSelector('subdir', recursive=True)
+    assert selector.base_dir == 'subdir'
+    assert selector.recursive is True
+
+    format = ParquetFileFormat()
+    assert format.name() == 'parquet'
+
+    options = FileSystemDiscoveryOptions('/')
+    assert options.partition_base_dir == '/'
+    assert options.ignore_prefixes == ['.', '_']
+    assert options.exclude_invalid_files is True
+
+    discovery = FileSystemDataSourceDiscovery(fs, selector, format, options)
+    assert isinstance(discovery.inspect(), pa.Schema)
+    # assert isinstance(discovery.schema(), pa.Schema)
+    assert isinstance(discovery.finish(), FileSystemDataSource)
+
+    scheme = SchemaPartitionScheme(
+        pa.schema([
+            pa.field('group', pa.int32()),
+            pa.field('key', pa.string())
+        ])
+    )
+    discovery.partition_scheme = scheme
+    assert isinstance(discovery.partition_scheme, SchemaPartitionScheme)
+
+    data_source = discovery.finish()
+    assert isinstance(data_source, DataSource)
+
+    inspected_schema = discovery.inspect()
+    dataset = Dataset([data_source], inspected_schema)
+
+    scanner = dataset.new_scan().finish()
+    for task in scanner.scan():
+        assert isinstance(task, ScanTask)
+        for record_batch in task.scan():
+            assert isinstance(record_batch, pa.RecordBatch)

python/pyarrow/tests/test_fs.py

@@ -26,7 +26,8 @@ import pytest
 import pyarrow as pa
 from pyarrow.tests.test_io import gzip_compress, gzip_decompress
 from pyarrow.fs import (FileType, Selector, FileSystem, LocalFileSystem,
-                        LocalFileSystemOptions, SubTreeFileSystem)
+                        LocalFileSystemOptions, SubTreeFileSystem,
+                        _MockFileSystem)
 
 
 @pytest.fixture
@@ -40,6 +41,16 @@ def localfs(request, tempdir):
     )
 
 
+@pytest.fixture
+def mockfs(request):
+    return dict(
+        fs=_MockFileSystem(),
+        pathfn=lambda p: p,
+        allow_move_dir=True,
+        allow_append_to_file=True,
+    )
+
+
 @pytest.fixture
 def localfs_with_mmap(request, tempdir):
     return dict(
@@ -141,6 +152,10 @@ def hdfs(request, hdfs_server):
         pytest.lazy_fixture('hdfs'),
         id='HadoopFileSystem'
     ),
+    pytest.param(
+        pytest.lazy_fixture('mockfs'),
+        id='_MockFileSystem()'
+    ),
 ])
 def filesystem_config(request):
     return request.param