fix: Refactor arithmetic serde and fix correctness issues with EvalMode::TRY (#2018)

Author: andygrove

native/core/src/execution/planner.rs

@@ -229,45 +229,78 @@ impl PhysicalPlanner {
         input_schema: SchemaRef,
     ) -> Result<Arc<dyn PhysicalExpr>, ExecutionError> {
         match spark_expr.expr_struct.as_ref().unwrap() {
-            ExprStruct::Add(expr) => self.create_binary_expr(
-                expr.left.as_ref().unwrap(),
-                expr.right.as_ref().unwrap(),
-                expr.return_type.as_ref(),
-                DataFusionOperator::Plus,
-                input_schema,
-            ),
-            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,
-            ),
-            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,
-            ),
-            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,
-            ),
-            ExprStruct::IntegralDivide(expr) => self.create_binary_expr_with_options(
-                expr.left.as_ref().unwrap(),
-                expr.right.as_ref().unwrap(),
-                expr.return_type.as_ref(),
-                DataFusionOperator::Divide,
-                input_schema,
-                BinaryExprOptions {
-                    is_integral_div: true,
-                },
-            ),
+            ExprStruct::Add(expr) => {
+                // TODO respect eval mode
+                // https://github.com/apache/datafusion-comet/issues/2021
+                // https://github.com/apache/datafusion-comet/issues/536
+                let _eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
+                self.create_binary_expr(
+                    expr.left.as_ref().unwrap(),
+                    expr.right.as_ref().unwrap(),
+                    expr.return_type.as_ref(),
+                    DataFusionOperator::Plus,
+                    input_schema,
+                )
+            }
+            ExprStruct::Subtract(expr) => {
+                // TODO respect eval mode
+                // https://github.com/apache/datafusion-comet/issues/2021
+                // https://github.com/apache/datafusion-comet/issues/535
+                let _eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
+                self.create_binary_expr(
+                    expr.left.as_ref().unwrap(),
+                    expr.right.as_ref().unwrap(),
+                    expr.return_type.as_ref(),
+                    DataFusionOperator::Minus,
+                    input_schema,
+                )
+            }
+            ExprStruct::Multiply(expr) => {
+                // TODO respect eval mode
+                // https://github.com/apache/datafusion-comet/issues/2021
+                // https://github.com/apache/datafusion-comet/issues/534
+                let _eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
+                self.create_binary_expr(
+                    expr.left.as_ref().unwrap(),
+                    expr.right.as_ref().unwrap(),
+                    expr.return_type.as_ref(),
+                    DataFusionOperator::Multiply,
+                    input_schema,
+                )
+            }
+            ExprStruct::Divide(expr) => {
+                // TODO respect eval mode
+                // https://github.com/apache/datafusion-comet/issues/2021
+                // https://github.com/apache/datafusion-comet/issues/533
+                let _eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
+                self.create_binary_expr(
+                    expr.left.as_ref().unwrap(),
+                    expr.right.as_ref().unwrap(),
+                    expr.return_type.as_ref(),
+                    DataFusionOperator::Divide,
+                    input_schema,
+                )
+            }
+            ExprStruct::IntegralDivide(expr) => {
+                // TODO respect eval mode
+                // https://github.com/apache/datafusion-comet/issues/2021
+                // https://github.com/apache/datafusion-comet/issues/533
+                let _eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
+                self.create_binary_expr_with_options(
+                    expr.left.as_ref().unwrap(),
+                    expr.right.as_ref().unwrap(),
+                    expr.return_type.as_ref(),
+                    DataFusionOperator::Divide,
+                    input_schema,
+                    BinaryExprOptions {
+                        is_integral_div: true,
+                    },
+                )
+            }
             ExprStruct::Remainder(expr) => {
+                let eval_mode = from_protobuf_eval_mode(expr.eval_mode)?;
+                // TODO add support for EvalMode::TRY
+                // https://github.com/apache/datafusion-comet/issues/2021
                 let left =
                     self.create_expr(expr.left.as_ref().unwrap(), Arc::clone(&input_schema))?;
                 let right =
@@ -278,7 +311,7 @@ impl PhysicalPlanner {
                     right,
                     expr.return_type.as_ref().map(to_arrow_datatype).unwrap(),
                     input_schema,
-                    expr.fail_on_error,
+                    eval_mode == EvalMode::Ansi,
                     &self.session_ctx.state(),
                 );
                 result.map_err(|e| GeneralError(e.to_string()))

native/proto/src/proto/expr.proto

@@ -220,19 +220,19 @@ message Literal {
    bool is_null = 12;
 }
 
-message MathExpr {
-  Expr left = 1;
-  Expr right = 2;
-  bool fail_on_error = 3;
-  DataType return_type = 4;
-}
-
 enum EvalMode {
   LEGACY = 0;
   TRY = 1;
   ANSI = 2;
 }
 
+message MathExpr {
+  Expr left = 1;
+  Expr right = 2;
+  DataType return_type = 4;
+  EvalMode eval_mode = 5;
+}
+
 message Cast {
   Expr child = 1;
   DataType datatype = 2;

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

@@ -23,7 +23,6 @@ import java.util.Locale
 
 import scala.collection.JavaConverters._
 import scala.collection.mutable.ListBuffer
-import scala.math.min
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.expressions._
@@ -67,6 +66,12 @@ object QueryPlanSerde extends Logging with CometExprShim {
    * Mapping of Spark expression class to Comet expression handler.
    */
   private val exprSerdeMap: Map[Class[_], CometExpressionSerde] = Map(
+    classOf[Add] -> CometAdd,
+    classOf[Subtract] -> CometSubtract,
+    classOf[Multiply] -> CometMultiply,
+    classOf[Divide] -> CometDivide,
+    classOf[IntegralDivide] -> CometIntegralDivide,
+    classOf[Remainder] -> CometRemainder,
     classOf[ArrayAppend] -> CometArrayAppend,
     classOf[ArrayContains] -> CometArrayContains,
     classOf[ArrayDistinct] -> CometArrayDistinct,
@@ -630,141 +635,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
       case c @ Cast(child, dt, timeZoneId, _) =>
         handleCast(expr, child, inputs, binding, dt, timeZoneId, evalMode(c))
 
-      case add @ Add(left, right, _) if supportedDataType(left.dataType) =>
-        createMathExpression(
-          expr,
-          left,
-          right,
-          inputs,
-          binding,
-          add.dataType,
-          add.evalMode == EvalMode.ANSI,
-          (builder, mathExpr) => builder.setAdd(mathExpr))
-
-      case add @ Add(left, _, _) if !supportedDataType(left.dataType) =>
-        withInfo(add, s"Unsupported datatype ${left.dataType}")
-        None
-
-      case sub @ Subtract(left, right, _) if supportedDataType(left.dataType) =>
-        createMathExpression(
-          expr,
-          left,
-          right,
-          inputs,
-          binding,
-          sub.dataType,
-          sub.evalMode == EvalMode.ANSI,
-          (builder, mathExpr) => builder.setSubtract(mathExpr))
-
-      case sub @ Subtract(left, _, _) if !supportedDataType(left.dataType) =>
-        withInfo(sub, s"Unsupported datatype ${left.dataType}")
-        None
-
-      case mul @ Multiply(left, right, _) if supportedDataType(left.dataType) =>
-        createMathExpression(
-          expr,
-          left,
-          right,
-          inputs,
-          binding,
-          mul.dataType,
-          mul.evalMode == EvalMode.ANSI,
-          (builder, mathExpr) => builder.setMultiply(mathExpr))
-
-      case mul @ Multiply(left, _, _) =>
-        if (!supportedDataType(left.dataType)) {
-          withInfo(mul, s"Unsupported datatype ${left.dataType}")
-        }
-        None
-
-      case div @ Divide(left, right, _) if supportedDataType(left.dataType) =>
-        // Datafusion now throws an exception for dividing by zero
-        // See https://github.com/apache/arrow-datafusion/pull/6792
-        // For now, use NullIf to swap zeros with nulls.
-        val rightExpr = nullIfWhenPrimitive(right)
-
-        createMathExpression(
-          expr,
-          left,
-          rightExpr,
-          inputs,
-          binding,
-          div.dataType,
-          div.evalMode == EvalMode.ANSI,
-          (builder, mathExpr) => builder.setDivide(mathExpr))
-
-      case div @ Divide(left, _, _) =>
-        if (!supportedDataType(left.dataType)) {
-          withInfo(div, s"Unsupported datatype ${left.dataType}")
-        }
-        None
-
-      case div @ IntegralDivide(left, right, _) if supportedDataType(left.dataType) =>
-        val rightExpr = nullIfWhenPrimitive(right)
-
-        val dataType = (left.dataType, right.dataType) match {
-          case (l: DecimalType, r: DecimalType) =>
-            // copy from IntegralDivide.resultDecimalType
-            val intDig = l.precision - l.scale + r.scale
-            DecimalType(min(if (intDig == 0) 1 else intDig, DecimalType.MAX_PRECISION), 0)
-          case _ => left.dataType
-        }
-
-        val divideExpr = createMathExpression(
-          expr,
-          left,
-          rightExpr,
-          inputs,
-          binding,
-          dataType,
-          div.evalMode == EvalMode.ANSI,
-          (builder, mathExpr) => builder.setIntegralDivide(mathExpr))
-
-        if (divideExpr.isDefined) {
-          val childExpr = if (dataType.isInstanceOf[DecimalType]) {
-            // check overflow for decimal type
-            val builder = ExprOuterClass.CheckOverflow.newBuilder()
-            builder.setChild(divideExpr.get)
-            builder.setFailOnError(div.evalMode == EvalMode.ANSI)
-            builder.setDatatype(serializeDataType(dataType).get)
-            Some(
-              ExprOuterClass.Expr
-                .newBuilder()
-                .setCheckOverflow(builder)
-                .build())
-          } else {
-            divideExpr
-          }
-
-          // cast result to long
-          castToProto(expr, None, LongType, childExpr.get, CometEvalMode.LEGACY)
-        } else {
-          None
-        }
-
-      case div @ IntegralDivide(left, _, _) =>
-        if (!supportedDataType(left.dataType)) {
-          withInfo(div, s"Unsupported datatype ${left.dataType}")
-        }
-        None
-
-      case rem @ Remainder(left, right, _) if supportedDataType(left.dataType) =>
-        createMathExpression(
-          expr,
-          left,
-          right,
-          inputs,
-          binding,
-          rem.dataType,
-          rem.evalMode == EvalMode.ANSI,
-          (builder, mathExpr) => builder.setRemainder(mathExpr))
-
-      case rem @ Remainder(left, _, _) =>
-        if (!supportedDataType(left.dataType)) {
-          withInfo(rem, s"Unsupported datatype ${left.dataType}")
-        }
-        None
-
       case EqualTo(left, right) =>
         createBinaryExpr(
           expr,
@@ -1962,42 +1832,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
     }
   }
 
-  private def createMathExpression(
-      expr: Expression,
-      left: Expression,
-      right: Expression,
-      inputs: Seq[Attribute],
-      binding: Boolean,
-      dataType: DataType,
-      failOnError: Boolean,
-      f: (ExprOuterClass.Expr.Builder, ExprOuterClass.MathExpr) => ExprOuterClass.Expr.Builder)
-      : Option[ExprOuterClass.Expr] = {
-    val leftExpr = exprToProtoInternal(left, inputs, binding)
-    val rightExpr = exprToProtoInternal(right, inputs, binding)
-
-    if (leftExpr.isDefined && rightExpr.isDefined) {
-      // create the generic MathExpr message
-      val builder = ExprOuterClass.MathExpr.newBuilder()
-      builder.setLeft(leftExpr.get)
-      builder.setRight(rightExpr.get)
-      builder.setFailOnError(failOnError)
-      serializeDataType(dataType).foreach { t =>
-        builder.setReturnType(t)
-      }
-      val inner = builder.build()
-      // call the user-supplied function to wrap MathExpr in a top-level Expr
-      // such as Expr.Add or Expr.Divide
-      Some(
-        f(
-          ExprOuterClass.Expr
-            .newBuilder(),
-          inner).build())
-    } else {
-      withInfo(expr, left, right)
-      None
-    }
-  }
-
   def in(
       expr: Expression,
       value: Expression,
@@ -2053,25 +1887,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
     Some(ExprOuterClass.Expr.newBuilder().setScalarFunc(builder).build())
   }
 
-  private def isPrimitive(expression: Expression): Boolean = expression.dataType match {
-    case _: ByteType | _: ShortType | _: IntegerType | _: LongType | _: FloatType |
-        _: DoubleType | _: TimestampType | _: DateType | _: BooleanType | _: DecimalType =>
-      true
-    case _ => false
-  }
-
-  private def nullIfWhenPrimitive(expression: Expression): Expression =
-    if (isPrimitive(expression)) {
-      val zero = Literal.default(expression.dataType)
-      expression match {
-        case _: Literal if expression != zero => expression
-        case _ =>
-          If(EqualTo(expression, zero), Literal.create(null, expression.dataType), expression)
-      }
-    } else {
-      expression
-    }
-
   private def nullIfNegative(expression: Expression): Expression = {
     val zero = Literal.default(expression.dataType)
     If(LessThanOrEqual(expression, zero), Literal.create(null, expression.dataType), expression)