Add merge join support

Summary:
This diff implements merge join as a new physical join algorithm in the Axiom optimizer. Merge join is a highly efficient join method for pre-sorted inputs, avoiding the hash table construction overhead of hash joins. When both join inputs are sorted on the join keys (or can be cheaply sorted), merge join streams through both inputs in lockstep, matching rows with equal keys. The applicability is limited to tables which are already copartitioned and where the join keys contain all the bucketing keys. For Hive  partitioned   tables, both sides must have a single hive partition or must have all hive partitioning columns on both sides in the join keys.

Main changes:

1. **Added joinByMerge() method** (Optimization.cpp, ~300 lines):
   - Core merge join candidate generation logic
   - Validates merge join preconditions:
     - Checks all join keys are columns (not expressions) - merge join requires direct column comparisons
     - Verifies left input has both partitioning and ordering (from distribution)
     - Confirms ordering is ascending (merge join requires monotonic order)
     - Validates all partition columns are in join keys (ensures copartitioning)
   - Identifies merge columns from left input's orderKeys that match join keys
   - Plans right side with matching partition and ordering distributions
   - Adds shuffle and sort operators to right side if needed
   - Computes join using merge method
   - Returns NextJoin candidate for cost comparison

2. **Merge join precondition checking**:
   - **Column-only keys**: Rejects if any join key is an expression (e.g., `CAST(orderkey AS BIGINT)`)
   - **Partitioning requirement**: Left input must be partitioned (not gathered)
   - **Ordering requirement**: Left input must have orderKeys specified
   - **Ascending order**: Only kAscNullsFirst and kAscNullsLast supported
   - **Partition subset**: All partition columns must appear in join keys
   - **Matching merge columns**: At least one orderKey must match a join key

3. **Right side preparation algorithm**:
   - Constructs Distribution for right side matching left's partition and ordering
   - For each left partition column, finds corresponding right join key
   - For each left merge column (from orderKeys), finds corresponding right join key
   - Creates `forRight` distribution with:
     - Same DistributionType as left (enables copartitioning)
     - rightPartition: right keys corresponding to left partition columns
     - rightOrderKeys: right keys corresponding to left merge columns
     - rightOrderTypes: ascending order (kAscNullsLast) to match left
   - Calls `makePlan()` with forRight distribution
   - If `needsShuffle` is true, adds Repartition operator
   - Checks if right input needs sorting (orderKeys don't match expected)
   - Adds OrderBy operator if needed, after shuffle or directly

4. **Merge join cost model** (RelationOp.cpp):
   - Added `setMergeJoinCost()` method to Join class
   - Cost formula: `3 * kKeyCompareCost * numKeys * min(1, fanout) + rightSideBytes + kHashExtractColumnCost * numRightSideColumns`
   - Rationale:
     - Key comparisons: Merge join compares keys 3 times on average per match (binary search in merge)
     - Scales with number of keys and fanout (more comparisons for multiple matches)
     - Data copying: Transfers right side bytes to output
     - Column extraction: Extracts columns from right side vectors
   - Significantly cheaper than hash join for large inputs (no hash table construction)
   - Cost difference grows with build side size (hash table cost is O(n log n), merge is O(n))

5. **Integration with join planning** (Optimization.cpp):
   - Modified `makeJoins()` to call `joinByMerge()` after `joinByIndex()`
   - Added `testingUseMergeJoin` option for testing:
     - `std::nullopt` (default): Normal cost-based selection among all join types
     - `true`: Prefer merge join - return immediately if joinByMerge produces a candidate
     - `false`: Disable merge join - skip calling joinByMerge entirely
   - If testing mode is off, merge join competes with hash join based on cost
   - If testing mode is on and merge join produced a candidate, skip hash join consideration

6. **Schema changes for lookup keys** (Schema.h, Schema.cpp):
   - Added `lookupColumns` field to ColumnGroup
   - Distinguished from `orderKeys` in Distribution:
     - `lookupColumns`: Columns used for index lookups (prefix of sort order)
     - `orderKeys`: Full sort order (may include additional sorting columns)
- The key point is that sortedness does not in and of itself make a table lookup-compatible.

- Modified `addIndex()` to accept both `columns` and `lookupColumns`
   - Updated `indexLookupCardinality()` to use lookupColumns for cardinality estimation
   - Enables accurate modeling of sorted table access patterns
   - Example: Table sorted on (orderkey, linenum) can be efficiently joined on just orderkey

7. **Velox plan translation** (ToVelox.cpp, ToVelox.h):
   - Added `makeMergeJoin()` method to create MergeJoinNode
   - Checks `join.method == JoinMethod::kMerge` to dispatch to merge join creation
   - Creates `velox::core::MergeJoinNode` with:
     - Join type (INNER, LEFT, RIGHT, FULL, SEMI, ANTI)
     - Left and right keys as field references
     - Filter expression (for non-equi join conditions)
     - Left and right child plan nodes
     - Output type from join columns
   - Registers prediction and history for cost feedback
   - MergeJoinNode relies on Velox runtime's merge join operator

8. **Bucketed sorted table creation** (ParquetTpchTest.cpp):
   - Added `makeBucketedSortedTables()` utility method
   - Creates `orders_bs`, `lineitem_bs`, `partsupp_bs`, `part_bs` tables
   - Uses 32 buckets (more buckets than `_b` versions for finer parallelism)
   - Specifies `sorted_by` property in addition to `bucketed_by`
   - Example: `orders_bs` is bucketed on `o_orderkey` and sorted on `o_orderkey` within each bucket
   - Parquet files maintain sort order within partitions
   - Used for testing merge join on realistic data

9. **Plan matcher support** (PlanMatcher.cpp, PlanMatcherGenerator.cpp):
   - Added `mergeJoin()` method to PlanMatcherBuilder
   - Signature: `mergeJoin(matcher, joinType)` similar to `hashJoin()`
   - Enables test assertions like:
     ```cpp
     auto matcher = PlanMatcherBuilder()
       .tableScan("orders_bs")
       .mergeJoin(rightMatcher, JoinType::kInner)
       .build();
     ```
   - Added merge join code generation in PlanMatcherGenerator
   - Generates proper `.mergeJoin()` calls when plan contains MergeJoinNode

10. **Testing infrastructure** (OptimizerOptions.h):
    - Added `testingUseMergeJoin` optional flag
    - Three modes for comprehensive testing:
      - `nullopt`: Production mode - cost-based selection
      - `true`: Force merge join - tests merge join implementation in isolation
      - `false`: Disable merge join - tests that hash join fallback works
    - Enables differential testing: run same query with and without merge join

The merge join selection algorithm in joinByMerge():

```
1. Check preconditions:
   - All join keys are columns
   - Left input partitioned and ordered
   - Order is ascending
   - Partition columns ⊆ join keys

2. Extract merge columns:
   - For each left orderKey that matches a join key
   - Build leftMergeColumns vector

3. Plan right side:
   - Construct matching Distribution (partition + order)
   - Call makePlan() to get right input plan
   - Check if shuffle/sort needed via needsShuffle flag

4. Add shuffle/sort if needed:
   - If needsShuffle:
     - Add Repartition on rightPartition
     - Add OrderBy on rightOrderKeys
   - Else if ordering doesn't match:
     - Add OrderBy on rightOrderKeys

5. Create Join operator:
   - method = JoinMethod::kMerge
   - Compute cost using setMergeJoinCost()
   - Return as NextJoin candidate
```

Differential Revision: D89875337

Author: oerling

axiom/optimizer/Optimization.cpp

@@ -313,6 +313,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/Optimization.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;

axiom/optimizer/OptimizerOptions.h

@@ -403,14 +403,16 @@ Join::Join(
     float fanout,
     float innerFanout,
     ColumnVector columns,
-    PlanState& state)
+    PlanState& state,
+    ExprVector discreteJoinFilterExprs)
     : RelationOp{RelType::kJoin, std::move(lhs), std::move(columns)},
       method{method},
       joinType{joinType},
       right{std::move(rhs)},
       leftKeys{std::move(lhsKeys)},
       rightKeys{std::move(rhsKeys)},
-      filter{std::move(filterExprs)} {
+      filter{std::move(filterExprs)},
+      discreteJoinFilter{std::move(discreteJoinFilterExprs)} {
   cost_.inputCardinality = inputCardinality();
 
   // Determine optionality for each side
@@ -471,18 +473,24 @@ Join::Join(
     }
   }
 
-  const float buildSize = right->resultCardinality();
-  const auto numKeys = leftKeys.size();
-  const auto probeCost = Costs::hashTableCost(buildSize) +
-      // Multiply by min(fanout, 1) because most misses will not compare and if
-      // fanout > 1, there is still only one compare.
-      (Costs::kKeyCompareCost * numKeys * std::min<float>(1, cost_.fanout)) +
-      numKeys * Costs::kHashColumnCost;
+  // Compute join cost based on method
+  if (method == JoinMethod::kMerge) {
+    setMergeJoinCost();
+  } else {
+    // Hash join costing
+    const float buildSize = right->resultCardinality();
+    const auto numKeys = leftKeys.size();
+    const auto probeCost = Costs::hashTableCost(buildSize) +
+        // Multiply by min(fanout, 1) because most misses will not compare and
+        // if fanout > 1, there is still only one compare.
+        (Costs::kKeyCompareCost * numKeys * std::min<float>(1, cost_.fanout)) +
+        numKeys * Costs::kHashColumnCost;
 
-  const auto rowBytes = byteSize(right->columns());
-  const auto rowCost = Costs::hashRowCost(buildSize, rowBytes);
+    const auto rowBytes = byteSize(right->columns());
+    const auto rowCost = Costs::hashRowCost(buildSize, rowBytes);
 
-  cost_.unitCost = probeCost + cost_.fanout * rowCost;
+    cost_.unitCost = probeCost + cost_.fanout * rowCost;
+  }
 
   // Add constraints for non-key columns from the optional side of an outer join
   if (leftOptional || rightOptional) {
@@ -532,6 +540,23 @@ Join::Join(
   }
 }
 
+void Join::setMergeJoinCost() {
+  const auto numKeys = leftKeys.size();
+
+  // Get right side columns for byte size calculation
+  const auto rightSideColumns = right->columns();
+  const auto rightSideBytes = byteSize(rightSideColumns);
+  const auto numRightSideColumns = rightSideColumns.size();
+
+  // Merge join cost formula:
+  // 3 * key compare cost * number of keys * min(1, fanout) +
+  // byteSize(rightSideColumns) +
+  // kHashExtractColumnCost * numRightSideColumns
+  cost_.unitCost =
+      3 * Costs::kKeyCompareCost * numKeys * std::min<float>(1, cost_.fanout) +
+      rightSideBytes + Costs::kHashExtractColumnCost * numRightSideColumns;
+}
+
 namespace {
 std::pair<std::string, std::string> joinKeysString(
     const ExprVector& left,

axiom/optimizer/RelationOp.cpp

@@ -491,7 +491,8 @@ struct Join : public RelationOp {
       float fanout,
       float innerFanout, // The fanout if this were an inner join
       ColumnVector columns,
-      PlanState& state);
+      PlanState& state,
+      ExprVector discreteJoinFilter = {});
 
   static Join* makeCrossJoin(
       RelationOpPtr input,
@@ -505,6 +506,7 @@ struct Join : public RelationOp {
   const ExprVector leftKeys;
   const ExprVector rightKeys;
   const ExprVector filter;
+  const ExprVector discreteJoinFilter;
 
   const QGString& historyKey() const override;
 
@@ -513,6 +515,9 @@ struct Join : public RelationOp {
   void accept(
       const RelationOpVisitor& visitor,
       RelationOpVisitorContext& context) const override;
+
+ private:
+  void setMergeJoinCost();
 };
 
 using JoinCP = const Join*;

axiom/optimizer/RelationOp.h

@@ -101,10 +101,15 @@ std::string Value::toString() const {
 ColumnGroupCP SchemaTable::addIndex(
     const connector::TableLayout& layout,
     Distribution distribution,
-    ColumnVector columns) {
+    ColumnVector columns,
+    ColumnVector lookupColumns) {
   return columnGroups.emplace_back(
       make<ColumnGroup>(
-          *this, layout, std::move(distribution), std::move(columns)));
+          *this,
+          layout,
+          std::move(distribution),
+          std::move(columns),
+          std::move(lookupColumns)));
 }
 
 ColumnCP SchemaTable::findColumn(Name name) const {
@@ -187,7 +192,15 @@ SchemaTableCP Schema::findTable(
 
     ColumnVector columns;
     appendColumns(layout->columns(), columns);
-    schemaTable->addIndex(*layout, std::move(distribution), std::move(columns));
+
+    ColumnVector lookupColumns;
+    appendColumns(layout->lookupKeys(), lookupColumns);
+
+    schemaTable->addIndex(
+        *layout,
+        std::move(distribution),
+        std::move(columns),
+        std::move(lookupColumns));
   }
   table = {std::move(connectorTable), schemaTable};
   return schemaTable;
@@ -281,26 +294,27 @@ IndexInfo SchemaTable::indexInfo(
 
   const auto& distribution = index->distribution;
 
-  const auto numSorting = distribution.orderTypes.size();
+  const auto numLookupKeys = index->lookupColumns.size();
   const auto numUnique = distribution.numKeysUnique;
 
   PlanObjectSet covered;
-  for (auto i = 0; i < numSorting || i < numUnique; ++i) {
-    auto orderKey = distribution.orderKeys[i];
-    auto part = findColumnByName(columnsSpan, orderKey->as<Column>()->name());
+  for (auto i = 0; i < numLookupKeys || i < numUnique; ++i) {
+    ExprCP lookupKey =
+        i < numLookupKeys ? index->lookupColumns[i] : distribution.orderKeys[i];
+    auto part = findColumnByName(columnsSpan, lookupKey->as<Column>()->name());
     if (!part) {
       break;
     }
 
     covered.add(part);
-    if (i < numSorting) {
+    if (i < numLookupKeys) {
       info.scanCardinality =
-          combine(info.scanCardinality, i, orderKey->value().cardinality);
+          combine(info.scanCardinality, i, lookupKey->value().cardinality);
       info.lookupKeys.push_back(part);
       info.joinCardinality = info.scanCardinality;
     } else {
       info.joinCardinality =
-          combine(info.joinCardinality, i, orderKey->value().cardinality);
+          combine(info.joinCardinality, i, lookupKey->value().cardinality);
     }
     if (i == numUnique - 1) {
       info.unique = true;

axiom/optimizer/Schema.cpp

@@ -231,16 +231,19 @@ struct ColumnGroup {
       const SchemaTable& table,
       const connector::TableLayout& layout,
       Distribution distribution,
-      ColumnVector columns)
+      ColumnVector columns,
+      ColumnVector lookupColumns)
       : table{&table},
         layout{&layout},
         distribution{std::move(distribution)},
-        columns{std::move(columns)} {}
+        columns{std::move(columns)},
+        lookupColumns{std::move(lookupColumns)} {}
 
   SchemaTableCP table;
   const connector::TableLayout* layout;
   const Distribution distribution;
   const ColumnVector columns;
+  const ColumnVector lookupColumns;
 
   /// Returns cost of next lookup when the hit is within 'range' rows
   /// of the previous hit. If lookups are not batched or not ordered,
@@ -302,7 +305,8 @@ struct SchemaTable {
   ColumnGroupCP addIndex(
       const connector::TableLayout& layout,
       Distribution distribution,
-      ColumnVector columns);
+      ColumnVector columns,
+      ColumnVector lookupColumns);
 
   ColumnCP findColumn(Name name) const;
 

axiom/optimizer/Schema.h

@@ -1112,6 +1112,10 @@ velox::core::PlanNodePtr ToVelox::makeJoin(
         nextId(), toAnd(join.filter), joinNode);
   }
 
+  if (join.method == JoinMethod::kMerge) {
+    return makeMergeJoin(join, fragment, stages, left, right);
+  }
+
   auto leftKeys = toFieldRefs(join.leftKeys);
   auto rightKeys = toFieldRefs(join.rightKeys);
 
@@ -1130,6 +1134,29 @@ velox::core::PlanNodePtr ToVelox::makeJoin(
   return joinNode;
 }
 
+velox::core::PlanNodePtr ToVelox::makeMergeJoin(
+    const Join& join,
+    runner::ExecutableFragment& fragment,
+    std::vector<runner::ExecutableFragment>& stages,
+    velox::core::PlanNodePtr left,
+    velox::core::PlanNodePtr right) {
+  auto leftKeys = toFieldRefs(join.leftKeys);
+  auto rightKeys = toFieldRefs(join.rightKeys);
+
+  auto joinNode = std::make_shared<velox::core::MergeJoinNode>(
+      nextId(),
+      join.joinType,
+      leftKeys,
+      rightKeys,
+      toAnd(join.filter),
+      left,
+      right,
+      makeOutputType(join.columns()));
+
+  makePredictionAndHistory(joinNode->id(), &join);
+  return joinNode;
+}
+
 velox::core::PlanNodePtr ToVelox::makeUnnest(
     const Unnest& op,
     runner::ExecutableFragment& fragment,

axiom/optimizer/ToVelox.cpp

@@ -179,6 +179,13 @@ class ToVelox {
       runner::ExecutableFragment& fragment,
       std::vector<runner::ExecutableFragment>& stages);
 
+  velox::core::PlanNodePtr makeMergeJoin(
+      const Join& join,
+      runner::ExecutableFragment& fragment,
+      std::vector<runner::ExecutableFragment>& stages,
+      velox::core::PlanNodePtr left,
+      velox::core::PlanNodePtr right);
+
   velox::core::PlanNodePtr makeRepartition(
       const Repartition& repartition,
       runner::ExecutableFragment& fragment,

axiom/optimizer/ToVelox.h

@@ -120,6 +120,7 @@ add_executable(
   HiveLimitQueriesTest.cpp
   HiveQueriesTest.cpp
   JoinTest.cpp
+  OrderedOpsTest.cpp
   ParquetTpchTest.cpp
   PrecomputeProjectionTest.cpp
   PlanTest.cpp