Refine the code and add tests

PHILO-HE · Mar 13, 2024 · c412300 · c412300
1 parent 1b998a2
commit c412300
Show file tree

Hide file tree

Showing 2 changed files with 67 additions and 20 deletions.
diff --git a/velox/exec/Window.cpp b/velox/exec/Window.cpp
@@ -320,24 +320,26 @@ void Window::updateKRowsFrameBounds(
       // Considers a very large int64 constantOffset is used.
       if (startValue < std::numeric_limits<int32_t>::min()) {
         std::fill_n(rawFrameBounds, numRows, startRow);
-      } else {
-        // Integer overflow cannot happen.
-        std::iota(rawFrameBounds, rawFrameBounds + numRows, startValue);
-      }
-    } else {
-      auto overflowStart = getOverflowStart(constantOffset);
-      if (overflowStart >= 0 && overflowStart < numRows) {
-        std::iota(rawFrameBounds, rawFrameBounds + overflowStart, startValue);
-        // For remaining, set with the largest index for this partition.
-        std::fill_n(
-            rawFrameBounds + overflowStart,
-            numRows - overflowStart,
-            startRow + numRows - 1);
-      } else {
-        // Integer overflow cannot happen.
-        std::iota(rawFrameBounds, rawFrameBounds + numRows, startValue);
+        return;
       }
+      // Integer overflow cannot happen.
+      std::iota(rawFrameBounds, rawFrameBounds + numRows, startValue);
+      return;
+    }
+    // KFollowing.
+    auto overflowStart = getOverflowStart(constantOffset);
+    if (overflowStart >= 0 && overflowStart < numRows) {
+      std::iota(rawFrameBounds, rawFrameBounds + overflowStart, startValue);
+      // For remaining, set with the largest index for this partition.
+      std::fill_n(
+          rawFrameBounds + overflowStart,
+          numRows - overflowStart,
+          startRow + numRows - 1);
+      return;
     }
+    // Integer overflow cannot happen.
+    std::iota(rawFrameBounds, rawFrameBounds + numRows, startValue);
+    return;
   } else {
     currentPartition_->extractColumn(
         frameArg.index, partitionOffset_, numRows, 0, frameArg.value);

diff --git a/velox/functions/prestosql/window/tests/AggregateWindowTest.cpp b/velox/functions/prestosql/window/tests/AggregateWindowTest.cpp
@@ -228,7 +228,7 @@ TEST_F(AggregateWindowTest, integerOverflowRowsFrame) {
   auto input = makeRowVector({c0, c1, c2, c3});
 
   std::string overClause = "partition by c0 order by c1 desc";
-  // Test with literal larger than INT32_MAX (2147483647).
+  // Test constant following larger than INT32_MAX (2147483647).
   std::string frameClause = "rows between 0 preceding and 2147483650 following";
   auto expected = makeRowVector(
       {c0,
@@ -239,7 +239,19 @@ TEST_F(AggregateWindowTest, integerOverflowRowsFrame) {
   WindowTestBase::testWindowFunction(
       {input}, "count(c1)", overClause, frameClause, expected);
 
-  // Test integer overflow with column-specified following (int32).
+  // Test overflow case that happens during the calculation for the middle
+  // partition.
+  frameClause = "rows between 0 preceding and 2147483645 following";
+  expected = makeRowVector(
+      {c0,
+       c1,
+       c2,
+       c3,
+       makeFlatVector<int64_t>({6, 5, 4, 3, 2, 1, 4, 3, 2, 1})});
+  WindowTestBase::testWindowFunction(
+      {input}, "count(c1)", overClause, frameClause, expected);
+
+  // Test with column-specified following (int32).
   frameClause = "rows between 0 preceding and c2 following";
   expected = makeRowVector(
       {c0,
@@ -250,7 +262,7 @@ TEST_F(AggregateWindowTest, integerOverflowRowsFrame) {
   WindowTestBase::testWindowFunction(
       {input}, "count(c1)", overClause, frameClause, expected);
 
-  // Test integer overflow with column-specified following (int64).
+  // Test with column-specified following (int64).
   frameClause = "rows between 0 preceding and c3 following";
   expected = makeRowVector(
       {c0,
@@ -261,7 +273,29 @@ TEST_F(AggregateWindowTest, integerOverflowRowsFrame) {
   WindowTestBase::testWindowFunction(
       {input}, "count(c1)", overClause, frameClause, expected);
 
-  // Test integer overflow with column-specified preceding (int64).
+  // Test constant preceding larger than INT32_MAX.
+  frameClause = "rows between 2147483650 preceding and 0 following";
+  expected = makeRowVector(
+      {c0,
+       c1,
+       c2,
+       c3,
+       makeFlatVector<int64_t>({1, 2, 3, 4, 5, 6, 1, 2, 3, 4})});
+  WindowTestBase::testWindowFunction(
+      {input}, "count(c1)", overClause, frameClause, expected);
+
+  // Test with column-specified preceding (int32).
+  frameClause = "rows between c2 preceding and 0 following";
+  expected = makeRowVector(
+      {c0,
+       c1,
+       c2,
+       c3,
+       makeFlatVector<int64_t>({1, 2, 3, 4, 2, 6, 1, 2, 3, 4})});
+  WindowTestBase::testWindowFunction(
+      {input}, "count(c1)", overClause, frameClause, expected);
+
+  // Test with column-specified preceding (int64).
   frameClause = "rows between c3 preceding and 0 following";
   expected = makeRowVector(
       {c0,
@@ -271,6 +305,17 @@ TEST_F(AggregateWindowTest, integerOverflowRowsFrame) {
        makeFlatVector<int64_t>({1, 2, 3, 4, 5, 6, 1, 2, 3, 4})});
   WindowTestBase::testWindowFunction(
       {input}, "count(c1)", overClause, frameClause, expected);
+
+  // Test constant preceding & following both larger than INT32_MAX.
+  frameClause = "rows between 2147483650 preceding and 2147483651 following";
+  expected = makeRowVector(
+      {c0,
+       c1,
+       c2,
+       c3,
+       makeFlatVector<int64_t>({6, 6, 6, 6, 6, 6, 4, 4, 4, 4})});
+  WindowTestBase::testWindowFunction(
+      {input}, "count(c1)", overClause, frameClause, expected);
 }
 
 }; // namespace