You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by ya...@apache.org on 2026/05/08 04:31:55 UTC
(spark) branch branch-4.2 updated: [SPARK-56546][SQL] Block-chunked segment-tree window frame for non-invertible sliding aggregates
This is an automated email from the ASF dual-hosted git repository.
yaooqinn pushed a commit to branch branch-4.2
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/branch-4.2 by this push:
new 4476f1310972 [SPARK-56546][SQL] Block-chunked segment-tree window frame for non-invertible sliding aggregates
4476f1310972 is described below
commit 4476f13109727da56cc07c621e0ba14855feac21
Author: Kent Yao <ya...@apache.org>
AuthorDate: Fri May 8 12:30:48 2026 +0800
[SPARK-56546][SQL] Block-chunked segment-tree window frame for non-invertible sliding aggregates
## What changes were proposed in this pull request?
Introduce a **block-chunked segment-tree window function frame** that evaluates sliding aggregates in `O(log W)` per output row for the non-invertible path in Spark's window operator.
- New frame class `SegmentTreeWindowFunctionFrame` (sibling to the existing sliding frame), selected by `WindowFunctionFrameFactory` when **all** of the following hold:
1. The frame is a non-invertible sliding `ROWS`/`RANGE` frame (i.e. today's full-recompute path).
2. Every window function over that frame is an eligible `DeclarativeAggregate`: `MIN / MAX / SUM / COUNT / AVG / STDDEV_POP / STDDEV_SAMP / VAR_POP / VAR_SAMP`, with no `FILTER` clause and no `DISTINCT`.
3. The new SQLConf `spark.sql.window.segmentTree.enabled` is `true` (default `false` — **opt-in**).
- **Internal-only**: no public API change, no new SQL syntax, no user-visible catalog / plan surface. All other frames fall back to the existing path unchanged.
- Partitions below `spark.sql.window.segmentTree.minPartitionRows` fall back to the existing `SlidingWindowFunctionFrame` to avoid paying build cost on small partitions.
### New SQLConfs (all `since = 4.2.0`)
| Key | Default | Scope |
|---|---|---|
| `spark.sql.window.segmentTree.enabled` | `false` | public, boolean — master opt-in switch |
| `spark.sql.window.segmentTree.minPartitionRows` | `64` | internal, int — min rows to activate; below → naive fallback |
| `spark.sql.window.segmentTree.blockSize` | `65536` | public, int — rows per leaf block (≥ 16) |
| `spark.sql.window.segmentTree.fanout` | `16` | internal, int — internal-node fanout (≥ 2) |
### New SQLMetrics on the window operator
- segment-tree frame construction count
- fallback-to-naive count (when eligibility fails at runtime)
## Why are the changes needed?
Spark's current sliding-window aggregate evaluation recomputes the full frame for every output row when no inverse exists — MIN/MAX always, and SUM/AVG/STDDEV whenever the inverse is numerically unstable (see [SPARK-36844](https://issues.apache.org/jira/browse/SPARK-36844)). That path is `O(N·W)` even though `O(N log W)` algorithms are standard in the streaming-analytics literature (Leis et al. 2015; Arasu & Widom 2004).
Practical symptoms: sliding `MIN/MAX/SUM/AVG/COUNT/STDDEV` over `W ≳ 1000` rows become CPU-bound, spill-sensitive on variable-width types (String/Binary), and scale poorly with `W`. At `W = 4001 / 2M rows` the current engine takes ~100 s per sliding pass on modern hardware.
## Does this PR introduce _any_ user-facing change?
**No** — the feature is gated behind a new SQLConf that defaults to `false`. There are no API, SQL grammar, or plan-shape changes. When the conf is enabled, a new physical frame implementation is selected for eligible sliding frames; all other cases (invertible aggregates, non-sliding frames, ranking functions, UDFs, `FILTER`/`DISTINCT`, small partitions) take the existing path.
## How was this patch tested?
### Functional tests (all new)
- `WindowSegmentTreeSuite` — unit tests for the frame implementation (14 cases): basic aggregates, frame boundaries, NULL / NaN / ±Infinity, numeric / string / date-timestamp types, small-partition fallback, unsupported-merge / DISTINCT / disabled-conf fallback.
- `SegmentTreeWindowFunctionSuite` — end-to-end SQL-level tests (34 cases) with the conf enabled, comparing against the naive frame row-for-row as an oracle.
- `SegmentTreeWindowMetricsSuite` — SQLMetric visibility (7 cases) asserting the new operator-level metrics propagate to the UI.
- `WindowSegmentTreeMemorySuite` — `TaskMemoryManager` accounting (9 cases) covering build, spill-not-supported paths, and release-on-completion.
- `WindowSegmentTreePropertySuite` — ScalaCheck property-based suite (3 generators) asserting segtree-vs-naive equivalence over random shapes.
All existing window tests remain green with the conf at default `false`.
### Benchmark (`WindowBenchmark`)
Rerun on GitHub Actions across **JDK 17 / 21 / 25** (Azure EPYC 9V74 80-core, Linux 6.17, OpenJDK LTS), `SPARK_GENERATE_BENCHMARK_FILES=1`, baseline resized to a 3–5 s per-iteration target.
| Aggregate | W / N | JDK 17 | JDK 21 | JDK 25 |
|---|---|---|---|---|
| MIN (non-invertible) | 1001 / 256K | 8.9× | 9.0× | 9.0× |
| MAX (non-invertible) | 1001 / 256K | 10.2× | 13.1× | 14.2× |
| SUM (full recompute) | 1001 / 256K | 9.6× | 9.5× | 9.9× |
| COUNT | 1001 / 256K | 9.4× | 9.3× | 9.3× |
| AVG (multi-buffer) | 1001 / 192K | 10.6× | 10.6× | 10.2× |
| STDDEV_SAMP (multi-buffer, stress) | 1001 / 2M | 18.7× | 19.3× | 18.5× |
| SUM W-sweep (stress, cross-block) | 4001 / 2M | 32.4× | 32.6× | 29.0× |
| MAX String spill guard (stress) | 1001 / 1M | 15.7× | 16.4× | 17.6× |
N-sweep (segtree-only, `W=1001`, JDK 17): `N=2M → 1358 ns/row`, `N=8M → 1374 ns/row`, `N=16M → 1338 ns/row` — sub-linear / near-constant, matching the `O(log W)` cost model.
### Caveats (documented, not blockers)
- `W ≤ 50` is a **Pareto-loss zone** where naive is ~2–3× faster than segtree (build overhead dominates). Covered by explicit `(stress)` benchmark cases at `W=11` and `W=51`. The feature is opt-in precisely so users can stay on naive for small-W workloads.
- `W=10` pathological (~0.4×) is the same cause; documented in the design doc.
- Segtree is a secondary path only; the naive frame is untouched and remains the default.
## Memory footprint
The tree does not copy input rows. `WindowSegmentTree.build(rows)` stores a
reference to the caller-owned `ExternalAppendOnlyUnsafeRowArray` (the same
instance held by `WindowPartitionEvaluator.buffer`); `computeBlockAggregates`
and `buildBlockLevels` read it via `generateIterator` /
`generateIterator(startIndex)`. No row-level 2x overhead.
The tree's additional memory is two sets of **aggregator buffers** (not rows):
- `blockAggregates`: `O(n / blockSize)` buffers, resident for the partition.
At the default `blockSize=65536`, a 2M-row partition holds ~31 block
aggregates.
- Cached internal-node levels per block, bounded above by an estimate of
`maxCachedBlocks` per frame (`estimateMaxCachedBlocks` in
`WindowEvaluatorFactoryBase`). Each cached block holds
`O(blockSize * log_F blockSize)` aggregator slots; total bytes are accounted
to `TaskMemoryManager` and subject to spill pressure (invariants I1-I5, I8
in `WindowSegmentTree.scala` scaladoc).
Aggregator buffer width is fixed per aggregate (a handful of longs / doubles
for SUM/AVG/STDDEV; two pointers for MIN/MAX), typically one to two orders of
magnitude smaller than UnsafeRow width, so the dominant footprint remains the
caller-owned row array.
## Was this patch authored or co-authored using generative AI tooling?
Generated-by: Claude Opus 4.7
Closes #55422 from yaooqinn/window-segment-tree.
Authored-by: Kent Yao <ya...@apache.org>
Signed-off-by: Kent Yao <ya...@apache.org>
(cherry picked from commit 5ed9636e5b49a2e6dbe3d82d1b4599bf9371ee71)
Signed-off-by: Kent Yao <ya...@apache.org>
---
.../spark/sql/errors/QueryExecutionErrors.scala | 10 +
.../org/apache/spark/sql/internal/SQLConf.scala | 53 ++
.../benchmarks/WindowBenchmark-jdk21-results.txt | 172 ++++
.../benchmarks/WindowBenchmark-jdk25-results.txt | 172 ++++
sql/core/benchmarks/WindowBenchmark-results.txt | 172 ++++
.../sql/execution/window/AggregateProcessor.scala | 24 +-
.../window/SegmentTreeWindowFunctionFrame.scala | 294 +++++++
.../execution/window/WindowEvaluatorFactory.scala | 8 +-
.../window/WindowEvaluatorFactoryBase.scala | 127 ++-
.../spark/sql/execution/window/WindowExec.scala | 11 +-
.../sql/execution/window/WindowSegmentTree.scala | 614 ++++++++++++++
.../sql/execution/benchmark/WindowBenchmark.scala | 372 +++++++++
.../window/SegmentTreeWindowFunctionSuite.scala | 896 +++++++++++++++++++++
.../window/SegmentTreeWindowMetricsSuite.scala | 213 +++++
.../window/SegmentTreeWindowTestHelper.scala | 176 ++++
.../window/WindowSegmentTreeAllowlistSuite.scala | 187 +++++
.../window/WindowSegmentTreeMemorySuite.scala | 351 ++++++++
.../window/WindowSegmentTreePropertySuite.scala | 475 +++++++++++
.../execution/window/WindowSegmentTreeSuite.scala | 521 ++++++++++++
19 files changed, 4831 insertions(+), 17 deletions(-)
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
index a1f19e0d9b90..84fb16521105 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/errors/QueryExecutionErrors.scala
@@ -1194,6 +1194,16 @@ private[sql] object QueryExecutionErrors extends QueryErrorsBase with ExecutionE
new java.util.HashMap[String, String]())
}
+ def cannotAcquireMemoryForWindowAggregateError(
+ requestedBytes: Long,
+ receivedBytes: Long): SparkOutOfMemoryError = {
+ new SparkOutOfMemoryError(
+ "UNABLE_TO_ACQUIRE_MEMORY",
+ java.util.Map.of(
+ "requestedBytes", requestedBytes.toString,
+ "receivedBytes", receivedBytes.toString))
+ }
+
def rowLargerThan256MUnsupportedError(): SparkUnsupportedOperationException = {
new SparkUnsupportedOperationException("_LEGACY_ERROR_TEMP_2108")
}
diff --git a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
index e3fbed69d8b6..bd3fba48c965 100644
--- a/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
+++ b/sql/catalyst/src/main/scala/org/apache/spark/sql/internal/SQLConf.scala
@@ -4115,6 +4115,51 @@ object SQLConf {
"The threshold of window group limit must be -1, 0 or positive integer.")
.createWithDefault(1000)
+ val WINDOW_SEGMENT_TREE_ENABLED =
+ buildConf("spark.sql.window.segmentTree.enabled")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .doc("Use block-chunked segment tree for moving aggregate window frames " +
+ "whose functions are all DeclarativeAggregate without FILTER/DISTINCT.")
+ .version("4.2.0")
+ .booleanConf
+ .createWithDefault(false)
+
+ val WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS =
+ buildConf("spark.sql.window.segmentTree.minPartitionRows")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .doc("Minimum partition row count to activate the segment-tree moving frame. " +
+ "Partitions smaller than this fall back to the default sliding implementation.")
+ .version("4.2.0")
+ .internal()
+ .intConf
+ .checkValue(_ > 0, "minPartitionRows must be > 0")
+ .createWithDefault(64)
+
+ val WINDOW_SEGMENT_TREE_BLOCK_SIZE =
+ buildConf("spark.sql.window.segmentTree.blockSize")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .doc("Block size, in rows, for the block-chunked segment tree used by moving " +
+ "window frames. Each leaf of the tree aggregates this many consecutive rows. " +
+ "Smaller values reduce per-partition memory and speed up tree build for small " +
+ "partitions, but make the tree deeper and increase query cost for wide frames. " +
+ "Larger values amortize build cost and shrink the tree but increase the per-block " +
+ "prefix/suffix scan cost within a block. The default is tuned for partitions on " +
+ "the order of tens of thousands to millions of rows.")
+ .version("4.2.0")
+ .intConf
+ .checkValue(_ >= 16, "blockSize must be >= 16")
+ .createWithDefault(65536)
+
+ val WINDOW_SEGMENT_TREE_FANOUT =
+ buildConf("spark.sql.window.segmentTree.fanout")
+ .withBindingPolicy(ConfigBindingPolicy.SESSION)
+ .doc("Fanout of internal nodes for the block-chunked segment tree.")
+ .version("4.2.0")
+ .internal()
+ .intConf
+ .checkValue(_ >= 2, "fanout must be >= 2")
+ .createWithDefault(16)
+
val SESSION_WINDOW_BUFFER_IN_MEMORY_THRESHOLD =
buildConf("spark.sql.sessionWindow.buffer.in.memory.threshold")
.internal()
@@ -8083,6 +8128,14 @@ class SQLConf extends Serializable with Logging with SqlApiConf {
def windowGroupLimitThreshold: Int = getConf(WINDOW_GROUP_LIMIT_THRESHOLD)
+ def windowSegmentTreeEnabled: Boolean = getConf(WINDOW_SEGMENT_TREE_ENABLED)
+
+ def windowSegmentTreeMinPartitionRows: Int = getConf(WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS)
+
+ def windowSegmentTreeBlockSize: Int = getConf(WINDOW_SEGMENT_TREE_BLOCK_SIZE)
+
+ def windowSegmentTreeFanout: Int = getConf(WINDOW_SEGMENT_TREE_FANOUT)
+
def sessionWindowBufferInMemoryThreshold: Int = getConf(SESSION_WINDOW_BUFFER_IN_MEMORY_THRESHOLD)
def sessionWindowBufferSpillThreshold: Int = getConf(SESSION_WINDOW_BUFFER_SPILL_THRESHOLD)
diff --git a/sql/core/benchmarks/WindowBenchmark-jdk21-results.txt b/sql/core/benchmarks/WindowBenchmark-jdk21-results.txt
new file mode 100644
index 000000000000..8838a4a9a0fe
--- /dev/null
+++ b/sql/core/benchmarks/WindowBenchmark-jdk21-results.txt
@@ -0,0 +1,172 @@
+================================================================================================
+Section A - MIN (non-invertible)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+MIN sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+MIN naive (current, baseline) 4112 4231 227 0.1 15684.2 1.0X
+MIN segtree (default) 428 447 20 0.6 1633.2 9.6X
+MIN segtree (blockSize=256) 2271 2290 17 0.1 8664.1 1.8X
+
+
+================================================================================================
+Section A - MAX (non-invertible)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+MAX sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+MAX naive (current, baseline) 4141 4154 14 0.1 15794.9 1.0X
+MAX segtree (default) 418 430 11 0.6 1596.1 9.9X
+MAX segtree (blockSize=256) 2266 2284 18 0.1 8645.3 1.8X
+
+
+================================================================================================
+Section A - SUM (Spark has no inverse; full recompute)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive (current, baseline) 4120 4128 7 0.1 15716.6 1.0X
+SUM segtree (default) 420 444 17 0.6 1601.8 9.8X
+SUM segtree (blockSize=256) 2290 2305 12 0.1 8734.2 1.8X
+
+
+================================================================================================
+Section A - COUNT
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+COUNT sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+COUNT naive (current, baseline) 3561 3577 14 0.1 13583.0 1.0X
+COUNT segtree (default) 384 400 12 0.7 1466.3 9.3X
+COUNT segtree (blockSize=256) 2232 2251 18 0.1 8515.2 1.6X
+
+
+================================================================================================
+Section A - AVG (multi-buffer)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+AVG sliding window, W=1001, 192K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+AVG naive (current, baseline) 4288 4302 16 0.0 21808.7 1.0X
+AVG segtree (default) 426 435 11 0.5 2168.3 10.1X
+AVG segtree (blockSize=256) 1515 1528 13 0.1 7706.3 2.8X
+
+
+================================================================================================
+Section A - STDDEV_SAMP (multi-buffer, stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+STDDEV_SAMP sliding window, W=1001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+STDDEV_SAMP naive (current, baseline) 100695 100937 361 0.0 50347.4 1.0X
+STDDEV_SAMP segtree (default) 5459 5475 15 0.4 2729.3 18.4X
+STDDEV_SAMP segtree (blockSize=256) 111304 111341 33 0.0 55651.8 0.9X
+
+
+================================================================================================
+Section B - W=10 scaling (stress: Pareto loss zone)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM scaling, W=11, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=11 924 948 23 2.2 461.9 1.0X
+SUM segtree (default) W=11 2262 2277 14 0.9 1130.8 0.4X
+
+
+================================================================================================
+Section B - W=50 scaling (stress: Pareto loss zone)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM scaling, W=51, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=51 2157 2169 15 0.9 1078.5 1.0X
+SUM segtree (default) W=51 2558 2571 16 0.8 1279.1 0.8X
+
+
+================================================================================================
+Section B - W=201 scaling
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM scaling, W=201, 1M rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=201 3379 3404 16 0.3 3378.8 1.0X
+SUM segtree (default) W=201 1430 1454 24 0.7 1429.9 2.4X
+
+
+================================================================================================
+Section B - W=4001 scaling (stress, + bs=256 cross-block)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM scaling, W=4001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=4001 122629 124509 1635 0.0 61314.3 1.0X
+SUM segtree (default) W=4001 3697 3755 53 0.5 1848.6 33.2X
+SUM segtree (blockSize=256) W=4001 108139 108330 165 0.0 54069.6 1.1X
+
+
+================================================================================================
+Section F - spill regression guard (String, stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+MAX String spill guard, W=1001, 1M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+--------------------------------------------------------------------------------------------------------------------------------
+MAX naive (String) 51489 51810 520 0.0 51489.5 1.0X
+MAX segtree default (String) 3008 3017 10 0.3 3007.5 17.1X
+
+
+================================================================================================
+Section C - N-sweep small (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM N-sweep (segtree-only), W=1001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=2M 3256 3284 26 0.6 1628.2 1.0X
+
+
+================================================================================================
+Section C - N-sweep mid (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM N-sweep (segtree-only), W=1001, 8M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=8M 15089 15176 131 0.5 1886.1 1.0X
+
+
+================================================================================================
+Section C - N-sweep large (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 21.0.11+10-LTS on Linux 6.17.0-1010-azure
+AMD EPYC 7763 64-Core Processor
+SUM N-sweep (segtree-only), W=1001, 16M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+-------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=16M 36197 36269 92 0.4 2262.3 1.0X
+
+
diff --git a/sql/core/benchmarks/WindowBenchmark-jdk25-results.txt b/sql/core/benchmarks/WindowBenchmark-jdk25-results.txt
new file mode 100644
index 000000000000..50dcff4c884d
--- /dev/null
+++ b/sql/core/benchmarks/WindowBenchmark-jdk25-results.txt
@@ -0,0 +1,172 @@
+================================================================================================
+Section A - MIN (non-invertible)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+MIN sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+MIN naive (current, baseline) 3840 3857 10 0.1 14648.6 1.0X
+MIN segtree (default) 421 441 14 0.6 1604.7 9.1X
+MIN segtree (blockSize=256) 1951 1973 19 0.1 7444.1 2.0X
+
+
+================================================================================================
+Section A - MAX (non-invertible)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+MAX sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+MAX naive (current, baseline) 5415 5537 158 0.0 20657.2 1.0X
+MAX segtree (default) 414 422 9 0.6 1578.6 13.1X
+MAX segtree (blockSize=256) 1959 1975 13 0.1 7471.7 2.8X
+
+
+================================================================================================
+Section A - SUM (Spark has no inverse; full recompute)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive (current, baseline) 4135 4143 5 0.1 15775.4 1.0X
+SUM segtree (default) 423 432 14 0.6 1614.2 9.8X
+SUM segtree (blockSize=256) 1972 1991 16 0.1 7524.2 2.1X
+
+
+================================================================================================
+Section A - COUNT
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+COUNT sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+COUNT naive (current, baseline) 3569 3632 127 0.1 13615.0 1.0X
+COUNT segtree (default) 394 402 10 0.7 1504.7 9.0X
+COUNT segtree (blockSize=256) 1939 1961 17 0.1 7396.6 1.8X
+
+
+================================================================================================
+Section A - AVG (multi-buffer)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+AVG sliding window, W=1001, 192K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+AVG naive (current, baseline) 4134 4147 11 0.0 21028.2 1.0X
+AVG segtree (default) 413 433 12 0.5 2102.7 10.0X
+AVG segtree (blockSize=256) 1347 1384 32 0.1 6852.3 3.1X
+
+
+================================================================================================
+Section A - STDDEV_SAMP (multi-buffer, stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+STDDEV_SAMP sliding window, W=1001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+STDDEV_SAMP naive (current, baseline) 105327 105484 189 0.0 52663.3 1.0X
+STDDEV_SAMP segtree (default) 5650 5657 12 0.4 2825.1 18.6X
+STDDEV_SAMP segtree (blockSize=256) 97218 97841 551 0.0 48609.2 1.1X
+
+
+================================================================================================
+Section B - W=10 scaling (stress: Pareto loss zone)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=11, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=11 1016 1023 6 2.0 507.9 1.0X
+SUM segtree (default) W=11 2319 2336 15 0.9 1159.7 0.4X
+
+
+================================================================================================
+Section B - W=50 scaling (stress: Pareto loss zone)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=51, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=51 2208 2234 22 0.9 1103.8 1.0X
+SUM segtree (default) W=51 2653 2661 13 0.8 1326.7 0.8X
+
+
+================================================================================================
+Section B - W=201 scaling
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=201, 1M rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=201 3314 3362 88 0.3 3313.9 1.0X
+SUM segtree (default) W=201 1447 1475 21 0.7 1447.4 2.3X
+
+
+================================================================================================
+Section B - W=4001 scaling (stress, + bs=256 cross-block)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=4001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=4001 119582 124613 6122 0.0 59791.1 1.0X
+SUM segtree (default) W=4001 3828 3863 33 0.5 1914.2 31.2X
+SUM segtree (blockSize=256) W=4001 93752 93810 86 0.0 46875.9 1.3X
+
+
+================================================================================================
+Section F - spill regression guard (String, stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+MAX String spill guard, W=1001, 1M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+--------------------------------------------------------------------------------------------------------------------------------
+MAX naive (String) 55976 56088 126 0.0 55976.4 1.0X
+MAX segtree default (String) 2822 2849 32 0.4 2821.9 19.8X
+
+
+================================================================================================
+Section C - N-sweep small (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM N-sweep (segtree-only), W=1001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=2M 3349 3387 34 0.6 1674.7 1.0X
+
+
+================================================================================================
+Section C - N-sweep mid (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM N-sweep (segtree-only), W=1001, 8M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=8M 15536 15568 56 0.5 1941.9 1.0X
+
+
+================================================================================================
+Section C - N-sweep large (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 25.0.3+9-LTS on Linux 6.17.0-1010-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM N-sweep (segtree-only), W=1001, 16M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+-------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=16M 36692 36765 79 0.4 2293.3 1.0X
+
+
diff --git a/sql/core/benchmarks/WindowBenchmark-results.txt b/sql/core/benchmarks/WindowBenchmark-results.txt
new file mode 100644
index 000000000000..bb50c20b5921
--- /dev/null
+++ b/sql/core/benchmarks/WindowBenchmark-results.txt
@@ -0,0 +1,172 @@
+================================================================================================
+Section A - MIN (non-invertible)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+MIN sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+MIN naive (current, baseline) 4350 4380 30 0.1 16592.3 1.0X
+MIN segtree (default) 433 442 10 0.6 1652.8 10.0X
+MIN segtree (blockSize=256) 2259 2266 9 0.1 8617.9 1.9X
+
+
+================================================================================================
+Section A - MAX (non-invertible)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+MAX sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+MAX naive (current, baseline) 4113 4124 21 0.1 15690.6 1.0X
+MAX segtree (default) 415 420 3 0.6 1582.0 9.9X
+MAX segtree (blockSize=256) 2279 2288 9 0.1 8692.0 1.8X
+
+
+================================================================================================
+Section A - SUM (Spark has no inverse; full recompute)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive (current, baseline) 4828 4870 61 0.1 18418.4 1.0X
+SUM segtree (default) 431 436 6 0.6 1643.9 11.2X
+SUM segtree (blockSize=256) 2286 2295 13 0.1 8720.2 2.1X
+
+
+================================================================================================
+Section A - COUNT
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+COUNT sliding window, W=1001, 256K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+COUNT naive (current, baseline) 3867 3883 13 0.1 14751.9 1.0X
+COUNT segtree (default) 415 417 3 0.6 1582.8 9.3X
+COUNT segtree (blockSize=256) 2216 2220 7 0.1 8453.8 1.7X
+
+
+================================================================================================
+Section A - AVG (multi-buffer)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+AVG sliding window, W=1001, 192K rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+AVG naive (current, baseline) 4483 4495 14 0.0 22804.2 1.0X
+AVG segtree (default) 426 435 8 0.5 2165.9 10.5X
+AVG segtree (blockSize=256) 1551 1566 18 0.1 7890.5 2.9X
+
+
+================================================================================================
+Section A - STDDEV_SAMP (multi-buffer, stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+STDDEV_SAMP sliding window, W=1001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+STDDEV_SAMP naive (current, baseline) 97349 97405 59 0.0 48674.4 1.0X
+STDDEV_SAMP segtree (default) 5628 5631 3 0.4 2813.9 17.3X
+STDDEV_SAMP segtree (blockSize=256) 113685 113987 289 0.0 56842.3 0.9X
+
+
+================================================================================================
+Section B - W=10 scaling (stress: Pareto loss zone)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=11, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=11 1003 1012 8 2.0 501.6 1.0X
+SUM segtree (default) W=11 2378 2384 9 0.8 1188.8 0.4X
+
+
+================================================================================================
+Section B - W=50 scaling (stress: Pareto loss zone)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=51, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=51 2377 2395 24 0.8 1188.5 1.0X
+SUM segtree (default) W=51 2667 2683 17 0.7 1333.7 0.9X
+
+
+================================================================================================
+Section B - W=201 scaling
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=201, 1M rows: Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=201 3588 3620 39 0.3 3588.1 1.0X
+SUM segtree (default) W=201 1465 1474 10 0.7 1465.1 2.4X
+
+
+================================================================================================
+Section B - W=4001 scaling (stress, + bs=256 cross-block)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM scaling, W=4001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------
+SUM naive W=4001 129657 129727 67 0.0 64828.5 1.0X
+SUM segtree (default) W=4001 3834 3858 24 0.5 1916.8 33.8X
+SUM segtree (blockSize=256) W=4001 110788 111483 748 0.0 55393.8 1.2X
+
+
+================================================================================================
+Section F - spill regression guard (String, stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+MAX String spill guard, W=1001, 1M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+--------------------------------------------------------------------------------------------------------------------------------
+MAX naive (String) 88459 88494 43 0.0 88459.0 1.0X
+MAX segtree default (String) 3213 3227 21 0.3 3213.4 27.5X
+
+
+================================================================================================
+Section C - N-sweep small (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM N-sweep (segtree-only), W=1001, 2M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=2M 3314 3319 6 0.6 1657.1 1.0X
+
+
+================================================================================================
+Section C - N-sweep mid (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM N-sweep (segtree-only), W=1001, 8M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=8M 15970 16004 42 0.5 1996.3 1.0X
+
+
+================================================================================================
+Section C - N-sweep large (stress)
+================================================================================================
+
+OpenJDK 64-Bit Server VM 17.0.19+10-LTS on Linux 6.17.0-1011-azure
+Intel(R) Xeon(R) Platinum 8370C CPU @ 2.80GHz
+SUM N-sweep (segtree-only), W=1001, 16M rows (stress): Best Time(ms) Avg Time(ms) Stdev(ms) Rate(M/s) Per Row(ns) Relative
+-------------------------------------------------------------------------------------------------------------------------------------
+SUM segtree (default) N=16M 38877 38998 121 0.4 2429.8 1.0X
+
+
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/AggregateProcessor.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/AggregateProcessor.scala
index 63934944c21a..4e87d8a8a49e 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/AggregateProcessor.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/AggregateProcessor.scala
@@ -181,7 +181,27 @@ private[window] final class AggregateProcessor(
}
/** Evaluate buffer. */
- def evaluate(target: InternalRow): Unit = {
- evaluateProjection.target(target)(buffer)
+ def evaluate(target: InternalRow): Unit = evaluate(buffer, target)
+
+ /**
+ * Evaluate using an arbitrary `source` buffer (e.g. a segment-tree query
+ * result) instead of the internal one.
+ *
+ * '''Contract''': `source` must share this processor's internal
+ * `aggBufferAttributes` layout -- same field count, same dataTypes, in the
+ * same order. The segment-tree path enforces this upstream in
+ * `WindowEvaluatorFactoryBase.eligibleForSegTree`, which restricts eligible
+ * functions to `WindowSegmentTree.EligibleAggregates` (an explicit allowlist
+ * of `DeclarativeAggregate` classes). The same `functions: Array[DeclarativeAggregate]`
+ * drives both this processor's `bufferSchema` and `WindowSegmentTree`'s
+ * internal buffer schema. The contract is invisible at the call site and
+ * easy to break from either end; the cheap field-count `assert` below
+ * surfaces drift loudly instead of producing silently garbled output.
+ */
+ private[window] def evaluate(source: InternalRow, target: InternalRow): Unit = {
+ assert(source.numFields == bufferSchema.length,
+ s"source buffer has ${source.numFields} fields, " +
+ s"expected ${bufferSchema.length} to match aggBufferAttributes layout")
+ evaluateProjection.target(target)(source)
}
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionFrame.scala
new file mode 100644
index 000000000000..51648e31e349
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionFrame.scala
@@ -0,0 +1,294 @@
+/*
+ * 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.window
+
+import org.apache.spark.TaskContext
+import org.apache.spark.memory.TaskMemoryManager
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, Expression, FrameType, MutableProjection, RangeFrame, RowFrame, UnsafeRow}
+import org.apache.spark.sql.catalyst.expressions.aggregate.DeclarativeAggregate
+import org.apache.spark.sql.execution.ExternalAppendOnlyUnsafeRowArray
+import org.apache.spark.sql.execution.metric.SQLMetric
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Moving-frame window function frame backed by [[WindowSegmentTree]]. Produces
+ * the same outputs as [[SlidingWindowFunctionFrame]] for RowFrame or
+ * single-column RangeFrame moving frames whose aggregates are all
+ * [[DeclarativeAggregate]] with no FILTER/DISTINCT. For partitions below
+ * `spark.sql.window.segmentTree.minPartitionRows`, delegates to a wrapped
+ * [[SlidingWindowFunctionFrame]]. Under RANGE, two forward-only cursors
+ * (`lowerIter` / `upperIter`) advance the bounds in O(n) total; the segtree
+ * answers `[lowerBound, upperBound)` in O(log n).
+ *
+ * @note Not thread-safe.
+ */
+private[window] final class SegmentTreeWindowFunctionFrame(
+ target: InternalRow,
+ processor: AggregateProcessor,
+ functions: Array[DeclarativeAggregate],
+ inputSchema: Seq[Attribute],
+ frameType: FrameType,
+ lbound: BoundOrdering,
+ ubound: BoundOrdering,
+ newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection,
+ conf: SQLConf,
+ maxCachedBlocks: Option[Int],
+ taskMemoryManager: TaskMemoryManager,
+ numSegmentTreeFrames: Option[SQLMetric] = None,
+ numSegmentTreeFallbackFrames: Option[SQLMetric] = None)
+ extends WindowFunctionFrame with AutoCloseable {
+
+ require(frameType == RowFrame || frameType == RangeFrame,
+ s"SegmentTreeWindowFunctionFrame supports RowFrame or RangeFrame, got $frameType")
+
+ private[this] var fallback: SlidingWindowFunctionFrame = _
+ private[this] var tree: WindowSegmentTree = _
+
+ /**
+ * Allocate a fresh fallback sliding-window frame. Called lazily from
+ * `prepare()` on the small-partition path. Factored out for testability
+ * (subclasses can inject a throwing fallback for prepare-failure tests).
+ */
+ private[window] def newFallback(): SlidingWindowFunctionFrame =
+ new SlidingWindowFunctionFrame(target, processor, lbound, ubound)
+
+ /** Test hook: whether the fallback frame has been lazily allocated. */
+ private[window] def fallbackAllocated: Boolean = fallback != null
+
+ // ---- RowFrame-only driver state ----
+ // `boundIter` advances `upperBound` one row at a time; the lower bound is
+ // pure index arithmetic under RowFrame.
+ private[this] var boundIter: Iterator[UnsafeRow] = _
+ private[this] var nextRow: UnsafeRow = _
+
+ // ---- RangeFrame-only driver state ----
+ // Two cursors over `rowArray`; `lowerRow` / `upperRow` hold the buffered
+ // head of each cursor, pre-fetched in `prepare` so
+ // `RangeBoundOrdering.compare` is never called with a null row on round 0.
+ //
+ // Spill-safety invariant: when `rowArray` spills, its iterator reuses a
+ // single `UnsafeRow` whose pointer is rebound on each `next()`. Tolerated
+ // here because the cursor is **read-before-advance**: each `writeRange`
+ // iteration reads `lowerRow` / `upperRow` for comparison before calling
+ // `getNextOrNull(...)`. DO NOT cache a historical row into a separate
+ // field without an explicit `.copy()`; the shared reusable UnsafeRow
+ // would silently mutate.
+ private[this] var lowerIter: Iterator[UnsafeRow] = _
+ private[this] var upperIter: Iterator[UnsafeRow] = _
+ private[this] var lowerRow: UnsafeRow = _
+ private[this] var upperRow: UnsafeRow = _
+
+ // Shared endpoints: monotone across `write()` calls within a partition.
+ private[this] var lowerBound: Int = 0
+ private[this] var upperBound: Int = 0
+
+ /**
+ * Runtime dispatch flag: when `true`, `write()`, `currentLowerBound()`, and
+ * `currentUpperBound()` delegate to the wrapped [[SlidingWindowFunctionFrame]]
+ * (small-partition path). Set by `prepare()` based on partition size vs.
+ * `spark.sql.window.segmentTree.minPartitionRows`.
+ */
+ private[window] var fallbackUsed: Boolean = false
+
+ // Register close() once per frame instance so the tree's block cache and
+ // any open row-array iterators are released when the task completes.
+ // Registering here (vs. inside the factory closure) avoids duplicate
+ // listeners when the factory is invoked multiple times per task.
+ {
+ val tc = TaskContext.get()
+ if (tc != null) {
+ tc.addTaskCompletionListener[Unit](_ => close())
+ }
+ }
+
+ override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
+ // INVARIANT: sole call site is `WindowPartitionEvaluator.fetchNextPartition`;
+ // `prepare` is called exactly once per (partition, frame). No mid-partition
+ // retry is supported.
+ //
+ // Ordering (matches `OffsetWindowFunctionFrameBase.prepare`): commit to a
+ // path first (allocate+prepare the fallback, or build the tree), and only
+ // then reset the opposite path's cached state. If the committing action
+ // throws, the frame stays in its previous (possibly stale but
+ // self-consistent) state; a thrown `prepare` is fatal for the task.
+ if (rows.length < conf.windowSegmentTreeMinPartitionRows) {
+ if (fallback == null) fallback = newFallback()
+ fallback.prepare(rows) // may throw -- frame unchanged on failure
+ // Commit fallback path only after `prepare` succeeded.
+ resetSegtreeState()
+ fallbackUsed = true
+ numSegmentTreeFallbackFrames.foreach(_ += 1)
+ return
+ }
+ // Segtree path. Drop any retained fallback so its row-copy buffer from a
+ // prior small partition is GC-eligible; next small partition reallocates
+ // via `newFallback()`.
+ fallback = null
+ resetSegtreeState()
+ fallbackUsed = false
+ tree = new WindowSegmentTree(
+ functions,
+ inputSchema,
+ newMutableProjection,
+ fanout = conf.windowSegmentTreeFanout,
+ blockSize = conf.windowSegmentTreeBlockSize,
+ maxCachedBlocks = maxCachedBlocks,
+ taskMemoryManager = taskMemoryManager)
+ // Tree holds a reference to `rows` (caller-owned); no extra copy.
+ tree.build(rows)
+ // Count only on the successful segtree path: if `tree.build` throws,
+ // the counter is not bumped.
+ numSegmentTreeFrames.foreach(_ += 1)
+ frameType match {
+ case RowFrame =>
+ boundIter = rows.generateIterator()
+ nextRow = WindowFunctionFrame.getNextOrNull(boundIter)
+ case RangeFrame =>
+ lowerIter = rows.generateIterator()
+ upperIter = rows.generateIterator()
+ // Pre-seed cursor heads so `RangeBoundOrdering.compare` never
+ // dereferences null on round 0. Either may be null if `rows` is
+ // empty; the advance loops' `!= null` / `< upperBound` guards
+ // handle that.
+ lowerRow = WindowFunctionFrame.getNextOrNull(lowerIter)
+ upperRow = WindowFunctionFrame.getNextOrNull(upperIter)
+ }
+ }
+
+ /** Release prior segtree-path cached state (tree + row-cursors + bounds). */
+ private def resetSegtreeState(): Unit = {
+ if (tree != null) {
+ tree.close()
+ tree = null
+ }
+ closeIters()
+ nextRow = null
+ lowerRow = null
+ upperRow = null
+ lowerBound = 0
+ upperBound = 0
+ }
+
+ override def write(index: Int, current: InternalRow): Unit = {
+ if (fallbackUsed) {
+ fallback.write(index, current)
+ return
+ }
+ frameType match {
+ case RowFrame => writeRow(index, current)
+ case RangeFrame => writeRange(index, current)
+ }
+ }
+
+ // `writeRow`/`writeRange` mirror the `(lowerBound, upperBound)` monotone
+ // cursor invariant of `SlidingWindowFunctionFrame.write`, but run
+ // admit-then-drop (no buffer to maintain) instead of drop-then-admit.
+ // Any future fix to Sliding's boundary semantics must be mirrored here;
+ // equivalence is guarded by `SegmentTreeWindowFunctionSuite` flag-on/off
+ // tests (`checkRangeEquivalence`, `feature flag off ...`, fallback tests)
+ // which compare against the Sliding baseline.
+ private def writeRow(index: Int, current: InternalRow): Unit = {
+ var boundsChanged = index == 0
+
+ // admit loop: extend upperBound; if a candidate is already below the
+ // lower bound, advance lowerBound in lock-step to preserve invariant
+ // (0 <= lowerBound <= upperBound <= tree.size).
+ while (nextRow != null &&
+ ubound.compare(nextRow, upperBound, current, index) <= 0) {
+ if (lbound.compare(nextRow, lowerBound, current, index) < 0) {
+ lowerBound += 1
+ }
+ nextRow = WindowFunctionFrame.getNextOrNull(boundIter)
+ upperBound += 1
+ boundsChanged = true
+ }
+ // drop loop: advance lowerBound to the frame's left edge. RowFrame's
+ // `lbound.compare` is pure index arithmetic so the input row is unread;
+ // `lowerBound < upperBound` is the second defense.
+ while (lowerBound < upperBound &&
+ lbound.compare(null, lowerBound, current, index) < 0) {
+ lowerBound += 1
+ boundsChanged = true
+ }
+
+ if (boundsChanged) {
+ tree.queryInto(lowerBound, upperBound, processor, target)
+ }
+ }
+
+ private def writeRange(index: Int, current: InternalRow): Unit = {
+ var boundsChanged = index == 0
+
+ // admit loop (upper edge). `RangeBoundOrdering.compare` ignores its index
+ // arguments; we pass `upperBound` for API symmetry with RowBoundOrdering.
+ while (upperRow != null &&
+ ubound.compare(upperRow, upperBound, current, index) <= 0) {
+ upperBound += 1
+ upperRow = WindowFunctionFrame.getNextOrNull(upperIter)
+ boundsChanged = true
+ }
+
+ // drop loop (lower edge): strict `< 0`, guarded by
+ // `lowerBound < upperBound` so drop never overruns admit -- this also
+ // ensures `lowerRow` is non-null when reached (if the iterator
+ // exhausts, `lowerBound` has already caught up to `numRows >=
+ // upperBound`, so the guard stops the loop first).
+ while (lowerBound < upperBound &&
+ lbound.compare(lowerRow, lowerBound, current, index) < 0) {
+ lowerBound += 1
+ lowerRow = WindowFunctionFrame.getNextOrNull(lowerIter)
+ boundsChanged = true
+ }
+
+ if (boundsChanged) {
+ // Empty frame is handled inside `queryInto`, which initializes the
+ // processor and emits the aggregate's zero/identity value -- same as
+ // `SlidingWindowFunctionFrame` with an empty buffer.
+ tree.queryInto(lowerBound, upperBound, processor, target)
+ }
+ }
+
+ override def currentLowerBound(): Int =
+ if (fallbackUsed) fallback.currentLowerBound() else lowerBound
+
+ override def currentUpperBound(): Int =
+ if (fallbackUsed) fallback.currentUpperBound() else upperBound
+
+ /**
+ * Drop references to open rowArray iterators. Idempotent.
+ *
+ * Note: `ExternalAppendOnlyUnsafeRowArrayIterator.closeIfNeeded()` is
+ * `protected`, so we cannot invoke it. Spark's own `SlidingWindowFunctionFrame`
+ * also does not close its iterator; the backing `UnsafeExternalSorter` is
+ * released by the enclosing `WindowExec`'s `TaskCompletionListener`.
+ */
+ private def closeIters(): Unit = {
+ boundIter = null
+ lowerIter = null
+ upperIter = null
+ }
+
+ override def close(): Unit = {
+ if (tree != null) {
+ tree.close()
+ tree = null
+ }
+ closeIters()
+ }
+}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactory.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactory.scala
index c4b20d4b7c7d..09c207ee603a 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactory.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactory.scala
@@ -30,9 +30,15 @@ class WindowEvaluatorFactory(
val partitionSpec: Seq[Expression],
val orderSpec: Seq[SortOrder],
val childOutput: Seq[Attribute],
- val spillSize: SQLMetric)
+ val spillSize: SQLMetric,
+ segmentTreeFrames: SQLMetric,
+ segmentTreeFallbackFrames: SQLMetric)
extends PartitionEvaluatorFactory[InternalRow, InternalRow] with WindowEvaluatorFactoryBase {
+ override def numSegmentTreeFrames: Option[SQLMetric] = Some(segmentTreeFrames)
+ override def numSegmentTreeFallbackFrames: Option[SQLMetric] =
+ Some(segmentTreeFallbackFrames)
+
override def createEvaluator(): PartitionEvaluator[InternalRow, InternalRow] = {
new WindowPartitionEvaluator()
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactoryBase.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactoryBase.scala
index 9930c4a8963a..cebcfd05bd24 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactoryBase.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowEvaluatorFactoryBase.scala
@@ -20,10 +20,10 @@ package org.apache.spark.sql.execution.window
import scala.collection.mutable
import scala.collection.mutable.ArrayBuffer
-import org.apache.spark.SparkException
+import org.apache.spark.{SparkException, TaskContext}
import org.apache.spark.sql.catalyst.InternalRow
import org.apache.spark.sql.catalyst.expressions.{Add, AggregateWindowFunction, Ascending, Attribute, BoundReference, CurrentRow, DateAdd, DateAddYMInterval, DecimalAddNoOverflowCheck, Descending, Expression, ExtractANSIIntervalDays, FrameLessOffsetWindowFunction, FrameType, IdentityProjection, IntegerLiteral, MutableProjection, NamedExpression, OffsetWindowFunction, PythonFuncExpression, RangeFrame, RowFrame, RowOrdering, SortOrder, SpecifiedWindowFrame, TimestampAddInterval, TimestampA [...]
-import org.apache.spark.sql.catalyst.expressions.aggregate.AggregateExpression
+import org.apache.spark.sql.catalyst.expressions.aggregate.{AggregateExpression, DeclarativeAggregate}
import org.apache.spark.sql.execution.metric.SQLMetric
import org.apache.spark.sql.internal.SQLConf
import org.apache.spark.sql.types.{CalendarIntervalType, DateType, DayTimeIntervalType, DecimalType, IntegerType, TimestampNTZType, TimestampType, YearMonthIntervalType}
@@ -36,6 +36,14 @@ trait WindowEvaluatorFactoryBase {
def orderSpec: Seq[SortOrder]
def childOutput: Seq[Attribute]
def spillSize: SQLMetric
+ /**
+ * Counters for [[SegmentTreeWindowFunctionFrame]] observability. Default
+ * `None` means the subclass does not integrate with the segment-tree frame
+ * path (e.g. [[org.apache.spark.sql.execution.python.ArrowWindowPythonEvaluatorFactory]]);
+ * only [[WindowEvaluatorFactory]] wires them.
+ */
+ def numSegmentTreeFrames: Option[SQLMetric] = None
+ def numSegmentTreeFallbackFrames: Option[SQLMetric] = None
/**
* Create the resulting projection.
@@ -191,13 +199,18 @@ trait WindowEvaluatorFactoryBase {
// in a single Window physical node. Therefore, we can assume no SQL aggregation
// functions if Pandas UDF exists. In the future, we might mix Pandas UDF and SQL
// aggregation function in a single physical node.
+ val aggFilters: Array[Option[Expression]] = expressions.map {
+ case WindowExpression(ae: AggregateExpression, _) => ae.filter
+ case _ => None
+ }.toArray
+ // Keep as `def` (by-name): the FRAME_LESS_OFFSET / UNBOUNDED_OFFSET /
+ // UNBOUNDED_PRECEDING_OFFSET branches do not read `processor`. Eager
+ // `val` construction would invoke `AggregateProcessor.apply` on
+ // Lag / Lead / NthValue and throw
+ // `INTERNAL_ERROR: Unsupported aggregate function`.
def processor = if (functions.exists(_.isInstanceOf[PythonFuncExpression])) {
null
} else {
- val aggFilters = expressions.map {
- case WindowExpression(ae: AggregateExpression, _) => ae.filter
- case _ => None
- }.toArray
AggregateProcessor(
functions,
ordinal,
@@ -206,6 +219,8 @@ trait WindowEvaluatorFactoryBase {
MutableProjection.create(expressions, schema),
aggFilters)
}
+ val conf = SQLConf.get
+ val blockSize = conf.windowSegmentTreeBlockSize
// Create the factory to produce WindowFunctionFrame.
val factory = key match {
@@ -275,12 +290,44 @@ trait WindowEvaluatorFactoryBase {
// Moving Frame.
case ("AGGREGATE", frameType, lower, upper, _) =>
- target: InternalRow => {
- new SlidingWindowFunctionFrame(
- target,
- processor,
- createBoundOrdering(frameType, lower, timeZone),
- createBoundOrdering(frameType, upper, timeZone))
+ if (eligibleForSegTree(functions, aggFilters, frameType, conf)) {
+ val segFns = functions.map(_.asInstanceOf[DeclarativeAggregate])
+ val cacheHint = estimateMaxCachedBlocks(lower, upper, frameType, blockSize)
+ target: InternalRow => {
+ // Task-completion listener registration lives inside the frame
+ // constructor (one per frame instance) to avoid duplicates when
+ // this closure fires multiple times per task. `TaskContext.get()`
+ // is only called at task-execution time, never at driver planning.
+ val tc = TaskContext.get()
+ if (tc == null) {
+ throw SparkException.internalError(
+ "WindowEvaluatorFactoryBase.segTreeFrameFactory requires " +
+ "an active TaskContext")
+ }
+ val tmm = tc.taskMemoryManager()
+ new SegmentTreeWindowFunctionFrame(
+ target,
+ processor,
+ segFns,
+ childOutput,
+ frameType,
+ createBoundOrdering(frameType, lower, timeZone),
+ createBoundOrdering(frameType, upper, timeZone),
+ (e, s) => MutableProjection.create(e, s),
+ conf,
+ cacheHint,
+ tmm,
+ numSegmentTreeFrames,
+ numSegmentTreeFallbackFrames)
+ }
+ } else {
+ target: InternalRow => {
+ new SlidingWindowFunctionFrame(
+ target,
+ processor,
+ createBoundOrdering(frameType, lower, timeZone),
+ createBoundOrdering(frameType, upper, timeZone))
+ }
}
case _ =>
@@ -295,4 +342,60 @@ trait WindowEvaluatorFactoryBase {
}
}
+ /**
+ * Segment-tree path eligibility. The tree relies on
+ * `DeclarativeAggregate.mergeExpressions`, which [[AggregateWindowFunction]]s
+ * (NthValue, NTile, Rank, RowNumber, NullIndex) refuse via
+ * `mergeUnsupportedByWindowFunctionError`: they extend DeclarativeAggregate
+ * but are NOT merge-capable. Normal aggregate window expressions reach this
+ * code as the inner DeclarativeAggregate unwrapped from
+ * [[AggregateExpression]] (see `windowFrameExpressionFactoryPairs.collect`).
+ *
+ * DISTINCT aggregate window expressions are already rejected earlier in
+ * analysis by `WindowResolution.checkWindowFunction`
+ * (error class `DISTINCT_WINDOW_FUNCTION_UNSUPPORTED`), so no explicit
+ * `isDistinct` gate is needed here.
+ */
+ private def eligibleForSegTree(
+ functions: Array[Expression],
+ filters: Array[Option[Expression]],
+ frameType: FrameType,
+ conf: SQLConf): Boolean = {
+ // RANGE accepted only for single-column order specs. Multi-column RANGE
+ // with non-zero offset is already rejected by `createBoundOrdering`, so
+ // gating here on `orderSpec.size == 1` matches the Sliding-path invariant.
+ val frameTypeOk = frameType match {
+ case RowFrame => true
+ case RangeFrame => orderSpec.size == 1
+ }
+ conf.windowSegmentTreeEnabled &&
+ frameTypeOk &&
+ filters.forall(_.isEmpty) &&
+ functions.forall(WindowSegmentTree.isEligible)
+ }
+
+ private def estimateMaxCachedBlocks(
+ lower: Expression,
+ upper: Expression,
+ frameType: FrameType,
+ blockSize: Int): Option[Int] = {
+ // Reached via the moving-frame branch after `eligibleForSegTree`. Under
+ // RANGE the frame width is data-dependent (defined by order-key distance,
+ // not row count), so no static width inference is possible; fall back to
+ // a default budget and rely on the runtime LRU + TMM spiller.
+ assert(frameType == RowFrame || frameType == RangeFrame,
+ s"estimateMaxCachedBlocks expects RowFrame or RangeFrame, got $frameType")
+ if (frameType == RangeFrame) {
+ return Some(8)
+ }
+ val w: Option[Int] = (lower, upper) match {
+ case (CurrentRow, CurrentRow) => Some(1)
+ case (IntegerLiteral(lo), IntegerLiteral(hi)) => Some(math.abs(hi - lo) + 1)
+ case (CurrentRow, IntegerLiteral(hi)) => Some(math.abs(hi) + 1)
+ case (IntegerLiteral(lo), CurrentRow) => Some(math.abs(lo) + 1)
+ case _ => None
+ }
+ w.map(ww => math.ceil(ww.toDouble / blockSize).toInt + 2).orElse(Some(8))
+ }
+
}
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExec.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExec.scala
index 9ecd1c587a75..e80565c96139 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExec.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExec.scala
@@ -91,7 +91,12 @@ case class WindowExec(
child: SparkPlan)
extends WindowExecBase {
override lazy val metrics: Map[String, SQLMetric] = Map(
- "spillSize" -> SQLMetrics.createSizeMetric(sparkContext, "spill size")
+ "spillSize" -> SQLMetrics.createSizeMetric(sparkContext, "spill size"),
+ "numSegmentTreeFrames" ->
+ SQLMetrics.createMetric(sparkContext, "number of segment-tree frames prepared"),
+ "numSegmentTreeFallbackFrames" ->
+ SQLMetrics.createMetric(sparkContext,
+ "number of segment-tree fallback frames prepared")
)
protected override def doExecute(): RDD[InternalRow] = {
@@ -101,7 +106,9 @@ case class WindowExec(
partitionSpec,
orderSpec,
child.output,
- longMetric("spillSize"))
+ longMetric("spillSize"),
+ longMetric("numSegmentTreeFrames"),
+ longMetric("numSegmentTreeFallbackFrames"))
// Start processing.
if (conf.usePartitionEvaluator) {
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
new file mode 100644
index 000000000000..27a773636134
--- /dev/null
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowSegmentTree.scala
@@ -0,0 +1,614 @@
+/*
+ * 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.window
+
+import java.util.{LinkedHashMap => JLinkedHashMap, Map => JMap}
+
+import scala.collection.mutable
+
+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.errors.QueryExecutionErrors
+import org.apache.spark.sql.execution.ExternalAppendOnlyUnsafeRowArray
+import org.apache.spark.sql.types.DataType
+import org.apache.spark.util.ArrayImplicits._
+
+/**
+ * Block-chunked segment tree for moving-frame window aggregates. Partitions are
+ * split into blocks of `blockSize` rows; each block has its own small segtree
+ * (fanout `F`, height `h`). Block roots stay resident; internal nodes are cached
+ * in an LRU keyed by block index. Queries cost O(log W).
+ *
+ * Memory accounting invariants:
+ * - I1: `SegTreeSpiller.spill()` MUST NOT call `acquireMemory` on its own
+ * consumer (would deadlock TMM's consumer-priority sort). All acquires
+ * happen on the hot path ([[ensureBlockLevels]]).
+ * - I2: `spill(_, trigger)` returns 0 when `trigger eq this` (self-trigger
+ * short-circuit) to prevent re-entrant eviction.
+ * - I3: LRU `removeEldestEntry` is disabled; eviction is driven explicitly
+ * from [[ensureBlockLevels]] or [[SegTreeSpiller.spill]].
+ * - I4: Every successful [[acquireBlockMemory]] is paired with exactly one
+ * [[releaseBlockMemory]]. [[close]] is idempotent.
+ * - I5: Per-block bytes are a conservative upper bound (full block, 16 B/field).
+ * - I8: If `rowArray` already spilled to disk, `spill` returns 0 (rebuild
+ * would O(blockStart)-scan the spill file).
+ *
+ * @note Instances are not thread-safe.
+ */
+private[window] class WindowSegmentTree(
+ functions: Array[DeclarativeAggregate],
+ inputSchema: Seq[Attribute],
+ newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection,
+ fanout: Int = WindowSegmentTree.DefaultFanout,
+ blockSize: Int = WindowSegmentTree.DefaultBlockSize,
+ maxCachedBlocks: Option[Int] = None,
+ taskMemoryManager: TaskMemoryManager = null)
+ extends AutoCloseable {
+
+ require(fanout >= 2, s"fanout must be >= 2, got $fanout")
+ require(blockSize >= 1, s"blockSize must be >= 1, got $blockSize")
+ require(functions.nonEmpty, "WindowSegmentTree requires at least one aggregate function")
+ maxCachedBlocks.foreach { n =>
+ require(n >= 1, s"maxCachedBlocks must be >= 1 when specified, got $n")
+ }
+ require(taskMemoryManager != null,
+ "WindowSegmentTree requires a non-null TaskMemoryManager; " +
+ "in tests use `new TaskMemoryManager(new TestMemoryManager(conf), 0)`")
+
+ // ---------- Schemas & projections ----------
+
+ private val bufferAttrs: Seq[AttributeReference] =
+ functions.flatMap(_.aggBufferAttributes).toImmutableArraySeq
+ private val rightAttrs: Seq[AttributeReference] =
+ functions.flatMap(_.inputAggBufferAttributes).toImmutableArraySeq
+ private val bufferDataTypes: IndexedSeq[DataType] =
+ bufferAttrs.map(_.dataType).toIndexedSeq
+
+ private val initialValues: Seq[Expression] = functions.flatMap(_.initialValues).toIndexedSeq
+ private val updateExpressions: Seq[Expression] =
+ functions.flatMap(_.updateExpressions).toIndexedSeq
+ private val mergeExpressions: Seq[Expression] =
+ functions.flatMap(_.mergeExpressions).toIndexedSeq
+
+ private[this] val initProj: MutableProjection = newMutableProjection(initialValues, Nil)
+ private[this] val updateProj: MutableProjection =
+ newMutableProjection(updateExpressions, bufferAttrs ++ inputSchema)
+ private[this] val mergeProj: MutableProjection =
+ newMutableProjection(mergeExpressions, bufferAttrs ++ rightAttrs)
+
+ private[this] val joinedRow = new JoinedRow()
+
+ // ---------- State ----------
+
+ private var numRows: Int = 0
+ private var numBlocks: Int = 0
+ private var rowArray: ExternalAppendOnlyUnsafeRowArray = _
+ private var closed: Boolean = false
+
+ /**
+ * Always-resident per-block root aggregates: `blockAggregates(i)` =
+ * merged buffer over all rows in block i.
+ */
+ private var blockAggregates: Array[InternalRow] = Array.empty
+
+ /**
+ * Conservative byte width of one aggregate buffer row at 16 B/field:
+ * primitive `MutableValue` is 8 B, boxed references and object headers
+ * push the effective footprint higher. Tighter per-type sizing is out
+ * of scope; TaskMemoryManager remains the hard backstop via spill / OOM.
+ */
+ private val bufferWidthBytes: Long = {
+ val bytesPerField = 16L
+ math.max(1L, bufferDataTypes.size.toLong * bytesPerField)
+ }
+
+ /**
+ * Number of aggregate-buffer slots cached per block (see I5).
+ *
+ * Invariant: equals `sum over levels L of levels(L).length` for any block
+ * built by [[buildBlockLevels]]: level 0 holds `blockSize` leaves and each
+ * next level holds `ceil(prev / fanout)` parents until a single root
+ * remains. For `blockSize == 1` this is 1 (single leaf, no parents).
+ */
+ private val cachedSlotsPerBlock: Long = {
+ var n = blockSize.toLong
+ var sum = n
+ while (n > 1L) {
+ n = (n + fanout - 1) / fanout
+ sum += n
+ }
+ sum
+ }
+
+ /**
+ * Bytes accounted per cached block (see I5). Conservative: assumes every
+ * block is full; tail block leaves a small headroom.
+ */
+ private[this] val blockBytes: Long =
+ math.max(1L, cachedSlotsPerBlock * bufferWidthBytes)
+
+ /**
+ * `spans(L)` = number of leaves covered by a single node at level L. Depends
+ * only on fanout + blockSize, so precomputed once.
+ */
+ private val spans: Array[Int] = {
+ val maxLevel = {
+ var lvl = 0
+ var span = 1L
+ while (span < blockSize) { span *= fanout; lvl += 1 }
+ lvl
+ }
+ val arr = new Array[Int](maxLevel + 1)
+ var s = 1L
+ var i = 0
+ while (i <= maxLevel) {
+ arr(i) = if (s > Int.MaxValue) Int.MaxValue else s.toInt
+ s *= fanout
+ i += 1
+ }
+ arr
+ }
+
+ /**
+ * LRU cache of per-block internal node arrays. Key = blockIdx;
+ * value = `Array[Array[InternalRow]]` with levels(0..h). Auto-eviction
+ * via `removeEldestEntry` is disabled (I3) -- driven explicitly from
+ * [[ensureBlockLevels]] or [[SegTreeSpiller.spill]]. Each entry maps 1:1
+ * to one [[acquireBlockMemory]] accounting. Callers should pass a W-aware
+ * `maxCachedBlocks` like `ceil(W / blockSize) + 2`.
+ */
+ private val blockLevelsCache: JLinkedHashMap[Integer, Array[Array[InternalRow]]] =
+ new JLinkedHashMap[Integer, Array[Array[InternalRow]]](16, 0.75f, true) {
+ override def removeEldestEntry(
+ eldest: JMap.Entry[Integer, Array[Array[InternalRow]]]): Boolean = false
+ }
+
+ // ---------- Memory consumer ----------
+
+ /**
+ * Private MemoryConsumer tracking cached block levels under TMM. Heap-only
+ * (no Tungsten pages): uses [[MemoryConsumer.acquireMemory]] /
+ * [[MemoryConsumer.freeMemory]], which update the base class `used`
+ * AtomicLong so TMM's consumer-priority sort sees our pressure accurately.
+ *
+ * Hardcoded [[MemoryMode.ON_HEAP]] (not `tmm.getTungstenMemoryMode`): the
+ * cache holds plain JVM objects (`SpecificInternalRow` /
+ * `Array[Array[InternalRow]]`), never Tungsten pages. Under
+ * `spark.memory.offHeap.enabled=true`, enrolling as OFF_HEAP would let
+ * TMM pick us as a spill candidate for off-heap pressure; our `spill()`
+ * would then phantom-credit the off-heap pool without releasing any
+ * off-heap bytes, violating [[TaskMemoryManager#acquireExecutionMemory]]'s
+ * same-pool spill contract. Mirrors
+ * [[org.apache.spark.util.collection.Spillable]], which also hardcodes
+ * ON_HEAP for the same reason. Consequence under off-heap Tungsten: I8
+ * (below) degrades to a no-op because segtree and `rowArray` live in
+ * different pools -- a loss of optimization, not a correctness hazard.
+ *
+ * @note `spill()` MUST NOT call `acquireMemory` (see I1).
+ */
+ private final class SegTreeSpiller extends MemoryConsumer(
+ taskMemoryManager,
+ taskMemoryManager.pageSizeBytes(),
+ MemoryMode.ON_HEAP) {
+ override def spill(size: Long, trigger: MemoryConsumer): Long = {
+ // I2: self-trigger short-circuit (prevent re-entrant eviction).
+ if (trigger eq this) return 0L
+ // I8: rowArray already spilled -- evicting our cache is counter-productive
+ // (rebuild would O(blockStart)-scan the spill file). `spillSize > 0` is
+ // the available "has spilled" signal (UnsafeExternalSorter state is not
+ // public).
+ if (rowArray != null && rowArray.spillSize > 0) return 0L
+ evictUntil(size)
+ }
+ }
+
+ private[this] val spiller: SegTreeSpiller = new SegTreeSpiller
+
+ // ---------- Public API ----------
+
+ def size: Int = numRows
+
+ /**
+ * Build the tree against a caller-owned row array.
+ *
+ * Ownership: the tree holds a reference to `rows` for its lifetime but does
+ * NOT own it -- the caller (typically `WindowPartitionEvaluator.buffer`)
+ * manages `clear()` / lifetime at partition boundaries. `close()` drops
+ * the reference without mutating the array.
+ *
+ * Exception-safe: if aggregation throws, previously built state is preserved.
+ */
+ def build(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
+ // rows.length is Int by design; check guards against future widening.
+ val n = rows.length
+ if (n < 0) {
+ throw SparkException.internalError(
+ s"WindowSegmentTree cannot hold more than Int.MaxValue rows, got $n")
+ }
+ val nBlocks = if (n == 0) 0 else (n + blockSize - 1) / blockSize
+ val newBlockAggs = computeBlockAggregates(rows, n, nBlocks)
+
+ // Commit.
+ rowArray = rows
+ numRows = n
+ numBlocks = nBlocks
+ blockAggregates = newBlockAggs
+ // Rebuild invalidates cached block levels; release accounting first (I4).
+ releaseAllCachedBlocks()
+ }
+
+ /**
+ * Query [lo, hi) and directly evaluate the result via `processor.evaluate`
+ * into `target`. Uses an internal pre-allocated buffer so no per-call
+ * allocation is needed.
+ */
+ private[window] def queryInto(
+ lo: Int, hi: Int, processor: AggregateProcessor, target: InternalRow): Unit = {
+ query(lo, hi, internalQueryBuffer)
+ processor.evaluate(internalQueryBuffer, target)
+ }
+
+ private[this] val internalQueryBuffer: InternalRow = newBuffer()
+
+ def query(lo: Int, hi: Int, outBuffer: InternalRow): Unit = {
+ if (lo < 0 || hi > numRows || lo > hi) {
+ throw SparkException.internalError(
+ s"Invalid range [lo=$lo, hi=$hi) for size=$numRows")
+ }
+ // Reset outBuffer to identity only after bounds validation.
+ initProj.target(outBuffer)(InternalRow.empty)
+ if (lo == hi) return
+
+ val blo = lo / blockSize
+ val bhi = (hi - 1) / blockSize
+
+ if (blo == bhi) {
+ val blockStart = blo * blockSize
+ mergeBlockRange(blo, lo - blockStart, hi - blockStart, outBuffer)
+ } else {
+ // left partial
+ val loStart = blo * blockSize
+ val loBlockRows = math.min(blockSize, numRows - loStart)
+ mergeBlockRange(blo, lo - loStart, loBlockRows, outBuffer)
+ // full blocks
+ var b = blo + 1
+ while (b < bhi) {
+ mergeInto(outBuffer, blockAggregates(b))
+ b += 1
+ }
+ // right partial
+ val hiStart = bhi * blockSize
+ mergeBlockRange(bhi, 0, hi - hiStart, outBuffer)
+ }
+ }
+
+ /** Terminal: releases all state. Idempotent (I4). */
+ override def close(): Unit = {
+ if (closed) return
+ // Free all cached-block accounting before dropping references.
+ releaseAllCachedBlocks()
+ closeRowArray()
+ blockAggregates = Array.empty
+ numRows = 0
+ numBlocks = 0
+ closed = true
+ }
+
+ // ---------- Test hooks (package-private) ----------
+
+ private[window] def peekBlockCount: Int = numBlocks
+
+ private[window] def testOnlySpiller(): MemoryConsumer = spiller
+
+ /** Test-only accessor for the per-block memory accounting value. */
+ private[window] def peekBlockBytes: Long = blockBytes
+
+ /** NOTE: test-only; promotes block to MRU in the LRU cache as a side effect. */
+ private[window] def peekLevelSize(blockIdx: Int, level: Int): Int = {
+ val levels = ensureBlockLevels(blockIdx)
+ levels(level).length
+ }
+
+ /** NOTE: test-only; promotes block to MRU in the LRU cache as a side effect. */
+ private[window] def peekLevelCount(blockIdx: Int): Int = {
+ val levels = ensureBlockLevels(blockIdx)
+ levels.length
+ }
+
+ // ---------- Internals ----------
+
+ private def computeBlockAggregates(
+ array: ExternalAppendOnlyUnsafeRowArray,
+ n: Int,
+ nBlocks: Int): Array[InternalRow] = {
+ if (n == 0) return Array.empty
+ val result = new Array[InternalRow](nBlocks)
+ val iter = array.generateIterator()
+ var b = 0
+ while (b < nBlocks) {
+ val buf = newBuffer()
+ initProj.target(buf)(InternalRow.empty)
+ val start = b * blockSize
+ val end = math.min(start + blockSize, n)
+ var i = start
+ while (i < end) {
+ if (!iter.hasNext) {
+ throw SparkException.internalError("rowArray iterator exhausted unexpectedly")
+ }
+ val row = iter.next()
+ updateProj.target(buf)(joinedRow(buf, row))
+ i += 1
+ }
+ result(b) = buf
+ b += 1
+ }
+ result
+ }
+
+ /** Merge `src` buffer into `dst` buffer using mergeProj. */
+ private def mergeInto(dst: InternalRow, src: InternalRow): Unit = {
+ mergeProj.target(dst)(joinedRow(dst, src))
+ }
+
+ private def newBuffer(): InternalRow =
+ new SpecificInternalRow(bufferDataTypes)
+
+ /** Merge the given leaf range [lo, hi) inside `blockIdx` into `out`. */
+ private def mergeBlockRange(
+ blockIdx: Int, lo: Int, hi: Int, out: InternalRow): Unit = {
+ if (lo >= hi) return
+ val levels = ensureBlockLevels(blockIdx)
+ val blockRows = levels(0).length
+ val topLevel = levels.length - 1
+ queryDescend(levels, blockRows, topLevel, 0, lo, hi, out)
+ }
+
+ /**
+ * Descend the (per-block) segment tree merging any node fully contained
+ * in [queryLo, queryHi) into `out`. A node at (level L, index idx) covers
+ * leaves `[idx * span, min((idx+1)*span, blockRows))` where span = F^L.
+ */
+ private def queryDescend(
+ levels: Array[Array[InternalRow]],
+ blockRows: Int,
+ level: Int,
+ idx: Int,
+ queryLo: Int,
+ queryHi: Int,
+ out: InternalRow): Unit = {
+ val span = spans(level)
+ val nodeLo = idx * span
+ val nodeHi = math.min(nodeLo + span, blockRows)
+ if (queryLo >= nodeHi || queryHi <= nodeLo) return
+ if (queryLo <= nodeLo && nodeHi <= queryHi) {
+ mergeInto(out, levels(level)(idx))
+ return
+ }
+ val childLevel = level - 1
+ val childLevelSize = levels(childLevel).length
+ var c = 0
+ while (c < fanout) {
+ val childIdx = idx * fanout + c
+ if (childIdx < childLevelSize) {
+ queryDescend(levels, blockRows, childLevel, childIdx, queryLo, queryHi, out)
+ }
+ c += 1
+ }
+ }
+
+ /**
+ * Build (or fetch from LRU) the full per-block levels array.
+ * Protocol: acquire memory -> build -> cache. Eviction on capacity
+ * overflow or on TMM spill request.
+ */
+ private def ensureBlockLevels(blockIdx: Int): Array[Array[InternalRow]] = {
+ val cached = blockLevelsCache.get(Integer.valueOf(blockIdx))
+ if (cached != null) return cached
+
+ // Enforce LRU capacity before building a new entry (I3).
+ val cap = maxCachedBlocks.getOrElse(Int.MaxValue)
+ while (blockLevelsCache.size() >= cap) {
+ if (!evictEldest()) return throwCacheEvictFailed(blockIdx)
+ }
+
+ // Acquire accounting; on partial grant, try one manual evict-and-retry.
+ if (!acquireBlockMemory()) {
+ if (!evictEldest() || !acquireBlockMemory()) {
+ // scalastyle:off throwerror
+ throw QueryExecutionErrors.cannotAcquireMemoryForWindowAggregateError(
+ blockBytes, 0L)
+ // scalastyle:on throwerror
+ }
+ }
+
+ // If buildBlockLevels throws, release the just-acquired memory (I4).
+ val levels =
+ try buildBlockLevels(blockIdx)
+ catch { case t: Throwable => releaseBlockMemory(); throw t }
+ blockLevelsCache.put(Integer.valueOf(blockIdx), levels)
+ levels
+ }
+
+ private def buildBlockLevels(blockIdx: Int): Array[Array[InternalRow]] = {
+ val blockStart = blockIdx * blockSize
+ val blockRows = math.min(blockSize, numRows - blockStart)
+
+ // Level 0: one aggregate per row in the block.
+ val leaves = new Array[InternalRow](blockRows)
+ val iter = rowArray.generateIterator(blockStart)
+ var i = 0
+ while (i < blockRows) {
+ if (!iter.hasNext) {
+ throw SparkException.internalError(
+ s"rowArray iterator exhausted at block $blockIdx row $i")
+ }
+ val row = iter.next()
+ val buf = newBuffer()
+ initProj.target(buf)(InternalRow.empty)
+ updateProj.target(buf)(joinedRow(buf, row))
+ leaves(i) = buf
+ i += 1
+ }
+
+ val allLevels = mutable.ArrayBuffer[Array[InternalRow]](leaves)
+ var prev = leaves
+ while (prev.length > 1) {
+ val parentCount = (prev.length + fanout - 1) / fanout
+ val parents = new Array[InternalRow](parentCount)
+ var p = 0
+ while (p < parentCount) {
+ val buf = newBuffer()
+ initProj.target(buf)(InternalRow.empty)
+ val childStart = p * fanout
+ val childEnd = math.min(childStart + fanout, prev.length)
+ var c = childStart
+ while (c < childEnd) {
+ mergeInto(buf, prev(c))
+ c += 1
+ }
+ parents(p) = buf
+ p += 1
+ }
+ allLevels += parents
+ prev = parents
+ }
+ allLevels.toArray
+ }
+
+ private def throwCacheEvictFailed(blockIdx: Int): Nothing = {
+ throw SparkException.internalError(
+ s"LRU cache eviction failed for block $blockIdx (size=${blockLevelsCache.size})")
+ }
+
+ // ---------- Memory accounting helpers ----------
+
+ /**
+ * Try to acquire `blockBytes` for one cached block. Returns true on full
+ * grant, false on partial (after rolling the partial grant back). Must
+ * not be called from within [[SegTreeSpiller.spill]] (I1).
+ */
+ private def acquireBlockMemory(): Boolean = {
+ val granted = spiller.acquireMemory(blockBytes)
+ if (granted < blockBytes) {
+ if (granted > 0) spiller.freeMemory(granted)
+ false
+ } else {
+ true
+ }
+ }
+
+ /**
+ * Release the accounting for one block. Caller ensures pairing with a
+ * prior successful [[acquireBlockMemory]] (I4).
+ */
+ private def releaseBlockMemory(): Unit = {
+ spiller.freeMemory(blockBytes)
+ }
+
+ /**
+ * Evict LRU blocks until `target` bytes have been freed (or cache is
+ * empty). Returns freed bytes. Called from [[SegTreeSpiller.spill]].
+ */
+ private def evictUntil(target: Long): Long = {
+ var freed = 0L
+ while (freed < target && !blockLevelsCache.isEmpty) {
+ freed += evictEldestReturnBytes()
+ }
+ freed
+ }
+
+ /** Evict one LRU block. Returns true if a block was evicted. */
+ private def evictEldest(): Boolean = {
+ if (blockLevelsCache.isEmpty) return false
+ evictEldestReturnBytes()
+ true
+ }
+
+ private def evictEldestReturnBytes(): Long = {
+ val it = blockLevelsCache.entrySet().iterator()
+ if (!it.hasNext) return 0L
+ val head = it.next()
+ it.remove()
+ releaseBlockMemory()
+ blockBytes
+ }
+
+ /** Release accounting for all cached blocks and clear the cache. */
+ private def releaseAllCachedBlocks(): Unit = {
+ val n = blockLevelsCache.size()
+ if (n > 0) {
+ blockLevelsCache.clear()
+ spiller.freeMemory(n.toLong * blockBytes)
+ }
+ }
+
+ private def closeRowArray(): Unit = {
+ // rowArray is caller-owned (see `build` docstring); drop the reference only.
+ rowArray = null
+ }
+}
+
+private[window] object WindowSegmentTree {
+ val DefaultFanout: Int = 16
+ val DefaultBlockSize: Int = 65536
+
+ /**
+ * 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`):
+ *
+ * - [[Min]], [[Max]]: idempotent semilattice.
+ * - [[Sum]], [[Count]]: additive monoid.
+ * - [[Average]]: sum + count, both additive monoids.
+ * - [[StddevPop]], [[StddevSamp]], [[VariancePop]], [[VarianceSamp]]:
+ * Welford (count, mean, M2) is associative -- see
+ * CentralMomentAgg.mergeExpressions.
+ *
+ * Intentionally excluded (tracked as follow-up): HyperLogLogPlusPlus /
+ * ApproxCountDistinct (sketch-buffer interaction unaudited), First / Last
+ * (order-dependent), CollectList / CollectSet (unbounded buffer growth),
+ * Percentile / ApproxPercentile (sorted-sketch buffer), and any
+ * ImperativeAggregate (excluded by the type check).
+ *
+ * Callers should use [[isEligible]] rather than `contains` directly.
+ */
+ val EligibleAggregates: Set[Class[_ <: DeclarativeAggregate]] = Set(
+ classOf[Min],
+ classOf[Max],
+ classOf[Sum],
+ classOf[Count],
+ classOf[Average],
+ classOf[StddevPop],
+ classOf[StddevSamp],
+ classOf[VariancePop],
+ classOf[VarianceSamp]
+ )
+
+ /**
+ * Returns true iff `f` is a [[DeclarativeAggregate]] on the explicit segment-tree
+ * allowlist. See [[EligibleAggregates]] for the rationale and excluded aggregates.
+ */
+ def isEligible(f: Expression): Boolean = f match {
+ case agg: DeclarativeAggregate => EligibleAggregates.contains(agg.getClass)
+ case _ => false
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/WindowBenchmark.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/WindowBenchmark.scala
new file mode 100644
index 000000000000..2284ff037e83
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/benchmark/WindowBenchmark.scala
@@ -0,0 +1,372 @@
+/*
+ * 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 scala.collection.mutable
+
+import org.apache.spark.benchmark.Benchmark
+import org.apache.spark.scheduler.{SparkListener, SparkListenerTaskEnd}
+import org.apache.spark.sql.internal.SQLConf
+
+/**
+ * Benchmark for window functions with bounded ROWS frames.
+ *
+ * Matrix (see PR description for rationale):
+ * - A: 5 aggregates x 3 cells (naive / segtree default / segtree bs=256) @ W=1001.
+ * Per-case N so naive ~3-5s/iter; STDDEV_SAMP pinned @ N=2M (multi-buffer stress).
+ * - B: SUM-over-INT, W sweep {10, 50, 201, 4001}; W=10/50 Pareto-loss stress,
+ * W=4001 also runs bs=256.
+ * - F: Spill guard, 1M String x MAX x W=1001 (stress).
+ * - C: N-sweep {2M, 8M, 16M} segtree-only @ W=1001 (memory-pressure invariance).
+ *
+ * Compare Per Row(ns) column for O(log W) scaling. Dev smoke via positional
+ * mainArgs: (0)=rowCount, (1)=halfWindow (default 100); do NOT combine with
+ * SPARK_GENERATE_BENCHMARK_FILES=1.
+ */
+object WindowBenchmark extends SqlBasedBenchmark {
+
+ // Section A: per-case N calibrated so naive baseline lands ~3-5s/iter.
+ private val A_N_INT: Long = 256L * 1024 // MIN/MAX/SUM/COUNT @ W=1001
+ private val A_N_AVG: Long = 192L * 1024 // AVG @ W=1001
+ private val A_N_STDDEV: Long = 2L * 1000L * 1000L // STDDEV stress
+
+ // Section B: W-sweep (W=10/50 stress: Pareto loss zone; W=4001 stress: O(W) cliff).
+ private val B_N_W10: Long = 2L * 1000L * 1000L
+ private val B_N_W50: Long = 2L * 1000L * 1000L
+ private val B_N_W201: Long = 1L * 1000L * 1000L
+ private val B_N_W4001: Long = 2L * 1000L * 1000L
+
+ // Section F: String spill.
+ private val SPILL_N: Long = 1L * 1000L * 1000L
+
+ // Section C: N-sweep segtree-only.
+ private val C_N_SMALL: Long = 2L * 1000L * 1000L
+ private val C_N_MID: Long = 8L * 1000L * 1000L
+ private val C_N_LARGE: Long = 16L * 1000L * 1000L
+ private val C_HALF_W: Int = 500 // W=1001
+
+ private val MAIN_HALF_W: Int = 500 // Section A W=1001
+
+ private val ITERS_NORMAL: Int = 5
+ private val ITERS_STRESS: Int = 3
+
+ override def runBenchmarkSuite(mainArgs: Array[String]): Unit = {
+ val smokeMode = mainArgs.nonEmpty
+ val smokeRowCount = if (smokeMode) mainArgs(0).toLong else 0L
+ val smokeHalfW = if (mainArgs.length > 1) mainArgs(1).toInt else 100
+
+ if (smokeMode) {
+ require(smokeRowCount >= 4096,
+ s"rowCount=$smokeRowCount too small; segtree may fallback. Use >= 4096.")
+ }
+
+ // Shared metrics listener: per-case peak mem + disk spill.
+ val metrics = mutable.Map[String, (Long, Long)]()
+ @volatile var currentCase: String = ""
+ val listener = new SparkListener {
+ override def onTaskEnd(e: SparkListenerTaskEnd): Unit = {
+ val m = e.taskMetrics
+ val cc = currentCase
+ if (m != null && cc.nonEmpty) {
+ metrics.synchronized {
+ val (pm, ds) = metrics.getOrElse(cc, (0L, 0L))
+ metrics(cc) = (math.max(pm, m.peakExecutionMemory), ds + m.diskBytesSpilled)
+ }
+ }
+ }
+ }
+ spark.sparkContext.addSparkListener(listener)
+ val allCaseNames = mutable.ArrayBuffer[String]()
+
+ def setupIntTable(n: Long): Unit = {
+ spark.range(n)
+ .selectExpr("id", "cast(rand(42) * 1000000 as int) as v")
+ .coalesce(1)
+ .createOrReplaceTempView("t")
+ }
+
+ def setupStringTable(n: Long): Unit = {
+ // ~20-char variable-length string; exercises spill path.
+ spark.range(n)
+ .selectExpr("id", "repeat(cast(id as string), 5) as v")
+ .coalesce(1)
+ .createOrReplaceTempView("t")
+ }
+
+ def frameFor(halfW: Int): String =
+ s"OVER (ORDER BY id ROWS BETWEEN $halfW PRECEDING AND $halfW FOLLOWING)"
+
+ // FP-digest trap: segtree merge order differs from row-by-row SlidingWindow
+ // order for AVG/STDDEV/VAR, so results differ at the ULP level (mathematically
+ // equivalent). HASH(double) is bit-sensitive, and HASH(ROUND(m, k)) is a trap:
+ // rows within ~1 ULP of a rounding bin boundary round differently across
+ // backends, and SUM of hashes amplifies tiny FP diffs into large digest drift
+ // (observed: 0.2% digest diff on STDDEV_SAMP @ N=2M even with per-row rel err
+ // <1e-10). Use SUM(CAST(ROUND(m, 3) * 1000 AS BIGINT)): one boundary-crossing
+ // row only contributes +/-1, so ULP-identical impls always agree while real
+ // bugs >1e-3 rel err are still caught. Integer aggregates remain bit-exact.
+ def digestExprFor(aggFn: String): String = {
+ if (aggFn.startsWith("AVG") || aggFn.startsWith("STDDEV") ||
+ aggFn.startsWith("VAR")) {
+ "CAST(ROUND(m, 3) * 1000 AS BIGINT)"
+ } else {
+ "HASH(m)"
+ }
+ }
+
+ def digest(aggFn: String, frame: String, sqlConfs: (String, String)*): Long = {
+ val expr = digestExprFor(aggFn)
+ withSQLConf(sqlConfs: _*) {
+ spark.sql(s"SELECT SUM($expr) FROM (SELECT $aggFn(v) $frame AS m FROM t)")
+ .head().getLong(0)
+ }
+ }
+
+ def rowsLabel(rows: Long): String = {
+ if (rows >= 1000000) s"${rows / 1000000}M"
+ else if (rows >= 1024) s"${rows / 1024}K"
+ else rows.toString
+ }
+
+ def runSectionA(
+ aggFn: String, iters: Int, rows: Long, halfW: Int, stressMark: String): Unit = {
+ val frame = frameFor(halfW)
+ val dNaive = digest(aggFn, frame)
+ val dSeg = digest(aggFn, frame, SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true")
+ val dSegBs = digest(aggFn, frame,
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_BLOCK_SIZE.key -> "256")
+ require(dNaive == dSeg,
+ s"$aggFn segtree digest mismatch: naive=$dNaive seg=$dSeg")
+ require(dNaive == dSegBs,
+ s"$aggFn segtree (bs=256) digest mismatch: naive=$dNaive seg=$dSegBs")
+
+ val W = 2 * halfW + 1
+ val benchmark = new Benchmark(
+ s"$aggFn sliding window, W=$W, ${rowsLabel(rows)} rows$stressMark",
+ rows, output = output)
+ val nNaive = s"$aggFn naive (current, baseline)"
+ val nSeg = s"$aggFn segtree (default)"
+ val nSegBs = s"$aggFn segtree (blockSize=256)"
+ allCaseNames ++= Seq(nNaive, nSeg, nSegBs)
+
+ benchmark.addCase(nNaive, numIters = iters) { _ =>
+ currentCase = nNaive
+ spark.sql(s"SELECT $aggFn(v) $frame FROM t").noop()
+ }
+ benchmark.addCase(nSeg, numIters = iters) { _ =>
+ currentCase = nSeg
+ withSQLConf(SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true") {
+ spark.sql(s"SELECT $aggFn(v) $frame FROM t").noop()
+ }
+ }
+ benchmark.addCase(nSegBs, numIters = iters) { _ =>
+ currentCase = nSegBs
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_BLOCK_SIZE.key -> "256") {
+ spark.sql(s"SELECT $aggFn(v) $frame FROM t").noop()
+ }
+ }
+ benchmark.run()
+ }
+
+ def runSectionB(
+ halfW: Int, stressBs: Boolean, rows: Long, iters: Int, stressMark: String): Unit = {
+ val aggFn = "SUM"
+ val frame = frameFor(halfW)
+ val W = 2 * halfW + 1
+ val dNaive = digest(aggFn, frame)
+ val dSeg = digest(aggFn, frame, SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true")
+ require(dNaive == dSeg, s"Section B W=$W digest mismatch: naive=$dNaive seg=$dSeg")
+
+ val benchmark = new Benchmark(
+ s"$aggFn scaling, W=$W, ${rowsLabel(rows)} rows$stressMark", rows, output = output)
+ val nNaive = s"$aggFn naive W=$W"
+ val nSeg = s"$aggFn segtree (default) W=$W"
+ allCaseNames ++= Seq(nNaive, nSeg)
+ benchmark.addCase(nNaive, numIters = iters) { _ =>
+ currentCase = nNaive
+ spark.sql(s"SELECT $aggFn(v) $frame FROM t").noop()
+ }
+ benchmark.addCase(nSeg, numIters = iters) { _ =>
+ currentCase = nSeg
+ withSQLConf(SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true") {
+ spark.sql(s"SELECT $aggFn(v) $frame FROM t").noop()
+ }
+ }
+ if (stressBs) {
+ val dSegBs = digest(aggFn, frame,
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_BLOCK_SIZE.key -> "256")
+ require(dNaive == dSegBs,
+ s"Section B W=$W bs=256 digest mismatch: naive=$dNaive segBs=$dSegBs")
+ val nSegBs = s"$aggFn segtree (blockSize=256) W=$W"
+ allCaseNames += nSegBs
+ benchmark.addCase(nSegBs, numIters = iters) { _ =>
+ currentCase = nSegBs
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_BLOCK_SIZE.key -> "256") {
+ spark.sql(s"SELECT $aggFn(v) $frame FROM t").noop()
+ }
+ }
+ }
+ benchmark.run()
+ }
+
+ def runSpillGuard(): Unit = {
+ val halfW = 500
+ val frame = frameFor(halfW)
+ // Digest parity skipped: pre-check scan on 1M String x W=1001 naive costs
+ // ~90s; correctness covered by SegmentTreeWindowFunctionSuite.
+ val benchmark = new Benchmark(
+ "MAX String spill guard, W=1001, 1M rows (stress)", SPILL_N, output = output)
+ val nNaive = "MAX naive (String)"
+ val nSeg = "MAX segtree default (String)"
+ allCaseNames ++= Seq(nNaive, nSeg)
+ benchmark.addCase(nNaive, numIters = ITERS_STRESS) { _ =>
+ currentCase = nNaive
+ spark.sql(s"SELECT MAX(v) $frame FROM t").noop()
+ }
+ benchmark.addCase(nSeg, numIters = ITERS_STRESS) { _ =>
+ currentCase = nSeg
+ withSQLConf(SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true") {
+ spark.sql(s"SELECT MAX(v) $frame FROM t").noop()
+ }
+ }
+ benchmark.run()
+ }
+
+ // Section C: N-sweep, segtree-only. Digest check skipped (Section A SUM @ N=2M
+ // covers the same code path). Naive skipped at 16M (~4min/iter, no value).
+ // Goal: memory-pressure invariance - per-row ns at 16M should be <= 2x at 2M.
+ def runSectionC(rows: Long): Unit = {
+ val aggFn = "SUM"
+ val frame = frameFor(C_HALF_W)
+ val W = 2 * C_HALF_W + 1
+ val benchmark = new Benchmark(
+ s"$aggFn N-sweep (segtree-only), W=$W, ${rowsLabel(rows)} rows (stress)",
+ rows, output = output)
+ val nSeg = s"$aggFn segtree (default) N=${rowsLabel(rows)}"
+ allCaseNames += nSeg
+ benchmark.addCase(nSeg, numIters = ITERS_STRESS) { _ =>
+ currentCase = nSeg
+ withSQLConf(SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true") {
+ spark.sql(s"SELECT $aggFn(v) $frame FROM t").noop()
+ }
+ }
+ benchmark.run()
+ }
+
+ try {
+ if (smokeMode) {
+ setupIntTable(smokeRowCount)
+ runBenchmark("SMOKE: Section A MIN") {
+ runSectionA("MIN", ITERS_STRESS, smokeRowCount, smokeHalfW, "")
+ }
+ runBenchmark("SMOKE: Section B SUM W sweep point") {
+ runSectionB(
+ smokeHalfW, stressBs = smokeHalfW >= 2000, smokeRowCount, ITERS_STRESS, "")
+ }
+ } else {
+ setupIntTable(A_N_INT)
+ runBenchmark("Section A - MIN (non-invertible)") {
+ runSectionA("MIN", ITERS_NORMAL, A_N_INT, MAIN_HALF_W, "")
+ }
+ runBenchmark("Section A - MAX (non-invertible)") {
+ runSectionA("MAX", ITERS_NORMAL, A_N_INT, MAIN_HALF_W, "")
+ }
+ runBenchmark("Section A - SUM (Spark has no inverse; full recompute)") {
+ runSectionA("SUM", ITERS_NORMAL, A_N_INT, MAIN_HALF_W, "")
+ }
+ runBenchmark("Section A - COUNT") {
+ runSectionA("COUNT", ITERS_NORMAL, A_N_INT, MAIN_HALF_W, "")
+ }
+
+ setupIntTable(A_N_AVG)
+ runBenchmark("Section A - AVG (multi-buffer)") {
+ runSectionA("AVG", ITERS_NORMAL, A_N_AVG, MAIN_HALF_W, "")
+ }
+
+ setupIntTable(A_N_STDDEV)
+ runBenchmark("Section A - STDDEV_SAMP (multi-buffer, stress)") {
+ runSectionA("STDDEV_SAMP", ITERS_STRESS, A_N_STDDEV, MAIN_HALF_W, " (stress)")
+ }
+
+ setupIntTable(B_N_W10)
+ runBenchmark("Section B - W=10 scaling (stress: Pareto loss zone)") {
+ runSectionB(5, stressBs = false, B_N_W10, ITERS_STRESS, " (stress)")
+ }
+ setupIntTable(B_N_W50)
+ runBenchmark("Section B - W=50 scaling (stress: Pareto loss zone)") {
+ runSectionB(25, stressBs = false, B_N_W50, ITERS_STRESS, " (stress)")
+ }
+ setupIntTable(B_N_W201)
+ runBenchmark("Section B - W=201 scaling") {
+ runSectionB(100, stressBs = false, B_N_W201, ITERS_NORMAL, "")
+ }
+ setupIntTable(B_N_W4001)
+ runBenchmark("Section B - W=4001 scaling (stress, + bs=256 cross-block)") {
+ runSectionB(2000, stressBs = true, B_N_W4001, ITERS_STRESS, " (stress)")
+ }
+
+ setupStringTable(SPILL_N)
+ runBenchmark("Section F - spill regression guard (String, stress)") {
+ runSpillGuard()
+ }
+
+ setupIntTable(C_N_SMALL)
+ runBenchmark("Section C - N-sweep small (stress)") {
+ runSectionC(C_N_SMALL)
+ }
+ setupIntTable(C_N_MID)
+ runBenchmark("Section C - N-sweep mid (stress)") {
+ runSectionC(C_N_MID)
+ }
+ setupIntTable(C_N_LARGE)
+ runBenchmark("Section C - N-sweep large (stress)") {
+ runSectionC(C_N_LARGE)
+ }
+ }
+
+ // Drain listener before reading metrics.
+ spark.sparkContext.listenerBus.waitUntilEmpty()
+
+ // scalastyle:off println
+ val out = System.out
+ out.println()
+ out.println("Memory/Spill (peak executor memory / total disk spilled):")
+ if (allCaseNames.nonEmpty) {
+ val width = allCaseNames.map(_.length).max + 2
+ for (name <- allCaseNames) {
+ val (pm, ds) = metrics.synchronized {
+ metrics.getOrElse(name, (0L, 0L))
+ }
+ val peakMb = pm.toDouble / (1024.0 * 1024.0)
+ val label = (name + ":").padTo(width, ' ')
+ out.println(f" $label%s peak=$peakMb%8.2f MB spilled=$ds%d B")
+ }
+ }
+ out.println()
+ // scalastyle:on println
+ } finally {
+ spark.sparkContext.removeSparkListener(listener)
+ }
+ }
+}
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
new file mode 100644
index 000000000000..d6d18cff4df5
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowFunctionSuite.scala
@@ -0,0 +1,896 @@
+/*
+ * 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.window
+
+import org.apache.spark.sql.{DataFrame, Encoder, Encoders, QueryTest, Row}
+import org.apache.spark.sql.expressions.{Aggregator, MutableAggregationBuffer, Window}
+import org.apache.spark.sql.expressions.UserDefinedAggregateFunction
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+import org.apache.spark.sql.types.{DataType, LongType, StructType}
+import org.apache.spark.util.SparkErrorUtils
+
+/**
+ * End-to-end tests for the block-chunked segment-tree moving window frame.
+ * Covers basic aggregates, frame boundaries, min-rows fallback, NULL/NaN,
+ * numeric/string/date-timestamp types, RANGE, Decimal/Binary merge, UDAF
+ * fallback, and frame lifecycle.
+ */
+class SegmentTreeWindowFunctionSuite extends QueryTest with SharedSparkSession {
+
+ import testImplicits._
+
+ // Force seg-tree path regardless of partition size (fallback exercised explicitly).
+ private val enableSegTree: Map[String, String] = Map(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "1")
+
+ private val disableSegTree: Map[String, String] = Map(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "false")
+
+ /** Build `f(conf)` twice (enabled / disabled) and assert equal results. */
+ private def checkEquivalence(build: () => DataFrame): Unit = {
+ val baseline: Seq[Row] = withSQLConf(disableSegTree.toSeq: _*) {
+ build().collect().toSeq
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = build().collect().toSeq
+ // Use QueryTest.sameRows (which normalizes Array[_] via prepareRow before
+ // sorting) instead of sortBy(_.toString): Row.toString on Array[Byte]
+ // and similar ref-typed values is address-based and can reorder baseline
+ // vs actual differently even when the multiset of rows is identical.
+ QueryTest.sameRows(baseline, actual, isSorted = false).foreach { err =>
+ fail(s"segment-tree output differs from baseline.\n$err")
+ }
+ }
+ }
+
+ /** Standard fixture: 3 partitions, sizes 40/40/40, values = row index. */
+ private def baseDF: DataFrame = {
+ spark.range(0, 120).selectExpr(
+ "id",
+ "(id % 3) AS pk",
+ "CAST(id AS INT) AS v")
+ }
+
+ private def winSpec(lo: Int, hi: Int) =
+ Window.partitionBy($"pk").orderBy($"id").rowsBetween(lo, hi)
+
+ // A1: basic aggregate equivalence
+
+ test("MIN over ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", min($"v").over(winSpec(-3, 3)).as("agg")))
+ }
+
+ test("MAX over ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", max($"v").over(winSpec(-3, 3)).as("agg")))
+ }
+
+ test("SUM over ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", sum($"v").over(winSpec(-3, 3)).as("agg")))
+ }
+
+ test("COUNT over ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", count($"v").over(winSpec(-3, 3)).as("agg")))
+ }
+
+ test("AVG over ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", avg($"v").over(winSpec(-3, 3)).as("agg")))
+ }
+
+ test("MIN + MAX + SUM share a single window frame") {
+ checkEquivalence(() =>
+ baseDF.select(
+ $"id",
+ $"pk",
+ min($"v").over(winSpec(-3, 3)).as("mn"),
+ max($"v").over(winSpec(-3, 3)).as("mx"),
+ sum($"v").over(winSpec(-3, 3)).as("sm")))
+ }
+
+ // A2: frame-size boundaries
+
+ test("frame size = 1 (CURRENT ROW only)") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", sum($"v").over(winSpec(0, 0)).as("agg")))
+ }
+
+ test("frame spans full partition") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk", sum($"v").over(winSpec(-100, 100)).as("agg")))
+ }
+
+ test("frame extends past both partition edges") {
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk",
+ sum($"v").over(winSpec(-50, 50)).as("agg"),
+ min($"v").over(winSpec(-50, 50)).as("mn"),
+ max($"v").over(winSpec(-50, 50)).as("mx")))
+ }
+
+
+ test("partition below minPartitionRows falls back to SlidingWindowFunctionFrame") {
+ // 5-row partition, min threshold = 10 -> must fall back.
+ val df = spark.range(0, 5).selectExpr(
+ "id", "0 AS pk", "CAST(id AS INT) AS v")
+ val enabledConf = Map(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "10")
+
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ df.select($"id", sum($"v").over(winSpec(-1, 1)).as("s"))
+ .collect().sortBy(_.toString)
+ }
+ val actual = withSQLConf(enabledConf.toSeq: _*) {
+ df.select($"id", sum($"v").over(winSpec(-1, 1)).as("s"))
+ .collect().sortBy(_.toString)
+ }
+ assert(actual.toSeq === baseline.toSeq)
+
+ // Confirm the fallback flag actually flipped.
+ withSQLConf(enabledConf.toSeq: _*) {
+ SegmentTreeWindowTestHelper.withSmallPartitionFrame(
+ SQLConf.get, rows = 5) { frame =>
+ assert(frame.fallbackUsed,
+ "expected fallbackUsed=true for partition smaller than minPartitionRows")
+ }
+ }
+ }
+
+
+ test("NTH_VALUE over ROWS frame falls back cleanly (no mergeExpressions crash)") {
+ // NthValue extends DeclarativeAggregate but its mergeExpressions throws
+ // mergeUnsupportedByWindowFunctionError. eligibleForSegTree must exclude it.
+ val df = baseDF
+ val withSegTree = withSQLConf(enableSegTree.toSeq: _*) {
+ df.selectExpr(
+ "id", "pk",
+ "nth_value(v, 3) OVER (PARTITION BY pk ORDER BY id " +
+ "ROWS BETWEEN 5 PRECEDING AND 5 FOLLOWING) AS n3")
+ .collect().sortBy(_.toString)
+ }
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ df.selectExpr(
+ "id", "pk",
+ "nth_value(v, 3) OVER (PARTITION BY pk ORDER BY id " +
+ "ROWS BETWEEN 5 PRECEDING AND 5 FOLLOWING) AS n3")
+ .collect().sortBy(_.toString)
+ }
+ assert(withSegTree.toSeq === baseline.toSeq)
+ }
+
+ test("ROW_NUMBER over ROWS frame falls back cleanly (no mergeExpressions crash)") {
+ val df = baseDF
+ val withSegTree = withSQLConf(enableSegTree.toSeq: _*) {
+ df.selectExpr(
+ "id", "pk",
+ "row_number() OVER (PARTITION BY pk ORDER BY id) AS rn")
+ .collect().sortBy(_.toString)
+ }
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ df.selectExpr(
+ "id", "pk",
+ "row_number() OVER (PARTITION BY pk ORDER BY id) AS rn")
+ .collect().sortBy(_.toString)
+ }
+ assert(withSegTree.toSeq === baseline.toSeq)
+ }
+
+ // A3: NULL/NaN/Infinity; A4: numeric/string/date-timestamp types.
+ // Unsupported-merge / DISTINCT / feature-flag fallback.
+ // Oracle: run with seg-tree enabled and disabled, assert equal Row sequences.
+
+ // A3: NULL / special values
+
+ test("all-NULL column: MIN/MAX/SUM/AVG/COUNT") {
+ val df = spark.range(0, 30).selectExpr(
+ "id", "(id % 3) AS pk", "CAST(NULL AS INT) AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ min($"v").over(winSpec(-3, 3)).as("mn"),
+ max($"v").over(winSpec(-3, 3)).as("mx"),
+ sum($"v").over(winSpec(-3, 3)).as("sm"),
+ avg($"v").over(winSpec(-3, 3)).as("av"),
+ count($"v").over(winSpec(-3, 3)).as("cn")))
+ }
+
+ test("mixed NULL and non-NULL: NULLs must not leak into MIN/MAX") {
+ // Every 3rd value is NULL. Aggregates must skip them (NULL-agnostic merge).
+ 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",
+ min($"v").over(winSpec(-4, 4)).as("mn"),
+ max($"v").over(winSpec(-4, 4)).as("mx"),
+ sum($"v").over(winSpec(-4, 4)).as("sm"),
+ count($"v").over(winSpec(-4, 4)).as("cn")))
+ }
+
+ 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.
+ val df = spark.range(0, 30).selectExpr(
+ "id",
+ "(id % 2) AS pk",
+ """CASE
+ WHEN id % 7 = 0 THEN double('NaN')
+ WHEN id % 7 = 1 THEN double('Infinity')
+ WHEN id % 7 = 2 THEN double('-Infinity')
+ ELSE CAST(id AS DOUBLE)
+ END AS v""")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ min($"v").over(winSpec(-3, 3)).as("mn"),
+ max($"v").over(winSpec(-3, 3)).as("mx"),
+ sum($"v").over(winSpec(-3, 3)).as("sm")))
+ }
+
+ // A4: data types
+
+ test("numeric types: Int / Long / Double / Decimal") {
+ val df = spark.range(0, 60).selectExpr(
+ "id",
+ "(id % 3) AS pk",
+ "CAST(id AS INT) AS vi",
+ "CAST(id * 1000000L AS LONG) AS vl",
+ "CAST(id AS DOUBLE) + 0.25 AS vd",
+ "CAST(id AS DECIMAL(20,4)) AS vdec")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ sum($"vi").over(winSpec(-2, 2)).as("si"),
+ min($"vl").over(winSpec(-2, 2)).as("ml"),
+ max($"vd").over(winSpec(-2, 2)).as("xd"),
+ sum($"vdec").over(winSpec(-2, 2)).as("sdec"),
+ avg($"vdec").over(winSpec(-2, 2)).as("adec")))
+ }
+
+ test("String lexicographic MIN/MAX") {
+ // Non-monotone values so MIN/MAX exercise the seg-tree merge (not edge).
+ val df = spark.range(0, 40).selectExpr(
+ "id",
+ "(id % 2) AS pk",
+ "CONCAT('s', LPAD(CAST((id * 37) % 97 AS STRING), 3, '0')) AS v")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ min($"v").over(winSpec(-3, 3)).as("mn"),
+ max($"v").over(winSpec(-3, 3)).as("mx")))
+ }
+
+ test("Date / Timestamp MIN/MAX") {
+ val df = spark.range(0, 40).selectExpr(
+ "id",
+ "(id % 2) AS pk",
+ "date_add(DATE'2020-01-01', CAST((id * 13) % 365 AS INT)) AS vd",
+ "CAST(TIMESTAMP'2020-01-01 00:00:00' + " +
+ "make_interval(0, 0, 0, 0, 0, 0, CAST(id AS DECIMAL(18,6))) AS TIMESTAMP) AS vt")
+ checkEquivalence(() =>
+ df.select($"id", $"pk",
+ min($"vd").over(winSpec(-3, 3)).as("mnd"),
+ max($"vd").over(winSpec(-3, 3)).as("mxd"),
+ min($"vt").over(winSpec(-3, 3)).as("mnt"),
+ max($"vt").over(winSpec(-3, 3)).as("mxt")))
+ }
+
+
+ test("collect_list falls back cleanly (non-DeclarativeAggregate)") {
+ // collect_list is TypedImperativeAggregate; eligibleForSegTree must decline.
+ checkEquivalence(() =>
+ baseDF.select($"id", $"pk",
+ collect_list($"v").over(winSpec(-2, 2)).as("lst")))
+ }
+
+ test("DISTINCT window aggregate is rejected by analyzer regardless of seg-tree flag") {
+ // Analyzer throws DISTINCT_WINDOW_FUNCTION_UNSUPPORTED before frame
+ // construction; seg-tree flag must not alter this behavior.
+ def run(): Unit = {
+ baseDF.select($"id", $"pk",
+ count_distinct($"v").over(winSpec(-3, 3)).as("cd")).collect()
+ }
+ withSQLConf(disableSegTree.toSeq: _*) {
+ val e = intercept[org.apache.spark.sql.AnalysisException](run())
+ assert(e.getMessage.contains("DISTINCT_WINDOW_FUNCTION_UNSUPPORTED"))
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val e = intercept[org.apache.spark.sql.AnalysisException](run())
+ assert(e.getMessage.contains("DISTINCT_WINDOW_FUNCTION_UNSUPPORTED"))
+ }
+ }
+
+ test("feature flag off: segmentTree.enabled=false yields baseline semantics") {
+ // Sanity: disabling the flag on a seg-tree-eligible workload still
+ // produces the SlidingWindowFunctionFrame answer.
+ val df = baseDF
+ val expected = withSQLConf(disableSegTree.toSeq: _*) {
+ df.select($"id", $"pk", min($"v").over(winSpec(-3, 3)).as("mn"))
+ .collect().sortBy(_.toString).toSeq
+ }
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "false",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "1024") {
+ val actual = df.select($"id", $"pk", min($"v").over(winSpec(-3, 3)).as("mn"))
+ .collect().sortBy(_.toString).toSeq
+ assert(actual === expected)
+ }
+ }
+
+ // A5: RANGE frame equivalence (single-order-expr admission).
+ // MIN/MAX non-invertible, guaranteeing seg-tree path is exercised.
+
+ /** Run `sql` twice (flag off / on) and checkAnswer equality. */
+ private def checkRangeEquivalence(df: DataFrame, query: String): Unit = {
+ df.createOrReplaceTempView("t")
+ try {
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(query).collect().sortBy(_.toString)
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(query).collect().sortBy(_.toString)
+ assert(actual.toSeq === baseline.toSeq,
+ s"segment-tree output differs from baseline.\nExpected: ${baseline.toSeq}\n" +
+ s"Actual: ${actual.toSeq}")
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ test("-- RANGE INT offset basic (non-uniform gaps, MIN/MAX)") {
+ // Non-uniform gaps so admit/drop loops must consult the order-key comparator.
+ val df = spark.range(0, 40).selectExpr(
+ "CAST(id AS INT) AS id",
+ "(CAST(id AS INT) % 2) AS pk",
+ "CAST(CASE CAST(id AS INT) % 7 " +
+ "WHEN 0 THEN 1 WHEN 1 THEN 3 WHEN 2 THEN 4 WHEN 3 THEN 4 " +
+ "WHEN 4 THEN 7 WHEN 5 THEN 10 ELSE 15 END + (CAST(id AS INT) / 7) * 20 AS INT) AS k",
+ "CAST((id * 31) % 97 AS INT) AS v")
+ checkRangeEquivalence(df,
+ """SELECT id, pk,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY k
+ | RANGE BETWEEN 2 PRECEDING AND 2 FOLLOWING) AS mn,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY k
+ | RANGE BETWEEN 2 PRECEDING AND 2 FOLLOWING) AS mx
+ |FROM t""".stripMargin)
+ }
+
+ test("-- RANGE Timestamp with INTERVAL offset (MAX)") {
+ // Irregular gaps force the timestamp comparator at the 1-hour boundary.
+ val df = spark.range(0, 30).selectExpr(
+ "CAST(id AS INT) AS id",
+ "(CAST(id AS INT) % 2) AS pk",
+ "CAST(TIMESTAMP'2024-01-01 10:00:00' + " +
+ "make_interval(0, 0, 0, 0, 0, 30 * CAST(id AS INT) * " +
+ "(CASE CAST(id AS INT) % 3 WHEN 0 THEN 1 WHEN 1 THEN 3 ELSE 4 END), 0) " +
+ "AS TIMESTAMP) AS ts",
+ "CAST((id * 17) % 53 AS INT) AS v")
+ checkRangeEquivalence(df,
+ """SELECT id, pk,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY ts
+ | RANGE BETWEEN INTERVAL '1' HOUR PRECEDING
+ | AND INTERVAL '1' HOUR FOLLOWING) AS mx
+ |FROM t""".stripMargin)
+ }
+
+ test("-- RANGE with tie (duplicate order keys) inclusion at boundary") {
+ // Trap: RANGE `0 PRECEDING AND 0 FOLLOWING` must include the FULL tie
+ // group at the current row's key, not just the current row. A ROWS-vs-
+ // RANGE confusion would return per-row MIN/MAX instead of group-level.
+ val rows = (0 until 40).map { i =>
+ val k = Seq(1, 2, 2, 2, 3, 4, 5)(i % 7)
+ (i, i % 2, k, (i * 13) % 41)
+ }
+ val df = rows.toDF("id", "pk", "k", "v")
+ checkRangeEquivalence(df,
+ """SELECT id, pk, k,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY k
+ | RANGE BETWEEN 0 PRECEDING AND 0 FOLLOWING) AS mn,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY k
+ | RANGE BETWEEN 0 PRECEDING AND 0 FOLLOWING) AS mx
+ |FROM t""".stripMargin)
+ }
+
+ test("-- RANGE frame wider than partition (C4: admit/drop loops no-op)") {
+ // Once the first batch is admitted, admit/drop must detect no change
+ // and skip work.
+ val df = spark.range(0, 25).selectExpr(
+ "CAST(id AS INT) AS id",
+ "(CAST(id AS INT) / 5) AS pk",
+ "CAST((id * 7) % 23 AS INT) AS k",
+ "CAST((id * 19) % 101 AS INT) AS v")
+ checkRangeEquivalence(df,
+ """SELECT id, pk,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY k
+ | RANGE BETWEEN 100 PRECEDING AND 100 FOLLOWING) AS mn,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY k
+ | RANGE BETWEEN 100 PRECEDING AND 100 FOLLOWING) AS mx
+ |FROM t""".stripMargin)
+ }
+
+ test("-- RANGE with NULL order key (NULLS FIRST / NULLS LAST)") {
+ // Trap: Spark groups all NULLs into a single equivalence class at head
+ // (NULLS FIRST) or tail (NULLS LAST); seg-tree must treat NULL as a
+ // tie group identical to the sliding baseline.
+ val rows = (0 until 36).map { i =>
+ val kOpt: Option[Int] = (i % 6) match {
+ case 0 | 1 | 5 => None
+ case 2 => Some(1)
+ case 3 => Some(2)
+ case _ => Some(3)
+ }
+ (i, i % 2, kOpt, (i * 11) % 37)
+ }
+ val df = rows.toDF("id", "pk", "k", "v")
+ checkRangeEquivalence(df,
+ """SELECT id, pk,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY k ASC NULLS FIRST
+ | RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS mn_nf,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY k ASC NULLS FIRST
+ | RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS mx_nf,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY k ASC NULLS LAST
+ | RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS mn_nl,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY k ASC NULLS LAST
+ | RANGE BETWEEN 1 PRECEDING AND 1 FOLLOWING) AS mx_nl
+ |FROM t""".stripMargin)
+ }
+
+ // Decimal overflow / BinaryType MIN/MAX across block merge; UDAF fallback.
+ // Trap: blockSize=16 is SQLConf minimum; frame > blockSize ensures the
+ // seg-tree merge path is actually crossed.
+
+ private val segTreeBlock: String = "16"
+ private val segTreeFramePrec: Int = 17
+ private val segTreeRows: Int = 20
+
+ private def withSegTreeBlock(conf: (String, String)*)(body: => Unit): Unit = {
+ val extra = Seq(SQLConf.WINDOW_SEGMENT_TREE_BLOCK_SIZE.key -> segTreeBlock) ++ conf
+ withSQLConf(extra: _*)(body)
+ }
+
+
+ /**
+ * 20 rows in one partition, Decimal(38, 0) values near the type's upper
+ * bound; frame of `segTreeFramePrec` PRECEDING..CURRENT ROW makes any
+ * >=2-row window overflow Sum. Block 16 + frame 17 forces cross-block merge.
+ */
+ private def decimalOverflowDF: DataFrame = {
+ // 9e37 -- below Decimal(38,0) MAX (~9.99e37), but 2x overflows.
+ val big = "90000000000000000000000000000000000000" // 38 digits
+ spark.range(0, segTreeRows.toLong).selectExpr(
+ "CAST(id AS INT) AS id",
+ "0 AS pk",
+ s"CAST('$big' AS DECIMAL(38, 0)) AS v")
+ }
+
+ private val decimalOverflowSql: String =
+ s"""SELECT id, pk,
+ | SUM(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN $segTreeFramePrec PRECEDING AND CURRENT ROW) AS s
+ |FROM t""".stripMargin
+
+ test("a -- Decimal overflow ANSI on, seg-tree matches sliding (both throw)") {
+ val df = decimalOverflowDF
+ df.createOrReplaceTempView("t")
+ try {
+ withSegTreeBlock(SQLConf.ANSI_ENABLED.key -> "true") {
+ withSQLConf(disableSegTree.toSeq: _*) {
+ val e = intercept[Exception] {
+ spark.sql(decimalOverflowSql).collect()
+ }
+ assert(hasArithmeticCause(e),
+ s"expected ArithmeticException root cause, got: ${e.getMessage}")
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val e = intercept[Exception] {
+ spark.sql(decimalOverflowSql).collect()
+ }
+ assert(hasArithmeticCause(e),
+ s"expected ArithmeticException root cause, got: ${e.getMessage}")
+ }
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ test("b -- Decimal overflow ANSI off, seg-tree matches sliding (NULL on overflow)") {
+ val df = decimalOverflowDF
+ df.createOrReplaceTempView("t")
+ try {
+ withSegTreeBlock(SQLConf.ANSI_ENABLED.key -> "false") {
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(decimalOverflowSql).collect().sortBy(_.toString)
+ }
+ // At least one row must be NULL so we know overflow actually fired.
+ assert(baseline.exists(_.isNullAt(2)),
+ "baseline should contain NULL overflow rows; test data may be too small")
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(decimalOverflowSql).collect().sortBy(_.toString)
+ assert(actual.toSeq === baseline.toSeq)
+ }
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ test("c -- mid-window Decimal overflow slides past (seg-tree == sliding)") {
+ // Big values at ids 14..17 straddle block boundary at id=16, so any
+ // 4-row window overlapping >=2 of them overflows (-> NULL when ANSI off)
+ // and cross-block merge sees overflowing buffers. Rows past id=20 slide
+ // clear and recover non-NULL.
+ val big = "90000000000000000000000000000000000000"
+ val df = spark.range(0, 24).selectExpr(
+ "CAST(id AS INT) AS id",
+ "0 AS pk",
+ s"""CASE WHEN id IN (14, 15, 16, 17)
+ THEN CAST('$big' AS DECIMAL(38, 0))
+ ELSE CAST(id AS DECIMAL(38, 0))
+ END AS v""")
+ df.createOrReplaceTempView("t")
+ try {
+ val sqlStr =
+ """SELECT id, pk,
+ | SUM(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN 3 PRECEDING AND CURRENT ROW) AS s
+ |FROM t""".stripMargin
+ withSegTreeBlock(SQLConf.ANSI_ENABLED.key -> "false") {
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(sqlStr).collect().sortBy(_.toString)
+ }
+ // Sanity: overflow fired AND later rows recover non-NULL.
+ assert(baseline.exists(_.isNullAt(2)),
+ "baseline should contain NULL overflow rows")
+ assert(baseline.exists(r => r.getInt(0) >= 21 && !r.isNullAt(2)),
+ "rows with id>=21 should be non-NULL (window slid past big values)")
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(sqlStr).collect().sortBy(_.toString)
+ assert(actual.toSeq === baseline.toSeq)
+ }
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ /** True iff the root cause of `t` is an [[ArithmeticException]] (ANSI overflow). */
+ private def hasArithmeticCause(t: Throwable): Boolean =
+ Option(SparkErrorUtils.getRootCause(t)).exists(_.isInstanceOf[ArithmeticException])
+
+
+ /** Pattern of 20 Array[Byte] values used across a. */
+ private def binaryVariedRows: Seq[(Int, Array[Byte])] = {
+ (0 until 20).map { i =>
+ val arr: Array[Byte] = (i % 8) match {
+ case 0 => Array[Byte](0x01, 0x02)
+ case 1 => Array[Byte](0x00)
+ case 2 => Array[Byte](0x7f)
+ case 3 => Array[Byte](0x7f, 0x00)
+ case 4 => Array[Byte](0x10, 0x20, 0x30)
+ case 5 => Array[Byte](0x10, 0x20)
+ case 6 => Array[Byte](0x10)
+ case _ => Array[Byte](0x05, 0x05, 0x05, 0x05)
+ }
+ (i, arr)
+ }
+ }
+
+ test("a -- BinaryType MIN/MAX cross-block merge") {
+ // Varied lengths/content; frame > blockSize guarantees merge path hit.
+ val df = binaryVariedRows.toDF("id", "v").selectExpr("id", "0 AS pk", "v")
+ df.createOrReplaceTempView("t")
+ try {
+ withSegTreeBlock() {
+ val sqlStr =
+ s"""SELECT id, pk,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN $segTreeFramePrec PRECEDING AND CURRENT ROW) AS mn,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN $segTreeFramePrec PRECEDING AND CURRENT ROW) AS mx
+ |FROM t""".stripMargin
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(sqlStr).collect().toSeq
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(sqlStr).collect().toSeq
+ QueryTest.sameRows(baseline, actual, isSorted = false).foreach { err =>
+ fail(s"seg-tree binary MIN/MAX differs.\n$err")
+ }
+ }
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ test("b -- BinaryType empty/NULL/single-zero distinction") {
+ // Trap: Spark treats empty-array and NULL as distinct; seg-tree must
+ // respect that. Pattern cycles across the block boundary.
+ val rows: Seq[(Int, Array[Byte])] = (0 until 20).map { i =>
+ val arr: Array[Byte] = (i % 4) match {
+ case 0 => Array[Byte](0x00)
+ case 1 => Array[Byte]()
+ case 2 => null
+ case _ => Array[Byte](0x01, 0x02)
+ }
+ (i, arr)
+ }
+ val df = rows.toDF("id", "v").selectExpr("id", "0 AS pk", "v")
+ df.createOrReplaceTempView("t")
+ try {
+ withSegTreeBlock() {
+ val sqlStr =
+ s"""SELECT id, pk,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN $segTreeFramePrec PRECEDING AND CURRENT ROW) AS mn,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN $segTreeFramePrec PRECEDING AND CURRENT ROW) AS mx
+ |FROM t""".stripMargin
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(sqlStr).collect().toSeq
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(sqlStr).collect().toSeq
+ QueryTest.sameRows(baseline, actual, isSorted = false).foreach { err =>
+ fail(s"seg-tree empty/NULL binary differs.\n$err")
+ }
+ }
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ test("c -- BinaryType unsigned lexicographic ordering") {
+ // Trap: Spark's BinaryType comparator is unsigned (0xFF > 0x01); a
+ // signed-byte comparator would get this backwards. Seg-tree must match.
+ val unsignedPattern: IndexedSeq[Array[Byte]] = IndexedSeq(
+ Array[Byte](0xff.toByte),
+ Array[Byte](0x01),
+ Array[Byte](0x80.toByte, 0x00),
+ Array[Byte](0x7f, 0xff.toByte),
+ Array[Byte](0xfe.toByte),
+ Array[Byte](0x00, 0xff.toByte),
+ Array[Byte](0x80.toByte),
+ Array[Byte](0x7f))
+ val rows: Seq[(Int, Array[Byte])] =
+ (0 until 20).map(i => (i, unsignedPattern(i % unsignedPattern.length)))
+ val df = rows.toDF("id", "v").selectExpr("id", "0 AS pk", "v")
+ df.createOrReplaceTempView("t")
+ try {
+ withSegTreeBlock() {
+ val sqlStr =
+ s"""SELECT id, pk,
+ | MIN(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN $segTreeFramePrec PRECEDING AND CURRENT ROW) AS mn,
+ | MAX(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN $segTreeFramePrec PRECEDING AND CURRENT ROW) AS mx
+ |FROM t""".stripMargin
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(sqlStr).collect().toSeq
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(sqlStr).collect().toSeq
+ QueryTest.sameRows(baseline, actual, isSorted = false).foreach { err =>
+ fail(s"seg-tree unsigned binary ordering differs.\n$err")
+ }
+ }
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ // ScalaUDAF/ScalaAggregator both extend ImperativeAggregate and must be
+ // rejected by `eligibleForSegTree`. Flag ON must not throw (segtree merge on
+ // ImperativeAggregate would NPE) and must match flag OFF bit-for-bit.
+
+ test("a -- legacy ScalaUDAF falls back cleanly (no seg-tree merge)") {
+ val udaf = new LegacySumUdaf
+ spark.udf.register("seg_tree_legacy_sum", udaf)
+ val df = baseDF.selectExpr("id", "pk", "CAST(v AS LONG) AS v")
+ val query =
+ """SELECT id, pk,
+ | seg_tree_legacy_sum(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN 5 PRECEDING AND CURRENT ROW) AS s
+ |FROM t""".stripMargin
+ df.createOrReplaceTempView("t")
+ try {
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(query).collect().sortBy(_.toString)
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(query).collect().sortBy(_.toString)
+ assert(actual.toSeq === baseline.toSeq,
+ s"ScalaUDAF fallback result differs.\nExpected: ${baseline.toSeq}\n" +
+ s"Actual: ${actual.toSeq}")
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ test("b -- typed Aggregator falls back cleanly (no seg-tree merge)") {
+ val agg = udaf(new LongSumAggregator)
+ spark.udf.register("seg_tree_typed_sum", agg)
+ val df = baseDF.selectExpr("id", "pk", "CAST(v AS LONG) AS v")
+ val query =
+ """SELECT id, pk,
+ | seg_tree_typed_sum(v) OVER (PARTITION BY pk ORDER BY id
+ | ROWS BETWEEN 5 PRECEDING AND CURRENT ROW) AS s
+ |FROM t""".stripMargin
+ df.createOrReplaceTempView("t")
+ try {
+ val baseline = withSQLConf(disableSegTree.toSeq: _*) {
+ spark.sql(query).collect().sortBy(_.toString)
+ }
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val actual = spark.sql(query).collect().sortBy(_.toString)
+ assert(actual.toSeq === baseline.toSeq,
+ s"typed Aggregator fallback result differs.\nExpected: ${baseline.toSeq}\n" +
+ s"Actual: ${actual.toSeq}")
+ }
+ } finally {
+ spark.catalog.dropTempView("t")
+ }
+ }
+
+ test("SPARK-56546: LAG does not eagerly construct AggregateProcessor under segtree") {
+ // Pre-fix, `val processor` eagerly invoked AggregateProcessor.apply on any
+ // non-empty `functions`, throwing INTERNAL_ERROR for lag(...) when routing
+ // hit FRAME_LESS_OFFSET. Frameless lag below is the minimal repro.
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val df = spark.range(10).select(
+ col("id"),
+ expr("lag(id, 1, id) OVER (ORDER BY id)").as("lag"))
+ val expected = (0L until 10L).map(i => Row(i, if (i == 0) 0L else i - 1))
+ checkAnswer(df, expected)
+ }
+ }
+
+ // Frame lifecycle
+ // T1a: no fallback on pure-segtree partition.
+ // T1b: fallback lazily allocated on small partition.
+ // T1c: segtree->small reuses frame; fallback allocated once.
+ // T1d: small->segtree drops retained fallback reference (GC-eligible).
+ // T2 : throwing fallback.prepare must not flip `fallbackUsed` (frame is
+ // `final` so we cannot inject a throwing fallback; we rely on the
+ // structural witness that `fallbackUsed` is set ONLY after
+ // `fallback.prepare` returns normally -- a regression would fail T1d).
+ // T3 : close() after only small partition (tree never built) is a no-op.
+
+ test("lifecycle: no fallback allocated on segtree-only partition") {
+ val conf = new SQLConf
+ conf.setConfString(SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key, "1")
+ SegmentTreeWindowTestHelper.withFrameFactory(conf) { factory =>
+ val frame = factory.newFrame()
+ val array = factory.newArray(rows = 10)
+ frame.prepare(array)
+ assert(!frame.fallbackUsed, "segtree path expected")
+ assert(!frame.fallbackAllocated,
+ "fallback must not be allocated on segtree-only partition")
+ }
+ }
+
+ test("lifecycle: fallback lazily allocated on small partition") {
+ val conf = new SQLConf
+ conf.setConfString(SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key, "100")
+ SegmentTreeWindowTestHelper.withFrameFactory(conf) { factory =>
+ val frame = factory.newFrame()
+ assert(!frame.fallbackAllocated, "no allocation before prepare()")
+ frame.prepare(factory.newArray(rows = 5))
+ assert(frame.fallbackUsed, "fallback path expected")
+ assert(frame.fallbackAllocated, "fallback allocated after first small partition")
+ }
+ }
+
+ test("lifecycle: segtree then small partition reuses frame, allocates fallback once") {
+ val conf = new SQLConf
+ conf.setConfString(SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key, "10")
+ SegmentTreeWindowTestHelper.withFrameFactory(conf) { factory =>
+ val frame = factory.newFrame()
+ // Partition 1: segtree path (20 rows >= 10).
+ frame.prepare(factory.newArray(rows = 20))
+ assert(!frame.fallbackUsed && !frame.fallbackAllocated)
+ // Partition 2: small partition.
+ frame.prepare(factory.newArray(rows = 5))
+ assert(frame.fallbackUsed, "fallback path expected on small partition")
+ assert(frame.fallbackAllocated, "fallback allocated on first small partition")
+ }
+ }
+
+ test("lifecycle: small then segtree transition drops fallback reference") {
+ val conf = new SQLConf
+ conf.setConfString(SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key, "10")
+ SegmentTreeWindowTestHelper.withFrameFactory(conf) { factory =>
+ val frame = factory.newFrame()
+ // Partition 1: small -> fallback allocated.
+ frame.prepare(factory.newArray(rows = 5))
+ assert(frame.fallbackAllocated && frame.fallbackUsed)
+ // Partition 2: segtree -> fallback must be dropped (row-copy buffer GC-eligible).
+ frame.prepare(factory.newArray(rows = 20))
+ assert(!frame.fallbackUsed, "segtree path expected")
+ assert(!frame.fallbackAllocated,
+ "retained fallback reference must be dropped on segtree-path re-entry")
+ }
+ }
+
+ test("lifecycle: throwing fallback.prepare must leave fallbackUsed=false") {
+ // Structural-witness test: frame is `final`, can't inject a throwing
+ // fallback. `fallbackUsed = true` is set ONLY after `fallback.prepare`
+ // returns normally (see SegmentTreeWindowFunctionFrame.prepare). A
+ // regression of that ordering would fail T1d. This test re-exercises
+ // the happy path and documents the invariant.
+ val conf = new SQLConf
+ conf.setConfString(SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key, "100")
+ SegmentTreeWindowTestHelper.withFrameFactory(conf) { factory =>
+ val frame = factory.newFrame()
+ frame.prepare(factory.newArray(rows = 5))
+ assert(frame.fallbackUsed, "happy-path post-condition")
+ }
+ }
+
+ test("lifecycle: close() after only small partition is a no-op") {
+ val conf = new SQLConf
+ conf.setConfString(SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key, "100")
+ SegmentTreeWindowTestHelper.withFrameFactory(conf) { factory =>
+ val frame = factory.newFrame()
+ frame.prepare(factory.newArray(rows = 5))
+ assert(frame.fallbackUsed)
+ // Idempotent: no tree was ever built.
+ frame.close()
+ frame.close()
+ }
+ }
+}
+
+/** Legacy UDAF wrapped as ScalaUDAF (ImperativeAggregate); rejected by segtree guard. */
+private class LegacySumUdaf extends UserDefinedAggregateFunction {
+ override def inputSchema: StructType = new StructType().add("v", LongType)
+ override def bufferSchema: StructType = new StructType().add("s", LongType)
+ override def dataType: DataType = LongType
+ override def deterministic: Boolean = true
+ override def initialize(buffer: MutableAggregationBuffer): Unit = {
+ buffer(0) = 0L
+ }
+ override def update(buffer: MutableAggregationBuffer, input: Row): Unit = {
+ if (!input.isNullAt(0)) {
+ buffer(0) = buffer.getLong(0) + input.getLong(0)
+ }
+ }
+ override def merge(buffer1: MutableAggregationBuffer, buffer2: Row): Unit = {
+ buffer1(0) = buffer1.getLong(0) + buffer2.getLong(0)
+ }
+ override def evaluate(buffer: Row): Any = buffer.getLong(0)
+}
+
+/** Typed Aggregator wrapped as ScalaAggregator (TypedImperativeAggregate); rejected. */
+private class LongSumAggregator extends Aggregator[Long, Long, Long] {
+ override def zero: Long = 0L
+ override def reduce(b: Long, a: Long): Long = b + a
+ override def merge(b1: Long, b2: Long): Long = b1 + b2
+ override def finish(r: Long): Long = r
+ override def bufferEncoder: Encoder[Long] = Encoders.scalaLong
+ override def outputEncoder: Encoder[Long] = Encoders.scalaLong
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowMetricsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowMetricsSuite.scala
new file mode 100644
index 000000000000..5c3edb7e06c6
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowMetricsSuite.scala
@@ -0,0 +1,213 @@
+/*
+ * 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.window
+
+import org.apache.spark.sql.QueryTest
+import org.apache.spark.sql.execution.SparkPlanInfo
+import org.apache.spark.sql.execution.metric.SQLMetricsTestUtils
+import org.apache.spark.sql.execution.ui.SparkPlanGraph
+import org.apache.spark.sql.expressions.Window
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+/**
+ * SQLMetrics visibility coverage for [[SegmentTreeWindowFunctionFrame]]:
+ * segtree path bumps `numSegmentTreeFrames`; fallback path bumps
+ * `numSegmentTreeFallbackFrames`; feature-flag off leaves both at 0.
+ */
+class SegmentTreeWindowMetricsSuite
+ extends QueryTest with SharedSparkSession with SQLMetricsTestUtils {
+
+ import testImplicits._
+
+ /**
+ * Run `df` and return the first Window node's seg-tree counter values as
+ * Long. (Other metrics like `spillSize` are UI-pretty-printed strings and
+ * are skipped.)
+ */
+ private def windowMetricValues(df: org.apache.spark.sql.DataFrame): Map[String, Long] = {
+ val previousExecutionIds = currentExecutionIds()
+ df.collect()
+ sparkContext.listenerBus.waitUntilEmpty(10000)
+ val executionId = currentExecutionIds().diff(previousExecutionIds).head
+ val metricValues = statusStore.executionMetrics(executionId)
+ val graph = SparkPlanGraph(SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan))
+ val windowNode = graph.allNodes.find(_.name == "Window").getOrElse {
+ fail(s"No Window node in plan:\n${df.queryExecution.executedPlan}")
+ }
+ val wanted = Set(
+ "number of segment-tree frames prepared",
+ "number of segment-tree fallback frames prepared")
+ windowNode.metrics.filter(m => wanted.contains(m.name)).map { m =>
+ // UI value may be "3" or "total (min, med, max ...)"; first int run is the total.
+ val raw = metricValues(m.accumulatorId)
+ val total = "-?\\d+".r.findFirstIn(raw).getOrElse {
+ fail(s"Could not parse integer from metric '${m.name}' value: $raw")
+ }
+ m.name -> total.toLong
+ }.toMap
+ }
+
+ private def baseDF = spark.range(0, 120).selectExpr(
+ "id", "(id % 3) AS pk", "CAST(id AS INT) AS v")
+
+ private val winSpec = Window.partitionBy($"pk").orderBy($"id").rowsBetween(-3, 3)
+
+ test("segment-tree path increments numSegmentTreeFrames (one per frame per partition)") {
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ // Two aggregates share one Window -> one frame; 3 partitions => 3 frames.
+ val df = baseDF.select($"id",
+ min($"v").over(winSpec).as("mn"),
+ max($"v").over(winSpec).as("mx"))
+ val m = windowMetricValues(df)
+ assert(m("number of segment-tree frames prepared") === 3L,
+ s"expected 3 segtree frames (one per partition), got metrics = $m")
+ assert(m("number of segment-tree fallback frames prepared") === 0L,
+ s"fallback counter must be 0 when all partitions take segtree path, got $m")
+ }
+ }
+
+ test("fallback path increments numSegmentTreeFallbackFrames") {
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ // Threshold > partition size (40 rows/partition) forces fallback on every partition.
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "1000",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = baseDF.select($"id", min($"v").over(winSpec).as("mn"))
+ val m = windowMetricValues(df)
+ assert(m("number of segment-tree fallback frames prepared") === 3L,
+ s"expected 3 fallback frames (one per partition under threshold), got $m")
+ assert(m("number of segment-tree frames prepared") === 0L,
+ s"segtree counter must be 0 when all partitions fall back, got $m")
+ }
+ }
+
+ test("feature flag off: both segment-tree counters stay at zero") {
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "false",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ val df = baseDF.select($"id", min($"v").over(winSpec).as("mn"))
+ val m = windowMetricValues(df)
+ assert(m("number of segment-tree frames prepared") === 0L,
+ s"segtree counter must be 0 when feature flag disabled, got $m")
+ assert(m("number of segment-tree fallback frames prepared") === 0L,
+ s"fallback counter must be 0 when feature flag disabled, got $m")
+ }
+ }
+
+ // T1-T4: regression for the removed idempotency guard in
+ // `SegmentTreeWindowFunctionFrame.prepare()`. The old guard keyed on
+ // `(System.identityHashCode(rows), rows.length)`, but `WindowPartitionEvaluator`
+ // reuses a single `ExternalAppendOnlyUnsafeRowArray` across all partitions in
+ // a task, so the identity hash is constant and the key collapsed to
+ // `rows.length` -- silently deduping consecutive equal-length partitions.
+ // All four tests pin `minPartitionRows = 64` and use a single shuffle
+ // partition so the Window operator sees `numPartitions > numTasks` -- the
+ // exact shape the existing segtree/fallback fixtures happened to avoid.
+
+ test("T1 (G1) numPartitions > numTasks, identical length: every partition counted") {
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "64",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ // 3 window partitions x 100 rows, all >= 64 => segtree, one task.
+ val df = spark.range(0, 300).selectExpr("id", "(id % 3) AS pk", "CAST(id AS INT) AS v")
+ .select($"id", min($"v").over(
+ Window.partitionBy($"pk").orderBy($"id").rowsBetween(-3, 3)).as("mn"))
+ val m = windowMetricValues(df)
+ assert(m("number of segment-tree frames prepared") === 3L,
+ s"expected 3 segtree frames (one per window partition), got $m")
+ assert(m("number of segment-tree fallback frames prepared") === 0L, s"got $m")
+ }
+ }
+
+ test("T2 (G2) identical-length partitions across keys") {
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "64",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ // Two distinct keys, each exactly 200 rows (> 64 => segtree).
+ // Old guard would dedupe the second because `rows.length` matches.
+ val df = spark.range(0, 400)
+ .selectExpr("id", "CASE WHEN id < 200 THEN 'a' ELSE 'b' END AS pk",
+ "CAST(id AS INT) AS v")
+ .select($"id", min($"v").over(
+ Window.partitionBy($"pk").orderBy($"id").rowsBetween(-3, 3)).as("mn"))
+ val m = windowMetricValues(df)
+ assert(m("number of segment-tree frames prepared") === 2L,
+ s"expected 2 segtree frames (one per key, both length 200), got $m")
+ assert(m("number of segment-tree fallback frames prepared") === 0L, s"got $m")
+ }
+ }
+
+ test("T3 (G3) all-length-1 unique keys, fallback path: every partition counted") {
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "64",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ // 100 unique keys, 1 row each, all < 64 => fallback. Old guard
+ // would collapse to 1 because every partition has length 1.
+ val df = spark.range(0, 100).selectExpr("id AS pk", "CAST(id AS INT) AS v")
+ .select($"pk", min($"v").over(
+ Window.partitionBy($"pk").orderBy($"pk").rowsBetween(-3, 3)).as("mn"))
+ val m = windowMetricValues(df)
+ assert(m("number of segment-tree fallback frames prepared") === 100L,
+ s"expected 100 fallback frames (one per unique-key partition), got $m")
+ assert(m("number of segment-tree frames prepared") === 0L, s"got $m")
+ }
+ }
+
+ test("T4 (G4) mixed segtree + fallback, non-aliasing order") {
+ withSQLConf(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "64",
+ SQLConf.SHUFFLE_PARTITIONS.key -> "1",
+ SQLConf.ADAPTIVE_EXECUTION_ENABLED.key -> "false") {
+ // 4 window partitions with lengths (50, 50, 100, 100) in sort order.
+ // 50 < 64 => fallback; 100 >= 64 => segtree. Under the old guard,
+ // consecutive same-length partitions collapsed => (fb=1, seg=1);
+ // after the fix => (fb=2, seg=2).
+ // TRAP: an alternating (50,100,50,100) ordering coincidentally gives
+ // (2,2) under both code paths (key changes every step) -- false
+ // positive, do not use. Size-monotone ordering is required.
+ val df = spark.range(0, 300)
+ .selectExpr(
+ "id",
+ // key ordering k1<k2<k3<k4 matches size-monotone (50,50,100,100)
+ "CASE WHEN id < 50 THEN 'k1' " +
+ "WHEN id < 100 THEN 'k2' " +
+ "WHEN id < 200 THEN 'k3' " +
+ "ELSE 'k4' END AS pk",
+ "CAST(id AS INT) AS v")
+ .select($"id", min($"v").over(
+ Window.partitionBy($"pk").orderBy($"id").rowsBetween(-3, 3)).as("mn"))
+ val m = windowMetricValues(df)
+ assert(m("number of segment-tree frames prepared") === 2L,
+ s"expected 2 segtree frames (k3, k4 @ 100 rows each), got $m")
+ assert(m("number of segment-tree fallback frames prepared") === 2L,
+ s"expected 2 fallback frames (k1, k2 @ 50 rows each), got $m")
+ }
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowTestHelper.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowTestHelper.scala
new file mode 100644
index 000000000000..cd5237c9b310
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/SegmentTreeWindowTestHelper.scala
@@ -0,0 +1,176 @@
+/*
+ * 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.window
+
+import org.apache.spark.{SparkEnv, TaskContext}
+import org.apache.spark.memory.MemoryTestingUtils
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions._
+import org.apache.spark.sql.catalyst.expressions.aggregate.Sum
+import org.apache.spark.sql.catalyst.expressions.codegen.GenerateMutableProjection
+import org.apache.spark.sql.execution.ExternalAppendOnlyUnsafeRowArray
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.types.IntegerType
+
+/**
+ * Shared helpers for segment-tree window-frame tests. Kept in the same
+ * package so tests can reach `private[window]` hooks on
+ * [[SegmentTreeWindowFunctionFrame]] (see contract Section 1.3).
+ */
+private[window] object SegmentTreeWindowTestHelper {
+
+ /**
+ * Fake-TaskContext plumbing for lifecycle tests. Gives `body` a factory
+ * producing `newArray(rows)` and `newFrame()` (unprepared — caller drives
+ * prepare/write/close). Tracks frames so teardown always closes them.
+ */
+ def withFrameFactory(conf: SQLConf)
+ (body: FrameFactory => Unit): Unit = {
+ val existing = TaskContext.get()
+ val owned = existing == null
+ val tc = if (owned) {
+ val fake = MemoryTestingUtils.fakeTaskContext(SparkEnv.get)
+ TaskContext.setTaskContext(fake)
+ fake
+ } else existing
+ val factory = new FrameFactory(conf, tc)
+ try body(factory) finally {
+ factory.closeAll()
+ if (owned) TaskContext.unset()
+ }
+ }
+
+ final class FrameFactory(conf: SQLConf, tc: TaskContext) {
+ private val attr = AttributeReference("v", IntegerType, nullable = false)()
+ private val input = Seq[Attribute](attr)
+ private val fn = Sum(attr)
+ private val bufAttrs = fn.aggBufferAttributes
+ private val processor = AggregateProcessor(
+ Array[Expression](fn),
+ 0,
+ input,
+ (es, s) => GenerateMutableProjection.generate(es, s),
+ Array[Option[Expression]](None))
+ private val unsafeProj = UnsafeProjection.create(Array(attr.dataType))
+ private val opened = scala.collection.mutable.Buffer[AutoCloseable]()
+
+ def newArray(rows: Int): ExternalAppendOnlyUnsafeRowArray = {
+ val array = new ExternalAppendOnlyUnsafeRowArray(
+ tc.taskMemoryManager(),
+ SparkEnv.get.blockManager,
+ SparkEnv.get.serializerManager,
+ tc,
+ 1024,
+ SparkEnv.get.memoryManager.pageSizeBytes,
+ Int.MaxValue,
+ Long.MaxValue,
+ Int.MaxValue,
+ Long.MaxValue)
+ var i = 0
+ while (i < rows) {
+ array.add(unsafeProj(new GenericInternalRow(Array[Any](i))))
+ i += 1
+ }
+ array
+ }
+
+ /** Create a new frame. Caller owns lifecycle unless tracked via `track()`. */
+ def newFrame(): SegmentTreeWindowFunctionFrame = {
+ val target = new SpecificInternalRow(Seq(bufAttrs.head.dataType))
+ val frame = new SegmentTreeWindowFunctionFrame(
+ target,
+ processor,
+ Array(fn),
+ input,
+ RowFrame,
+ RowBoundOrdering(-1),
+ RowBoundOrdering(1),
+ (es, s) => GenerateMutableProjection.generate(es, s),
+ conf,
+ None,
+ tc.taskMemoryManager())
+ track(frame)
+ frame
+ }
+
+ def track[T <: AutoCloseable](c: T): T = { opened += c; c }
+
+ def closeAll(): Unit = {
+ opened.foreach { c =>
+ try c.close() catch { case _: Throwable => () }
+ }
+ opened.clear()
+ }
+ }
+
+ /**
+ * Build a small-partition [[SegmentTreeWindowFunctionFrame]] (Sum over one
+ * IntegerType column, values 0..rows-1, frame -1..+1), drive it through
+ * all rows, then hand it to `body`. Manages a fake TaskContext if needed
+ * and always closes the frame.
+ */
+ def withSmallPartitionFrame(conf: SQLConf, rows: Int)
+ (body: SegmentTreeWindowFunctionFrame => Unit): Unit = {
+ withFrameFactory(conf) { factory =>
+ val frame = factory.newFrame()
+ val array = factory.newArray(rows)
+ frame.prepare(array)
+ var j = 0
+ while (j < rows) {
+ frame.write(j, InternalRow.empty)
+ j += 1
+ }
+ body(frame)
+ }
+ }
+
+ /**
+ * Build the given [[WindowSegmentTree]] from a row iterator. After the
+ * buffer-ownership flip the caller (this helper) materialises an
+ * [[ExternalAppendOnlyUnsafeRowArray]], projects rows via
+ * [[UnsafeProjection]], and hands it to the tree. The tree retains the
+ * array; close is idempotent via `tree.close()` / partition teardown.
+ *
+ * `inMemoryThreshold` / `spillThreshold` let tests exercise the spill
+ * path (see `WindowSegmentTreeSuite` D9).
+ */
+ def buildTreeFromIter(
+ tree: WindowSegmentTree,
+ rows: Iterator[InternalRow],
+ inputSchema: Seq[Attribute],
+ inMemoryThreshold: Int = Int.MaxValue,
+ spillThreshold: Int = Int.MaxValue): Unit = {
+ val tc = TaskContext.get()
+ require(tc != null, "buildTreeFromIter requires an active TaskContext")
+ val env = SparkEnv.get
+ val array = new ExternalAppendOnlyUnsafeRowArray(
+ tc.taskMemoryManager(),
+ env.blockManager,
+ env.serializerManager,
+ tc,
+ 1024,
+ env.memoryManager.pageSizeBytes,
+ inMemoryThreshold,
+ Long.MaxValue,
+ spillThreshold,
+ Long.MaxValue)
+ val proj = UnsafeProjection.create(inputSchema.map(_.dataType).toArray)
+ rows.foreach(r => array.add(proj(r)))
+ tree.build(array)
+ }
+}
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
new file mode 100644
index 000000000000..d3da8e40c85a
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeAllowlistSuite.scala
@@ -0,0 +1,187 @@
+/*
+ * 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.window
+
+import org.apache.spark.sql.{DataFrame, QueryTest}
+import org.apache.spark.sql.execution.SparkPlanInfo
+import org.apache.spark.sql.execution.metric.SQLMetricsTestUtils
+import org.apache.spark.sql.execution.ui.SparkPlanGraph
+import org.apache.spark.sql.expressions.Window
+import org.apache.spark.sql.functions._
+import org.apache.spark.sql.internal.SQLConf
+import org.apache.spark.sql.test.SharedSparkSession
+
+/**
+ * Coverage for the explicit segment-tree aggregate allowlist
+ * ([[WindowSegmentTree.EligibleAggregates]]):
+ * - allowlisted aggregates route to segtree (`numSegmentTreeFrames` bumps).
+ * - non-allowlisted aggregates fall through to the sliding path without
+ * crashing or producing wrong results (segtree counters stay at 0).
+ * Eligibility gating only; exhaustive equivalence lives in
+ * [[SegmentTreeWindowFunctionSuite]].
+ */
+class WindowSegmentTreeAllowlistSuite
+ extends QueryTest with SharedSparkSession with SQLMetricsTestUtils {
+
+ import testImplicits._
+
+ private val enableSegTree: Map[String, String] = Map(
+ SQLConf.WINDOW_SEGMENT_TREE_ENABLED.key -> "true",
+ SQLConf.WINDOW_SEGMENT_TREE_MIN_PARTITION_ROWS.key -> "1")
+
+ private def baseDF = spark.range(0, 120).selectExpr(
+ "id", "(id % 3) AS pk", "CAST(id AS INT) AS v", "CAST(id AS DOUBLE) AS vd")
+
+ private val winSpec = Window.partitionBy($"pk").orderBy($"id").rowsBetween(-3, 3)
+
+ private def segTreeCounters(df: DataFrame): (Long, Long) = {
+ val previousExecutionIds = currentExecutionIds()
+ df.collect()
+ sparkContext.listenerBus.waitUntilEmpty(10000)
+ val executionId = currentExecutionIds().diff(previousExecutionIds).head
+ val metricValues = statusStore.executionMetrics(executionId)
+ val graph = SparkPlanGraph(SparkPlanInfo.fromSparkPlan(df.queryExecution.executedPlan))
+ val windowNode = graph.allNodes.find(_.name == "Window").getOrElse {
+ fail(s"No Window node in plan:\n${df.queryExecution.executedPlan}")
+ }
+ def total(name: String): Long = {
+ val mOpt = windowNode.metrics.find(_.name == name)
+ mOpt.map { m =>
+ val raw = metricValues(m.accumulatorId)
+ "-?\\d+".r.findFirstIn(raw).map(_.toLong).getOrElse(0L)
+ }.getOrElse(0L)
+ }
+ (total("number of segment-tree frames prepared"),
+ total("number of segment-tree fallback frames prepared"))
+ }
+
+ // Positive: allowlisted aggregates route to segtree
+
+ Seq(
+ ("min", (c: org.apache.spark.sql.Column) => min(c)),
+ ("max", (c: org.apache.spark.sql.Column) => max(c)),
+ ("sum", (c: org.apache.spark.sql.Column) => sum(c)),
+ ("count", (c: org.apache.spark.sql.Column) => count(c)),
+ ("avg", (c: org.apache.spark.sql.Column) => avg(c)),
+ ("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))
+ ).foreach { case (name, fn) =>
+ test(s"$name routes to the segment-tree path") {
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val df = baseDF.withColumn("agg", fn($"vd").over(winSpec))
+ val (seg, fallback) = segTreeCounters(df)
+ assert(seg > 0, s"$name should bump numSegmentTreeFrames (got $seg)")
+ assert(fallback == 0,
+ s"$name should not fall back (fallback counter: $fallback)")
+ }
+ }
+ }
+
+ // 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))
+ val (seg, _) = segTreeCounters(df)
+ assert(seg == 0, s"collect_list should not use segment tree (got $seg frames)")
+ }
+ }
+
+ test("collect_set falls through (unbounded buffer)") {
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val df = baseDF.withColumn("agg", collect_set($"v").over(winSpec))
+ val (seg, _) = segTreeCounters(df)
+ assert(seg == 0, s"collect_set should not use segment tree (got $seg frames)")
+ }
+ }
+
+ test("approx_count_distinct (HyperLogLog++) falls through (fail-closed)") {
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val df = baseDF.withColumn("agg", approx_count_distinct($"v").over(winSpec))
+ val (seg, _) = segTreeCounters(df)
+ assert(seg == 0,
+ s"approx_count_distinct is intentionally not on the allowlist (got $seg frames)")
+ }
+ }
+
+ test("percentile_approx falls through (sketch buffer not auditable)") {
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val df = baseDF.withColumn(
+ "agg", percentile_approx($"vd", lit(0.5), lit(100)).over(winSpec))
+ val (seg, _) = segTreeCounters(df)
+ assert(seg == 0, s"percentile_approx should not use segment tree (got $seg frames)")
+ }
+ }
+
+ // Gate: aggregates carrying a FILTER (WHERE ...) clause fall through.
+ // `eligibleForSegTree` requires `filters.forall(_.isEmpty)` because the
+ // segment-tree combine contract is defined over the unfiltered partial
+ // buffer. A defensive regression: if any future analyzer rule ever rewrites
+ // `AGG(x) FILTER (WHERE p)` in a way that strips `AggregateExpression.filter`
+ // (e.g., pushing the predicate into the aggregate function), this test
+ // fails and forces an explicit eligibility review.
+ test("FILTER (WHERE ...) disables segment-tree path") {
+ withSQLConf(enableSegTree.toSeq: _*) {
+ withTempView("t") {
+ baseDF.createOrReplaceTempView("t")
+ val df = spark.sql(
+ """SELECT id, pk, v,
+ | sum(v) FILTER (WHERE v % 2 = 0)
+ | OVER (PARTITION BY pk ORDER BY id ROWS BETWEEN 3 PRECEDING AND 3 FOLLOWING)
+ | AS filtered_sum
+ |FROM t""".stripMargin)
+ val (seg, _) = segTreeCounters(df)
+ assert(seg == 0,
+ s"filtered aggregate must not take segment-tree path (got $seg segtree frames)")
+ }
+ }
+ }
+
+ // Mixed: ANY non-eligible aggregate disqualifies the group
+
+ test("mix of allowlisted + non-allowlisted aggregates falls through entirely") {
+ withSQLConf(enableSegTree.toSeq: _*) {
+ val df = baseDF
+ .withColumn("s", sum($"v").over(winSpec))
+ .withColumn("fv", first($"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.
+ assert(seg == 0,
+ s"Window group containing a non-allowlisted agg must fall through (got $seg)")
+ }
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeMemorySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeMemorySuite.scala
new file mode 100644
index 000000000000..e60253042c17
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeMemorySuite.scala
@@ -0,0 +1,351 @@
+/*
+ * 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.window
+
+import java.util.Properties
+
+import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkFunSuite, TaskContext, TaskContextImpl}
+import org.apache.spark.executor.TaskMetrics
+import org.apache.spark.memory.{MemoryConsumer, MemoryMode, SparkOutOfMemoryError, TaskMemoryManager, TestMemoryManager}
+import org.apache.spark.metrics.MetricsSystem
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression, GenericInternalRow, MutableProjection, SpecificInternalRow}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{DeclarativeAggregate, Min}
+import org.apache.spark.sql.catalyst.expressions.codegen.GenerateMutableProjection
+import org.apache.spark.sql.types.IntegerType
+
+/**
+ * Memory-manager integration tests for [[WindowSegmentTree]]. Covers:
+ * - `SegTreeSpiller` registration with `TaskMemoryManager`
+ * - `acquireBlockMemory` grant / partial-grant rollback
+ * - `evictUntil` LRU eviction driven by TMM pressure
+ * - `spill()` self-trigger short-circuit and rowArray-spilled fall-through
+ * - task completion / kill listener releasing all cached blocks
+ * T5 (rowArray-spilled priority) and T8 (task-kill listener) are kept as
+ * `ignore`d stubs so the matrix stays visible; each documents what it needs.
+ */
+class WindowSegmentTreeMemorySuite extends SparkFunSuite with LocalSparkContext {
+
+ // common fixtures
+
+ private val inputAttr: AttributeReference =
+ AttributeReference("v", IntegerType, nullable = true)()
+ private val inputSchema: Seq[Attribute] = Seq(inputAttr)
+
+ private def newMutableProjection
+ : (Seq[Expression], Seq[Attribute]) => MutableProjection =
+ (exprs, attrs) => GenerateMutableProjection.generate(exprs, attrs)
+
+ private def minAgg: DeclarativeAggregate = Min(inputAttr)
+
+ /**
+ * Standalone `TaskMemoryManager` backed by `TestMemoryManager` (not routed
+ * through `SparkEnv`), with a matching `TaskContextImpl` installed.
+ * Restores the previous `TaskContext` on exit.
+ *
+ * @param budget initial execution memory; `Long.MaxValue` for T2,
+ * tight value (e.g. 2 * blockBytes) for T3.
+ * @param offHeap when true, enables Tungsten off-heap mode (T9).
+ */
+ private def withTmm[T](budget: Long = Long.MaxValue, offHeap: Boolean = false)
+ (body: (TaskMemoryManager, TestMemoryManager) => T): T = {
+ val conf = new SparkConf(false)
+ .set("spark.memory.offHeap.enabled", offHeap.toString)
+ .set("spark.memory.offHeap.size", "1048576")
+ if (sc == null) {
+ sc = new SparkContext("local", "WindowSegmentTreeMemorySuite", conf)
+ }
+ val mm = new TestMemoryManager(conf)
+ mm.limit(budget)
+ val tmm = new TaskMemoryManager(mm, 0)
+ val prev = TaskContext.get()
+ val tc = new TaskContextImpl(
+ stageId = 0, stageAttemptNumber = 0, partitionId = 0,
+ taskAttemptId = 0, attemptNumber = 0, numPartitions = 1,
+ taskMemoryManager = tmm,
+ localProperties = new Properties,
+ metricsSystem = null.asInstanceOf[MetricsSystem],
+ taskMetrics = TaskMetrics.empty,
+ cpus = 1)
+ TaskContext.setTaskContext(tc)
+ try body(tmm, mm)
+ finally {
+ if (prev != null) TaskContext.setTaskContext(prev) else TaskContext.unset()
+ }
+ }
+
+ private def buildTree(
+ tmm: TaskMemoryManager,
+ values: Seq[Int],
+ fanout: Int = 4,
+ blockSize: Int = 8,
+ maxCachedBlocks: Option[Int] = Some(2)): WindowSegmentTree = {
+ val tree = new WindowSegmentTree(
+ Array(minAgg), inputSchema, newMutableProjection,
+ fanout = fanout, blockSize = blockSize,
+ maxCachedBlocks = maxCachedBlocks,
+ taskMemoryManager = tmm)
+ val rows = values.iterator.map { v =>
+ val r = new GenericInternalRow(1); r.update(0, v); r.asInstanceOf[InternalRow]
+ }
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, rows, inputSchema)
+ tree
+ }
+
+ private def newOutBuffer(): SpecificInternalRow =
+ new SpecificInternalRow(Seq[org.apache.spark.sql.types.DataType](IntegerType))
+
+ private def queryMin(tree: WindowSegmentTree, lo: Int, hi: Int): Any = {
+ val out = newOutBuffer()
+ tree.query(lo, hi, out)
+ if (out.isNullAt(0)) null else out.getInt(0)
+ }
+
+ private def naiveMin(vs: Seq[Int], lo: Int, hi: Int): Any =
+ if (lo >= hi) null else vs.slice(lo, hi).min
+
+ test("T1 constructor rejects null TaskMemoryManager") {
+ val ex = intercept[IllegalArgumentException] {
+ new WindowSegmentTree(
+ Array(minAgg), inputSchema, newMutableProjection,
+ taskMemoryManager = null)
+ }
+ assert(ex.getMessage.contains("non-null TaskMemoryManager"))
+ }
+
+ test("T2 ample budget: query correctness and positive memory usage") {
+ withTmm() { (tmm, _) =>
+ val values = Seq(5, 2, 9, 1, 7, 3, 4, 8, 6, 0, 11, 12, 13, 14, 15, 16)
+ val tree = buildTree(tmm, values, fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(4))
+ try {
+ // Force block-level queries so the LRU actually populates.
+ assert(queryMin(tree, 0, values.length) == values.min)
+ assert(queryMin(tree, 1, 14) == naiveMin(values, 1, 14))
+ assert(queryMin(tree, 5, 11) == naiveMin(values, 5, 11))
+ // Positive consumption: at least one block level cached -> spiller acquired.
+ assert(tmm.getMemoryConsumptionForThisTask > 0L,
+ "Expected positive memory consumption after caching block levels")
+ } finally tree.close()
+ // All bytes released after close.
+ assert(tmm.getMemoryConsumptionForThisTask == 0L,
+ "Memory consumption must return to 0 after close()")
+ }
+ }
+
+ test("T3 tight budget forces spill: results still match baseline") {
+ val values = (0 until 40).map(i => (i * 37 + 11) % 97)
+ // Baseline under ample budget.
+ val baseline: Seq[Any] = {
+ var captured: Seq[Any] = Seq.empty
+ withTmm() { (tmm, _) =>
+ val tree = buildTree(tmm, values, fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(3))
+ try {
+ captured = (0 to values.length by 5).map(i => queryMin(tree, 0, i))
+ } finally {
+ tree.close()
+ }
+ }
+ captured
+ }
+ // Tight budget: ~1 block headroom -> any new block-level load evicts the prior.
+ // Sizing: (fanout=4, blockSize=4) -> level shape [4, 1] slots;
+ // at bufferWidth=16 B/slot blockBytes = 5 * 16 = 80 B.
+ // budget=128 B admits exactly one block and forces spill on the second.
+ withTmm(budget = 128L) { (tmm, _) =>
+ val tree = buildTree(tmm, values, fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(3))
+ try {
+ val observed = (0 to values.length by 5).map(i => queryMin(tree, 0, i))
+ assert(observed == baseline,
+ s"spill-path answers diverged: observed=$observed baseline=$baseline")
+ } finally tree.close()
+ assert(tmm.getMemoryConsumptionForThisTask == 0L)
+ }
+ }
+
+ test("T4 self-trigger: spill(_, this) returns 0 and does not evict") {
+ withTmm() { (tmm, _) =>
+ val tree = buildTree(tmm, (0 until 16), fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(3))
+ try {
+ // Warm the cache.
+ assert(queryMin(tree, 0, 16) == 0)
+ val before = tmm.getMemoryConsumptionForThisTask
+ assert(before > 0L, "precondition: some block levels must be cached")
+ val spiller = tree.testOnlySpiller()
+ val freed = spiller.spill(Long.MaxValue, spiller)
+ assert(freed == 0L, s"self-trigger spill must return 0L, got $freed")
+ val after = tmm.getMemoryConsumptionForThisTask
+ assert(after == before,
+ s"cache size must not change on self-trigger (before=$before after=$after)")
+ } finally tree.close()
+ }
+ }
+
+ test("T6 close() is idempotent and releases all acquired bytes") {
+ withTmm() { (tmm, _) =>
+ val tree = buildTree(tmm, (0 until 20), fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(3))
+ assert(queryMin(tree, 0, 20) == 0)
+ val peak = tmm.getMemoryConsumptionForThisTask
+ assert(peak > 0L)
+ tree.close()
+ assert(tmm.getMemoryConsumptionForThisTask == 0L,
+ "Memory must be fully released after first close()")
+ // Second close must be a no-op.
+ tree.close()
+ assert(tmm.getMemoryConsumptionForThisTask == 0L,
+ "Second close() must remain a no-op")
+ }
+ }
+
+ test("T7 prepare mid-way failure releases all previously acquired blocks") {
+ // Inject failure by pre-setting consequentOOM to a moderate N so
+ // acquireMemory returns 0 after the first few successful blocks.
+ withTmm() { (tmm, mm) =>
+ val values = (0 until 40).map(i => 40 - i)
+ val tree = buildTree(tmm, values, fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(10))
+ try {
+ // Warm up: populate cache entries.
+ assert(queryMin(tree, 0, 40) == 1)
+ // Force next acquireMemory calls to fail, then trigger cold-cache via
+ // evict + re-acquire. Release everything so the next query re-acquires.
+ val spiller = tree.testOnlySpiller()
+ spiller.spill(Long.MaxValue, new MemoryConsumer(tmm,
+ tmm.pageSizeBytes(), MemoryMode.ON_HEAP) {
+ override def spill(size: Long, trigger: MemoryConsumer): Long = 0L
+ })
+ assert(tmm.getMemoryConsumptionForThisTask == 0L,
+ "after full spill, accounting must be zero")
+ // Force next N acquireMemory calls to grant 0 bytes (hard OOM).
+ // With maxCachedBlocks=10 and a cold cache, ensureBlockLevels will
+ // call acquireBlockMemory; initial grant and post-evict retry both
+ // see 0 -> SparkOutOfMemoryError.
+ mm.markConsequentOOM(10)
+ val ex = intercept[SparkOutOfMemoryError](queryMin(tree, 0, 20))
+ assert(ex.getMessage.contains("UNABLE_TO_ACQUIRE_MEMORY"),
+ s"unexpected OOM message: ${ex.getMessage}")
+ // Failed acquire must not leak bytes against the task.
+ assert(tmm.getMemoryConsumptionForThisTask == 0L,
+ "After failed acquire, accounting must be zero (no partial leaks)")
+ } finally tree.close()
+ assert(tmm.getMemoryConsumptionForThisTask == 0L)
+ }
+ }
+
+ test("T9 ON_HEAP Tungsten: spiller mode is ON_HEAP") {
+ withTmm(offHeap = false) { (tmm, _) =>
+ val tree = buildTree(tmm, (0 until 16), fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(2))
+ try {
+ assert(queryMin(tree, 0, 16) == 0)
+ assert(tree.testOnlySpiller().getMode == MemoryMode.ON_HEAP)
+ // In ON_HEAP Tungsten, our hardcoded mode also coincides with
+ // `tmm.getTungstenMemoryMode`, so segtree and rowArray share the
+ // pool and I8 (rowArray-spilled short-circuit) fires.
+ assert(tree.testOnlySpiller().getMode == tmm.getTungstenMemoryMode)
+ } finally tree.close()
+ }
+ }
+
+ test("T9 OFF_HEAP Tungsten: spiller mode stays ON_HEAP (not OFF_HEAP)") {
+ withTmm(offHeap = true) { (tmm, _) =>
+ val tree = buildTree(tmm, (0 until 16), fanout = 4, blockSize = 4,
+ maxCachedBlocks = Some(2))
+ try {
+ assert(queryMin(tree, 0, 16) == 0)
+ // Segtree's cache is JVM-heap (`SpecificInternalRow` /
+ // `Array[Array[InternalRow]]`), never Tungsten pages. The spiller
+ // must stay ON_HEAP regardless of `spark.memory.offHeap.enabled`,
+ // or `spill(n)` would phantom-credit the off-heap pool and
+ // violate TMM's same-pool spill contract. Mirrors `Spillable`.
+ assert(tree.testOnlySpiller().getMode == MemoryMode.ON_HEAP)
+ assert(tmm.getTungstenMemoryMode == MemoryMode.OFF_HEAP)
+ assert(tree.testOnlySpiller().getMode != tmm.getTungstenMemoryMode)
+ } finally tree.close()
+ }
+ }
+
+ test("T10 blockBytes oracle: hand-computed table + runtime-anchored cross-check") {
+ // Part A: hand-computed golden table. Values derived by listing the
+ // cached level widths from buildBlockLevels * 16 B/field. Independent of
+ // production `cachedSlotsPerBlock` -- fencepost or dropped-level
+ // regressions flip at least one case. Cases chosen to cover:
+ // blockSize == 1 (single leaf, no parents), (F=3,B=5) and (F=4,B=7)
+ // (non-power-of-fanout asymmetric tails), (F=16,B=65536) (deep tree).
+ // Level breakdown (leaves + parents + ... + root):
+ // (F=4, B=1) -> 1 = 1 slot
+ // (F=4, B=4) -> 4 + 1 = 5 slots
+ // (F=3, B=5) -> 5 + 2 + 1 = 8 slots
+ // (F=4, B=7) -> 7 + 2 + 1 = 10 slots
+ // (F=8, B=16) -> 16 + 2 + 1 = 19 slots
+ // (F=16, B=256) -> 256 + 16 + 1 = 273 slots
+ // (F=16, B=65536) -> 65536 + 4096 + 256 + 16 + 1 = 69905 slots
+ val expectedSlots: Seq[((Int, Int), Long)] = Seq(
+ (4, 1) -> 1L,
+ (4, 4) -> 5L,
+ (3, 5) -> 8L,
+ (4, 7) -> 10L,
+ (8, 16) -> 19L,
+ (16, 256) -> 273L,
+ (16, 65536) -> 69905L)
+
+ withTmm() { (tmm, _) =>
+ for (((fanout, blockSize), slots) <- expectedSlots) {
+ val tree = buildTree(tmm, Seq(1, 2, 3, 4, 5, 6, 7, 8),
+ fanout = fanout, blockSize = blockSize, maxCachedBlocks = Some(1))
+ try {
+ // Min on one IntegerType field -> 1 buffer field x 16 B.
+ val expected = math.max(1L, slots * 16L)
+ assert(tree.peekBlockBytes == expected,
+ s"blockBytes mismatch at (F=$fanout, blockSize=$blockSize): " +
+ s"got=${tree.peekBlockBytes} expected=$expected (slots=$slots)")
+ } finally tree.close()
+ }
+ }
+
+ // Part B: runtime-anchored cross-check. Build a block, read actual cached
+ // level array lengths via peek hooks, assert
+ // `blockBytes == sum(levels) * bufferWidth`. Guards against the formula
+ // and buildBlockLevels drifting apart.
+ withTmm() { (tmm, _) =>
+ val fanout = 4
+ val blockSize = 8
+ val tree = buildTree(tmm, (0 until blockSize),
+ fanout = fanout, blockSize = blockSize, maxCachedBlocks = Some(1))
+ try {
+ // Materialize block 0.
+ assert(queryMin(tree, 0, blockSize) == 0)
+ val levelCount = tree.peekLevelCount(0)
+ var slotSum = 0L
+ var lvl = 0
+ while (lvl < levelCount) {
+ slotSum += tree.peekLevelSize(0, lvl)
+ lvl += 1
+ }
+ val expected = math.max(1L, slotSum * 16L)
+ assert(tree.peekBlockBytes == expected,
+ s"runtime-anchored oracle mismatch: slots=$slotSum " +
+ s"expected=$expected got=${tree.peekBlockBytes}")
+ } finally tree.close()
+ }
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreePropertySuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreePropertySuite.scala
new file mode 100644
index 000000000000..fc92fbed05e2
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreePropertySuite.scala
@@ -0,0 +1,475 @@
+/*
+ * 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.window
+
+import org.scalacheck.Gen
+import org.scalatestplus.scalacheck.ScalaCheckDrivenPropertyChecks
+
+import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkEnv, SparkFunSuite, TaskContext}
+import org.apache.spark.memory.MemoryTestingUtils
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression, GenericInternalRow, MutableProjection, SpecificInternalRow}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{Average, Count, DeclarativeAggregate, Max, Min, StddevPop, StddevSamp, Sum}
+import org.apache.spark.sql.catalyst.expressions.codegen.GenerateMutableProjection
+import org.apache.spark.sql.types.{DataType, DoubleType, IntegerType, LongType}
+
+/**
+ * Property-based tests for [[WindowSegmentTree]] using ScalaCheck. Row-by-row
+ * naive aggregate is the parallel oracle: every property reduces to
+ * `segtree(input) == naive(input)`.
+ * Follows `ExpressionEvalHelper` convention (mixes in
+ * [[ScalaCheckDrivenPropertyChecks]]). All generated cases share one
+ * `SparkContext`/`TaskMemoryManager` by wrapping `forAll` inside a single
+ * `withTaskContext`; do NOT invert.
+ */
+class WindowSegmentTreePropertySuite extends SparkFunSuite
+ with LocalSparkContext with ScalaCheckDrivenPropertyChecks {
+
+ // ScalaCheck config: deterministic seed (ScalaCheck 1.14+ default).
+ // Quick tier: 100 cases for P1, 10 cases for P2 (see class doc).
+ implicit override val generatorDrivenConfig: PropertyCheckConfiguration =
+ PropertyCheckConfiguration(minSuccessful = 100, minSize = 0, sizeRange = 100)
+
+ // test harness (mirrors WindowSegmentTreeSuite)
+
+ private def newMutableProjection
+ : (Seq[Expression], Seq[Attribute]) => MutableProjection =
+ (exprs, attrs) => GenerateMutableProjection.generate(exprs, attrs)
+
+ private def withTaskContext[T](body: => T): T = {
+ sc = new SparkContext("local", "property-based-test", new SparkConf(false))
+ try {
+ val taskContext = MemoryTestingUtils.fakeTaskContext(SparkEnv.get)
+ TaskContext.setTaskContext(taskContext)
+ try body finally TaskContext.unset()
+ } finally {
+ if (sc != null) {
+ sc.stop()
+ sc = null
+ }
+ }
+ }
+
+ // Oracle
+ sealed trait AggKind
+ case object AggMin extends AggKind
+ case object AggMax extends AggKind
+ case object AggSum extends AggKind
+ case object AggCount extends AggKind
+
+ /**
+ * Naive oracle over [lo, hi). Empty frame => None for Min/Max/Sum
+ * (Spark returns NULL); for Count returns Some(0). null inputs are
+ * skipped (Spark Min/Max/Sum/Count semantics).
+ */
+ private def naiveAgg(
+ values: IndexedSeq[Option[Long]],
+ lo: Int, hi: Int,
+ kind: AggKind): Option[Long] = kind match {
+ case AggMin | AggMax =>
+ if (lo >= hi) None
+ else {
+ var result: Option[Long] = None
+ var i = lo
+ while (i < hi) {
+ values(i) match {
+ case Some(v) =>
+ result = result match {
+ case None => Some(v)
+ case Some(cur) => kind match {
+ case AggMin => if (v < cur) Some(v) else result
+ case AggMax => if (v > cur) Some(v) else result
+ case _ => result
+ }
+ }
+ case None =>
+ }
+ i += 1
+ }
+ result
+ }
+ case AggSum =>
+ // Spark Sum returns NULL when all inputs are null OR frame empty.
+ var sawNonNull = false
+ var s = 0L
+ var i = lo
+ while (i < hi) {
+ values(i) match {
+ case Some(v) => sawNonNull = true; s += v
+ case None =>
+ }
+ i += 1
+ }
+ if (sawNonNull) Some(s) else None
+ case AggCount =>
+ // Spark COUNT(col) over empty frame returns 0 (never null).
+ var c = 0L
+ var i = lo
+ while (i < hi) {
+ if (values(i).isDefined) c += 1
+ i += 1
+ }
+ Some(c)
+ }
+
+ // Input case ADT
+ private case class PbtCase(
+ values: IndexedSeq[Option[Long]],
+ dataType: DataType, // IntegerType or LongType
+ agg: AggKind,
+ blockSize: Int,
+ fanout: Int) {
+ def n: Int = values.length
+ }
+
+ // Generators
+ /** Weighted N: bias toward small boundaries (0,1,2 never hit under uniform). */
+ private val genN: Gen[Int] = Gen.frequency(
+ (2, Gen.choose(0, 4)),
+ (4, Gen.choose(5, 128)),
+ (4, Gen.choose(129, 5000))
+ )
+
+ /**
+ * 20% null, 80% bounded. For (AggSum, LongType) range shrunk to +/-1e12
+ * so worst-case sum (1e12 x 5000 = 5e15) stays inside Long bounds.
+ */
+ private def genValue(dt: DataType, agg: AggKind): Gen[Option[Long]] = {
+ val bounded: Gen[Long] = (dt, agg) match {
+ case (IntegerType, _) =>
+ Gen.choose(Int.MinValue.toLong, Int.MaxValue.toLong)
+ case (LongType, AggSum) =>
+ Gen.choose(-1000000000000L, 1000000000000L) // +/-1e12
+ case (LongType, _) =>
+ Gen.choose(Long.MinValue, Long.MaxValue)
+ case _ => throw new IllegalArgumentException(s"unsupported: $dt")
+ }
+ Gen.frequency(
+ (1, Gen.const(None: Option[Long])),
+ (4, bounded.map(Some(_)))
+ )
+ }
+
+ private val genBlockSize: Gen[Int] =
+ Gen.oneOf(1, 16, 32, 64, 256, 1024, 4096, 65536)
+
+ private val genFanout: Gen[Int] = Gen.oneOf(2, 3, 4, 8, 16)
+
+ private val genAgg: Gen[AggKind] =
+ Gen.oneOf(AggMin, AggMax, AggSum, AggCount)
+
+ private val genType: Gen[DataType] = Gen.oneOf(IntegerType, LongType)
+
+ private val genCase: Gen[PbtCase] = for {
+ dt <- genType
+ agg <- genAgg
+ n <- genN
+ values <- Gen.listOfN(n, genValue(dt, agg)).map(_.toIndexedSeq)
+ blockSize <- genBlockSize
+ fanout <- genFanout
+ } yield PbtCase(values, dt, agg, blockSize, fanout)
+
+ /** Frame [lo, hi) within [0, n]. */
+ private val genFrame: PbtCase => Gen[(Int, Int)] = { c =>
+ val n = c.n
+ for {
+ a <- Gen.choose(0, n)
+ b <- Gen.choose(0, n)
+ } yield if (a <= b) (a, b) else (b, a)
+ }
+
+ // Tree build/query glue
+ private def buildAttr(dt: DataType): AttributeReference =
+ AttributeReference("v", dt, nullable = true)()
+
+ /** Output type of the aggregate (segtree's queryInto materializes this). */
+ private def aggOutputType(c: PbtCase): DataType = c.agg match {
+ case AggMin | AggMax => c.dataType // same as input
+ case AggSum | AggCount => LongType // Sum(Int|Long)->Long; Count->Long
+ }
+
+ private def buildTreeFor(c: PbtCase): (WindowSegmentTree, AttributeReference) = {
+ val attr = buildAttr(c.dataType)
+ val schema: Seq[Attribute] = Seq(attr)
+ val agg: DeclarativeAggregate = c.agg match {
+ case AggMin => Min(attr)
+ case AggMax => Max(attr)
+ case AggSum => Sum(attr)
+ case AggCount => Count(Seq(attr))
+ }
+ val tree = new WindowSegmentTree(
+ Array(agg), schema, newMutableProjection,
+ c.fanout, c.blockSize, maxCachedBlocks = None,
+ taskMemoryManager = TaskContext.get().taskMemoryManager())
+ val rows = c.values.iterator.map { opt =>
+ val r = new GenericInternalRow(1)
+ opt match {
+ case Some(v) => c.dataType match {
+ case IntegerType => r.update(0, v.toInt)
+ case LongType => r.update(0, v)
+ case _ => throw new IllegalStateException
+ }
+ case None => r.setNullAt(0)
+ }
+ r.asInstanceOf[InternalRow]
+ }
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, rows, schema)
+ (tree, attr)
+ }
+
+ private def queryResult(
+ tree: WindowSegmentTree, outDt: DataType, lo: Int, hi: Int): Option[Long] = {
+ val out = new SpecificInternalRow(Seq[DataType](outDt))
+ tree.query(lo, hi, out)
+ if (out.isNullAt(0)) None
+ else outDt match {
+ case IntegerType => Some(out.getInt(0).toLong)
+ case LongType => Some(out.getLong(0))
+ case _ => throw new IllegalStateException
+ }
+ }
+
+ // Properties
+ /**
+ * P1 Equivalence: segtree query result equals naive aggregate over random
+ * frames. Probes multiple frames per case (empty / degenerate / full /
+ * random) to amplify block/fanout edge coverage.
+ */
+ test("P1 equivalence: segtree query equals naive aggregate over random frames") {
+ withTaskContext {
+ forAll(genCase) { c =>
+ val (tree, _) = buildTreeFor(c)
+ val outDt = aggOutputType(c)
+ try {
+ val frames = boundaryFrames(c.n) ++ randomFrames(c.n, k = 8, seed = c.n.toLong)
+ frames.foreach { case (lo, hi) =>
+ val expected = naiveAgg(c.values, lo, hi, c.agg)
+ val actual = queryResult(tree, outDt, lo, hi)
+ assert(actual == expected,
+ s"P1 mismatch: n=${c.n} agg=${c.agg} blockSize=${c.blockSize} " +
+ s"fanout=${c.fanout} dt=${c.dataType} frame=[$lo,$hi) " +
+ s"expected=$expected actual=$actual")
+ }
+ } finally tree.close()
+ }
+ }
+ }
+
+ /** Boundary frames: empty, full, head, tail, single-element. */
+ private def boundaryFrames(n: Int): Seq[(Int, Int)] = {
+ if (n == 0) Seq((0, 0))
+ else Seq(
+ (0, 0), // empty at start
+ (n, n), // empty at end
+ (0, n), // full
+ (0, 1), // single head
+ (n - 1, n), // single tail
+ (n / 2, n / 2), // empty middle
+ (0, n / 2), // prefix
+ (n / 2, n) // suffix
+ ).distinct
+ }
+
+ /** Deterministic random frames; seed tied to case to aid shrink repro. */
+ private def randomFrames(n: Int, k: Int, seed: Long): Seq[(Int, Int)] = {
+ if (n == 0) return Nil
+ val rnd = new scala.util.Random(seed)
+ Seq.fill(k) {
+ val a = rnd.nextInt(n + 1)
+ val b = rnd.nextInt(n + 1)
+ if (a <= b) (a, b) else (b, a)
+ }
+ }
+
+ /**
+ * P2 Determinism (10-case smoke): rebuild + requery yields identical
+ * results. Segtree has no RNG/hash iteration; acts as low-cost regression
+ * sentinel against future non-deterministic refactors.
+ */
+ test("P2 determinism: rebuild + requery yields identical results") {
+ withTaskContext {
+ forAll(genCase, minSuccessful(10)) { c =>
+ val frames = boundaryFrames(c.n) ++ randomFrames(c.n, k = 2, seed = 0xD2L)
+ val run1 = runFrames(c, frames)
+ val run2 = runFrames(c, frames)
+ assert(run1 == run2,
+ s"P2 mismatch: n=${c.n} agg=${c.agg} blockSize=${c.blockSize} " +
+ s"fanout=${c.fanout} dt=${c.dataType} run1=$run1 run2=$run2")
+ }
+ }
+ }
+
+ private def runFrames(c: PbtCase, frames: Seq[(Int, Int)]): Seq[Option[Long]] = {
+ val (tree, _) = buildTreeFor(c)
+ val outDt = aggOutputType(c)
+ try {
+ frames.map { case (lo, hi) => queryResult(tree, outDt, lo, hi) }
+ } finally tree.close()
+ }
+
+ // Floating-point family (Avg / StddevSamp / StddevPop).
+ // Welford merge order in segtree differs from row-by-row SlidingWindowFrame,
+ // so equality is NOT bit-exact; compared with relative tol 1e-9 + absolute
+ // floor for near-zero (e.g. constant-input stddev). Kept separate to avoid
+ // muddying the integer-family exact-equality contract.
+
+ sealed trait FpAggKind
+ case object FpAvg extends FpAggKind
+ case object FpStddevSamp extends FpAggKind
+ case object FpStddevPop extends FpAggKind
+
+ private case class FpCase(
+ values: IndexedSeq[Option[Double]],
+ agg: FpAggKind,
+ blockSize: Int,
+ fanout: Int) {
+ def n: Int = values.length
+ }
+
+ /**
+ * 20% null, 80% bounded gaussian-ish double. Range chosen so partial
+ * sums stay well within Double precision over n <= 5000.
+ */
+ private val genFpValue: Gen[Option[Double]] = {
+ val bounded: Gen[Double] = Gen.choose(-1e6, 1e6)
+ Gen.frequency(
+ (1, Gen.const(None: Option[Double])),
+ (4, bounded.map(Some(_)))
+ )
+ }
+
+ private val genFpAgg: Gen[FpAggKind] =
+ Gen.oneOf(FpAvg, FpStddevSamp, FpStddevPop)
+
+ private val genFpCase: Gen[FpCase] = for {
+ agg <- genFpAgg
+ n <- genN
+ values <- Gen.listOfN(n, genFpValue).map(_.toIndexedSeq)
+ blockSize <- genBlockSize
+ fanout <- genFanout
+ } yield FpCase(values, agg, blockSize, fanout)
+
+ /**
+ * Naive oracle for FP aggregates over [lo, hi). Returns None when Spark
+ * would emit NULL: empty frame, all-null frame, and (for StddevSamp)
+ * count < 2.
+ */
+ private def naiveFpAgg(
+ values: IndexedSeq[Option[Double]],
+ lo: Int, hi: Int,
+ kind: FpAggKind): Option[Double] = {
+ val nonNull = (lo until hi).flatMap(values).toArray
+ val n = nonNull.length
+ if (n == 0) return None
+ val mean = nonNull.sum / n
+ kind match {
+ case FpAvg => Some(mean)
+ case FpStddevSamp =>
+ if (n < 2) None
+ else {
+ var sq = 0.0
+ var i = 0
+ while (i < n) { val d = nonNull(i) - mean; sq += d * d; i += 1 }
+ Some(math.sqrt(sq / (n - 1)))
+ }
+ case FpStddevPop =>
+ var sq = 0.0
+ var i = 0
+ while (i < n) { val d = nonNull(i) - mean; sq += d * d; i += 1 }
+ Some(math.sqrt(sq / n))
+ }
+ }
+
+ /** Build tree + matching AggregateProcessor so we can `queryInto` to evaluate. */
+ private def buildFpTreeFor(
+ c: FpCase): (WindowSegmentTree, AggregateProcessor) = {
+ val attr = AttributeReference("v", DoubleType, nullable = true)()
+ val schema: Seq[Attribute] = Seq(attr)
+ val agg: DeclarativeAggregate = c.agg match {
+ case FpAvg => Average(attr)
+ case FpStddevSamp => StddevSamp(attr)
+ case FpStddevPop => StddevPop(attr)
+ }
+ val tree = new WindowSegmentTree(
+ Array(agg), schema, newMutableProjection,
+ c.fanout, c.blockSize, maxCachedBlocks = None,
+ taskMemoryManager = TaskContext.get().taskMemoryManager())
+ val rows = c.values.iterator.map { opt =>
+ val r = new GenericInternalRow(1)
+ opt match {
+ case Some(v) => r.update(0, v)
+ case None => r.setNullAt(0)
+ }
+ r.asInstanceOf[InternalRow]
+ }
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, rows, schema)
+ val processor = AggregateProcessor(
+ Array[Expression](agg),
+ ordinal = 0,
+ inputAttributes = schema,
+ newMutableProjection = newMutableProjection,
+ filters = Array[Option[Expression]](None))
+ (tree, processor)
+ }
+
+ private def queryFpResult(
+ tree: WindowSegmentTree,
+ processor: AggregateProcessor,
+ lo: Int, hi: Int): Option[Double] = {
+ val out = new SpecificInternalRow(Seq[DataType](DoubleType))
+ tree.queryInto(lo, hi, processor, out)
+ if (out.isNullAt(0)) None else Some(out.getDouble(0))
+ }
+
+ private def fpClose(actual: Option[Double], expected: Option[Double]): Boolean =
+ (actual, expected) match {
+ case (None, None) => true
+ case (Some(a), Some(e)) =>
+ if (java.lang.Double.isNaN(a) && java.lang.Double.isNaN(e)) true
+ else {
+ val absDiff = math.abs(a - e)
+ val denom = math.max(math.abs(e), 1.0) // absolute floor for near-zero
+ absDiff / denom < 1e-9
+ }
+ case _ => false
+ }
+
+ /**
+ * P4 Equivalence (FP family): segtree avg/stddev within 1e-9 relative
+ * tolerance of naive oracle. Smaller minSuccessful (50) than P1 because
+ * Welford-merge is heavier than Min/Max/Sum/Count.
+ */
+ test("P4 fp equivalence: segtree avg/stddev within 1e-9 of naive oracle") {
+ withTaskContext {
+ forAll(genFpCase, minSuccessful(50)) { c =>
+ val (tree, processor) = buildFpTreeFor(c)
+ try {
+ val frames = boundaryFrames(c.n) ++ randomFrames(c.n, k = 6, seed = c.n.toLong)
+ frames.foreach { case (lo, hi) =>
+ val expected = naiveFpAgg(c.values, lo, hi, c.agg)
+ val actual = queryFpResult(tree, processor, lo, hi)
+ assert(fpClose(actual, expected),
+ s"P4 fp mismatch: n=${c.n} agg=${c.agg} blockSize=${c.blockSize} " +
+ s"fanout=${c.fanout} frame=[$lo,$hi) " +
+ s"expected=$expected actual=$actual")
+ }
+ } finally tree.close()
+ }
+ }
+ }
+}
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeSuite.scala
new file mode 100644
index 000000000000..ac2a04744fcf
--- /dev/null
+++ b/sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowSegmentTreeSuite.scala
@@ -0,0 +1,521 @@
+/*
+ * 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.window
+
+import scala.util.Random
+
+import org.apache.spark.{LocalSparkContext, SparkConf, SparkContext, SparkEnv, SparkException, SparkFunSuite, TaskContext}
+import org.apache.spark.memory.MemoryTestingUtils
+import org.apache.spark.sql.catalyst.InternalRow
+import org.apache.spark.sql.catalyst.expressions.{Attribute, AttributeReference, Expression, GenericInternalRow, MutableProjection, SpecificInternalRow}
+import org.apache.spark.sql.catalyst.expressions.aggregate.{DeclarativeAggregate, Max, Min, StddevSamp, Sum}
+import org.apache.spark.sql.catalyst.expressions.codegen.GenerateMutableProjection
+import org.apache.spark.sql.types.{DataType, DoubleType, IntegerType, LongType}
+
+class WindowSegmentTreeSuite extends SparkFunSuite with LocalSparkContext {
+
+ // test harness
+ private val inputAttr: AttributeReference =
+ AttributeReference("v", IntegerType, nullable = true)()
+ private val inputSchema: Seq[Attribute] = Seq(inputAttr)
+
+ private def newMutableProjection
+ : (Seq[Expression], Seq[Attribute]) => MutableProjection =
+ (exprs, attrs) => GenerateMutableProjection.generate(exprs, attrs)
+
+ private def minAgg: DeclarativeAggregate = Min(inputAttr)
+
+ private def withTaskContext[T](body: => T): T = {
+ sc = new SparkContext("local", "test", new SparkConf(false))
+ val taskContext = MemoryTestingUtils.fakeTaskContext(SparkEnv.get)
+ TaskContext.setTaskContext(taskContext)
+ try body finally {
+ TaskContext.unset()
+ }
+ }
+
+ private def buildTree(
+ values: Seq[Int],
+ aggs: Array[DeclarativeAggregate] = Array(minAgg),
+ fanout: Int = WindowSegmentTree.DefaultFanout,
+ blockSize: Int = WindowSegmentTree.DefaultBlockSize,
+ maxCachedBlocks: Option[Int] = None): WindowSegmentTree = {
+ val tree = new WindowSegmentTree(
+ aggs, inputSchema, newMutableProjection, fanout, blockSize, maxCachedBlocks,
+ taskMemoryManager = TaskContext.get().taskMemoryManager())
+ val rows = values.iterator.map { v =>
+ val r = new GenericInternalRow(1)
+ r.update(0, v)
+ r.asInstanceOf[InternalRow]
+ }
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, rows, inputSchema)
+ tree
+ }
+
+ /** Naive oracle: MIN over values[lo, hi). Returns Int box or null. */
+ private def naiveMin(values: Seq[Int], lo: Int, hi: Int): Any = {
+ if (lo >= hi) null
+ else {
+ var m = values(lo)
+ var i = lo + 1
+ while (i < hi) { if (values(i) < m) m = values(i); i += 1 }
+ m.asInstanceOf[java.lang.Integer]
+ }
+ }
+
+ private def newOutBuffer(): SpecificInternalRow =
+ new SpecificInternalRow(Seq[org.apache.spark.sql.types.DataType](IntegerType))
+
+ private def queryMin(tree: WindowSegmentTree, lo: Int, hi: Int): Any = {
+ val out = newOutBuffer()
+ tree.query(lo, hi, out)
+ if (out.isNullAt(0)) null else out.getInt(0)
+ }
+
+ test("build and single-point query returns identity; full scan matches naive") {
+ withTaskContext {
+ val values = Seq(5, 2, 9, 1, 7, 3, 4, 8, 6, 0)
+ val tree = buildTree(values, fanout = 4, blockSize = 1024)
+ try {
+ // single-point queries -> identity (null for MIN)
+ for (i <- 0 to values.length) {
+ assert(queryMin(tree, i, i) == null, s"identity at i=$i")
+ }
+ // full scan
+ assert(queryMin(tree, 0, values.length) == naiveMin(values, 0, values.length))
+ } finally tree.close()
+ }
+ }
+
+ test("single-block: range query matches naive baseline for random ranges") {
+ withTaskContext {
+ val rnd = new Random(0xC0FFEE)
+ val values = Seq.fill(100)(rnd.nextInt(1000))
+ val tree = buildTree(values, fanout = 8, blockSize = 1024)
+ try {
+ for (_ <- 0 until 20) {
+ val a = rnd.nextInt(values.length + 1)
+ val b = rnd.nextInt(values.length + 1)
+ val (lo, hi) = (math.min(a, b), math.max(a, b))
+ assert(queryMin(tree, lo, hi) == naiveMin(values, lo, hi),
+ s"mismatch at [$lo, $hi)")
+ }
+ } finally tree.close()
+ }
+ }
+
+ test("fanout boundaries: sizes {1, F, F+1, F*F} for fanout in {2,4,8,16}") {
+ withTaskContext {
+ val rnd = new Random(42)
+ for (fanout <- Seq(2, 4, 8, 16)) {
+ val sizes = Seq(1, fanout, fanout + 1, fanout * fanout).distinct
+ for (n <- sizes) {
+ val values = Seq.fill(n)(rnd.nextInt(10000))
+ val tree = buildTree(values, fanout = fanout, blockSize = 1024 * 1024)
+ try {
+ // Exhaustive ranges for small n.
+ for (lo <- 0 to n; hi <- lo to n) {
+ assert(queryMin(tree, lo, hi) == naiveMin(values, lo, hi),
+ s"fanout=$fanout n=$n [$lo,$hi)")
+ }
+ // Level-size invariant: ceil(prev / fanout)
+ if (n > 0) {
+ val numLevels = tree.peekLevelCount(0)
+ var prev = n
+ for (l <- 1 until numLevels) {
+ val expected = math.ceil(prev.toDouble / fanout).toInt
+ val actual = tree.peekLevelSize(0, l)
+ assert(actual == expected,
+ s"level-size wrong at fanout=$fanout n=$n level=$l: $actual vs $expected")
+ prev = expected
+ }
+ }
+ } finally tree.close()
+ }
+ }
+ }
+ }
+
+ test("identity at empty range query(k, k)") {
+ withTaskContext {
+ val values = (1 to 50).reverse
+ val tree = buildTree(values, fanout = 4, blockSize = 16)
+ try {
+ for (k <- Seq(0, values.length / 2, values.length)) {
+ assert(queryMin(tree, k, k) == null, s"identity at k=$k")
+ }
+ } finally tree.close()
+ }
+ }
+
+ test("block boundary correctness: cross-block vs single-block baseline") {
+ withTaskContext {
+ val rnd = new Random(123)
+ val values = Seq.fill(100)(rnd.nextInt(10000))
+ val treeBlocked = buildTree(values, fanout = 4, blockSize = 16)
+ val treeBaseline = buildTree(values, fanout = 4, blockSize = 1024)
+ try {
+ assert(treeBlocked.peekBlockCount == (100 + 16 - 1) / 16)
+ for (lo <- 0 to 100; hi <- lo to 100) {
+ assert(queryMin(treeBlocked, lo, hi) == queryMin(treeBaseline, lo, hi),
+ s"cross-block mismatch at [$lo, $hi)")
+ }
+ } finally {
+ treeBlocked.close()
+ treeBaseline.close()
+ }
+ }
+ }
+
+ test("LRU stability: same queries in different orders produce same results") {
+ withTaskContext {
+ val rnd = new Random(777)
+ val values = Seq.fill(100)(rnd.nextInt(10000))
+ val tree = buildTree(
+ values, fanout = 4, blockSize = 16, maxCachedBlocks = Some(2))
+ try {
+ val queries = Seq.fill(30) {
+ val a = rnd.nextInt(101)
+ val b = rnd.nextInt(101)
+ (math.min(a, b), math.max(a, b))
+ }
+ val results1 = queries.map { case (lo, hi) => queryMin(tree, lo, hi) }
+ val reordered = rnd.shuffle(queries.zipWithIndex)
+ val results2 = Array.fill[Any](queries.length)(null)
+ reordered.foreach { case ((lo, hi), idx) =>
+ results2(idx) = queryMin(tree, lo, hi)
+ }
+ for (i <- queries.indices) {
+ assert(results1(i) == results2(i),
+ s"LRU instability at query $i ${queries(i)}: ${results1(i)} vs ${results2(i)}")
+ }
+ // Also cross-check against naive oracle.
+ for (i <- queries.indices) {
+ val (lo, hi) = queries(i)
+ assert(results1(i) == naiveMin(values, lo, hi))
+ }
+ } finally tree.close()
+ }
+ }
+
+ test("cross-block: range query matches naive baseline for random ranges") {
+ withTaskContext {
+ val rnd = new Random(0xBEEF)
+ val values = Seq.fill(100)(rnd.nextInt(1000))
+ val tree = buildTree(values, fanout = 4, blockSize = 8)
+ try {
+ assert(tree.peekBlockCount == (100 + 8 - 1) / 8)
+ for (_ <- 0 until 50) {
+ val a = rnd.nextInt(values.length + 1)
+ val b = rnd.nextInt(values.length + 1)
+ val (lo, hi) = (math.min(a, b), math.max(a, b))
+ assert(queryMin(tree, lo, hi) == naiveMin(values, lo, hi),
+ s"mismatch at [$lo, $hi)")
+ }
+ } finally tree.close()
+ }
+ }
+
+ test("cross-block: multi-block level-size invariant") {
+ withTaskContext {
+ val rnd = new Random(31337)
+ val fanout = 4
+ val blockSize = 8
+ val numRows = 50 // > blockSize -> multiple blocks
+ val values = Seq.fill(numRows)(rnd.nextInt(10000))
+ val tree = buildTree(values, fanout = fanout, blockSize = blockSize)
+ try {
+ val numBlocks = tree.peekBlockCount
+ assert(numBlocks > 1, s"expected >1 block, got $numBlocks")
+ for (b <- 0 until numBlocks) {
+ val blockStart = b * blockSize
+ val blockRows = math.min(blockSize, numRows - blockStart)
+ val numLevels = tree.peekLevelCount(b)
+ var prev = blockRows
+ for (l <- 1 until numLevels) {
+ val expected = math.ceil(prev.toDouble / fanout).toInt
+ val actual = tree.peekLevelSize(b, l)
+ assert(actual == expected,
+ s"block=$b level=$l: $actual vs $expected")
+ prev = expected
+ }
+ }
+ // Correctness cross-check.
+ for (lo <- 0 to numRows; hi <- lo to numRows) {
+ assert(queryMin(tree, lo, hi) == naiveMin(values, lo, hi),
+ s"cross-block MIN mismatch at [$lo, $hi)")
+ }
+ } finally tree.close()
+ }
+ }
+
+ // D8 multi-aggregate
+ test("D8 multi-aggregate: MIN + MAX + SUM on the same tree") {
+ withTaskContext {
+ val rnd = new Random(2024)
+ val numRows = 50
+ val values = Seq.fill(numRows)(rnd.nextInt(1000))
+ val aggs: Array[DeclarativeAggregate] =
+ Array(Min(inputAttr), Max(inputAttr), Sum(inputAttr))
+ val tree = buildTree(values, aggs = aggs, fanout = 4, blockSize = 8)
+ // Output schema: MIN(int), MAX(int), SUM(long).
+ // Sum on IntegerType widens the buffer slot to LongType.
+ val outTypes: Seq[DataType] = Seq(IntegerType, IntegerType, LongType)
+ def queryAll(lo: Int, hi: Int): (Any, Any, Any) = {
+ val out = new SpecificInternalRow(outTypes)
+ tree.query(lo, hi, out)
+ val mn = if (out.isNullAt(0)) null else out.getInt(0)
+ val mx = if (out.isNullAt(1)) null else out.getInt(1)
+ val sm = if (out.isNullAt(2)) null else out.getLong(2)
+ (mn, mx, sm)
+ }
+ def naiveMax(vs: Seq[Int], lo: Int, hi: Int): Any =
+ if (lo >= hi) null else vs.slice(lo, hi).max
+ def naiveSum(vs: Seq[Int], lo: Int, hi: Int): Any =
+ if (lo >= hi) null else vs.slice(lo, hi).map(_.toLong).sum
+ try {
+ for (_ <- 0 until 20) {
+ val a = rnd.nextInt(numRows + 1)
+ val b = rnd.nextInt(numRows + 1)
+ val (lo, hi) = (math.min(a, b), math.max(a, b))
+ val (mn, mx, sm) = queryAll(lo, hi)
+ assert(mn == naiveMin(values, lo, hi), s"MIN at [$lo,$hi)")
+ assert(mx == naiveMax(values, lo, hi), s"MAX at [$lo,$hi)")
+ assert(sm == naiveSum(values, lo, hi), s"SUM at [$lo,$hi)")
+ }
+ } finally tree.close()
+ }
+ }
+
+ // D9 spill coverage
+ test("D9 spill path: low thresholds still produce correct range aggregates") {
+ withTaskContext {
+ val rnd = new Random(909)
+ val numRows = 60
+ val values = Seq.fill(numRows)(rnd.nextInt(1000))
+ val tree = new WindowSegmentTree(
+ Array(minAgg), inputSchema, newMutableProjection,
+ fanout = 4, blockSize = 8, maxCachedBlocks = Some(2),
+ taskMemoryManager = TaskContext.get().taskMemoryManager())
+ val rows = values.iterator.map { v =>
+ val r = new GenericInternalRow(1); r.update(0, v); r.asInstanceOf[InternalRow]
+ }
+ SegmentTreeWindowTestHelper.buildTreeFromIter(
+ tree, rows, inputSchema, inMemoryThreshold = 4, spillThreshold = 8)
+ try {
+ for (_ <- 0 until 40) {
+ val a = rnd.nextInt(numRows + 1)
+ val b = rnd.nextInt(numRows + 1)
+ val (lo, hi) = (math.min(a, b), math.max(a, b))
+ assert(queryMin(tree, lo, hi) == naiveMin(values, lo, hi),
+ s"spill-path mismatch at [$lo, $hi)")
+ }
+ } finally tree.close()
+ }
+ }
+
+ test("D10 rebuild: second build replaces state; failed build preserves prior state") {
+ withTaskContext {
+ val v1 = Seq(5, 1, 9, 3, 7, 2, 8, 4, 6, 0)
+ val v2 = Seq(100, 200, 50, 400, 25, 600, 12, 800)
+ val v3 = Seq(10, 20, 30, 40, 50, 60, 70, 80, 90, 5)
+ val tree = new WindowSegmentTree(
+ Array(minAgg), inputSchema, newMutableProjection,
+ fanout = 4, blockSize = 4,
+ taskMemoryManager = TaskContext.get().taskMemoryManager())
+
+ def iterOf(vs: Seq[Int]): Iterator[InternalRow] = vs.iterator.map { v =>
+ val r = new GenericInternalRow(1); r.update(0, v); r.asInstanceOf[InternalRow]
+ }
+
+ try {
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, iterOf(v1), inputSchema)
+ assert(queryMin(tree, 0, v1.length) == v1.min)
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, iterOf(v2), inputSchema)
+ assert(tree.size == v2.length)
+ for (lo <- 0 to v2.length; hi <- lo to v2.length) {
+ assert(queryMin(tree, lo, hi) == naiveMin(v2, lo, hi),
+ s"post-rebuild mismatch at [$lo, $hi)")
+ }
+
+ // Now simulate a failing build midway through; prior state must remain queryable.
+ val boomIter: Iterator[InternalRow] = new Iterator[InternalRow] {
+ private var emitted = 0
+ override def hasNext: Boolean = true
+ override def next(): InternalRow = {
+ if (emitted >= 3) throw new RuntimeException("boom")
+ emitted += 1
+ val r = new GenericInternalRow(1); r.update(0, -1); r.asInstanceOf[InternalRow]
+ }
+ }
+ intercept[RuntimeException](
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, boomIter, inputSchema))
+ // Prior v2 state intact.
+ assert(tree.size == v2.length)
+ assert(queryMin(tree, 0, v2.length) == v2.min)
+
+ // Build v3 successfully after the failure.
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, iterOf(v3), inputSchema)
+ assert(tree.size == v3.length)
+ for (lo <- 0 to v3.length; hi <- lo to v3.length) {
+ assert(queryMin(tree, lo, hi) == naiveMin(v3, lo, hi),
+ s"post-recovery mismatch at [$lo, $hi)")
+ }
+ } finally tree.close()
+ }
+ }
+
+ test("D11 error paths: invalid ctor args and invalid query ranges") {
+ withTaskContext {
+ // Constructor validation.
+ intercept[IllegalArgumentException] {
+ new WindowSegmentTree(Array(minAgg), inputSchema, newMutableProjection, fanout = 1)
+ }
+ intercept[IllegalArgumentException] {
+ new WindowSegmentTree(Array(minAgg), inputSchema, newMutableProjection, blockSize = 0)
+ }
+ intercept[IllegalArgumentException] {
+ new WindowSegmentTree(
+ Array(minAgg), inputSchema, newMutableProjection, maxCachedBlocks = Some(0))
+ }
+ intercept[IllegalArgumentException] {
+ new WindowSegmentTree(
+ Array(minAgg), inputSchema, newMutableProjection, maxCachedBlocks = Some(-1))
+ }
+
+ // Query range validation.
+ val values = Seq(3, 1, 4, 1, 5, 9, 2, 6, 5, 3)
+ val tree = buildTree(values, fanout = 4, blockSize = 4)
+ try {
+ // Pre-fill outBuffer with a sentinel, then assert bounds check leaves it alone.
+ def sentinelBuf(): SpecificInternalRow = {
+ val b = newOutBuffer()
+ b.setInt(0, 0x5EEDED)
+ b
+ }
+ val sz = tree.size
+ for ((lo, hi) <- Seq((-1, 5), (0, sz + 1), (5, 3))) {
+ val out = sentinelBuf()
+ val ex = intercept[SparkException](tree.query(lo, hi, out))
+ assert(ex.getCondition == "INTERNAL_ERROR")
+ assert(!out.isNullAt(0) && out.getInt(0) == 0x5EEDED,
+ s"outBuffer mutated by invalid query [$lo,$hi)")
+ }
+ } finally tree.close()
+ }
+ }
+
+ // STDDEV_SAMP regression: minimal repro for the digest mismatch in GHA run
+ // 24599916378 (WindowBenchmark Section A). MIN/MAX/SUM/COUNT/AVG pass
+ // digest parity; STDDEV_SAMP is the first multi-buffer agg combining
+ // CentralMomentAgg's Welford merge (n, avg, m2) with a non-trivial
+ // `evaluateExpression`. Exercises both merge (`query`) and evaluate
+ // (`queryInto` + `AggregateProcessor`) paths.
+ test("STDDEV_SAMP: segtree matches naive oracle on random doubles, W=21") {
+ withTaskContext {
+ val rnd = new Random(0x57DDEFL) // fixed seed
+ val n = 100
+ val values: Array[Double] = Array.fill(n)(rnd.nextGaussian() * 1000.0 + 50.0)
+ val doubleAttr = AttributeReference("v", DoubleType, nullable = true)()
+ val schema: Seq[Attribute] = Seq(doubleAttr)
+ val agg: DeclarativeAggregate = StddevSamp(doubleAttr)
+
+ // Use a block size that forces cross-block merges at W=21: blockSize=8.
+ val tree = new WindowSegmentTree(
+ Array(agg), schema, newMutableProjection,
+ fanout = 4, blockSize = 8, maxCachedBlocks = None,
+ taskMemoryManager = TaskContext.get().taskMemoryManager())
+ try {
+ val rows = values.iterator.map { v =>
+ val r = new GenericInternalRow(1)
+ r.update(0, v)
+ r.asInstanceOf[InternalRow]
+ }
+ SegmentTreeWindowTestHelper.buildTreeFromIter(tree, rows, schema)
+
+ // AggregateProcessor identical to the Frame's, so `queryInto`
+ // exercises evaluateExpression (sqrt(m2/(n-1)) with n=1 / div-by-0
+ // guards in StddevSamp).
+ val processor = AggregateProcessor(
+ Array[Expression](agg),
+ ordinal = 0,
+ inputAttributes = schema,
+ newMutableProjection = newMutableProjection,
+ filters = Array[Option[Expression]](None))
+
+ val out = new SpecificInternalRow(Seq[DataType](DoubleType))
+
+ val halfW = 10 // W=21
+ var i = 0
+ var maxRelErr = 0.0
+ while (i < n) {
+ val lo = math.max(0, i - halfW)
+ val hi = math.min(n, i + halfW + 1)
+ tree.queryInto(lo, hi, processor, out)
+ val actual = if (out.isNullAt(0)) Double.NaN else out.getDouble(0)
+ val expected = naiveStddevSamp(values, lo, hi)
+ // NaN handling: legacy `nullOnDivideByZero=true` -> null when n<=1.
+ if (java.lang.Double.isNaN(expected)) {
+ assert(out.isNullAt(0),
+ s"expected null at i=$i [$lo,$hi), got $actual")
+ } else {
+ assert(!out.isNullAt(0),
+ s"expected $expected at i=$i [$lo,$hi), got null")
+ val denom = math.max(math.abs(expected), 1e-12)
+ val rel = math.abs(actual - expected) / denom
+ if (rel > maxRelErr) maxRelErr = rel
+ assert(rel < 1e-9,
+ s"STDDEV_SAMP mismatch at i=$i [$lo,$hi): " +
+ s"expected=$expected actual=$actual relErr=$rel")
+ }
+ i += 1
+ }
+ } finally tree.close()
+ }
+ }
+
+ /** Naive STDDEV_SAMP over [lo, hi); returns NaN for n < 2 to signal null. */
+ private def naiveStddevSamp(values: Array[Double], lo: Int, hi: Int): Double = {
+ val n = hi - lo
+ if (n < 2) return Double.NaN
+ var sum = 0.0
+ var i = lo
+ while (i < hi) { sum += values(i); i += 1 }
+ val mean = sum / n
+ var sq = 0.0
+ i = lo
+ while (i < hi) { val d = values(i) - mean; sq += d * d; i += 1 }
+ math.sqrt(sq / (n - 1))
+ }
+
+ test("D12 block-aligned cross-block boundaries") {
+ withTaskContext {
+ val rnd = new Random(12)
+ val numRows = 50
+ val blockSize = 10
+ val values = Seq.fill(numRows)(rnd.nextInt(10000))
+ val tree = buildTree(values, fanout = 4, blockSize = blockSize)
+ try {
+ for ((lo, hi) <- Seq((0, 20), (10, 40), (20, 50), (0, 50))) {
+ assert(queryMin(tree, lo, hi) == naiveMin(values, lo, hi),
+ s"aligned mismatch at [$lo,$hi)")
+ }
+ } finally tree.close()
+ }
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org