Fix finalSelection when null value's evalution (facebookincubator#2073)

barsondei · Jul 22, 2022 · cd8ba50 · cd8ba50
1 parent bc4da9a
commit cd8ba50
Show file tree

Hide file tree

Showing 3 changed files with 66 additions and 0 deletions.
diff --git a/velox/exec/tests/TableScanTest.cpp b/velox/exec/tests/TableScanTest.cpp
@@ -1830,6 +1830,28 @@ TEST_F(TableScanTest, structLazy) {
   assertQuery(op, {filePath}, "select c0 % 3 from tmp");
 }
 
+TEST_F(TableScanTest, lazyVectorAccessTwiceWithDifferentRows) {
+  auto data = makeRowVector({
+      makeNullableFlatVector<int64_t>({1, 1, 1, std::nullopt}),
+      makeNullableFlatVector<int64_t>({0, 1, 2, 3}),
+  });
+
+  auto filePath = TempFilePath::create();
+  writeToFile(filePath->path, {data});
+  createDuckDbTable({data});
+
+  auto plan =
+      PlanBuilder()
+          .tableScan(asRowType(data->type()))
+          .filter(
+              "element_at(array_constructor(c0 + c1, if(c1 >= 0, c1, 0)), 1) > 0")
+          .planNode();
+  assertQuery(
+      plan,
+      {filePath},
+      "SELECT c0, c1 from tmp where ([c0 + c1, if(c1 >= 0, c1, 0)])[1] > 0");
+}
+
 TEST_F(TableScanTest, structInArrayOrMap) {
   vector_size_t size = 1'000;
 

diff --git a/velox/expression/Expr.cpp b/velox/expression/Expr.cpp
@@ -774,6 +774,8 @@ void Expr::evalWithNulls(
       if (removeSureNulls(rows, context, nonNullHolder)) {
         VarSetter noMoreNulls(context.mutableNullsPruned(), true);
         if (nonNullHolder.get()->hasSelections()) {
+          // No need fix finalSelection here, LazyVector already loaded due to
+          // removeSureNulls method
           evalAll(*nonNullHolder.get(), context, result);
         }
         auto rawNonNulls = nonNullHolder.get()->asRange().bits();
@@ -988,6 +990,14 @@ void Expr::evalAll(
   bool defaultNulls = vectorFunction_->isDefaultNullBehavior();
   inputValues_.resize(inputs_.size());
   for (int32_t i = 0; i < inputs_.size(); ++i) {
+    // Fix finalSelection at "rows" if missingRows is a strict subset.
+    // "rows" may be used to evaluate exprs outside of current expr node.
+    bool updateFinalSelection = context.isFinalSelection() &&
+        (remainingRows->countSelected() < rows.countSelected());
+    VarSetter isFinalSelection(
+        context.mutableIsFinalSelection(), false, updateFinalSelection);
+    VarSetter finalSelection(
+        context.mutableFinalSelection(), &rows, updateFinalSelection);
     inputs_[i]->eval(*remainingRows, context, inputValues_[i]);
     tryPeelArgs = tryPeelArgs && isPeelable(inputValues_[i]->encoding());
     if (defaultNulls && inputValues_[i]->mayHaveNulls()) {

diff --git a/velox/expression/tests/ExprTest.cpp b/velox/expression/tests/ExprTest.cpp
@@ -1649,6 +1649,40 @@ TEST_F(ExprTest, selectiveLazyLoadingOr) {
   assertEqualVectors(expected, result);
 }
 
+TEST_F(ExprTest, lazyVectorAccessTwiceWithDifferentRows) {
+  const vector_size_t size = 4;
+
+  // [1, 1, 1, null]
+  auto inputC0 = makeNullableFlatVector<int64_t>({1, 1, 1, std::nullopt});
+  // [0, 1, 2, 3] if fully loaded
+  std::vector<vector_size_t> loadedRows;
+  auto valueAt = [](auto row) { return row; };
+  VectorPtr inputC1 = std::make_shared<LazyVector>(
+      pool_.get(),
+      BIGINT(),
+      size,
+      std::make_unique<test::SimpleVectorLoader>([&](auto rows) {
+        for (auto row : rows) {
+          loadedRows.push_back(row);
+        }
+        return makeFlatVector<int64_t>(rows.back() + 1, valueAt);
+      }));
+
+  // isFinalSelection_ == true
+  auto result = evaluate(
+      "row_constructor(c0 + c1, if (c1 >= 0, c1, 0))",
+      makeRowVector({inputC0, inputC1}));
+
+  // [1, 2, 3, null]
+  auto outputCol0 = makeNullableFlatVector<int64_t>({1, 2, 3, std::nullopt});
+  // [0, 1, 2, 3]
+  auto outputCol1 = makeNullableFlatVector<int64_t>({0, 1, 2, 3});
+  // [(1, 0), (2, 1), (3, 2), (null, 3)]
+  auto expected = ExprTest::makeRowVector({outputCol0, outputCol1});
+
+  assertEqualVectors(expected, result);
+}
+
 TEST_F(ExprTest, selectiveLazyLoadingIf) {
   const vector_size_t size = 1'000;