Expose Scalar/Comparison/Boolean expressions [skip ci]

python/pyarrow/_dataset.pyx

@@ -24,6 +24,7 @@ import six
 
 from cython.operator cimport dereference as deref
 
+import pyarrow as pa
 from pyarrow.compat import frombytes, tobytes
 from pyarrow.lib cimport *
 # from pyarrow.lib import ArrowException
@@ -823,20 +824,29 @@ cdef class ScalarExpression(Expression):
 
     cdef CScalarExpression* scalar
 
-    def __init__(self, Scalar value not None):
+    def __init__(self, value):
         cdef:
             shared_ptr[CScalar] scalar
             shared_ptr[CScalarExpression] expression
-        scalar = pyarrow_unwrap_scalar(value)
-        expression = make_shared[CScalarExpression](scalar)
+
+        if isinstance(value, bool):
+            scalar = MakeScalar(<c_bool>value)
+        elif isinstance(value, float):
+            scalar = MakeScalar(<double>value)
+        elif isinstance(value, int):
+            scalar = MakeScalar(<int64_t>value)
+        else:
+            raise TypeError('Not yet supported scalar value: {}'.format(value))
+
+        expression.reset(new CScalarExpression(scalar))
         self.init(<shared_ptr[CExpression]> expression)
 
     cdef void init(self, const shared_ptr[CExpression]& sp):
         Expression.init(self, sp)
         self.scalar = <CScalarExpression*> sp.get()
 
-    def value(self):
-        return pyarrow_wrap_scalar(self.scalar.value())
+    # def value(self):
+    #     return pyarrow_wrap_scalar(self.scalar.value())
 
 
 cdef class FieldExpression(Expression):
@@ -851,10 +861,32 @@ cdef class FieldExpression(Expression):
         return frombytes(self.scalar.name())
 
 
+cpdef enum CompareOperator:
+    Equal = <int8_t> CCompareOperator_EQUAL
+    NotEqual = <int8_t> CCompareOperator_NOT_EQUAL
+    Greater = <int8_t> CCompareOperator_GREATER
+    GreaterEqual = <int8_t> CCompareOperator_GREATER_EQUAL
+    Less = <int8_t> CCompareOperator_LESS
+    LessEqual = <int8_t> CCompareOperator_LESS_EQUAL
+
+
 cdef class ComparisonExpression(BinaryExpression):
 
     cdef CComparisonExpression* comparison
 
+    def __init__(self, CompareOperator op,
+                 Expression left_operand not None,
+                 Expression right_operand not None):
+        cdef shared_ptr[CComparisonExpression] expression
+        expression.reset(
+            new CComparisonExpression(
+                <CCompareOperator>op,
+                left_operand.unwrap(),
+                right_operand.unwrap()
+            )
+        )
+        self.init(<shared_ptr[CExpression]> expression)
+
     cdef void init(self, const shared_ptr[CExpression]& sp):
         BinaryExpression.init(self, sp)
         self.comparison = <CComparisonExpression*> sp.get()
@@ -863,13 +895,39 @@ cdef class ComparisonExpression(BinaryExpression):
     #     return ...
 
 
-cdef class AndExpression(ComparisonExpression):
-    pass
+cdef class NotExpression(UnaryExpression):
+
+    def __init__(self, Expression operand not None):
+        cdef shared_ptr[CNotExpression] expression
+        expression = MakeNotExpression(operand.unwrap())
+        self.init(<shared_ptr[CExpression]> expression)
 
 
-cdef class OrExpression(ComparisonExpression):
-    pass
+cdef class AndExpression(BinaryExpression):
+
+    def __init__(self, Expression left_operand not None,
+                 Expression right_operand not None,
+                 *additional_operands):
+        cdef:
+            Expression operand
+            vector[shared_ptr[CExpression]] exprs
+        exprs.push_back(left_operand.unwrap())
+        exprs.push_back(right_operand.unwrap())
+        for operand in additional_operands:
+            exprs.push_back(operand.unwrap())
+        self.init(MakeAndExpression(exprs))
 
 
-cdef class NotExpression(ComparisonExpression):
-    pass
+cdef class OrExpression(BinaryExpression):
+
+    def __init__(self, Expression left_operand not None,
+                 Expression right_operand not None,
+                 *additional_operands):
+        cdef:
+            Expression operand
+            vector[shared_ptr[CExpression]] exprs
+        exprs.push_back(left_operand.unwrap())
+        exprs.push_back(right_operand.unwrap())
+        for operand in additional_operands:
+            exprs.push_back(operand.unwrap())
+        self.init(MakeOrExpression(exprs))

python/pyarrow/dataset.py

@@ -45,5 +45,7 @@ from pyarrow._dataset import (
     FieldExpression,
     AndExpression,
     OrExpression,
-    NotExpression
+    NotExpression,
+    CompareOperator,
+    ComparisonExpression,
 )

python/pyarrow/includes/libarrow.pxd

@@ -791,6 +791,8 @@ cdef extern from "arrow/api.h" namespace "arrow" nogil:
     cdef cppclass CDoubleScalar" arrow::DoubleScalar"(CScalar):
         double value
 
+    shared_ptr[CScalar] MakeScalar[Value](Value value)
+
     CStatus ConcatenateTables(const vector[shared_ptr[CTable]]& tables,
                               shared_ptr[CTable]* result)
 
@@ -1420,6 +1422,14 @@ cdef extern from "arrow/compute/api.h" namespace "arrow::compute" nogil:
         CFunctionContext* context, const CDatum& values,
         const CDatum& filter, CDatum* out)
 
+    enum CCompareOperator" arrow::compute::CompareOperator":
+        CCompareOperator_EQUAL" arrow::compute::CompareOperator::EQUAL"
+        CCompareOperator_NOT_EQUAL" arrow::compute::CompareOperator::NOT_EQUAL"
+        CCompareOperator_GREATER" arrow::compute::CompareOperator::GREATER"
+        CCompareOperator_GREATER_EQUAL" arrow::compute::CompareOperator::GREATER_EQUAL"
+        CCompareOperator_LESS" arrow::compute::CompareOperator::LESS"
+        CCompareOperator_LESS_EQUAL" arrow::compute::CompareOperator::LESS_EQUAL"
+
 
 cdef extern from "arrow/python/api.h" namespace "arrow::py":
     # Requires GIL

python/pyarrow/includes/libarrow_dataset.pxd

@@ -74,6 +74,7 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         #                         compute::CastOptions options = compute::CastOptions()) const;
 
     ctypedef shared_ptr[CExpression] CExpressionPtr "arrow::dataset::ExpressionPtr"
+    ctypedef vector[CExpressionPtr] CExpressionVector "arrow::dataset::ExpressionVector"
 
     cdef cppclass CUnaryExpression "arrow::dataset::UnaryExpression"(CExpression):
         const CExpressionPtr& operand() const
@@ -90,22 +91,25 @@ cdef extern from "arrow/dataset/api.h" namespace "arrow::dataset" nogil:
         c_string name() const
 
     cdef cppclass CComparisonExpression "arrow::dataset::ComparisonExpression"(CBinaryExpression):
-        # ComparisonExpression(compute::CompareOperator op, ExpressionPtr left_operand,
-        #                      ExpressionPtr right_operand)
-        # compute::CompareOperator op() const { return op_; }
+        CComparisonExpression(CCompareOperator op, CExpressionPtr left_operand,
+                              CExpressionPtr right_operand)
+        CCompareOperator op() const
+
+    cdef cppclass CAndExpression "arrow::dataset::AndExpression"(CBinaryExpression):
         pass
 
-    cdef cppclass CAndExpression "arrow::dataset::AndExpression"(CComparisonExpression):
+    cdef cppclass COrExpression "arrow::dataset::OrExpression"(CBinaryExpression):
         pass
 
-    cdef cppclass COrExpression "arrow::dataset::OrExpression"(CComparisonExpression):
+    cdef cppclass CNotExpression "arrow::dataset::NotExpression"(CBinaryExpression):
         pass
 
-    cdef cppclass CNotExpression "arrow::dataset::NotExpression"(CComparisonExpression):
+    cdef cppclass CIsValidExpression "arrow::dataset::IsValidExpression"(CBinaryExpression):
         pass
 
-    cdef cppclass CIsValidExpression "arrow::dataset::IsValidExpression"(CComparisonExpression):
-        pass
+    cdef shared_ptr[CNotExpression] MakeNotExpression "arrow::dataset::not_"(CExpressionPtr operand)
+    cdef CExpressionPtr MakeAndExpression "arrow::dataset::and_"(const CExpressionVector& subexpressions)
+    cdef CExpressionPtr MakeOrExpression "arrow::dataset::or_"(const CExpressionVector& subexpressions)
 
     # InExpression
     # CastExpression

python/pyarrow/tests/test_dataset.py

@@ -50,7 +50,9 @@ try:
         FieldExpression,
         AndExpression,
         OrExpression,
-        NotExpression
+        NotExpression,
+        CompareOperator,
+        ComparisonExpression,
     )
 except ImportError as e:
     ds = None
@@ -235,6 +237,14 @@ def test_file_system_data_source():
 
 def test_expression():
     a = ScalarExpression(1)
+    b = ScalarExpression(1.1)
+    c = ScalarExpression(True)
+    ComparisonExpression(CompareOperator.Equal, a, b)
+    AndExpression(a, b)
+    AndExpression(a, b, c)
+    OrExpression(a, b)
+    OrExpression(a, b, c)
+    NotExpression(OrExpression(a, b, c))
 
 
 def test_file_system_discovery(fs):