Fix null value check for kRange frame column

facebookincubator · Sep 24, 2024 · ba41a62 · ba41a62
1 parent 4e45bc5
commit ba41a62
Show file tree

Hide file tree

Showing 2 changed files with 47 additions and 2 deletions.
diff --git a/velox/exec/WindowPartition.cpp b/velox/exec/WindowPartition.cpp
@@ -363,8 +363,11 @@ void WindowPartition::updateKRangeFrameBounds(
 
   vector_size_t start = 0;
   vector_size_t end;
+  // frameColumn is a column index into the original input rows, while
+  // orderByColumn is a column index into rows in data_ after the columns are
+  // reordered as per inputMapping_.
   RowColumn frameRowColumn = columns_[frameColumn];
-  RowColumn orderByRowColumn = columns_[inputMapping_[orderByColumn]];
+  RowColumn orderByRowColumn = data_->columnAt(orderByColumn);
   for (auto i = 0; i < numRows; i++) {
     auto currentRow = startRow + i;
     auto* partitionRow = partition_[currentRow];
@@ -378,7 +381,9 @@ void WindowPartition::updateKRangeFrameBounds(
         RowContainer::isNullAt(
             partitionRow,
             orderByRowColumn.nullByte(),
-            orderByRowColumn.nullMask()));
+            orderByRowColumn.nullMask()),
+        "Frame column should be NULL iff ORDER BY value is NULL for row: {}",
+        partitionRow);
 
     // If the frame is NULL or 0 preceding or 0 following then the current row
     // has same values for order by and frame column. In that case

diff --git a/velox/exec/tests/WindowTest.cpp b/velox/exec/tests/WindowTest.cpp
@@ -487,5 +487,45 @@ TEST_F(WindowTest, nagativeFrameArg) {
   }
 }
 
+DEBUG_ONLY_TEST_F(WindowTest, frameColumnNullCheck) {
+  // Null values in order-by column 's0' and frame column 'off0' do not match,
+  // so exception is expected.
+  auto inputThrow = {makeRowVector(
+      {"c0", "p0", "s0", "off0"},
+      {
+          makeNullableFlatVector<int64_t>({1, std::nullopt, 1, 2, 2}),
+          makeFlatVector<int64_t>({1, 2, 1, 2, 1}),
+          makeNullableFlatVector<int64_t>({1, 2, 3, std::nullopt, 5}),
+          makeNullableFlatVector<int64_t>({2, std::nullopt, 4, 5, 6}),
+      })};
+  auto planNode =
+      PlanBuilder(std::make_shared<core::PlanNodeIdGenerator>())
+          .values(inputThrow)
+          .window(
+              {"sum(c0) OVER (PARTITION BY p0 ORDER BY s0 RANGE BETWEEN UNBOUNDED PRECEDING AND off0 FOLLOWING)"})
+          .planNode();
+  VELOX_ASSERT_THROW(
+      AssertQueryBuilder(planNode).copyResults(pool()),
+      "Frame column should be NULL iff ORDER BY value is NULL");
+
+  // Null values in order-by column 's0' and frame column 'off0' match, so no
+  // exception should be thrown.
+  auto inputNoThrow = {makeRowVector(
+      {"c0", "p0", "s0", "off0"},
+      {
+          makeNullableFlatVector<int64_t>({1, 1, 2, std::nullopt, 2}),
+          makeFlatVector<int64_t>({1, 1, 1, 2, 2}),
+          makeNullableFlatVector<int64_t>({1, std::nullopt, 2, 3, 5}),
+          makeNullableFlatVector<int64_t>({2, std::nullopt, 3, 4, 6}),
+      })};
+  planNode =
+      PlanBuilder(std::make_shared<core::PlanNodeIdGenerator>())
+          .values(inputNoThrow)
+          .window(
+              {"sum(c0) OVER (PARTITION BY p0 ORDER BY s0 RANGE BETWEEN UNBOUNDED PRECEDING AND off0 FOLLOWING)"})
+          .planNode();
+  ASSERT_NO_THROW(AssertQueryBuilder(planNode).copyResults(pool()));
+}
+
 } // namespace
 } // namespace facebook::velox::exec