Java Code Examples for org.apache.calcite.plan.RelTraitSet#replace()

public static RelNode getExchange(RelOptCluster cluster, boolean isSingleton, boolean isExchangeRequired,
                                  RelTraitSet traits, DrillDistributionTrait distributionTrait,
                                  IndexCallContext indexContext, RelNode input) {
  if (!isExchangeRequired) {
    return input;
  }

  if (isSingleton) {
    return new SingleMergeExchangePrel(cluster,
            traits.replace(DrillDistributionTrait.SINGLETON),
            input, indexContext.getCollation());
  } else {
    return new HashToMergeExchangePrel(cluster,
            traits.replace(distributionTrait),
            input, distributionTrait.getFields(), indexContext.getCollation(),
            PrelUtil.getSettings(cluster).numEndPoints());
  }
}
 

public RelNode convert(RelNode rel) {
  final LogicalAggregate agg = (LogicalAggregate) rel;
  if (!Aggregate.isSimple(agg)) {
    return null;
  }
  final RelTraitSet inputTraits = rel.getCluster()
      .traitSet().replace(EnumerableConvention.INSTANCE)
      .replace(
          RelCollations.of(
              ImmutableIntList.copyOf(
          agg.getGroupSet().asList())));
  final RelTraitSet selfTraits = inputTraits.replace(
      RelCollations.of(
      ImmutableIntList.identity(agg.getGroupSet().cardinality())));
  return new EnumerableSortedAggregate(
      rel.getCluster(),
      selfTraits,
      convert(agg.getInput(), inputTraits),
      agg.getGroupSet(),
      agg.getGroupSets(),
      agg.getAggCallList());
}
 

/**
 * This function can be reused when a Join's traits pass-down shall only
 * pass through collation to left input.
 *
 * @param required required trait set for the join
 * @param joinType the join type
 * @param leftInputFieldCount number of field count of left join input
 * @param joinTraitSet trait set of the join
 */
static Pair<RelTraitSet, List<RelTraitSet>> passThroughTraitsForJoin(
    RelTraitSet required, JoinRelType joinType,
    int leftInputFieldCount, RelTraitSet joinTraitSet) {
  RelCollation collation = required.getCollation();
  if (collation == null
      || collation == RelCollations.EMPTY
      || joinType == JoinRelType.FULL
      || joinType == JoinRelType.RIGHT) {
    return null;
  }

  for (RelFieldCollation fc : collation.getFieldCollations()) {
    // If field collation belongs to right input: cannot push down collation.
    if (fc.getFieldIndex() >= leftInputFieldCount) {
      return null;
    }
  }

  RelTraitSet passthroughTraitSet = joinTraitSet.replace(collation);
  return Pair.of(passthroughTraitSet,
      ImmutableList.of(
          passthroughTraitSet,
          passthroughTraitSet.replace(RelCollations.EMPTY)));
}
 

@Test void testRelDistributionSatisfy() {
  RelDistribution distribution1 = RelDistributions.hash(ImmutableList.of(0));
  RelDistribution distribution2 = RelDistributions.hash(ImmutableList.of(1));

  RelTraitSet traitSet = RelTraitSet.createEmpty();
  RelTraitSet simpleTrait1 = traitSet.plus(distribution1);
  RelTraitSet simpleTrait2 = traitSet.plus(distribution2);
  RelTraitSet compositeTrait =
      traitSet.replace(RelDistributionTraitDef.INSTANCE,
          ImmutableList.of(distribution1, distribution2));

  assertThat(compositeTrait.satisfies(simpleTrait1), is(true));
  assertThat(compositeTrait.satisfies(simpleTrait2), is(true));

  assertThat(distribution1.compareTo(distribution2), is(-1));
  assertThat(distribution2.compareTo(distribution1), is(1));
  //noinspection EqualsWithItself
  assertThat(distribution2.compareTo(distribution2), is(0));
}
 

public static RelNode convert(RelNode rel, RelTraitSet toTraits) {
  toTraits = toTraits.simplify();

  PlannerSettings settings = PrelUtil.getSettings(rel.getCluster());
  if(settings.isSingleMode()){
    toTraits = toTraits.replace(DrillDistributionTrait.ANY);
  }

  return RelOptRule.convert(rel, toTraits.simplify());
}
 

public static RelTraitSet fixTraits(RelOptPlanner cluster, RelTraitSet set) {
    if (getPlannerSettings(cluster).isSingleMode()) {
        return set.replace(DrillDistributionTrait.ANY);
    } else {
        return set;
    }
}
 

public static RelNode convert(RelNode rel, RelTraitSet toTraits){
  toTraits = toTraits.simplify();

  PlannerSettings settings = PrelUtil.getSettings(rel.getCluster());
  if(settings.isSingleMode()){
    toTraits = toTraits.replace(DistributionTrait.ANY);
  }

  return RelOptRule.convert(rel, toTraits.simplify());
}
 

public static RelTraitSet fixTraits(RelOptPlanner cluster, RelTraitSet set) {
  if (getPlannerSettings(cluster).isSingleMode()) {
    return set.replace(DistributionTrait.ANY);
  } else {
    return set;
  }
}
 

public void onMatch(RelOptRuleCall call) {
  SingleRel singleRel = call.rel(0);
  RelNode emptyValues = call.builder().push(singleRel).empty().build();
  RelTraitSet traits = singleRel.getTraitSet();
  // propagate all traits (except convention) from the original singleRel into the empty values
  if (emptyValues.getConvention() != null) {
    traits = traits.replace(emptyValues.getConvention());
  }
  emptyValues = emptyValues.copy(traits, Collections.emptyList());
  call.transformTo(emptyValues);
}
 

public static EnumerableMergeJoin create(RelNode left, RelNode right,
    RexNode condition, ImmutableIntList leftKeys,
    ImmutableIntList rightKeys, JoinRelType joinType) {
  final RelOptCluster cluster = right.getCluster();
  RelTraitSet traitSet = cluster.traitSet();
  if (traitSet.isEnabled(RelCollationTraitDef.INSTANCE)) {
    final RelMetadataQuery mq = cluster.getMetadataQuery();
    final List<RelCollation> collations =
        RelMdCollation.mergeJoin(mq, left, right, leftKeys, rightKeys, joinType);
    traitSet = traitSet.replace(collations);
  }
  return new EnumerableMergeJoin(cluster, traitSet, left, right, condition,
      ImmutableSet.of(), joinType);
}
 

public EmptyRel(RelOptCluster cluster, RelTraitSet traitSet, RelDataType rowType, BatchSchema schema) {
  super(cluster, traitSet.replace(LOGICAL));
  this.schema = schema;
  this.rowType = rowType;
  this.digest = schema.toString();
}
 

static protected RelTraitSet adjustTraits(RelTraitSet traitSet) {
  return traitSet.replace(RelCollationTraitDef.INSTANCE, ImmutableList.of());
}
 

private PlannerResult runPlanner(String defaultTableSpace, String query) throws RelConversionException,
        SqlParseException, ValidationException, MetadataStorageManagerException, StatementExecutionException {
    SchemaPlus subSchema = getSchemaForTableSpace(defaultTableSpace);
    if (subSchema == null) {
        clearCache();
        throw new StatementExecutionException("tablespace " + defaultTableSpace + " is not available");
    }
    Properties props = new Properties();
    props.put(CalciteConnectionProperty.TIME_ZONE.camelName(), TimeZone.getDefault().getID());
    props.put(CalciteConnectionProperty.LOCALE.camelName(), Locale.ROOT.toString());
    final CalciteConnectionConfigImpl calciteRuntimeContextConfig = new CalciteConnectionConfigImpl(props);

    final FrameworkConfig config = Frameworks.newConfigBuilder()
            .parserConfig(SQL_PARSER_CONFIG)
            .defaultSchema(subSchema)
            .traitDefs(TRAITS)
            .context(new Context() {
                @Override
                public <C> C unwrap(Class<C> aClass) {
                    if (aClass == CalciteConnectionConfigImpl.class
                            || aClass == CalciteConnectionConfig.class) {
                        return (C) calciteRuntimeContextConfig;
                    }
                    return null;
                }
            })
            // define the rules you want to apply
            .programs(Programs.ofRules(Programs.RULE_SET))
            .build();
    Planner planner = Frameworks.getPlanner(config);
    if (LOG.isLoggable(Level.FINER)) {
        LOG.log(Level.FINER, "Query: {0}", query);
    }
    try {
        SqlNode n = planner.parse(query);
        n = planner.validate(n);
        RelNode logicalPlan = planner.rel(n).project();
        if (LOG.isLoggable(DUMP_QUERY_LEVEL)) {
            LOG.log(DUMP_QUERY_LEVEL, "Query: {0} {1}", new Object[]{query,
                RelOptUtil.dumpPlan("-- Logical Plan", logicalPlan, SqlExplainFormat.TEXT,
                SqlExplainLevel.ALL_ATTRIBUTES)});
        }
        RelDataType originalRowType = logicalPlan.getRowType();
        RelOptCluster cluster = logicalPlan.getCluster();
        final RelOptPlanner optPlanner = cluster.getPlanner();

        optPlanner.addRule(ReduceExpressionsRule.FILTER_INSTANCE);
        RelTraitSet desiredTraits =
                cluster.traitSet()
                        .replace(EnumerableConvention.INSTANCE);
        final RelCollation collation =
                logicalPlan instanceof Sort
                        ? ((Sort) logicalPlan).collation
                        : null;
        if (collation != null) {
            desiredTraits = desiredTraits.replace(collation);
        }
        final RelNode newRoot = optPlanner.changeTraits(logicalPlan, desiredTraits);
        optPlanner.setRoot(newRoot);
        RelNode bestExp = optPlanner.findBestExp();
        if (LOG.isLoggable(DUMP_QUERY_LEVEL)) {
            LOG.log(DUMP_QUERY_LEVEL, "Query: {0} {1}", new Object[]{query,
                RelOptUtil.dumpPlan("-- Best  Plan", bestExp, SqlExplainFormat.TEXT,
                SqlExplainLevel.ALL_ATTRIBUTES)});
        }

        return new PlannerResult(bestExp, originalRowType, logicalPlan, n);
    } catch (AssertionError err) {
        throw new StatementExecutionException("Internal Calcite error " + err, err);
    }
}
 

@Override public RelNode convert(RelNode rel) {
  LogicalJoin join = (LogicalJoin) rel;
  final JoinInfo info = join.analyzeCondition();
  if (!EnumerableMergeJoin.isMergeJoinSupported(join.getJoinType())) {
    // EnumerableMergeJoin only supports certain join types.
    return null;
  }
  if (info.pairs().size() == 0) {
    // EnumerableMergeJoin CAN support cartesian join, but disable it for now.
    return null;
  }
  final List<RelNode> newInputs = new ArrayList<>();
  final List<RelCollation> collations = new ArrayList<>();
  int offset = 0;
  for (Ord<RelNode> ord : Ord.zip(join.getInputs())) {
    RelTraitSet traits = ord.e.getTraitSet()
        .replace(EnumerableConvention.INSTANCE);
    if (!info.pairs().isEmpty()) {
      final List<RelFieldCollation> fieldCollations = new ArrayList<>();
      for (int key : info.keys().get(ord.i)) {
        fieldCollations.add(
            new RelFieldCollation(key, RelFieldCollation.Direction.ASCENDING,
                RelFieldCollation.NullDirection.LAST));
      }
      final RelCollation collation = RelCollations.of(fieldCollations);
      collations.add(RelCollations.shift(collation, offset));
      traits = traits.replace(collation);
    }
    newInputs.add(convert(ord.e, traits));
    offset += ord.e.getRowType().getFieldCount();
  }
  final RelNode left = newInputs.get(0);
  final RelNode right = newInputs.get(1);
  final RelOptCluster cluster = join.getCluster();
  RelNode newRel;

  RelTraitSet traitSet = join.getTraitSet()
      .replace(EnumerableConvention.INSTANCE);
  if (!collations.isEmpty()) {
    traitSet = traitSet.replace(collations);
  }
  // Re-arrange condition: first the equi-join elements, then the non-equi-join ones (if any);
  // this is not strictly necessary but it will be useful to avoid spurious errors in the
  // unit tests when verifying the plan.
  final RexBuilder rexBuilder = join.getCluster().getRexBuilder();
  final RexNode equi = info.getEquiCondition(left, right, rexBuilder);
  final RexNode condition;
  if (info.isEqui()) {
    condition = equi;
  } else {
    final RexNode nonEqui = RexUtil.composeConjunction(rexBuilder, info.nonEquiConditions);
    condition = RexUtil.composeConjunction(rexBuilder, Arrays.asList(equi, nonEqui));
  }
  newRel = new EnumerableMergeJoin(cluster,
      traitSet,
      left,
      right,
      condition,
      join.getVariablesSet(),
      join.getJoinType());
  return newRel;
}