[CALCITE-7442] Getting Wrong index of Correlated variable inside Subquery after FilterJoinRule

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableAsofJoin.java

@@ -102,6 +102,14 @@
     throw new RuntimeException("This method should not be called");
   }
 
+  @Override public EnumerableAsofJoin copy(RelTraitSet traitSet, RexNode condition,
+                                           RelNode left, RelNode right, JoinRelType joinType,
+                                           boolean semiJoinDone,
+                                           Set<CorrelationId> variablesSet) {
+    // This method does not know about the matchCondition, so it should not be called
+    throw new RuntimeException("This method should not be called");
+  }
+
   @Override public Join copy(RelTraitSet traitSet, List<RelNode> inputs) {
     assert inputs.size() == 2;
     return new EnumerableAsofJoin(

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableBatchNestedLoopJoin.java

@@ -119,6 +119,13 @@ public static EnumerableBatchNestedLoopJoin create(
         left, right, condition, variablesSet, requiredColumns, joinType);
   }
 
+  @Override public EnumerableBatchNestedLoopJoin copy(RelTraitSet traitSet,
+      RexNode condition, RelNode left, RelNode right, JoinRelType joinType,
+      boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+    return new EnumerableBatchNestedLoopJoin(getCluster(), traitSet,
+        left, right, condition, variablesSet, requiredColumns, joinType);
+  }
+
   @Override public @Nullable RelOptCost computeSelfCost(
       final RelOptPlanner planner,
       final RelMetadataQuery mq) {

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableHashJoin.java

@@ -106,6 +106,13 @@ public static EnumerableHashJoin create(
         condition, variablesSet, joinType);
   }
 
+  @Override public EnumerableHashJoin copy(RelTraitSet traitSet, RexNode condition,
+      RelNode left, RelNode right, JoinRelType joinType,
+      boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+    return new EnumerableHashJoin(getCluster(), traitSet, left, right,
+        condition, variablesSet, joinType);
+  }
+
   @Override public @Nullable Pair<RelTraitSet, List<RelTraitSet>> passThroughTraits(
       final RelTraitSet required) {
     return EnumerableTraitsUtils.passThroughTraitsForJoin(

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeJoin.java

@@ -432,6 +432,13 @@ public static EnumerableMergeJoin create(RelNode left, RelNode right,
         condition, variablesSet, joinType);
   }
 
+  @Override public EnumerableMergeJoin copy(RelTraitSet traitSet,
+      RexNode condition, RelNode left, RelNode right, JoinRelType joinType,
+      boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+    return new EnumerableMergeJoin(getCluster(), traitSet, left, right,
+        condition, variablesSet, joinType);
+  }
+
   @Override public @Nullable RelOptCost computeSelfCost(RelOptPlanner planner,
       RelMetadataQuery mq) {
     // We assume that the inputs are sorted. The price of sorting them has

core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableNestedLoopJoin.java

@@ -70,6 +70,13 @@ protected EnumerableNestedLoopJoin(RelOptCluster cluster, RelTraitSet traits,
         condition, variablesSet, joinType);
   }
 
+  @Override public EnumerableNestedLoopJoin copy(RelTraitSet traitSet,
+      RexNode condition, RelNode left, RelNode right, JoinRelType joinType,
+      boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+    return new EnumerableNestedLoopJoin(getCluster(), traitSet, left, right,
+        condition, variablesSet, joinType);
+  }
+
   /** Creates an EnumerableNestedLoopJoin. */
   public static EnumerableNestedLoopJoin create(
       RelNode left,

core/src/main/java/org/apache/calcite/adapter/jdbc/JdbcRules.java

@@ -414,6 +414,19 @@ protected JdbcJoin(
       }
     }
 
+    @Override public JdbcJoin copy(RelTraitSet traitSet, RexNode condition,
+        RelNode left, RelNode right, JoinRelType joinType,
+        boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+      try {
+        return new JdbcJoin(getCluster(), traitSet, left, right,
+            condition, variablesSet, joinType);
+      } catch (InvalidRelException e) {
+        // Semantic error not possible. Must be a bug. Convert to
+        // internal error.
+        throw new AssertionError(e);
+      }
+    }
+
     @Override public @Nullable RelOptCost computeSelfCost(RelOptPlanner planner,
         RelMetadataQuery mq) {
       // We always "build" the

core/src/main/java/org/apache/calcite/interpreter/Bindables.java

@@ -556,6 +556,13 @@ protected BindableJoin(RelOptCluster cluster, RelTraitSet traitSet,
           conditionExpr, variablesSet, joinType);
     }
 
+    @Override public BindableJoin copy(RelTraitSet traitSet, RexNode conditionExpr,
+        RelNode left, RelNode right, JoinRelType joinType,
+        boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+      return new BindableJoin(getCluster(), traitSet, left, right,
+          conditionExpr, variablesSet, joinType);
+    }
+
     @Override public Class<Object[]> getElementType() {
       return Object[].class;
     }

core/src/main/java/org/apache/calcite/plan/RelOptUtil.java

@@ -2934,15 +2934,39 @@ public static boolean classifyFilters(
     ImmutableBitSet rightBitmap =
         ImmutableBitSet.range(nSysFields + nFieldsLeft, nTotalFields);
 
+    // Correlation variables introduced by this join itself: i.e. ids whose
+    // binding is established by *this* join. A predicate that references any
+    // of these cannot be pushed to either input -- the binder lives on the
+    // join, so pushing the reference below it would strand the variable.
+    // Such predicates must stay on the join itself.
+    final Set<CorrelationId> joinCorrelationIds = joinRel instanceof Join
+        ? joinRel.getVariablesSet()
+        : ImmutableSet.of();
+
     final List<RexNode> filtersToRemove = new ArrayList<>();
     for (RexNode filter : filters) {
-      final InputFinder inputFinder = InputFinder.analyze(filter);
+
+      // Only consider correlation ids bound by *this* join when computing
+      // the input bitmap of a sub-query inside the predicate. Foreign
+      // correlation ids are bound by an outer scope and their
+      // correlationColumns indices would otherwise alias onto unrelated
+      // columns of this join's row type, mis-classifying the predicate.
+      final InputFinder inputFinder = InputFinder.analyze(filter, joinCorrelationIds);
       final ImmutableBitSet inputBits = inputFinder.build();
 
+      // Block pushing to either input for filters that reference a
+      // CorrelationId bound by this join; they must remain on the join.
+      // pushing down correlated subqueries carries risks and involves extremely complex logic,
+      // and therefore pushing down is prohibited.
+      final boolean blockPush =
+          RexUtil.containsCorrelation(filter, joinCorrelationIds);
+      final boolean effectivePushLeft = pushLeft && !blockPush && leftBitmap.contains(inputBits);
+      final boolean effectivePushRight = pushRight && !blockPush && rightBitmap.contains(inputBits);
+
       // REVIEW - are there any expressions that need special handling
       // and therefore cannot be pushed?
 
-      if (pushLeft && leftBitmap.contains(inputBits)) {
+      if (effectivePushLeft) {
         // ignore filters that always evaluate to true
         if (!filter.isAlwaysTrue()) {
           // adjust the field references in the filter to reflect
@@ -2962,7 +2986,7 @@ public static boolean classifyFilters(
           leftFilters.add(shiftedFilter);
         }
         filtersToRemove.add(filter);
-      } else if (pushRight && rightBitmap.contains(inputBits)) {
+      } else if (effectivePushRight) {
         if (!filter.isAlwaysTrue()) {
           // adjust the field references in the filter to reflect
           // that fields in the right now shift over to the left
@@ -4678,12 +4702,29 @@ public RexCorrelVariableMapShuttle(final CorrelationId correlationId,
   public static class InputFinder extends RexVisitorImpl<Void> {
     private final ImmutableBitSet.Builder bitBuilder;
     private final @Nullable Set<RelDataTypeField> extraFields;
+    /** Correlation ids whose binder is the current scope. When non-null,
+     * {@link #visitSubQuery} projects bits for each id in this set by looking
+     * up its {@code correlationColumns} against the sub-query's inner plan
+     * and adding those column indices to the bitmap. Correlation ids bound
+     * by an outer scope are skipped, since their column indices are relative
+     * to a foreign row type and would otherwise alias onto unrelated columns
+     * of the current scope. When null, {@link #visitSubQuery} contributes no
+     * correlation-related bits and simply descends into the sub-query's
+     * operands (legacy behaviour). */
+    private final @Nullable Set<CorrelationId> localCorrelationIds;
 
     private InputFinder(@Nullable Set<RelDataTypeField> extraFields,
-        ImmutableBitSet.Builder bitBuilder) {
+        ImmutableBitSet.Builder bitBuilder,
+        @Nullable Set<CorrelationId> localCorrelationIds) {
       super(true);
       this.bitBuilder = bitBuilder;
       this.extraFields = extraFields;
+      this.localCorrelationIds = localCorrelationIds;
+    }
+
+    private InputFinder(@Nullable Set<RelDataTypeField> extraFields,
+        ImmutableBitSet.Builder bitBuilder) {
+      this(extraFields, bitBuilder, null);
     }
 
     public InputFinder() {
@@ -4706,6 +4747,22 @@ public static InputFinder analyze(RexNode node) {
       return inputFinder;
     }
 
+    /** Returns an input finder that has analyzed a given expression,
+     * treating {@code localCorrelationIds} as the set of correlation ids
+     * bound by the current scope. For each nested {@link RexSubQuery},
+     * any correlation id used inside it that belongs to this set
+     * contributes its {@code correlationColumns} indices to the bitmap;
+     * correlation ids bound by an outer scope are ignored, because their
+     * indices are relative to a foreign row type and would otherwise
+     * alias onto unrelated columns of the current scope. */
+    public static InputFinder analyze(RexNode node,
+        Set<CorrelationId> localCorrelationIds) {
+      final InputFinder inputFinder =
+          new InputFinder(null, ImmutableBitSet.builder(), localCorrelationIds);
+      node.accept(inputFinder);
+      return inputFinder;
+    }
+
     /**
      * Returns a bit set describing the inputs used by an expression.
      */
@@ -4752,6 +4809,28 @@ public ImmutableBitSet build() {
       }
       return super.visitCall(call);
     }
+
+    @Override public Void visitSubQuery(RexSubQuery subQuery) {
+      if (localCorrelationIds == null) {
+        return super.visitSubQuery(subQuery);
+      }
+
+      final Set<CorrelationId> variablesSet = RelOptUtil.getVariablesUsed(subQuery.rel);
+      for (CorrelationId id : variablesSet) {
+        // Skip correlation ids that are not bound by the *current* scope.
+        // Their requiredColumns indices are relative to whichever outer
+        // RelNode produces them and would otherwise alias onto unrelated
+        // columns of the current row type.
+        if (!localCorrelationIds.contains(id)) {
+          continue;
+        }
+        ImmutableBitSet requiredColumns = RelOptUtil.correlationColumns(id, subQuery.rel);
+        for (int index : requiredColumns) {
+          bitBuilder.set(index);
+        }
+      }
+      return super.visitSubQuery(subQuery);
+    }
   }
 
   /**

core/src/main/java/org/apache/calcite/rel/core/Join.java

@@ -355,6 +355,28 @@ public static RelDataType createJoinType(
   public abstract Join copy(RelTraitSet traitSet, RexNode conditionExpr,
       RelNode left, RelNode right, JoinRelType joinType, boolean semiJoinDone);
 
+  /**
+   * Creates a copy of this join, overriding condition, system fields and
+   * inputs.
+   *
+   * <p>General contract as {@link RelNode#copy}.
+   *
+   * @param traitSet      Traits
+   * @param conditionExpr Condition
+   * @param left          Left input
+   * @param right         Right input
+   * @param joinType      Join type
+   * @param semiJoinDone  Whether this join has been translated to a
+   *                      semi-join
+   * @param variablesSet  Set of variables that are set by the
+   *                      LHS and used by the RHS and are not available to
+   *                      nodes above this LogicalJoin in the tree
+   * @return Copy of this join
+   */
+  public abstract Join copy(RelTraitSet traitSet, RexNode conditionExpr,
+      RelNode left, RelNode right, JoinRelType joinType, boolean semiJoinDone,
+      Set<CorrelationId> variablesSet);
+
   /**
    * Analyzes the join condition.
    *

core/src/main/java/org/apache/calcite/rel/logical/LogicalAsofJoin.java

@@ -23,6 +23,7 @@
 import org.apache.calcite.rel.RelNode;
 import org.apache.calcite.rel.RelShuttle;
 import org.apache.calcite.rel.core.AsofJoin;
+import org.apache.calcite.rel.core.CorrelationId;
 import org.apache.calcite.rel.core.Join;
 import org.apache.calcite.rel.core.JoinRelType;
 import org.apache.calcite.rel.hint.RelHint;
@@ -37,6 +38,7 @@
 import java.util.ArrayList;
 import java.util.List;
 import java.util.Objects;
+import java.util.Set;
 
 import static java.util.Objects.requireNonNull;
 
@@ -144,6 +146,13 @@
     throw new RuntimeException("This method should not be called");
   }
 
+  @Override public Join copy(
+      RelTraitSet traitSet, RexNode conditionExpr, RelNode left, RelNode right,
+      JoinRelType joinType, boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+    // This method does not provide the matchCondition as an argument, so it should never be called
+    throw new RuntimeException("This method should not be called");
+  }
+
   @Override public Join copy(RelTraitSet traitSet, List<RelNode> inputs) {
     assert inputs.size() == 2;
     return new LogicalAsofJoin(getCluster(), traitSet, hints,

core/src/main/java/org/apache/calcite/rel/logical/LogicalJoin.java

@@ -181,6 +181,14 @@
         variablesSet, joinType, semiJoinDone, systemFieldList);
   }
 
+  @Override public LogicalJoin copy(RelTraitSet traitSet, RexNode conditionExpr, RelNode left,
+      RelNode right, JoinRelType joinType, boolean semiJoinDone, Set<CorrelationId> variablesSet) {
+    assert traitSet.containsIfApplicable(Convention.NONE);
+    return new LogicalJoin(getCluster(),
+        getCluster().traitSetOf(Convention.NONE), hints, left, right, conditionExpr,
+        variablesSet, joinType, semiJoinDone, systemFieldList);
+  }
+
   @Override public RelNode accept(RelShuttle shuttle) {
     return shuttle.visit(this);
   }

core/src/main/java/org/apache/calcite/rel/rules/CoreRules.java

@@ -284,7 +284,7 @@ private CoreRules() {}
    * {@link org.apache.calcite.rel.rules.FilterProjectTransposeRule}.
    *
    * <p>It does not allow a Filter to be pushed past the Project if
-   * {@link RexUtil#containsCorrelation there is a correlation condition}
+   * {@link RexUtil#containsCorrelation(org.apache.calcite.rex.RexNode) there is a correlation condition}
    * anywhere in the Filter, since in some cases it can prevent a
    * {@link Correlate} from being de-correlated.
    */