github-actions[bot] commented on code in PR #67886: URL: https://github.com/apache/doris/pull/67886#discussion_r4036276929
########## fe/fe-core/src/main/java/org/apache/doris/nereids/properties/UnionDataTraitUtils.java: ########## @@ -0,0 +1,491 @@ +// 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.doris.nereids.properties; + +import org.apache.doris.nereids.CascadesContext; +import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext; +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.NullSafeEqual; +import org.apache.doris.nereids.trees.expressions.Slot; +import org.apache.doris.nereids.trees.expressions.SlotReference; +import org.apache.doris.nereids.trees.expressions.literal.BooleanLiteral; +import org.apache.doris.nereids.trees.expressions.literal.DateLiteral; +import org.apache.doris.nereids.trees.expressions.literal.Literal; +import org.apache.doris.nereids.trees.expressions.literal.NullLiteral; +import org.apache.doris.nereids.trees.expressions.literal.NumericLiteral; +import org.apache.doris.nereids.trees.expressions.literal.StringLikeLiteral; +import org.apache.doris.nereids.trees.plans.Plan; +import org.apache.doris.nereids.trees.plans.algebra.Union; +import org.apache.doris.nereids.util.ExpressionUtils; +import org.apache.doris.nereids.util.TypeCoercionUtils; +import org.apache.doris.qe.ConnectContext; + +import java.math.BigDecimal; +import java.util.ArrayList; +import java.util.HashMap; +import java.util.LinkedHashMap; +import java.util.List; +import java.util.Map; +import java.util.Objects; +import java.util.Optional; +import java.util.Set; + +/** + * Derives equal sets shared by logical and physical union plans. + * + * <p>Two union output slots are equal only when the corresponding values are equal in every regular + * child and every constant row. Regular children contribute equality information through their data + * traits, while constant rows contribute equality information through constant folding and + * null-safe equality semantics. + */ +public final class UnionDataTraitUtils { + + /** Utility class; it must not be instantiated. */ + private UnionDataTraitUtils() { + } + + /** + * Computes union output equalities that hold for every row source and adds them to {@code builder}. + * + * <p>An output ordinal identifies the same union column across {@code outputs}, every entry in + * {@code regularChildrenOutputs}, and every constant row. For regular children, this method keeps + * only groups of ordinals whose mapped child slots belong to the same equality class in every + * child. It then refines those groups with every constant row. A constant row retains a pair only + * when folding their expressions and evaluating their null-safe equality produces {@code TRUE}. + * + * <p>The union is assumed to satisfy the structural invariants established during analysis: each + * regular child has one output mapping and every regular or constant input has the union output + * width. This method does not validate those invariants. It leaves existing entries in + * {@code builder} unchanged and only adds equal pairs proven by all union inputs. + * + * @param union union metadata that supplies regular-child output mappings and constant rows + * @param unionPlan concrete logical or physical union plan that supplies children and output slots + * @param builder destination to which proven equal pairs between union output slots are added + */ + public static void computeEqualSet(Union union, Plan unionPlan, DataTrait.Builder builder) { + List<Slot> outputs = unionPlan.getOutput(); + List<Plan> children = unionPlan.children(); + List<List<SlotReference>> childrenOutputs = union.getRegularChildrenOutputs(); + List<List<NamedExpression>> constantRows = union.getConstantExprsList(); + if (outputs.size() < 2 || (children.isEmpty() && constantRows.isEmpty())) { + return; + } + + List<List<Integer>> equalGroups = children.isEmpty() + ? oneGroupForAllOutputs(outputs.size()) + : intersectChildEqualGroups(children, childrenOutputs, outputs.size()); + + if (!constantRows.isEmpty() && !equalGroups.isEmpty()) { + Optional<ExpressionRewriteContext> context = createRewriteContext(unionPlan); + for (List<NamedExpression> row : constantRows) { + equalGroups = refineByConstantRow(equalGroups, row, context, outputs.size()); + if (equalGroups.isEmpty()) { + return; + } + } + } + + for (List<Integer> equalGroup : equalGroups) { + int first = equalGroup.get(0); + for (int i = 1; i < equalGroup.size(); i++) { + builder.addEqualPair(outputs.get(first), outputs.get(equalGroup.get(i))); + } + } + } + + /** + * Intersects the equality partitions of all regular union children by union output ordinal. + * + * <p>For each output ordinal, this method builds a signature containing that ordinal's equality + * class ID in every child. Two output ordinals have the same signature exactly when their mapped + * child slots are equal in every regular child. Singleton signature groups are omitted because + * they do not describe an equality between different union outputs. + * + * @param children regular union children; child {@code i} corresponds to mapping {@code i} + * @param childrenOutputs mapped child slots indexed first by child and then by union output ordinal + * @param outputSize number of union output ordinals represented by every child mapping + * @return groups of at least two output ordinals that are equal in every regular child + */ + private static List<List<Integer>> intersectChildEqualGroups(List<Plan> children, + List<List<SlotReference>> childrenOutputs, int outputSize) { + List<List<Integer>> classIdsByChild = new ArrayList<>(children.size()); + for (int childIndex = 0; childIndex < children.size(); childIndex++) { + classIdsByChild.add(equalClassIds(children.get(childIndex), childrenOutputs.get(childIndex))); + } + + Map<List<Integer>, List<Integer>> ordinalsBySignature = new LinkedHashMap<>(); + for (int outputIndex = 0; outputIndex < outputSize; outputIndex++) { + List<Integer> signature = new ArrayList<>(children.size()); + for (List<Integer> childClassIds : classIdsByChild) { + signature.add(childClassIds.get(outputIndex)); + } + ordinalsBySignature.computeIfAbsent(signature, key -> new ArrayList<>()).add(outputIndex); + } + return onlyNonTrivialGroups(ordinalsBySignature.values()); + } + + /** + * Encodes one child's mapped output slots as equality-class IDs. + * + * <p>Slots in the same child data-trait equal set receive the same ID. A mapped slot not present in + * any equal set receives its own ID, so it cannot accidentally compare equal to a different slot. + * Repeated occurrences of the same mapped slot reuse the same ID. + * + * @param child child plan whose logical data trait defines slot equalities + * @param childOutputs child slots in union output-ordinal order + * @return class IDs in union output-ordinal order; equal IDs denote equal mapped child slots + */ + private static List<Integer> equalClassIds(Plan child, List<SlotReference> childOutputs) { + DataTrait childTrait = child.getLogicalProperties().getTrait(); + Map<Slot, Integer> classIdBySlot = new HashMap<>(); + int nextClassId = 0; + for (Set<Slot> equalSet : childTrait.calAllEqualSet()) { + for (Slot slot : equalSet) { + classIdBySlot.put(slot, nextClassId); + } + nextClassId++; + } + + List<Integer> classIds = new ArrayList<>(childOutputs.size()); + for (Slot childOutput : childOutputs) { + Integer classId = classIdBySlot.get(childOutput); + if (classId == null) { + classId = nextClassId++; + classIdBySlot.put(childOutput, classId); + } + classIds.add(classId); + } + return classIds; + } + + /** + * Refines candidate output equality groups using one constant row. + * + * <p>Each expression is first folded to a literal when possible. Candidate ordinals are bucketed by + * a normalized {@link ConstantValueKey} to avoid comparing values that clearly differ. Every pair + * in a multi-ordinal bucket is then checked independently, and proven pairs are merged into + * equality components. All typed NULL literals use one shared key, so compatible NULL expressions + * can reach that final proof without an incompatible pair discarding the entire bucket. An ordinal + * whose expression cannot be folded, cannot be normalized, or cannot be connected to another + * ordinal by a proven null-safe equality is omitted from the returned groups. + * + * @param equalGroups candidate output-ordinal groups proven equal by inputs processed so far + * @param row constant expressions in union output-ordinal order + * @param context optional rewrite context used while folding constants and comparisons + * @param outputSize number of union outputs, used to size the per-ordinal literal list + * @return non-singleton subgroups whose expressions are also proven equal in this constant row + */ + private static List<List<Integer>> refineByConstantRow(List<List<Integer>> equalGroups, + List<NamedExpression> row, Optional<ExpressionRewriteContext> context, int outputSize) { + List<Optional<Literal>> literals = new ArrayList<>(outputSize); + for (NamedExpression expression : row) { + literals.add(foldConstant(unwrapAlias(expression), context)); + } + + List<List<Integer>> refinedGroups = new ArrayList<>(); + for (List<Integer> equalGroup : equalGroups) { + Map<ConstantValueKey, List<Integer>> ordinalsByValue = new LinkedHashMap<>(); + for (int outputIndex : equalGroup) { + Optional<ConstantValueKey> key = literals.get(outputIndex).flatMap( + UnionDataTraitUtils::constantValueKey); + key.ifPresent(valueKey -> ordinalsByValue + .computeIfAbsent(valueKey, ignored -> new ArrayList<>()).add(outputIndex)); + } + for (List<Integer> sameValueOrdinals : ordinalsByValue.values()) { + if (sameValueOrdinals.size() <= 1) { + continue; + } + refinedGroups.addAll(splitByProvenEquality(sameValueOrdinals, row, context)); + } + } + return refinedGroups; + } + + /** + * Splits one normalized-value bucket into independently proven equality components. + * + * <p>Each bucket position starts in its own disjoint-set component. This method evaluates every + * unordered pair of output ordinals and merges their components only when null-safe comparison + * folds to {@code TRUE}. Evaluating all pairs makes the result independent of ordinal order and + * preserves a compatible subgroup even when another member, such as an ARRAY-typed NULL, cannot + * be coerced with it. Connected pairs may share a component because proven value equality is + * transitive; singleton components are omitted because they publish no output equality. + * + * @param sameValueOrdinals output ordinals that share one normalized constant-value key + * @param row constant expressions in union output-ordinal order + * @param context optional rewrite context used for coercion and constant evaluation + * @return non-singleton ordinal components connected by proven null-safe equality pairs + */ + private static List<List<Integer>> splitByProvenEquality(List<Integer> sameValueOrdinals, + List<NamedExpression> row, Optional<ExpressionRewriteContext> context) { + int[] parents = new int[sameValueOrdinals.size()]; + for (int i = 0; i < parents.length; i++) { + parents[i] = i; + } + + for (int left = 0; left < sameValueOrdinals.size(); left++) { Review Comment: [P2] Avoid quadratic folding in wide constant unions This loop performs coercion and constant folding before consulting the union-find state. For a legal two-arm, 1,000-column union where every expression is `CAST(NULL AS INT)`, the first 999 comparisons connect each row's bucket, but the method still evaluates the remaining 498,501 pairs; that is 999,000 `NullSafeEqual` coercion/fold calls across the two rows. Mixed NULL buckets with repeated incompatible types can likewise repeat the same failed type-pair proof quadratically. Because this work runs on the FE planning path, a moderately wide constant UNION can consume disproportionate planner CPU. Please skip pairs whose roots are already equal and reuse proof/failure for repeated proof-equivalent folded-value/type/expression signatures (or otherwise bound proof work near-linearly in ordinals plus distinct proof signatures) while preserving the current mixed-type correctness. -- This is an automated message from the Apache Git Service. To respond to the message, please log on to GitHub and use the URL above to go to the specific comment. To unsubscribe, e-mail: [email protected] For queries about this service, please contact Infrastructure at: [email protected] --------------------------------------------------------------------- To unsubscribe, e-mail: [email protected] For additional commands, e-mail: [email protected]
