Optimize code

apache · Oct 27, 2020 · 72e7805 · 72e7805
1 parent fd59f6e
commit 72e7805
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 6 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
@@ -200,7 +200,7 @@ trait WindowExecBase extends UnaryExecNode {
                 child.output,
                 (expressions, schema) =>
                   MutableProjection.create(expressions, schema),
-                offset - 1)
+                offset)
             }
           case ("UNBOUNDED_PRECEDING_OFFSET", _, IntegerLiteral(offset), _) =>
             target: InternalRow => {
@@ -212,7 +212,7 @@ trait WindowExecBase extends UnaryExecNode {
                 child.output,
                 (expressions, schema) =>
                   MutableProjection.create(expressions, schema),
-                offset - 1)
+                offset)
             }
 
           // 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
@@ -196,13 +196,13 @@ class UnboundedOffsetWindowFunctionFrame(
 
   override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
     input = rows
-    if (offset >= input.length) {
+    if (offset > input.length) {
       fillDefaultValue(EmptyRow)
     } else {
       inputIterator = input.generateIterator()
-      // drain the first few rows if offset is larger than zero
+      // drain the first few rows if offset is larger than one
       inputIndex = 0
-      while (inputIndex < offset) {
+      while (inputIndex < offset - 1) {
         if (inputIterator.hasNext) inputIterator.next()
         inputIndex += 1
       }
@@ -240,7 +240,14 @@ class UnboundedPrecedingOffsetWindowFunctionFrame(
   var selectedRow: UnsafeRow = null
 
   override def prepare(rows: ExternalAppendOnlyUnsafeRowArray): Unit = {
-    super.prepare(rows)
+    input = rows
+    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 >= 0 && inputIndex < input.length) {
       selectedRow = WindowFunctionFrame.getNextOrNull(inputIterator)
     }