Fix the semi join result mismatch issue with filter and multi duplicated rows

Author: JkSelf

velox/exec/MergeJoin.cpp

@@ -90,7 +90,9 @@ void MergeJoin::initialize() {
     initializeFilter(joinNode_->filter(), leftType, rightType);
 
     if (joinNode_->isLeftJoin() || joinNode_->isAntiJoin() ||
-        joinNode_->isRightJoin() || joinNode_->isFullJoin()) {
+        joinNode_->isRightJoin() || joinNode_->isFullJoin() ||
+        joinNode_->isLeftSemiFilterJoin() ||
+        joinNode_->isRightSemiFilterJoin()) {
       joinTracker_ = JoinTracker(outputBatchSize_, pool());
     }
   } else if (joinNode_->isAntiJoin()) {
@@ -421,7 +423,7 @@ bool MergeJoin::tryAddOutputRow(
         filterRightInputProjections_);
 
     if (joinTracker_) {
-      if (isRightJoin(joinType_)) {
+      if (isRightJoin(joinType_) || isRightSemiFilterJoin(joinType_)) {
         // Record right-side row with a match on the left-side.
         joinTracker_->addMatch(rightBatch, rightRow, outputSize_);
       } else {
@@ -613,16 +615,17 @@ bool MergeJoin::addToOutputForLeftJoin() {
       // one match on the other side, we could explore specialized algorithms
       // or data structures that short-circuit the join process once a match
       // is found.
-      for (size_t r = isLeftSemiFilterJoin(joinType_) ? numRightBatches - 1
-                                                      : firstRightBatch;
+      for (size_t r = (isLeftSemiFilterJoin(joinType_) && !filter_)
+               ? numRightBatches - 1
+               : firstRightBatch;
            r < numRightBatches;
            ++r) {
         const auto rightBatch = rightMatch_->inputs[r];
         auto rightStartRow = r == firstRightBatch ? rightStartRowIndex : 0;
         const auto rightEndRow = r == numRightBatches - 1
             ? rightMatch_->endRowIndex
             : rightBatch->size();
-        if (isLeftSemiFilterJoin(joinType_)) {
+        if (isLeftSemiFilterJoin(joinType_) && !filter_) {
           rightStartRow = rightEndRow - 1;
         }
         if (prepareOutput(leftBatch, rightBatch)) {
@@ -693,16 +696,17 @@ bool MergeJoin::addToOutputForRightJoin() {
       // one match on the other side, we could explore specialized algorithms
       // or data structures that short-circuit the join process once a match
       // is found.
-      for (size_t l = isRightSemiFilterJoin(joinType_) ? numLeftBatches - 1
-                                                       : firstLeftBatch;
+      for (size_t l = (isRightSemiFilterJoin(joinType_) && !filter_)
+               ? numLeftBatches - 1
+               : firstLeftBatch;
            l < numLeftBatches;
            ++l) {
         const auto leftBatch = leftMatch_->inputs[l];
         auto leftStartRow = l == firstLeftBatch ? leftStartRowIndex : 0;
         const auto leftEndRow = l == numLeftBatches - 1
             ? leftMatch_->endRowIndex
             : leftBatch->size();
-        if (isRightSemiFilterJoin(joinType_)) {
+        if (isRightSemiFilterJoin(joinType_) && !filter_) {
           // RightSemiFilter produce each row from the right at most once.
           leftStartRow = leftEndRow - 1;
         }
@@ -818,7 +822,7 @@ RowVectorPtr MergeJoin::getOutput() {
         continue;
       } else if (isAntiJoin(joinType_)) {
         output = filterOutputForAntiJoin(output);
-        if (output) {
+        if (output != nullptr && output->size() > 0) {
           return output;
         }
 
@@ -1274,7 +1278,8 @@ RowVectorPtr MergeJoin::applyFilter(const RowVectorPtr& output) {
     // If all matches for a given left-side row fail the filter, add a row to
     // the output with nulls for the right-side columns.
     const auto onMiss = [&](auto row) {
-      if (isAntiJoin(joinType_)) {
+      if (isAntiJoin(joinType_) || isLeftSemiFilterJoin(joinType_) ||
+          isRightSemiFilterJoin(joinType_)) {
         return;
       }
       rawIndices[numPassed++] = row;
@@ -1346,18 +1351,26 @@ RowVectorPtr MergeJoin::applyFilter(const RowVectorPtr& output) {
       }
     };
 
+    auto onMatch = [&](auto row) {
+      if (isLeftSemiFilterJoin(joinType_) || isRightSemiFilterJoin(joinType_)) {
+        rawIndices[numPassed++] = row;
+      }
+    };
+
     for (auto i = 0; i < numRows; ++i) {
       if (filterRows.isValid(i)) {
         const bool passed = !decodedFilterResult_.isNullAt(i) &&
             decodedFilterResult_.valueAt<bool>(i);
 
-        joinTracker_->processFilterResult(i, passed, onMiss);
+        joinTracker_->processFilterResult(i, passed, onMiss, onMatch);
 
         if (isAntiJoin(joinType_)) {
           if (!passed) {
             rawIndices[numPassed++] = i;
           }
-        } else {
+        } else if (
+            !isLeftSemiFilterJoin(joinType_) &&
+            !isRightSemiFilterJoin(joinType_)) {
           if (passed) {
             rawIndices[numPassed++] = i;
           }
@@ -1371,26 +1384,27 @@ RowVectorPtr MergeJoin::applyFilter(const RowVectorPtr& output) {
 
     // Every time we start a new left key match, `processFilterResult()` will
     // check if at least one row from the previous match passed the filter. If
-    // none did, it calls onMiss to add a record with null right projections to
-    // the output.
+    // none did, it calls onMiss to add a record with null right projections
+    // to the output.
     //
     // Before we leave the current buffer, since we may not have seen the next
-    // left key match yet, the last key match may still be pending to produce a
-    // row (because `processFilterResult()` was not called yet).
+    // left key match yet, the last key match may still be pending to produce
+    // a row (because `processFilterResult()` was not called yet).
     //
     // To handle this, we need to call `noMoreFilterResults()` unless the
-    // same current left key match may continue in the next buffer. So there are
-    // two cases to check:
+    // same current left key match may continue in the next buffer. So there
+    // are two cases to check:
     //
-    // 1. If leftMatch_ is nullopt, there for sure the next buffer will contain
-    // a different key match.
+    // 1. If leftMatch_ is nullopt, there for sure the next buffer will
+    // contain a different key match.
     //
     // 2. leftMatch_ may not be nullopt, but may be related to a different
     // (subsequent) left key. So we check if the last row in the batch has the
     // same left row number as the last key match.
     if (!leftMatch_ || !joinTracker_->isCurrentLeftMatch(numRows - 1)) {
       joinTracker_->noMoreFilterResults(onMiss);
     }
+
   } else {
     filterRows_.resize(numRows);
     filterRows_.setAll();

velox/exec/MergeJoin.h

@@ -394,11 +394,12 @@ class MergeJoin : public Operator {
     // rows that correspond to a single left-side row. Use
     // 'noMoreFilterResults' to make sure 'onMiss' is called for the last
     // left-side row.
-    template <typename TOnMiss>
+    template <typename TOnMiss, typename TOnMatch>
     void processFilterResult(
         vector_size_t outputIndex,
         bool passed,
-        TOnMiss onMiss) {
+        TOnMiss onMiss,
+        TOnMatch onMatch) {
       const auto rowNumber = rawLeftRowNumbers_[outputIndex];
       if (currentLeftRowNumber_ != rowNumber) {
         if (currentRow_ != -1 && !currentRowPassed_) {
@@ -407,12 +408,18 @@ class MergeJoin : public Operator {
         currentRow_ = outputIndex;
         currentLeftRowNumber_ = rowNumber;
         currentRowPassed_ = false;
+        firstMatched_ = false;
       } else {
         currentRow_ = outputIndex;
       }
 
       if (passed) {
         currentRowPassed_ = true;
+
+        if (!firstMatched_) {
+          onMatch(outputIndex);
+          firstMatched_ = true;
+        }
       }
     }
 
@@ -434,6 +441,7 @@ class MergeJoin : public Operator {
 
       currentRow_ = -1;
       currentRowPassed_ = false;
+      firstMatched_ = false;
     }
 
     void reset();
@@ -470,6 +478,10 @@ class MergeJoin : public Operator {
     // True if at least one row in a block of output rows corresponding a single
     // left-side row identified by 'currentRowNumber' passed the filter.
     bool currentRowPassed_{false};
+
+    // Retains only the first matching record for a semi join in scenarios
+    // involving filters.
+    bool firstMatched_{false};
   };
 
   /// Used to record both left and right join.

velox/exec/tests/MergeJoinTest.cpp

@@ -1016,6 +1016,58 @@ TEST_F(MergeJoinTest, semiJoinWithMultipleMatchVectors) {
       core::JoinType::kLeftSemiFilter);
 }
 
+TEST_F(MergeJoinTest, semiJoinWithMultiMatchedRowsWithFilter) {
+  auto left = makeRowVector(
+      {"t0", "t1"},
+      {makeNullableFlatVector<int64_t>({2, 2, 2, 2, 2}),
+       makeNullableFlatVector<int64_t>({3, 2, 3, 2, 2})});
+
+  auto right = makeRowVector(
+      {"u0", "u1"},
+      {makeNullableFlatVector<int64_t>({2, 2, 2, 2, 2, 2}),
+       makeNullableFlatVector<int64_t>({2, 2, 2, 2, 2, 4})});
+
+  createDuckDbTable("t", {left});
+  createDuckDbTable("u", {right});
+
+  auto testSemiJoin = [&](const std::string& filter,
+                          const std::string& sql,
+                          const std::vector<std::string>& outputLayout,
+                          core::JoinType joinType) {
+    auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+    auto plan = PlanBuilder(planNodeIdGenerator)
+                    .values(split(left, 2))
+                    .mergeJoin(
+                        {"t0"},
+                        {"u0"},
+                        PlanBuilder(planNodeIdGenerator)
+                            .values(split(right, 2))
+                            .planNode(),
+                        filter,
+                        outputLayout,
+                        joinType)
+                    .planNode();
+    AssertQueryBuilder(plan, duckDbQueryRunner_)
+        .config(core::QueryConfig::kPreferredOutputBatchRows, "2")
+        .config(core::QueryConfig::kMaxOutputBatchRows, "2")
+        .assertResults(sql);
+  };
+
+  // Left Semi join With filter
+  testSemiJoin(
+      "t1 > u1",
+      "SELECT t0, t1 FROM t where t0 IN (SELECT u0 from u where t1 > u1)",
+      {"t0", "t1"},
+      core::JoinType::kLeftSemiFilter);
+
+  // Right Semi join With filter
+  testSemiJoin(
+      "u1 > t1",
+      "SELECT u0, u1 FROM u where u0 IN (SELECT t0 from t where u1 > t1)",
+      {"u0", "u1"},
+      core::JoinType::kRightSemiFilter);
+}
+
 TEST_F(MergeJoinTest, rightJoin) {
   auto left = makeRowVector(
       {"t0"},