lazy_coalesce_fallback_case_statement

Author: coderfender

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

@@ -191,7 +191,8 @@ object QueryPlanSerde extends Logging with CometExprShim {
     classOf[Log2] -> CometLog2,
     classOf[Pow] -> CometScalarFunction[Pow]("pow"),
     classOf[If] -> CometIf,
-    classOf[CaseWhen] -> CometCaseWhen)
+    classOf[CaseWhen] -> CometCaseWhen,
+    classOf[Coalesce] -> CometCoalesce)
 
   /**
    * Mapping of Spark aggregate expression class to Comet expression handler.
@@ -1078,10 +1079,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/main/scala/org/apache/comet/serde/conditional.scala

@@ -21,7 +21,7 @@ package org.apache.comet.serde
 
 import scala.collection.JavaConverters._
 
-import org.apache.spark.sql.catalyst.expressions.{Attribute, CaseWhen, Expression, If}
+import org.apache.spark.sql.catalyst.expressions.{Attribute, CaseWhen, Coalesce, Expression, If, IsNotNull}
 
 import org.apache.comet.CometSparkSessionExtensions.withInfo
 import org.apache.comet.serde.QueryPlanSerde.exprToProtoInternal
@@ -91,3 +91,48 @@ object CometCaseWhen extends CometExpressionSerde[CaseWhen] {
     }
   }
 }
+
+object CometCoalesce extends CometExpressionSerde[Coalesce] {
+  override def convert(
+      expr: Coalesce,
+      inputs: Seq[Attribute],
+      binding: Boolean): Option[ExprOuterClass.Expr] = {
+    val branches = expr.children.dropRight(1).map { child =>
+      (IsNotNull(child), child)
+    }
+    val elseValue = Some(expr.children.last)
+    var allBranches: Seq[Expression] = Seq()
+    val whenSeq = branches.map(elements => {
+      allBranches = allBranches :+ elements._1
+      exprToProtoInternal(elements._1, inputs, binding)
+    })
+    val thenSeq = branches.map(elements => {
+      allBranches = allBranches :+ elements._2
+      exprToProtoInternal(elements._2, inputs, binding)
+    })
+    assert(whenSeq.length == thenSeq.length)
+    if (whenSeq.forall(_.isDefined) && thenSeq.forall(_.isDefined)) {
+      val builder = ExprOuterClass.CaseWhen.newBuilder()
+      builder.addAllWhen(whenSeq.map(_.get).asJava)
+      builder.addAllThen(thenSeq.map(_.get).asJava)
+      if (elseValue.isDefined) {
+        val elseValueExpr =
+          exprToProtoInternal(elseValue.get, inputs, binding)
+        if (elseValueExpr.isDefined) {
+          builder.setElseExpr(elseValueExpr.get)
+        } else {
+          withInfo(expr, elseValue.get)
+          return None
+        }
+      }
+      Some(
+        ExprOuterClass.Expr
+          .newBuilder()
+          .setCaseWhen(builder)
+          .build())
+    } else {
+      withInfo(expr, allBranches: _*)
+      None
+    }
+  }
+}

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

@@ -395,13 +395,15 @@ class CometExpressionSuite extends CometTestBase with AdaptiveSparkPlanHelper {
   }
 
   test("test coalesce lazy eval") {
-    withSQLConf(SQLConf.ANSI_ENABLED.key -> "true") {
-      val data = Seq((100, 0))
+    withSQLConf(
+      SQLConf.ANSI_ENABLED.key -> "true",
+      CometConf.COMET_CAST_ALLOW_INCOMPATIBLE.key -> "true") {
+      val data = Seq((9999999999999L, 0))
       withParquetTable(data, "t1") {
         val res = spark.sql("""
-            |SELECT coalesce(_1 , 1/0) from t1;
+            |SELECT coalesce(_1, CAST(_1 AS TINYINT)) from t1;
             |  """.stripMargin)
-        checkSparkAnswer(res)
+        checkSparkAnswerAndOperator(res)
       }
     }
   }