Java Code Examples for org.apache.calcite.plan.RelOptUtil#getAllFields()

public static RelNode apply(RelOptRuleCall call, Aggregate aggregate,
    Project project) {
  // Find all fields which we need to be straightforward field projections.
  final Set<Integer> interestingFields = RelOptUtil.getAllFields(aggregate);

  // Build the map from old to new; abort if any entry is not a
  // straightforward field projection.
  final Map<Integer, Integer> map = new HashMap<>();
  for (int source : interestingFields) {
    final RexNode rex = project.getProjects().get(source);
    if (!(rex instanceof RexInputRef)) {
      return null;
    }
    map.put(source, ((RexInputRef) rex).getIndex());
  }

  final ImmutableBitSet newGroupSet = aggregate.getGroupSet().permute(map);
  ImmutableList<ImmutableBitSet> newGroupingSets = null;
  if (aggregate.getGroupType() != Group.SIMPLE) {
    newGroupingSets =
        ImmutableBitSet.ORDERING.immutableSortedCopy(
            ImmutableBitSet.permute(aggregate.getGroupSets(), map));
  }

  final ImmutableList.Builder<AggregateCall> aggCalls =
      ImmutableList.builder();
  final int sourceCount = aggregate.getInput().getRowType().getFieldCount();
  final int targetCount = project.getInput().getRowType().getFieldCount();
  final Mappings.TargetMapping targetMapping =
      Mappings.target(map, sourceCount, targetCount);
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    aggCalls.add(aggregateCall.transform(targetMapping));
  }

  final Aggregate newAggregate =
      aggregate.copy(aggregate.getTraitSet(), project.getInput(),
          newGroupSet, newGroupingSets, aggCalls.build());

  // Add a project if the group set is not in the same order or
  // contains duplicates.
  final RelBuilder relBuilder = call.builder();
  relBuilder.push(newAggregate);
  final List<Integer> newKeys =
      Lists.transform(aggregate.getGroupSet().asList(), map::get);
  if (!newKeys.equals(newGroupSet.asList())) {
    final List<Integer> posList = new ArrayList<>();
    for (int newKey : newKeys) {
      posList.add(newGroupSet.indexOf(newKey));
    }
    for (int i = newAggregate.getGroupCount();
         i < newAggregate.getRowType().getFieldCount(); i++) {
      posList.add(i);
    }
    relBuilder.project(relBuilder.fields(posList));
  }

  return relBuilder.build();
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final Join join = call.rel(1);
  boolean isLeftJoin = join.getJoinType() == JoinRelType.LEFT;
  int lower = isLeftJoin
      ? join.getLeft().getRowType().getFieldCount() - 1 : 0;
  int upper = isLeftJoin ? join.getRowType().getFieldCount()
      : join.getLeft().getRowType().getFieldCount();

  // Check whether the aggregate uses columns whose index is between
  // lower(included) and upper(excluded).
  final Set<Integer> allFields = RelOptUtil.getAllFields(aggregate);
  if (allFields.stream().anyMatch(i -> i >= lower && i < upper)) {
    return;
  }

  if (aggregate.getAggCallList().stream().anyMatch(
      aggregateCall -> !aggregateCall.isDistinct())) {
    return;
  }

  RelNode node;
  if (isLeftJoin) {
    node = aggregate.copy(aggregate.getTraitSet(), join.getLeft(),
        aggregate.getGroupSet(), aggregate.getGroupSets(),
        aggregate.getAggCallList());
  } else {
    final Map<Integer, Integer> map = new HashMap<>();
    allFields.forEach(index -> map.put(index, index - upper));
    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)));

    final RelBuilder relBuilder = call.builder();
    node = relBuilder.push(join.getRight())
        .aggregate(relBuilder.groupKey(groupSet), aggCalls.build())
        .build();
  }
  call.transformTo(node);
}
 

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