Skip to content

Commit

Permalink
[GLUTEN-4039] [VL] Add ntile window function support in Gluten (#4776)
Browse files Browse the repository at this point in the history
* Add ntile window function support in Gluten

* Update document for ntile
  • Loading branch information
JkSelf committed Feb 27, 2024
1 parent 5185666 commit 420ffd1
Show file tree
Hide file tree
Showing 4 changed files with 24 additions and 3 deletions.
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
}
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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 ") {
Expand Down
2 changes: 1 addition & 1 deletion docs/velox-backend-support-progress.md
Original file line number Diff line number Diff line change
Expand Up @@ -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 | | | | | | | | | | | | | |
Expand Down
Original file line number Diff line number Diff line change
Expand Up @@ -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: " +
Expand Down

0 comments on commit 420ffd1

Please sign in to comment.