misc: Add option to IndexLookupJoinNode to enable/disable splitting of output (#17507)

Summary:
Pull Request resolved: #17507

Similar to #16762 which added a per-operator splitOutput flag to UnnestNode, this adds the same option to IndexLookupJoinNode.

The existing QueryConfig `index_lookup_join_split_output` applies to all IndexLookupJoin operators in the plan. This is too coarse — the decision to split output depends on downstream operators and should be made per-operator.

This adds a new `splitOutput` parameter (`std::optional<bool>`) to IndexLookupJoinNode. When set, the operator uses this value instead of the QueryConfig. When unset (default), the operator falls back to the QueryConfig, preserving existing behavior.

Reviewed By: xiaoxmeng

Differential Revision: D105021498

fbshipit-source-id: 848510099a2789e29c39252d322b526bb27cf104

Author: kevinwilfong

velox/core/PlanNode.cpp

@@ -1820,9 +1820,35 @@ IndexLookupJoinNode::IndexLookupJoinNode(
     const std::vector<IndexLookupConditionPtr>& joinConditions,
     TypedExprPtr filter,
     bool hasMarker,
+    std::optional<bool> splitOutput,
     PlanNodePtr left,
     TableScanNodePtr right,
     RowTypePtr outputType)
+    : IndexLookupJoinNode(
+          id,
+          joinType,
+          leftKeys,
+          rightKeys,
+          joinConditions,
+          std::move(filter),
+          hasMarker,
+          std::move(left),
+          std::move(right),
+          std::move(outputType),
+          splitOutput) {}
+
+IndexLookupJoinNode::IndexLookupJoinNode(
+    const PlanNodeId& id,
+    JoinType joinType,
+    const std::vector<FieldAccessTypedExprPtr>& leftKeys,
+    const std::vector<FieldAccessTypedExprPtr>& rightKeys,
+    const std::vector<IndexLookupConditionPtr>& joinConditions,
+    TypedExprPtr filter,
+    bool hasMarker,
+    PlanNodePtr left,
+    TableScanNodePtr right,
+    RowTypePtr outputType,
+    std::optional<bool> splitOutput)
     : AbstractJoinNode(
           id,
           joinType,
@@ -1834,7 +1860,8 @@ IndexLookupJoinNode::IndexLookupJoinNode(
           outputType),
       lookupSourceNode_(std::move(right)),
       joinConditions_(joinConditions),
-      hasMarker_(hasMarker) {
+      hasMarker_(hasMarker),
+      splitOutput_(splitOutput) {
   VELOX_USER_CHECK(
       !leftKeys.empty(),
       "The index lookup join node requires at least one join key");
@@ -1921,6 +1948,11 @@ PlanNodePtr IndexLookupJoinNode::create(
 
   const bool hasMarker = obj["hasMarker"].asBool();
 
+  std::optional<bool> splitOutput = std::nullopt;
+  if (obj.count("splitOutput")) {
+    splitOutput = obj["splitOutput"].asBool();
+  }
+
   auto outputType = deserializeRowType(obj["outputType"]);
 
   return std::make_shared<IndexLookupJoinNode>(
@@ -1931,6 +1963,7 @@ PlanNodePtr IndexLookupJoinNode::create(
       std::move(joinConditions),
       filter,
       hasMarker,
+      splitOutput,
       sources[0],
       std::move(lookupSource),
       std::move(outputType));
@@ -1949,6 +1982,9 @@ folly::dynamic IndexLookupJoinNode::serialize() const {
     obj["filter"] = filter_->serialize();
   }
   obj["hasMarker"] = hasMarker_;
+  if (splitOutput_.has_value()) {
+    obj["splitOutput"] = splitOutput_.value();
+  }
   return obj;
 }
 

velox/core/PlanNode.h

@@ -3711,6 +3711,26 @@ class IndexLookupJoinNode : public AbstractJoinNode {
       bool hasMarker,
       PlanNodePtr left,
       TableScanNodePtr right,
+      RowTypePtr outputType,
+      std::optional<bool> splitOutput = std::nullopt);
+
+  /// @param splitOutput Optional flag to control whether the operator should
+  /// split output batches if they are too large. If true, output is split into
+  /// batches according to Operator's outputBatchRows logic. If false, output is
+  /// not split and output batches match input batches 1:1. If not set, defaults
+  /// to the value of the index_lookup_join_split_output config in the
+  /// QueryConfig.
+  IndexLookupJoinNode(
+      const PlanNodeId& id,
+      JoinType joinType,
+      const std::vector<FieldAccessTypedExprPtr>& leftKeys,
+      const std::vector<FieldAccessTypedExprPtr>& rightKeys,
+      const std::vector<IndexLookupConditionPtr>& joinConditions,
+      TypedExprPtr filter,
+      bool hasMarker,
+      std::optional<bool> splitOutput,
+      PlanNodePtr left,
+      TableScanNodePtr right,
       RowTypePtr outputType);
 
   class Builder
@@ -3723,6 +3743,7 @@ class IndexLookupJoinNode : public AbstractJoinNode {
       joinConditions_ = other.joinConditions();
       filter_ = other.filter();
       hasMarker_ = other.hasMarker();
+      splitOutput_ = other.splitOutput();
     }
 
     /// Set lookup conditions for index lookup that can't be converted into
@@ -3745,6 +3766,13 @@ class IndexLookupJoinNode : public AbstractJoinNode {
       return *this;
     }
 
+    /// Set whether to split large output into multiple batches, std::nullopt
+    /// means respect index_lookup_join_split_output in the QueryConfig.
+    Builder& splitOutput(std::optional<bool> splitOutput) {
+      splitOutput_ = splitOutput;
+      return *this;
+    }
+
     std::shared_ptr<IndexLookupJoinNode> build() const {
       VELOX_USER_CHECK(id_.has_value(), "IndexLookupJoinNode id is not set");
       VELOX_USER_CHECK(
@@ -3768,6 +3796,7 @@ class IndexLookupJoinNode : public AbstractJoinNode {
           joinConditions_,
           filter_.value_or(nullptr),
           hasMarker_,
+          splitOutput_,
           left_.value(),
           std::dynamic_pointer_cast<const TableScanNode>(right_.value()),
           outputType_.value());
@@ -3776,6 +3805,7 @@ class IndexLookupJoinNode : public AbstractJoinNode {
    private:
     std::vector<IndexLookupConditionPtr> joinConditions_;
     bool hasMarker_{false};
+    std::optional<bool> splitOutput_;
   };
 
   bool supportsBarrier() const override {
@@ -3805,6 +3835,10 @@ class IndexLookupJoinNode : public AbstractJoinNode {
     return hasMarker_;
   }
 
+  const std::optional<bool>& splitOutput() const {
+    return splitOutput_;
+  }
+
   void accept(const PlanNodeVisitor& visitor, PlanNodeVisitorContext& context)
       const override;
 
@@ -3828,6 +3862,10 @@ class IndexLookupJoinNode : public AbstractJoinNode {
 
   /// Whether to include a marker column for left joins to indicate matches.
   const bool hasMarker_;
+
+  /// Optional flag to control whether to split output batches. When set,
+  /// overrides the index_lookup_join_split_output QueryConfig.
+  const std::optional<bool> splitOutput_;
 };
 
 using IndexLookupJoinNodePtr = std::shared_ptr<const IndexLookupJoinNode>;

velox/core/QueryConfig.h

@@ -1335,7 +1335,9 @@ class QueryConfig {
       0,
       "Maximum input batches to prefetch for index lookup. 0 disables.")
 
-  /// If true, index join operator may split output per input batch.
+  /// If true, index join operator may split output per input batch. This can be
+  /// overridden on a per operator basis by the splitOutput parameter in the
+  /// IndexLookupJoinNode.
   VELOX_QUERY_CONFIG(
       kIndexLookupJoinSplitOutput,
       indexLookupJoinSplitOutput,

velox/core/tests/PlanNodeBuilderTest.cpp

@@ -776,6 +776,7 @@ TEST_F(PlanNodeBuilderTest, indexLookupJoinNode) {
           .assignments({{"c1", std::make_shared<DummyColumnHandle>()}})
           .build();
   const auto outputType = ROW({"c0"}, {BIGINT()});
+  std::optional<bool> splitOutput = true;
 
   const auto verify =
       [&](const std::shared_ptr<const IndexLookupJoinNode>& node) {
@@ -790,6 +791,7 @@ TEST_F(PlanNodeBuilderTest, indexLookupJoinNode) {
         EXPECT_EQ(node->sources()[0], left);
         EXPECT_EQ(node->sources()[1], right);
         EXPECT_EQ(node->outputType(), outputType);
+        EXPECT_EQ(node->splitOutput(), splitOutput);
       };
 
   const auto node = IndexLookupJoinNode::Builder()
@@ -801,6 +803,7 @@ TEST_F(PlanNodeBuilderTest, indexLookupJoinNode) {
                         .left(left)
                         .right(right)
                         .outputType(outputType)
+                        .splitOutput(splitOutput)
                         .build();
   verify(node);
 

velox/exec/IndexLookupJoin.cpp

@@ -302,13 +302,16 @@ IndexLookupJoin::IndexLookupJoin(
           operatorId,
           joinNode->id(),
           OperatorType::kIndexLookupJoin),
-      splitOutput_{driverCtx->queryConfig().indexLookupJoinSplitOutput()},
+      splitOutput_{
+          (joinNode->splitOutput().has_value() &&
+           joinNode->splitOutput().value()) ||
+          (!joinNode->splitOutput().has_value() &&
+           driverCtx->queryConfig().indexLookupJoinSplitOutput())},
       // TODO: support to update output batch size with output size stats during
       // the lookup processing.
       outputBatchSize_{
-          driverCtx->queryConfig().indexLookupJoinSplitOutput()
-              ? outputBatchRows()
-              : std::numeric_limits<vector_size_t>::max()},
+          splitOutput_ ? outputBatchRows()
+                       : std::numeric_limits<vector_size_t>::max()},
       joinType_{joinNode->joinType()},
       hasMarker_(joinNode->hasMarker()),
       probeType_{joinNode->sources()[0]->outputType()},

velox/exec/tests/IndexLookupJoinTest.cpp

@@ -492,6 +492,28 @@ TEST_P(IndexLookupJoinTest, planNodeAndSerde) {
     testSerde(plan);
   }
 
+  // with splitOutput.
+  for (const auto splitOutput :
+       {std::optional<bool>(true),
+        std::optional<bool>(false),
+        std::optional<bool>(std::nullopt)}) {
+    auto plan = PlanBuilder(planNodeIdGenerator)
+                    .values({left})
+                    .startIndexLookupJoin()
+                    .leftKeys({"t0"})
+                    .rightKeys({"u0"})
+                    .indexSource(indexTableScan)
+                    .outputLayout({"t0", "u1", "t2", "t1"})
+                    .joinType(core::JoinType::kInner)
+                    .splitOutput(splitOutput)
+                    .endIndexLookupJoin()
+                    .planNode();
+    auto indexLookupJoinNode =
+        std::dynamic_pointer_cast<const core::IndexLookupJoinNode>(plan);
+    ASSERT_EQ(indexLookupJoinNode->splitOutput(), splitOutput);
+    testSerde(plan);
+  }
+
   // bad join type.
   {
     VELOX_ASSERT_USER_THROW(
@@ -2664,6 +2686,130 @@ TEST_P(IndexLookupJoinTest, outputBatchSizeWithLeftJoin) {
   }
 }
 
+TEST_P(IndexLookupJoinTest, splitOutputNodeOverride) {
+  IndexTableData tableData;
+  generateIndexTableData({3'000, 1, 1}, tableData, pool_);
+
+  const int numProbeBatches = 10;
+  const int numRowsPerProbeBatch = 100;
+  const int maxBatchRows = 10;
+
+  const auto probeVectors = generateProbeInput(
+      numProbeBatches,
+      numRowsPerProbeBatch,
+      1,
+      tableData,
+      pool_,
+      {"t0", "t1", "t2"},
+      GetParam().hasNullKeys,
+      {},
+      {},
+      /*equalMatchPct=*/100);
+  std::vector<std::shared_ptr<TempFilePath>> probeFiles =
+      createProbeFiles(probeVectors);
+
+  createDuckDbTable("t", probeVectors);
+  createDuckDbTable("u", {tableData.tableVectors});
+
+  const auto indexTable = TestIndexTable::create(
+      /*numEqualJoinKeys=*/3,
+      tableData.keyVectors,
+      tableData.valueVectors,
+      *pool());
+  const auto indexTableHandle = makeIndexTableHandle(
+      indexTable, GetParam().asyncLookup, GetParam().needsIndexSplit);
+
+  struct {
+    std::optional<bool> nodeSplitOutput;
+    bool configSplitOutput;
+    bool expectSplit;
+
+    std::string debugString() const {
+      return fmt::format(
+          "nodeSplitOutput: {}, configSplitOutput: {}, expectSplit: {}",
+          nodeSplitOutput.has_value()
+              ? (nodeSplitOutput.value() ? "true" : "false")
+              : "nullopt",
+          configSplitOutput,
+          expectSplit);
+    }
+  } testSettings[] = {
+      // Node splitOutput not set, use config.
+      {std::nullopt, true, true},
+      {std::nullopt, false, false},
+      // Node splitOutput=true overrides config.
+      {true, true, true},
+      {true, false, true},
+      // Node splitOutput=false overrides config.
+      {false, true, false},
+      {false, false, false},
+  };
+
+  for (const auto& testData : testSettings) {
+    SCOPED_TRACE(testData.debugString());
+
+    auto planNodeIdGenerator = std::make_shared<core::PlanNodeIdGenerator>();
+
+    const auto indexScanNode = makeIndexScanNode(
+        planNodeIdGenerator,
+        indexTableHandle,
+        makeScanOutputType({"u0", "u1", "u2", "u5"}),
+        makeIndexColumnHandles({"u0", "u1", "u2", "u5"}));
+
+    core::PlanNodeId joinNodeId;
+    auto plan = PlanBuilder(planNodeIdGenerator, pool_.get())
+                    .startTableScan()
+                    .outputType(probeType_)
+                    .endTableScan()
+                    .captureScanNodeId(probeScanNodeId_)
+                    .startIndexLookupJoin()
+                    .leftKeys({"t0", "t1", "t2"})
+                    .rightKeys({"u0", "u1", "u2"})
+                    .indexSource(indexScanNode)
+                    .outputLayout({"t4", "u5"})
+                    .joinType(core::JoinType::kInner)
+                    .splitOutput(testData.nodeSplitOutput)
+                    .endIndexLookupJoin()
+                    .capturePlanNodeId(joinNodeId)
+                    .planNode();
+
+    const int expectedNumOutputVectors = testData.expectSplit
+        ? ((numRowsPerProbeBatch + maxBatchRows - 1) / maxBatchRows) *
+            numProbeBatches
+        : numProbeBatches;
+
+    AssertQueryBuilder queryBuilder(duckDbQueryRunner_);
+    queryBuilder.plan(plan)
+        .config(
+            core::QueryConfig::kIndexLookupJoinMaxPrefetchBatches,
+            std::to_string(GetParam().numPrefetches))
+        .config(
+            core::QueryConfig::kPreferredOutputBatchRows,
+            std::to_string(maxBatchRows))
+        .config(
+            core::QueryConfig::kPreferredOutputBatchBytes,
+            std::to_string(1ULL << 30))
+        .config(
+            core::QueryConfig::kIndexLookupJoinSplitOutput,
+            testData.configSplitOutput ? "true" : "false")
+        .splits(probeScanNodeId_, makeHiveConnectorSplits(probeFiles))
+        .serialExecution(GetParam().serialExecution)
+        .barrierExecution(GetParam().serialExecution);
+    if (GetParam().needsIndexSplit) {
+      queryBuilder.split(
+          indexScanNodeId_,
+          Split(
+              std::make_shared<TestIndexConnectorSplit>(
+                  kTestIndexConnectorName)));
+    }
+    const auto task = queryBuilder.assertResults(
+        "SELECT t.c4, u.c5 FROM t, u WHERE t.c0 = u.c0 AND t.c1 = u.c1 AND t.c2 = u.c2");
+    ASSERT_EQ(
+        toPlanStats(task->taskStats()).at(joinNodeId).outputVectors,
+        expectedNumOutputVectors);
+  }
+}
+
 DEBUG_ONLY_TEST_P(IndexLookupJoinTest, statsSplitter) {
   IndexTableData tableData;
   generateIndexTableData({100, 1, 1}, tableData, pool_);

velox/exec/tests/utils/PlanBuilder.cpp

@@ -2725,6 +2725,7 @@ core::PlanNodePtr PlanBuilder::IndexLookupJoinBuilder::build(
       std::move(joinConditionPtrs),
       filterExpr,
       hasMarker_,
+      splitOutput_,
       std::move(planBuilder_.planNode_),
       indexSource_,
       std::move(outputType));