ARROW-7413: [Python] Expose and test the partioning discovery

https://issues.apache.org/jira/browse/ARROW-7413

Closes #6151 from jorisvandenbossche/ARROW-7413-partition-discovery and squashes the following commits:

0a2e215 <Joris Van den Bossche> ARROW-7413:  Expose and test the partioning discovery

Authored-by: Joris Van den Bossche <[email protected]>
Signed-off-by: Krisztián Szűcs <[email protected]>

Author: jorisvandenbossche

python/pyarrow/_dataset.pyx

@@ -128,12 +128,16 @@ cdef class PartitionSchemeDiscovery:
     def __init__(self):
         _forbid_instantiation(self.__class__)
 
-    @staticmethod
-    cdef wrap(const shared_ptr[CPartitionSchemeDiscovery]& sp):
-        cdef PartitionSchemeDiscovery self
-        self = PartitionSchemeDiscovery()
+    cdef init(self, const shared_ptr[CPartitionSchemeDiscovery]& sp):
         self.wrapped = sp
         self.discovery = sp.get()
+
+    @staticmethod
+    cdef wrap(const shared_ptr[CPartitionSchemeDiscovery]& sp):
+        cdef PartitionSchemeDiscovery self = PartitionSchemeDiscovery.__new__(
+            PartitionSchemeDiscovery
+        )
+        self.init(sp)
         return self
 
     cdef inline shared_ptr[CPartitionSchemeDiscovery] unwrap(self):
@@ -196,6 +200,29 @@ cdef class SchemaPartitionScheme(PartitionScheme):
         PartitionScheme.init(self, sp)
         self.schema_scheme = <CSchemaPartitionScheme*> sp.get()
 
+    @staticmethod
+    def discover(field_names):
+        """
+        Discover a SchemaPartitionScheme.
+
+        Parameters
+        ----------
+        field_names : list of str
+            The names to associate with the values from the subdirectory names.
+
+        Returns
+        -------
+        PartionSchemeDiscovery
+            To be used in the FileSystemDiscoveryOptions.
+        """
+        cdef:
+            PartitionSchemeDiscovery discovery
+            vector[c_string] c_field_names
+        c_field_names = [tobytes(s) for s in field_names]
+        discovery = PartitionSchemeDiscovery.wrap(
+            CSchemaPartitionScheme.MakeDiscovery(c_field_names))
+        return discovery
+
 
 cdef class HivePartitionScheme(PartitionScheme):
     """
@@ -244,6 +271,22 @@ cdef class HivePartitionScheme(PartitionScheme):
         PartitionScheme.init(self, sp)
         self.hive_scheme = <CHivePartitionScheme*> sp.get()
 
+    @staticmethod
+    def discover():
+        """
+        Discover a HivePartitionScheme.
+
+        Returns
+        -------
+        PartionSchemeDiscovery
+            To be used in the FileSystemDiscoveryOptions.
+        """
+        cdef:
+            PartitionSchemeDiscovery discovery
+        discovery = PartitionSchemeDiscovery.wrap(
+            CHivePartitionScheme.MakeDiscovery())
+        return discovery
+
 
 cdef class FileSystemDiscoveryOptions:
     """

python/pyarrow/dataset.py

@@ -45,6 +45,7 @@
     OrExpression,
     ParquetFileFormat,
     PartitionScheme,
+    PartitionSchemeDiscovery,
     ScalarExpression,
     Scanner,
     ScannerBuilder,

python/pyarrow/includes/libarrow_dataset.pxd

@@ -301,10 +301,15 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     cdef cppclass CSchemaPartitionScheme \
             "arrow::dataset::SchemaPartitionScheme"(CPartitionScheme):
         CSchemaPartitionScheme(shared_ptr[CSchema] schema)
+        @staticmethod
+        shared_ptr[CPartitionSchemeDiscovery] MakeDiscovery(
+            vector[c_string] field_names)
 
     cdef cppclass CHivePartitionScheme \
             "arrow::dataset::HivePartitionScheme"(CPartitionScheme):
         CHivePartitionScheme(shared_ptr[CSchema] schema)
+        @staticmethod
+        shared_ptr[CPartitionSchemeDiscovery] MakeDiscovery()
 
     cdef cppclass CPartitionSchemeOrDiscovery \
             "arrow::dataset::PartitionSchemeOrDiscovery":

python/pyarrow/tests/test_dataset.py

@@ -352,3 +352,29 @@ def test_file_system_discovery(mockfs, paths_or_selector):
     assert isinstance(table, pa.Table)
     assert len(table) == 10
     assert table.num_columns == 4
+
+
+def test_partition_scheme_discovery(mockfs):
+    paths_or_selector = fs.FileSelector('subdir', recursive=True)
+    format = ds.ParquetFileFormat()
+
+    options = ds.FileSystemDiscoveryOptions('subdir')
+    schema_discovery = ds.SchemaPartitionScheme.discover(['group', 'key'])
+    assert isinstance(schema_discovery, ds.PartitionSchemeDiscovery)
+    options.partition_scheme_discovery = schema_discovery
+
+    discovery = ds.FileSystemDataSourceDiscovery(
+        mockfs, paths_or_selector, format, options
+    )
+    inspected_schema = discovery.inspect()
+    # i64/f64 from data, group/key from "/1/xxx" and "/2/yyy" paths
+    expected_schema = pa.schema([
+        ("i64", pa.int64()),
+        ("f64", pa.float64()),
+        ("group", pa.int32()),
+        ("key", pa.string()),
+    ])
+    assert inspected_schema.remove_metadata().equals(expected_schema)
+
+    hive_discovery = ds.HivePartitionScheme.discover()
+    assert isinstance(hive_discovery, ds.PartitionSchemeDiscovery)