[Gluten-692] Support spark check_overflow function. facebookincubator…

…#231 (facebookincubator#691) add check_overflow transform support. test by CH[[231]]
rui-mo · Dec 13, 2022 · 31db423 · 31db423
1 parent 4d12a03
commit 31db423
Show file tree

Hide file tree

Showing 3 changed files with 62 additions and 3 deletions.
diff --git a/gluten-core/src/main/scala/io/glutenproject/expression/ExpressionMappings.scala b/gluten-core/src/main/scala/io/glutenproject/expression/ExpressionMappings.scala
@@ -161,6 +161,8 @@ object ExpressionMappings {
   final val IN_SET = "in_set"
   final val SCALAR_SUBQUERY = "scalar_subquery"
   final val EXPLODE = "explode"
+  final val CHECK_OVERFLOW = "check_overflow"
+  final val PROMOTE_PRECISION = "promote_precision"
 
   // Directly use child expression transformer
   final val KNOWN_FLOATING_POINT_NORMALIZED = "known_floating_point_normalized"
@@ -287,7 +289,9 @@ object ExpressionMappings {
     Sig[CaseWhen](CASE_WHEN),
     Sig[In](IN),
     Sig[InSet](IN_SET),
-    Sig[ScalarSubquery](SCALAR_SUBQUERY)
+    Sig[ScalarSubquery](SCALAR_SUBQUERY),
+    Sig[CheckOverflow](CHECK_OVERFLOW),
+    Sig[PromotePrecision](PROMOTE_PRECISION)
   ) ++ SparkShimLoader.getSparkShims.expressionMappings
 
   /**

diff --git a/gluten-core/src/main/scala/io/glutenproject/expression/UnaryExpressionTransformer.scala b/gluten-core/src/main/scala/io/glutenproject/expression/UnaryExpressionTransformer.scala
@@ -20,7 +20,7 @@ package io.glutenproject.expression
 import com.google.common.collect.Lists
 
 import io.glutenproject.expression.ConverterUtils.FunctionConfig
-import io.glutenproject.substrait.expression.{ExpressionBuilder, ExpressionNode}
+import io.glutenproject.substrait.expression.{BooleanLiteralNode, ExpressionBuilder, ExpressionNode}
 
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.catalyst.expressions._
@@ -81,6 +81,38 @@ class ExplodeTransformer(substraitExprName: String, child: ExpressionTransformer
   }
 }
 
+class PromotePrecisionTransformer(child: ExpressionTransformer, original: PromotePrecision)
+  extends ExpressionTransformer with Logging {
+
+  override def doTransform(args: java.lang.Object): ExpressionNode = {
+    child.doTransform(args)
+  }
+}
+
+class CheckOverflowTransformer(
+                                substraitExprName: String,
+                                child: ExpressionTransformer,
+                                original: CheckOverflow)
+  extends ExpressionTransformer
+    with Logging {
+
+  override def doTransform(args: java.lang.Object): ExpressionNode = {
+    val childNode = child.doTransform(args)
+    val functionMap = args.asInstanceOf[java.util.HashMap[String, java.lang.Long]]
+    val functionId = ExpressionBuilder.newScalarFunction(
+      functionMap,
+      ConverterUtils.makeFuncName(
+        substraitExprName,
+        Seq(original.dataType, BooleanType),
+        FunctionConfig.OPT))
+
+    val expressionNodes = Lists.newArrayList(childNode,
+      new BooleanLiteralNode(original.nullOnOverflow))
+    val typeNode = ConverterUtils.getTypeNode(original.dataType, original.nullable)
+    ExpressionBuilder.makeScalarFunction(functionId, expressionNodes, typeNode)
+  }
+}
+
 /**
  * Transformer for the normal unary expression
  */
@@ -109,6 +141,13 @@ object UnaryExpressionTransformer {
   def apply(substraitExprName: String,
             child: ExpressionTransformer,
             original: Expression): ExpressionTransformer = {
-    new UnaryExpressionTransformer(substraitExprName, child, original)
+    original match {
+      case c: CheckOverflow =>
+        new CheckOverflowTransformer(substraitExprName, child, c)
+      case p: PromotePrecision =>
+        new PromotePrecisionTransformer(child, p)
+      case _ =>
+        new UnaryExpressionTransformer(substraitExprName, child, original)
+    }
   }
 }
diff --git a/gluten-ut/src/test/scala/io/glutenproject/utils/clickhouse/ClickHouseTestSettings.scala b/gluten-ut/src/test/scala/io/glutenproject/utils/clickhouse/ClickHouseTestSettings.scala
@@ -19,6 +19,7 @@ package io.glutenproject.utils.clickhouse
 
 import io.glutenproject.utils.BackendTestSettings
 import org.apache.spark.sql._
+import org.apache.spark.sql.catalyst.expressions.GlutenArithmeticExpressionSuite
 
 object ClickHouseTestSettings extends BackendTestSettings {
 
@@ -103,4 +104,19 @@ object ClickHouseTestSettings extends BackendTestSettings {
     )
 
   enableSuite[GlutenComplexTypesSuite]
+
+  enableSuite[GlutenArithmeticExpressionSuite]
+    .exclude(
+      "- (UnaryMinus)",
+      "/ (Divide) basic",
+      "/ (Divide) for Long and Decimal type",
+      "% (Remainder)",
+      "SPARK-17617: % (Remainder) double % double on super big double",
+      "pmod",
+      "function least",
+      "function greatest",
+      "SPARK-28322: IntegralDivide supports decimal type",
+      "SPARK-33008: division by zero on divide-like operations returns incorrect result",
+      "SPARK-34920: error class"
+    )
 }