Java Code Examples for org.apache.calcite.rex.RexUtil#shift()

@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());
}
 

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

  // Check whether the aggregate uses columns in the right input of
  // bottom join.
  final Set<Integer> allFields = RelOptUtil.getAllFields(aggregate);
  if (allFields.stream().anyMatch(i -> i >= leftBottomChildSize
      && i < bottomJoin.getRowType().getFieldCount())) {
    return;
  }

  if (aggregate.getAggCallList().stream().anyMatch(aggregateCall ->
      !aggregateCall.isDistinct())) {
    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<>();
  RelOptUtil.splitJoinCondition(bottomJoin.getLeft(), bottomJoin.getRight(),
      bottomJoin.getCondition(), leftChildKeys, new ArrayList<>(),
      new ArrayList<>());
  if (!leftKeys.equals(leftChildKeys)) {
    return;
  }

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

  final Map<Integer, Integer> map = new HashMap<>();
  allFields.forEach(
      index ->
          map.put(index,
              index < leftBottomChildSize ? index : index - offset));
  final ImmutableBitSet groupSet = aggregate.getGroupSet().permute(map);

  final ImmutableList.Builder<AggregateCall> aggCalls =
      ImmutableList.builder();
  final int sourceCount = aggregate.getInput().getRowType().getFieldCount();
  final Mappings.TargetMapping targetMapping =
      Mappings.target(map, sourceCount, sourceCount);
  aggregate.getAggCallList().forEach(
      aggregateCall ->
          aggCalls.add(aggregateCall.transform(targetMapping)));

  RelNode newAggregate = relBuilder.push(join)
      .aggregate(relBuilder.groupKey(groupSet), aggCalls.build())
      .build();

  call.transformTo(newAggregate);
}
 

@Override protected UnifyResult apply(UnifyRuleCall call) {
  final MutableJoin query = (MutableJoin) call.query;
  final MutableRel qInput0 = query.getLeft();
  final MutableCalc qInput1 = (MutableCalc) query.getRight();
  final Pair<RexNode, List<RexNode>> qInput1Explained = explainCalc(qInput1);
  final RexNode qInput1Cond = qInput1Explained.left;
  final List<RexNode> qInput1Projs = qInput1Explained.right;

  final MutableJoin target = (MutableJoin) call.target;

  final RexBuilder rexBuilder = call.getCluster().getRexBuilder();

  // Try pulling up MutableCalc only when:
  // 1. it's inner join.
  // 2. it's outer join but no filtering condition from MutableCalc.
  final JoinRelType joinRelType = sameJoinType(query.joinType, target.joinType);
  if (joinRelType == null) {
    return null;
  }
  if (joinRelType != JoinRelType.INNER
      && !(joinRelType.isOuterJoin() && qInput1Cond.isAlwaysTrue())) {
    return null;
  }
  // Try pulling up MutableCalc only when Join condition references mapping.
  final List<RexNode> identityProjects =
      (List<RexNode>) rexBuilder.identityProjects(qInput0.rowType);
  if (!referenceByMapping(query.condition, identityProjects, qInput1Projs)) {
    return null;
  }

  final RexNode newQueryJoinCond = new RexShuttle() {
    @Override public RexNode visitInputRef(RexInputRef inputRef) {
      final int idx = inputRef.getIndex();
      if (idx < fieldCnt(qInput0)) {
        return inputRef;
      } else {
        final int newIdx = ((RexInputRef) qInput1Projs.get(idx - fieldCnt(qInput0)))
            .getIndex() + fieldCnt(qInput0);
        return new RexInputRef(newIdx, inputRef.getType());
      }
    }
  }.apply(query.condition);

  final RexNode splitted =
      splitFilter(call.getSimplify(), newQueryJoinCond, target.condition);
  // MutableJoin matches only when the conditions are analyzed to be same.
  if (splitted != null && splitted.isAlwaysTrue()) {
    final RexNode compenCond =
        RexUtil.shift(qInput1Cond, qInput0.rowType.getFieldCount());
    final List<RexNode> compenProjs = new ArrayList<>();
    for (int i = 0; i < query.rowType.getFieldCount(); i++) {
      if (i < fieldCnt(qInput0)) {
        compenProjs.add(
            new RexInputRef(i, query.rowType.getFieldList().get(i).getType()));
      } else {
        final RexNode shifted = RexUtil.shift(qInput1Projs.get(i - fieldCnt(qInput0)),
            qInput0.rowType.getFieldCount());
        compenProjs.add(shifted);
      }
    }
    final RexProgram compensatingRexProgram = RexProgram.create(
        target.rowType, compenProjs, compenCond,
        query.rowType, rexBuilder);
    final MutableCalc compenCalc = MutableCalc.of(target, compensatingRexProgram);
    return tryMergeParentCalcAndGenResult(call, compenCalc);
  }
  return null;
}