check_upstream_json_enrichments

Author: coderfender

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

@@ -1248,13 +1248,36 @@ object QueryPlanSerde extends Logging with CometExprShim {
           None
         }
 
-      case CaseWhen(branches, elseValue) =>
+//      case a @ Coalesce(_) =>
+//        val branches = a.children.dropRight(1).map { child =>
+//          (IsNotNull(child), child)
+//        }
+//        val elseValue = Some(a.children.last)
+//
+//        exprToProtoInternal(CaseWhen(branches, elseValue), inputs, binding)
+      //
+      //        val exprChildren = a.children.map(exprToProtoInternal(_, inputs, binding))
+      //        scalarFunctionExprToProto("coalesce", exprChildren: _*)
+
+      case c @ (CaseWhen(_, _) | Coalesce(_)) =>
+        val (finalBranches, finalElse) = c match {
+          case CaseWhen(branches, elseValue) =>
+            (branches, elseValue)
+
+          case Coalesce(children) =>
+            val branches = children.dropRight(1).map { child =>
+              (IsNotNull(child), child)
+            }
+            val elseValue = Some(children.last)
+            (branches, elseValue)
+        }
+
         var allBranches: Seq[Expression] = Seq()
-        val whenSeq = branches.map(elements => {
+        val whenSeq = finalBranches.map(elements => {
           allBranches = allBranches :+ elements._1
           exprToProtoInternal(elements._1, inputs, binding)
         })
-        val thenSeq = branches.map(elements => {
+        val thenSeq = finalBranches.map(elements => {
           allBranches = allBranches :+ elements._2
           exprToProtoInternal(elements._2, inputs, binding)
         })
@@ -1263,13 +1286,13 @@ object QueryPlanSerde extends Logging with CometExprShim {
           val builder = ExprOuterClass.CaseWhen.newBuilder()
           builder.addAllWhen(whenSeq.map(_.get).asJava)
           builder.addAllThen(thenSeq.map(_.get).asJava)
-          if (elseValue.isDefined) {
+          if (finalElse.isDefined) {
             val elseValueExpr =
-              exprToProtoInternal(elseValue.get, inputs, binding)
+              exprToProtoInternal(finalElse.get, inputs, binding)
             if (elseValueExpr.isDefined) {
               builder.setElseExpr(elseValueExpr.get)
             } else {
-              withInfo(expr, elseValue.get)
+              withInfo(expr, finalElse.get)
               return None
             }
           }
@@ -1319,10 +1342,6 @@ object QueryPlanSerde extends Logging with CometExprShim {
           None
         }
 
-      case a @ Coalesce(_) =>
-        val exprChildren = a.children.map(exprToProtoInternal(_, inputs, binding))
-        scalarFunctionExprToProto("coalesce", exprChildren: _*)
-
       // With Spark 3.4, CharVarcharCodegenUtils.readSidePadding gets called to pad spaces for
       // char types.
       // See https://github.com/apache/spark/pull/38151

spark/src/test/scala/org/apache/comet/CometExpressionSuite.scala

@@ -44,6 +44,7 @@ import org.apache.spark.sql.types._
 import org.apache.comet.CometSparkSessionExtensions.isSpark40Plus
 
 class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
+
   import testImplicits._
 
   override protected def test(testName: String, testTags: Tag*)(testFun: => Any)(implicit
@@ -367,15 +368,15 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     withParquetTable(data, "tbl") {
       checkSparkAnswerAndOperator("SELECT try_divide(_1, _2) FROM tbl")
       checkSparkAnswerAndOperator("""
-            |SELECT
-            |  try_divide(10, 0),
-            |  try_divide(NULL, 5),
-            |  try_divide(5, NULL),
-            |  try_divide(-2147483648, -1),
-            |  try_divide(-9223372036854775808, -1),
-            |  try_divide(DECIMAL('9999999999999999999999999999'), 0.1)
-            |  from tbl
-            |""".stripMargin)
+          |SELECT
+          |  try_divide(10, 0),
+          |  try_divide(NULL, 5),
+          |  try_divide(5, NULL),
+          |  try_divide(-2147483648, -1),
+          |  try_divide(-9223372036854775808, -1),
+          |  try_divide(DECIMAL('9999999999999999999999999999'), 0.1)
+          |  from tbl
+          |""".stripMargin)
     }
   }
 
@@ -384,13 +385,28 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
     withParquetTable(data, "tbl") {
       checkSparkAnswerAndOperator("SELECT try_divide(_1, _2) FROM tbl")
       checkSparkAnswerAndOperator("""
-                                    |SELECT try_divide(-128, -1),
-                                    |try_divide(-32768, -1),
-                                    |try_divide(-2147483648, -1),
-                                    |try_divide(-9223372036854775808, -1),
-                                    |try_divide(CAST(99999 AS DECIMAL(5,0)), CAST(0.0001 AS DECIMAL(5,4)))
-                                    |from tbl
-                                    |""".stripMargin)
+          |SELECT try_divide(-128, -1),
+          |try_divide(-32768, -1),
+          |try_divide(-2147483648, -1),
+          |try_divide(-9223372036854775808, -1),
+          |try_divide(CAST(99999 AS DECIMAL(5,0)), CAST(0.0001 AS DECIMAL(5,4)))
+          |from tbl
+          |""".stripMargin)
+    }
+  }
+
+  test("test coalesce lazy eval") {
+    withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
+      val data = Seq((100, 0))
+      withParquetTable(data, "t1") {
+        val res = spark.sql("""
+             |SELECT coalesce(_1 , 1/0) from t1;
+             |  """.stripMargin)
+
+        res.explain(true)
+
+        checkSparkAnswer(res)
+      }
     }
   }