org.apache.calcite.plan.RelTrait Java Examples

/** Ensures that the subset that is the root relational expression contains
 * converters to all other subsets in its equivalence set.
 *
 * <p>Thus the planner tries to find cheap implementations of those other
 * subsets, which can then be converted to the root. This is the only place
 * in the plan where explicit converters are required; elsewhere, a consumer
 * will be asking for the result in a particular convention, but the root has
 * no consumers. */
void ensureRootConverters() {
  final Set<RelSubset> subsets = new HashSet<>();
  for (RelNode rel : root.getRels()) {
    if (rel instanceof AbstractConverter) {
      subsets.add((RelSubset) ((AbstractConverter) rel).getInput());
    }
  }
  for (RelSubset subset : root.set.subsets) {
    final ImmutableList<RelTrait> difference =
        root.getTraitSet().difference(subset.getTraitSet());
    if (difference.size() == 1 && subsets.add(subset)) {
      register(
          new AbstractConverter(subset.getCluster(), subset,
              difference.get(0).getTraitDef(), root.getTraitSet()),
          root);
    }
  }
}
 

@Override
public boolean satisfies(RelTrait trait) {

  if (trait instanceof DistributionTrait) {
    DistributionType requiredDist = ((DistributionTrait) trait).getType();
    if (requiredDist == DistributionType.ANY) {
      return true;
    }

    if (this.type == DistributionType.HASH_DISTRIBUTED) {
      if (requiredDist == DistributionType.HASH_DISTRIBUTED) {
        // A subset of the required distribution columns can satisfy (subsume) the requirement
        // e.g: required distribution: {a, b, c}
        // Following can satisfy the requirements: {a}, {b}, {c}, {a, b}, {b, c}, {a, c} or {a, b, c}

        // New: Use equals for subsumes check of hash distribution. If we uses subsumes,
        // a join may end up with hash-distributions using different keys. This would
        // cause incorrect query result.
        return this.equals(trait);
      }
    }
  }

  return this.equals(trait);
}
 

public boolean removeRule(RelOptRule rule) {
  // Remove description.
  if (!super.removeRule(rule)) {
    // Rule was not present.
    return false;
  }

  // Remove operands.
  classOperands.values().removeIf(entry -> entry.getRule().equals(rule));

  // Remove trait mappings. (In particular, entries from conversion
  // graph.)
  if (rule instanceof ConverterRule) {
    ConverterRule converterRule = (ConverterRule) rule;
    final RelTrait ruleTrait = converterRule.getInTrait();
    final RelTraitDef ruleTraitDef = ruleTrait.getTraitDef();
    if (traitDefs.contains(ruleTraitDef)) {
      ruleTraitDef.deregisterConverterRule(this, converterRule);
    }
  }
  return true;
}
 

private boolean doesConverterApply(
    ConverterRule converterRule,
    HepRelVertex vertex) {
  RelTrait outTrait = converterRule.getOutTrait();
  List<HepRelVertex> parents = Graphs.predecessorListOf(graph, vertex);
  for (HepRelVertex parent : parents) {
    RelNode parentRel = parent.getCurrentRel();
    if (parentRel instanceof Converter) {
      // We don't support converter chains.
      continue;
    }
    if (parentRel.getTraitSet().contains(outTrait)) {
      // This parent wants the traits produced by the converter.
      return true;
    }
  }
  return (vertex == root)
      && (requestedRootTraits != null)
      && requestedRootTraits.contains(outTrait);
}
 

@Override
public void write(final Kryo kryo, final Output output, final RelTraitSet traitSet) {
  final int size = traitSet.size();
  kryo.writeObject(output, size);

  for (int i = 0; i < size; i++) {
    final RelTrait trait = traitSet.getTrait(i);
    kryo.writeClassAndObject(output, trait);
  }
  kryo.writeObject(output, traitSet.toString());

  try {
    final Object cache = getField(NAME_CACHE).get(traitSet);
    kryo.writeClassAndObject(output, cache);
  } catch (final NoSuchFieldException|IllegalAccessException e) {
    throw new RuntimeException("unable to read TraitSet", e);
  }
}
 

/**
 * Returns a guess as to which subset (that is equivalence class of
 * relational expressions combined with a set of physical traits) the result
 * of this rule will belong to.
 *
 * @return expected subset, or null if we cannot guess
 */
private RelSubset guessSubset() {
  if (targetSubset != null) {
    return targetSubset;
  }
  final RelTrait targetTrait = getRule().getOutTrait();
  if ((targetSet != null) && (targetTrait != null)) {
    final RelTraitSet targetTraitSet =
        rels[0].getTraitSet().replace(targetTrait);

    // Find the subset in the target set which matches the expected
    // set of traits. It may not exist yet.
    targetSubset = targetSet.getSubset(targetTraitSet);
    return targetSubset;
  }

  // The target subset doesn't exist yet.
  return null;
}
 

/** Ensures that the subset that is the root relational expression contains
 * converters to all other subsets in its equivalence set.
 *
 * <p>Thus the planner tries to find cheap implementations of those other
 * subsets, which can then be converted to the root. This is the only place
 * in the plan where explicit converters are required; elsewhere, a consumer
 * will be asking for the result in a particular convention, but the root has
 * no consumers. */
void ensureRootConverters() {
  final Set<RelSubset> subsets = new HashSet<>();
  for (RelNode rel : root.getRels()) {
    if (rel instanceof AbstractConverter) {
      subsets.add((RelSubset) ((AbstractConverter) rel).getInput());
    }
  }
  for (RelSubset subset : root.set.subsets) {
    final ImmutableList<RelTrait> difference =
        root.getTraitSet().difference(subset.getTraitSet());
    if (difference.size() == 1 && subsets.add(subset)) {
      register(
          new AbstractConverter(subset.getCluster(), subset,
              difference.get(0).getTraitDef(), root.getTraitSet()),
          root);
    }
  }
}
 

public boolean removeRule(RelOptRule rule) {
  if (!ruleSet.remove(rule)) {
    // Rule was not present.
    return false;
  }

  // Remove description.
  unmapRuleDescription(rule);

  // Remove operands.
  classOperands.values().removeIf(entry -> entry.getRule().equals(rule));

  // Remove trait mappings. (In particular, entries from conversion
  // graph.)
  if (rule instanceof ConverterRule) {
    ConverterRule converterRule = (ConverterRule) rule;
    final RelTrait ruleTrait = converterRule.getInTrait();
    final RelTraitDef ruleTraitDef = ruleTrait.getTraitDef();
    if (traitDefs.contains(ruleTraitDef)) {
      ruleTraitDef.deregisterConverterRule(this, converterRule);
    }
  }
  return true;
}
 

public boolean satisfies(RelTrait trait) {
  if (trait == this || trait == ANY) {
    return true;
  }
  if (trait instanceof RelDistributionImpl) {
    RelDistributionImpl distribution = (RelDistributionImpl) trait;
    if (type == distribution.type) {
      switch (type) {
      case HASH_DISTRIBUTED:
        // The "leading edge" property of Range does not apply to Hash.
        // Only Hash[x, y] satisfies Hash[x, y].
        return keys.equals(distribution.keys);
      case RANGE_DISTRIBUTED:
        // Range[x, y] satisfies Range[x, y, z] but not Range[x]
        return Util.startsWith(distribution.keys, keys);
      default:
        return true;
      }
    }
  }
  if (trait == RANDOM_DISTRIBUTED) {
    // RANDOM is satisfied by HASH, ROUND-ROBIN, RANDOM, RANGE;
    // we've already checked RANDOM
    return type == Type.HASH_DISTRIBUTED
        || type == Type.ROUND_ROBIN_DISTRIBUTED
        || type == Type.RANGE_DISTRIBUTED;
  }
  return false;
}
 

public static RelTraitSet removeCollation(RelTraitSet traitSet, RelOptRuleCall call) {
    RelTraitSet newTraitSet = call.getPlanner().emptyTraitSet();

    for (RelTrait trait : traitSet) {
        if (!trait.getTraitDef().getTraitClass().equals(RelCollation.class)) {
            newTraitSet = newTraitSet.plus(trait);
        }
    }

    return newTraitSet;
}
 

public RelWriter explainTerms(RelWriter pw) {
  super.explainTerms(pw);
  for (RelTrait trait : traitSet) {
    pw.item(trait.getTraitDef().getSimpleName(), trait);
  }
  return pw;
}
 

@SuppressWarnings("unchecked")
@Deprecated // to be removed before 2.0
JdbcConverterRule(Class<? extends RelNode> clazz, RelTrait in,
    JdbcConvention out, String description) {
  this(clazz, (Predicate<RelNode>) r -> true, in, out,
      RelFactories.LOGICAL_BUILDER, description);
}
 

@Override public RelWriter done(RelNode node) {
  StringBuilder sb = new StringBuilder();
  sb.append(node.getRelTypeName());

  for (RelTrait trait : node.getTraitSet()) {
    sb.append('.');
    sb.append(trait.toString());
  }

  sb.append('(');
  int j = 0;
  for (Pair<String, Object> value : values) {
    if (j++ > 0) {
      sb.append(',');
    }
    sb.append(value.left);
    sb.append('=');
    if (value.right instanceof RelNode) {
      RelNode input = (RelNode) value.right;
      sb.append(input.getRelTypeName());
      sb.append('#');
      sb.append(input.getId());
    } else {
      sb.append(value.right);
    }
  }
  sb.append(')');
  digest = sb.toString();
  return this;
}
 

private AltTraitConverter(
    RelOptCluster cluster,
    RelNode child,
    RelTrait toTrait) {
  super(
      cluster,
      toTrait.getTraitDef(),
      child.getTraitSet().replace(toTrait),
      child);

  this.toTrait = toTrait;
}
 

/**
 * Creates a <code>ConverterRule</code> with a predicate.
 *
 * @param clazz       Type of relational expression to consider converting
 * @param predicate   Predicate on the relational expression
 * @param in          Trait of relational expression to consider converting
 * @param out         Trait which is converted to
 * @param relBuilderFactory Builder for relational expressions
 * @param description Description of rule
 */
public <R extends RelNode> ConverterRule(Class<R> clazz,
    Predicate<? super R> predicate, RelTrait in, RelTrait out,
    RelBuilderFactory relBuilderFactory, String description) {
  super(convertOperand(clazz, predicate, in),
      relBuilderFactory,
      description == null
          ? "ConverterRule<in=" + in + ",out=" + out + ">"
          : description);
  this.inTrait = Objects.requireNonNull(in);
  this.outTrait = Objects.requireNonNull(out);

  // Source and target traits must have same type
  assert in.getTraitDef() == out.getTraitDef();
}
 

@SuppressWarnings("Guava")
@Deprecated // to be removed before 2.0
public <R extends RelNode> ConverterRule(Class<R> clazz,
    com.google.common.base.Predicate<? super R> predicate, RelTrait in,
    RelTrait out, RelBuilderFactory relBuilderFactory, String description) {
  this(clazz, (Predicate<? super R>) predicate::apply, in, out,
      relBuilderFactory, description);
}
 

@Override protected String computeDigest() {
  StringBuilder digest = new StringBuilder("Subset#");
  digest.append(set.id);
  for (RelTrait trait : traitSet) {
    digest.append('.').append(trait);
  }
  return digest.toString();
}
 

public RelNode convert(
    RelOptPlanner planner,
    RelNode rel,
    AltTrait toTrait,
    boolean allowInfiniteCostConverters) {
  RelTrait fromTrait = rel.getTraitSet().getTrait(this);

  if (conversionMap.containsKey(fromTrait)) {
    final RelMetadataQuery mq = rel.getCluster().getMetadataQuery();
    for (Pair<RelTrait, ConverterRule> traitAndRule
        : conversionMap.get(fromTrait)) {
      RelTrait trait = traitAndRule.left;
      ConverterRule rule = traitAndRule.right;

      if (trait == toTrait) {
        RelNode converted = rule.convert(rel);
        if ((converted != null)
            && (!planner.getCost(converted, mq).isInfinite()
            || allowInfiniteCostConverters)) {
          return converted;
        }
      }
    }
  }

  return null;
}
 

@Override
public boolean satisfies(RelTrait trait) {

  if (trait instanceof DrillDistributionTrait) {
    DistributionType requiredDist = ((DrillDistributionTrait) trait).getType();
    if (requiredDist == DistributionType.ANY) {
      return true;
    }

    if (this.type == DistributionType.HASH_DISTRIBUTED) {
      if (requiredDist == DistributionType.HASH_DISTRIBUTED) {
        // A subset of the required distribution columns can satisfy (subsume) the requirement
        // e.g: required distribution: {a, b, c}
        // Following can satisfy the requirements: {a}, {b}, {c}, {a, b}, {b, c}, {a, c} or {a, b, c}

        // New: Use equals for subsumes check of hash distribution. If we uses subsumes,
        // a join may end up with hash-distributions using different keys. This would
        // cause incorrect query result.
        return this.equals(trait);
      }
      else if (requiredDist == DistributionType.RANDOM_DISTRIBUTED) {
        return true; // hash distribution subsumes random distribution and ANY distribution
      }
    }

    if(this.type == DistributionType.RANGE_DISTRIBUTED) {
      if (requiredDist == DistributionType.RANDOM_DISTRIBUTED) {
        return true; // RANGE_DISTRIBUTED distribution subsumes random distribution and ANY distribution
      }
    }
  }

  return this.equals(trait);
}
 

public boolean addRule(RelOptRule rule) {
  if (locked) {
    return false;
  }

  if (!super.addRule(rule)) {
    return false;
  }

  // Each of this rule's operands is an 'entry point' for a rule call.
  // Register each operand against all concrete sub-classes that could match
  // it.
  for (RelOptRuleOperand operand : rule.getOperands()) {
    for (Class<? extends RelNode> subClass
        : subClasses(operand.getMatchedClass())) {
      if (PhysicalNode.class.isAssignableFrom(subClass)
          && rule instanceof TransformationRule) {
        continue;
      }
      classOperands.put(subClass, operand);
    }
  }

  // If this is a converter rule, check that it operates on one of the
  // kinds of trait we are interested in, and if so, register the rule
  // with the trait.
  if (rule instanceof ConverterRule) {
    ConverterRule converterRule = (ConverterRule) rule;

    final RelTrait ruleTrait = converterRule.getInTrait();
    final RelTraitDef ruleTraitDef = ruleTrait.getTraitDef();
    if (traitDefs.contains(ruleTraitDef)) {
      ruleTraitDef.registerConverterRule(this, converterRule);
    }
  }

  return true;
}
 

@SuppressWarnings({"Guava", "unchecked"})
@Deprecated // to be removed before 2.0
<R extends RelNode> JdbcConverterRule(Class<R> clazz,
    com.google.common.base.Predicate<? super R> predicate,
    RelTrait in, JdbcConvention out,
    RelBuilderFactory relBuilderFactory, String description) {
  this(clazz, (Predicate<R>) predicate, in, out, relBuilderFactory,
      description);
}
 

@SuppressWarnings("Guava")
@Deprecated // to be removed before 2.0
public <R extends RelNode> ConverterRule(Class<R> clazz,
    com.google.common.base.Predicate<? super R> predicate,
    RelTrait in, RelTrait out, String description) {
  this(clazz, predicate, in, out, RelFactories.LOGICAL_BUILDER, description);
}
 

public boolean satisfies(RelTrait trait) {
  if (trait == this || trait == ANY) {
    return true;
  }
  if (trait instanceof RelDistributionImpl) {
    RelDistributionImpl distribution = (RelDistributionImpl) trait;
    if (type == distribution.type) {
      switch (type) {
      case HASH_DISTRIBUTED:
        // The "leading edge" property of Range does not apply to Hash.
        // Only Hash[x, y] satisfies Hash[x, y].
        return keys.equals(distribution.keys);
      case RANGE_DISTRIBUTED:
        // Range[x, y] satisfies Range[x, y, z] but not Range[x]
        return Util.startsWith(distribution.keys, keys);
      default:
        return true;
      }
    }
  }
  if (trait == RANDOM_DISTRIBUTED) {
    // RANDOM is satisfied by HASH, ROUND-ROBIN, RANDOM, RANGE;
    // we've already checked RANDOM
    return type == Type.HASH_DISTRIBUTED
        || type == Type.ROUND_ROBIN_DISTRIBUTED
        || type == Type.RANGE_DISTRIBUTED;
  }
  return false;
}
 

public boolean canConvert(
    RelOptPlanner planner,
    AltTrait fromTrait,
    AltTrait toTrait) {
  if (conversionMap.containsKey(fromTrait)) {
    for (Pair<RelTrait, ConverterRule> traitAndRule
        : conversionMap.get(fromTrait)) {
      if (traitAndRule.left == toTrait) {
        return true;
      }
    }
  }

  return false;
}
 

@SuppressWarnings("Guava")
@Deprecated // to be removed before 2.0
public <R extends RelNode> ConverterRule(Class<R> clazz,
    com.google.common.base.Predicate<? super R> predicate, RelTrait in,
    RelTrait out, RelBuilderFactory relBuilderFactory, String description) {
  this(clazz, (Predicate<? super R>) predicate::apply, in, out,
      relBuilderFactory, description);
}
 

/**
 * Creates a <code>ConverterRule</code> with a predicate.
 *
 * @param clazz       Type of relational expression to consider converting
 * @param predicate   Predicate on the relational expression
 * @param in          Trait of relational expression to consider converting
 * @param out         Trait which is converted to
 * @param relBuilderFactory Builder for relational expressions
 * @param descriptionPrefix Description prefix of rule
 */
public <R extends RelNode> ConverterRule(Class<R> clazz,
    Predicate<? super R> predicate, RelTrait in, RelTrait out,
    RelBuilderFactory relBuilderFactory, String descriptionPrefix) {
  super(convertOperand(clazz, predicate, in),
      relBuilderFactory,
      createDescription(descriptionPrefix, in, out));
  this.inTrait = Objects.requireNonNull(in);
  this.outTrait = Objects.requireNonNull(out);

  // Source and target traits must have same type
  assert in.getTraitDef() == out.getTraitDef();
}
 

@Override
public RelTraitSet read(final Kryo kryo, final Input input, final Class<RelTraitSet> type) {

  final int size = kryo.readObject(input, Integer.class);
  final RelTrait[] traits = new RelTrait[size];
  for (int i = 0; i < size; i++) {
    final RelTrait trait = (RelTrait) kryo.readClassAndObject(input);
    // normalize trait so that stupid calcite won't complain about == checks.
    traits[i] = trait.getTraitDef().canonize(trait);
  }

  final String digest = kryo.readObject(input, String.class);
  final Object cache = kryo.readClassAndObject(input);

  final RelTraitSet traitSet = kryo.newInstance(type);

  try {
    getField(NAME_CACHE).set(traitSet, cache);
    getField(NAME_TRAITS).set(traitSet, traits);
    getField(NAME_STRING).set(traitSet, digest);
  } catch (final NoSuchFieldException|IllegalAccessException e) {
    throw new RuntimeException("unable to deserialize TraitSet", e);
  }

  kryo.reference(traitSet);
  return traitSet;
}
 

public RelTraitSet newTraitSet(RelTrait... traits) {
  RelTraitSet set = call.getPlanner().emptyTraitSet();
  for (RelTrait t : traits) {
    set = set.plus(t);
  }
  return set;

}
 

@Test
public void testDistributionInPlan() throws Exception {
  String query = String.format("create table %s.\"%s\" hash partition by (\"position\") as select \"position\", isFte, " +
    "count(rating) as \"agg-7-0\", " +
    "sum(rating) as \"agg-7-1\", " +
    "min(rating) as \"agg-7-2\", " +
    "max(rating) as \"agg-7-3\", " +
    "count(1) as convert_count_star " +
    "from cp.acceleration.\"employees.json\" " +
    "group by \"position\", isFte", DFS_TEST_PLUGIN_NAME, "hashDistribute");
  final RelNode physical = getPlan(query);

  final AtomicBoolean hashDistributedWriter = new AtomicBoolean(false);
  physical.accept(new StatelessRelShuttleImpl() {

    @Override
    public RelNode visit(final RelNode other) {
      if (other instanceof WriterPrel) {
        for (RelTrait trait : other.getTraitSet()) {
          if (trait instanceof DistributionTrait) {
            List<DistributionField> distributionFieldList = ((DistributionTrait) trait).getFields();
            if (distributionFieldList.size() != 1) {
              continue;
            }
            int fieldId = distributionFieldList.get(0).getFieldId();
            String fieldName = ((WriterPrel) other).getInput().getRowType().getFieldNames().get(fieldId);
            if ("position".equals(fieldName)) {
              hashDistributedWriter.set(true);
            }
          }
        }
      }
      return super.visit(other);
    }
  });

  Assert.assertTrue("Physical plan must have a HashExchange", hashDistributedWriter.get());
}
 

public GremlinConverterRule(
        Class<? extends RelNode> clazz,
        RelTrait in,
        Convention out,
        String description) {
    super(clazz, in, out, description);
    this.out = out;
}