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

gaoyangxiaozhu · 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