Add ANSI support for Subtract #535

core/src/execution/datafusion/expressions/binary.rs

@@ -0,0 +1,114 @@
+use std::any::Any;
+use std::fmt::{Display, Formatter};
+use std::hash::Hasher;
+use std::sync::Arc;
+
+use arrow_array::{BooleanArray, RecordBatch};
+use arrow_schema::{DataType, Schema};
+use datafusion_common::Result;
+use datafusion_expr::{ColumnarValue, Operator};
+use datafusion_expr::interval_arithmetic::Interval;
+use datafusion_expr::sort_properties::ExprProperties;
+use datafusion_physical_expr::expressions::BinaryExpr;
+use datafusion_physical_expr_common::physical_expr::{down_cast_any_ref, PhysicalExpr};
+
+use crate::execution::datafusion::expressions::EvalMode;
+
+#[derive(Debug, Hash, Clone)]
+pub struct CometBinaryExpr {
+    left: Arc<dyn PhysicalExpr>,
+    op: Operator,
+    right: Arc<dyn PhysicalExpr>,
+    eval_mode: EvalMode,
+    inner: Arc<BinaryExpr>,
+}
+
+impl CometBinaryExpr {
+    pub fn new(
+        left: Arc<dyn PhysicalExpr>,
+        op: Operator,
+        right: Arc<dyn PhysicalExpr>,
+        eval_mode: EvalMode,
+    ) -> Self {
+        Self {
+            left: Arc::clone(&left),
+            op,
+            right: Arc::clone(&right),
+            eval_mode,
+            inner: Arc::new(BinaryExpr::new(left, op, right)),
+        }
+    }
+}
+
+impl Display for CometBinaryExpr {
+    fn fmt(&self, f: &mut Formatter<'_>) -> std::fmt::Result {
+        self.inner.fmt(f)
+    }
+}
+
+impl PhysicalExpr for CometBinaryExpr {
+    fn as_any(&self) -> &dyn Any {
+        self.inner.as_any()
+    }
+
+    fn data_type(&self, input_schema: &Schema) -> Result<DataType> {
+        self.inner.data_type(input_schema)
+    }
+
+    fn nullable(&self, input_schema: &Schema) -> Result<bool> {
+        self.inner.nullable(input_schema)
+    }
+
+    fn evaluate(&self, batch: &RecordBatch) -> Result<ColumnarValue> {
+        // TODO: Do some work here
+        self.inner.evaluate(batch)
+    }
+
+    fn evaluate_selection(
+        &self,
+        batch: &RecordBatch,
+        selection: &BooleanArray,
+    ) -> Result<ColumnarValue> {
+        self.inner.evaluate_selection(batch, selection)
+    }
+
+    fn children(&self) -> Vec<&Arc<dyn PhysicalExpr>> {
+        self.inner.children()
+    }
+
+    fn with_new_children(
+        self: Arc<Self>,
+        children: Vec<Arc<dyn PhysicalExpr>>,
+    ) -> Result<Arc<dyn PhysicalExpr>> {
+        Arc::clone(&self.inner).with_new_children(children)
+    }
+
+    fn evaluate_bounds(&self, children: &[&Interval]) -> Result<Interval> {
+        self.inner.evaluate_bounds(children)
+    }
+
+    fn propagate_constraints(
+        &self,
+        interval: &Interval,
+        children: &[&Interval],
+    ) -> Result<Option<Vec<Interval>>> {
+        self.inner.propagate_constraints(interval, children)
+    }
+
+    fn dyn_hash(&self, state: &mut dyn Hasher) {
+        self.inner.dyn_hash(state)
+    }
+
+    fn get_properties(&self, children: &[ExprProperties]) -> Result<ExprProperties> {
+        self.inner.get_properties(children)
+    }
+}
+
+impl PartialEq<dyn Any> for CometBinaryExpr {
+    fn eq(&self, other: &dyn Any) -> bool {
+        down_cast_any_ref(other)
+            .downcast_ref::<Self>()
+            .map(|x| self.left.eq(&x.left) && self.op == x.op && self.right.eq(&x.right))
+            .unwrap_or(false)
+    }
+}

core/src/execution/datafusion/expressions/mod.rs

@@ -30,6 +30,7 @@ use crate::{errors::CometError, execution::spark_expression};
 pub mod abs;
 pub mod avg;
 pub mod avg_decimal;
+pub mod binary;
 pub mod bloom_filter_might_contain;
 pub mod correlation;
 pub mod covariance;

core/src/execution/datafusion/planner.rs

@@ -56,6 +56,7 @@ use itertools::Itertools;
 use jni::objects::GlobalRef;
 use num::{BigInt, ToPrimitive};
 
+use crate::execution::datafusion::expressions::binary::CometBinaryExpr;
 use crate::{
     errors::ExpressionError,
     execution::{
@@ -166,34 +167,39 @@ impl PhysicalPlanner {
                 expr.return_type.as_ref(),
                 DataFusionOperator::Plus,
                 input_schema,
+                EvalMode::Legacy,
             ),
             ExprStruct::Subtract(expr) => self.create_binary_expr(
                 expr.left.as_ref().unwrap(),
                 expr.right.as_ref().unwrap(),
                 expr.return_type.as_ref(),
                 DataFusionOperator::Minus,
                 input_schema,
+                EvalMode::try_from(expr.eval_mode)?,
             ),
             ExprStruct::Multiply(expr) => self.create_binary_expr(
                 expr.left.as_ref().unwrap(),
                 expr.right.as_ref().unwrap(),
                 expr.return_type.as_ref(),
                 DataFusionOperator::Multiply,
                 input_schema,
+                EvalMode::Legacy,
             ),
             ExprStruct::Divide(expr) => self.create_binary_expr(
                 expr.left.as_ref().unwrap(),
                 expr.right.as_ref().unwrap(),
                 expr.return_type.as_ref(),
                 DataFusionOperator::Divide,
                 input_schema,
+                EvalMode::Legacy,
             ),
             ExprStruct::Remainder(expr) => self.create_binary_expr(
                 expr.left.as_ref().unwrap(),
                 expr.right.as_ref().unwrap(),
                 expr.return_type.as_ref(),
                 DataFusionOperator::Modulo,
                 input_schema,
+                EvalMode::Legacy,
             ),
             ExprStruct::Eq(expr) => {
                 let left = self.create_expr(expr.left.as_ref().unwrap(), input_schema.clone())?;
@@ -627,6 +633,7 @@ impl PhysicalPlanner {
         return_type: Option<&spark_expression::DataType>,
         op: DataFusionOperator,
         input_schema: SchemaRef,
+        eval_mode: EvalMode,
     ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
         let left = self.create_expr(left, input_schema.clone())?;
         let right = self.create_expr(right, input_schema.clone())?;
@@ -681,7 +688,7 @@ impl PhysicalPlanner {
                     data_type,
                 )))
             }
-            _ => Ok(Arc::new(BinaryExpr::new(left, op, right))),
+            _ => Ok(Arc::new(CometBinaryExpr::new(left, op, right, eval_mode))),
         }
     }
 

core/src/execution/proto/expr.proto

@@ -219,6 +219,7 @@ message Subtract {
   Expr right = 2;
   bool fail_on_error = 3;
   DataType return_type = 4;
+  EvalMode eval_mode = 5;
 }
 
 message Multiply {

spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala

@@ -736,9 +736,13 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
 
           if (leftExpr.isDefined && rightExpr.isDefined) {
             val builder = ExprOuterClass.Subtract.newBuilder()
+            val failOnErr = getFailOnError(sub)
+            val evalModeStr =
+              if (failOnErr) ExprOuterClass.EvalMode.ANSI else ExprOuterClass.EvalMode.LEGACY
             builder.setLeft(leftExpr.get)
             builder.setRight(rightExpr.get)
-            builder.setFailOnError(getFailOnError(sub))
+            builder.setFailOnError(failOnErr)
+            builder.setEvalMode(evalModeStr)
             serializeDataType(sub.dataType).foreach { t =>
               builder.setReturnType(t)
             }