| /* |
| * 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.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.util.mapping.Mappings; |
| |
| import com.google.common.collect.ImmutableList; |
| |
| import org.checkerframework.checker.nullness.qual.Nullable; |
| |
| import java.util.AbstractList; |
| import java.util.Arrays; |
| import java.util.HashMap; |
| import java.util.List; |
| import java.util.Map; |
| import java.util.function.Supplier; |
| |
| /** |
| * RelTraitSet represents an ordered set of {@link RelTrait}s. |
| */ |
| public final class RelTraitSet extends AbstractList<RelTrait> { |
| private static final RelTrait[] EMPTY_TRAITS = new RelTrait[0]; |
| |
| //~ Instance fields -------------------------------------------------------- |
| |
| private final Cache cache; |
| private final RelTrait[] traits; |
| private @Nullable String string; |
| /** Caches the hash code for the traits. */ |
| private int hash; // Default to 0 |
| |
| //~ Constructors ----------------------------------------------------------- |
| |
| /** |
| * Constructs a RelTraitSet with the given set of RelTraits. |
| * |
| * @param cache Trait set cache (and indirectly cluster) that this set |
| * belongs to |
| * @param traits Traits |
| */ |
| private RelTraitSet(Cache cache, RelTrait[] traits) { |
| // NOTE: We do not copy the array. It is important that the array is not |
| // shared. However, since this constructor is private, we assume that |
| // the caller has made a copy. |
| this.cache = cache; |
| this.traits = traits; |
| } |
| |
| //~ Methods ---------------------------------------------------------------- |
| |
| /** |
| * Creates an empty trait set. |
| * |
| * <p>It has a new cache, which will be shared by any trait set created from |
| * it. Thus each empty trait set is the start of a new ancestral line. |
| */ |
| public static RelTraitSet createEmpty() { |
| return new RelTraitSet(new Cache(), EMPTY_TRAITS); |
| } |
| |
| /** |
| * Retrieves a RelTrait from the set. |
| * |
| * @param index 0-based index into ordered RelTraitSet |
| * @return the RelTrait |
| * @throws ArrayIndexOutOfBoundsException if index greater than or equal to |
| * {@link #size()} or less than 0. |
| */ |
| public RelTrait getTrait(int index) { |
| return traits[index]; |
| } |
| |
| /** |
| * Retrieves a list of traits from the set. |
| * |
| * @param index 0-based index into ordered RelTraitSet |
| * @return the RelTrait |
| * @throws ArrayIndexOutOfBoundsException if index greater than or equal to |
| * {@link #size()} or less than 0. |
| */ |
| public <E extends RelMultipleTrait> List<E> getTraits(int index) { |
| final RelTrait trait = traits[index]; |
| if (trait instanceof RelCompositeTrait) { |
| //noinspection unchecked |
| return ((RelCompositeTrait<E>) trait).traitList(); |
| } else { |
| //noinspection unchecked |
| return ImmutableList.of((E) trait); |
| } |
| } |
| |
| @Override public RelTrait get(int index) { |
| return getTrait(index); |
| } |
| |
| /** |
| * Returns whether a given kind of trait is enabled. |
| */ |
| public <T extends RelTrait> boolean isEnabled(RelTraitDef<T> traitDef) { |
| return getTrait(traitDef) != null; |
| } |
| |
| /** |
| * Retrieves a RelTrait of the given type from the set. |
| * |
| * @param traitDef the type of RelTrait to retrieve |
| * @return the RelTrait, or null if not found |
| */ |
| public <T extends RelTrait> @Nullable T getTrait(RelTraitDef<T> traitDef) { |
| int index = findIndex(traitDef); |
| if (index >= 0) { |
| //noinspection unchecked |
| return (T) getTrait(index); |
| } |
| |
| return null; |
| } |
| |
| /** |
| * Retrieves a list of traits of the given type from the set. |
| * |
| * <p>Only valid for traits that support multiple entries. (E.g. collation.) |
| * |
| * @param traitDef the type of RelTrait to retrieve |
| * @return the RelTrait, or null if not found |
| */ |
| public <T extends RelMultipleTrait> @Nullable List<T> getTraits( |
| RelTraitDef<T> traitDef) { |
| int index = findIndex(traitDef); |
| if (index >= 0) { |
| //noinspection unchecked |
| return (List<T>) getTraits(index); |
| } |
| |
| return null; |
| } |
| |
| /** |
| * Replaces an existing RelTrait in the set. |
| * Returns a different trait set; does not modify this trait set. |
| * |
| * @param index 0-based index into ordered RelTraitSet |
| * @param trait the new RelTrait |
| * @return the old RelTrait at the index |
| */ |
| public RelTraitSet replace(int index, RelTrait trait) { |
| assert traits[index].getTraitDef() == trait.getTraitDef() |
| : "RelTrait has different RelTraitDef than replacement"; |
| |
| RelTrait canonizedTrait = canonize(trait); |
| if (traits[index] == canonizedTrait) { |
| return this; |
| } |
| RelTrait[] newTraits = traits.clone(); |
| newTraits[index] = canonizedTrait; |
| return cache.getOrAdd(new RelTraitSet(cache, newTraits)); |
| } |
| |
| /** |
| * Returns a trait set consisting of the current set plus a new trait. |
| * |
| * <p>If the set does not contain a trait of the same {@link RelTraitDef}, |
| * the trait is ignored, and this trait set is returned. |
| * |
| * @param trait the new trait |
| * @return New set |
| * @see #plus(RelTrait) |
| */ |
| public RelTraitSet replace( |
| RelTrait trait) { |
| // Quick check for common case |
| if (containsShallow(traits, trait)) { |
| return this; |
| } |
| final RelTraitDef traitDef = trait.getTraitDef(); |
| int index = findIndex(traitDef); |
| if (index < 0) { |
| // Trait is not present. Ignore it. |
| return this; |
| } |
| |
| return replace(index, trait); |
| } |
| |
| /** Returns whether an element occurs within an array. |
| * |
| * <p>Uses {@code ==}, not {@link #equals}. Nulls are allowed. */ |
| private static <T> boolean containsShallow(T[] ts, RelTrait seek) { |
| for (T t : ts) { |
| if (t == seek) { |
| return true; |
| } |
| } |
| return false; |
| } |
| |
| /** Replaces the trait(s) of a given type with a list of traits of the same |
| * type. |
| * |
| * <p>The list must not be empty, and all traits must be of the same type. |
| */ |
| public <T extends RelMultipleTrait> RelTraitSet replace(List<T> traits) { |
| assert !traits.isEmpty(); |
| final RelTraitDef def = traits.get(0).getTraitDef(); |
| return replace(RelCompositeTrait.of(def, traits)); |
| } |
| |
| /** Replaces the trait(s) of a given type with a list of traits of the same |
| * type. |
| * |
| * <p>The list must not be empty, and all traits must be of the same type. |
| */ |
| public <T extends RelMultipleTrait> RelTraitSet replace(RelTraitDef<T> def, |
| List<T> traits) { |
| return replace(RelCompositeTrait.of(def, traits)); |
| } |
| |
| /** If a given multiple trait is enabled, replaces it by calling the given |
| * function. */ |
| public <T extends RelMultipleTrait> RelTraitSet replaceIfs(RelTraitDef<T> def, |
| Supplier<? extends @Nullable List<T>> traitSupplier) { |
| int index = findIndex(def); |
| if (index < 0) { |
| return this; // trait is not enabled; ignore it |
| } |
| final List<T> traitList = traitSupplier.get(); |
| if (traitList == null) { |
| return replace(index, def.getDefault()); |
| } |
| return replace(index, RelCompositeTrait.of(def, traitList)); |
| } |
| |
| /** If a given trait is enabled, replaces it by calling the given function. */ |
| public <T extends RelTrait> RelTraitSet replaceIf(RelTraitDef<T> def, |
| Supplier<? extends @Nullable T> traitSupplier) { |
| int index = findIndex(def); |
| if (index < 0) { |
| return this; // trait is not enabled; ignore it |
| } |
| T traitList = traitSupplier.get(); |
| if (traitList == null) { |
| traitList = def.getDefault(); |
| } |
| return replace(index, traitList); |
| } |
| |
| /** |
| * Applies a mapping to this traitSet. |
| * |
| * @param mapping Mapping |
| * @return traitSet with mapping applied |
| */ |
| public RelTraitSet apply(Mappings.TargetMapping mapping) { |
| RelTrait[] newTraits = new RelTrait[traits.length]; |
| for (int i = 0; i < traits.length; i++) { |
| newTraits[i] = traits[i].apply(mapping); |
| } |
| return cache.getOrAdd(new RelTraitSet(cache, newTraits)); |
| } |
| |
| /** |
| * 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; |
| } |
| if (this.size() != other.size()) { |
| return false; |
| } |
| for (int i = 0; i < traits.length; i++) { |
| if (traits[i].getTraitDef() == ConventionTraitDef.INSTANCE) { |
| continue; |
| } |
| // each trait should be canonized already |
| if (traits[i] != 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 @Nullable 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. |
| */ |
| @SuppressWarnings("unchecked") |
| public <T extends RelDistribution> @Nullable T getDistribution() { |
| return (@Nullable 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. |
| */ |
| @SuppressWarnings("unchecked") |
| public <T extends RelCollation> @Nullable T getCollation() { |
| return (@Nullable T) getTrait(RelCollationTraitDef.INSTANCE); |
| } |
| |
| /** |
| * Returns the size of the RelTraitSet. |
| * |
| * @return the size of the RelTraitSet. |
| */ |
| @Override public int size() { |
| return traits.length; |
| } |
| |
| /** |
| * Converts a trait to canonical form. |
| * |
| * <p>After canonization, t1.equals(t2) if and only if t1 == t2. |
| * |
| * @param trait Trait |
| * @return Trait in canonical form |
| */ |
| public <T extends RelTrait> T canonize(T trait) { |
| if (trait == null) { |
| // Return "trait" makes the input type to be the same as the output type, |
| // so checkerframework is happy |
| return trait; |
| } |
| |
| if (trait instanceof RelCompositeTrait) { |
| // Composite traits are canonized on creation |
| //noinspection unchecked |
| return trait; |
| } |
| |
| //noinspection unchecked |
| return (T) trait.getTraitDef().canonize(trait); |
| } |
| |
| /** |
| * Compares two RelTraitSet objects for equality. |
| * |
| * @param obj another RelTraitSet |
| * @return true if traits are equal and in the same order, false otherwise |
| */ |
| @Override public boolean equals(@Nullable Object obj) { |
| if (this == obj) { |
| return true; |
| } |
| if (!(obj instanceof RelTraitSet)) { |
| return false; |
| } |
| RelTraitSet that = (RelTraitSet) obj; |
| if (this.hash != 0 |
| && that.hash != 0 |
| && this.hash != that.hash) { |
| return false; |
| } |
| if (traits.length != that.traits.length) { |
| return false; |
| } |
| for (int i = 0; i < traits.length; i++) { |
| if (traits[i] != that.traits[i]) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| @Override public int hashCode() { |
| if (hash == 0) { |
| hash = Arrays.hashCode(traits); |
| } |
| return hash; |
| } |
| |
| /** |
| * Returns whether this trait set satisfies another trait set. |
| * |
| * <p>For that to happen, each trait satisfies the corresponding trait in the |
| * other set. In particular, each trait set satisfies itself, because each |
| * trait subsumes itself. |
| * |
| * <p>Intuitively, if a relational expression is needed that has trait set |
| * S (A, B), and trait set S1 (A1, B1) subsumes S, then any relational |
| * expression R in S1 meets that need. |
| * |
| * <p>For example, if we need a relational expression that has |
| * trait set S = {enumerable convention, sorted on [C1 asc]}, and R |
| * has {enumerable convention, sorted on [C3], [C1, C2]}. R has two |
| * sort keys, but one them [C1, C2] satisfies S [C1], and that is enough. |
| * |
| * @param that another RelTraitSet |
| * @return whether this trait set satisfies other trait set |
| * |
| * @see org.apache.calcite.plan.RelTrait#satisfies(RelTrait) |
| */ |
| public boolean satisfies(RelTraitSet that) { |
| if (this == that) { |
| return true; |
| } |
| final int n = |
| Math.min( |
| this.size(), |
| that.size()); |
| for (int i = 0; i < n; i++) { |
| RelTrait thisTrait = this.traits[i]; |
| RelTrait thatTrait = that.traits[i]; |
| if (!thisTrait.satisfies(thatTrait)) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| /** |
| * Compares two RelTraitSet objects to see if they match for the purposes of |
| * firing a rule. A null RelTrait within a RelTraitSet indicates a wildcard: |
| * any RelTrait in the other RelTraitSet will match. If one RelTraitSet is |
| * smaller than the other, comparison stops when the last RelTrait from the |
| * smaller set has been examined and the remaining RelTraits in the larger |
| * set are assumed to match. |
| * |
| * @param that another RelTraitSet |
| * @return true if the RelTraitSets match, false otherwise |
| */ |
| public boolean matches(RelTraitSet that) { |
| final int n = |
| Math.min( |
| this.size(), |
| that.size()); |
| |
| for (int i = 0; i < n; i++) { |
| RelTrait thisTrait = this.traits[i]; |
| RelTrait thatTrait = that.traits[i]; |
| |
| if ((thisTrait == null) || (thatTrait == null)) { |
| continue; |
| } |
| |
| if (thisTrait != thatTrait) { |
| return false; |
| } |
| } |
| |
| return true; |
| } |
| |
| /** |
| * Returns whether this trait set contains a given trait. |
| * |
| * @param trait Sought trait |
| * @return Whether set contains given trait |
| */ |
| public boolean contains(RelTrait trait) { |
| for (RelTrait relTrait : traits) { |
| if (trait == relTrait) { |
| return true; |
| } |
| } |
| return false; |
| } |
| |
| /** |
| * Returns whether this trait set contains the given trait, or whether the |
| * trait is not present because its {@link RelTraitDef} is not enabled. |
| * Returns false if another trait of the same {@code RelTraitDef} is |
| * present. |
| * |
| * @param trait Trait |
| * @return Whether trait is present, or is absent because disabled |
| */ |
| 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; |
| } |
| |
| /** |
| * Returns whether this trait set comprises precisely the list of given |
| * traits. |
| * |
| * @param relTraits Traits |
| * @return Whether this trait set's traits are the same as the argument |
| */ |
| public boolean comprises(RelTrait... relTraits) { |
| return Arrays.equals(traits, relTraits); |
| } |
| |
| @Override public String toString() { |
| if (string == null) { |
| string = computeString(); |
| } |
| return string; |
| } |
| |
| /** |
| * Outputs the traits of this set as a String. Traits are output in order, |
| * separated by periods. |
| */ |
| String computeString() { |
| StringBuilder s = new StringBuilder(); |
| for (int i = 0; i < traits.length; i++) { |
| final RelTrait trait = traits[i]; |
| if (i > 0) { |
| s.append('.'); |
| } |
| if ((trait == null) |
| && (traits.length == 1)) { |
| // Special format for a list containing a single null trait; |
| // otherwise its string appears as "null", which is the same |
| // as if the whole trait set were null, and so confusing. |
| s.append("{null}"); |
| } else { |
| s.append(trait); |
| } |
| } |
| return s.toString(); |
| } |
| |
| /** |
| * Finds the index of a trait of a given type in this set. |
| * |
| * @param traitDef Sought trait definition |
| * @return index of trait, or -1 if not found |
| */ |
| private int findIndex(RelTraitDef traitDef) { |
| for (int i = 0; i < traits.length; i++) { |
| RelTrait trait = traits[i]; |
| if ((trait != null) && (trait.getTraitDef() == traitDef)) { |
| return i; |
| } |
| } |
| |
| return -1; |
| } |
| |
| /** |
| * Returns this trait set with a given trait added or overridden. Does not |
| * modify this trait set. |
| * |
| * @param trait Trait |
| * @return Trait set with given trait |
| */ |
| public RelTraitSet plus(RelTrait trait) { |
| if (contains(trait)) { |
| return this; |
| } |
| int i = findIndex(trait.getTraitDef()); |
| if (i >= 0) { |
| return replace(i, trait); |
| } |
| final RelTrait canonizedTrait = canonize(trait); |
| assert canonizedTrait != null; |
| RelTrait[] newTraits = new RelTrait[traits.length + 1]; |
| System.arraycopy(traits, 0, newTraits, 0, traits.length); |
| newTraits[traits.length] = canonizedTrait; |
| return cache.getOrAdd(new RelTraitSet(cache, newTraits)); |
| } |
| |
| public RelTraitSet plusAll(RelTrait[] traits) { |
| RelTraitSet t = this; |
| for (RelTrait trait : traits) { |
| t = t.plus(trait); |
| } |
| return t; |
| } |
| |
| public RelTraitSet merge(RelTraitSet additionalTraits) { |
| return plusAll(additionalTraits.traits); |
| } |
| |
| /** Returns a list of traits that are in {@code traitSet} but not in this |
| * RelTraitSet. */ |
| public ImmutableList<RelTrait> difference(RelTraitSet traitSet) { |
| final ImmutableList.Builder<RelTrait> builder = ImmutableList.builder(); |
| final int n = |
| Math.min( |
| this.size(), |
| traitSet.size()); |
| |
| for (int i = 0; i < n; i++) { |
| RelTrait thisTrait = this.traits[i]; |
| RelTrait thatTrait = traitSet.traits[i]; |
| if (thisTrait != thatTrait) { |
| builder.add(thatTrait); |
| } |
| } |
| return builder.build(); |
| } |
| |
| /** Returns whether there are any composite traits in this set. */ |
| public boolean allSimple() { |
| for (RelTrait trait : traits) { |
| if (trait instanceof RelCompositeTrait) { |
| return false; |
| } |
| } |
| return true; |
| } |
| |
| /** Returns a trait set similar to this one but with all composite traits |
| * flattened. */ |
| public RelTraitSet simplify() { |
| RelTraitSet x = this; |
| for (int i = 0; i < traits.length; i++) { |
| final RelTrait trait = traits[i]; |
| if (trait instanceof RelCompositeTrait) { |
| x = x.replace(i, |
| ((RelCompositeTrait) trait).size() == 1 |
| ? ((RelCompositeTrait) trait).trait(0) |
| : trait.getTraitDef().getDefault()); |
| } |
| } |
| return x; |
| } |
| |
| /** Cache of trait sets. */ |
| private static class Cache { |
| final Map<RelTraitSet, RelTraitSet> map = new HashMap<>(); |
| |
| Cache() { |
| } |
| |
| RelTraitSet getOrAdd(RelTraitSet t) { |
| RelTraitSet exist = map.putIfAbsent(t, t); |
| return exist == null ? t : exist; |
| } |
| } |
| } |