From 91fb95d1fe2a15f32d1e5ec33b815423ec844d1e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?=E9=AB=98=E9=98=B3=E9=98=B3?= Date: Tue, 19 Mar 2024 08:49:41 +0800 Subject: [PATCH] [VL] Add support for spark_partition_id function (#4969) --- .../scala/io/glutenproject/utils/CHExpressionUtil.scala | 1 + .../execution/VeloxFunctionsValidateSuite.scala | 9 ++++++++- .../io/glutenproject/expression/ExpressionMappings.scala | 1 + .../io/glutenproject/expression/ExpressionNames.scala | 1 + 4 files changed, 11 insertions(+), 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 index cb607da4c4c9..86fd425f438d 100644 --- 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 index b558cc77bf21..187c93514e13 100644 --- 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 index 9f3c2ed81583..c0ff0e8707eb 100644 --- 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 index 7998d7642d3a..91c68e903bbd 100644 --- 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