HyukjinKwon commented on code in PR #59248:
URL: https://github.com/apache/spark/pull/59248#discussion_r4223934558
##########
sql/core/src/test/scala/org/apache/spark/sql/execution/window/WindowFunctionFrameSuite.scala:
##########
@@ -37,28 +42,83 @@ class WindowFunctionFrameSuite extends SparkFunSuite {
override def childOutput: Seq[Attribute] = Nil
override def spillSize: SQLMetric = null
- def check(numRows: Long): Unit = checkPartitionSizeLimit(numRows)
+ def check(numRows: Long, frames: WindowFunctionFrame*): Unit =
+ checkPartitionSizeLimit(numRows, frames.toArray)
}
private val maxRows = Int.MaxValue.toLong
+ private val attr = AttributeReference("v", IntegerType, nullable = true)()
+ private val newProjection =
+ (expressions: Seq[Expression], schema: Seq[Attribute]) =>
+ MutableProjection.create(expressions, schema)
+
+ private def target = new SpecificInternalRow(Seq(IntegerType))
+
+ /** LAG(v, 1): a frameless offset frame. */
+ private def lagFrame: WindowFunctionFrame =
+ new FrameLessOffsetWindowFunctionFrame(
+ target,
+ ordinal = 0,
+ Array[OffsetWindowFunction](Lag(attr, Literal(1), Literal(null,
IntegerType), false)),
+ Seq[Attribute](attr),
+ newProjection,
+ offset = -1)
+
+ /** NTH_VALUE(v, 1) over the whole partition. */
+ private def unboundedNthValueFrame: WindowFunctionFrame =
+ new UnboundedOffsetWindowFunctionFrame(
+ target,
+ ordinal = 0,
+ Array[OffsetWindowFunction](NthValue(attr, Literal(1), ignoreNulls =
false)),
+ Seq[Attribute](attr),
+ newProjection,
+ offset = 1)
+
+ /** NTH_VALUE(v, 1) over UNBOUNDED PRECEDING AND CURRENT ROW; its `write`
reads the row index. */
+ private def unboundedPrecedingNthValueFrame: WindowFunctionFrame =
+ new UnboundedPrecedingOffsetWindowFunctionFrame(
+ target,
+ ordinal = 0,
+ Array[OffsetWindowFunction](NthValue(attr, Literal(1), ignoreNulls =
false)),
+ Seq[Attribute](attr),
+ newProjection,
+ offset = 1)
+
+ /** An aggregate over the whole partition; inherits the default
`supportsLargePartition`. */
+ private def unboundedAggregateFrame: WindowFunctionFrame =
+ new UnboundedWindowFunctionFrame(target, processor = null)
+
+ test("frames declare whether they support partitions larger than
Int.MaxValue rows") {
+ assert(lagFrame.supportsLargePartition)
+ assert(unboundedNthValueFrame.supportsLargePartition)
+ assert(!unboundedPrecedingNthValueFrame.supportsLargePartition)
+ // Fails if the `WindowFunctionFrame` default ever becomes `true`.
+ assert(!unboundedAggregateFrame.supportsLargePartition)
+ }
+
+ test("partitions up to Int.MaxValue rows are allowed for every frame") {
+ val factory = new TestFactory
+ factory.check(0L, unboundedAggregateFrame)
+ factory.check(maxRows, unboundedAggregateFrame,
unboundedPrecedingNthValueFrame)
+ }
- test("partitions up to Int.MaxValue rows are allowed") {
+ test("larger partitions are allowed when every frame supports them") {
Review Comment:
**Nit (P3):** These tests cover the opt-in flags and the guard decision, but
nothing drives an opted-in frame across Int.MaxValue rows. The frame-side half
of the fix is `OffsetWindowFunctionFrameBase.inputIndex` being a `Long`. If it
went back to `protected var inputIndex = 0`, every use (`inputIndex <
input.length`, `inputIndex = offset`, `inputIndex += 1`, ...) would still
compile, both frames would still report `supportsLargePartition = true`, and
this suite and the sorter test would still pass. At runtime, LAG's cursor would
wrap negative after 2^31 rows, `inputIndex >= 0` would become false, and every
later row would silently get the default value: exactly the wrong result the
guard exists to prevent.
A focused test can close this without materializing 2^31 rows.
`FrameLessOffsetWindowFunctionFrame` is non-final with a `protected` cursor,
and `ExternalAppendOnlyUnsafeRowArray` is a plain class, so a test can prepare
the frame over a stubbed array whose `length` exceeds `Int.MaxValue`, move the
cursor just below `Int.MaxValue`, write a few rows, and check that the rows
after the boundary hold real input values rather than the default.
**Verification:**
- **Regression:** With the cursor seeded just below Int.MaxValue on a
partition reported as longer than Int.MaxValue rows, writes across the boundary
return real input rows; the same assertion fails when the cursor is an Int.
- **Behavior:** The existing frame-flag and partition-size-guard assertions
still pass unchanged.
--
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]
