[CALCITE-6744] RelMetadataQuery.getColumnOrigins should return null w…

…hen column origin includes correlation variables
apache · Jan 17, 2025 · 891ffc4 · 891ffc4
1 parent b1c09c6
commit 891ffc4
Show file tree

Hide file tree

Showing 2 changed files with 53 additions and 4 deletions.
diff --git a/core/src/main/java/org/apache/calcite/rel/metadata/RelMdColumnOrigins.java b/core/src/main/java/org/apache/calcite/rel/metadata/RelMdColumnOrigins.java
@@ -32,12 +32,15 @@
 import org.apache.calcite.rel.core.TableFunctionScan;
 import org.apache.calcite.rel.core.TableModify;
 import org.apache.calcite.rel.core.TableScan;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexFieldAccess;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLocalRef;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexShuttle;
 import org.apache.calcite.rex.RexVisitor;
 import org.apache.calcite.rex.RexVisitorImpl;
+import org.apache.calcite.util.Util;
 
 import org.checkerframework.checker.nullness.qual.Nullable;
 import org.checkerframework.checker.nullness.qual.PolyNull;
@@ -84,9 +87,10 @@ private RelMdColumnOrigins() {}
       Set<RelColumnOrigin> inputSet =
           mq.getColumnOrigins(rel.getInput(), iInput);
       inputSet = createDerivedColumnOrigins(inputSet);
-      if (inputSet != null) {
-        set.addAll(inputSet);
+      if (inputSet == null) {
+        return null;
       }
+      set.addAll(inputSet);
     }
     return set;
   }
@@ -290,7 +294,7 @@ private RelMdColumnOrigins() {}
     return set;
   }
 
-  private static Set<RelColumnOrigin> getMultipleColumns(RexNode rexNode, RelNode input,
+  private static @Nullable Set<RelColumnOrigin> getMultipleColumns(RexNode rexNode, RelNode input,
       final RelMetadataQuery mq) {
     final Set<RelColumnOrigin> set = new HashSet<>();
     final RexVisitor<Void> visitor =
@@ -303,8 +307,19 @@ private static Set<RelColumnOrigin> getMultipleColumns(RexNode rexNode, RelNode
             }
             return null;
           }
+
+          @Override public Void visitFieldAccess(RexFieldAccess fieldAccess) {
+            if (fieldAccess.getReferenceExpr() instanceof RexCorrelVariable) {
+              throw Util.FoundOne.NULL;
+            }
+            return null;
+          }
         };
-    rexNode.accept(visitor);
+    try {
+      rexNode.accept(visitor);
+    } catch (Util.FoundOne one) {
+      return null;
+    }
     return set;
   }
 }
diff --git a/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java b/core/src/test/java/org/apache/calcite/test/RelMetadataTest.java
@@ -143,6 +143,7 @@
 import static com.google.common.collect.ImmutableList.toImmutableList;
 
 import static org.apache.calcite.test.Matchers.hasFieldNames;
+import static org.apache.calcite.test.Matchers.hasTree;
 import static org.apache.calcite.test.Matchers.isAlmost;
 import static org.apache.calcite.test.Matchers.sortsAs;
 
@@ -458,6 +459,39 @@ void testColumnOriginsUnion() {
         equalTo("SAL"));
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-6744">[CALCITE-6744]
+   * RelMetadataQuery.getColumnOrigins should return null when column origin
+   * includes correlation variables</a>. */
+  @Test void testColumnOriginsForCorrelate() {
+    final String sql = "select (select max(dept.name || '_' || emp.ename)"
+        + "from dept where emp.deptno = dept.deptno) from emp";
+    final RelMetadataFixture fixture = sql(sql);
+
+    final HepProgramBuilder programBuilder = HepProgram.builder();
+    programBuilder.addRuleInstance(CoreRules.PROJECT_SUB_QUERY_TO_CORRELATE);
+    final HepPlanner planner = new HepPlanner(programBuilder.build());
+    planner.setRoot(fixture.toRel());
+    final RelNode relNode = planner.findBestExp();
+
+    String expect = "LogicalProject(EXPR$0=[$9])\n"
+        + "  LogicalCorrelate(correlation=[$cor1], joinType=[left], requiredColumns=[{1, 7}])\n"
+        + "    LogicalTableScan(table=[[CATALOG, SALES, EMP]])\n"
+        + "    LogicalAggregate(group=[{}], EXPR$0=[MAX($0)])\n"
+        + "      LogicalProject($f0=[||(||($1, '_'), $cor1.ENAME)])\n"
+        + "        LogicalFilter(condition=[=($cor1.DEPTNO, $0)])\n"
+        + "          LogicalTableScan(table=[[CATALOG, SALES, DEPT]])\n";
+    assertThat(relNode, hasTree(expect));
+
+    // check correlate input column origins
+    final RelMetadataFixture.MetadataConfig metadataConfig = fixture.metadataConfig;
+    final RelMetadataQuery mq =
+        new RelMetadataQuery(metadataConfig.getDefaultHandlerProvider());
+    Aggregate aggregate = (Aggregate) relNode.getInput(0).getInput(1);
+    Set<RelColumnOrigin> origins = mq.getColumnOrigins(aggregate, 0);
+    assertNull(origins);
+  }
+
   // ----------------------------------------------------------------------
   // Tests for getRowCount, getMinRowCount, getMaxRowCount
   // ----------------------------------------------------------------------