Java Code Examples for org.apache.calcite.util.mapping.Mappings#target()

protected void apply(RelOptRuleCall call, Filter filter, TableScan scan) {
  final ImmutableIntList projects;
  final ImmutableList.Builder<RexNode> filters = ImmutableList.builder();
  if (scan instanceof Bindables.BindableTableScan) {
    final Bindables.BindableTableScan bindableScan =
        (Bindables.BindableTableScan) scan;
    filters.addAll(bindableScan.filters);
    projects = bindableScan.projects;
  } else {
    projects = scan.identity();
  }

  final Mapping mapping = Mappings.target(projects,
      scan.getTable().getRowType().getFieldCount());
  filters.add(
      RexUtil.apply(mapping.inverse(), filter.getCondition()));

  call.transformTo(
      Bindables.BindableTableScan.create(scan.getCluster(), scan.getTable(),
          filters.build(), projects));
}
 

private Mappings.TargetMapping buildMapping(boolean left2Right) {
  ImmutableIntList sourceKeys = left2Right ? joinInfo.leftKeys : joinInfo.rightKeys;
  ImmutableIntList targetKeys = left2Right ? joinInfo.rightKeys : joinInfo.leftKeys;
  Map<Integer, Integer> keyMap = new HashMap<>();
  for (int i = 0; i < joinInfo.leftKeys.size(); i++) {
    keyMap.put(sourceKeys.get(i), targetKeys.get(i));
  }

  Mappings.TargetMapping mapping = Mappings.target(keyMap,
      (left2Right ? left : right).getRowType().getFieldCount(),
      (left2Right ? right : left).getRowType().getFieldCount());
  return mapping;
}
 

/**
 * Returns a mapping between old and new fields.
 *
 * @param oldRel Old relational expression
 * @return Mapping between fields of old and new
 */
private Mappings.TargetMapping getNewForOldInputMapping(RelNode oldRel) {
  final RelNode newRel = getNewForOldRel(oldRel);
  return Mappings.target(
      this::getNewForOldInput,
      oldRel.getRowType().getFieldCount(),
      newRel.getRowType().getFieldCount());
}
 

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

@Override protected UnifyResult apply(UnifyRuleCall call) {
  final MutableAggregate query = (MutableAggregate) call.query;
  final MutableCalc qInput = (MutableCalc) query.getInput();
  final Pair<RexNode, List<RexNode>> qInputExplained = explainCalc(qInput);
  final RexNode qInputCond = qInputExplained.left;
  final List<RexNode> qInputProjs = qInputExplained.right;

  final MutableAggregate target = (MutableAggregate) call.target;

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

  final Mappings.TargetMapping mapping =
      Project.getMapping(fieldCnt(qInput.getInput()), qInputProjs);
  if (mapping == null) {
    return null;
  }

  if (!qInputCond.isAlwaysTrue()) {
    try {
      // Fail the matching when filtering condition references
      // non-grouping columns in target.
      qInputCond.accept(new RexVisitorImpl<Void>(true) {
        @Override public Void visitInputRef(RexInputRef inputRef) {
          if (!target.groupSets.stream()
              .allMatch(groupSet -> groupSet.get(inputRef.getIndex()))) {
            throw Util.FoundOne.NULL;
          }
          return super.visitInputRef(inputRef);
        }
      });
    } catch (Util.FoundOne one) {
      return null;
    }
  }

  final Mapping inverseMapping = mapping.inverse();
  final MutableAggregate aggregate2 =
      permute(query, qInput.getInput(), inverseMapping);

  final Mappings.TargetMapping mappingForQueryCond = Mappings.target(
      target.groupSet::indexOf,
      target.getInput().rowType.getFieldCount(),
      target.groupSet.cardinality());
  final RexNode targetCond = RexUtil.apply(mappingForQueryCond, qInputCond);

  final MutableRel unifiedAggregate =
      unifyAggregates(aggregate2, targetCond, target);
  if (unifiedAggregate == null) {
    return null;
  }
  // Add Project if the mapping breaks order of fields in GroupSet
  if (!Mappings.keepsOrdering(mapping)) {
    final List<Integer> posList = new ArrayList<>();
    final int fieldCount = aggregate2.rowType.getFieldCount();
    final List<Pair<Integer, Integer>> pairs = new ArrayList<>();
    final List<Integer> groupings = aggregate2.groupSet.toList();
    for (int i = 0; i < groupings.size(); i++) {
      pairs.add(Pair.of(mapping.getTarget(groupings.get(i)), i));
    }
    Collections.sort(pairs);
    pairs.forEach(pair -> posList.add(pair.right));
    for (int i = posList.size(); i < fieldCount; i++) {
      posList.add(i);
    }
    final List<RexNode> compenProjs =
        MutableRels.createProjectExprs(unifiedAggregate, posList);
    final RexProgram compensatingRexProgram = RexProgram.create(
        unifiedAggregate.rowType, compenProjs, null,
        query.rowType, rexBuilder);
    final MutableCalc compenCalc =
        MutableCalc.of(unifiedAggregate, compensatingRexProgram);
    if (unifiedAggregate instanceof MutableCalc) {
      final MutableCalc newCompenCalc =
          mergeCalc(rexBuilder, compenCalc, (MutableCalc) unifiedAggregate);
      return tryMergeParentCalcAndGenResult(call, newCompenCalc);
    } else {
      return tryMergeParentCalcAndGenResult(call, compenCalc);
    }
  } else {
    return tryMergeParentCalcAndGenResult(call, unifiedAggregate);
  }
}
 

private static Mapping mapping(int sourceCount, int... sources) {
  return Mappings.target(ImmutableIntList.of(sources), sourceCount);
}
 

private static Mapping mapping(int sourceCount, int... sources) {
  return Mappings.target(ImmutableIntList.of(sources), sourceCount);
}
 

/**
 * Returns a mapping between old and new fields.
 *
 * @param oldRel Old relational expression
 * @return Mapping between fields of old and new
 */
private Mappings.TargetMapping getNewForOldInputMapping(RelNode oldRel) {
    final RelNode newRel = getNewForOldRel(oldRel);
    return Mappings.target(this::getNewForOldInput, oldRel.getRowType().getFieldCount(),
            newRel.getRowType().getFieldCount());
}