You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@calcite.apache.org by hy...@apache.org on 2020/05/11 03:43:50 UTC

[calcite] 02/03: [CALCITE-3983] Add utility methods to RelTraitSet

This is an automated email from the ASF dual-hosted git repository.

hyuan pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/calcite.git

commit ef28c9b8bfdd59f0a6ab0584ac07c9d362e2c67e
Author: Haisheng Yuan <h....@alibaba-inc.com>
AuthorDate: Sun May 10 08:40:27 2020 -0500

    [CALCITE-3983] Add utility methods to RelTraitSet
---
 .../java/org/apache/calcite/plan/RelTraitSet.java  | 114 +++++++++++++++++++++
 .../java/org/apache/calcite/plan/RelTraitTest.java |  37 ++++++-
 2 files changed, 150 insertions(+), 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 563aae3..09fb10a 100644
--- a/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java
+++ b/core/src/main/java/org/apache/calcite/plan/RelTraitSet.java
@@ -16,6 +16,10 @@
  */
 package org.apache.calcite.plan;
 
+import org.apache.calcite.rel.RelCollation;
+import org.apache.calcite.rel.RelCollationTraitDef;
+import org.apache.calcite.rel.RelDistribution;
+import org.apache.calcite.rel.RelDistributionTraitDef;
 import org.apache.calcite.runtime.FlatLists;
 import org.apache.calcite.util.mapping.Mappings;
 
@@ -265,6 +269,116 @@ public final class RelTraitSet extends AbstractList<RelTrait> {
   }
 
   /**
+   * Returns whether all the traits are default trait value.
+   */
+  public boolean isDefault() {
+    for (final RelTrait trait : traits) {
+      if (trait != trait.getTraitDef().getDefault()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /**
+   * Returns whether all the traits except {@link Convention}
+   * are default trait value.
+   */
+  public boolean isDefaultSansConvention() {
+    for (final RelTrait trait : traits) {
+      if (trait.getTraitDef() == ConventionTraitDef.INSTANCE) {
+        continue;
+      }
+      if (trait != trait.getTraitDef().getDefault()) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /**
+   * Returns whether all the traits except {@link Convention}
+   * equals with traits in {@code other} traitSet.
+   */
+  public boolean equalsSansConvention(RelTraitSet other) {
+    if (this == other) {
+      return true;
+    }
+    for (int i = 0; i < traits.length; i++) {
+      if (traits[i].getTraitDef() == ConventionTraitDef.INSTANCE) {
+        continue;
+      }
+      if (!traits[i].equals(other.traits[i])) {
+        return false;
+      }
+    }
+    return true;
+  }
+
+  /**
+   * Returns a new traitSet with same traitDefs with
+   * current traitSet, but each trait is the default
+   * trait value.
+   */
+  public RelTraitSet getDefault() {
+    RelTrait[] newTraits = new RelTrait[traits.length];
+    for (int i = 0; i < traits.length; i++) {
+      newTraits[i] = traits[i].getTraitDef().getDefault();
+    }
+    return cache.getOrAdd(new RelTraitSet(cache, newTraits));
+  }
+
+  /**
+   * Returns a new traitSet with same traitDefs with
+   * current traitSet, but each trait except {@link Convention}
+   * is the default trait value. {@link Convention} trait
+   * remains the same with current traitSet.
+   */
+  public RelTraitSet getDefaultSansConvention() {
+    RelTrait[] newTraits = new RelTrait[traits.length];
+    for (int i = 0; i < traits.length; i++) {
+      if (traits[i].getTraitDef() == ConventionTraitDef.INSTANCE) {
+        newTraits[i] = traits[i];
+      } else {
+        newTraits[i] = traits[i].getTraitDef().getDefault();
+      }
+    }
+    return cache.getOrAdd(new RelTraitSet(cache, newTraits));
+  }
+
+  /**
+   * Returns {@link Convention} trait defined by
+   * {@link ConventionTraitDef#INSTANCE}, or null if the
+   * {@link ConventionTraitDef#INSTANCE} is not registered
+   * in this traitSet.
+   */
+  public Convention getConvention() {
+    return getTrait(ConventionTraitDef.INSTANCE);
+  }
+
+  /**
+   * Returns {@link RelDistribution} trait defined by
+   * {@link RelDistributionTraitDef#INSTANCE}, or null if the
+   * {@link RelDistributionTraitDef#INSTANCE} is not registered
+   * in this traitSet.
+   */
+  public <T extends RelDistribution> T getDistribution() {
+    //noinspection unchecked
+    return (T) getTrait(RelDistributionTraitDef.INSTANCE);
+  }
+
+  /**
+   * Returns {@link RelCollation} trait defined by
+   * {@link RelCollationTraitDef#INSTANCE}, or null if the
+   * {@link RelCollationTraitDef#INSTANCE} is not registered
+   * in this traitSet.
+   */
+  public <T extends RelCollation> T getCollation() {
+    //noinspection unchecked
+    return (T) getTrait(RelCollationTraitDef.INSTANCE);
+  }
+
+  /**
    * Returns the size of the RelTraitSet.
    *
    * @return the size of the RelTraitSet.
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 de38704..af272ea 100644
--- a/core/src/test/java/org/apache/calcite/plan/RelTraitTest.java
+++ b/core/src/test/java/org/apache/calcite/plan/RelTraitTest.java
@@ -16,6 +16,7 @@
  */
 package org.apache.calcite.plan;
 
+import org.apache.calcite.adapter.enumerable.EnumerableConvention;
 import org.apache.calcite.rel.RelCollation;
 import org.apache.calcite.rel.RelCollationTraitDef;
 import org.apache.calcite.rel.RelCollations;
@@ -28,12 +29,15 @@ import java.util.List;
 import java.util.function.Supplier;
 
 import static org.junit.jupiter.api.Assertions.assertEquals;
+import static org.junit.jupiter.api.Assertions.assertFalse;
+import static org.junit.jupiter.api.Assertions.assertNotEquals;
+import static org.junit.jupiter.api.Assertions.assertTrue;
 
 import static java.lang.Integer.toHexString;
 import static java.lang.System.identityHashCode;
 
 /**
- * Test to verify {@link RelCompositeTrait}.
+ * Test to verify {@link RelCompositeTrait} and {@link RelTraitSet}.
  */
 class RelTraitTest {
   private static final RelCollationTraitDef COLLATION = RelCollationTraitDef.INSTANCE;
@@ -61,4 +65,35 @@ class RelTraitTest {
     assertCanonical("composite with two elements",
         () -> ImmutableList.of(RelCollations.of(0), RelCollations.of(1)));
   }
+
+  @Test void testTraitSetDefault() {
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    traits = traits.plus(Convention.NONE).plus(RelCollations.EMPTY);
+    assertEquals(traits.size(), 2);
+    assertTrue(traits.isDefault());
+    traits = traits.replace(EnumerableConvention.INSTANCE);
+    assertFalse(traits.isDefault());
+    assertTrue(traits.isDefaultSansConvention());
+    traits = traits.replace(RelCollations.of(0));
+    assertFalse(traits.isDefault());
+    assertFalse(traits.replace(Convention.NONE).isDefaultSansConvention());
+    assertTrue(traits.getDefault().isDefault());
+    traits = traits.getDefaultSansConvention();
+    assertFalse(traits.isDefault());
+    assertEquals(traits.getConvention(), EnumerableConvention.INSTANCE);
+    assertTrue(traits.isDefaultSansConvention());
+    assertEquals(traits.toString(), "ENUMERABLE.[]");
+  }
+
+  @Test void testTraitSetEqual() {
+    RelTraitSet traits = RelTraitSet.createEmpty();
+    RelTraitSet traits1 = traits.plus(Convention.NONE).plus(RelCollations.of(0));
+    assertEquals(traits1.size(), 2);
+    RelTraitSet traits2 = traits1.replace(EnumerableConvention.INSTANCE);
+    assertEquals(traits2.size(), 2);
+    assertNotEquals(traits1, traits2);
+    assertTrue(traits1.equalsSansConvention(traits2));
+    RelTraitSet traits3 = traits2.replace(RelCollations.of(1));
+    assertFalse(traits3.equalsSansConvention(traits2));
+  }
 }