Support KnownNullable in 332db

NVIDIA · May 17, 2023 · 8068205 · 8068205
1 parent abe8662
commit 8068205
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 2 deletions.
diff --git a/...spark/rapids/shims/GpuKnownNullable.scala → ...spark/rapids/shims/GpuKnownNullable.scala b/...spark/rapids/shims/GpuKnownNullable.scala → ...spark/rapids/shims/GpuKnownNullable.scala
@@ -15,6 +15,7 @@
  */
 
 /*** spark-rapids-shim-json-lines
+{"spark": "332db"}
 {"spark": "340"}
 spark-rapids-shim-json-lines ***/
 package com.nvidia.spark.rapids.shims

diff --git a/sql-plugin/src/main/spark332db/scala/com/nvidia/spark/rapids/shims/SparkShims.scala b/sql-plugin/src/main/spark332db/scala/com/nvidia/spark/rapids/shims/SparkShims.scala
@@ -55,10 +55,18 @@ object SparkShimImpl extends Spark321PlusDBShims {
   }
 
   override def getExprs: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = {
-    val elementAtExpr: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = Seq(
+    val shimExprs: Map[Class[_ <: Expression], ExprRule[_ <: Expression]] = Seq(
+      GpuOverrides.expr[KnownNullable](
+        "Tags the expression as being nullable",
+        ExprChecks.unaryProjectInputMatchesOutput(
+          TypeSig.all, TypeSig.all),
+        (a, conf, p, r) => new UnaryExprMeta[KnownNullable](a, conf, p, r) {
+          override def convertToGpu(child: Expression): GpuExpression = GpuKnownNullable(child)
+        }
+      ),
       GpuElementAtMeta.elementAtRule(true)
     ).map(r => (r.getClassFor.asSubclass(classOf[Expression]), r)).toMap
-    super.getExprs ++ DayTimeIntervalShims.exprs ++ RoundingShims.exprs ++ elementAtExpr
+    super.getExprs ++ shimExprs ++ DayTimeIntervalShims.exprs ++ RoundingShims.exprs
   }
 
   private val shimExecs: Map[Class[_ <: SparkPlan], ExecRule[_ <: SparkPlan]] = Seq(