feature: Add merge join support

Differential Revision: D89875337

Author: oerling

axiom/optimizer/Optimization.cpp

@@ -1674,6 +1674,302 @@ void Optimization::joinByHashRight(
   state.addNextJoin(&candidate, join, toTry);
 }
 
+void Optimization::joinByMerge(
+    const RelationOpPtr& plan,
+    const JoinCandidate& candidate,
+    PlanState& state,
+    std::vector<NextJoin>& toTry) {
+  checkTables(candidate);
+
+  // Merge join requires: [right, left] assignment
+  auto [right, left] = candidate.joinSides();
+
+  // Check if all join keys are columns (not expressions)
+  for (auto* key : left.keys) {
+    if (!key->isColumn()) {
+      return;  // Cannot do merge join with non-column keys
+    }
+  }
+  for (auto* key : right.keys) {
+    if (!key->isColumn()) {
+      return;  // Cannot do merge join with non-column keys
+    }
+  }
+
+  // Check if left side (plan) has partitioning and ordering
+  const auto& distribution = plan->distribution();
+  if (distribution.partition.empty() || distribution.orderKeys.empty()) {
+    return;  // Need both partitioning and ordering for merge join
+  }
+
+  // Check that all ordering has ascending order
+  for (const auto& orderType : distribution.orderTypes) {
+    if (orderType != OrderType::kAscNullsFirst &&
+        orderType != OrderType::kAscNullsLast) {
+      return;  // Merge join requires ascending order
+    }
+  }
+
+  // Check if all left partition columns are in left.keys
+  ExprVector leftKeyExprs;
+  for (auto* col : left.keys) {
+    leftKeyExprs.push_back(col);
+  }
+
+  PlanObjectSet leftKeySet;
+  leftKeySet.unionObjects(leftKeyExprs);
+
+  PlanObjectSet partitionSet;
+  partitionSet.unionObjects(distribution.partition);
+
+  if (!partitionSet.isSubset(leftKeySet)) {
+    return;  // Not all partition columns are in join keys
+  }
+
+  // Find the left keys that correspond to ordering columns
+  ColumnVector leftMergeColumns;
+  for (const auto& orderKey : distribution.orderKeys) {
+    // Find matching column in left.keys
+    ColumnCP matchingKey = nullptr;
+    for (size_t i = 0; i < left.keys.size(); ++i) {
+      if (left.keys[i] == orderKey) {
+        matchingKey = dynamic_cast<ColumnCP>(left.keys[i]);
+        break;
+      }
+    }
+
+    if (!matchingKey) {
+      break;  // Stop looking if we don't find a match
+    }
+
+    leftMergeColumns.push_back(matchingKey);
+  }
+
+  if (leftMergeColumns.empty()) {
+    return;  // No ordering columns match join keys
+  }
+
+  // At this point, left input (plan) is a candidate for merge join
+  PlanStateSaver save(state, candidate);
+
+  // Make the right side plan similar to hash join build side
+  PlanObjectSet rightFilterColumns;
+  rightFilterColumns.unionColumns(candidate.join->filter());
+  rightFilterColumns.intersect(availableColumns(candidate.tables[0]));
+
+  PlanObjectSet rightTables;
+  PlanObjectSet rightColumns;
+  for (auto rightTable : candidate.tables) {
+    rightColumns.unionSet(availableColumns(rightTable));
+    rightTables.add(rightTable);
+  }
+
+  auto joinColumnMapping = makeJoinColumnMapping(candidate.join);
+
+  state.placed.unionSet(rightTables);
+  rightColumns.intersect(
+      translateToJoinInput(state.downstreamColumns(), joinColumnMapping));
+
+  rightColumns.unionColumns(right.keys);
+  rightColumns.unionSet(rightFilterColumns);
+  state.columns.unionSet(rightColumns);
+
+  MemoKey memoKey{
+      candidate.tables[0], rightColumns, rightTables, candidate.existences};
+
+  // Distribution for right side matching left partition and ordering
+  ExprVector rightPartition;
+  for (const auto& leftPartCol : distribution.partition) {
+    // Find corresponding right key
+    for (size_t i = 0; i < left.keys.size(); ++i) {
+      if (left.keys[i] == leftPartCol) {
+        rightPartition.push_back(right.keys[i]);
+        break;
+      }
+    }
+  }
+
+  ExprVector rightOrderKeys;
+  OrderTypeVector rightOrderTypes;
+  for (const auto* leftMergeCol : leftMergeColumns) {
+    // Find corresponding right key
+    for (size_t i = 0; i < left.keys.size(); ++i) {
+      if (left.keys[i] == leftMergeCol) {
+        rightOrderKeys.push_back(right.keys[i]);
+        // Use ascending order to match left side (default to kAscNullsLast)
+        rightOrderTypes.push_back(OrderType::kAscNullsLast);
+        break;
+      }
+    }
+  }
+
+  Distribution forRight{
+      distribution.distributionType,
+      rightPartition,
+      rightOrderKeys,
+      rightOrderTypes};
+
+  PlanObjectSet empty;
+  bool needsShuffle = false;
+  auto rightPlan = makePlan(
+      *state.dt,
+      memoKey,
+      forRight,
+      empty,
+      candidate.existsFanout,
+      needsShuffle);
+
+  PlanState rightState(state.optimization, state.dt, rightPlan);
+  RelationOpPtr rightInput = rightPlan->op;
+  RelationOpPtr probeInput = plan;
+
+  // Handle shuffle and sort for right side
+  if (!isSingleWorker_) {
+    if (needsShuffle) {
+      // Add shuffle
+      Distribution distribution{
+          plan->distribution().distributionType, rightPartition};
+      auto* repartition = make<Repartition>(
+          rightInput, std::move(distribution), rightInput->columns());
+      rightState.addCost(*repartition);
+      rightInput = repartition;
+
+      // Add sort after shuffle
+      PrecomputeProjection precomputeSort(rightInput, state.dt);
+      auto orderKeys = precomputeSort.toColumns(rightOrderKeys);
+      rightInput = std::move(precomputeSort).maybeProject(rightState);
+
+      auto* orderBy = make<OrderBy>(
+          rightInput,
+          std::move(orderKeys),
+          rightOrderTypes,
+          /*limit=*/-1,
+          /*offset=*/0);
+      rightState.addCost(*orderBy);
+      rightInput = orderBy;
+    } else {
+      // Check if ordering columns match expected
+      const auto& rightDist = rightInput->distribution();
+      bool needsSort = rightDist.orderKeys.size() != rightOrderKeys.size();
+      if (!needsSort) {
+        for (size_t i = 0; i < rightOrderKeys.size(); ++i) {
+          if (rightDist.orderKeys[i] != rightOrderKeys[i]) {
+            needsSort = true;
+            break;
+          }
+        }
+      }
+
+      if (needsSort) {
+        // Add sort without shuffle
+        PrecomputeProjection precomputeSort(rightInput, state.dt);
+        auto orderKeys = precomputeSort.toColumns(rightOrderKeys);
+        rightInput = std::move(precomputeSort).maybeProject(rightState);
+
+        auto* orderBy = make<OrderBy>(
+            rightInput,
+            std::move(orderKeys),
+            rightOrderTypes,
+            /*limit=*/-1,
+            /*offset=*/0);
+        rightState.addCost(*orderBy);
+        rightInput = orderBy;
+      }
+    }
+  }
+
+  PrecomputeProjection precomputeRight(rightInput, state.dt);
+  auto rightKeys = precomputeRight.toColumns(right.keys);
+  rightInput = std::move(precomputeRight).maybeProject(rightState);
+
+  auto joinType = right.leftJoinType();
+  const bool probeOnly = joinType == velox::core::JoinType::kLeftSemiFilter ||
+      joinType == velox::core::JoinType::kLeftSemiProject ||
+      joinType == velox::core::JoinType::kAnti;
+
+  PlanObjectSet probeColumns;
+  probeColumns.unionObjects(plan->columns());
+
+  ColumnCP mark = nullptr;
+
+  auto* joinEdge = candidate.join;
+
+  PlanObjectSet joinColumns;
+  joinColumns.unionObjects(joinEdge->leftColumns());
+  joinColumns.unionObjects(joinEdge->rightColumns());
+
+  ProjectionBuilder projectionBuilder;
+  bool needsProjection = false;
+
+  state.downstreamColumns().forEach<Column>([&](auto column) {
+    if (column == right.markColumn) {
+      mark = column;
+      return;
+    }
+
+    if (joinColumns.contains(column)) {
+      projectionBuilder.add(column, joinColumnMapping.at(column));
+      needsProjection = true;
+      return;
+    }
+
+    if ((probeOnly || !rightColumns.contains(column)) &&
+        !probeColumns.contains(column)) {
+      return;
+    }
+
+    projectionBuilder.add(column, column);
+  });
+
+  if (mark) {
+    setMarkTrueFraction(
+        mark, joinType, candidate.fanout, candidate.join->rlFanout());
+  }
+
+  tryOptimizeSemiProject(joinType, mark, state, negation_);
+
+  if (mark) {
+    projectionBuilder.add(mark, mark);
+  }
+
+  state.columns = projectionBuilder.outputColumns();
+
+  const auto fanout = fanoutJoinTypeLimit(
+      joinType,
+      candidate.fanout,
+      candidate.join->rlFanout(),
+      rightState.cost.cardinality / state.cost.cardinality);
+
+  PrecomputeProjection precomputeProbe(probeInput, state.dt);
+  auto probeKeys = precomputeProbe.toColumns(left.keys);
+  probeInput = std::move(precomputeProbe).maybeProject(state);
+
+  const auto innerFanout = candidate.fanout;
+
+  // Create merge join
+  RelationOp* join = make<Join>(
+      JoinMethod::kMerge,
+      joinType,
+      probeInput,
+      rightInput,
+      std::move(probeKeys),
+      std::move(rightKeys),
+      candidate.join->filter(),
+      fanout,
+      innerFanout,
+      projectionBuilder.inputColumns(),
+      state);
+
+  state.addCost(*join);
+  state.cost.cost += rightState.cost.cost;
+
+  if (needsProjection) {
+    join = projectionBuilder.build(join, state);
+  }
+
+  state.addNextJoin(&candidate, join, toTry);
+}
+
 void Optimization::crossJoin(
     const RelationOpPtr& plan,
     const JoinCandidate& candidate,
@@ -1746,6 +2042,21 @@ void Optimization::addJoin(
   std::vector<NextJoin> toTry;
   joinByIndex(plan, candidate, state, toTry);
 
+  // For testing: if testingUseMergeJoin is false, skip merge join entirely.
+  if (!options_.testingUseMergeJoin.has_value() ||
+      options_.testingUseMergeJoin.value()) {
+    const auto sizeBeforeMerge = toTry.size();
+    joinByMerge(plan, candidate, state, toTry);
+
+    // For testing: if testingUseMergeJoin is true and joinByMerge produced a
+    // result, return immediately without trying other join methods.
+    if (options_.testingUseMergeJoin.has_value() &&
+        options_.testingUseMergeJoin.value() && toTry.size() > sizeBeforeMerge) {
+      result.insert(result.end(), toTry.begin(), toTry.end());
+      return;
+    }
+  }
+
   const auto sizeAfterIndex = toTry.size();
   joinByHash(plan, candidate, state, toTry);
 

axiom/optimizer/Optimization.h

@@ -314,6 +314,16 @@ class Optimization {
       PlanState& state,
       std::vector<NextJoin>& toTry);
 
+  // Adds 'candidate' on top of 'plan' as a merge join. Checks if the left
+  // input (plan) is partitioned and ordered, and if join keys match the
+  // ordering. Prepares the right side with appropriate partitioning and
+  // ordering, adding shuffle and sort operators as needed.
+  void joinByMerge(
+      const RelationOpPtr& plan,
+      const JoinCandidate& candidate,
+      PlanState& state,
+      std::vector<NextJoin>& toTry);
+
   void crossJoin(
       const RelationOpPtr& plan,
       const JoinCandidate& candidate,

axiom/optimizer/OptimizerOptions.h

@@ -75,6 +75,13 @@ struct OptimizerOptions {
   /// partial + final or not.
   bool alwaysPlanPartialAggregation = false;
 
+  /// For testing: control merge join behavior.
+  /// - std::nullopt (default): normal cost-based selection among all join types
+  /// - true: prefer merge joins - return immediately if joinByMerge produces a
+  ///   candidate
+  /// - false: disable merge joins - skip calling joinByMerge
+  std::optional<bool> testingUseMergeJoin{std::nullopt};
+
   bool isMapAsStruct(std::string_view table, std::string_view column) const {
     if (allMapsAsStruct) {
       return true;