RelTraitSet.java
/*
* 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;
import static java.util.Objects.requireNonNull;
/**
* 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);
requireNonNull(canonizedTrait, "canonizedTrait");
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;
}
}
}