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

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/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.
    */

core/src/main/java/org/apache/calcite/rex/RexUtil.java

@@ -2412,6 +2412,23 @@ public static boolean containsCorrelation(RexNode condition) {
     }
   }
 
+  /** Returns whether an expression references a {@link RexCorrelVariable}
+   * whose id is in {@code ids}, either directly (typically through a
+   * {@link RexFieldAccess}) or transitively via the inner plan of a
+   * {@link RexSubQuery}. */
+  public static boolean containsCorrelation(RexNode condition,
+      Set<CorrelationId> ids) {
+    if (ids.isEmpty()) {
+      return false;
+    }
+    try {
+      condition.accept(new CorrelationFinder(ids));
+      return false;
+    } catch (Util.FoundOne e) {
+      return true;
+    }
+  }
+
   /**
    * Given an expression, it will swap the table references contained in its
    * {@link RexTableInputRef} using the contents in the map.
@@ -3116,14 +3133,28 @@ private static class RexShiftShuttle extends RexShuttle {
   /** Visitor that throws {@link org.apache.calcite.util.Util.FoundOne} if
    * applied to an expression that contains a {@link RexCorrelVariable}. */
   private static class CorrelationFinder extends RexVisitorImpl<Void> {
-    static final CorrelationFinder INSTANCE = new CorrelationFinder();
+    static final CorrelationFinder INSTANCE = new CorrelationFinder(null);
 
-    private CorrelationFinder() {
+    /** Optional filter: when non-null, only correlation ids in this set
+     * trigger a match; when null, every correlation id matches. */
+    private final @Nullable Set<CorrelationId> ids;
+
+    /**
+     * Creates a CorrelationFinder.
+     *
+     * @param ids correlation ids to look for; pass {@code null} to match any
+     *            {@link RexCorrelVariable} regardless of its id
+     */
+    private CorrelationFinder(@Nullable Set<CorrelationId> ids) {
       super(true);
+      this.ids = ids;
     }
 
     @Override public Void visitCorrelVariable(RexCorrelVariable var) {
-      throw Util.FoundOne.NULL;
+      if (ids == null || ids.contains(var.id)) {
+        throw Util.FoundOne.NULL;
+      }
+      return null;
     }
 
     @Override public Void visitSubQuery(RexSubQuery subQuery) {
@@ -3135,8 +3166,16 @@ private CorrelationFinder() {
         operand.accept(this);
       }
 
-      if (!RelOptUtil.getVariablesUsed(subQuery.rel).isEmpty()) {
-        throw Util.FoundOne.NULL;
+      Set<CorrelationId> used = RelOptUtil.getVariablesUsed(subQuery.rel);
+      if (!used.isEmpty()) {
+        if (ids == null) {
+          throw Util.FoundOne.NULL;
+        }
+        for (CorrelationId id : used) {
+          if (ids.contains(id)) {
+            throw Util.FoundOne.NULL;
+          }
+        }
       }
       return null;
     }