[CALCITE-7439] Qualify GROUP BY keys for DISTINCT over joins

Author: bvolpato

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

@@ -858,12 +858,21 @@ private List<SqlNode> generateGroupList(Builder builder,
         + aggregate.getGroupSet() + ", just possibly a different order";
 
     final List<SqlNode> groupKeys = new ArrayList<>();
+    final Join aggregateJoinInput =
+        aggregate.getInput() instanceof Join ? (Join) aggregate.getInput() : null;
+    final SqlJoin fromJoin =
+        builder.select.getFrom() instanceof SqlJoin ? (SqlJoin) builder.select.getFrom() : null;
+    final int leftFieldCount = aggregateJoinInput == null
+        ? -1
+        : aggregateJoinInput.getLeft().getRowType().getFieldCount();
     for (int key : groupList) {
-      final SqlNode field = builder.context.field(key);
+      SqlNode field = builder.context.field(key);
+      field = maybeQualifyJoinKey(field, key, fromJoin, leftFieldCount);
       groupKeys.add(field);
     }
     for (int key : sortedGroupList) {
-      final SqlNode field = builder.context.field(key);
+      SqlNode field =
+          maybeQualifyJoinKey(builder.context.field(key), key, fromJoin, leftFieldCount);
       addSelect(selectList, field, aggregate.getRowType());
     }
     switch (aggregate.getGroupType()) {
@@ -905,6 +914,31 @@ private List<SqlNode> generateGroupList(Builder builder,
     }
   }
 
+  private SqlNode maybeQualifyJoinKey(SqlNode field, int key,
+      @Nullable SqlJoin fromJoin, int leftFieldCount) {
+    if (!(field instanceof SqlIdentifier)
+        || ((SqlIdentifier) field).names.size() != 1
+        || fromJoin == null) {
+      return field;
+    }
+
+    final SqlNode side;
+    if (leftFieldCount < 0) {
+      if (key != 0) {
+        return field;
+      }
+      side = fromJoin.getLeft();
+    } else {
+      side = key < leftFieldCount ? fromJoin.getLeft() : fromJoin.getRight();
+    }
+    final String sideAlias = SqlValidatorUtil.alias(side);
+    if (sideAlias == null) {
+      return field;
+    }
+
+    return new SqlIdentifier(ImmutableList.of(sideAlias, ((SqlIdentifier) field).getSimple()), POS);
+  }
+
   private static SqlNode groupItem(List<SqlNode> groupKeys,
       ImmutableBitSet groupSet, ImmutableBitSet wholeGroupSet) {
     final List<SqlNode> nodes = groupSet.asList().stream()

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

@@ -127,7 +127,9 @@
 import static org.hamcrest.CoreMatchers.is;
 import static org.hamcrest.CoreMatchers.notNullValue;
 import static org.hamcrest.MatcherAssert.assertThat;
+import static org.hamcrest.Matchers.containsString;
 import static org.hamcrest.Matchers.hasToString;
+import static org.hamcrest.Matchers.not;
 import static org.junit.jupiter.api.Assertions.assertFalse;
 import static org.junit.jupiter.api.Assertions.assertTrue;
 
@@ -11887,6 +11889,109 @@ 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-7439">[CALCITE-7439]
+   * RelToSqlConverter emits ambiguous GROUP BY after LEFT JOIN USING with
+   * semi-join rewrite.</a>. */
+  @Test void testPostgresqlRoundTripDistinctLeftJoinInSubqueryWithSemiJoinRules() {
+    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.PROJECT_SUB_QUERY_TO_CORRELATE,
+            CoreRules.FILTER_SUB_QUERY_TO_CORRELATE,
+            CoreRules.JOIN_SUB_QUERY_TO_CORRELATE,
+            CoreRules.PROJECT_SUB_QUERY_TO_MARK_CORRELATE,
+            CoreRules.FILTER_SUB_QUERY_TO_MARK_CORRELATE,
+            CoreRules.MARK_TO_SEMI_OR_ANTI_JOIN_RULE,
+            CoreRules.PROJECT_TO_SEMI_JOIN);
+
+    final String generated = sql(query).withPostgresql().optimize(rules, null).exec();
+    assertThat(generated, containsString("GROUP BY \"t2\".\"product_id\""));
+  }
+
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-7439">[CALCITE-7439]
+   * RelToSqlConverter should not emit ambiguous GROUP BY after RIGHT JOIN USING
+   * with semi-join rewrite.</a>. */
+  @Test void testPostgresqlRoundTripDistinctRightJoinInSubqueryWithSemiJoinRules() {
+    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"
+        + "RIGHT 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.PROJECT_SUB_QUERY_TO_CORRELATE,
+            CoreRules.FILTER_SUB_QUERY_TO_CORRELATE,
+            CoreRules.JOIN_SUB_QUERY_TO_CORRELATE,
+            CoreRules.PROJECT_SUB_QUERY_TO_MARK_CORRELATE,
+            CoreRules.FILTER_SUB_QUERY_TO_MARK_CORRELATE,
+            CoreRules.MARK_TO_SEMI_OR_ANTI_JOIN_RULE,
+            CoreRules.PROJECT_TO_SEMI_JOIN);
+
+    final String generated = sql(query).withPostgresql().optimize(rules, null).exec();
+    assertThat(generated, containsString("GROUP BY "));
+    assertThat(generated, containsString(".\"product_id\""));
+    assertThat(generated, not(containsString("GROUP BY \"product_id\"")));
+  }
+
+  /** Test case for
+   * <a href="https://issues.apache.org/jira/browse/CALCITE-7439">[CALCITE-7439]
+   * RelToSqlConverter should not emit ambiguous GROUP BY after FULL JOIN USING
+   * with semi-join rewrite.</a>. */
+  @Test void testPostgresqlRoundTripDistinctFullJoinInSubqueryWithSemiJoinRules() {
+    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"
+        + "FULL 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.PROJECT_SUB_QUERY_TO_CORRELATE,
+            CoreRules.FILTER_SUB_QUERY_TO_CORRELATE,
+            CoreRules.JOIN_SUB_QUERY_TO_CORRELATE,
+            CoreRules.PROJECT_SUB_QUERY_TO_MARK_CORRELATE,
+            CoreRules.FILTER_SUB_QUERY_TO_MARK_CORRELATE,
+            CoreRules.MARK_TO_SEMI_OR_ANTI_JOIN_RULE,
+            CoreRules.PROJECT_TO_SEMI_JOIN);
+
+    final String generated = sql(query).withPostgresql().optimize(rules, null).exec();
+    assertThat(generated, containsString("GROUP BY "));
+    assertThat(generated, containsString(".\"product_id\""));
+    assertThat(generated, not(containsString("GROUP BY \"product_id\"")));
+  }
+
   @Test void testNotBetween() {
     Sql f = fixture().withConvertletTable(new SqlRexConvertletTable() {
       @Override public @Nullable SqlRexConvertlet get(SqlCall call) {