changed the logic, modified nullifwhenprimitive function

apache · Jun 22, 2024 · b658eac · b658eac
1 parent d82ffbf
commit b658eac
Show file tree

Hide file tree

Showing 5 changed files with 14 additions and 202 deletions.
diff --git a/core/src/execution/datafusion/expressions/mod.rs b/core/src/execution/datafusion/expressions/mod.rs
@@ -33,7 +33,6 @@ pub mod avg_decimal;
 pub mod bloom_filter_might_contain;
 pub mod correlation;
 pub mod covariance;
-pub mod modulo;
 pub mod negative;
 pub mod stats;
 pub mod stddev;

diff --git a/core/src/execution/datafusion/expressions/modulo.rs b/core/src/execution/datafusion/expressions/modulo.rs
diff --git a/core/src/execution/datafusion/planner.rs b/core/src/execution/datafusion/planner.rs
@@ -97,7 +97,7 @@ use crate::{
     },
 };
 
-use super::expressions::{abs::CometAbsFunc, modulo::ModuloExpr, EvalMode};
+use super::expressions::{abs::CometAbsFunc, EvalMode};
 
 // For clippy error on type_complexity.
 type ExecResult<T> = Result<T, ExecutionError>;
@@ -681,13 +681,7 @@ impl PhysicalPlanner {
                     data_type,
                 )))
             }
-            _ => {
-                // Improves compatibility with Spark
-                if op == DataFusionOperator::Modulo {
-                    return Ok(Arc::new(ModuloExpr::new(left, right)));
-                }
-                Ok(Arc::new(BinaryExpr::new(left, op, right)))
-            }
+            _ => Ok(Arc::new(BinaryExpr::new(left, op, right))),
         }
     }
 

diff --git a/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala b/spark/src/main/scala/org/apache/comet/serde/QueryPlanSerde.scala
@@ -2240,7 +2240,15 @@ object QueryPlanSerde extends Logging with ShimQueryPlanSerde with CometExprShim
 
   def nullIfWhenPrimitive(expression: Expression): Expression = if (isPrimitive(expression)) {
     val zero = Literal.default(expression.dataType)
-    If(EqualTo(expression, zero), Literal.create(null, expression.dataType), expression)
+    val negZero = Literal.create(-0.0, DoubleType)
+    if (expression.dataType == DoubleType) {
+      If(
+        Or(EqualTo(expression, zero), EqualTo(expression, negZero)),
+        Literal.create(null, expression.dataType),
+        expression)
+    } else {
+      If(EqualTo(expression, zero), Literal.create(null, expression.dataType), expression)
+    }
   } else {
     expression
   }

diff --git a/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala b/spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala
@@ -858,7 +858,9 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
 
   test("remainder") {
     withTempDir { dir =>
-      val df = Seq((21840, -0.0)).toDF("c90", "c1")
+      val df =
+        Seq((21840, -0.0), (21840, 5.0))
+          .toDF("c90", "c1")
       val path = new Path(dir.toURI.toString, "remainder_test.parquet").toString
       df.write.mode("overwrite").parquet(path)