[CALCITE-5894] Add SortRemoveRedundantRule to remove redundant sort fields if they are functionally dependent on other sort fields

core/src/main/java/org/apache/calcite/rel/metadata/BuiltInMetadata.java

@@ -833,10 +833,32 @@ interface Handler extends MetadataHandler<Memory> {
     }
   }
 
+  /** Metadata about the functional dependency of columns. */
+  public interface FunctionalDependency extends Metadata {
+    MetadataDef<FunctionalDependency> DEF =
+        MetadataDef.of(FunctionalDependency.class, FunctionalDependency.Handler.class,
+            BuiltInMethod.FUNCTIONAL_DEPENDENCY.method);
+
+    /**
+     * Returns whether column is functionally dependent on columns.
+     */
+    @Nullable Boolean functionallyDetermine(ImmutableBitSet columns, int column);
+
+    /** Handler API. */
+    interface Handler extends MetadataHandler<FunctionalDependency> {
+      @Nullable Boolean functionallyDetermine(RelNode r, RelMetadataQuery mq,
+          ImmutableBitSet columns, int column);
+
+      @Override default MetadataDef<FunctionalDependency> getDef() {
+        return DEF;
+      }
+    }
+  }
+
   /** The built-in forms of metadata. */
   interface All extends Selectivity, UniqueKeys, RowCount, DistinctRowCount,
       PercentageOriginalRows, ColumnUniqueness, ColumnOrigin, Predicates,
       Collation, Distribution, Size, Parallelism, Memory, AllPredicates,
-      ExpressionLineage, TableReferences, NodeTypes {
+      ExpressionLineage, TableReferences, NodeTypes, FunctionalDependency {
   }
 }

core/src/main/java/org/apache/calcite/rel/metadata/DefaultRelMetadataProvider.java

@@ -62,6 +62,7 @@ protected DefaultRelMetadataProvider() {
             RelMdExplainVisibility.SOURCE,
             RelMdPredicates.SOURCE,
             RelMdAllPredicates.SOURCE,
-            RelMdCollation.SOURCE));
+            RelMdCollation.SOURCE,
+            RelMdFunctionalDependency.SOURCE));
   }
 }

core/src/main/java/org/apache/calcite/rel/metadata/RelMdFunctionalDependency.java

@@ -0,0 +1,68 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF 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.
+ */
+package org.apache.calcite.rel.metadata;
+
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.util.ImmutableBitSet;
+
+import org.checkerframework.checker.nullness.qual.Nullable;
+
+/**
+ * Default implementation of
+ * {@link RelMetadataQuery#functionallyDetermine(RelNode, ImmutableBitSet, int)}
+ * for the standard logical algebra.
+ *
+ * <p>The goal of this provider is to determine whether
+ * column is functionally dependent on columns.
+ *
+ * <p>If the functional dependency cannot be determined, we return false.
+ */
+public class RelMdFunctionalDependency
+    implements MetadataHandler<BuiltInMetadata.FunctionalDependency> {
+  public static final RelMetadataProvider SOURCE =
+      ReflectiveRelMetadataProvider.reflectiveSource(
+          new RelMdFunctionalDependency(), BuiltInMetadata.FunctionalDependency.Handler.class);
+
+  //~ Constructors -----------------------------------------------------------
+
+  protected RelMdFunctionalDependency() {}
+
+  //~ Methods ----------------------------------------------------------------
+
+  @Override public MetadataDef<BuiltInMetadata.FunctionalDependency> getDef() {
+    return BuiltInMetadata.FunctionalDependency.DEF;
+  }
+
+  /** Catch-all implementation for
+   * {@link BuiltInMetadata.FunctionalDependency#functionallyDetermine(ImmutableBitSet, int)},
+   * invoked using reflection.
+   *
+   * @see org.apache.calcite.rel.metadata.RelMetadataQuery#areColumnsUnique(
+   * RelNode, ImmutableBitSet, boolean)
+   */
+  public @Nullable Boolean functionallyDetermine(RelNode rel, RelMetadataQuery mq,
+      ImmutableBitSet columns, int column) {
+    // This is a minimal implementation that needs to be gradually improved in the future.
+    // If columns contains a unique key, the column is functionally dependent on columns.
+    for (ImmutableBitSet subColumns : columns.powerSet()) {
+      if (Boolean.TRUE.equals(mq.areColumnsUnique(rel, subColumns, false))) {
+        return true;
+      }
+    }
+    return false;
+  }
+}

core/src/main/java/org/apache/calcite/rel/metadata/RelMetadataQuery.java

@@ -107,6 +107,7 @@ public class RelMetadataQuery extends RelMetadataQueryBase {
   private BuiltInMetadata.Size.Handler sizeHandler;
   private BuiltInMetadata.UniqueKeys.Handler uniqueKeysHandler;
   private BuiltInMetadata.LowerBoundCost.Handler lowerBoundCostHandler;
+  private BuiltInMetadata.FunctionalDependency.Handler functionalDependencyHandler;
 
   /**
    * Creates the instance with {@link JaninoRelMetadataProvider} instance
@@ -151,6 +152,8 @@ public RelMetadataQuery(MetadataHandlerProvider provider) {
     this.sizeHandler = provider.handler(BuiltInMetadata.Size.Handler.class);
     this.uniqueKeysHandler = provider.handler(BuiltInMetadata.UniqueKeys.Handler.class);
     this.lowerBoundCostHandler = provider.handler(BuiltInMetadata.LowerBoundCost.Handler.class);
+    this.functionalDependencyHandler =
+        provider.handler(BuiltInMetadata.FunctionalDependency.Handler.class);
   }
 
   /** Creates and initializes the instance that will serve as a prototype for
@@ -183,6 +186,8 @@ private RelMetadataQuery(@SuppressWarnings("unused") boolean dummy) {
     this.sizeHandler = initialHandler(BuiltInMetadata.Size.Handler.class);
     this.uniqueKeysHandler = initialHandler(BuiltInMetadata.UniqueKeys.Handler.class);
     this.lowerBoundCostHandler = initialHandler(BuiltInMetadata.LowerBoundCost.Handler.class);
+    this.functionalDependencyHandler =
+        initialHandler(BuiltInMetadata.FunctionalDependency.Handler.class);
   }
 
   private RelMetadataQuery(
@@ -213,6 +218,7 @@ private RelMetadataQuery(
     this.sizeHandler = prototype.sizeHandler;
     this.uniqueKeysHandler = prototype.uniqueKeysHandler;
     this.lowerBoundCostHandler = prototype.lowerBoundCostHandler;
+    this.functionalDependencyHandler = prototype.functionalDependencyHandler;
   }
 
   //~ Methods ----------------------------------------------------------------
@@ -920,4 +926,24 @@ public Boolean isVisibleInExplain(RelNode rel,
       }
     }
   }
+
+  /**
+   * Determines whether column is functionally dependent on columns.
+   *
+   * @param rel the relational expression
+   * @param columns the target columns
+   * @param column the source column
+   * @return whether column is functionally dependent on columns.
+   */
+  public @Nullable Boolean functionallyDetermine(RelNode rel,
+      ImmutableBitSet columns, int column) {
+    for (;;) {
+      try {
+        return functionalDependencyHandler.functionallyDetermine(rel, this,
+            columns, column);
+      } catch (MetadataHandlerProvider.NoHandler e) {
+        functionalDependencyHandler = revise(BuiltInMetadata.FunctionalDependency.Handler.class);
+      }
+    }
+  }
 }

core/src/main/java/org/apache/calcite/rel/rules/CoreRules.java

@@ -732,6 +732,11 @@ private CoreRules() {}
   public static final SortRemoveRedundantRule SORT_REMOVE_REDUNDANT =
       SortRemoveRedundantRule.Config.DEFAULT.toRule();
 
+  /** Rule that removes keys from a {@link Sort}
+   * if those keys are known to be functionally dependent on other sort keys. */
+  public static final SortRemoveRedundantKeysRule SORT_REMOVE_REDUNDANT_KEYS =
+      SortRemoveRedundantKeysRule.Config.DEFAULT.toRule();
+
   /** Rule that pushes a {@link Sort} past a {@link Join}. */
   public static final SortJoinTransposeRule SORT_JOIN_TRANSPOSE =
       SortJoinTransposeRule.Config.DEFAULT.toRule();

core/src/main/java/org/apache/calcite/rel/rules/SortRemoveRedundantKeysRule.java

@@ -0,0 +1,99 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one or more
+ * contributor license agreements.  See the NOTICE file distributed with
+ * this work for additional information regarding copyright ownership.
+ * The ASF 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.
+ */
+package org.apache.calcite.rel.rules;
+
+import org.apache.calcite.plan.RelOptRuleCall;
+import org.apache.calcite.plan.RelRule;
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollations;
+import org.apache.calcite.rel.RelFieldCollation;
+import org.apache.calcite.rel.RelNode;
+import org.apache.calcite.rel.core.Sort;
+import org.apache.calcite.rel.metadata.RelMetadataQuery;
+import org.apache.calcite.util.ImmutableBitSet;
+import org.apache.calcite.util.ImmutableIntList;
+
+import org.immutables.value.Value;
+
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.stream.Collectors;
+
+/**
+ * Planner rule that removes keys from a
+ * a {@link org.apache.calcite.rel.core.Sort} if those keys functionally depend on
+ * other sort field.
+ *
+ * <p>Requires {@link org.apache.calcite.rel.metadata.RelMdFunctionalDependency}.
+ */
+@Value.Enclosing
+public class SortRemoveRedundantKeysRule extends RelRule<SortRemoveRedundantKeysRule.Config> {
+
+  /**
+   * Creates a SortRemoveRedundantRule.
+   */
+  protected SortRemoveRedundantKeysRule(Config config) {
+    super(config);
+  }
+
+  @Override public void onMatch(RelOptRuleCall call) {
+    final Sort sort = call.rel(0);
+    final RelNode input = sort.getInput();
+    final RelMetadataQuery mq = call.getMetadataQuery();
+    final RelCollation collation = sort.getCollation();
+    final ImmutableIntList keys = collation.getKeys();
+    if (keys.isEmpty()) {
+      return;
+    }
+    ImmutableBitSet sortKeySet = ImmutableBitSet.of(keys);
+    final Set<Integer> needRemoveKeys = new HashSet<>();
+    // Remove redundant sort field by functional dependency
+    for (int keyIndex = keys.size() - 1; keyIndex >= 0; keyIndex--) {
+      Integer key = keys.get(keyIndex);
+      final ImmutableBitSet remainingSortKeySet = sortKeySet.except(ImmutableBitSet.of(key));
+      if (Boolean.TRUE.equals(mq.functionallyDetermine(input, remainingSortKeySet, key))) {
+        needRemoveKeys.add(key);
+        sortKeySet = remainingSortKeySet;
+      }
+    }
+    if (needRemoveKeys.isEmpty()) {
+      // No functional dependency, bail out.
+      return;
+    }
+    List<RelFieldCollation> remainingFieldCollationList = collation.getFieldCollations()
+        .stream()
+        .filter(fieldCollation -> !needRemoveKeys.contains(fieldCollation.getFieldIndex()))
+        .collect(Collectors.toList());
+    final Sort result =
+        sort.copy(sort.getTraitSet(), input, RelCollations.of(remainingFieldCollationList));
+    call.transformTo(result);
+  }
+
+  /**
+   * Rule configuration.
+   */
+  @Value.Immutable
+  public interface Config extends RelRule.Config {
+    SortRemoveRedundantKeysRule.Config DEFAULT = ImmutableSortRemoveRedundantKeysRule.Config.of()
+        .withOperandSupplier(b -> b.operand(Sort.class).anyInputs());
+
+    @Override default SortRemoveRedundantKeysRule toRule() {
+      return new SortRemoveRedundantKeysRule(this);
+    }
+  }
+}

core/src/main/java/org/apache/calcite/util/BuiltInMethod.java

@@ -62,6 +62,7 @@
 import org.apache.calcite.rel.metadata.BuiltInMetadata.Distribution;
 import org.apache.calcite.rel.metadata.BuiltInMetadata.ExplainVisibility;
 import org.apache.calcite.rel.metadata.BuiltInMetadata.ExpressionLineage;
+import org.apache.calcite.rel.metadata.BuiltInMetadata.FunctionalDependency;
 import org.apache.calcite.rel.metadata.BuiltInMetadata.LowerBoundCost;
 import org.apache.calcite.rel.metadata.BuiltInMetadata.MaxRowCount;
 import org.apache.calcite.rel.metadata.BuiltInMetadata.Memory;
@@ -781,6 +782,8 @@ public enum BuiltInMethod {
       "cumulativeMemoryWithinPhaseSplit"),
   COLUMN_UNIQUENESS(ColumnUniqueness.class, "areColumnsUnique",
       ImmutableBitSet.class, boolean.class),
+  FUNCTIONAL_DEPENDENCY(FunctionalDependency.class, "functionallyDetermine",
+      ImmutableBitSet.class, int.class),
   COLLATIONS(Collation.class, "collations"),
   DISTRIBUTION(Distribution.class, "distribution"),
   NODE_TYPES(NodeTypes.class, "getNodeTypes"),