infer col name when set

seawinde · Dec 17, 2024 · a1cc3be · a1cc3be
1 parent 0c97e04
commit a1cc3be
Show file tree

Hide file tree

Showing 12 changed files with 197 additions and 34 deletions.
diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushCountIntoUnionAll.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/rules/rewrite/PushCountIntoUnionAll.java
@@ -110,7 +110,8 @@ private Plan doPush(LogicalAggregate<LogicalUnion> agg) {
         List<NamedExpression> newLogicalUnionOutputs = Lists.newArrayList();
         for (NamedExpression ce : upperOutputExpressions) {
             if (ce instanceof Alias) {
-                newLogicalUnionOutputs.add(new SlotReference(ce.getName(), ce.getDataType(), ce.nullable()));
+                newLogicalUnionOutputs.add(new SlotReference(ce.getName(), ce.getDataType(), ce.nullable(),
+                        ce.isNameFromChild()));
             } else if (ce instanceof SlotReference) {
                 newLogicalUnionOutputs.add(ce);
             } else {

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Alias.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/Alias.java
@@ -51,6 +51,10 @@ public Alias(Expression child, String name) {
         this(StatementScopeIdGenerator.newExprId(), child, name, false);
     }
 
+    public Alias(Expression child, String name, boolean nameFromChild) {
+        this(StatementScopeIdGenerator.newExprId(), child, name, nameFromChild);
+    }
+
     public Alias(Expression child) {
         this(StatementScopeIdGenerator.newExprId(), ImmutableList.of(child),
                 Suppliers.memoize(child::toSql), ImmutableList.of(), true);
@@ -99,7 +103,8 @@ public Slot toSlot() throws UnboundException {
                 internalName,
                 slotReference != null
                         ? slotReference.getSubPath()
-                        : ImmutableList.of(), Optional.empty()
+                        : ImmutableList.of(), Optional.empty(),
+                nameFromChild
         );
     }
 
@@ -168,6 +173,7 @@ public <R, C> R accept(ExpressionVisitor<R, C> visitor, C context) {
         return visitor.visitAlias(this, context);
     }
 
+    @Override
     public boolean isNameFromChild() {
         return nameFromChild;
     }

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ArrayItemReference.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/ArrayItemReference.java
@@ -144,7 +144,7 @@ public static class ArrayItemSlot extends SlotReference implements SlotNotFromCh
          */
         public ArrayItemSlot(ExprId exprId, String name, DataType dataType, boolean nullable) {
             super(exprId, name, dataType, nullable, ImmutableList.of(),
-                    null, null, Optional.empty(), ImmutableList.of());
+                    null, null, Optional.empty(), ImmutableList.of(), false);
         }
 
         @Override

diff --git a/...-core/src/main/java/org/apache/doris/nereids/trees/expressions/MarkJoinSlotReference.java b/...-core/src/main/java/org/apache/doris/nereids/trees/expressions/MarkJoinSlotReference.java
@@ -30,12 +30,12 @@ public class MarkJoinSlotReference extends SlotReference {
     final boolean existsHasAgg;
 
     public MarkJoinSlotReference(String name) {
-        super(name, BooleanType.INSTANCE, true);
+        super(name, BooleanType.INSTANCE, true, false);
         this.existsHasAgg = false;
     }
 
     public MarkJoinSlotReference(String name, boolean existsHasAgg) {
-        super(name, BooleanType.INSTANCE, true);
+        super(name, BooleanType.INSTANCE, true, false);
         this.existsHasAgg = existsHasAgg;
     }
 

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/NamedExpression.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/NamedExpression.java
@@ -48,6 +48,10 @@ public List<String> getQualifier() throws UnboundException {
         throw new UnboundException("qualifier");
     }
 
+    public boolean isNameFromChild() {
+        return false;
+    }
+
     /**
      * Get qualified name of NamedExpression.
      *

diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SlotReference.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/expressions/SlotReference.java
@@ -56,41 +56,47 @@ public class SlotReference extends Slot {
 
     private final TableIf table;
     private final Column column;
+    private final boolean nameFromChild;
 
     public SlotReference(String name, DataType dataType) {
         this(StatementScopeIdGenerator.newExprId(), name, dataType, true, ImmutableList.of(),
-                null, null, Optional.empty(), ImmutableList.of());
+                null, null, Optional.empty(), ImmutableList.of(), false);
     }
 
-    public SlotReference(String name, DataType dataType, boolean nullable) {
+    public SlotReference(String name, DataType dataType, boolean nullable, boolean nameFromChild) {
         this(StatementScopeIdGenerator.newExprId(), name, dataType, nullable, ImmutableList.of(),
-                null, null, Optional.empty(), ImmutableList.of());
+                null, null, Optional.empty(), ImmutableList.of(), nameFromChild);
     }
 
     public SlotReference(String name, DataType dataType, boolean nullable, List<String> qualifier) {
         this(StatementScopeIdGenerator.newExprId(), name, dataType, nullable,
-                qualifier, null, null, Optional.empty(), ImmutableList.of());
+                qualifier, null, null, Optional.empty(), ImmutableList.of(),
+                false);
     }
 
     public SlotReference(ExprId exprId, String name, DataType dataType, boolean nullable, List<String> qualifier) {
-        this(exprId, name, dataType, nullable, qualifier, null, null, Optional.empty(), ImmutableList.of());
+        this(exprId, name, dataType, nullable, qualifier, null, null, Optional.empty(),
+                ImmutableList.of(), false);
     }
 
     public SlotReference(ExprId exprId, String name, DataType dataType, boolean nullable,
             List<String> qualifier, @Nullable TableIf table, @Nullable Column column) {
-        this(exprId, name, dataType, nullable, qualifier, table, column, Optional.empty(), ImmutableList.of());
+        this(exprId, name, dataType, nullable, qualifier, table, column, Optional.empty(), ImmutableList.of(),
+                false);
     }
 
     public SlotReference(ExprId exprId, String name, DataType dataType, boolean nullable,
             List<String> qualifier, @Nullable TableIf table, @Nullable Column column, Optional<String> internalName) {
-        this(exprId, name, dataType, nullable, qualifier, table, column, internalName, ImmutableList.of());
+        this(exprId, name, dataType, nullable, qualifier, table, column, internalName, ImmutableList.of(),
+                false);
     }
 
     public SlotReference(ExprId exprId, String name, DataType dataType, boolean nullable,
             List<String> qualifier, @Nullable TableIf table, @Nullable Column column,
-            Optional<String> internalName, List<String> subColLabels) {
+            Optional<String> internalName, List<String> subColLabels, boolean nameFromChild) {
         this(exprId, () -> name, dataType, nullable, qualifier, table, column,
-                buildInternalName(() -> name, subColLabels, internalName), subColLabels, Optional.empty());
+                buildInternalName(() -> name, subColLabels, internalName), subColLabels, Optional.empty(),
+                nameFromChild);
     }
 
     /**
@@ -108,7 +114,8 @@ public SlotReference(ExprId exprId, String name, DataType dataType, boolean null
     public SlotReference(ExprId exprId, Supplier<String> name, DataType dataType, boolean nullable,
             List<String> qualifier, @Nullable TableIf table, @Nullable Column column,
             Supplier<Optional<String>> internalName, List<String> subPath,
-            Optional<Pair<Integer, Integer>> indexInSql) {
+            Optional<Pair<Integer, Integer>> indexInSql,
+            boolean nameFromChild) {
         super(indexInSql);
         this.exprId = exprId;
         this.name = name;
@@ -120,6 +127,7 @@ public SlotReference(ExprId exprId, Supplier<String> name, DataType dataType, bo
         this.column = column;
         this.subPath = Objects.requireNonNull(subPath, "subPath can not be null");
         this.internalName = internalName;
+        this.nameFromChild = nameFromChild;
     }
 
     public static SlotReference of(String name, DataType type) {
@@ -135,13 +143,15 @@ public static SlotReference fromColumn(TableIf table, Column column, List<String
         DataType dataType = DataType.fromCatalogType(column.getType());
         return new SlotReference(StatementScopeIdGenerator.newExprId(), column::getName, dataType,
                 column.isAllowNull(), qualifier, table, column,
-                () -> Optional.of(column.getName()), ImmutableList.of(), Optional.empty());
+                () -> Optional.of(column.getName()), ImmutableList.of(), Optional.empty(),
+                false);
     }
 
     public static SlotReference fromColumn(TableIf table, Column column, String name, List<String> qualifier) {
         DataType dataType = DataType.fromCatalogType(column.getType());
         return new SlotReference(StatementScopeIdGenerator.newExprId(), name, dataType,
-            column.isAllowNull(), qualifier, table, column, Optional.empty(), ImmutableList.of());
+            column.isAllowNull(), qualifier, table, column, Optional.empty(), ImmutableList.of(),
+                false);
     }
 
     @Override
@@ -259,7 +269,7 @@ public SlotReference withNullable(boolean nullable) {
             return this;
         }
         return new SlotReference(exprId, name, dataType, nullable,
-                qualifier, table, column, internalName, subPath, indexInSqlString);
+                qualifier, table, column, internalName, subPath, indexInSqlString, nameFromChild);
     }
 
     @Override
@@ -268,13 +278,13 @@ public Slot withNullableAndDataType(boolean nullable, DataType dataType) {
             return this;
         }
         return new SlotReference(exprId, name, dataType, nullable,
-                qualifier, table, column, internalName, subPath, indexInSqlString);
+                qualifier, table, column, internalName, subPath, indexInSqlString, nameFromChild);
     }
 
     @Override
     public SlotReference withQualifier(List<String> qualifier) {
         return new SlotReference(exprId, name, dataType, nullable, qualifier, table, column, internalName, subPath,
-                indexInSqlString);
+                indexInSqlString, nameFromChild);
     }
 
     @Override
@@ -284,29 +294,29 @@ public SlotReference withName(String name) {
         }
         return new SlotReference(
                 exprId, () -> name, dataType, nullable, qualifier, table, column, internalName, subPath,
-                indexInSqlString);
+                indexInSqlString, nameFromChild);
     }
 
     @Override
     public SlotReference withExprId(ExprId exprId) {
         return new SlotReference(exprId, name, dataType, nullable, qualifier, table, column, internalName, subPath,
-                indexInSqlString);
+                indexInSqlString, nameFromChild);
     }
 
     public SlotReference withSubPath(List<String> subPath) {
         return new SlotReference(exprId, name, dataType, !subPath.isEmpty() || nullable,
-                qualifier, table, column, internalName, subPath, indexInSqlString);
+                qualifier, table, column, internalName, subPath, indexInSqlString, nameFromChild);
     }
 
     @Override
     public Slot withIndexInSql(Pair<Integer, Integer> index) {
         return new SlotReference(exprId, name, dataType, nullable, qualifier, table, column, internalName, subPath,
-                Optional.ofNullable(index));
+                Optional.ofNullable(index), nameFromChild);
     }
 
     public SlotReference withColumn(Column column) {
         return new SlotReference(exprId, name, dataType, nullable, qualifier, table, column, internalName, subPath,
-                indexInSqlString);
+                indexInSqlString, nameFromChild);
     }
 
     public boolean isVisible() {
@@ -337,4 +347,9 @@ private static Supplier<Optional<String>> buildInternalName(
     public String getQualifiedNameWithBackquote() throws UnboundException {
         return Utils.qualifiedNameWithBackquote(getQualifier(), getName());
     }
+
+    @Override
+    public boolean isNameFromChild() {
+        return nameFromChild;
+    }
 }
diff --git a/...java/org/apache/doris/nereids/trees/expressions/functions/AggCombinerFunctionBuilder.java b/...java/org/apache/doris/nereids/trees/expressions/functions/AggCombinerFunctionBuilder.java
@@ -95,7 +95,7 @@ private AggregateFunction buildForEach(String nestedName, List<? extends Object>
                         "foreach must be input array type: '" + nestedName);
             }
             DataType itemType = ((ArrayType) arrayType).getItemType();
-            return new SlotReference("mocked", itemType, (((ArrayType) arrayType).containsNull()));
+            return new SlotReference("mocked", itemType, (((ArrayType) arrayType).containsNull()), false);
         }).collect(Collectors.toList());
         return (AggregateFunction) nestedBuilder.build(nestedName, forEachargs).first;
     }

diff --git a/...-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalSetOperation.java b/...-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalSetOperation.java
@@ -116,7 +116,8 @@ public List<NamedExpression> buildNewOutputs() {
         List<Slot> slots = resetNullableForLeftOutputs();
         ImmutableList.Builder<NamedExpression> newOutputs = ImmutableList.builderWithExpectedSize(slots.size());
         for (Slot slot : slots) {
-            newOutputs.add(new SlotReference(slot.toSql(), slot.getDataType(), slot.nullable()));
+            newOutputs.add(new SlotReference(slot.toSql(), slot.getDataType(), slot.nullable(),
+                    slot.isNameFromChild()));
         }
         return newOutputs.build();
     }
@@ -151,10 +152,10 @@ private List<List<NamedExpression>> castCommonDataTypeOutputs() {
             Expression newLeft = TypeCoercionUtils.castIfNotSameTypeStrict(left, compatibleType);
             Expression newRight = TypeCoercionUtils.castIfNotSameTypeStrict(right, compatibleType);
             if (newLeft instanceof Cast) {
-                newLeft = new Alias(newLeft, left.getName());
+                newLeft = new Alias(newLeft, left.getName(), left.isNameFromChild());
             }
             if (newRight instanceof Cast) {
-                newRight = new Alias(newRight, right.getName());
+                newRight = new Alias(newRight, right.getName(), right.isNameFromChild());
             }
             newLeftOutputs.add((NamedExpression) newLeft);
             newRightOutputs.add((NamedExpression) newRight);

diff --git a/...core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/InferPlanOutputAlias.java b/...core/src/main/java/org/apache/doris/nereids/trees/plans/visitor/InferPlanOutputAlias.java
@@ -22,22 +22,30 @@
 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.plans.Plan;
+import org.apache.doris.nereids.trees.plans.logical.LogicalSetOperation;
 
 import com.google.common.collect.ImmutableCollection;
 import com.google.common.collect.ImmutableMultimap;
 import com.google.common.collect.ImmutableSet;
+import org.apache.logging.log4j.LogManager;
+import org.apache.logging.log4j.Logger;
 
 import java.util.ArrayList;
+import java.util.HashMap;
 import java.util.HashSet;
 import java.util.List;
+import java.util.Map;
 import java.util.Set;
 
 /**
  * Infer output column name when it refers an expression and not has an alias manually.
  */
 public class InferPlanOutputAlias {
 
+    public static final Logger LOG = LogManager.getLogger(InferPlanOutputAlias.class);
+
     private final List<Slot> currentOutputs;
     private final List<NamedExpression> finalOutputs;
     private final Set<Integer> shouldProcessOutputIndex;
@@ -55,22 +63,38 @@ public InferPlanOutputAlias(List<Slot> currentOutputs) {
     /** infer */
     public List<NamedExpression> infer(Plan plan, ImmutableMultimap<ExprId, Integer> currentExprIdAndIndexMap) {
         ImmutableSet<ExprId> currentOutputExprIdSet = currentExprIdAndIndexMap.keySet();
+        Map<ExprId, ExprId> childToParentMapping = new HashMap<>();
         // Breath First Search
         plan.foreachBreath(childPlan -> {
             if (shouldProcessOutputIndex.isEmpty()) {
                 return true;
             }
+            if (childPlan instanceof LogicalSetOperation) {
+                List<SlotReference> regularChildOutputs = ((LogicalSetOperation) childPlan).getRegularChildOutput(0);
+                List<NamedExpression> currentOutputs = ((LogicalSetOperation) childPlan).getOutputs();
+                if (regularChildOutputs.size() != currentOutputs.size()) {
+                    // set current output size is different from currentOutputs, not excepted
+                    LOG.error("InferPlanOutputAlias infer regularChildOutputs is different from currentOutputs,"
+                            + "child plan is {}", ((LogicalSetOperation) childPlan).treeString());
+                    return true;
+                }
+                for (int index = 0; index < regularChildOutputs.size(); index++) {
+                    childToParentMapping.put(regularChildOutputs.get(index).getExprId(),
+                            currentOutputs.get(index).getExprId());
+                }
+            }
             for (Expression expression : ((Plan) childPlan).getExpressions()) {
                 if (!(expression instanceof Alias)) {
                     continue;
                 }
                 Alias projectItem = (Alias) expression;
                 ExprId exprId = projectItem.getExprId();
                 // Infer name when alias child is expression and alias's name is from child
-                if (currentOutputExprIdSet.contains(projectItem.getExprId())
-                        && projectItem.isNameFromChild()) {
+                if (contains(currentOutputExprIdSet, exprId, childToParentMapping)
+                        && (projectItem.isNameFromChild())) {
                     String inferredAliasName = projectItem.child().getExpressionName();
-                    ImmutableCollection<Integer> outputExprIndexes = currentExprIdAndIndexMap.get(exprId);
+                    ImmutableCollection<Integer> outputExprIndexes = getOutputSlotIndex(currentExprIdAndIndexMap,
+                            exprId, childToParentMapping);
                     // replace output name by inferred name
                     for (Integer index : outputExprIndexes) {
                         Slot slot = currentOutputs.get(index);
@@ -89,4 +113,50 @@ public List<NamedExpression> infer(Plan plan, ImmutableMultimap<ExprId, Integer>
         });
         return finalOutputs;
     }
+
+    /**
+     * Such as LogicalIntersect, targetExprIdSet is ['back'#41, 'we'#42], but child is ['back'#38, 'we'#40]
+     * should construct childToParentMapping {
+     *     'back'#38 : 'back'#41,
+     *     'we'#40 : 'we'#42
+     * }
+     * LogicalIntersect ( qualifier=DISTINCT,
+     * outputs=['back'#41, 'we'#42],
+     * regularChildrenOutputs=[['back'#38, 'we'#40],
+     * [col_char_25__undef_signed#39, col_varchar_25__undef_signed_not_null#27]] )
+     */
+    private static boolean contains(ImmutableSet<ExprId> targetExprIdSet, ExprId exprId,
+            Map<ExprId, ExprId> childToParentMapping) {
+        if (targetExprIdSet.contains(exprId)) {
+            return true;
+        }
+        return targetExprIdSet.contains(childToParentMapping.get(exprId));
+    }
+
+    /**
+     * Such as LogicalIntersect, currentExprIdAndIndexMap is
+     * {
+     *     'back'#38 : [0],
+     *     'we'#40 : [1]
+     * }
+     * LogicalIntersect ( qualifier=DISTINCT,
+     * outputs=['back'#41, 'we'#42],
+     * regularChildrenOutputs=[['back'#38, 'we'#40],
+     * [col_char_25__undef_signed#39, col_varchar_25__undef_signed_not_null#27]] )
+     * exprId is 'back'#38
+     * childToParentMapping {
+     *     'back'#38 : 'back'#41,
+     *     'we'#40 : 'we'#42
+     * }
+     */
+    private static ImmutableCollection<Integer> getOutputSlotIndex(
+            ImmutableMultimap<ExprId, Integer> currentExprIdAndIndexMap,
+            ExprId exprId,
+            Map<ExprId, ExprId> childToParentMapping) {
+        ImmutableCollection<Integer> indexes = currentExprIdAndIndexMap.get(exprId);
+        if (!indexes.isEmpty()) {
+            return indexes;
+        }
+        return currentExprIdAndIndexMap.get(childToParentMapping.get(exprId));
+    }
 }