More expressions and testing

python/pyarrow/_dataset.pyx

@@ -27,7 +27,7 @@ from cython.operator cimport dereference as deref
 from pyarrow.lib cimport *
 from pyarrow.includes.libarrow_dataset cimport *
 from pyarrow.compat import frombytes, tobytes
-from pyarrow._fs cimport FileSystem, Selector
+from pyarrow._fs cimport FileSystem, FileStats, Selector
 
 
 def _forbid_instantiation(klass):
@@ -130,6 +130,11 @@ cdef class PartitionScheme:
     cdef inline shared_ptr[CPartitionScheme] unwrap(self):
         return self.wrapped
 
+    def parse(self, path):
+        cdef CResult[shared_ptr[CExpression]] result
+        result = self.scheme.Parse(tobytes(path))
+        return Expression.wrap(GetResultValue(result))
+
 
 cdef class SchemaPartitionScheme(PartitionScheme):
 
@@ -242,7 +247,12 @@ cdef class DataSourceDiscovery:
 
     @property
     def partition_scheme(self):
-        return PartitionScheme.wrap(self.discovery.partition_scheme())
+        cdef shared_ptr[CPartitionScheme] scheme
+        scheme = self.discovery.partition_scheme()
+        if scheme.get() == nullptr:
+            return None
+        else:
+            return PartitionScheme.wrap(scheme)
 
     @partition_scheme.setter
     def partition_scheme(self, PartitionScheme scheme not None):
@@ -272,9 +282,7 @@ cdef class FileSystemDataSourceDiscovery(DataSourceDiscovery):
     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
-
         options = options or FileSystemDiscoveryOptions()
         result = CFileSystemDataSourceDiscovery.Make(
             filesystem.unwrap(),
@@ -531,20 +539,33 @@ cdef class FileSystemDataSource(DataSource):
         CFileSystemDataSource* filesystem_source
 
     def __init__(self, FileSystem filesystem not None, file_stats,
-                 source_partition, path_partitions, file_format):
+                 Expression source_partition, dict path_partitions,
+                 FileFormat file_format not None):
         cdef:
-            DataSource child
-            CDataSourceVector children
-            shared_ptr[CFileSystemDataSource] filesystem_source
+            FileStats stats
+            Expression expression
+            vector[CFileStats] c_file_stats
+            shared_ptr[CExpression] c_source_partition
+            unordered_map[c_string, shared_ptr[CExpression]] c_path_partitions
+            CResult[shared_ptr[CDataSource]] result
 
-        # TODO(kszucs)
-        # filesystem_source = CFileSystemDataSource.Make(
-        #     fs::FileSystem* filesystem, std::vector<fs::FileStats> stats,
-        #     std::shared_ptr<Expression> source_partition,
-        #     PathPartitions partitions, std::shared_ptr<FileFormat> format,
-        #     std::shared_ptr<DataSource>* out
-        # )
-        self.init(<shared_ptr[CDataSource]> filesystem_source)
+        for stats in file_stats:
+            c_file_stats.push_back(stats.unwrap())
+
+        for path, expression in path_partitions.items():
+            c_path_partitions[tobytes(path)] = expression.unwrap()
+
+        if source_partition is not None:
+            c_source_partition = source_partition.unwrap()
+
+        result = CFileSystemDataSource.Make(
+            filesystem.unwrap(),
+            c_file_stats,
+            c_source_partition,
+            c_path_partitions,
+            file_format.unwrap()
+        )
+        self.init(GetResultValue(result))
 
     cdef void init(self, const shared_ptr[CDataSource]& sp):
         DataSource.init(self, sp)
@@ -815,18 +836,19 @@ cdef class Expression:
             self = ScalarExpression.__new__(ScalarExpression)
         elif typ == CExpressionType_NOT:
             self = NotExpression.__new__(NotExpression)
-        # elif typ == CExpressionType_CAST:
-        #     self = CastExpression.__new__(CastExpression)
+        elif typ == CExpressionType_CAST:
+            self = CastExpression.__new__(CastExpression)
         elif typ == CExpressionType_AND:
             self = AndExpression.__new__(AndExpression)
         elif typ == CExpressionType_OR:
             self = OrExpression.__new__(OrExpression)
         elif typ == CExpressionType_COMPARISON:
             self = ComparisonExpression.__new__(ComparisonExpression)
-        # elif typ == CExpressionType_IS_VALID:
-        #     self = IsValidExpression.__new__(IsValidExpression)
-        # elif typ == CExpressionType_IN:
-        #     self = InExpression.__new__(InExpression)
+        elif typ == CExpressionType_IS_VALID:
+            self = IsValidExpression.__new__(IsValidExpression)
+        elif typ == CExpressionType_IN:
+            self = InExpression.__new__(InExpression)
+        # TODO(kszucs): implement it
         # elif typ == CExpressionType_CUSTOM:
         #     self = CustomExpression.__new__(CustomExpression)
         else:
@@ -844,12 +866,6 @@ cdef class Expression:
     def __str__(self):
         return frombytes(self.expression.ToString())
 
-    def __eq__(self, other):
-        try:
-            return self.equals(other)
-        except TypeError:
-            return NotImplemented
-
     def validate(self, Schema schema not None):
         cdef:
             shared_ptr[CSchema] sp_schema
@@ -913,6 +929,7 @@ cdef class ScalarExpression(Expression):
         Expression.init(self, sp)
         self.scalar = <CScalarExpression*> sp.get()
 
+    # TODO(kszucs): implement once we have proper Scalar bindings
     # @property
     # def value(self):
     #     return pyarrow_wrap_scalar(self.scalar.value())
@@ -971,6 +988,117 @@ cdef class ComparisonExpression(BinaryExpression):
         return <CompareOperator> self.comparison.op()
 
 
+cdef class IsValidExpression(UnaryExpression):
+
+    def __init__(self, Expression operand not None):
+        cdef shared_ptr[CIsValidExpression] expression
+        expression = make_shared[CIsValidExpression](operand.unwrap())
+        self.init(<shared_ptr[CExpression]> expression)
+
+
+cdef class CastOptions:
+
+    cdef:
+        CCastOptions options
+
+    __slots__ = ()  # avoid mistakingly creating attributes
+
+    def __init__(self, allow_int_overflow=None,
+                 allow_time_truncate=None, allow_time_overflow=None,
+                 allow_float_truncate=None, allow_invalid_utf8=None):
+        if allow_int_overflow is not None:
+            self.allow_int_overflow = allow_int_overflow
+        if allow_time_truncate is not None:
+            self.allow_time_truncate = allow_time_truncate
+        if allow_time_overflow is not None:
+            self.allow_time_overflow = allow_time_overflow
+        if allow_float_truncate is not None:
+            self.allow_float_truncate = allow_float_truncate
+        if allow_invalid_utf8 is not None:
+            self.allow_invalid_utf8 = allow_invalid_utf8
+
+    @staticmethod
+    cdef wrap(CCastOptions options):
+        cdef CastOptions self = CastOptions.__new__(CastOptions)
+        self.options = options
+        return self
+
+    @staticmethod
+    def safe():
+        return CastOptions.wrap(CCastOptions.Safe())
+
+    @staticmethod
+    def unsafe():
+        return CastOptions.wrap(CCastOptions.Unsafe())
+
+    cdef inline CCastOptions unwrap(self):
+        return self.options
+
+    @property
+    def allow_int_overflow(self):
+        return self.options.allow_int_overflow
+
+    @allow_int_overflow.setter
+    def allow_int_overflow(self, bint flag):
+        self.options.allow_int_overflow = flag
+
+    @property
+    def allow_time_truncate(self):
+        return self.options.allow_time_truncate
+
+    @allow_time_truncate.setter
+    def allow_time_truncate(self, bint flag):
+        self.options.allow_time_truncate = flag
+
+    @property
+    def allow_time_overflow(self):
+        return self.options.allow_time_overflow
+
+    @allow_time_overflow.setter
+    def allow_time_overflow(self, bint flag):
+        self.options.allow_time_overflow = flag
+
+    @property
+    def allow_float_truncate(self):
+        return self.options.allow_float_truncate
+
+    @allow_float_truncate.setter
+    def allow_float_truncate(self, bint flag):
+        self.options.allow_float_truncate = flag
+
+    @property
+    def allow_invalid_utf8(self):
+        return self.options.allow_invalid_utf8
+
+    @allow_invalid_utf8.setter
+    def allow_invalid_utf8(self, bint flag):
+        self.options.allow_invalid_utf8 = flag
+
+
+cdef class CastExpression(UnaryExpression):
+
+    def __init__(self, Expression operand not None, DataType to not None,
+                 CastOptions options=None):
+        cdef shared_ptr[CExpression] expression
+        options = options or CastOptions.safe()
+        expression.reset(new CCastExpression(
+            operand.unwrap(),
+            pyarrow_unwrap_data_type(to),
+            options.unwrap()
+        ))
+        self.init(expression)
+
+
+cdef class InExpression(UnaryExpression):
+
+    def __init__(self, Expression operand not None, Array haystack not None):
+        cdef shared_ptr[CExpression] expression
+        expression.reset(
+            new CInExpression(operand.unwrap(), pyarrow_unwrap_array(haystack))
+        )
+        self.init(expression)
+
+
 cdef class NotExpression(UnaryExpression):
 
     def __init__(self, Expression operand not None):

python/pyarrow/_fs.pxd

@@ -41,6 +41,8 @@ cdef class FileStats:
     @staticmethod
     cdef wrap(CFileStats stats)
 
+    cdef inline CFileStats unwrap(self)
+
 
 cdef class Selector:
     cdef:

python/pyarrow/_fs.pyx

@@ -53,6 +53,9 @@ cdef class FileStats:
         self.stats = stats
         return self
 
+    cdef inline CFileStats unwrap(self):
+        return self.stats
+
     def __repr__(self):
         def getvalue(attr):
             try:

python/pyarrow/dataset.py

@@ -24,6 +24,8 @@ if sys.version_info < (3,):
 
 from pyarrow._dataset import (  # noqa
     AndExpression,
+    CastExpression,
+    CastOptions,
     CompareOperator,
     ComparisonExpression,
     DataFragment,
@@ -37,6 +39,8 @@ from pyarrow._dataset import (  # noqa
     FileSystemDataSourceDiscovery,
     FileSystemDiscoveryOptions,
     HivePartitionScheme,
+    InExpression,
+    IsValidExpression,
     NotExpression,
     OrExpression,
     ParquetFileFormat,

python/pyarrow/includes/libarrow.pxd

@@ -1365,11 +1365,15 @@ cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
     cdef cppclass CCastOptions" arrow::compute::CastOptions":
         CCastOptions()
         CCastOptions(c_bool safe)
+        @staticmethod
         CCastOptions Safe()
+        @staticmethod
         CCastOptions Unsafe()
         c_bool allow_int_overflow
         c_bool allow_time_truncate
+        c_bool allow_time_overflow
         c_bool allow_float_truncate
+        c_bool allow_invalid_utf8
 
     cdef cppclass CTakeOptions" arrow::compute::TakeOptions":
         pass

python/pyarrow/includes/libarrow_dataset.pxd

@@ -17,7 +17,7 @@
 
 # distutils: language = c++
 
-from libcpp.functional cimport function
+from libcpp.unordered_map cimport unordered_map
 
 from pyarrow.includes.common cimport *
 from pyarrow.includes.libarrow cimport *
@@ -126,6 +126,14 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
             CBinaryExpression):
         pass
 
+    cdef cppclass CCastExpression "arrow::dataset::CastExpression"(
+            CExpression):
+        CCastExpression(CExpressionPtr operand, shared_ptr[CDataType] to,
+                        CCastOptions options)
+
+    cdef cppclass CInExpression "arrow::dataset::InExpression"(CExpression):
+        CInExpression(CExpressionPtr operand, shared_ptr[CArray] set)
+
     cdef shared_ptr[CNotExpression] MakeNotExpression "arrow::dataset::not_"(
         CExpressionPtr operand)
     cdef CExpressionPtr MakeAndExpression "arrow::dataset::and_"(
@@ -275,14 +283,17 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         CFileFormatPtr format()
         CScanOptionsPtr scan_options()
 
+    ctypedef unordered_map[c_string, CExpressionPtr] CPathPartitions \
+        "arrow::dataset::PathPartitions"
+
     cdef cppclass CFileSystemDataSource \
             "arrow::dataset::FileSystemDataSource"(CDataSource):
         @staticmethod
-        CStatus Make(CFileSystem* filesystem,
-                     const CSelector& selector,
-                     CFileFormatPtr format,
-                     CScanOptionsPtr scan_options,
-                     shared_ptr[CFileSystemDataSource]* out)
+        CResult[CDataSourcePtr] Make(CFileSystemPtr filesystem,
+                                     CFileStatsVector stats,
+                                     CExpressionPtr source_partition,
+                                     CPathPartitions partitions,
+                                     CFileFormatPtr format)
         c_string type()
         shared_ptr[CDataFragmentIterator] GetFragments(CScanOptionsPtr options)
 

python/pyarrow/tests/test_dataset.py

@@ -15,7 +15,6 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from collections.abc import Generator
 import pytest
 
 import pyarrow as pa
@@ -109,7 +108,6 @@ def test_simple_data_fragment(record_batch):
     assert fragment.splittable is False
     assert isinstance(fragment.scan_options, ds.ScanOptions)
 
-    assert isinstance(fragment.scan(), Generator)
     tasks = list(fragment.scan())
     assert len(tasks) == 5
     for task, batch in zip(tasks, batches):
@@ -125,8 +123,6 @@ def test_simple_data_source(simple_data_fragment):
     assert source.partition_expression is None
 
     result = source.fragments()
-    assert isinstance(result, Generator)
-
     fragments = list(result)
     assert len(fragments) == 4
     for fragment in fragments:
@@ -138,14 +134,58 @@ def test_tree_data_source(simple_data_source):
     assert isinstance(source, ds.TreeDataSource)
 
     result = source.fragments()
-    assert isinstance(result, Generator)
-
     fragments = list(result)
     assert len(fragments) == 8
     for fragment in fragments:
         assert isinstance(fragment, ds.DataFragment)
 
 
+def test_filesystem_data_source(mockfs):
+    file_stats = mockfs.get_target_stats([
+        'subdir/1/xxx/file0.parquet',
+        'subdir/2/yyy/file1.parquet',
+    ])
+    file_format = ds.ParquetFileFormat()
+    source_partition = None
+    path_partitions = {}
+
+    source = ds.FileSystemDataSource(
+        mockfs,
+        file_stats,
+        source_partition=source_partition,
+        path_partitions=path_partitions,
+        file_format=file_format
+    )
+    assert len(list(source.fragments())) == 2
+
+    source_partition = ds.ComparisonExpression(
+        ds.CompareOperator.Equal,
+        ds.FieldExpression('source'),
+        ds.ScalarExpression(1337)
+    )
+    path_partitions = {
+        'subdir/1/xxx/file0.parquet': ds.ComparisonExpression(
+            ds.CompareOperator.Equal,
+            ds.FieldExpression('part'),
+            ds.ScalarExpression(1)
+        ),
+        'subdir/2/xxx/file1.parquet': ds.ComparisonExpression(
+            ds.CompareOperator.Equal,
+            ds.FieldExpression('part'),
+            ds.ScalarExpression(2)
+        )
+    }
+    source = ds.FileSystemDataSource(
+        mockfs,
+        file_stats,
+        source_partition=source_partition,
+        path_partitions=path_partitions,
+        file_format=file_format
+    )
+    assert len(list(source.fragments())) == 2
+    assert source.partition_expression.equals(source_partition)
+
+
 def test_dataset(simple_data_source, tree_data_source, schema):
     dataset = ds.Dataset([simple_data_source, tree_data_source], schema)
 
@@ -231,6 +271,53 @@ def test_partition_scheme(schema):
         scheme = klass(schema)
         assert isinstance(scheme, ds.PartitionScheme)
 
+    scheme = ds.SchemaPartitionScheme(
+        pa.schema([
+            pa.field('group', pa.int64()),
+            pa.field('key', pa.float64())
+        ])
+    )
+    expr = scheme.parse('/3/3.14')
+    assert isinstance(expr, ds.Expression)
+
+    expected = ds.AndExpression(
+        ds.ComparisonExpression(
+            ds.CompareOperator.Equal,
+            ds.FieldExpression('group'),
+            ds.ScalarExpression(3)
+        ),
+        ds.ComparisonExpression(
+            ds.CompareOperator.Equal,
+            ds.FieldExpression('key'),
+            ds.ScalarExpression(3.14)
+        )
+    )
+    assert expr.equals(expected)
+
+    with pytest.raises(pa.ArrowInvalid):
+        scheme.parse('/prefix/3/aaa')
+
+    scheme = ds.HivePartitionScheme(
+        pa.schema([
+            pa.field('alpha', pa.int64()),
+            pa.field('beta', pa.int64())
+        ])
+    )
+    expr = scheme.parse('/alpha=0/beta=3')
+    expected = ds.AndExpression(
+        ds.ComparisonExpression(
+            ds.CompareOperator.Equal,
+            ds.FieldExpression('alpha'),
+            ds.ScalarExpression(0)
+        ),
+        ds.ComparisonExpression(
+            ds.CompareOperator.Equal,
+            ds.FieldExpression('beta'),
+            ds.ScalarExpression(3)
+        )
+    )
+    assert expr.equals(expected)
+
 
 def test_expression(schema):
     a = ds.ScalarExpression(1)
@@ -245,13 +332,15 @@ def test_expression(schema):
     assert isinstance(and_.right_operand, ds.Expression)
     assert and_.equals(ds.AndExpression(a, b))
     assert and_.equals(and_)
-    assert and_ == ds.AndExpression(a, b)
-    assert and_ != 'other object'
 
     ds.AndExpression(a, b, c)
     ds.OrExpression(a, b)
     ds.OrExpression(a, b, c)
     ds.NotExpression(ds.OrExpression(a, b, c))
+    ds.IsValidExpression(a)
+    ds.CastExpression(a, pa.int32())
+    ds.CastExpression(a, pa.int32(), ds.CastOptions.unsafe())
+    ds.InExpression(a, pa.array([1, 2, 3]))
 
     condition = ds.ComparisonExpression(
         ds.CompareOperator.Greater,
@@ -270,10 +359,8 @@ def test_expression(schema):
         ds.FieldExpression('i64'),
         ds.ScalarExpression(7)
     )
-    true_ = ds.ScalarExpression(True)
-    false_ = ds.ScalarExpression(False)
-    assert condition.assume(i64_is_5) == false_
-    assert condition.assume(i64_is_7) == true_
+    assert condition.assume(i64_is_5).equals(ds.ScalarExpression(False))
+    assert condition.assume(i64_is_7).equals(ds.ScalarExpression(True))
     assert str(condition) == "(i64 > 5:int64)"
 
 
@@ -293,8 +380,9 @@ def test_file_system_discovery(mockfs):
     discovery = ds.FileSystemDataSourceDiscovery(mockfs, selector, format,
                                                  options)
     assert isinstance(discovery.inspect(), pa.Schema)
-    # assert isinstance(discovery.schema(), pa.Schema)
+    assert isinstance(discovery.schema(), pa.Schema)
     assert isinstance(discovery.finish(), ds.FileSystemDataSource)
+    assert discovery.partition_scheme is None
 
     scheme = ds.SchemaPartitionScheme(
         pa.schema([