HivePartitionScheme

python/pyarrow/_dataset.pyx

@@ -104,6 +104,9 @@ cdef class PartitionScheme:
         shared_ptr[CPartitionScheme] wrapped
         CPartitionScheme* scheme
 
+    def __init__(self):
+        _forbid_instantiation(self.__class__)
+
     cdef init(self, const shared_ptr[CPartitionScheme]& sp):
         self.wrapped = sp
         self.scheme = sp.get()
@@ -116,6 +119,8 @@ cdef class PartitionScheme:
         typ = frombytes(sp.get().name())
         if typ == 'schema_partition_scheme':
             self = SchemaPartitionScheme.__new__(SchemaPartitionScheme)
+        elif typ == 'hive_partition_scheme':
+            self = HivePartitionScheme.__new__(HivePartitionScheme)
         else:
             raise TypeError(typ)
 
@@ -131,7 +136,7 @@ cdef class SchemaPartitionScheme(PartitionScheme):
     cdef:
         CSchemaPartitionScheme* schema_scheme  # hmmm...
 
-    def __init__(self, Schema schema):
+    def __init__(self, Schema schema not None):
         cdef shared_ptr[CSchemaPartitionScheme] scheme
         scheme = make_shared[CSchemaPartitionScheme](
             pyarrow_unwrap_schema(schema)
@@ -147,6 +152,27 @@ cdef class SchemaPartitionScheme(PartitionScheme):
         return pyarrow_wrap_schema(self.schema_scheme.schema())
 
 
+cdef class HivePartitionScheme(PartitionScheme):
+
+    cdef:
+        CHivePartitionScheme* hive_scheme
+
+    def __init__(self, Schema schema not None):
+        cdef shared_ptr[CHivePartitionScheme] scheme
+        scheme = make_shared[CHivePartitionScheme](
+            pyarrow_unwrap_schema(schema)
+        )
+        self.init(<shared_ptr[CPartitionScheme]> scheme)
+
+    cdef init(self, const shared_ptr[CPartitionScheme]& sp):
+        PartitionScheme.init(self, sp)
+        self.hive_scheme = <CHivePartitionScheme*> sp.get()
+
+    @property
+    def schema(self):
+        return pyarrow_wrap_schema(self.hive_scheme.schema())
+
+
 cdef class FileSystemDiscoveryOptions:
 
     cdef:
@@ -446,6 +472,15 @@ cdef class DataSource:
             else:
                 yield DataFragment.wrap(fragment)
 
+    @property
+    def partition_expression(self):
+        cdef shared_ptr[CExpression] expression
+        expression = self.source.partition_expression()
+        if expression.get() == nullptr:
+            return None
+        else:
+            return Expression.wrap(expression)
+
 
 cdef class SimpleDataSource(DataSource):
 

python/pyarrow/dataset.py

@@ -40,9 +40,11 @@ from pyarrow._dataset import (  # noqa
     SimpleDataSource,
     TreeDataSource,
     FileSystemDataSource,
+    HivePartitionScheme,
     FileSystemDiscoveryOptions,
     FileSystemDataSourceDiscovery,
     SchemaPartitionScheme,
+    PartitionScheme,
     ScalarExpression,
     FieldExpression,
     AndExpression,

python/pyarrow/includes/libarrow_dataset.pxd

@@ -329,6 +329,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
     cdef cppclass CHivePartitionScheme \
             "arrow::dataset::HivePartitionScheme"(CPartitionScheme):
         CHivePartitionScheme(shared_ptr[CSchema] schema)
+        const shared_ptr[CSchema]& schema()
         # vector[CUnconvertedKey] GetUnconvertedKeys(const c_string& path) const;
 
     ############################### Discovery #################################

python/pyarrow/tests/test_dataset.py

@@ -122,6 +122,7 @@ def test_simple_data_fragment(record_batch):
 def test_simple_data_source(simple_data_fragment):
     source = ds.SimpleDataSource([simple_data_fragment] * 4)
     assert isinstance(source, ds.SimpleDataSource)
+    assert source.partition_expression is None
 
     result = source.fragments()
     assert isinstance(result, Generator)
@@ -216,13 +217,23 @@ def test_file_source(mockfs):
     ds.FileFormat,
     ds.DataFragment,
     ds.DataSource,
-    ds.Expression
+    ds.Expression,
+    ds.PartitionScheme,
 ])
 def test_abstract_classes(klass):
     with pytest.raises(TypeError):
         klass()
 
 
+@pytest.mark.parametrize('klass', [
+    ds.SchemaPartitionScheme,
+    ds.HivePartitionScheme
+])
+def test_partition_scheme(schema, klass):
+    scheme = klass(schema)
+    assert isinstance(scheme, ds.PartitionScheme)
+
+
 def test_expression(schema):
     a = ds.ScalarExpression(1)
     b = ds.ScalarExpression(1.1)