[VL] Add support for spark_partition_id function (#4969)

apache · Mar 19, 2024 · 91fb95d · 91fb95d
1 parent 81b023e
commit 91fb95d
Show file tree

Hide file tree

Showing 4 changed files with 11 additions and 1 deletion.
diff --git a/backends-clickhouse/src/main/scala/io/glutenproject/utils/CHExpressionUtil.scala b/backends-clickhouse/src/main/scala/io/glutenproject/utils/CHExpressionUtil.scala
@@ -176,6 +176,7 @@ object CHExpressionUtil {
     ARRAY_REPEAT -> DefaultValidator(),
     DATE_FROM_UNIX_DATE -> DefaultValidator(),
     MONOTONICALLY_INCREASING_ID -> DefaultValidator(),
+    SPARK_PARTITION_ID -> DefaultValidator(),
     SKEWNESS -> DefaultValidator()
   )
 }
diff --git a/backends-velox/src/test/scala/io/glutenproject/execution/VeloxFunctionsValidateSuite.scala b/backends-velox/src/test/scala/io/glutenproject/execution/VeloxFunctionsValidateSuite.scala
@@ -446,7 +446,14 @@ class VeloxFunctionsValidateSuite extends VeloxWholeStageTransformerSuite {
 
   test("Test monotonically_increasing_id function") {
     runQueryAndCompare("""SELECT monotonically_increasing_id(), l_orderkey
-                         | from lineitem limit 1""".stripMargin) {
+                         | from lineitem limit 100""".stripMargin) {
+      checkOperatorMatch[ProjectExecTransformer]
+    }
+  }
+
+  test("Test spark_partition_id function") {
+    runQueryAndCompare("""SELECT spark_partition_id(), l_orderkey
+                         | from lineitem limit 100""".stripMargin) {
       checkOperatorMatch[ProjectExecTransformer]
     }
   }

diff --git a/gluten-core/src/main/scala/io/glutenproject/expression/ExpressionMappings.scala b/gluten-core/src/main/scala/io/glutenproject/expression/ExpressionMappings.scala
@@ -243,6 +243,7 @@ object ExpressionMappings {
     Sig[MakeDecimal](MAKE_DECIMAL),
     Sig[PromotePrecision](PROMOTE_PRECISION),
     Sig[MonotonicallyIncreasingID](MONOTONICALLY_INCREASING_ID),
+    Sig[SparkPartitionID](SPARK_PARTITION_ID),
     // Decimal
     Sig[UnscaledValue](UNSCALED_VALUE)
   ) ++ SparkShimLoader.getSparkShims.expressionMappings

diff --git a/shims/common/src/main/scala/io/glutenproject/expression/ExpressionNames.scala b/shims/common/src/main/scala/io/glutenproject/expression/ExpressionNames.scala
@@ -262,6 +262,7 @@ object ExpressionNames {
   final val CHECK_OVERFLOW = "check_overflow"
   final val MAKE_DECIMAL = "make_decimal"
   final val PROMOTE_PRECISION = "promote_precision"
+  final val SPARK_PARTITION_ID = "spark_partition_id"
   final val MONOTONICALLY_INCREASING_ID = "monotonically_increasing_id"
 
   // Directly use child expression transformer