[refactor](Nereids): compute unique and uniform property respectively (…

…apache#32908)

fe/fe-core/src/main/java/org/apache/doris/nereids/properties/FunctionalDependencies.java

@@ -136,6 +136,10 @@ public void addUniqueSlot(ImmutableSet<Slot> slotSet) {
             uniqueSet.add(slotSet);
         }
 
+        public void addUniqueSlot(FunctionalDependencies functionalDependencies) {
+            uniqueSet.add(functionalDependencies.uniqueSet);
+        }
+
         public void addFdItems(ImmutableSet<FdItem> items) {
             fdItems = ImmutableSet.copyOf(items);
         }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/AbstractPlan.java

@@ -179,7 +179,7 @@ public LogicalProperties computeLogicalProperties() {
         } else {
             Supplier<List<Slot>> outputSupplier = Suppliers.memoize(this::computeOutput);
             Supplier<FunctionalDependencies> fdSupplier = () -> this instanceof LogicalPlan
-                    ? ((LogicalPlan) this).computeFuncDeps(outputSupplier)
+                    ? ((LogicalPlan) this).computeFuncDeps()
                     : FunctionalDependencies.EMPTY_FUNC_DEPS;
             return new LogicalProperties(outputSupplier, fdSupplier);
         }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/BlockFuncDepsPropagation.java

@@ -19,25 +19,31 @@
 
 import org.apache.doris.nereids.properties.FdItem;
 import org.apache.doris.nereids.properties.FunctionalDependencies;
-import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 
 import com.google.common.collect.ImmutableSet;
 
-import java.util.List;
-import java.util.function.Supplier;
-
 /**
  * Block fd propagation, it always returns an empty fd
  */
 public interface BlockFuncDepsPropagation extends LogicalPlan {
     @Override
-    default FunctionalDependencies computeFuncDeps(Supplier<List<Slot>> outputSupplier) {
+    default FunctionalDependencies computeFuncDeps() {
         return FunctionalDependencies.EMPTY_FUNC_DEPS;
     }
 
     @Override
-    default ImmutableSet<FdItem> computeFdItems(Supplier<List<Slot>> outputSupplier) {
+    default ImmutableSet<FdItem> computeFdItems() {
         return ImmutableSet.of();
     }
+
+    @Override
+    default void computeUnique(FunctionalDependencies.Builder fdBuilder) {
+        // don't generate
+    }
+
+    @Override
+    default void computeUniform(FunctionalDependencies.Builder fdBuilder) {
+        // don't generate
+    }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/PropagateFuncDeps.java

@@ -19,20 +19,16 @@
 
 import org.apache.doris.nereids.properties.FdItem;
 import org.apache.doris.nereids.properties.FunctionalDependencies;
-import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.plans.logical.LogicalPlan;
 
 import com.google.common.collect.ImmutableSet;
 
-import java.util.List;
-import java.util.function.Supplier;
-
 /**
  * Propagate fd, keep children's fd
  */
 public interface PropagateFuncDeps extends LogicalPlan {
     @Override
-    default FunctionalDependencies computeFuncDeps(Supplier<List<Slot>> outputSupplier) {
+    default FunctionalDependencies computeFuncDeps() {
         if (children().size() == 1) {
             // Note when changing function dependencies, we always clone it.
             // So it's safe to return a reference
@@ -42,13 +38,13 @@ default FunctionalDependencies computeFuncDeps(Supplier<List<Slot>> outputSuppli
         children().stream()
                 .map(p -> p.getLogicalProperties().getFunctionalDependencies())
                 .forEach(builder::addFunctionalDependencies);
-        ImmutableSet<FdItem> fdItems = computeFdItems(outputSupplier);
+        ImmutableSet<FdItem> fdItems = computeFdItems();
         builder.addFdItems(fdItems);
         return builder.build();
     }
 
     @Override
-    default ImmutableSet<FdItem> computeFdItems(Supplier<List<Slot>> outputSupplier) {
+    default ImmutableSet<FdItem> computeFdItems() {
         if (children().size() == 1) {
             // Note when changing function dependencies, we always clone it.
             // So it's safe to return a reference
@@ -60,4 +56,14 @@ default ImmutableSet<FdItem> computeFdItems(Supplier<List<Slot>> outputSupplier)
                 .forEach(builder::addAll);
         return builder.build();
     }
+
+    @Override
+    default void computeUnique(FunctionalDependencies.Builder fdBuilder) {
+        fdBuilder.addUniqueSlot(child(0).getLogicalProperties().getFunctionalDependencies());
+    }
+
+    @Override
+    default void computeUniform(FunctionalDependencies.Builder fdBuilder) {
+        fdBuilder.addUniformSlot(child(0).getLogicalProperties().getFunctionalDependencies());
+    }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalAggregate.java

@@ -21,17 +21,13 @@
 import org.apache.doris.nereids.properties.FdFactory;
 import org.apache.doris.nereids.properties.FdItem;
 import org.apache.doris.nereids.properties.FunctionalDependencies;
-import org.apache.doris.nereids.properties.FunctionalDependencies.Builder;
 import org.apache.doris.nereids.properties.LogicalProperties;
 import org.apache.doris.nereids.properties.TableFdItem;
-import org.apache.doris.nereids.trees.expressions.Alias;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.NamedExpression;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.expressions.SlotReference;
-import org.apache.doris.nereids.trees.expressions.VirtualSlotReference;
 import org.apache.doris.nereids.trees.expressions.functions.ExpressionTrait;
-import org.apache.doris.nereids.trees.expressions.functions.agg.AggregateFunction;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Count;
 import org.apache.doris.nereids.trees.expressions.functions.agg.Ndv;
 import org.apache.doris.nereids.trees.plans.Plan;
@@ -45,12 +41,9 @@
 import com.google.common.collect.ImmutableList;
 import com.google.common.collect.ImmutableSet;
 
-import java.util.HashSet;
 import java.util.List;
 import java.util.Objects;
 import java.util.Optional;
-import java.util.Set;
-import java.util.function.Supplier;
 import java.util.stream.Collectors;
 
 /**
@@ -300,80 +293,85 @@ public LogicalAggregate<Plan> withNormalized(List<Expression> normalizedGroupBy,
                 hasPushed, sourceRepeat, Optional.empty(), Optional.empty(), normalizedChild);
     }
 
-    private void updateFuncDepsGroupByUnique(NamedExpression namedExpression, Builder fdBuilder) {
-        if (ExpressionUtils.isInjective(namedExpression)) {
-            fdBuilder.addUniqueSlot(ImmutableSet.copyOf(namedExpression.getInputSlots()));
-            return;
+    private boolean isUniqueGroupByUnique(NamedExpression namedExpression) {
+        if (namedExpression.children().size() != 1) {
+            return false;
         }
+        Expression agg = namedExpression.child(0);
+        return ExpressionUtils.isInjectiveAgg(agg)
+                && child().getLogicalProperties().getFunctionalDependencies().isUniqueAndNotNull(agg.getInputSlots());
+    }
 
-        if (!(namedExpression instanceof Alias && namedExpression.child(0) instanceof AggregateFunction)) {
-            return;
+    private boolean isUniformGroupByUnique(NamedExpression namedExpression) {
+        if (namedExpression.children().size() != 1) {
+            return false;
         }
+        Expression agg = namedExpression.child(0);
+        return agg instanceof Count || agg instanceof Ndv;
+    }
 
-        AggregateFunction agg = (AggregateFunction) namedExpression.child(0);
-        if (agg instanceof Count || agg instanceof Ndv) {
-            fdBuilder.addUniformSlot(namedExpression.toSlot());
+    @Override
+    public void computeUnique(FunctionalDependencies.Builder fdBuilder) {
+        if (this.sourceRepeat.isPresent()) {
+            // roll up may generate new data
+            return;
+        }
+        FunctionalDependencies childFd = child(0).getLogicalProperties().getFunctionalDependencies();
+        ImmutableSet<Slot> groupByKeys = groupByExpressions.stream()
+                .map(s -> (Slot) s)
+                .collect(ImmutableSet.toImmutableSet());
+        // when group by all tuples, the result only have one row
+        if (groupByExpressions.isEmpty() || childFd.isUniformAndNotNull(groupByKeys)) {
+            getOutput().forEach(fdBuilder::addUniqueSlot);
             return;
         }
 
-        if (ExpressionUtils.isInjectiveAgg(agg)
-                && child().getLogicalProperties().getFunctionalDependencies().isUniqueAndNotNull(agg.getInputSlots())) {
-            fdBuilder.addUniqueSlot(namedExpression.toSlot());
+        // propagate all unique slots
+        fdBuilder.addUniqueSlot(childFd);
+
+        // group by keys is unique
+        fdBuilder.addUniqueSlot(groupByKeys);
+
+        // group by unique may has unique aggregate result
+        if (childFd.isUniqueAndNotNull(groupByKeys)) {
+            for (NamedExpression namedExpression : getOutputExpressions()) {
+                if (isUniqueGroupByUnique(namedExpression)) {
+                    fdBuilder.addUniqueSlot(namedExpression.toSlot());
+                }
+            }
         }
     }
 
     @Override
-    public FunctionalDependencies computeFuncDeps(Supplier<List<Slot>> outputSupplier) {
+    public void computeUniform(FunctionalDependencies.Builder fdBuilder) {
+        // always propagate uniform
         FunctionalDependencies childFd = child(0).getLogicalProperties().getFunctionalDependencies();
-        Set<Slot> outputSet = new HashSet<>(outputSupplier.get());
-        Builder fdBuilder = new Builder();
-        // when group by all tuples, the result only have one row
-        if (groupByExpressions.isEmpty()) {
-            outputSet.forEach(s -> {
-                fdBuilder.addUniformSlot(s);
-                fdBuilder.addUniqueSlot(s);
-            });
-            return fdBuilder.build();
-        }
+        fdBuilder.addUniformSlot(childFd);
 
-        // when group by complicate expression or virtual slot, just propagate uniform slots
-        if (groupByExpressions.stream()
-                .anyMatch(s -> !(s instanceof SlotReference) || s instanceof VirtualSlotReference)) {
-            fdBuilder.addUniformSlot(childFd);
-            fdBuilder.pruneSlots(outputSet);
-            return fdBuilder.build();
+        if (this.sourceRepeat.isPresent()) {
+            // roll up may generate new data
+            return;
         }
-
-        // when group by uniform slot, the result only have one row
         ImmutableSet<Slot> groupByKeys = groupByExpressions.stream()
                 .map(s -> (Slot) s)
                 .collect(ImmutableSet.toImmutableSet());
-        if (childFd.isUniformAndNotNull(groupByKeys)) {
-            outputSupplier.get().forEach(s -> {
-                fdBuilder.addUniformSlot(s);
-                fdBuilder.addUniqueSlot(s);
-            });
+        // when group by all tuples, the result only have one row
+        if (groupByExpressions.isEmpty() || childFd.isUniformAndNotNull(groupByKeys)) {
+            getOutput().forEach(fdBuilder::addUniformSlot);
+            return;
         }
 
-        // when group by unique slot, the result depends on agg func
         if (childFd.isUniqueAndNotNull(groupByKeys)) {
             for (NamedExpression namedExpression : getOutputExpressions()) {
-                updateFuncDepsGroupByUnique(namedExpression, fdBuilder);
+                if (isUniformGroupByUnique(namedExpression)) {
+                    fdBuilder.addUniformSlot(namedExpression.toSlot());
+                }
             }
         }
-
-        // group by keys is unique
-        fdBuilder.addUniqueSlot(groupByKeys);
-        fdBuilder.pruneSlots(outputSet);
-
-        ImmutableSet<FdItem> fdItems = computeFdItems(outputSupplier);
-        fdBuilder.addFdItems(fdItems);
-
-        return fdBuilder.build();
     }
 
     @Override
-    public ImmutableSet<FdItem> computeFdItems(Supplier<List<Slot>> outputSupplier) {
+    public ImmutableSet<FdItem> computeFdItems() {
         ImmutableSet.Builder<FdItem> builder = ImmutableSet.builder();
 
         ImmutableSet<SlotReference> groupByExprs = getGroupByExpressions().stream()

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalAssertNumRows.java

@@ -18,18 +18,21 @@
 package org.apache.doris.nereids.trees.plans.logical;
 
 import org.apache.doris.nereids.memo.GroupExpression;
+import org.apache.doris.nereids.properties.FdItem;
+import org.apache.doris.nereids.properties.FunctionalDependencies.Builder;
 import org.apache.doris.nereids.properties.LogicalProperties;
 import org.apache.doris.nereids.trees.expressions.AssertNumRowsElement;
+import org.apache.doris.nereids.trees.expressions.AssertNumRowsElement.Assertion;
 import org.apache.doris.nereids.trees.expressions.Expression;
 import org.apache.doris.nereids.trees.expressions.Slot;
 import org.apache.doris.nereids.trees.plans.Plan;
 import org.apache.doris.nereids.trees.plans.PlanType;
-import org.apache.doris.nereids.trees.plans.PropagateFuncDeps;
 import org.apache.doris.nereids.trees.plans.visitor.PlanVisitor;
 import org.apache.doris.nereids.util.Utils;
 
 import com.google.common.base.Preconditions;
 import com.google.common.collect.ImmutableList;
+import com.google.common.collect.ImmutableSet;
 
 import java.util.List;
 import java.util.Objects;
@@ -42,8 +45,7 @@
  * If the number of rows is more than the desired num of rows, the query will be cancelled.
  * The cancelled reason will be reported by Backend and displayed back to the user.
  */
-public class LogicalAssertNumRows<CHILD_TYPE extends Plan> extends LogicalUnary<CHILD_TYPE> implements
-        PropagateFuncDeps {
+public class LogicalAssertNumRows<CHILD_TYPE extends Plan> extends LogicalUnary<CHILD_TYPE> {
 
     private final AssertNumRowsElement assertNumRowsElement;
 
@@ -118,4 +120,29 @@ public Plan withGroupExprLogicalPropChildren(Optional<GroupExpression> groupExpr
     public List<Slot> computeOutput() {
         return child().getOutput().stream().map(o -> o.withNullable(true)).collect(Collectors.toList());
     }
+
+    @Override
+    public ImmutableSet<FdItem> computeFdItems() {
+        return ImmutableSet.of();
+    }
+
+    @Override
+    public void computeUnique(Builder fdBuilder) {
+        if (assertNumRowsElement.getDesiredNumOfRows() == 1
+                && (assertNumRowsElement.getAssertion() == Assertion.EQ
+                || assertNumRowsElement.getAssertion() == Assertion.LT
+                || assertNumRowsElement.getAssertion() == Assertion.LE)) {
+            getOutput().forEach(fdBuilder::addUniqueSlot);
+        }
+    }
+
+    @Override
+    public void computeUniform(Builder fdBuilder) {
+        if (assertNumRowsElement.getDesiredNumOfRows() == 1
+                && (assertNumRowsElement.getAssertion() == Assertion.EQ
+                || assertNumRowsElement.getAssertion() == Assertion.LT
+                || assertNumRowsElement.getAssertion() == Assertion.LE)) {
+            getOutput().forEach(fdBuilder::addUniformSlot);
+        }
+    }
 }

fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalCatalogRelation.java

@@ -49,7 +49,6 @@
 import java.util.Objects;
 import java.util.Optional;
 import java.util.Set;
-import java.util.function.Supplier;
 
 /**
  * abstract class catalog relation for logical relation
@@ -128,9 +127,16 @@ public String qualifiedName() {
     }
 
     @Override
-    public FunctionalDependencies computeFuncDeps(Supplier<List<Slot>> outputSupplier) {
+    public FunctionalDependencies computeFuncDeps() {
         Builder fdBuilder = new Builder();
-        Set<Slot> outputSet = Utils.fastToImmutableSet(outputSupplier.get());
+        computeUnique(fdBuilder);
+        fdBuilder.addFdItems(computeFdItems(Utils.fastToImmutableSet(getOutputSet())));
+        return fdBuilder.build();
+    }
+
+    @Override
+    public void computeUnique(FunctionalDependencies.Builder fdBuilder) {
+        Set<Slot> outputSet = Utils.fastToImmutableSet(getOutputSet());
         if (table instanceof OlapTable && ((OlapTable) table).getKeysType().isAggregationFamily()) {
             ImmutableSet.Builder<Slot> uniqSlots = ImmutableSet.builderWithExpectedSize(outputSet.size());
             for (Slot slot : outputSet) {
@@ -154,13 +160,16 @@ public FunctionalDependencies computeFuncDeps(Supplier<List<Slot>> outputSupplie
             Set<Column> columns = c.getUniqueKeys(table);
             fdBuilder.addUniqueSlot((ImmutableSet) findSlotsByColumn(outputSet, columns));
         }
-        fdBuilder.addFdItems(computeFdItems(outputSet));
-        return fdBuilder.build();
     }
 
     @Override
-    public ImmutableSet<FdItem> computeFdItems(Supplier<List<Slot>> outputSupplier) {
-        return computeFdItems(Utils.fastToImmutableSet(outputSupplier.get()));
+    public void computeUniform(FunctionalDependencies.Builder fdBuilder) {
+        // No uniform slot for catalog relation
+    }
+
+    @Override
+    public ImmutableSet<FdItem> computeFdItems() {
+        return computeFdItems(Utils.fastToImmutableSet(getOutputSet()));
     }
 
     private ImmutableSet<FdItem> computeFdItems(Set<Slot> outputSet) {