Copilot commented on code in PR #13050:
URL: https://github.com/apache/gluten/pull/13050#discussion_r4231586947
##########
backends-velox/src/test/scala/org/apache/gluten/functions/WindowFunctionsValidateSuite.scala:
##########
@@ -92,4 +105,118 @@ class WindowFunctionsValidateSuite extends
FunctionsValidateSuite {
}
}
}
+
+ test("window partition/order by TIMESTAMP_NTZ goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+
+ runQueryAndCompare(
+ "select k, ts, rank() over (partition by k order by ts) as rnk " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("lag(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, lag(ts, 1) over (partition by k order by v) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
Review Comment:
PR description says: “Not covered: window functions taking TIMESTAMP_NTZ
directly (lag, lead, first_value, etc.) still fall back.” However, this PR adds
tests asserting these now go native (and the validator whitelists them). Please
align the PR description with the implemented behavior (or, if the intent is
still to fallback for NTZ-typed window-function arguments, update the
validator/tests accordingly).
##########
backends-velox/src/test/scala/org/apache/gluten/functions/WindowFunctionsValidateSuite.scala:
##########
@@ -92,4 +105,118 @@ class WindowFunctionsValidateSuite extends
FunctionsValidateSuite {
}
}
}
+
+ test("window partition/order by TIMESTAMP_NTZ goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+
+ runQueryAndCompare(
+ "select k, ts, rank() over (partition by k order by ts) as rnk " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("lag(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, lag(ts, 1) over (partition by k order by v) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("lead(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, lead(ts, 1) over (partition by k order by v) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("nth_value(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, nth_value(ts, 1) over (partition by k order by v) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("first_value(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, first_value(ts) over (partition by k order by v) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("last_value(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, last_value(ts) over (partition by k order by v) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("count(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, count(ts) over (partition by k) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("min(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, min(ts) over (partition by k) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("max(TIMESTAMP_NTZ) window function goes native") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+ runQueryAndCompare(
+ "select k, max(ts) over (partition by k) " +
+ "from ntz_window_data") {
+ checkGlutenPlan[WindowExecTransformer]
+ }
+ }
+ }
+
+ test("unaudited window function argument with TIMESTAMP_NTZ falls back") {
+ withTable("ntz_window_data") {
+ createNtzWindowTable()
+
+ val df = spark.sql(
+ "select k, collect_list(ts) over (partition by k order by v) " +
+ "from ntz_window_data")
+ // This plan never goes native at all, so the helper's count - 1 comes
+ // out negative. collect_list isn't whitelisted yet (unaudited for NTZ
+ // support), so this must always fall back.
+ checkFallbackOperators(df, -1)
Review Comment:
Asserting `checkFallbackOperators(df, -1)` is brittle and encodes an
internal implementation detail (“count - 1 goes negative”). This can easily
break if the helper or plan shape changes. Prefer asserting fallback in a way
that doesn’t rely on negative counts (e.g., assert no `WindowExecTransformer`
appears in the plan, or add/extend a helper that explicitly verifies the query
is handled by Spark for this operator).
##########
gluten-substrait/src/main/scala/org/apache/gluten/extension/columnar/validator/Validators.scala:
##########
@@ -299,6 +300,46 @@ object Validators {
case _ => false
}
}
+ case w: WindowExec =>
+ w.windowExpression.forall {
+ expr =>
+ expr.collectFirst { case WindowExpression(fn, _) => fn } match
{
+ case Some(windowFunction) =>
+ (!containsNTZ(windowFunction.dataType) &&
+ !windowFunction.references.exists(a =>
containsNTZ(a.dataType))) ||
+ windowFunction.exists {
+ case Lag(input, _, _, _) => containsNTZ(input.dataType)
+ case Lead(input, _, _, _) => containsNTZ(input.dataType)
+ case NthValue(input, _, _) => containsNTZ(input.dataType)
+ case First(child, _) => containsNTZ(child.dataType)
+ case Last(child, _) => containsNTZ(child.dataType)
+ case Count(children) => children.exists(c =>
containsNTZ(c.dataType))
+ case Min(child) => containsNTZ(child.dataType)
+ case Max(child) => containsNTZ(child.dataType)
+ // RankLike's children mirror orderSpec; only ever
compared, not computed.
+ case _: RankLike => true
+ case _ => false
+ }
+ case None => false
+ }
+ }
Review Comment:
The new `WindowExec` NTZ validation only inspects `windowExpression` (the
window function tree). It does not validate `WindowExec.partitionSpec` /
`orderSpec` directly, which are the actual partition/order-by keys this PR aims
to permit. If Spark produces a `WindowExec` without a preceding
`SortExec`/`ShuffleExchangeExec` (e.g., input already satisfies
distribution/order), this can allow native execution without checking the NTZ
keys at all. Recommend adding explicit checks for `w.partitionSpec` and
`w.orderSpec.map(_.child)` (similar to the `SortExec` / `ShuffleExchangeExec`
rules: allow NTZ only for bare `Attribute` keys, otherwise fallback).
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
