markj-db opened a new pull request, #59248: URL: https://github.com/apache/spark/pull/59248
### What changes were proposed in this pull request? Follow-up to [SPARK-59929](https://issues.apache.org/jira/browse/SPARK-59929) (#59190), which made window execution fail fast with `WINDOW_FUNCTION_PARTITION_SIZE_EXCEEDS_LIMIT` for any partition larger than `Int.MaxValue` rows. That was necessary because the spill-backed sorter underneath the window buffer counted rows in `int`. This PR fixes that accounting so LEAD/LAG and unbounded-offset NTH_VALUE return correct results over such partitions instead of failing. **Core: `UnsafeExternalSorter` row accounting** - `getIterator(startIndex)` summed the record counts of its spill files into an `int i`. Once the spill files together held more than `Int.MaxValue` records, `i` wrapped negative: the check `i + recordsSpilled() > startIndex` went false, so later spill files were skipped (rows silently dropped), and the in-memory tail was asked to skip `startIndex - i` (about 2^31) records, throwing `ArrayIndexOutOfBoundsException` from `moveOver`. The running count is now a `long`, and `moveOver` takes a `long` step count. - `UnsafeSorterIterator.getNumRecords()` now returns `long`, and the aggregates that sum it across spill files (`ChainedIterator`, `UnsafeSorterSpillMerger`) are `long`. Per-spill-file and per-in-memory-sorter counts are inherently `int`-bounded (a spill file's count is written with `writeInt` and comes from one in-memory sorter), so those stay `int` internally; `SpillableIterator` narrows with `Math.toIntExact`. All `UnsafeSorterIterator` subclasses are in this package, and `org.apache.spark.util.collection.unsafe.*` is already excluded from MiMa. - Adds a package-private `@VisibleForTesting addSpillWriterForTesting` hook, alongside the class's existing test hooks, so a test can simulate spill files with huge record counts. **SQL: window partition-size guard** - Reintroduces `WindowFunctionFrame.supportsLargePartition` (defaults to `false`). `WindowEvaluatorFactoryBase.checkPartitionSizeLimit` now rejects a partition larger than `Int.MaxValue` rows only if some active frame does not support it. - `FrameLessOffsetWindowFunctionFrame` (LEAD/LAG) and `UnboundedOffsetWindowFunctionFrame` (NTH_VALUE over the whole partition) return `true`: their `write` ignores the driver's `Int` row index, they do not use the `Int` window bounds, and `OffsetWindowFunctionFrameBase.inputIndex` is widened to `Long`. - The error message now says the limit comes from at least one of the query's window functions. Frames that still track the row index or window bounds as `Int` (sliding, running aggregates, ranking, segment-tree, and NTH_VALUE over UNBOUNDED PRECEDING AND CURRENT ROW) keep failing fast with the clear error. Supporting them would require widening the `WindowFunctionFrame.write` index, the `currentLowerBound`/`currentUpperBound` API and `getIterator(int startIndex)`, which is out of scope here. ### Why are the changes needed? A window query whose single partition spills more than `Int.MaxValue` rows hit this overflow. Before SPARK-59929 it showed up as either silently wrong results (e.g. `lag` returning `NULL` for every row) or `ArrayIndexOutOfBoundsException: Failed to move the iterator ... steps forward`. SPARK-59929 turned both into a clear error; this PR makes LEAD/LAG and unbounded NTH_VALUE actually work for those partitions. ### Does this PR introduce _any_ user-facing change? Yes. LEAD/LAG and NTH_VALUE over the whole partition now return correct results for window partitions larger than `Int.MaxValue` rows, where master currently fails with `WINDOW_FUNCTION_PARTITION_SIZE_EXCEEDS_LIMIT`. Other window functions over such partitions still fail with that error, whose message is reworded slightly. Partitions up to `Int.MaxValue` rows are unaffected. ### How was this patch tested? - New `UnsafeExternalSorterSuite.testGetIteratorWithMoreThanIntMaxSpilledRecords` (also runs in `UnsafeExternalSorterRadixSortSuite`): registers mocked spill writers reporting `Int.MaxValue` and 10 records plus 5 real in-memory records, then checks that `getIterator(0)` reads every spill file, reports `Int.MaxValue + 15` records, and returns the in-memory tail from its first record. With the previous `int` accumulator this test fails with `ArrayIndexOutOfBoundsException: Failed to move the iterator 2147483639 steps forward`. - `WindowFunctionFrameSuite` rewritten against real frames: LEAD/LAG and unbounded NTH_VALUE declare support, unbounded-preceding NTH_VALUE and an aggregate frame (which inherits the default) do not, and the guard allows or rejects partitions above `Int.MaxValue` rows accordingly. - Existing `UnsafeExternalSorterSuite`, `UnsafeExternalSorterRadixSortSuite` and other `org.apache.spark.util.collection.unsafe.sort` suites, `DataFrameWindowFunctionsSuite`, `WindowSegmentTreeSuite`, `ExternalAppendOnlyUnsafeRowArraySuite` and `SparkThrowableSuite` pass. A real end-to-end run with more than 2^31 rows in one partition is not practical as a unit test (about 30 minutes for 2.2B rows on one core), so the overflow is covered by the targeted tests above. ### Was this patch authored or co-authored using generative AI tooling? Generated-by: Isaac This pull request and its description were written by Isaac. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
