This is an automated email from the ASF dual-hosted git repository.
cloud-fan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new ed150f9df74a [SPARK-57424][SQL] Add First/Last to segment-tree window
aggregate allowlist
ed150f9df74a is described below
commit ed150f9df74a6a8ee9caf72ea8d9fcd1729cc4d3
Author: Anupam Yadav <[email protected]>
AuthorDate: Thu Jul 9 10:10:03 2026 +0800
[SPARK-57424][SQL] Add First/Last to segment-tree window aggregate allowlist
### What changes were proposed in this pull request?
Add `classOf[First]` and `classOf[Last]` to
`WindowSegmentTree.EligibleAggregates`,
routing First/Last window aggregates through the segment-tree path
established
by SPARK-56546 (sliding) and SPARK-57220 (shrinking) instead of the legacy
O(N x W) sliding / O(N^2) shrinking frame implementations. No new frame
class,
no new SQLConf, no dispatcher changes -- the existing dispatcher branches in
`WindowEvaluatorFactoryBase` already gate on `eligibleForSegTree`, which
calls
`WindowSegmentTree.isEligible`.
### Why are the changes needed?
`First` and `Last` were previously denylisted as "order-dependent". This was
over-conservative: order-dependence in row-traversal order is exactly what
`WindowSegmentTree.query` provides. The query walks left-to-right (left
partial -> full blocks ascending -> right partial; within a block,
`queryDescend` walks children in ascending index order).
`First.mergeExpressions`
and `Last.mergeExpressions` are correct under that traversal -- they pick
the
row-order extreme across any contiguous range. For IGNORE NULLS the same
merge
is mode-agnostic: per-row `updateExpressions` only set `valueSet=true` on
non-null values, so a per-block partial of `(null, false)` for an all-NULL
block is correctly skipped when merged with a later non-null block.
JIRA: https://issues.apache.org/jira/browse/SPARK-57424
### Does this PR introduce _any_ user-facing change?
Yes -- when `spark.sql.window.segmentTree.enabled=true`, FIRST/LAST window
aggregates over sliding or shrinking ROWS/RANGE frames execute through the
segment-tree path instead of the legacy frame implementations. Same opt-in
conf (default off), same eligibility allowlist mechanism, same fallback
below
`minPartitionRows`, same SQLMetrics. No public API changes.
### How was this patch tested?
New tests, all differential against the legacy frame (segment-tree on vs
off):
* `WindowSegmentTreeAllowlistSuite`: routing tests for
`first / last / first_ignore_nulls / last_ignore_nulls`; the previous
"first/last falls through" negative tests are flipped; the mixed-allowlist
test now uses `collect_list` (still on the denylist).
* `SegmentTreeWindowFunctionSuite`: sliding First/Last respect-nulls and
ignore-nulls, all-NULL columns in both modes, a
stretches-of-consecutive-NULLs
case for the IGNORE NULLS merge path, and a multi-block case (block size
16, a
40-row partition with an all-NULL middle block, wide frame) that forces
the
cross-block left-to-right combine spine and the all-NULL block partial the
correctness argument depends on.
* `UnboundedFollowingSegmentTreeSuite`: shrinking First/Last respect-nulls
and
ignore-nulls plus an all-NULL column boundary case.
### Benchmark
`FirstLastSegmentTreeWindowBenchmark` (results checked in at
`sql/core/benchmarks/FirstLastSegmentTreeWindowBenchmark-results.txt`; Linux
x86_64, Intel Xeon Platinum 8259CL 2.50GHz, OpenJDK 17):
Sliding frame `[-1000, +1000]` at N=10K:
| Aggregate | Naive | Segtree | Speedup |
|---|---|---|---|
| FIRST respect-nulls | 439 ms | 97 ms | 4.5x |
| LAST respect-nulls | 540 ms | 86 ms | 6.3x |
| FIRST ignore-nulls | 535 ms | 88 ms | 6.1x |
| LAST ignore-nulls | 729 ms | 83 ms | 8.8x |
Shrinking frame `[CURRENT ROW, UNBOUNDED FOLLOWING]` at N=10K:
| Aggregate | Naive | Segtree | Speedup |
|---|---|---|---|
| FIRST respect-nulls | 2,190 ms | 78 ms | 28.0x |
| LAST respect-nulls | 2,175 ms | 86 ms | 25.4x |
| FIRST ignore-nulls | 2,433 ms | 71 ms | 34.5x |
| LAST ignore-nulls | 2,887 ms | 72 ms | 39.9x |
N-sweep on FIRST shrinking:
| N | Naive | Segtree | Speedup |
|---|---|---|---|
| 5K | 584 ms | 66 ms | 8.8x |
| 25K | 13,473 ms | 96 ms | 140.3x |
| 50K | 53,593 ms | 154 ms | 347.5x |
| 100K | -- | 224 ms | -- |
Naive at N=100K is omitted (extrapolated cost ~3-4 min/iter); segtree path
stays sub-second.
### Was this patch authored or co-authored using generative AI tooling?
Yes.
Closes #56485 from yadavay-amzn/firstlast-segtree.
Authored-by: Anupam Yadav <[email protected]>
Signed-off-by: Wenchen Fan <[email protected]>
---
...FirstLastSegmentTreeWindowBenchmark-results.txt | 143 ++++++++++++++
.../sql/execution/window/WindowSegmentTree.scala | 29 ++-
.../FirstLastSegmentTreeWindowBenchmark.scala | 212 +++++++++++++++++++++
.../window/SegmentTreeWindowFunctionSuite.scala | 101 ++++++++++
.../UnboundedFollowingSegmentTreeSuite.scala | 53 ++++++
.../window/WindowSegmentTreeAllowlistSuite.scala | 28 +--
6 files changed, 538 insertions(+), 28 deletions(-)
diff --git
a/sql/core/benchmarks/FirstLastSegmentTreeWindowBenchmark-results.txt
b/sql/core/benchmarks/FirstLastSegmentTreeWindowBenchmark-results.txt
new file mode 100644
index 000000000000..9337dee15e6d
--- /dev/null
+++ b/sql/core/benchmarks/FirstLastSegmentTreeWindowBenchmark-results.txt
@@ -0,0 +1,143 @@
+================================================================================================
+Section A - FIRST sliding respect-nulls
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST sliding respect-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+FIRST sliding respect-nulls naive 439 450
9 0.0 42903.0 1.0X
+FIRST sliding respect-nulls segtree 97 99
1 0.1 9489.9 4.5X
+
+
+================================================================================================
+Section A - LAST sliding respect-nulls
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+LAST sliding respect-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+LAST sliding respect-nulls naive 540 542
2 0.0 52725.5 1.0X
+LAST sliding respect-nulls segtree 86 91
4 0.1 8412.8 6.3X
+
+
+================================================================================================
+Section A - FIRST sliding IGNORE NULLS
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST sliding ignore-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+FIRST sliding ignore-nulls naive 535 551
18 0.0 52204.9 1.0X
+FIRST sliding ignore-nulls segtree 88 88
1 0.1 8582.2 6.1X
+
+
+================================================================================================
+Section A - LAST sliding IGNORE NULLS
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+LAST sliding ignore-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+LAST sliding ignore-nulls naive 729 733
3 0.0 71238.0 1.0X
+LAST sliding ignore-nulls segtree 83 83
0 0.1 8077.1 8.8X
+
+
+================================================================================================
+Section B - FIRST shrinking respect-nulls
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST shrinking respect-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+-------------------------------------------------------------------------------------------------------------------------
+FIRST shrinking respect-nulls naive 2190 2212
20 0.0 213893.5 1.0X
+FIRST shrinking respect-nulls segtree 78 79
1 0.1 7648.5 28.0X
+
+
+================================================================================================
+Section B - LAST shrinking respect-nulls
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+LAST shrinking respect-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+LAST shrinking respect-nulls naive 2175 2187
8 0.0 212398.8 1.0X
+LAST shrinking respect-nulls segtree 86 90
3 0.1 8377.1 25.4X
+
+
+================================================================================================
+Section B - FIRST shrinking IGNORE NULLS
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST shrinking ignore-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+FIRST shrinking ignore-nulls naive 2433 2454
13 0.0 237596.6 1.0X
+FIRST shrinking ignore-nulls segtree 71 73
3 0.1 6885.1 34.5X
+
+
+================================================================================================
+Section B - LAST shrinking IGNORE NULLS
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+LAST shrinking ignore-nulls, N=10K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+LAST shrinking ignore-nulls naive 2887 2897
11 0.0 281963.0 1.0X
+LAST shrinking ignore-nulls segtree 72 74
1 0.1 7065.8 39.9X
+
+
+================================================================================================
+Section C - N=5K
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST shrinking frame, N=5K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+FIRST naive N=5K 584 592
8 0.0 114141.5 1.0X
+FIRST segtree N=5K 66 67
0 0.1 12925.3 8.8X
+
+
+================================================================================================
+Section C - N=25K (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST shrinking frame, N=25K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+FIRST naive N=25K 13473 13494
27 0.0 526301.8 1.0X
+FIRST segtree N=25K 96 99
5 0.3 3750.5 140.3X
+
+
+================================================================================================
+Section C - N=50K (stress, last naive run)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST shrinking frame, N=50K rows: Best Time(ms) Avg Time(ms)
Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+FIRST naive N=50K 53593 54105
846 0.0 1046735.1 1.0X
+FIRST segtree N=50K 154 156
1 0.3 3012.0 347.5X
+
+
+================================================================================================
+Section C - N=100K (segtree-only, stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux
5.10.258-261.1043.amzn2int.x86_64
+Intel(R) Xeon(R) Platinum 8259CL CPU @ 2.50GHz
+FIRST shrinking frame, N=100K rows (segtree-only): Best Time(ms) Avg
Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+---------------------------------------------------------------------------------------------------------------------------------
+FIRST segtree N=100K 224
224 0 0.5 2188.5 1.0X
+
+
diff --git
a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala
b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala
index 27a773636134..cdc6556f1862 100644
---
a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala
+++
b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala
@@ -25,7 +25,7 @@ import org.apache.spark.SparkException
import org.apache.spark.memory.{MemoryConsumer, MemoryMode, TaskMemoryManager}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions._
-import org.apache.spark.sql.catalyst.expressions.aggregate.{Average, Count,
DeclarativeAggregate, Max, Min, StddevPop, StddevSamp, Sum, VariancePop,
VarianceSamp}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Average, Count,
DeclarativeAggregate, First, Last, Max, Min, StddevPop, StddevSamp, Sum,
VariancePop, VarianceSamp}
import org.apache.spark.sql.errors.QueryExecutionErrors
import org.apache.spark.sql.execution.ExternalAppendOnlyUnsafeRowArray
import org.apache.spark.sql.types.DataType
@@ -572,20 +572,29 @@ private[window] object WindowSegmentTree {
/**
* Explicit allowlist of [[DeclarativeAggregate]] subclasses safe for
- * segment-tree execution. Safe iff combine semantics form a commutative
- * monoid on the partial-buffer representation (associativity +
- * compatibility with `mergeExpressions`):
+ * segment-tree execution. Safe iff combine semantics are correct under the
+ * left-to-right combine order produced by [[WindowSegmentTree.query]]
+ * (left partial -> full blocks ascending -> right partial; within a block,
+ * `queryDescend` walks children in ascending index order).
*
- * - [[Min]], [[Max]]: idempotent semilattice.
- * - [[Sum]], [[Count]]: additive monoid.
+ * - [[Min]], [[Max]]: idempotent semilattice (associative + commutative).
+ * - [[Sum]], [[Count]]: additive monoid (associative + commutative).
* - [[Average]]: sum + count, both additive monoids.
* - [[StddevPop]], [[StddevSamp]], [[VariancePop]], [[VarianceSamp]]:
* Welford (count, mean, M2) is associative -- see
* CentralMomentAgg.mergeExpressions.
+ * - [[First]], [[Last]]: order-dependent but correct under left-to-right
+ * combine. `First.mergeExpressions` is `if(valueSet.left, left, right)`
+ * and `Last.mergeExpressions` is `if(valueSet.right, right, left)`;
+ * under the left-to-right traversal both pick the row-order extreme
+ * across any contiguous range. `IGNORE NULLS` is also handled: per-row
+ * `updateExpressions` only sets `valueSet=true` on non-null values, so
+ * a per-block partial of `(null, false)` for an all-NULL block is
+ * correctly skipped when merged with a later non-null block.
*
* Intentionally excluded (tracked as follow-up): HyperLogLogPlusPlus /
- * ApproxCountDistinct (sketch-buffer interaction unaudited), First / Last
- * (order-dependent), CollectList / CollectSet (unbounded buffer growth),
+ * ApproxCountDistinct (sketch-buffer interaction unaudited),
+ * CollectList / CollectSet (unbounded buffer growth),
* Percentile / ApproxPercentile (sorted-sketch buffer), and any
* ImperativeAggregate (excluded by the type check).
*
@@ -600,7 +609,9 @@ private[window] object WindowSegmentTree {
classOf[StddevPop],
classOf[StddevSamp],
classOf[VariancePop],
- classOf[VarianceSamp]
+ classOf[VarianceSamp],
+ classOf[First],
+ classOf[Last]
)
/**
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/FirstLastSegmentTreeWindowBenchmark.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/FirstLastSegmentTreeWindowBenchmark.scala
new file mode 100644
index 000000000000..f1287419704a
--- /dev/null
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/FirstLastSegmentTreeWindowBenchmark.scala
@@ -0,0 +1,212 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements. See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF licenses this file to You under the Apache License, Version 2.0
+ * (the "License"); you may not use this file except in compliance with
+ * the License. You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+package org.apache.spark.sql.execution.benchmark
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Benchmark for FIRST/LAST window aggregates over sliding and shrinking
+ * ROWS frames, comparing the legacy O(N x W) / O(N^2) frame paths against
+ * the segment-tree path enabled by adding `classOf[First]` / `classOf[Last]`
+ * to `WindowSegmentTree.EligibleAggregates`.
+ *
+ * Today's slow paths:
+ * - Sliding: `SlidingWindowFunctionFrame.write` rebuilds the per-row
+ * buffer aggregate by iterating `processor.update` over every row in
+ * the buffer (O(W) per output row, O(N*W) total).
+ * - Shrinking: `UnboundedFollowingWindowFunctionFrame.write` walks the
+ * remaining suffix on every output row (O(N^2) total; class scaladoc
+ * literally says O(n*(n-1)/2)).
+ *
+ * Sections:
+ * - A: FIRST/LAST per-mode at N=10K, sliding wide frame.
+ * - B: FIRST/LAST per-mode at N=10K, shrinking frame.
+ * - C: N-sweep for FIRST shrinking, naive vs segtree, demonstrating the
+ * algorithmic gap. Mirrors UnboundedFollowingWindowBenchmark layout.
+ */
+object FirstLastSegmentTreeWindowBenchmark extends SqlBasedBenchmark {
+
+ // Section A/B: per-mode per-frame-shape at calibrated N
+ private val AB_N: Long = 10L * 1024L
+
+ // Section C: N-sweep for shrinking FIRST
+ private val C_N_SMALL: Long = 5L * 1024L
+ private val C_N_MID: Long = 25L * 1024L
+ private val C_N_LARGE: Long = 50L * 1024L
+ private val C_N_HUGE: Long = 100L * 1024L
+
+ private val ITERS_NORMAL: Int = 5
+ private val ITERS_STRESS: Int = 3
+
+ // Sliding frame width tuned so the O(N*W) baseline is observable but not
+ // catastrophic. With N=10K and W=2001 the legacy SlidingWindowFunctionFrame
+ // performs ~20M update calls (10K rows x ~2K-row buffer rebuild on each
+ // boundary change) which is enough to expose the gap without dominating
+ // wall-clock at calibration time.
+ private val SLIDING_FRAME =
+ "OVER (ORDER BY id ROWS BETWEEN 1000 PRECEDING AND 1000 FOLLOWING)"
+ private val SHRINKING_FRAME =
+ "OVER (ORDER BY id ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING)"
+
+ override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+ val smokeMode = mainArgs.nonEmpty
+ val smokeRowCount = if (smokeMode) mainArgs(0).toLong else 0L
+
+ def setupIntTable(n: Long): Unit = {
+ // Sprinkle ~10% NULLs so IGNORE NULLS exercises the merge path
+ // distinctly from respect-nulls; integer values otherwise.
+ spark.range(n)
+ .selectExpr("id",
+ "CASE WHEN rand(7) < 0.1 THEN NULL " +
+ "ELSE cast(rand(42) * 1000000 as int) END as v")
+ .coalesce(1)
+ .createOrReplaceTempView("t")
+ }
+
+ def rowsLabel(rows: Long): String = {
+ if (rows >= 1000000) s"${rows / 1000000}M"
+ else if (rows >= 1024) s"${rows / 1024}K"
+ else rows.toString
+ }
+
+ /**
+ * Equivalence digest. FIRST/LAST in respect-nulls mode are bit-exact
+ * across naive and segtree paths; in IGNORE NULLS the per-block merge
+ * also yields the same result row-for-row. We hash the result column
+ * directly (not COALESCE'd) so a NULL row hashes to NULL and a NULL
+ * sum is propagated, but the comparison still distinguishes shapes.
+ */
+ def digest(sql: String, sqlConfs: (String, String)*): Long = {
+ withSQLConf(sqlConfs: _*) {
+ val r = spark.sql(s"SELECT SUM(HASH(m)) FROM (SELECT $sql AS m FROM
t)")
+ .head().get(0)
+ if (r == null) 0L else r.asInstanceOf[Long]
+ }
+ }
+
+ def runCase(label: String, sql: String, iters: Int, rows: Long): Unit = {
+ val dNaive = digest(sql)
+ val dSeg = digest(sql, SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true")
+ require(dNaive == dSeg,
+ s"$label digest mismatch: naive=$dNaive seg=$dSeg")
+
+ val benchmark = new Benchmark(
+ s"$label, N=${rowsLabel(rows)} rows", rows, output = output)
+ benchmark.addCase(s"$label naive", numIters = iters) { _ =>
+ spark.sql(s"SELECT $sql FROM t").noop()
+ }
+ benchmark.addCase(s"$label segtree", numIters = iters) { _ =>
+ withSQLConf(SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true") {
+ spark.sql(s"SELECT $sql FROM t").noop()
+ }
+ }
+ benchmark.run()
+ }
+
+ def runSweepCase(rows: Long, includeNaive: Boolean, iters: Int): Unit = {
+ val sql = s"FIRST(v) $SHRINKING_FRAME"
+ if (includeNaive) {
+ val dNaive = digest(sql)
+ val dSeg = digest(sql, SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key ->
"true")
+ require(dNaive == dSeg,
+ s"Section C N=${rowsLabel(rows)} digest mismatch: naive=$dNaive
seg=$dSeg")
+ }
+ val benchmark = new Benchmark(
+ s"FIRST shrinking frame, N=${rowsLabel(rows)} rows" +
+ (if (!includeNaive) " (segtree-only)" else ""),
+ rows, output = output)
+ if (includeNaive) {
+ benchmark.addCase(s"FIRST naive N=${rowsLabel(rows)}", numIters =
iters) { _ =>
+ spark.sql(s"SELECT $sql FROM t").noop()
+ }
+ }
+ benchmark.addCase(s"FIRST segtree N=${rowsLabel(rows)}", numIters =
iters) { _ =>
+ withSQLConf(SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true") {
+ spark.sql(s"SELECT $sql FROM t").noop()
+ }
+ }
+ benchmark.run()
+ }
+
+ if (smokeMode) {
+ setupIntTable(smokeRowCount)
+ runBenchmark("SMOKE Section A FIRST sliding") {
+ runCase("FIRST sliding respect-nulls",
+ s"FIRST(v) $SLIDING_FRAME", ITERS_STRESS, smokeRowCount)
+ }
+ } else {
+ setupIntTable(AB_N)
+
+ // Section A: sliding frame, all four mode/function combinations.
+ runBenchmark("Section A - FIRST sliding respect-nulls") {
+ runCase("FIRST sliding respect-nulls",
+ s"FIRST(v) $SLIDING_FRAME", ITERS_NORMAL, AB_N)
+ }
+ runBenchmark("Section A - LAST sliding respect-nulls") {
+ runCase("LAST sliding respect-nulls",
+ s"LAST(v) $SLIDING_FRAME", ITERS_NORMAL, AB_N)
+ }
+ runBenchmark("Section A - FIRST sliding IGNORE NULLS") {
+ runCase("FIRST sliding ignore-nulls",
+ s"FIRST(v) IGNORE NULLS $SLIDING_FRAME", ITERS_NORMAL, AB_N)
+ }
+ runBenchmark("Section A - LAST sliding IGNORE NULLS") {
+ runCase("LAST sliding ignore-nulls",
+ s"LAST(v) IGNORE NULLS $SLIDING_FRAME", ITERS_NORMAL, AB_N)
+ }
+
+ // Section B: shrinking frame, all four mode/function combinations.
+ runBenchmark("Section B - FIRST shrinking respect-nulls") {
+ runCase("FIRST shrinking respect-nulls",
+ s"FIRST(v) $SHRINKING_FRAME", ITERS_NORMAL, AB_N)
+ }
+ runBenchmark("Section B - LAST shrinking respect-nulls") {
+ runCase("LAST shrinking respect-nulls",
+ s"LAST(v) $SHRINKING_FRAME", ITERS_NORMAL, AB_N)
+ }
+ runBenchmark("Section B - FIRST shrinking IGNORE NULLS") {
+ runCase("FIRST shrinking ignore-nulls",
+ s"FIRST(v) IGNORE NULLS $SHRINKING_FRAME", ITERS_NORMAL, AB_N)
+ }
+ runBenchmark("Section B - LAST shrinking IGNORE NULLS") {
+ runCase("LAST shrinking ignore-nulls",
+ s"LAST(v) IGNORE NULLS $SHRINKING_FRAME", ITERS_NORMAL, AB_N)
+ }
+
+ // Section C: shrinking-frame N-sweep on FIRST, demonstrating O(N^2)
+ // legacy vs O(N log N) segtree gap that widens with N.
+ setupIntTable(C_N_SMALL)
+ runBenchmark("Section C - N=5K") {
+ runSweepCase(C_N_SMALL, includeNaive = true, ITERS_NORMAL)
+ }
+ setupIntTable(C_N_MID)
+ runBenchmark("Section C - N=25K (stress)") {
+ runSweepCase(C_N_MID, includeNaive = true, ITERS_STRESS)
+ }
+ setupIntTable(C_N_LARGE)
+ runBenchmark("Section C - N=50K (stress, last naive run)") {
+ runSweepCase(C_N_LARGE, includeNaive = true, ITERS_STRESS)
+ }
+ setupIntTable(C_N_HUGE)
+ runBenchmark("Section C - N=100K (segtree-only, stress)") {
+ runSweepCase(C_N_HUGE, includeNaive = false, ITERS_STRESS)
+ }
+ }
+ }
+}
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionSuite.scala
index 5e644cceb142..5bf74c351c08 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionSuite.scala
@@ -99,6 +99,19 @@ class SegmentTreeWindowFunctionSuite extends
SharedSparkSession {
baseDF.select($"id", $"pk", avg($"v").over(winSpec(-3, 3)).as("agg")))
}
+ // First / Last basic equivalence (respect-nulls; the default for
+ // first()/last()). Order-correctness depends on the segment-tree
+ // combine being left-to-right; see WindowSegmentTree.EligibleAggregates.
+ test("FIRST over ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", first($"v").over(winSpec(-3, 3)).as("agg")))
+ }
+
+ test("LAST over ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", last($"v").over(winSpec(-3, 3)).as("agg")))
+ }
+
test("MIN + MAX + SUM share a single window frame") {
checkEquivalence(() =>
baseDF.select(
@@ -229,6 +242,94 @@ class SegmentTreeWindowFunctionSuite extends
SharedSparkSession {
count($"v").over(winSpec(-4, 4)).as("cn")))
}
+ // First / Last respect-nulls: NULL is a valid value. If the first row in the
+ // frame is NULL, FIRST returns NULL. The seg-tree merge must preserve this.
+ test("FIRST/LAST respect-nulls with mixed NULL frame contents") {
+ val df = spark.range(0, 60).selectExpr(
+ "id",
+ "(id % 3) AS pk",
+ "CASE WHEN id % 3 = 0 THEN NULL ELSE CAST(id AS INT) END AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v").over(winSpec(-4, 4)).as("fv"),
+ last($"v").over(winSpec(-4, 4)).as("lv")))
+ }
+
+ // First / Last IGNORE NULLS: per-row updates only set valueSet on non-null
+ // values. A per-block partial of (null, false) for an all-NULL block must
+ // be correctly skipped when merged with a later non-null block.
+ test("FIRST/LAST ignore-nulls with mixed NULL frame contents") {
+ val df = spark.range(0, 60).selectExpr(
+ "id",
+ "(id % 3) AS pk",
+ "CASE WHEN id % 3 = 0 THEN NULL ELSE CAST(id AS INT) END AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v", ignoreNulls = true).over(winSpec(-4, 4)).as("fv_ign"),
+ last($"v", ignoreNulls = true).over(winSpec(-4, 4)).as("lv_ign")))
+ }
+
+ // All-NULL column edge case for First/Last in both modes.
+ // Respect-nulls: returns NULL. Ignore-nulls: also returns NULL (no
+ // non-null candidate ever sets valueSet).
+ test("all-NULL column: FIRST/LAST in both modes") {
+ val df = spark.range(0, 30).selectExpr(
+ "id", "(id % 3) AS pk", "CAST(NULL AS INT) AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v").over(winSpec(-3, 3)).as("fv"),
+ last($"v").over(winSpec(-3, 3)).as("lv"),
+ first($"v", ignoreNulls = true).over(winSpec(-3, 3)).as("fv_ign"),
+ last($"v", ignoreNulls = true).over(winSpec(-3, 3)).as("lv_ign")))
+ }
+
+ // Adversarial NULL distribution for IGNORE NULLS: per-block aggregates need
+ // to compose correctly when an entire block is all-NULL. With block size
+ // 65536 and partition size 120 we cannot literally produce a fully-NULL
+ // block via the standard fixture, but a long stretch of consecutive NULLs
+ // exercises the same merge path (per-row updates produce intermediate
+ // valueSet=false buffers which then merge with a later valueSet=true buffer
+ // via mergeExpressions). Combined with a wide frame to force tree queries
+ // crossing the all-NULL stretch.
+ test("FIRST/LAST ignore-nulls: stretches of consecutive NULLs cross-merge
correctly") {
+ val df = spark.range(0, 90).selectExpr(
+ "id",
+ "0 AS pk",
+ // First 30 rows non-null, next 30 all NULL, last 30 non-null again.
+ "CASE WHEN id BETWEEN 30 AND 59 THEN NULL ELSE CAST(id AS INT) END AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v", ignoreNulls = true).over(winSpec(-20, 20)).as("fv_ign"),
+ last($"v", ignoreNulls = true).over(winSpec(-20, 20)).as("lv_ign")))
+ }
+
+ // Multi-block First/Last: the correctness of First/Last on the segment tree
+ // relies on the combine being applied strictly left-to-right (First keeps
the
+ // left operand, Last the right). The default 65536 block size answers every
+ // small-partition query inside a single block (queryDescend only), so the
+ // cross-block combine spine (left partial -> ascending full blocks -> right
+ // partial in WindowSegmentTree.query) is never exercised by the fixtures
+ // above. Force blockSize=16 on a 40-row partition so a wide frame crosses
+ // three blocks, and make the middle block (rows 16..31) entirely NULL so the
+ // all-NULL (null, false) block partial is produced and merged for real --
the
+ // exact path the IGNORE NULLS rationale depends on. Differential vs the
+ // legacy frame in both respect-nulls and ignore-nulls modes.
+ test("FIRST/LAST multi-block combine with an all-NULL middle block") {
+ val df = spark.range(0, 40).selectExpr(
+ "id",
+ "0 AS pk",
+ // Blocks (size 16): [0,16) non-null, [16,32) all NULL, [32,40) non-null.
+ "CASE WHEN id BETWEEN 16 AND 31 THEN NULL ELSE CAST(id AS INT) END AS v")
+ withSegTreeBlock() {
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v").over(winSpec(-20, 20)).as("fv"),
+ last($"v").over(winSpec(-20, 20)).as("lv"),
+ first($"v", ignoreNulls = true).over(winSpec(-20, 20)).as("fv_ign"),
+ last($"v", ignoreNulls = true).over(winSpec(-20, 20)).as("lv_ign")))
+ }
+ }
+
test("Double NaN and +/-Infinity propagate correctly through MIN/MAX/SUM") {
// Trap: NaN > +Inf in Spark's MIN/MAX ordering; +Inf + -Inf = NaN in SUM.
// Seg-tree uses DeclarativeAggregate.merge; behavior must match baseline.
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/UnboundedFollowingSegmentTreeSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/UnboundedFollowingSegmentTreeSuite.scala
index 90df2240927b..67316ba7a048 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/UnboundedFollowingSegmentTreeSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/UnboundedFollowingSegmentTreeSuite.scala
@@ -117,6 +117,22 @@ class UnboundedFollowingSegmentTreeSuite extends
SharedSparkSession {
baseDF.select($"id", $"pk",
avg($"v").over(shrinkingRowsFrame(0)).as("agg")))
}
+ // First / Last over a shrinking frame: in respect-nulls mode, FIRST is just
+ // `rows[lower]` (the first row of the suffix). LAST advances with the
+ // shrinking lower bound but always sees the partition's end. Both modes
+ // exercise the segment-tree merge path through a series of `[lower, n)`
+ // queries; correctness depends on the same left-to-right combine that
+ // makes First/Last safe in the sliding case.
+ test("FIRST over ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk",
first($"v").over(shrinkingRowsFrame(0)).as("agg")))
+ }
+
+ test("LAST over ROWS BETWEEN CURRENT ROW AND UNBOUNDED FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk",
last($"v").over(shrinkingRowsFrame(0)).as("agg")))
+ }
+
// ============================================================
// ROWS frame: lower-bound variations
// ============================================================
@@ -240,6 +256,43 @@ class UnboundedFollowingSegmentTreeSuite extends
SharedSparkSession {
count($"v").over(shrinkingRowsFrame(0)).as("c")))
}
+ // First / Last over a shrinking frame with NULL distribution. Mirrors the
+ // sliding-suite NULL tests; verifies the merge path is correct when the
+ // lower-edge partial block of the segtree query crosses a NULL/non-NULL
+ // boundary.
+ test("FIRST/LAST over shrinking frame: respect-nulls with mixed NULLs") {
+ val df = spark.range(0, 60).selectExpr(
+ "id",
+ "(id % 3) AS pk",
+ "CASE WHEN id % 3 = 0 THEN NULL ELSE CAST(id AS INT) END AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v").over(shrinkingRowsFrame(0)).as("fv"),
+ last($"v").over(shrinkingRowsFrame(0)).as("lv")))
+ }
+
+ test("FIRST/LAST over shrinking frame: ignore-nulls with mixed NULLs") {
+ val df = spark.range(0, 60).selectExpr(
+ "id",
+ "(id % 3) AS pk",
+ "CASE WHEN id % 3 = 0 THEN NULL ELSE CAST(id AS INT) END AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v", ignoreNulls =
true).over(shrinkingRowsFrame(0)).as("fv_ign"),
+ last($"v", ignoreNulls =
true).over(shrinkingRowsFrame(0)).as("lv_ign")))
+ }
+
+ test("all-NULL column: FIRST/LAST shrinking frame in both modes") {
+ val df = spark.range(0, 30).selectExpr("id", "(id % 3) AS pk",
+ "CAST(NULL AS INT) AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ first($"v").over(shrinkingRowsFrame(0)).as("fv"),
+ last($"v").over(shrinkingRowsFrame(0)).as("lv"),
+ first($"v", ignoreNulls =
true).over(shrinkingRowsFrame(0)).as("fv_ign"),
+ last($"v", ignoreNulls =
true).over(shrinkingRowsFrame(0)).as("lv_ign")))
+ }
+
test("mixed NULL and non-NULL: NULLs must not leak into MIN/MAX") {
val df = (0 until 60).map { i =>
val v: Option[Int] = if (i % 4 == 0) None else Some(i)
diff --git
a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
index 236d38cc6a91..8dbb4ceecf1b 100644
---
a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
+++
b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
@@ -81,7 +81,13 @@ class WindowSegmentTreeAllowlistSuite
("stddev_pop", (c: org.apache.spark.sql.Column) => stddev_pop(c)),
("stddev_samp", (c: org.apache.spark.sql.Column) => stddev_samp(c)),
("var_pop", (c: org.apache.spark.sql.Column) => var_pop(c)),
- ("var_samp", (c: org.apache.spark.sql.Column) => var_samp(c))
+ ("var_samp", (c: org.apache.spark.sql.Column) => var_samp(c)),
+ ("first", (c: org.apache.spark.sql.Column) => first(c)),
+ ("last", (c: org.apache.spark.sql.Column) => last(c)),
+ ("first_ignore_nulls",
+ (c: org.apache.spark.sql.Column) => first(c, ignoreNulls = true)),
+ ("last_ignore_nulls",
+ (c: org.apache.spark.sql.Column) => last(c, ignoreNulls = true))
).foreach { case (name, fn) =>
test(s"$name routes to the segment-tree path") {
withSQLConf(enableSegTree.toSeq: _*) {
@@ -96,22 +102,6 @@ class WindowSegmentTreeAllowlistSuite
// Negative: non-allowlisted aggregates fall through
- test("first_value falls through (order-dependent aggregate)") {
- withSQLConf(enableSegTree.toSeq: _*) {
- val df = baseDF.withColumn("agg", first($"v").over(winSpec))
- val (seg, _) = segTreeCounters(df)
- assert(seg == 0, s"first_value should not use segment tree (got $seg
frames)")
- }
- }
-
- test("last_value falls through (order-dependent aggregate)") {
- withSQLConf(enableSegTree.toSeq: _*) {
- val df = baseDF.withColumn("agg", last($"v").over(winSpec))
- val (seg, _) = segTreeCounters(df)
- assert(seg == 0, s"last_value should not use segment tree (got $seg
frames)")
- }
- }
-
test("collect_list falls through (unbounded buffer)") {
withSQLConf(enableSegTree.toSeq: _*) {
val df = baseDF.withColumn("agg", collect_list($"v").over(winSpec))
@@ -176,10 +166,10 @@ class WindowSegmentTreeAllowlistSuite
withSQLConf(enableSegTree.toSeq: _*) {
val df = baseDF
.withColumn("s", sum($"v").over(winSpec))
- .withColumn("fv", first($"v").over(winSpec))
+ .withColumn("cl", collect_list($"v").over(winSpec))
val (seg, _) = segTreeCounters(df)
// Both aggregates share the same Window node; gating is
forall(isEligible),
- // so `first_value` drops the whole group.
+ // so `collect_list` (unbounded-buffer denylist) drops the whole group.
assert(seg == 0,
s"Window group containing a non-allowlisted agg must fall through (got
$seg)")
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
