[CALCITE-7440] Preserve correlated scope mapping in RelToSqlConverter

Author: bvolpato

core/src/main/java/org/apache/calcite/rel/rel2sql/RelToSqlConverter.java

@@ -54,11 +54,14 @@
 import org.apache.calcite.rel.type.RelDataTypeField;
 import org.apache.calcite.rex.RexBuilder;
 import org.apache.calcite.rex.RexCall;
+import org.apache.calcite.rex.RexCorrelVariable;
+import org.apache.calcite.rex.RexFieldAccess;
 import org.apache.calcite.rex.RexInputRef;
 import org.apache.calcite.rex.RexLiteral;
 import org.apache.calcite.rex.RexLocalRef;
 import org.apache.calcite.rex.RexNode;
 import org.apache.calcite.rex.RexProgram;
+import org.apache.calcite.rex.RexSubQuery;
 import org.apache.calcite.sql.JoinConditionType;
 import org.apache.calcite.sql.JoinType;
 import org.apache.calcite.sql.SqlAsofJoin;
@@ -557,6 +560,7 @@ public Result visit(Filter e) {
           visitInput(e, 0, isAnon(), ignoreClauses,
               ImmutableSet.of(Clause.HAVING));
       parseCorrelTable(e, x);
+      parseCorrelVariables(e.getCondition());
       final Builder builder = x.builder(e);
       x.asSelect().setHaving(
           SqlUtil.andExpressions(x.asSelect().getHaving(),
@@ -569,6 +573,7 @@ public Result visit(Filter e) {
       }
       parseCorrelTable(e, x);
       final Builder builder = x.builder(e);
+      parseCorrelVariables(e.getCondition());
       if (input instanceof Join) {
         final Context context = x.qualifiedContext();
         if (selectListRequired(context)) {
@@ -604,6 +609,31 @@ private static boolean selectListRequired(Context context) {
     return false;
   }
 
+  private void parseCorrelVariables(RexNode rexNode) {
+    if (rexNode instanceof RexCorrelVariable) {
+      final RexCorrelVariable correl = (RexCorrelVariable) rexNode;
+      correlTableMap.putIfAbsent(
+          correl.id,
+          aliasContext(ImmutableMap.of(correl.id.getName(), correl.getType()), true));
+      return;
+    }
+    if (rexNode instanceof RexSubQuery) {
+      for (RexNode operand : ((RexSubQuery) rexNode).operands) {
+        parseCorrelVariables(operand);
+      }
+      return;
+    }
+    if (rexNode instanceof RexFieldAccess) {
+      parseCorrelVariables(((RexFieldAccess) rexNode).getReferenceExpr());
+      return;
+    }
+    if (rexNode instanceof RexCall) {
+      for (RexNode operand : ((RexCall) rexNode).operands) {
+        parseCorrelVariables(operand);
+      }
+    }
+  }
+
   /**
    * Extracts the table name from a SqlNode if it represents a simple table reference.
    * Returns null for complex nodes like subqueries or joins.
@@ -772,6 +802,7 @@ private Builder visitAggregate(Aggregate e, List<Integer> groupKeyList,
     // "select a, b, sum(x) from ( ... ) group by a, b"
     final boolean ignoreClauses = e.getInput() instanceof Project;
     final Result x = visitInput(e, 0, isAnon(), ignoreClauses, clauseSet);
+    parseCorrelTable(e, x);
     final Builder builder = x.builder(e);
     final List<SqlNode> selectList = new ArrayList<>();
     final List<SqlNode> groupByList =

core/src/test/java/org/apache/calcite/rel/rel2sql/RelToSqlConverterTest.java

@@ -11887,6 +11887,49 @@ public Sql schema(CalciteAssert.SchemaSpec schemaSpec) {
     sql(sql).schema(CalciteAssert.SchemaSpec.JDBC_SCOTT).ok(expected);
   }
 
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-7440">[CALCITE-7440]
+   * RelToSqlConverter throws NPE when correlation scope is missing after
+   * semi-join rewrites.</a>. */
+  @Test void testPostgresqlRoundTripCorrelatedProjectWithSemiJoinRules() {
+    final String query = "WITH product_keys AS (\n"
+        + "  SELECT p.\"product_id\",\n"
+        + "         (SELECT MAX(p3.\"product_id\")\n"
+        + "          FROM \"foodmart\".\"product\" p3\n"
+        + "          WHERE p3.\"product_id\" = p.\"product_id\") AS \"mx\"\n"
+        + "  FROM \"foodmart\".\"product\" p\n"
+        + ")\n"
+        + "SELECT DISTINCT pk.\"product_id\"\n"
+        + "FROM product_keys pk\n"
+        + "LEFT JOIN \"foodmart\".\"product\" p2 USING (\"product_id\")\n"
+        + "WHERE pk.\"product_id\" IN (\n"
+        + "  SELECT p4.\"product_id\"\n"
+        + "  FROM \"foodmart\".\"product\" p4\n"
+        + ")";
+
+    final RuleSet rules =
+        RuleSets.ofList(CoreRules.FILTER_SUB_QUERY_TO_CORRELATE,
+        CoreRules.PROJECT_SUB_QUERY_TO_CORRELATE,
+        CoreRules.JOIN_SUB_QUERY_TO_CORRELATE,
+        CoreRules.FILTER_SUB_QUERY_TO_MARK_CORRELATE,
+        CoreRules.PROJECT_SUB_QUERY_TO_MARK_CORRELATE,
+        CoreRules.MARK_TO_SEMI_OR_ANTI_JOIN_RULE,
+        CoreRules.PROJECT_TO_SEMI_JOIN,
+        CoreRules.JOIN_TO_SEMI_JOIN,
+        CoreRules.SEMI_JOIN_FILTER_TRANSPOSE,
+        CoreRules.SEMI_JOIN_JOIN_TRANSPOSE);
+
+    final String generated = sql(query).withPostgresql().optimize(rules, null).exec();
+    try {
+      sql(generated).withPostgresql().exec();
+    } catch (Exception e) {
+      throw new AssertionError(
+          "Generated SQL failed PostgreSQL round-trip validation:\n"
+              + generated,
+          e);
+    }
+  }
+
   @Test void testNotBetween() {
     Sql f = fixture().withConvertletTable(new SqlRexConvertletTable() {
       @Override public @Nullable SqlRexConvertlet get(SqlCall call) {