You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@spark.apache.org by we...@apache.org on 2021/01/21 07:32:05 UTC
[spark] branch master updated: [SPARK-34096][SQL] Improve
performance for nth_value ignore nulls over offset window
This is an automated email from the ASF dual-hosted git repository.
wenchen pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/spark.git
The following commit(s) were added to refs/heads/master by this push:
new 140538e [SPARK-34096][SQL] Improve performance for nth_value ignore nulls over offset window
140538e is described below
commit 140538ea5bf2a87734491c02892a9826972afbd9
Author: beliefer <be...@163.com>
AuthorDate: Thu Jan 21 07:31:36 2021 +0000
[SPARK-34096][SQL] Improve performance for nth_value ignore nulls over offset window
### What changes were proposed in this pull request?
The current implement of `UnboundedOffsetWindowFunctionFrame` and `UnboundedPrecedingOffsetWindowFunctionFrame` only support `nth_value` that respect nulls. So nth_value will execute with `UnboundedWindowFunctionFrame` and `UnboundedPrecedingWindowFunctionFrame`.
`UnboundedWindowFunctionFrame` and `UnboundedPrecedingWindowFunctionFrame` will call `updateExpressions` of `nth_value` multiple times.
### Why are the changes needed?
Improve performance for nth_value ignore nulls over offset window
### Does this PR introduce _any_ user-facing change?
No
### How was this patch tested?
Jenkins test
Closes #31178 from beliefer/SPARK-34096.
Authored-by: beliefer <be...@163.com>
Signed-off-by: Wenchen Fan <we...@databricks.com>
---
.../sql/execution/window/WindowExecBase.scala | 23 +++---
.../sql/execution/window/WindowFunctionFrame.scala | 92 +++++++++++++---------
.../spark/sql/DataFrameWindowFunctionsSuite.scala | 35 ++++++++
3 files changed, 101 insertions(+), 49 deletions(-)
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExecBase.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExecBase.scala
index 5d999cb..ed1d24b 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExecBase.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowExecBase.scala
@@ -126,11 +126,12 @@ trait WindowExecBase extends UnaryExecNode {
// Add a function and its function to the map for a given frame.
def collect(tpe: String, fr: SpecifiedWindowFrame, e: Expression, fn: Expression): Unit = {
val key = fn match {
- // This branch is used for Lead/Lag to support ignoring null.
- // All window frames move in rows. If there are multiple Leads or Lags acting on a row
- // and operating on different input expressions, they should not be moved uniformly
+ // This branch is used for Lead/Lag to support ignoring null and optimize the performance
+ // for NthValue ignoring null.
+ // All window frames move in rows. If there are multiple Leads, Lags or NthValues acting on
+ // a row and operating on different input expressions, they should not be moved uniformly
// by row. Therefore, we put these functions in different window frames.
- case f: FrameLessOffsetWindowFunction if f.ignoreNulls =>
+ case f: OffsetWindowFunction if f.ignoreNulls =>
(tpe, fr.frameType, fr.lower, fr.upper, f.children.map(_.canonicalized))
case _ => (tpe, fr.frameType, fr.lower, fr.upper, Nil)
}
@@ -149,8 +150,8 @@ trait WindowExecBase extends UnaryExecNode {
case AggregateExpression(f, _, _, _, _) => collect("AGGREGATE", frame, e, f)
case f: FrameLessOffsetWindowFunction =>
collect("FRAME_LESS_OFFSET", f.fakeFrame, e, f)
- case f: OffsetWindowFunction if !f.ignoreNulls &&
- frame.frameType == RowFrame && frame.lower == UnboundedPreceding =>
+ case f: OffsetWindowFunction if frame.frameType == RowFrame &&
+ frame.lower == UnboundedPreceding =>
frame.upper match {
case UnboundedFollowing => collect("UNBOUNDED_OFFSET", f.fakeFrame, e, f)
case CurrentRow => collect("UNBOUNDED_PRECEDING_OFFSET", f.fakeFrame, e, f)
@@ -203,7 +204,7 @@ trait WindowExecBase extends UnaryExecNode {
MutableProjection.create(expressions, schema),
offset,
expr.nonEmpty)
- case ("UNBOUNDED_OFFSET", _, IntegerLiteral(offset), _, _) =>
+ case ("UNBOUNDED_OFFSET", _, IntegerLiteral(offset), _, expr) =>
target: InternalRow => {
new UnboundedOffsetWindowFunctionFrame(
target,
@@ -213,9 +214,10 @@ trait WindowExecBase extends UnaryExecNode {
child.output,
(expressions, schema) =>
MutableProjection.create(expressions, schema),
- offset)
+ offset,
+ expr.nonEmpty)
}
- case ("UNBOUNDED_PRECEDING_OFFSET", _, IntegerLiteral(offset), _, _) =>
+ case ("UNBOUNDED_PRECEDING_OFFSET", _, IntegerLiteral(offset), _, expr) =>
target: InternalRow => {
new UnboundedPrecedingOffsetWindowFunctionFrame(
target,
@@ -225,7 +227,8 @@ trait WindowExecBase extends UnaryExecNode {
child.output,
(expressions, schema) =>
MutableProjection.create(expressions, schema),
- offset)
+ offset,
+ expr.nonEmpty)
}
// Entire Partition Frame.
diff --git a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowFunctionFrame.scala b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowFunctionFrame.scala
index 0408deb..e4ad72c 100644
--- a/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowFunctionFrame.scala
+++ b/sql/core/src/main/scala/org/apache/spark/sql/execution/window/WindowFunctionFrame.scala
@@ -130,6 +130,31 @@ abstract class OffsetWindowFunctionFrameBase(
newMutableProjection(boundExpressions, Nil).target(target)
}
+ /** Holder the UnsafeRow where the input operator by function is not null. */
+ protected var nextSelectedRow = EmptyRow
+
+ // The number of rows skipped to get the next UnsafeRow where the input operator by function
+ // is not null.
+ protected var skippedNonNullCount = 0
+
+ /** Create the projection to determine whether input is null. */
+ protected val project = UnsafeProjection.create(Seq(IsNull(expressions.head.input)), inputSchema)
+
+ /** Check if the output value of the first index is null. */
+ protected def nullCheck(row: InternalRow): Boolean = project(row).getBoolean(0)
+
+ /** find the offset row whose input is not null */
+ protected def findNextRowWithNonNullInput(): Unit = {
+ while (skippedNonNullCount < offset && inputIndex < input.length) {
+ val r = WindowFunctionFrame.getNextOrNull(inputIterator)
+ if (!nullCheck(r)) {
+ nextSelectedRow = r
+ skippedNonNullCount += 1
+ }
+ inputIndex += 1
+ }
+ }
+
override def currentLowerBound(): Int = throw new UnsupportedOperationException()
override def currentUpperBound(): Int = throw new UnsupportedOperationException()
@@ -153,31 +178,6 @@ class FrameLessOffsetWindowFunctionFrame(
extends OffsetWindowFunctionFrameBase(
target, ordinal, expressions, inputSchema, newMutableProjection, offset) {
- /** Holder the UnsafeRow where the input operator by function is not null. */
- private var nextSelectedRow = EmptyRow
-
- // The number of rows skipped to get the next UnsafeRow where the input operator by function
- // is not null.
- private var skippedNonNullCount = 0
-
- /** Create the projection to determine whether input is null. */
- private val project = UnsafeProjection.create(Seq(IsNull(expressions.head.input)), inputSchema)
-
- /** Check if the output value of the first index is null. */
- private def nullCheck(row: InternalRow): Boolean = project(row).getBoolean(0)
-
- /** find the offset row whose input is not null */
- private def findNextRowWithNonNullInput(): Unit = {
- while (skippedNonNullCount < offset && inputIndex < input.length) {
- val r = WindowFunctionFrame.getNextOrNull(inputIterator)
- if (!nullCheck(r)) {
- nextSelectedRow = r
- skippedNonNullCount += 1
- }
- inputIndex += 1
- }
- }
-
override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
input = rows
inputIterator = input.generateIterator()
@@ -292,7 +292,8 @@ class UnboundedOffsetWindowFunctionFrame(
expressions: Array[OffsetWindowFunction],
inputSchema: Seq[Attribute],
newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection,
- offset: Int)
+ offset: Int,
+ ignoreNulls: Boolean = false)
extends OffsetWindowFunctionFrameBase(
target, ordinal, expressions, inputSchema, newMutableProjection, offset) {
assert(offset > 0)
@@ -305,12 +306,22 @@ class UnboundedOffsetWindowFunctionFrame(
inputIterator = input.generateIterator()
// drain the first few rows if offset is larger than one
inputIndex = 0
- while (inputIndex < offset - 1) {
- if (inputIterator.hasNext) inputIterator.next()
- inputIndex += 1
+ if (ignoreNulls) {
+ findNextRowWithNonNullInput()
+ if (nextSelectedRow == EmptyRow) {
+ // Use default values since the offset row whose input value is not null does not exist.
+ fillDefaultValue(EmptyRow)
+ } else {
+ projection(nextSelectedRow)
+ }
+ } else {
+ var selectedRow: UnsafeRow = null
+ while (inputIndex < offset) {
+ selectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
+ inputIndex += 1
+ }
+ projection(selectedRow)
}
- val r = WindowFunctionFrame.getNextOrNull(inputIterator)
- projection(r)
}
}
@@ -336,7 +347,8 @@ class UnboundedPrecedingOffsetWindowFunctionFrame(
expressions: Array[OffsetWindowFunction],
inputSchema: Seq[Attribute],
newMutableProjection: (Seq[Expression], Seq[Attribute]) => MutableProjection,
- offset: Int)
+ offset: Int,
+ ignoreNulls: Boolean = false)
extends OffsetWindowFunctionFrameBase(
target, ordinal, expressions, inputSchema, newMutableProjection, offset) {
assert(offset > 0)
@@ -348,17 +360,19 @@ class UnboundedPrecedingOffsetWindowFunctionFrame(
inputIterator = input.generateIterator()
// drain the first few rows if offset is larger than one
inputIndex = 0
- while (inputIndex < offset - 1) {
- if (inputIterator.hasNext) inputIterator.next()
- inputIndex += 1
- }
- if (inputIndex < input.length) {
- selectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
+ if (ignoreNulls) {
+ findNextRowWithNonNullInput()
+ selectedRow = nextSelectedRow.asInstanceOf[UnsafeRow]
+ } else {
+ while (inputIndex < offset) {
+ selectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
+ inputIndex += 1
+ }
}
}
override def write(index: Int, current: InternalRow): Unit = {
- if (index >= inputIndex && selectedRow != null) {
+ if (index >= inputIndex - 1 && selectedRow != null) {
projection(selectedRow)
} else {
fillDefaultValue(EmptyRow)
diff --git a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
index 3568ad3..e1d0247 100644
--- a/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
+++ b/sql/core/src/test/scala/org/apache/spark/sql/DataFrameWindowFunctionsSuite.scala
@@ -669,6 +669,41 @@ class DataFrameWindowFunctionsSuite extends QueryTest
Row("b", 2, null, null, null, null)))
}
+ test("nth_value with ignoreNulls over offset window frame") {
+ val nullStr: String = null
+ val df = Seq(
+ ("a", 0, nullStr),
+ ("a", 1, "x"),
+ ("a", 2, "y"),
+ ("a", 3, "z"),
+ ("a", 4, nullStr),
+ ("b", 1, nullStr),
+ ("b", 2, nullStr)).
+ toDF("key", "order", "value")
+ val window1 = Window.partitionBy($"key").orderBy($"order")
+ .rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)
+ val window2 = Window.partitionBy($"key").orderBy($"order")
+ .rowsBetween(Window.unboundedPreceding, Window.currentRow)
+ checkAnswer(
+ df.select(
+ $"key",
+ $"order",
+ nth_value($"value", 2).over(window1),
+ nth_value($"value", 2, ignoreNulls = false).over(window1),
+ nth_value($"value", 2, ignoreNulls = true).over(window1),
+ nth_value($"value", 2).over(window2),
+ nth_value($"value", 2, ignoreNulls = false).over(window2),
+ nth_value($"value", 2, ignoreNulls = true).over(window2)),
+ Seq(
+ Row("a", 0, "x", "x", "y", null, null, null),
+ Row("a", 1, "x", "x", "y", "x", "x", null),
+ Row("a", 2, "x", "x", "y", "x", "x", "y"),
+ Row("a", 3, "x", "x", "y", "x", "x", "y"),
+ Row("a", 4, "x", "x", "y", "x", "x", "y"),
+ Row("b", 1, null, null, null, null, null, null),
+ Row("b", 2, null, null, null, null, null, null)))
+ }
+
test("nth_value on descending ordered window") {
val nullStr: String = null
val df = Seq(
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@spark.apache.org
For additional commands, e-mail: commits-help@spark.apache.org