Merge nest loop join build vectors (facebookincubator#11428)

Summary:
Pull Request resolved: facebookincubator#11428

When build vectors in nested loop join are small, we should merge them to get better performance.  In some extreme case, the performance difference can be more than 900 times.

Reviewed By: pedroerp

Differential Revision: D65450017

fbshipit-source-id: 370655e418327c6bd33b191b3190d3322b9bd9af

Author: Yuhta

velox/exec/NestedLoopJoinBuild.cpp

@@ -70,6 +70,38 @@ BlockingReason NestedLoopJoinBuild::isBlocked(ContinueFuture* future) {
   return BlockingReason::kWaitForJoinBuild;
 }
 
+// Merge adjacent vectors to larger vectors as long as the result do not exceed
+// the size limit.  This is important for performance because each small vector
+// here would be duplicated by the number of rows on probe side, result in huge
+// number of small vectors in the output.
+std::vector<RowVectorPtr> NestedLoopJoinBuild::mergeDataVectors() const {
+  const auto maxBatchRows =
+      operatorCtx_->task()->queryCtx()->queryConfig().maxOutputBatchRows();
+  std::vector<RowVectorPtr> merged;
+  for (int i = 0; i < dataVectors_.size();) {
+    auto batchSize = dataVectors_[i]->size();
+    auto j = i + 1;
+    while (j < dataVectors_.size() &&
+           batchSize + dataVectors_[j]->size() <= maxBatchRows) {
+      batchSize += dataVectors_[j++]->size();
+    }
+    if (j == i + 1) {
+      merged.push_back(dataVectors_[i++]);
+    } else {
+      auto batch = BaseVector::create<RowVector>(
+          dataVectors_[i]->type(), batchSize, pool());
+      batchSize = 0;
+      while (i < j) {
+        auto* source = dataVectors_[i++].get();
+        batch->copy(source, batchSize, 0, source->size());
+        batchSize += source->size();
+      }
+      merged.push_back(std::move(batch));
+    }
+  }
+  return merged;
+}
+
 void NestedLoopJoinBuild::noMoreInput() {
   Operator::noMoreInput();
   std::vector<ContinuePromise> promises;
@@ -105,6 +137,7 @@ void NestedLoopJoinBuild::noMoreInput() {
     }
   }
 
+  dataVectors_ = mergeDataVectors();
   operatorCtx_->task()
       ->getNestedLoopJoinBridge(
           operatorCtx_->driverCtx()->splitGroupId, planNodeId())

velox/exec/NestedLoopJoinBuild.h

@@ -59,6 +59,8 @@ class NestedLoopJoinBuild : public Operator {
   }
 
  private:
+  std::vector<RowVectorPtr> mergeDataVectors() const;
+
   std::vector<RowVectorPtr> dataVectors_;
 
   // Future for synchronizing with other Drivers of the same pipeline. All build

velox/exec/tests/NestedLoopJoinTest.cpp

@@ -569,5 +569,58 @@ TEST_F(NestedLoopJoinTest, outputOrder) {
   assertEqualVectors(expectedLeft, results);
 }
 
+TEST_F(NestedLoopJoinTest, mergeBuildVectors) {
+  const std::vector<RowVectorPtr> buildVectors = {
+      makeRowVector({makeFlatVector<int64_t>({1, 2})}),
+      makeRowVector({makeFlatVector<int64_t>({3, 4})}),
+      makeRowVector(
+          {makeFlatVector<int64_t>(20, [](auto i) { return 5 + i; })}),
+      makeRowVector(
+          {makeFlatVector<int64_t>(20, [](auto i) { return 25 + i; })}),
+      makeRowVector({makeFlatVector<int64_t>({45, 46})}),
+  };
+  const std::vector<RowVectorPtr> probeVectors = {
+      makeRowVector({makeFlatVector<int64_t>({1, 2})}),
+  };
+  auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+  CursorParameters params;
+  params.planNode = PlanBuilder(planNodeIdGenerator)
+                        .values(probeVectors)
+                        .nestedLoopJoin(
+                            PlanBuilder(planNodeIdGenerator)
+                                .values(buildVectors)
+                                .project({"c0 as r0"})
+                                .planNode(),
+                            {"c0", "r0"})
+                        .planNode();
+  params.queryConfigs[core::QueryConfig::kMaxOutputBatchRows] = "10";
+  auto cursor = TaskCursor::create(params);
+  // Expect the first 2 build side vectors are merged together since they are
+  // under the limit after merge.  Others are left along.
+  for (int i = 0; i < 2; ++i) {
+    auto makeExpected = [&](vector_size_t size, vector_size_t buildOffset) {
+      return makeRowVector({
+          makeConstant<int64_t>(1 + i, size),
+          makeFlatVector<int64_t>(
+              size, [&](auto i) { return buildOffset + i; }),
+      });
+    };
+    ASSERT_TRUE(cursor->moveNext());
+    ASSERT_EQ(cursor->current()->size(), 4);
+    assertEqualVectors(makeExpected(4, 1), cursor->current());
+    ASSERT_TRUE(cursor->moveNext());
+    ASSERT_EQ(cursor->current()->size(), 20);
+    assertEqualVectors(makeExpected(20, 5), cursor->current());
+    ASSERT_TRUE(cursor->moveNext());
+    ASSERT_EQ(cursor->current()->size(), 20);
+    assertEqualVectors(makeExpected(20, 25), cursor->current());
+    ASSERT_TRUE(cursor->moveNext());
+    ASSERT_EQ(cursor->current()->size(), 2);
+    assertEqualVectors(makeExpected(2, 45), cursor->current());
+  }
+  ASSERT_FALSE(cursor->moveNext());
+  ASSERT_TRUE(waitForTaskCompletion(cursor->task().get()));
+}
+
 } // namespace
 } // namespace facebook::velox::exec::test