Optimize away redundant DISTINCT for primary key (replay of crate/crate#20505)

docs/appendices/release-notes/6.6.0.rst

@@ -109,6 +109,11 @@ Performance and Resilience Improvements
   Previously, the resulting logical plan would contain 1 inner join and 2 cross
   joins. Now, the logical plan has 2 inner joins and 1 cross join.
 
+- Optimized queries with ``DISTINCT`` when the grouping columns include all
+  columns of an explicit primary key. In these cases, we can skip
+  unnecessary duplicate elimination because the primary key guarantees that
+  each row is unique.
+
 Administration and Operations
 -----------------------------
 

server/src/main/java/io/crate/planner/operators/HashAggregate.java

@@ -194,6 +194,10 @@ public List<Function> aggregates() {
         return aggregates;
     }
 
+    public boolean rewriteDistinct() {
+        return rewriteDistinct;
+    }
+
     @Override
     public List<Symbol> outputs() {
         return new ArrayList<>(aggregates);

server/src/main/java/io/crate/planner/operators/LogicalPlanner.java

@@ -128,6 +128,7 @@
 import io.crate.planner.optimizer.rule.MoveOrderBeneathUnion;
 import io.crate.planner.optimizer.rule.OptimizeCollectWhereClauseAccess;
 import io.crate.planner.optimizer.rule.RemoveOrderBeneathInsert;
+import io.crate.planner.optimizer.rule.RemoveRedundantDistinct;
 import io.crate.planner.optimizer.rule.RemoveRedundantEval;
 import io.crate.planner.optimizer.rule.ReorderHashJoin;
 import io.crate.planner.optimizer.rule.ReorderNestedLoopJoin;
@@ -185,6 +186,7 @@ public class LogicalPlanner {
         new MoveOrderBeneathRename(),
         new DeduplicateOrder(),
         new OptimizeCollectWhereClauseAccess(),
+        new RemoveRedundantDistinct(),
         new RewriteGroupByKeysLimitToLimitDistinct(),
         new RewriteDistinctAggToGroupBy(),
         new MoveConstantJoinConditionsBeneathJoin(),

server/src/main/java/io/crate/planner/optimizer/rule/RemoveRedundantDistinct.java

@@ -0,0 +1,223 @@
+/*
+ * Licensed to Crate.io GmbH ("Crate") under one or more contributor
+ * license agreements.  See the NOTICE file distributed with this work for
+ * additional information regarding copyright ownership.  Crate licenses
+ * this file to you under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.  You may
+ * obtain a copy of the License at
+ *
+ *   http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+ * WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.  See the
+ * License for the specific language governing permissions and limitations
+ * under the License.
+ *
+ * However, if you have executed another commercial license agreement
+ * with Crate these terms will supersede the license and you may use the
+ * software solely pursuant to the terms of the relevant commercial agreement.
+ */
+
+package io.crate.planner.optimizer.rule;
+
+import static io.crate.planner.optimizer.matcher.Pattern.typeOf;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.function.UnaryOperator;
+
+import io.crate.expression.symbol.AliasSymbol;
+import io.crate.expression.symbol.ScopedColumn;
+import io.crate.expression.symbol.Symbol;
+import io.crate.metadata.ColumnIdent;
+import io.crate.metadata.DocTableInfo;
+import io.crate.metadata.RelationName;
+import io.crate.metadata.doc.SysColumns;
+import io.crate.planner.operators.Collect;
+import io.crate.planner.operators.Eval;
+import io.crate.planner.operators.Filter;
+import io.crate.planner.operators.GroupHashAggregate;
+import io.crate.planner.operators.Limit;
+import io.crate.planner.operators.LogicalPlan;
+import io.crate.planner.operators.Order;
+import io.crate.planner.operators.Rename;
+import io.crate.planner.optimizer.Rule;
+import io.crate.planner.optimizer.matcher.Captures;
+import io.crate.planner.optimizer.matcher.Pattern;
+
+/**
+ * Removes a redundant DISTINCT operation when its grouping keys
+ * contain every column of the source table's primary key.
+ */
+public final class RemoveRedundantDistinct implements Rule<GroupHashAggregate> {
+
+    private final Pattern<GroupHashAggregate> pattern;
+
+    public RemoveRedundantDistinct() {
+        this.pattern = typeOf(GroupHashAggregate.class)
+            .with(groupBy -> groupBy.aggregates().isEmpty());
+    }
+
+    @Override
+    public Pattern<GroupHashAggregate> pattern() {
+        return pattern;
+    }
+
+    @Override
+    public LogicalPlan apply(
+        GroupHashAggregate groupBy,
+        Captures captures,
+        Rule.Context context
+    ) {
+        if (!isUniqueByPrimaryKey(groupBy.source(), groupBy.groupKeys(), context.resolvePlan())) {
+            return null;
+        }
+
+        // The grouping cannot eliminate any rows because the keys already
+        // contain the complete primary key. Preserve the requested outputs
+        // without performing the redundant grouping operation.
+        return Eval.create(groupBy.source(), groupBy.outputs());
+    }
+
+    public static boolean isUniqueByPrimaryKey(
+        LogicalPlan plan,
+        List<Symbol> groupKeys,
+        UnaryOperator<LogicalPlan> resolvePlan
+    ) {
+        if (plan == null || groupKeys == null || groupKeys.isEmpty()) {
+            return false;
+        }
+
+        LogicalPlan current = resolvePlan.apply(plan);
+        List<Symbol> currentKeys = groupKeys;
+
+        while (current != null) {
+            if (current instanceof Filter || current instanceof Limit || current instanceof Order) {
+                // These operators do not alter the column identities or duplicate rows in a way
+                // that breaks uniqueness, so we just pass through to their source.
+                if (current.sources().size() != 1) {
+                    return false;
+                }
+                current = resolvePlan.apply(current.sources().getFirst());
+            } else if (current instanceof Rename rename) {
+                // A Rename operator alters the symbol names. We map our current group keys
+                // back to the original symbols produced by the Rename's source.
+                if (rename.sources().size() != 1) {
+                    return false;
+                }
+                LogicalPlan source = resolvePlan.apply(rename.source());
+                List<Symbol> renameOutputs = rename.outputs();
+                List<Symbol> sourceOutputs = source.outputs();
+                if (renameOutputs.size() != sourceOutputs.size()) {
+                    return false;
+                }
+
+                List<Symbol> mappedKeys = new ArrayList<>(currentKeys.size());
+                for (Symbol key : currentKeys) {
+                    // Match the unwrapped expression so aliases do not prevent identifying
+                    // a grouping key that originates from the source plan.
+                    Symbol unwrappedKey = unwrapAlias(key);
+                    int idx = -1;
+                    for (int i = 0; i < renameOutputs.size(); i++) {
+                        if (unwrapAlias(renameOutputs.get(i)).equals(unwrappedKey)) {
+                            idx = i;
+                            break;
+                        }
+                    }
+                    if (idx < 0) {
+                        return false;
+                    }
+                    mappedKeys.add(sourceOutputs.get(idx));
+                }
+                currentKeys = mappedKeys;
+                current = source;
+            } else if (current instanceof Eval eval) {
+                if (eval.sources().size() != 1) {
+                    return false;
+                }
+                LogicalPlan source = resolvePlan.apply(eval.source());
+                List<Symbol> evalOutputs = eval.outputs();
+                List<Symbol> sourceOutputs = source.outputs();
+
+                // Only propagate keys that are preserved by Eval as direct source
+                // outputs. Computed expressions are not considered when proving
+                // primary-key uniqueness.
+                List<Symbol> mappedKeys = new ArrayList<>(currentKeys.size());
+                for (Symbol key : currentKeys) {
+                    Symbol unwrappedKey = unwrapAlias(key);
+                    boolean foundInEval = false;
+                    for (Symbol evalOut : evalOutputs) {
+                        if (unwrapAlias(evalOut).equals(unwrappedKey)) {
+                            foundInEval = true;
+                            break;
+                        }
+                    }
+                    if (!foundInEval) {
+                        return false;
+                    }
+
+                    // Ensure the key can also be traced to an actual source output.
+                    // Otherwise, do not assume that the expression preserves uniqueness.
+                    Symbol matchingSourceOutput = null;
+                    for (Symbol srcOut : sourceOutputs) {
+                        if (unwrapAlias(srcOut).equals(unwrappedKey)) {
+                            matchingSourceOutput = srcOut;
+                            break;
+                        }
+                    }
+                    if (matchingSourceOutput == null) {
+                        return false;
+                    }
+                    mappedKeys.add(matchingSourceOutput);
+                }
+                currentKeys = mappedKeys;
+                current = source;
+            } else if (current instanceof Collect collect) {
+                if (!(collect.relation().tableInfo() instanceof DocTableInfo tableInfo)) {
+                    return false;
+                }
+
+                // If the PK is auto-generated (like _id), user columns can't guarantee uniqueness
+                if (tableInfo.hasAutoGeneratedPrimaryKey()) {
+                    return false;
+                }
+
+                var primaryKey = tableInfo.primaryKey();
+                if (primaryKey.isEmpty() || primaryKey.contains(SysColumns.ID.COLUMN)) {
+                    return false;
+                }
+                RelationName tableIdent = tableInfo.ident();
+                Set<ColumnIdent> groupedColumns = new HashSet<>();
+
+                for (Symbol key : currentKeys) {
+                    Symbol unwrapped = unwrapAlias(key);
+                    if (unwrapped instanceof ScopedColumn scopedColumn) {
+                        if (scopedColumn.relation().equals(tableIdent)) {
+                            groupedColumns.add(scopedColumn.column());
+                        }
+                    }
+                }
+
+                // The rows are uniquely identifiable if every column of the primary key
+                // is present in our grouping columns.
+                return groupedColumns.containsAll(primaryKey);
+            } else {
+                // Joins, unions, aggregates and other unsupported operators may change
+                // row multiplicity or key identity. Retain DISTINCT unless uniqueness
+                // can be proven
+                return false;
+            }
+        }
+        return false;
+    }
+
+    private static Symbol unwrapAlias(Symbol symbol) {
+        while (symbol instanceof AliasSymbol alias) {
+            symbol = alias.symbol();
+        }
+        return symbol;
+    }
+}

server/src/main/java/io/crate/planner/optimizer/rule/RewriteDistinctAggToGroupBy.java

@@ -22,12 +22,14 @@
 package io.crate.planner.optimizer.rule;
 
 import static io.crate.planner.optimizer.matcher.Pattern.typeOf;
+import static io.crate.planner.optimizer.rule.RemoveRedundantDistinct.isUniqueByPrimaryKey;
 
 import java.util.ArrayList;
 import java.util.LinkedHashMap;
 import java.util.List;
 import java.util.Map;
 import java.util.Set;
+import java.util.function.UnaryOperator;
 
 import io.crate.common.annotations.VisibleForTesting;
 import io.crate.execution.engine.aggregation.impl.CountAggregation;
@@ -131,7 +133,8 @@ public RewriteDistinctAggToGroupBy() {
     }
 
     private static boolean matches(HashAggregate agg) {
-        return agg.aggregates().stream()
+        return agg.rewriteDistinct() &&
+        agg.aggregates().stream()
             .allMatch(fn ->
                 fn.distinct() &&
                     hasNoFilter(fn) &&
@@ -187,6 +190,10 @@ private static LogicalPlan applyToSameColumn(HashAggregate aggregate, Rule.Conte
             return null;
         }
 
+        if (isPrimaryKeyColumn(resolvedSource, groupKey, context.resolvePlan())) {
+            return null;
+        }
+
         GroupHashAggregate dedup = new GroupHashAggregate(source, List.of(groupKey), List.of());
         return new HashAggregate(dedup, aggregate.aggregates(), false);
     }
@@ -236,4 +243,8 @@ private static boolean meetsNumRowsRequirement(Stats stats) {
 
         return stats.numDocs() <= MAX_TABLE_SIZE;
     }
+
+    private static boolean isPrimaryKeyColumn(LogicalPlan plan, Symbol groupKey, UnaryOperator<LogicalPlan> resolvePlan) {
+        return isUniqueByPrimaryKey(plan, List.of(groupKey), resolvePlan);
+    }
 }

server/src/main/java/io/crate/planner/optimizer/rule/RewriteGroupByKeysLimitToLimitDistinct.java

@@ -27,6 +27,7 @@
 import org.elasticsearch.Version;
 
 import io.crate.expression.symbol.Literal;
+import io.crate.metadata.settings.CoordinatorSessionSettings;
 import io.crate.planner.operators.GroupHashAggregate;
 import io.crate.planner.operators.Limit;
 import io.crate.planner.operators.LimitDistinct;
@@ -166,6 +167,15 @@ public LogicalPlan apply(Limit limit,
                              Captures captures,
                              Rule.Context context) {
         GroupHashAggregate groupBy = captures.get(groupCapture);
+        // Defer to RemoveRedundantDistinct when it is enabled and PK uniqueness
+        // makes the grouping redundant. Otherwise, allow LimitDistinct to act
+        // as a fallback optimization.
+        boolean removeDistinctExcluded = context.txnCtx().sessionSettings() instanceof CoordinatorSessionSettings cs
+            && cs.excludedOptimizerRules().contains(RemoveRedundantDistinct.class);
+        if (!removeDistinctExcluded
+            && RemoveRedundantDistinct.isUniqueByPrimaryKey(groupBy.source(), groupBy.groupKeys(), context.resolvePlan())) {
+            return null;
+        }
         if (!eagerTerminateIsLikely(limit, groupBy, context.planStats())) {
             return null;
         }

server/src/test/java/io/crate/integrationtests/GroupByAggregateTest.java

@@ -1435,7 +1435,7 @@ public void testSelectCollectSetAndGroupBy() {
     @UseRandomizedOptimizerRules(0)
     @Test
     public void test_optimized_limit_distinct_returns_2_unique_items() throws Throwable {
-        execute("create table m.tbl (id int primary key)");
+        execute("create table m.tbl (id int)");
         Object[][] ids = new Object[100][];
         for (int i = 0; i < ids.length; i++) {
             ids[i] = new Object[] { i };

server/src/test/java/io/crate/integrationtests/PgCatalogITest.java

@@ -295,6 +295,7 @@ public void testPgSettingsTable() throws Exception {
             "optimizer_move_order_beneath_union| true| Indicates if the optimizer rule MoveOrderBeneathUnion is activated.| NULL| NULL",
             "optimizer_optimize_collect_where_clause_access| true| Indicates if the optimizer rule OptimizeCollectWhereClauseAccess is activated.| NULL| NULL",
             "optimizer_remove_order_beneath_insert| true| Indicates if the optimizer rule RemoveOrderBeneathInsert is activated.| NULL| NULL",
+            "optimizer_remove_redundant_distinct| true| Indicates if the optimizer rule RemoveRedundantDistinct is activated.| NULL| NULL",
             "optimizer_remove_redundant_eval| true| Indicates if the optimizer rule RemoveRedundantEval is activated.| NULL| NULL",
             "optimizer_reorder_hash_join| true| Indicates if the optimizer rule ReorderHashJoin is activated.| NULL| NULL",
             "optimizer_reorder_nested_loop_join| false| Indicates if the optimizer rule ReorderNestedLoopJoin is activated.| NULL| NULL",

server/src/test/java/io/crate/integrationtests/ShowIntegrationTest.java

@@ -439,6 +439,7 @@ public void testShowAll() {
             "optimizer_move_order_beneath_union| true| Indicates if the optimizer rule MoveOrderBeneathUnion is activated.",
             "optimizer_optimize_collect_where_clause_access| true| Indicates if the optimizer rule OptimizeCollectWhereClauseAccess is activated.",
             "optimizer_remove_order_beneath_insert| true| Indicates if the optimizer rule RemoveOrderBeneathInsert is activated.",
+            "optimizer_remove_redundant_distinct| true| Indicates if the optimizer rule RemoveRedundantDistinct is activated.",
             "optimizer_remove_redundant_eval| true| Indicates if the optimizer rule RemoveRedundantEval is activated.",
             "optimizer_reorder_hash_join| true| Indicates if the optimizer rule ReorderHashJoin is activated.",
             "optimizer_reorder_nested_loop_join| false| Indicates if the optimizer rule ReorderNestedLoopJoin is activated.",

server/src/test/java/io/crate/planner/GroupByScalarPlannerTest.java

@@ -38,6 +38,7 @@
 import io.crate.expression.symbol.Function;
 import io.crate.metadata.RowGranularity;
 import io.crate.planner.node.dql.Collect;
+import io.crate.planner.optimizer.rule.RemoveRedundantDistinct;
 import io.crate.test.integration.CrateDummyClusterServiceUnitTest;
 import io.crate.testing.SQLExecutor;
 import io.crate.types.DataTypes;
@@ -52,6 +53,7 @@ public void prepare() throws IOException {
             .setNumNodes(2)
             .build()
             .addTable(TableDefinitions.USER_TABLE_DEFINITION);
+        e.getSessionSettings().excludedOptimizerRules().add(RemoveRedundantDistinct.class);
     }
 
     @Test

server/src/test/java/io/crate/planner/InsertPlannerTest.java

@@ -61,6 +61,7 @@
 import io.crate.planner.node.dql.join.Join;
 import io.crate.planner.operators.InsertFromValues;
 import io.crate.planner.operators.LogicalPlan;
+import io.crate.planner.optimizer.rule.RemoveRedundantDistinct;
 import io.crate.test.integration.CrateDummyClusterServiceUnitTest;
 import io.crate.testing.SQLExecutor;
 import io.crate.types.DataTypes;
@@ -94,6 +95,7 @@ public void prepare() throws IOException {
                 ") clustered into 4 shards")
             .addTable("create table source (id int primary key, name string)")
             .addTable("CREATE TABLE double_parted(x int, y int) PARTITIONED BY (x, y)");
+        e.getSessionSettings().excludedOptimizerRules().add(RemoveRedundantDistinct.class);
     }
 
     @Test