markj-db commented on code in PR #59248:
URL: https://github.com/apache/spark/pull/59248#discussion_r4225478450
##########
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:
Fixed in 9e2f421. Added `LAG keeps returning real rows when its cursor
crosses Int.MaxValue`, along the lines you suggested: a real
`FrameLessOffsetWindowFunctionFrame` is prepared over a small
`ExternalAppendOnlyUnsafeRowArray` subclass that reports `length = Int.MaxValue
+ 10`. The cursor is started at `Int.MaxValue - 1` through a small test
subclass, and four writes cross the boundary. Each write must return the lagged
input row, not the default.
I checked that it catches the regression: with `inputIndex` as an `Int`
again, the test fails at the third write (`result.isNullAt(0) was true write 2
returned the default value`). With the `Long` it passes, as do the other tests
in the suite and `DataFrameWindowFunctionsSuite`.
I also updated the PR description. The testing section now describes this
test, and it no longer says one end-to-end run is covered by the targeted
tests. Since `inputIndex` lives in the shared base class, this protects the
field type for the unbounded-offset NTH_VALUE frame as well; that frame's own
`prepare` scan is not driven past `Int.MaxValue`, and the description says so.
--
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]
