[GLUTEN-4039] [VL] Add ntile window function support in Gluten (#4776)

* Add ntile window function support in Gluten * Update document for ntile
apache · Feb 27, 2024 · 420ffd1 · 420ffd1
1 parent 5185666
commit 420ffd1
Show file tree

Hide file tree

Showing 4 changed files with 24 additions and 3 deletions.
diff --git a/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/VeloxBackend.scala b/backends-velox/src/main/scala/io/glutenproject/backendsapi/velox/VeloxBackend.scala
@@ -26,7 +26,7 @@ import io.glutenproject.substrait.rel.LocalFilesNode.ReadFileFormat
 import io.glutenproject.substrait.rel.LocalFilesNode.ReadFileFormat.{DwrfReadFormat, OrcReadFormat, ParquetReadFormat}
 
 import org.apache.spark.sql.catalyst.catalog.BucketSpec
-import org.apache.spark.sql.catalyst.expressions.{Alias, CumeDist, DenseRank, Descending, Expression, Literal, NamedExpression, NthValue, PercentRank, Rand, RangeFrame, Rank, RowNumber, SortOrder, SpecialFrameBoundary, SpecifiedWindowFrame}
+import org.apache.spark.sql.catalyst.expressions.{Alias, CumeDist, DenseRank, Descending, Expression, Literal, NamedExpression, NthValue, NTile, PercentRank, Rand, RangeFrame, Rank, RowNumber, SortOrder, SpecialFrameBoundary, SpecifiedWindowFrame}
 import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, Count, Sum}
 import org.apache.spark.sql.catalyst.plans.JoinType
 import org.apache.spark.sql.catalyst.util.CharVarcharUtils
@@ -286,7 +286,7 @@ object BackendSettings extends BackendSettingsApi {
           }
           windowExpression.windowFunction match {
             case _: RowNumber | _: AggregateExpression | _: Rank | _: CumeDist | _: DenseRank |
-                _: PercentRank | _: NthValue =>
+                _: PercentRank | _: NthValue | _: NTile =>
             case _ =>
               allSupported = false
           }

diff --git a/backends-velox/src/test/scala/io/glutenproject/execution/TestOperator.scala b/backends-velox/src/test/scala/io/glutenproject/execution/TestOperator.scala
@@ -212,6 +212,12 @@ class TestOperator extends VeloxWholeStageTransformerSuite with AdaptiveSparkPla
     Seq("sort", "streaming").foreach {
       windowType =>
         withSQLConf("spark.gluten.sql.columnar.backend.velox.window.type" -> windowType) {
+          runQueryAndCompare(
+            "select ntile(4) over" +
+              " (partition by l_suppkey order by l_orderkey) from lineitem ") {
+            assertWindowOffloaded
+          }
+
           runQueryAndCompare(
             "select row_number() over" +
               " (partition by l_suppkey order by l_orderkey) from lineitem ") {

diff --git a/docs/velox-backend-support-progress.md b/docs/velox-backend-support-progress.md
@@ -387,7 +387,7 @@ Gluten supports 199 functions. (Draw to right to see all data types)
 | lag                           |                        |                       |        |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
 | lead                          |                        |                       |        |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
 | nth_value                     | nth_value              | nth_value             | PS     |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
-| ntile                         |                        |                       | S      |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
+| ntile                         |   ntile                     |     ntile                  | S      |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
 | percent_rank                  | percent_rank           |                       | S      |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
 | rank                          | rank                   |                       | S      |                        |         |      |       |     |      |       |        |      |           |        |         |      |        |          |       |      |        |      |
 | row_number                    | row_number             |                       | S      |                        |         |      | S     | S   | S    |       |        |      |           |        |         |      |        |          |       |      |        |      |

diff --git a/gluten-core/src/main/scala/io/glutenproject/backendsapi/SparkPlanExecApi.scala b/gluten-core/src/main/scala/io/glutenproject/backendsapi/SparkPlanExecApi.scala
@@ -551,6 +551,21 @@ trait SparkPlanExecApi {
               frame.frameType.sql
             )
             windowExpressionNodes.add(windowFunctionNode)
+          case wf @ NTile(buckets: Expression) =>
+            val frame = wExpression.windowSpec.frameSpecification.asInstanceOf[SpecifiedWindowFrame]
+            val childrenNodeList = new JArrayList[ExpressionNode]()
+            val literal = buckets.asInstanceOf[Literal]
+            childrenNodeList.add(LiteralTransformer(literal).doTransform(args))
+            val windowFunctionNode = ExpressionBuilder.makeWindowFunction(
+              WindowFunctionsBuilder.create(args, wf).toInt,
+              childrenNodeList,
+              columnName,
+              ConverterUtils.getTypeNode(wf.dataType, wf.nullable),
+              frame.upper.sql,
+              frame.lower.sql,
+              frame.frameType.sql
+            )
+            windowExpressionNodes.add(windowFunctionNode)
           case _ =>
             throw new UnsupportedOperationException(
               "unsupported window function type: " +