diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/UnionDataTraitUtils.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/UnionDataTraitUtils.java new file mode 100644 index 00000000000000..0f00bde0e2aea9 --- /dev/null +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/properties/UnionDataTraitUtils.java @@ -0,0 +1,288 @@ +// 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.EqualTo; +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.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.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 com.google.common.base.Preconditions; + +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; + +/** Shared equal-set derivation for logical and physical union plans. */ +public final class UnionDataTraitUtils { + + private UnionDataTraitUtils() { + } + + /** Compute output equal pairs that hold for every row source of a union. */ + public static void computeEqualSet(Union union, Plan unionPlan, DataTrait.Builder builder) { + List outputs = unionPlan.getOutput(); + List children = unionPlan.children(); + List> childrenOutputs = union.getRegularChildrenOutputs(); + List> constantRows = union.getConstantExprsList(); + Preconditions.checkState(children.size() == childrenOutputs.size(), + "Union child count %s does not match regular child output mapping count %s", + children.size(), childrenOutputs.size()); + for (int childIndex = 0; childIndex < childrenOutputs.size(); childIndex++) { + List childOutputs = childrenOutputs.get(childIndex); + Preconditions.checkState(childOutputs.size() == outputs.size(), + "Union child output mapping at index %s has width %s, expected %s", + childIndex, childOutputs.size(), outputs.size()); + } + for (int rowIndex = 0; rowIndex < constantRows.size(); rowIndex++) { + List row = constantRows.get(rowIndex); + Preconditions.checkState(row.size() == outputs.size(), + "Union constant row at index %s has width %s, expected %s", + rowIndex, row.size(), outputs.size()); + } + if (outputs.size() < 2 || (children.isEmpty() && constantRows.isEmpty())) { + return; + } + + List> equalGroups = children.isEmpty() + ? oneGroupForAllOutputs(outputs.size()) + : intersectChildEqualGroups(children, childrenOutputs, outputs.size()); + + if (!constantRows.isEmpty() && !equalGroups.isEmpty()) { + Optional context = createRewriteContext(unionPlan); + for (List row : constantRows) { + equalGroups = refineByConstantRow(equalGroups, row, context, outputs.size()); + if (equalGroups.isEmpty()) { + return; + } + } + } + + for (List 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))); + } + } + } + + private static List> intersectChildEqualGroups(List children, + List> childrenOutputs, int outputSize) { + List> classIdsByChild = new ArrayList<>(children.size()); + for (int childIndex = 0; childIndex < children.size(); childIndex++) { + classIdsByChild.add(equalClassIds(children.get(childIndex), childrenOutputs.get(childIndex))); + } + + Map, List> ordinalsBySignature = new LinkedHashMap<>(); + for (int outputIndex = 0; outputIndex < outputSize; outputIndex++) { + List signature = new ArrayList<>(children.size()); + for (List childClassIds : classIdsByChild) { + signature.add(childClassIds.get(outputIndex)); + } + ordinalsBySignature.computeIfAbsent(signature, key -> new ArrayList<>()).add(outputIndex); + } + return onlyNonTrivialGroups(ordinalsBySignature.values()); + } + + private static List equalClassIds(Plan child, List childOutputs) { + DataTrait childTrait = child.getLogicalProperties().getTrait(); + Map classIdBySlot = new HashMap<>(); + int nextClassId = 0; + for (Set equalSet : childTrait.calAllEqualSet()) { + for (Slot slot : equalSet) { + classIdBySlot.put(slot, nextClassId); + } + nextClassId++; + } + + List 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; + } + + private static List> refineByConstantRow(List> equalGroups, + List row, Optional context, int outputSize) { + List> literals = new ArrayList<>(outputSize); + for (NamedExpression expression : row) { + literals.add(foldConstant(unwrapAlias(expression), context)); + } + + List> refinedGroups = new ArrayList<>(); + for (List equalGroup : equalGroups) { + Map> ordinalsByValue = new LinkedHashMap<>(); + for (int outputIndex : equalGroup) { + Optional key = literals.get(outputIndex).flatMap( + UnionDataTraitUtils::constantValueKey); + key.ifPresent(valueKey -> ordinalsByValue + .computeIfAbsent(valueKey, ignored -> new ArrayList<>()).add(outputIndex)); + } + for (List sameValueOrdinals : ordinalsByValue.values()) { + if (sameValueOrdinals.size() <= 1) { + continue; + } + int first = sameValueOrdinals.get(0); + boolean allProvenEqual = true; + for (int i = 1; i < sameValueOrdinals.size(); i++) { + if (!isEqualInConstantRow(row, first, sameValueOrdinals.get(i), context)) { + allProvenEqual = false; + break; + } + } + if (allProvenEqual) { + refinedGroups.add(sameValueOrdinals); + } + } + } + return refinedGroups; + } + + private static boolean isEqualInConstantRow(List row, int left, int right, + Optional context) { + try { + Expression leftExpression = unwrapAlias(row.get(left)); + Expression rightExpression = unwrapAlias(row.get(right)); + Expression equality = TypeCoercionUtils.processComparisonPredicate( + new EqualTo(leftExpression, rightExpression)); + Optional result = ExpressionUtils.checkConstantExpr(equality, context); + // NULL = NULL is UNKNOWN, so only a folded TRUE is a proof of equality. + return result.isPresent() && BooleanLiteral.TRUE.equals(result.get()); + } catch (RuntimeException e) { + // Unsupported coercion or an expression that cannot be folded is not proof. + return false; + } + } + + private static Optional foldConstant(Expression expression, + Optional context) { + try { + return ExpressionUtils.checkConstantExpr(expression, context); + } catch (RuntimeException e) { + return Optional.empty(); + } + } + + private static Optional constantValueKey(Literal literal) { + if (literal.isNullLiteral()) { + return Optional.empty(); + } + try { + if (literal instanceof NumericLiteral) { + BigDecimal value = ((NumericLiteral) literal).getBigDecimalValue().stripTrailingZeros(); + return Optional.of(new ConstantValueKey(NumericLiteral.class, value)); + } else if (literal instanceof StringLikeLiteral) { + return Optional.of(new ConstantValueKey(StringLikeLiteral.class, literal.getStringValue())); + } else if (literal instanceof DateLiteral) { + return Optional.of(new ConstantValueKey(DateLiteral.class, literal.getStringValue())); + } + return Optional.of(new ConstantValueKey(literal.getClass(), literal)); + } catch (RuntimeException e) { + return Optional.empty(); + } + } + + private static Expression unwrapAlias(NamedExpression expression) { + Expression unwrapped = expression; + while (unwrapped instanceof Alias) { + unwrapped = unwrapped.child(0); + } + return unwrapped; + } + + private static Optional createRewriteContext(Plan plan) { + ConnectContext connectContext = ConnectContext.get(); + if (connectContext == null || connectContext.getStatementContext() == null) { + return Optional.empty(); + } + return Optional.of(new ExpressionRewriteContext(plan, CascadesContext.initContext( + connectContext.getStatementContext(), plan, PhysicalProperties.ANY))); + } + + private static List> oneGroupForAllOutputs(int outputSize) { + List allOutputs = new ArrayList<>(outputSize); + for (int outputIndex = 0; outputIndex < outputSize; outputIndex++) { + allOutputs.add(outputIndex); + } + List> groups = new ArrayList<>(1); + groups.add(allOutputs); + return groups; + } + + private static List> onlyNonTrivialGroups(Iterable> groups) { + List> nonTrivialGroups = new ArrayList<>(); + for (List group : groups) { + if (group.size() > 1) { + nonTrivialGroups.add(group); + } + } + return nonTrivialGroups; + } + + private static final class ConstantValueKey { + private final Class family; + private final Object value; + + private ConstantValueKey(Class family, Object value) { + this.family = family; + this.value = value; + } + + @Override + public boolean equals(Object object) { + if (this == object) { + return true; + } + if (!(object instanceof ConstantValueKey)) { + return false; + } + ConstantValueKey that = (ConstantValueKey) object; + return family.equals(that.family) && value.equals(that.value); + } + + @Override + public int hashCode() { + return Objects.hash(family, value); + } + } +} diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalUnion.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalUnion.java index 9e276035787e3e..1fbecf167469f6 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalUnion.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/logical/LogicalUnion.java @@ -23,6 +23,7 @@ import org.apache.doris.nereids.properties.DataTrait; import org.apache.doris.nereids.properties.LogicalProperties; import org.apache.doris.nereids.properties.PhysicalProperties; +import org.apache.doris.nereids.properties.UnionDataTraitUtils; import org.apache.doris.nereids.rules.expression.ExpressionRewriteContext; import org.apache.doris.nereids.trees.expressions.Expression; import org.apache.doris.nereids.trees.expressions.NamedExpression; @@ -45,14 +46,9 @@ import com.google.common.collect.ImmutableSet; import com.google.common.collect.Lists; -import java.util.ArrayList; -import java.util.BitSet; -import java.util.HashMap; import java.util.List; -import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.Set; /** * Logical Union. @@ -290,75 +286,9 @@ public boolean hasUnboundExpression() { return super.hasUnboundExpression(); } - private List mapSlotToIndex(Plan plan, List> equalSlotsList) { - Map slotToIndex = new HashMap<>(); - for (int i = 0; i < plan.getOutput().size(); i++) { - slotToIndex.put(plan.getOutput().get(i), i); - } - List equalSlotIndicesList = new ArrayList<>(); - for (Set equalSlots : equalSlotsList) { - BitSet equalSlotIndices = new BitSet(); - for (Slot slot : equalSlots) { - if (slotToIndex.containsKey(slot)) { - equalSlotIndices.set(slotToIndex.get(slot)); - } - } - if (equalSlotIndices.cardinality() > 1) { - equalSlotIndicesList.add(equalSlotIndices); - } - } - return equalSlotIndicesList; - } - @Override public void computeEqualSet(DataTrait.Builder builder) { - if (children.isEmpty()) { - return; - } - - // Get the list of equal slot sets and their corresponding index mappings for the first child - List> childEqualSlotsList = child(0).getLogicalProperties() - .getTrait().calAllEqualSet(); - List childEqualSlotsIndicesList = mapSlotToIndex(child(0), childEqualSlotsList); - List unionEqualSlotIndicesList = new ArrayList<>(childEqualSlotsIndicesList); - - // Traverse all children and find the equal sets that exist in all children - for (int i = 1; i < children.size(); i++) { - Plan child = children.get(i); - - // Get the equal slot sets for the current child - childEqualSlotsList = child.getLogicalProperties().getTrait().calAllEqualSet(); - - // Map slots to indices for the current child - childEqualSlotsIndicesList = mapSlotToIndex(child, childEqualSlotsList); - - // Only keep the equal pairs that exist in all children of the union - // This is done by calculating the intersection of all children's equal slot indices - for (BitSet unionEqualSlotIndices : unionEqualSlotIndicesList) { - BitSet intersect = new BitSet(); - for (BitSet childEqualSlotIndices : childEqualSlotsIndicesList) { - if (unionEqualSlotIndices.intersects(childEqualSlotIndices)) { - intersect = childEqualSlotIndices; - break; - } - } - unionEqualSlotIndices.and(intersect); - } - } - - // Build the functional dependencies for the output slots - List outputList = getOutput(); - for (BitSet equalSlotIndices : unionEqualSlotIndicesList) { - if (equalSlotIndices.cardinality() <= 1) { - continue; - } - int first = equalSlotIndices.nextSetBit(0); - int next = equalSlotIndices.nextSetBit(first + 1); - while (next > 0) { - builder.addEqualPair(outputList.get(first), outputList.get(next)); - next = equalSlotIndices.nextSetBit(next + 1); - } - } + UnionDataTraitUtils.computeEqualSet(this, this, builder); } @Override diff --git a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalUnion.java b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalUnion.java index 24cad14f45a3f1..ee6f5a4d024696 100644 --- a/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalUnion.java +++ b/fe/fe-core/src/main/java/org/apache/doris/nereids/trees/plans/physical/PhysicalUnion.java @@ -21,9 +21,9 @@ import org.apache.doris.nereids.properties.DataTrait; import org.apache.doris.nereids.properties.LogicalProperties; import org.apache.doris.nereids.properties.PhysicalProperties; +import org.apache.doris.nereids.properties.UnionDataTraitUtils; 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.AbstractPlan; import org.apache.doris.nereids.trees.plans.Plan; @@ -38,14 +38,9 @@ import com.google.common.collect.ImmutableList; import com.google.common.collect.ImmutableSet; -import java.util.ArrayList; -import java.util.BitSet; -import java.util.HashMap; import java.util.List; -import java.util.Map; import java.util.Objects; import java.util.Optional; -import java.util.Set; import java.util.stream.Collectors; /** @@ -189,75 +184,9 @@ public void computeUniform(DataTrait.Builder builder) { // don't propagate uniform slots } - private List mapSlotToIndex(Plan plan, List> equalSlotsList) { - Map slotToIndex = new HashMap<>(); - for (int i = 0; i < plan.getOutput().size(); i++) { - slotToIndex.put(plan.getOutput().get(i), i); - } - List equalSlotIndicesList = new ArrayList<>(); - for (Set equalSlots : equalSlotsList) { - BitSet equalSlotIndices = new BitSet(); - for (Slot slot : equalSlots) { - if (slotToIndex.containsKey(slot)) { - equalSlotIndices.set(slotToIndex.get(slot)); - } - } - if (equalSlotIndices.cardinality() > 1) { - equalSlotIndicesList.add(equalSlotIndices); - } - } - return equalSlotIndicesList; - } - @Override public void computeEqualSet(DataTrait.Builder builder) { - if (children.isEmpty()) { - return; - } - - // Get the list of equal slot sets and their corresponding index mappings for the first child - List> childEqualSlotsList = child(0).getLogicalProperties() - .getTrait().calAllEqualSet(); - List childEqualSlotsIndicesList = mapSlotToIndex(child(0), childEqualSlotsList); - List unionEqualSlotIndicesList = new ArrayList<>(childEqualSlotsIndicesList); - - // Traverse all children and find the equal sets that exist in all children - for (int i = 1; i < children.size(); i++) { - Plan child = children.get(i); - - // Get the equal slot sets for the current child - childEqualSlotsList = child.getLogicalProperties().getTrait().calAllEqualSet(); - - // Map slots to indices for the current child - childEqualSlotsIndicesList = mapSlotToIndex(child, childEqualSlotsList); - - // Only keep the equal pairs that exist in all children of the union - // This is done by calculating the intersection of all children's equal slot indices - for (BitSet unionEqualSlotIndices : unionEqualSlotIndicesList) { - BitSet intersect = new BitSet(); - for (BitSet childEqualSlotIndices : childEqualSlotsIndicesList) { - if (unionEqualSlotIndices.intersects(childEqualSlotIndices)) { - intersect = childEqualSlotIndices; - break; - } - } - unionEqualSlotIndices.and(intersect); - } - } - - // Build the functional dependencies for the output slots - List outputList = getOutput(); - for (BitSet equalSlotIndices : unionEqualSlotIndicesList) { - if (equalSlotIndices.cardinality() <= 1) { - continue; - } - int first = equalSlotIndices.nextSetBit(0); - int next = equalSlotIndices.nextSetBit(first + 1); - while (next > 0) { - builder.addEqualPair(outputList.get(first), outputList.get(next)); - next = equalSlotIndices.nextSetBit(next + 1); - } - } + UnionDataTraitUtils.computeEqualSet(this, this, builder); } @Override diff --git a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/EqualSetTest.java b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/EqualSetTest.java index d33dd556a8c437..f2d3b60b2f4ea8 100644 --- a/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/EqualSetTest.java +++ b/fe/fe-core/src/test/java/org/apache/doris/nereids/properties/EqualSetTest.java @@ -17,12 +17,21 @@ package org.apache.doris.nereids.properties; +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.algebra.Union; +import org.apache.doris.nereids.trees.plans.logical.LogicalUnion; +import org.apache.doris.nereids.trees.plans.physical.PhysicalUnion; +import org.apache.doris.nereids.types.IntegerType; import org.apache.doris.nereids.util.PlanChecker; import org.apache.doris.utframe.TestWithFeService; +import com.google.common.collect.ImmutableList; import org.junit.jupiter.api.Assertions; import org.junit.jupiter.api.Test; +import org.mockito.Mockito; class EqualSetTest extends TestWithFeService { @Override @@ -98,6 +107,132 @@ void testSetOp() { .isEmpty()); } + @Test + void testUnionEqualSetUsesRegularChildOutputMapping() { + String sql = "select name, id, id2 from agg where id = id2 " + + "union all select name, id, id2 from agg where id = id2"; + LogicalUnion union = analyzeLogicalUnion(sql); + Assertions.assertNotEquals(union.child(0).getOutput(), union.getRegularChildOutput(0), + "the test must exercise an ordinal mapping that differs from child.getOutput()"); + assertUnionEqualPair(sql, union, 1, 2, true); + assertUnionEqualPair(sql, union, 0, 1, false); + } + + @Test + void testUnionEqualSetChecksEveryConstantRow() { + assertUnionEqualPair( + "select id, id2 from agg where id = id2 " + + "union all select 1, 1 union all select 2, 2", + 0, 1, true); + assertUnionEqualPair( + "select id, id2 from agg where id = id2 " + + "union all select 1, 1 union all select 2, 3", + 0, 1, false); + assertUnionEqualPair( + "select id, id2 from agg where id = id2 " + + "union all select 1, 2 union all select 2, 1", + 0, 1, false); + } + + @Test + void testConstantOnlyUnionEqualSetUsesSqlEquality() { + assertUnionEqualPair( + "select cast(1 as int), cast(1 as bigint) " + + "union all select cast(2 as int), cast(2 as bigint)", + 0, 1, true); + assertUnionEqualPair( + "select 1 + 1, cast(2 as bigint) " + + "union all select 2 * 2, cast(4 as bigint)", + 0, 1, true); + assertUnionEqualPair( + "select 1, 1 union all select 2, 3", + 0, 1, false); + assertUnionEqualPair( + "select cast(null as int), cast(null as bigint) union all select 1, 1", + 0, 1, false); + } + + @Test + void testMalformedUnionMappingsFailFast() { + SlotReference output0 = SlotReference.of("output0", IntegerType.INSTANCE); + SlotReference output1 = SlotReference.of("output1", IntegerType.INSTANCE); + SlotReference childOutput = SlotReference.of("childOutput", IntegerType.INSTANCE); + Union union = Mockito.mock(Union.class); + Plan unionPlan = Mockito.mock(Plan.class); + Plan child = Mockito.mock(Plan.class); + + Mockito.when(unionPlan.getOutput()).thenReturn(ImmutableList.of(output0, output1)); + Mockito.when(unionPlan.children()).thenReturn(ImmutableList.of(child)); + Mockito.when(union.getRegularChildrenOutputs()).thenReturn(ImmutableList.of()); + Mockito.when(union.getConstantExprsList()).thenReturn(ImmutableList.of()); + Assertions.assertThrows(IllegalStateException.class, + () -> UnionDataTraitUtils.computeEqualSet(union, unionPlan, new DataTrait.Builder())); + + Mockito.when(union.getRegularChildrenOutputs()) + .thenReturn(ImmutableList.of(ImmutableList.of(childOutput))); + Assertions.assertThrows(IllegalStateException.class, + () -> UnionDataTraitUtils.computeEqualSet(union, unionPlan, new DataTrait.Builder())); + + Mockito.when(unionPlan.children()).thenReturn(ImmutableList.of()); + Mockito.when(union.getRegularChildrenOutputs()).thenReturn(ImmutableList.of()); + Mockito.when(union.getConstantExprsList()) + .thenReturn(ImmutableList.of(ImmutableList.of(childOutput))); + Assertions.assertThrows(IllegalStateException.class, + () -> UnionDataTraitUtils.computeEqualSet(union, unionPlan, new DataTrait.Builder())); + } + + private void assertUnionEqualPair(String sql, int leftIndex, int rightIndex, boolean expected) { + assertUnionEqualPair(sql, analyzeLogicalUnion(sql), leftIndex, rightIndex, expected); + } + + private void assertUnionEqualPair(String sql, LogicalUnion logicalUnion, + int leftIndex, int rightIndex, boolean expected) { + Assertions.assertEquals(expected, logicalUnion.getLogicalProperties().getTrait().isNullSafeEqual( + logicalUnion.getOutput().get(leftIndex), logicalUnion.getOutput().get(rightIndex))); + + Plan physicalPlan = PlanChecker.from(connectContext).analyze(sql).rewrite().implement().getPhysicalPlan(); + PhysicalUnion physicalUnion = findPhysicalUnion(physicalPlan); + Assertions.assertNotNull(physicalUnion, "expected a PhysicalUnion in: " + physicalPlan.treeString()); + DataTrait.Builder builder = new DataTrait.Builder(); + physicalUnion.computeEqualSet(builder); + DataTrait physicalTrait = builder.build(); + Assertions.assertEquals(expected, physicalTrait.isNullSafeEqual( + physicalUnion.getOutput().get(leftIndex), physicalUnion.getOutput().get(rightIndex))); + } + + private LogicalUnion analyzeLogicalUnion(String sql) { + Plan rewritten = PlanChecker.from(connectContext).analyze(sql).rewrite().getPlan(); + LogicalUnion logicalUnion = findLogicalUnion(rewritten); + Assertions.assertNotNull(logicalUnion, "expected a LogicalUnion in: " + rewritten.treeString()); + return logicalUnion; + } + + private LogicalUnion findLogicalUnion(Plan plan) { + if (plan instanceof LogicalUnion) { + return (LogicalUnion) plan; + } + for (Plan child : plan.children()) { + LogicalUnion union = findLogicalUnion(child); + if (union != null) { + return union; + } + } + return null; + } + + private PhysicalUnion findPhysicalUnion(Plan plan) { + if (plan instanceof PhysicalUnion) { + return (PhysicalUnion) plan; + } + for (Plan child : plan.children()) { + PhysicalUnion union = findPhysicalUnion(child); + if (union != null) { + return union; + } + } + return null; + } + @Test void testFilterHaving() { Plan plan = PlanChecker.from(connectContext) diff --git a/regression-test/data/nereids_rules_p0/union_equal_set/union_equal_set.out b/regression-test/data/nereids_rules_p0/union_equal_set/union_equal_set.out new file mode 100644 index 00000000000000..64ac6a00a6ba53 --- /dev/null +++ b/regression-test/data/nereids_rules_p0/union_equal_set/union_equal_set.out @@ -0,0 +1,28 @@ +-- This file is automatically generated. You should know what you did if you want to edit this +-- !reordered_regular_child_mapping -- +10 1 1 1 +10 1 1 1 +10 2 2 3 +10 2 2 3 +20 1 1 5 +20 1 1 5 + +-- !mixed_regular_and_constant_rows -- +1 1 1 +1 2 2 +2 1 3 +2 2 4 + +-- !one_constant_row_breaks_equality -- +1 1 1 +1 1 1 +2 2 3 +2 3 4 + +-- !all_constant_rows_equal_after_coercion -- +1 1 1 +2 2 2 + +-- !constant_null_breaks_equality -- +\N \N 1 +1 1 2 diff --git a/regression-test/suites/nereids_rules_p0/union_equal_set/union_equal_set.groovy b/regression-test/suites/nereids_rules_p0/union_equal_set/union_equal_set.groovy new file mode 100644 index 00000000000000..b26d64e674f595 --- /dev/null +++ b/regression-test/suites/nereids_rules_p0/union_equal_set/union_equal_set.groovy @@ -0,0 +1,121 @@ +// 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. + +suite("union_equal_set") { + sql "drop table if exists union_equal_reordered_t" + sql "drop table if exists union_equal_constant_t" + + sql """ + create table union_equal_reordered_t ( + a int not null, + b int not null, + c int not null + ) duplicate key(a) + distributed by hash(a) buckets 1 + properties("replication_num" = "1") + """ + sql """ + insert into union_equal_reordered_t values + (1, 1, 10), + (2, 2, 10), + (1, 1, 20) + """ + + explain { + sql """ + select c, a, b, rank() over(order by c, a) rk + from ( + select c, a, b from union_equal_reordered_t where a = b + union all + select c, a, b from union_equal_reordered_t where a = b + ) u + """ + contains "functions: [rank()]" + contains "ASC NULLS FIRST, a[#" + } + + qt_reordered_regular_child_mapping """ + select c, a, b, rank() over(order by c, a) rk + from ( + select c, a, b from union_equal_reordered_t where a = b + union all + select c, a, b from union_equal_reordered_t where a = b + ) u + order by c, a, b, rk + """ + + sql """ + create table union_equal_constant_t ( + a int not null, + b int not null + ) duplicate key(a) + distributed by hash(a) buckets 1 + properties("replication_num" = "1") + """ + sql "insert into union_equal_constant_t values (1, 1), (2, 2)" + + explain { + sql """ + select a, b, rank() over(order by a, b) rk + from ( + select a, b from union_equal_constant_t where a = b + union all select 1, 2 + union all select 2, 1 + ) u + """ + contains "functions: [rank()]" + contains "ASC NULLS FIRST, b[#" + } + + qt_mixed_regular_and_constant_rows """ + select a, b, rank() over(order by a, b) rk + from ( + select a, b from union_equal_constant_t where a = b + union all select 1, 2 + union all select 2, 1 + ) u + order by a, b, rk + """ + + qt_one_constant_row_breaks_equality """ + select a, b, rank() over(order by a, b) rk + from ( + select a, b from union_equal_constant_t where a = b + union all select 1, 1 + union all select 2, 3 + ) u + order by a, b, rk + """ + + qt_all_constant_rows_equal_after_coercion """ + select a, b, rank() over(order by a, b) rk + from ( + select cast(1 as int) a, cast(1 as bigint) b + union all select cast(2 as int), cast(2 as bigint) + ) u + order by a, b, rk + """ + + qt_constant_null_breaks_equality """ + select a, b, rank() over(order by a, b) rk + from ( + select cast(null as int) a, cast(null as bigint) b + union all select cast(1 as int), cast(1 as bigint) + ) u + order by a, b, rk + """ +}