Java Code Examples for org.apache.calcite.plan.RelOptRuleCall#transformTo()

public void onMatch(RelOptRuleCall call) {
  SemiJoin semiJoin = call.rel(0);
  LogicalFilter filter = call.rel(1);

  RelNode newSemiJoin =
      SemiJoin.create(filter.getInput(),
          semiJoin.getRight(),
          semiJoin.getCondition(),
          semiJoin.getLeftKeys(),
          semiJoin.getRightKeys());

  final RelFactories.FilterFactory factory =
      RelFactories.DEFAULT_FILTER_FACTORY;
  RelNode newFilter =
      factory.createFilter(newSemiJoin, filter.getCondition());

  call.transformTo(newFilter);
}
 

public void onMatch(RelOptRuleCall call) {
  final Sort sort = call.rel(0);
  final DruidQuery query = call.rel(1);
  if (!DruidQuery.isValidSignature(query.signature() + 'l')) {
    return;
  }
  // Either it is:
  // - a pure limit above a query of type scan
  // - a sort and limit on a dimension/metric part of the druid group by query
  if (sort.offset != null && RexLiteral.intValue(sort.offset) != 0) {
    // offset not supported by Druid
    return;
  }
  if (query.getQueryType() == QueryType.SCAN && !RelOptUtil.isPureLimit(sort)) {
    return;
  }

  final RelNode newSort = sort
      .copy(sort.getTraitSet(), ImmutableList.of(Util.last(query.rels)));
  call.transformTo(DruidQuery.extendQuery(query, newSort));
}
 

public void onMatch(RelOptRuleCall call) {
  final Filter topFilter = call.rel(0);
  final Filter bottomFilter = call.rel(1);

  // use RexPrograms to merge the two FilterRels into a single program
  // so we can convert the two LogicalFilter conditions to directly
  // reference the bottom LogicalFilter's child
  RexBuilder rexBuilder = topFilter.getCluster().getRexBuilder();
  RexProgram bottomProgram = createProgram(bottomFilter);
  RexProgram topProgram = createProgram(topFilter);

  RexProgram mergedProgram = RexProgramBuilder.mergePrograms(topProgram, bottomProgram, rexBuilder);

  RexNode newCondition = mergedProgram.expandLocalRef(mergedProgram.getCondition());

  final RelBuilder relBuilder = call.builder();
  relBuilder.push(bottomFilter.getInput()).filter(newCondition);

  call.transformTo(relBuilder.build());
}
 

public void onMatch(RelOptRuleCall call) {
  final Project project = call.rel(0);
  final Sort sort = call.rel(1);
  if (sort.getClass() != Sort.class) {
    return;
  }
  RelNode newProject =
      project.copy(
          project.getTraitSet(), ImmutableList.of(sort.getInput()));
  final Sort newSort =
      sort.copy(
          sort.getTraitSet(),
          newProject,
          sort.getCollation(),
          sort.offset,
          sort.fetch);
  call.transformTo(newSort);
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LimitPrel limit = (LimitPrel) call.rel(0);
  final SingleMergeExchangePrel smex = (SingleMergeExchangePrel) call.rel(1);
  final SortPrel sort = (SortPrel) call.rel(2);

  // First offset to include into results (inclusive). Null implies it is starting from offset 0
  int offset = limit.getOffset() != null ? Math.max(0, RexLiteral.intValue(limit.getOffset())) : 0;
  int fetch = limit.getFetch() != null?  Math.max(0, RexLiteral.intValue(limit.getFetch())) : 0;

  final TopNPrel topN = new TopNPrel(limit.getCluster(), sort.getTraitSet(), sort.getInput(), offset + fetch, sort.getCollation());
  final LimitPrel newLimit = new LimitPrel(limit.getCluster(), limit.getTraitSet(),
      new SingleMergeExchangePrel(smex.getCluster(), smex.getTraitSet(), topN, sort.getCollation()),
      limit.getOffset(), limit.getFetch());

  call.transformTo(newLimit);
}
 

public void onMatch(RelOptRuleCall call) {
  final Join join = call.rel(0);

  if (join.getJoinType() != JoinRelType.INNER) {
    return;
  }

  if (join.getCondition().isAlwaysTrue()) {
    return;
  }

  if (!join.getSystemFieldList().isEmpty()) {
    // FIXME Enable this rule for joins with system fields
    return;
  }

  final RelBuilder builder = call.builder();

  // NOTE jvs 14-Mar-2006:  See JoinCommuteRule for why we
  // preserve attribute semiJoinDone here.

  final RelNode cartesianJoin =
      join.copy(
          join.getTraitSet(),
          builder.literal(true),
          join.getLeft(),
          join.getRight(),
          join.getJoinType(),
          join.isSemiJoinDone());

  builder.push(cartesianJoin)
      .filter(join.getCondition());

  call.transformTo(builder.build());
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  DrillLimitRel limitRel = call.rel(0);
  DrillProjectRel projectRel = call.rel(1);
  RelNode child = projectRel.getInput();
  final RelNode limitUnderProject = limitRel.copy(limitRel.getTraitSet(), ImmutableList.of(child));
  final RelNode newProject = projectRel.copy(projectRel.getTraitSet(), ImmutableList.of(limitUnderProject));
  call.transformTo(newProject);
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LogicalCorrelate correlate = call.rel(0);
  final RelNode left = correlate.getLeft();
  final RelNode right = correlate.getRight();
  final RelNode convertedLeft = convert(left, left.getTraitSet().plus(DrillRel.DRILL_LOGICAL).simplify());
  final RelNode convertedRight = convert(right, right.getTraitSet().plus(DrillRel.DRILL_LOGICAL).simplify());

  final RelTraitSet traits = correlate.getTraitSet().plus(DrillRel.DRILL_LOGICAL);
  DrillLateralJoinRel lateralJoinRel = new DrillLateralJoinRel(correlate.getCluster(),
      traits, convertedLeft, convertedRight, false, correlate.getCorrelationId(),
      correlate.getRequiredColumns(), correlate.getJoinType());
  call.transformTo(lateralJoinRel);
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Delta delta = call.rel(0);
  final TableScan scan = call.rel(1);
  final RelOptTable relOptTable = scan.getTable();
  final StreamableTable streamableTable =
      relOptTable.unwrap(StreamableTable.class);
  final RelBuilder builder = call.builder();
  if (streamableTable == null) {
    call.transformTo(builder.values(delta.getRowType()).build());
  }
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final WriterRel writer = call.rel(0);
  final RelNode input = call.rel(1);

  final RelTraitSet requestedTraits = writer.getCreateTableEntry()
      .getOptions()
      .inferTraits(input.getTraitSet(), input.getRowType());
  final RelNode convertedInput = convert(input, requestedTraits);

  if (!new WriteTraitPull(call).go(writer, convertedInput)) {
    call.transformTo(convertWriter(writer, convertedInput));
  }
}
 

public void onMatch(RelOptRuleCall call) {
  LogicalJoin join = call.rel(0);
  assertThat(join.getHints().size(), is(1));
  call.transformTo(
      LogicalJoin.create(join.getLeft(),
          join.getRight(),
          join.getHints(),
          join.getCondition(),
          join.getVariablesSet(),
          join.getJoinType()));
}
 

public void onMatch(RelOptRuleCall call) {
  final Union union = call.rel(0);
  final RelBuilder relBuilder = call.builder();
  relBuilder.pushAll(union.getInputs());
  relBuilder.union(true, union.getInputs().size());
  relBuilder.distinct();
  call.transformTo(relBuilder.build());
}
 

public void onMatch(RelOptRuleCall call) {
  Aggregate singleAggregate = call.rel(0);
  Project project = call.rel(1);
  Aggregate aggregate = call.rel(2);

  // check singleAggRel is single_value agg
  if ((!singleAggregate.getGroupSet().isEmpty())
      || (singleAggregate.getAggCallList().size() != 1)
      || !(singleAggregate.getAggCallList().get(0).getAggregation()
      instanceof SqlSingleValueAggFunction)) {
    return;
  }

  // check projRel only projects one expression
  // check this project only projects one expression, i.e. scalar
  // sub-queries.
  List<RexNode> projExprs = project.getProjects();
  if (projExprs.size() != 1) {
    return;
  }

  // check the input to project is an aggregate on the entire input
  if (!aggregate.getGroupSet().isEmpty()) {
    return;
  }

  // singleAggRel produces a nullable type, so create the new
  // projection that casts proj expr to a nullable type.
  final RelBuilder relBuilder = call.builder();
  final RelDataType type =
      relBuilder.getTypeFactory()
          .createTypeWithNullability(projExprs.get(0).getType(), true);
  final RexNode cast =
      relBuilder.getRexBuilder().makeCast(type, projExprs.get(0));
  relBuilder.push(aggregate)
      .project(cast);
  call.transformTo(relBuilder.build());
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  LogicalFilter filter = (LogicalFilter) call.rels[0];
  RexNode newCondition = RexRewriter.rewrite(filter.getCondition(), getRules(filter.getCluster().getRexBuilder()));
  if (newCondition != filter.getCondition()) {
    call.transformTo(LogicalFilter.create(filter.getInput(), newCondition));
  }
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final DrillAggregateRel agg = (DrillAggregateRel) call.rel(0);
  final DrillScanRel scan = (DrillScanRel) call.rel(call.rels.length - 1);
  final DrillProjectRel project = call.rels.length == 3 ? (DrillProjectRel) call.rel(1) : null;

  final GroupScan oldGrpScan = scan.getGroupScan();
  final PlannerSettings settings = PrelUtil.getPlannerSettings(call.getPlanner());

  // Only apply the rule when:
  //    1) scan knows the exact row count in getSize() call,
  //    2) No GroupBY key,
  //    3) No distinct agg call.
  if (!(oldGrpScan.getScanStats(settings).getGroupScanProperty().hasExactRowCount()
      && agg.getGroupCount() == 0
      && !agg.containsDistinctCall())) {
    return;
  }

  Map<String, Long> result = collectCounts(settings, agg, scan, project);
  logger.trace("Calculated the following aggregate counts: ", result);
  // if could not determine the counts, rule won't be applied
  if (result.isEmpty()) {
    return;
  }

  final RelDataType scanRowType = CountToDirectScanUtils.constructDataType(agg, result.keySet());

  final DynamicPojoRecordReader<Long> reader = new DynamicPojoRecordReader<>(
      CountToDirectScanUtils.buildSchema(scanRowType.getFieldNames()),
      Collections.singletonList((List<Long>) new ArrayList<>(result.values())));

  final ScanStats scanStats = new ScanStats(ScanStats.GroupScanProperty.EXACT_ROW_COUNT, 1, 1, scanRowType.getFieldCount());
  final GroupScan directScan = new MetadataDirectGroupScan(reader, oldGrpScan.getFiles(), scanStats, false);

  final DirectScanPrel newScan = DirectScanPrel.create(scan, scan.getTraitSet().plus(Prel.DRILL_PHYSICAL)
      .plus(DrillDistributionTrait.SINGLETON), directScan, scanRowType);

  final ProjectPrel newProject = new ProjectPrel(agg.getCluster(), agg.getTraitSet().plus(Prel.DRILL_PHYSICAL)
      .plus(DrillDistributionTrait.SINGLETON), newScan, CountToDirectScanUtils.prepareFieldExpressions(scanRowType), agg.getRowType());

  call.transformTo(newProject);
}
 

public void onMatch(RelOptRuleCall call) {
  NoneLeafRel leafRel = call.rel(0);
  call.transformTo(new LeafRel(leafRel.getCluster(), leafRel.label));
}
 

/**
 * Similar to {@link #onMatch}, but swaps the upper sibling with the left
 * of the two lower siblings, rather than the right.
 */
private void onMatchLeft(RelOptRuleCall call) {
  final Join topJoin = call.rel(0);
  final Join bottomJoin = call.rel(1);
  final RelNode relC = call.rel(2);
  final RelNode relA = bottomJoin.getLeft();
  final RelNode relB = bottomJoin.getRight();
  final RelOptCluster cluster = topJoin.getCluster();

  //        topJoin
  //        /     \
  //   bottomJoin  C
  //    /    \
  //   A      B

  final int aCount = relA.getRowType().getFieldCount();
  final int bCount = relB.getRowType().getFieldCount();
  final int cCount = relC.getRowType().getFieldCount();
  final ImmutableBitSet aBitSet = ImmutableBitSet.range(aCount);

  // becomes
  //
  //        newTopJoin
  //        /        \
  //   newBottomJoin  A
  //    /    \
  //   C      B

  // If either join is not inner, we cannot proceed.
  // (Is this too strict?)
  if (topJoin.getJoinType() != JoinRelType.INNER
      || bottomJoin.getJoinType() != JoinRelType.INNER) {
    return;
  }

  // Split the condition of topJoin into a conjunction. Each of the
  // parts that does not use columns from A can be pushed down.
  final List<RexNode> intersecting = new ArrayList<>();
  final List<RexNode> nonIntersecting = new ArrayList<>();
  split(topJoin.getCondition(), aBitSet, intersecting, nonIntersecting);

  // If there's nothing to push down, it's not worth proceeding.
  if (nonIntersecting.isEmpty()) {
    return;
  }

  // Split the condition of bottomJoin into a conjunction. Each of the
  // parts that use columns from A will need to be pulled up.
  final List<RexNode> bottomIntersecting = new ArrayList<>();
  final List<RexNode> bottomNonIntersecting = new ArrayList<>();
  split(
      bottomJoin.getCondition(), aBitSet, bottomIntersecting,
      bottomNonIntersecting);

  // target: | C      | B |
  // source: | A       | B | C      |
  final Mappings.TargetMapping bottomMapping =
      Mappings.createShiftMapping(
          aCount + bCount + cCount,
          cCount, aCount, bCount,
          0, aCount + bCount, cCount);
  final List<RexNode> newBottomList = new ArrayList<>();
  new RexPermuteInputsShuttle(bottomMapping, relC, relB)
      .visitList(nonIntersecting, newBottomList);
  new RexPermuteInputsShuttle(bottomMapping, relC, relB)
      .visitList(bottomNonIntersecting, newBottomList);
  final RexBuilder rexBuilder = cluster.getRexBuilder();
  RexNode newBottomCondition =
      RexUtil.composeConjunction(rexBuilder, newBottomList);
  final Join newBottomJoin =
      bottomJoin.copy(bottomJoin.getTraitSet(), newBottomCondition, relC,
          relB, bottomJoin.getJoinType(), bottomJoin.isSemiJoinDone());

  // target: | C      | B | A       |
  // source: | A       | B | C      |
  final Mappings.TargetMapping topMapping =
      Mappings.createShiftMapping(
          aCount + bCount + cCount,
          cCount + bCount, 0, aCount,
          cCount, aCount, bCount,
          0, aCount + bCount, cCount);
  final List<RexNode> newTopList = new ArrayList<>();
  new RexPermuteInputsShuttle(topMapping, newBottomJoin, relA)
      .visitList(intersecting, newTopList);
  new RexPermuteInputsShuttle(topMapping, newBottomJoin, relA)
      .visitList(bottomIntersecting, newTopList);
  RexNode newTopCondition =
      RexUtil.composeConjunction(rexBuilder, newTopList);
  @SuppressWarnings("SuspiciousNameCombination")
  final Join newTopJoin =
      topJoin.copy(topJoin.getTraitSet(), newTopCondition, newBottomJoin,
          relA, topJoin.getJoinType(), topJoin.isSemiJoinDone());

  final RelBuilder relBuilder = call.builder();
  relBuilder.push(newTopJoin);
  relBuilder.project(relBuilder.fields(topMapping));
  call.transformTo(relBuilder.build());
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Project project = call.rel(0);
  final Join topJoin = call.rel(1);
  final Join bottomJoin = call.rel(2);
  int leftBottomChildSize = bottomJoin.getLeft().getRowType().getFieldCount();

  // Check whether the project uses columns in the right input of bottom join.
  for (RexNode expr: project.getProjects()) {
    if (RelOptUtil.InputFinder.bits(expr).asList().stream().anyMatch(
        i -> i >= leftBottomChildSize
            && i < bottomJoin.getRowType().getFieldCount())) {
      return;
    }
  }

  // Check whether the top join uses columns in the right input of bottom join.
  final List<Integer> leftKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(topJoin.getLeft(), topJoin.getRight(),
      topJoin.getCondition(), leftKeys, new ArrayList<>(),
      new ArrayList<>());
  if (leftKeys.stream().anyMatch(s -> s >= leftBottomChildSize)) {
    return;
  }

  // Check whether left join keys in top join and bottom join are equal.
  final List<Integer> leftChildKeys = new ArrayList<>();
  final List<Integer> rightChildKeys = new ArrayList<>();
  RelOptUtil.splitJoinCondition(bottomJoin.getLeft(), bottomJoin.getRight(),
      bottomJoin.getCondition(), leftChildKeys, rightChildKeys,
      new ArrayList<>());
  if (!leftKeys.equals(leftChildKeys)) {
    return;
  }

  // Make sure that right keys of bottom join are unique.
  final ImmutableBitSet.Builder columns = ImmutableBitSet.builder();
  rightChildKeys.forEach(key -> columns.set(key));
  final RelMetadataQuery mq = call.getMetadataQuery();
  if (!mq.areColumnsUnique(bottomJoin.getRight(), columns.build())) {
    return;
  }

  int offset = bottomJoin.getRight().getRowType().getFieldCount();
  final RelBuilder relBuilder = call.builder();

  final RexNode condition = RexUtil.shift(topJoin.getCondition(),
      leftBottomChildSize, -offset);
  final RelNode join = relBuilder.push(bottomJoin.getLeft())
      .push(topJoin.getRight())
      .join(topJoin.getJoinType(), condition)
      .build();

  final List<RexNode> newExprs = project.getProjects().stream()
      .map(expr -> RexUtil.shift(expr, leftBottomChildSize, -offset))
      .collect(Collectors.toList());
  relBuilder.push(join).project(newExprs);
  call.transformTo(relBuilder.build());
}
 

private void onMatch2(
    RelOptRuleCall call,
    LogicalCorrelate correlate,
    RelNode leftInput,
    LogicalProject aggOutputProject,
    LogicalAggregate aggregate) {
  if (generatedCorRels.contains(correlate)) {
    // This Correlate was generated by a previous invocation of
    // this rule. No further work to do.
    return;
  }

  setCurrent(call.getPlanner().getRoot(), correlate);

  // check for this pattern
  // The pattern matching could be simplified if rules can be applied
  // during decorrelation,
  //
  // CorrelateRel(left correlation, condition = true)
  //   leftInput
  //   Project-A (a RexNode)
  //     Aggregate (groupby (0), agg0(), agg1()...)

  // check aggOutputProj projects only one expression
  List<RexNode> aggOutputProjExprs = aggOutputProject.getProjects();
  if (aggOutputProjExprs.size() != 1) {
    return;
  }

  JoinRelType joinType = correlate.getJoinType();
  // corRel.getCondition was here, however Correlate was updated so it
  // never includes a join condition. The code was not modified for brevity.
  RexNode joinCond = relBuilder.literal(true);
  if ((joinType != JoinRelType.LEFT)
      || (joinCond != relBuilder.literal(true))) {
    return;
  }

  // check that the agg is on the entire input
  if (!aggregate.getGroupSet().isEmpty()) {
    return;
  }

  List<AggregateCall> aggCalls = aggregate.getAggCallList();
  Set<Integer> isCount = new HashSet<>();

  // remember the count() positions
  int i = -1;
  for (AggregateCall aggCall : aggCalls) {
    ++i;
    if (aggCall.getAggregation() instanceof SqlCountAggFunction) {
      isCount.add(i);
    }
  }

  // now rewrite the plan to
  //
  // Project-A' (all LHS plus transformed original projections,
  //             replacing references to count() with case statement)
  //   Correlate(left correlation, condition = true)
  //     leftInput
  //     Aggregate(groupby (0), agg0(), agg1()...)
  //
  LogicalCorrelate newCorrelate =
      LogicalCorrelate.create(leftInput, aggregate,
          correlate.getCorrelationId(), correlate.getRequiredColumns(),
          correlate.getJoinType());

  // remember this rel so we don't fire rule on it again
  // REVIEW jhyde 29-Oct-2007: rules should not save state; rule
  // should recognize patterns where it does or does not need to do
  // work
  generatedCorRels.add(newCorrelate);

  // need to update the mapCorToCorRel Update the output position
  // for the corVars: only pass on the corVars that are not used in
  // the join key.
  if (cm.mapCorToCorRel.get(correlate.getCorrelationId()) == correlate) {
    cm.mapCorToCorRel.put(correlate.getCorrelationId(), newCorrelate);
  }

  RelNode newOutput =
      aggregateCorrelatorOutput(newCorrelate, aggOutputProject, isCount);

  call.transformTo(newOutput);
}
 

protected void perform(RelOptRuleCall call, Project project,
    Join join, RelNode left, Aggregate aggregate) {
  final RelOptCluster cluster = join.getCluster();
  final RexBuilder rexBuilder = cluster.getRexBuilder();
  if (project != null) {
    final ImmutableBitSet bits =
        RelOptUtil.InputFinder.bits(project.getProjects(), null);
    final ImmutableBitSet rightBits =
        ImmutableBitSet.range(left.getRowType().getFieldCount(),
            join.getRowType().getFieldCount());
    if (bits.intersects(rightBits)) {
      return;
    }
  } else {
    if (join.getJoinType().projectsRight()
        && !IS_EMPTY_AGGREGATE.test(aggregate)) {
      return;
    }
  }
  final JoinInfo joinInfo = join.analyzeCondition();
  if (!joinInfo.rightSet().equals(
      ImmutableBitSet.range(aggregate.getGroupCount()))) {
    // Rule requires that aggregate key to be the same as the join key.
    // By the way, neither a super-set nor a sub-set would work.
    return;
  }
  if (!joinInfo.isEqui()) {
    return;
  }
  final RelBuilder relBuilder = call.builder();
  relBuilder.push(left);
  switch (join.getJoinType()) {
  case SEMI:
  case INNER:
    final List<Integer> newRightKeyBuilder = new ArrayList<>();
    final List<Integer> aggregateKeys = aggregate.getGroupSet().asList();
    for (int key : joinInfo.rightKeys) {
      newRightKeyBuilder.add(aggregateKeys.get(key));
    }
    final ImmutableIntList newRightKeys = ImmutableIntList.copyOf(newRightKeyBuilder);
    relBuilder.push(aggregate.getInput());
    final RexNode newCondition =
        RelOptUtil.createEquiJoinCondition(relBuilder.peek(2, 0),
            joinInfo.leftKeys, relBuilder.peek(2, 1), newRightKeys,
            rexBuilder);
    relBuilder.semiJoin(newCondition);
    break;

  case LEFT:
    // The right-hand side produces no more than 1 row (because of the
    // Aggregate) and no fewer than 1 row (because of LEFT), and therefore
    // we can eliminate the semi-join.
    break;

  default:
    throw new AssertionError(join.getJoinType());
  }
  if (project != null) {
    relBuilder.project(project.getProjects(), project.getRowType().getFieldNames());
  }
  final RelNode relNode = relBuilder.build();
  call.transformTo(relNode);
}