Skip to content

Commit 3d6e31a

Browse files
beliefercloud-fan
authored andcommitted
[SPARK-45649][SQL] Unify the prepare framework for OffsetWindowFunctionFrame
### What changes were proposed in this pull request? Currently, the implementation of the `prepare` of all the `OffsetWindowFunctionFrame` have the same code logic show below. ``` override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = { if (offset > rows.length) { fillDefaultValue(EmptyRow) } else { resetStates(rows) if (ignoreNulls) { ... } else { ... } } } ``` This PR want unify the prepare framework for `OffsetWindowFunctionFrame` ### Why are the changes needed? Unify the prepare framework for `OffsetWindowFunctionFrame` ### Does this PR introduce _any_ user-facing change? 'No'. Inner update. ### How was this patch tested? Exists test cases. ### Was this patch authored or co-authored using generative AI tooling? 'No'. Closes apache#43507 from beliefer/SPARK-45430_followup. Authored-by: Jiaan Geng <beliefer@163.com> Signed-off-by: Wenchen Fan <wenchen@databricks.com>
1 parent f1d1dc1 commit 3d6e31a

1 file changed

Lines changed: 61 additions & 57 deletions

File tree

sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowFunctionFrame.scala

Lines changed: 61 additions & 57 deletions
Original file line numberDiff line numberDiff line change
@@ -86,7 +86,8 @@ abstract class OffsetWindowFunctionFrameBase(
8686
expressions: Array[OffsetWindowFunction],
8787
inputSchema: Seq[Attribute],
8888
newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection,
89-
offset: Int)
89+
offset: Int,
90+
ignoreNulls: Boolean)
9091
extends WindowFunctionFrame {
9192

9293
/** Rows of the partition currently being processed. */
@@ -140,6 +141,8 @@ abstract class OffsetWindowFunctionFrameBase(
140141
// is not null.
141142
protected var skippedNonNullCount = 0
142143

144+
protected val absOffset = Math.abs(offset)
145+
143146
// Reset the states by the data of the new partition.
144147
protected def resetStates(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
145148
input = rows
@@ -175,6 +178,31 @@ abstract class OffsetWindowFunctionFrameBase(
175178
}
176179
}
177180

181+
override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
182+
if (absOffset > rows.length) {
183+
fillDefaultValue(EmptyRow)
184+
} else {
185+
resetStates(rows)
186+
if (ignoreNulls) {
187+
prepareForIgnoreNulls()
188+
} else {
189+
prepareForRespectNulls()
190+
}
191+
}
192+
}
193+
194+
protected def prepareForIgnoreNulls(): Unit = findNextRowWithNonNullInput()
195+
196+
protected def prepareForRespectNulls(): Unit
197+
198+
override def write(index: Int, current: InternalRow): Unit = {
199+
if (input != null) {
200+
doWrite(index, current)
201+
}
202+
}
203+
204+
protected def doWrite(index: Int, current: InternalRow): Unit
205+
178206
override def currentLowerBound(): Int = throw new UnsupportedOperationException()
179207

180208
override def currentUpperBound(): Int = throw new UnsupportedOperationException()
@@ -196,24 +224,15 @@ class FrameLessOffsetWindowFunctionFrame(
196224
offset: Int,
197225
ignoreNulls: Boolean = false)
198226
extends OffsetWindowFunctionFrameBase(
199-
target, ordinal, expressions, inputSchema, newMutableProjection, offset) {
227+
target, ordinal, expressions, inputSchema, newMutableProjection, offset, ignoreNulls) {
200228

201-
override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
202-
resetStates(rows)
203-
if (ignoreNulls) {
204-
if (Math.abs(offset) > rows.length) {
205-
fillDefaultValue(EmptyRow)
206-
} else {
207-
findNextRowWithNonNullInput()
208-
}
209-
} else {
210-
// drain the first few rows if offset is larger than zero
211-
while (inputIndex < offset) {
212-
if (inputIterator.hasNext) inputIterator.next()
213-
inputIndex += 1
214-
}
215-
inputIndex = offset
229+
override def prepareForRespectNulls(): Unit = {
230+
// drain the first few rows if offset is larger than zero
231+
while (inputIndex < offset) {
232+
if (inputIterator.hasNext) inputIterator.next()
233+
inputIndex += 1
216234
}
235+
inputIndex = offset
217236
}
218237

219238
private val doWrite = if (ignoreNulls && offset > 0) {
@@ -260,7 +279,6 @@ class FrameLessOffsetWindowFunctionFrame(
260279
// 7. current row -> z, next selected row -> y, output: y;
261280
// 8. current row -> v, next selected row -> z, output: z;
262281
// 9. current row -> null, next selected row -> v, output: v;
263-
val absOffset = Math.abs(offset)
264282
(current: InternalRow) =>
265283
if (skippedNonNullCount == absOffset) {
266284
nextSelectedRow = EmptyRow
@@ -294,7 +312,7 @@ class FrameLessOffsetWindowFunctionFrame(
294312
inputIndex += 1
295313
}
296314

297-
override def write(index: Int, current: InternalRow): Unit = {
315+
protected def doWrite(index: Int, current: InternalRow): Unit = {
298316
doWrite(current)
299317
}
300318
}
@@ -317,35 +335,30 @@ class UnboundedOffsetWindowFunctionFrame(
317335
offset: Int,
318336
ignoreNulls: Boolean = false)
319337
extends OffsetWindowFunctionFrameBase(
320-
target, ordinal, expressions, inputSchema, newMutableProjection, offset) {
338+
target, ordinal, expressions, inputSchema, newMutableProjection, offset, ignoreNulls) {
321339
assert(offset > 0)
322340

323-
override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
324-
if (offset > rows.length) {
341+
override def prepareForIgnoreNulls(): Unit = {
342+
findNextRowWithNonNullInput()
343+
if (nextSelectedRow == EmptyRow) {
344+
// Use default values since the offset row whose input value is not null does not exist.
325345
fillDefaultValue(EmptyRow)
326346
} else {
327-
resetStates(rows)
328-
if (ignoreNulls) {
329-
findNextRowWithNonNullInput()
330-
if (nextSelectedRow == EmptyRow) {
331-
// Use default values since the offset row whose input value is not null does not exist.
332-
fillDefaultValue(EmptyRow)
333-
} else {
334-
projection(nextSelectedRow)
335-
}
336-
} else {
337-
var selectedRow: UnsafeRow = null
338-
// drain the first few rows if offset is larger than one
339-
while (inputIndex < offset) {
340-
selectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
341-
inputIndex += 1
342-
}
343-
projection(selectedRow)
344-
}
347+
projection(nextSelectedRow)
345348
}
346349
}
347350

348-
override def write(index: Int, current: InternalRow): Unit = {
351+
override def prepareForRespectNulls(): Unit = {
352+
var selectedRow: UnsafeRow = null
353+
// drain the first few rows if offset is larger than one
354+
while (inputIndex < offset) {
355+
selectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
356+
inputIndex += 1
357+
}
358+
projection(selectedRow)
359+
}
360+
361+
protected def doWrite(index: Int, current: InternalRow): Unit = {
349362
// The results are the same for each row in the partition, and have been evaluated in prepare.
350363
// Don't need to recalculate here.
351364
}
@@ -370,27 +383,18 @@ class UnboundedPrecedingOffsetWindowFunctionFrame(
370383
offset: Int,
371384
ignoreNulls: Boolean = false)
372385
extends OffsetWindowFunctionFrameBase(
373-
target, ordinal, expressions, inputSchema, newMutableProjection, offset) {
386+
target, ordinal, expressions, inputSchema, newMutableProjection, offset, ignoreNulls) {
374387
assert(offset > 0)
375388

376-
override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
377-
if (offset > rows.length) {
378-
fillDefaultValue(EmptyRow)
379-
} else {
380-
resetStates(rows)
381-
if (ignoreNulls) {
382-
findNextRowWithNonNullInput()
383-
} else {
384-
// drain the first few rows if offset is larger than one
385-
while (inputIndex < offset) {
386-
nextSelectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
387-
inputIndex += 1
388-
}
389-
}
389+
override def prepareForRespectNulls(): Unit = {
390+
// drain the first few rows if offset is larger than one
391+
while (inputIndex < offset) {
392+
nextSelectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
393+
inputIndex += 1
390394
}
391395
}
392396

393-
override def write(index: Int, current: InternalRow): Unit = {
397+
protected def doWrite(index: Int, current: InternalRow): Unit = {
394398
if (index >= inputIndex - 1 && nextSelectedRow != null) {
395399
projection(nextSelectedRow)
396400
} else {

0 commit comments

Comments
 (0)