xiedeyantu commented on code in PR #4827:
URL: https://github.com/apache/calcite/pull/4827#discussion_r2929601833
##########
core/src/main/java/org/apache/calcite/plan/RelTraitSet.java:
##########
@@ -375,37 +389,86 @@ public RelTraitSet getDefaultSansConvention() {
* {@link RelDistributionTraitDef#INSTANCE}, or null if the
* {@link RelDistributionTraitDef#INSTANCE} is not registered
* in this traitSet.
+ *
+ * <p>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 <T extends RelDistribution> @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) {
+ 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}.
+ *
+ * <p>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<RelDistribution> getDistributions() {
+ int index = findIndex(RelDistributionTraitDef.INSTANCE);
+ if (index < 0) {
+ return ImmutableList.of();
+ }
+ return (List<RelDistribution>) (List<?>) getTraits(index);
}
/**
* Returns {@link RelCollation} trait defined by
* {@link RelCollationTraitDef#INSTANCE}, or null if the
* {@link RelCollationTraitDef#INSTANCE} is not registered
* in this traitSet.
+ *
+ * <p>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 <T extends RelCollation> @Nullable T getCollation() {
- return (@Nullable T) getTrait(RelCollationTraitDef.INSTANCE);
+ int index = findIndex(RelCollationTraitDef.INSTANCE);
Review Comment:
Done!
--
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]
