From 48bcd13b1ffb10e2330f3cf5f68736bf417cfaa8 Mon Sep 17 00:00:00 2001 From: xiedeyantu Date: Sun, 8 Mar 2026 20:35:44 +0800 Subject: [PATCH 1/6] [CALCITE-7431] RelTraitSet#getTrait seems to mishandle RelCompositeTrait --- .../enumerable/EnumerableMergeUnion.java | 14 ++-- .../org/apache/calcite/plan/RelTraitSet.java | 82 ++++++++++++++++--- .../org/apache/calcite/plan/RelTraitTest.java | 35 ++++++++ 3 files changed, 113 insertions(+), 18 deletions(-) diff --git a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeUnion.java b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeUnion.java index d06fa31661a1..09d5a9e221e8 100644 --- a/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeUnion.java +++ b/core/src/main/java/org/apache/calcite/adapter/enumerable/EnumerableMergeUnion.java @@ -22,10 +22,8 @@ import org.apache.calcite.linq4j.tree.Expressions; import org.apache.calcite.linq4j.tree.ParameterExpression; import org.apache.calcite.plan.RelOptCluster; -import org.apache.calcite.plan.RelTrait; import org.apache.calcite.plan.RelTraitSet; import org.apache.calcite.rel.RelCollation; -import org.apache.calcite.rel.RelCollationTraitDef; import org.apache.calcite.rel.RelNode; import org.apache.calcite.util.BuiltInMethod; import org.apache.calcite.util.Pair; @@ -49,13 +47,17 @@ protected EnumerableMergeUnion(RelOptCluster cluster, RelTraitSet traitSet, throw new IllegalArgumentException("EnumerableMergeUnion with no collation"); } for (RelNode input : inputs) { - final RelTrait inputCollationTrait = - input.getTraitSet().getTrait(RelCollationTraitDef.INSTANCE); + // Use getCollations() rather than getTrait() so that we handle the case + // where the input's collation slot holds a RelCompositeTrait (multiple + // collations). For each required collation, at least one of the input's + // collations must satisfy it. + final List inputCollations = input.getTraitSet().getCollations(); for (RelCollation collation : collations) { - if (inputCollationTrait == null || !inputCollationTrait.satisfies(collation)) { + boolean satisfied = inputCollations.stream().anyMatch(ic -> ic.satisfies(collation)); + if (!satisfied) { throw new IllegalArgumentException("EnumerableMergeUnion input does " + "not satisfy collation. EnumerableMergeUnion collation: " - + collation + ". Input collation: " + inputCollationTrait + ". Input: " + + collation + ". Input collations: " + inputCollations + ". Input: " + input); } } diff --git a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java index 65747426062c..e13a9fa330f4 100644 --- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java +++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java @@ -117,20 +117,34 @@ public List getTraits(int index) { * Returns whether a given kind of trait is enabled. */ public boolean isEnabled(RelTraitDef traitDef) { - return getTrait(traitDef) != null; + return findIndex(traitDef) >= 0; } /** * Retrieves a RelTrait of the given type from the set. * + *

If this trait def supports multiple values (i.e. its trait implements + * {@link RelMultipleTrait}), the underlying slot may contain a + * {@link RelCompositeTrait} when more than one value is present. In that + * case this method throws {@link IllegalStateException}; use + * {@link #getTraits(RelTraitDef)} instead. + * * @param traitDef the type of RelTrait to retrieve * @return the RelTrait, or null if not found + * @throws IllegalStateException if the slot holds a composite (multiple) + * trait; use {@link #getTraits(RelTraitDef)} in that case */ public @Nullable T getTrait(RelTraitDef traitDef) { int index = findIndex(traitDef); if (index >= 0) { + final RelTrait trait = getTrait(index); + if (trait instanceof RelCompositeTrait) { + throw new IllegalStateException("TraitDef " + traitDef + + " has multiple values in this trait set; " + + "use getTraits(RelTraitDef) instead of getTrait(RelTraitDef)"); + } //noinspection unchecked - return (T) getTrait(index); + return (T) trait; } return null; @@ -375,10 +389,23 @@ public RelTraitSet getDefaultSansConvention() { * {@link RelDistributionTraitDef#INSTANCE}, or null if the * {@link RelDistributionTraitDef#INSTANCE} is not registered * in this traitSet. + * + *

If multiple distribution values are present (a composite trait), the + * first one is returned. Use {@link #getTraits(RelTraitDef)} with + * {@link RelDistributionTraitDef#INSTANCE} to retrieve all values. */ @SuppressWarnings("unchecked") public @Nullable T getDistribution() { - return (@Nullable T) getTrait(RelDistributionTraitDef.INSTANCE); + int index = findIndex(RelDistributionTraitDef.INSTANCE); + if (index < 0) { + return null; + } + final RelTrait trait = getTrait(index); + if (trait instanceof RelCompositeTrait) { + // Return the first distribution; caller can use getTraits() for all. + return (T) ((RelCompositeTrait) trait).trait(0); + } + return (T) trait; } /** @@ -386,26 +413,43 @@ public RelTraitSet getDefaultSansConvention() { * {@link RelCollationTraitDef#INSTANCE}, or null if the * {@link RelCollationTraitDef#INSTANCE} is not registered * in this traitSet. + * + *

If this trait set contains multiple collations (a composite trait), + * this method throws {@link IllegalStateException}. Use + * {@link #getCollations()} to handle both the single and multi-collation + * cases uniformly. */ @SuppressWarnings("unchecked") public @Nullable T getCollation() { - return (@Nullable T) getTrait(RelCollationTraitDef.INSTANCE); + int index = findIndex(RelCollationTraitDef.INSTANCE); + if (index < 0) { + return null; + } + final RelTrait trait = getTrait(index); + if (trait instanceof RelCompositeTrait) { + throw new IllegalStateException( + "This trait set contains multiple collations; " + + "use getCollations() instead of getCollation()"); + } + return (T) trait; } /** * Returns {@link RelCollation} traits defined by * {@link RelCollationTraitDef#INSTANCE}. + * + *

Returns an empty list when the trait def is not registered, a + * singleton list for the common single-collation case, and a list with + * more than one element when a {@link RelCompositeTrait} is present. */ @SuppressWarnings("unchecked") public List getCollations() { - RelCollation trait = getTrait(RelCollationTraitDef.INSTANCE); - if (trait == null) { + int index = findIndex(RelCollationTraitDef.INSTANCE); + if (index < 0) { return ImmutableList.of(); } - if (trait instanceof RelCompositeTrait) { - return ((RelCompositeTrait) trait).traitList(); - } - return ImmutableList.of(trait); + // getTraits(int) already unwraps RelCompositeTrait transparently. + return (List) (List) getTraits(index); } /** @@ -577,8 +621,22 @@ public boolean contains(RelTrait trait) { */ public boolean containsIfApplicable(RelTrait trait) { // Note that '==' is sufficient, because trait should be canonized. - final RelTrait trait1 = getTrait(trait.getTraitDef()); - return trait1 == null || trait1 == trait; + int index = findIndex(trait.getTraitDef()); + if (index < 0) { + // TraitDef not registered in this set → treat as "not applicable" → true + return true; + } + final RelTrait stored = getTrait(index); + if (stored instanceof RelCompositeTrait) { + // Any member of the composite matching the sought trait counts. + for (Object t : ((RelCompositeTrait) stored).traitList()) { + if (t == trait) { + return true; + } + } + return false; + } + return stored == trait; } /** diff --git a/core/src/test/java/org/apache/calcite/plan/RelTraitTest.java b/core/src/test/java/org/apache/calcite/plan/RelTraitTest.java index 6a9ddbaa49a5..2760b1e157e7 100644 --- a/core/src/test/java/org/apache/calcite/plan/RelTraitTest.java +++ b/core/src/test/java/org/apache/calcite/plan/RelTraitTest.java @@ -20,6 +20,10 @@ import org.apache.calcite.rel.RelCollation; import org.apache.calcite.rel.RelCollationTraitDef; import org.apache.calcite.rel.RelCollations; +import org.apache.calcite.rel.RelNode; +import org.apache.calcite.test.RelBuilderTest; +import org.apache.calcite.tools.FrameworkConfig; +import org.apache.calcite.tools.RelBuilder; import com.google.common.collect.ImmutableList; @@ -33,6 +37,7 @@ import static org.hamcrest.Matchers.hasSize; import static org.junit.jupiter.api.Assertions.assertFalse; import static org.junit.jupiter.api.Assertions.assertNotEquals; +import static org.junit.jupiter.api.Assertions.assertThrows; import static org.junit.jupiter.api.Assertions.assertTrue; import static java.lang.Integer.toHexString; @@ -98,4 +103,34 @@ private void assertCanonical(String message, Supplier> collat RelTraitSet traits3 = traits2.replace(RelCollations.of(1)); assertFalse(traits3.equalsSansConvention(traits2)); } + + /** Test for + * [CALCITE-7431] + * RelTraitSet#getTrait seems to mishandle RelCompositeTrait. */ + @Test void testRelCompositeTrait() { + // Build: EMP -> Sort(MGR asc) -> Project(MGR, MGR as MGR2) + // The project maps both output columns 0 and 1 back to input column 3 + // (MGR), so the planner derives two collations: [0 ASC] and [1 ASC], which + // are stored as a RelCompositeTrait in the output trait set. + final FrameworkConfig config = RelBuilderTest.config().build(); + final RelBuilder b = RelBuilder.create(config); + final RelNode in = b + .scan("EMP") + .sort(3) // MGR asc + .project(b.field(3), b.alias(b.field(3), "MGR2")) // MGR, MGR as MGR2 + .build(); + + final RelTraitSet traitSet = in.getTraitSet(); + + final List collations = traitSet.getCollations(); + assertTrue(collations.size() >= 2, + "getCollations() should expose all composite collations"); + + assertThrows(IllegalStateException.class, traitSet::getCollation, + "getCollation() should throw when a RelCompositeTrait is present"); + + assertThrows(IllegalStateException.class, + () -> traitSet.getTrait(RelCollationTraitDef.INSTANCE), + "getTrait() should throw when a RelCompositeTrait is present"); + } } From f6edd7217442ae4b298227fe656909a444cbabce Mon Sep 17 00:00:00 2001 From: xiedeyantu Date: Mon, 9 Mar 2026 15:54:30 +0800 Subject: [PATCH 2/6] Fix comments --- .../org/apache/calcite/plan/RelTraitSet.java | 29 +++++++++++++++---- 1 file changed, 24 insertions(+), 5 deletions(-) diff --git a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java index e13a9fa330f4..36d03f31665f 100644 --- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java +++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java @@ -390,9 +390,10 @@ public RelTraitSet getDefaultSansConvention() { * {@link RelDistributionTraitDef#INSTANCE} is not registered * in this traitSet. * - *

If multiple distribution values are present (a composite trait), the - * first one is returned. Use {@link #getTraits(RelTraitDef)} with - * {@link RelDistributionTraitDef#INSTANCE} to retrieve all values. + *

If this trait set contains multiple distributions (a composite trait), + * this method throws {@link IllegalStateException}. Use + * {@link #getDistributions()} to handle both the single and multi-distribution + * cases uniformly. */ @SuppressWarnings("unchecked") public @Nullable T getDistribution() { @@ -402,12 +403,30 @@ public RelTraitSet getDefaultSansConvention() { } final RelTrait trait = getTrait(index); if (trait instanceof RelCompositeTrait) { - // Return the first distribution; caller can use getTraits() for all. - return (T) ((RelCompositeTrait) trait).trait(0); + throw new IllegalStateException( + "This trait set contains multiple distributions; " + + "use getDistributions() instead of getDistribution()"); } return (T) trait; } + /** + * Returns {@link RelDistribution} traits defined by + * {@link RelDistributionTraitDef#INSTANCE}. + * + *

Returns an empty list when the trait def is not registered, a + * singleton list for the common single-distribution case, and a list with + * more than one element when a {@link RelCompositeTrait} is present. + */ + @SuppressWarnings("unchecked") + public List getDistributions() { + int index = findIndex(RelDistributionTraitDef.INSTANCE); + if (index < 0) { + return ImmutableList.of(); + } + return (List) (List) getTraits(index); + } + /** * Returns {@link RelCollation} trait defined by * {@link RelCollationTraitDef#INSTANCE}, or null if the From 2c62f7b32a399ec5f6be1e8e74461f86e27366fd Mon Sep 17 00:00:00 2001 From: xiedeyantu Date: Fri, 13 Mar 2026 15:36:08 +0800 Subject: [PATCH 3/6] Change to use getTrait --- .../org/apache/calcite/plan/RelTraitSet.java | 24 ++----------------- 1 file changed, 2 insertions(+), 22 deletions(-) diff --git a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java index 36d03f31665f..48161b8e9599 100644 --- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java +++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java @@ -397,17 +397,7 @@ public RelTraitSet getDefaultSansConvention() { */ @SuppressWarnings("unchecked") public @Nullable T getDistribution() { - int index = findIndex(RelDistributionTraitDef.INSTANCE); - if (index < 0) { - return null; - } - final RelTrait trait = getTrait(index); - if (trait instanceof RelCompositeTrait) { - throw new IllegalStateException( - "This trait set contains multiple distributions; " - + "use getDistributions() instead of getDistribution()"); - } - return (T) trait; + return (T) getTrait(RelDistributionTraitDef.INSTANCE); } /** @@ -440,17 +430,7 @@ public List getDistributions() { */ @SuppressWarnings("unchecked") public @Nullable T getCollation() { - int index = findIndex(RelCollationTraitDef.INSTANCE); - if (index < 0) { - return null; - } - final RelTrait trait = getTrait(index); - if (trait instanceof RelCompositeTrait) { - throw new IllegalStateException( - "This trait set contains multiple collations; " - + "use getCollations() instead of getCollation()"); - } - return (T) trait; + return (T) getTrait(RelCollationTraitDef.INSTANCE); } /** From ff5549dea45e6b00ae11f6a9b0c105cd6e53ee48 Mon Sep 17 00:00:00 2001 From: xiedeyantu Date: Fri, 13 Mar 2026 17:30:58 +0800 Subject: [PATCH 4/6] Fix getTrait(int index) --- .../org/apache/calcite/plan/RelTraitSet.java | 16 ++++++++-------- 1 file changed, 8 insertions(+), 8 deletions(-) diff --git a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java index 48161b8e9599..9091422c3c28 100644 --- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java +++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java @@ -87,7 +87,13 @@ public static RelTraitSet createEmpty() { * {@link #size()} or less than 0. */ public RelTrait getTrait(int index) { - return traits[index]; + final RelTrait trait = traits[index]; + if (trait instanceof RelCompositeTrait) { + throw new IllegalStateException("Trait index " + index + + " has multiple values in this trait set; " + + "use getTraits(RelTraitDef) instead of getTrait(RelTraitDef)"); + } + return trait; } /** @@ -137,14 +143,8 @@ public boolean isEnabled(RelTraitDef traitDef) { public @Nullable T getTrait(RelTraitDef traitDef) { int index = findIndex(traitDef); if (index >= 0) { - final RelTrait trait = getTrait(index); - if (trait instanceof RelCompositeTrait) { - throw new IllegalStateException("TraitDef " + traitDef - + " has multiple values in this trait set; " - + "use getTraits(RelTraitDef) instead of getTrait(RelTraitDef)"); - } //noinspection unchecked - return (T) trait; + return (T) getTrait(index); } return null; From f4839e3f497e74515d3e58766187aa08c1af5bdb Mon Sep 17 00:00:00 2001 From: xiedeyantu Date: Fri, 13 Mar 2026 20:00:51 +0800 Subject: [PATCH 5/6] Fix CI error --- core/src/main/java/org/apache/calcite/plan/RelTraitSet.java | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java index 9091422c3c28..3854b4bf8906 100644 --- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java +++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java @@ -116,7 +116,7 @@ public List getTraits(int index) { } @Override public RelTrait get(int index) { - return getTrait(index); + return traits[index]; } /** From a3981f1594e8006f5c7c4d742040a00fa9b90210 Mon Sep 17 00:00:00 2001 From: xiedeyantu Date: Fri, 13 Mar 2026 20:34:39 +0800 Subject: [PATCH 6/6] Fix CI error --- core/src/main/java/org/apache/calcite/plan/RelTraitSet.java | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java index 3854b4bf8906..79b4ebbc411a 100644 --- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java +++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java @@ -397,7 +397,7 @@ public RelTraitSet getDefaultSansConvention() { */ @SuppressWarnings("unchecked") public @Nullable T getDistribution() { - return (T) getTrait(RelDistributionTraitDef.INSTANCE); + return (@Nullable T) getTrait(RelDistributionTraitDef.INSTANCE); } /** @@ -430,7 +430,7 @@ public List getDistributions() { */ @SuppressWarnings("unchecked") public @Nullable T getCollation() { - return (T) getTrait(RelCollationTraitDef.INSTANCE); + return (@Nullable T) getTrait(RelCollationTraitDef.INSTANCE); } /**