Java Code Examples for org.apache.calcite.rel.core.Aggregate#copy()

/**
 * Convert aggregate with AUXILIARY_GROUP to regular aggregate.
 * Return original aggregate and null project if the given aggregate does not contain AUXILIARY_GROUP,
 * else new aggregate without AUXILIARY_GROUP and a project to permute output columns if needed.
 */
private Pair<Aggregate, List<RexNode>> toRegularAggregate(Aggregate aggregate) {
	Tuple2<int[], Seq<AggregateCall>> auxGroupAndRegularAggCalls = AggregateUtil.checkAndSplitAggCalls(aggregate);
	final int[] auxGroup = auxGroupAndRegularAggCalls._1;
	final Seq<AggregateCall> regularAggCalls = auxGroupAndRegularAggCalls._2;
	if (auxGroup.length != 0) {
		int[] fullGroupSet = AggregateUtil.checkAndGetFullGroupSet(aggregate);
		ImmutableBitSet newGroupSet = ImmutableBitSet.of(fullGroupSet);
		List<AggregateCall> aggCalls = JavaConverters.seqAsJavaListConverter(regularAggCalls).asJava();
		final Aggregate newAgg = aggregate.copy(
				aggregate.getTraitSet(),
				aggregate.getInput(),
				aggregate.indicator,
				newGroupSet,
				com.google.common.collect.ImmutableList.of(newGroupSet),
				aggCalls);
		final List<RelDataTypeField> aggFields = aggregate.getRowType().getFieldList();
		final List<RexNode> projectAfterAgg = new ArrayList<>();
		for (int i = 0; i < fullGroupSet.length; ++i) {
			int group = fullGroupSet[i];
			int index = newGroupSet.indexOf(group);
			projectAfterAgg.add(new RexInputRef(index, aggFields.get(i).getType()));
		}
		int fieldCntOfAgg = aggFields.size();
		for (int i = fullGroupSet.length; i < fieldCntOfAgg; ++i) {
			projectAfterAgg.add(new RexInputRef(i, aggFields.get(i).getType()));
		}
		Preconditions.checkArgument(projectAfterAgg.size() == fieldCntOfAgg);
		return new Pair<>(newAgg, projectAfterAgg);
	} else {
		return new Pair<>(aggregate, null);
	}
}
 

private Aggregate mergeAggregate(Aggregate aggregate1, Aggregate aggregate2) {
  //Support only simple groups
  if (aggregate1.getGroupType() != Aggregate.Group.SIMPLE
      || aggregate2.getGroupType() != Aggregate.Group.SIMPLE) {
    return null;
  }

  final int callLen1 = aggregate1.getAggCallList().size();
  final int callLen2 = aggregate2.getAggCallList().size();
  final List<AggregateCall> newAggCalls = Lists.newArrayList();
  if (callLen1 <= callLen2) {
    //Create new Call list
    for (AggregateCall call : aggregate1.getAggCallList()) {
      AggregateCall newAggCall = getMergedAggCall(aggregate2, call);
      if (newAggCall == null) {
        return null;
      } else {
        newAggCalls.add(newAggCall);
      }
    }

    //Create new groupSets
    ImmutableBitSet.Builder groupSetsBuilder = ImmutableBitSet.builder();
    for (int key : aggregate1.getGroupSet()) {
      try {
        groupSetsBuilder.set(aggregate2.getGroupSet().nth(key));
      } catch (IndexOutOfBoundsException e) {
        return null;
      }
    }
    final ImmutableBitSet newGroupSets = groupSetsBuilder.build();
    return aggregate1.copy(aggregate1.getTraitSet(), aggregate2.getInput(),
        aggregate1.indicator, newGroupSets, ImmutableList.of(newGroupSets), newAggCalls);
  } else {
    return null;
  }
}
 

/**
 * Pushes a {@link org.apache.calcite.rel.core.Filter}
 * past a {@link org.apache.calcite.rel.core.Aggregate}.
 */
RelNode filterAggregateTranspose(RelOptRuleCall call,
                                 Filter filterRel,
                                 Aggregate aggRel) {
  final List<RexNode> conditions =
      RelOptUtil.conjunctions(filterRel.getCondition());
  final RexBuilder rexBuilder = filterRel.getCluster().getRexBuilder();
  final List<RelDataTypeField> origFields =
      aggRel.getRowType().getFieldList();
  final int[] adjustments = new int[origFields.size()];
  int i = 0;
  for (int key : aggRel.getGroupSet()) {
    adjustments[i] = key - i;
    i++;
  }
  final List<RexNode> pushedConditions = Lists.newArrayList();

  for (RexNode condition : conditions) {
    ImmutableBitSet rCols = RelOptUtil.InputFinder.bits(condition);
    if (canPush(aggRel, rCols)) {
      pushedConditions.add(
          condition.accept(
              new RelOptUtil.RexInputConverter(rexBuilder, origFields,
                  aggRel.getInput(0).getRowType().getFieldList(),
                  adjustments)));
    } else {
      return null;
    }
  }

  final RelBuilder builder = call.builder();
  RelNode rel =
      builder.push(aggRel.getInput()).filter(pushedConditions).build();
  if (rel == aggRel.getInput(0)) {
    return null;
  }
  rel = aggRel.copy(aggRel.getTraitSet(), ImmutableList.of(rel));
  return rel;
}
 

/**
 * Convert aggregate with AUXILIARY_GROUP to regular aggregate.
 * Return original aggregate and null project if the given aggregate does not contain AUXILIARY_GROUP,
 * else new aggregate without AUXILIARY_GROUP and a project to permute output columns if needed.
 */
private Pair<Aggregate, List<RexNode>> toRegularAggregate(Aggregate aggregate) {
	Tuple2<int[], Seq<AggregateCall>> auxGroupAndRegularAggCalls = AggregateUtil.checkAndSplitAggCalls(aggregate);
	final int[] auxGroup = auxGroupAndRegularAggCalls._1;
	final Seq<AggregateCall> regularAggCalls = auxGroupAndRegularAggCalls._2;
	if (auxGroup.length != 0) {
		int[] fullGroupSet = AggregateUtil.checkAndGetFullGroupSet(aggregate);
		ImmutableBitSet newGroupSet = ImmutableBitSet.of(fullGroupSet);
		List<AggregateCall> aggCalls = JavaConverters.seqAsJavaListConverter(regularAggCalls).asJava();
		final Aggregate newAgg = aggregate.copy(
				aggregate.getTraitSet(),
				aggregate.getInput(),
				aggregate.indicator,
				newGroupSet,
				com.google.common.collect.ImmutableList.of(newGroupSet),
				aggCalls);
		final List<RelDataTypeField> aggFields = aggregate.getRowType().getFieldList();
		final List<RexNode> projectAfterAgg = new ArrayList<>();
		for (int i = 0; i < fullGroupSet.length; ++i) {
			int group = fullGroupSet[i];
			int index = newGroupSet.indexOf(group);
			projectAfterAgg.add(new RexInputRef(index, aggFields.get(i).getType()));
		}
		int fieldCntOfAgg = aggFields.size();
		for (int i = fullGroupSet.length; i < fieldCntOfAgg; ++i) {
			projectAfterAgg.add(new RexInputRef(i, aggFields.get(i).getType()));
		}
		Preconditions.checkArgument(projectAfterAgg.size() == fieldCntOfAgg);
		return new Pair<>(newAgg, projectAfterAgg);
	} else {
		return new Pair<>(aggregate, null);
	}
}
 

public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final DruidQuery query = call.rel(1);
  final RelNode topDruidNode = query.getTopNode();
  final Project project = topDruidNode instanceof Project ? (Project) topDruidNode : null;
  if (!DruidQuery.isValidSignature(query.signature() + 'a')) {
    return;
  }

  if (aggregate.getGroupSets().size() != 1) {
    return;
  }
  if (DruidQuery
      .computeProjectGroupSet(project, aggregate.getGroupSet(), query.table.getRowType(), query)
      == null) {
    return;
  }
  final List<String> aggNames = Util
      .skip(aggregate.getRowType().getFieldNames(), aggregate.getGroupSet().cardinality());
  if (DruidQuery.computeDruidJsonAgg(aggregate.getAggCallList(), aggNames, project, query)
      == null) {
    return;
  }
  final RelNode newAggregate = aggregate
      .copy(aggregate.getTraitSet(), ImmutableList.of(query.getTopNode()));
  call.transformTo(DruidQuery.extendQuery(query, newAggregate));
}
 

public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final Project project = call.rel(1);
  final DruidQuery query = call.rel(2);
  if (!DruidQuery.isValidSignature(query.signature() + 'p' + 'a')) {
    return;
  }
  if (aggregate.getGroupSets().size() != 1) {
    return;
  }
  if (DruidQuery
      .computeProjectGroupSet(project, aggregate.getGroupSet(), query.table.getRowType(), query)
      == null) {
    return;
  }
  final List<String> aggNames = Util
      .skip(aggregate.getRowType().getFieldNames(), aggregate.getGroupSet().cardinality());
  if (DruidQuery.computeDruidJsonAgg(aggregate.getAggCallList(), aggNames, project, query)
      == null) {
    return;
  }
  final RelNode newProject = project.copy(project.getTraitSet(),
          ImmutableList.of(Util.last(query.rels)));
  final RelNode newAggregate = aggregate.copy(aggregate.getTraitSet(),
          ImmutableList.of(newProject));
  List<Integer> filterRefs = getFilterRefs(aggregate.getAggCallList());
  final DruidQuery query2;
  if (filterRefs.size() > 0) {
    query2 = optimizeFilteredAggregations(call, query, (Project) newProject,
        (Aggregate) newAggregate);
  } else {
    final DruidQuery query1 = DruidQuery.extendQuery(query, newProject);
    query2 = DruidQuery.extendQuery(query1, newAggregate);
  }
  call.transformTo(query2);
}
 

@Override
protected RelNode rewriteQuery(RelBuilder relBuilder, RexBuilder rexBuilder, RexSimplify simplify,
        RelMetadataQuery mq, RexNode compensationColumnsEquiPred, RexNode otherCompensationPred,
        Project topProject, RelNode node, BiMap<RelTableRef, RelTableRef> queryToViewTableMapping,
        EquivalenceClasses viewEC, EquivalenceClasses queryEC) {
    Aggregate aggregate = (Aggregate) node;

    // Our target node is the node below the root, which should have the maximum
    // number of available expressions in the tree in order to maximize our
    // number of rewritings.
    // If the program is available, we execute it to maximize rewriting opportunities.
    // For instance, a program might pull up all the expressions that are below the
    // aggregate so we can introduce compensation filters easily. This is important
    // depending on the planner strategy.
    RelNode newAggregateInput = aggregate.getInput(0);
    RelNode target = aggregate.getInput(0);
    if (unionRewritingPullProgram != null) {
        final HepPlanner tmpPlanner = new HepPlanner(unionRewritingPullProgram);
        tmpPlanner.setRoot(newAggregateInput);
        newAggregateInput = tmpPlanner.findBestExp();
        target = newAggregateInput.getInput(0);
    }

    // We need to check that all columns required by compensating predicates
    // are contained in the query.
    List<RexNode> queryExprs = extractReferences(rexBuilder, target);
    if (!compensationColumnsEquiPred.isAlwaysTrue()) {
        compensationColumnsEquiPred = rewriteExpression(rexBuilder, mq, target, target, queryExprs,
                queryToViewTableMapping, queryEC, false, compensationColumnsEquiPred);
        if (compensationColumnsEquiPred == null) {
            // Skip it
            return null;
        }
    }
    // For the rest, we use the query equivalence classes
    if (!otherCompensationPred.isAlwaysTrue()) {
        otherCompensationPred = rewriteExpression(rexBuilder, mq, target, target, queryExprs,
                queryToViewTableMapping, viewEC, true, otherCompensationPred);
        if (otherCompensationPred == null) {
            // Skip it
            return null;
        }
    }
    final RexNode queryCompensationPred = RexUtil.not(RexUtil.composeConjunction(rexBuilder,
            ImmutableList.of(compensationColumnsEquiPred, otherCompensationPred)));

    // Generate query rewriting.
    RelNode rewrittenPlan = relBuilder.push(target)
            .filter(simplify.simplifyUnknownAsFalse(queryCompensationPred)).build();
    if (unionRewritingPullProgram != null) {
        return aggregate.copy(aggregate.getTraitSet(), ImmutableList
                .of(newAggregateInput.copy(newAggregateInput.getTraitSet(), ImmutableList.of(rewrittenPlan))));
    }
    return aggregate.copy(aggregate.getTraitSet(), ImmutableList.of(rewrittenPlan));
}
 

@Override protected RelNode rewriteQuery(
    RelBuilder relBuilder,
    RexBuilder rexBuilder,
    RexSimplify simplify,
    RelMetadataQuery mq,
    RexNode compensationColumnsEquiPred,
    RexNode otherCompensationPred,
    Project topProject,
    RelNode node,
    BiMap<RelTableRef, RelTableRef> queryToViewTableMapping,
    EquivalenceClasses viewEC, EquivalenceClasses queryEC) {
  Aggregate aggregate = (Aggregate) node;

  // Our target node is the node below the root, which should have the maximum
  // number of available expressions in the tree in order to maximize our
  // number of rewritings.
  // If the program is available, we execute it to maximize rewriting opportunities.
  // For instance, a program might pull up all the expressions that are below the
  // aggregate so we can introduce compensation filters easily. This is important
  // depending on the planner strategy.
  RelNode newAggregateInput = aggregate.getInput(0);
  RelNode target = aggregate.getInput(0);
  if (unionRewritingPullProgram != null) {
    final HepPlanner tmpPlanner = new HepPlanner(unionRewritingPullProgram);
    tmpPlanner.setRoot(newAggregateInput);
    newAggregateInput = tmpPlanner.findBestExp();
    target = newAggregateInput.getInput(0);
  }

  // We need to check that all columns required by compensating predicates
  // are contained in the query.
  List<RexNode> queryExprs = extractReferences(rexBuilder, target);
  if (!compensationColumnsEquiPred.isAlwaysTrue()) {
    compensationColumnsEquiPred = rewriteExpression(rexBuilder, mq,
        target, target, queryExprs, queryToViewTableMapping, queryEC, false,
        compensationColumnsEquiPred);
    if (compensationColumnsEquiPred == null) {
      // Skip it
      return null;
    }
  }
  // For the rest, we use the query equivalence classes
  if (!otherCompensationPred.isAlwaysTrue()) {
    otherCompensationPred = rewriteExpression(rexBuilder, mq,
        target, target, queryExprs, queryToViewTableMapping, viewEC, true,
        otherCompensationPred);
    if (otherCompensationPred == null) {
      // Skip it
      return null;
    }
  }
  final RexNode queryCompensationPred = RexUtil.not(
      RexUtil.composeConjunction(rexBuilder,
          ImmutableList.of(compensationColumnsEquiPred,
              otherCompensationPred)));

  // Generate query rewriting.
  RelNode rewrittenPlan = relBuilder
      .push(target)
      .filter(simplify.simplifyUnknownAsFalse(queryCompensationPred))
      .build();
  if (unionRewritingPullProgram != null) {
    return aggregate.copy(aggregate.getTraitSet(),
        ImmutableList.of(
            newAggregateInput.copy(newAggregateInput.getTraitSet(),
                ImmutableList.of(rewrittenPlan))));
  }
  return aggregate.copy(aggregate.getTraitSet(), ImmutableList.of(rewrittenPlan));
}
 

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

public void onMatch(RelOptRuleCall call) {
  final Aggregate topAgg = call.rel(0);
  final Aggregate bottomAgg = call.rel(1);
  if (topAgg.getGroupCount() > bottomAgg.getGroupCount()) {
    return;
  }

  final ImmutableBitSet bottomGroupSet = bottomAgg.getGroupSet();
  final Map<Integer, Integer> map = new HashMap<>();
  bottomGroupSet.forEach(v -> map.put(map.size(), v));
  for (int k : topAgg.getGroupSet()) {
    if (!map.containsKey(k)) {
      return;
    }
  }

  // top aggregate keys must be subset of lower aggregate keys
  final ImmutableBitSet topGroupSet = topAgg.getGroupSet().permute(map);
  if (!bottomGroupSet.contains(topGroupSet)) {
    return;
  }

  boolean hasEmptyGroup = topAgg.getGroupSets()
      .stream().anyMatch(n -> n.isEmpty());

  final List<AggregateCall> finalCalls = new ArrayList<>();
  for (AggregateCall topCall : topAgg.getAggCallList()) {
    if (!isAggregateSupported(topCall)
        || topCall.getArgList().size() == 0) {
      return;
    }
    // Make sure top aggregate argument refers to one of the aggregate
    int bottomIndex = topCall.getArgList().get(0) - bottomGroupSet.cardinality();
    if (bottomIndex >= bottomAgg.getAggCallList().size()
        || bottomIndex < 0) {
      return;
    }
    AggregateCall bottomCall = bottomAgg.getAggCallList().get(bottomIndex);
    // Should not merge if top agg with empty group keys and the lower agg
    // function is COUNT, because in case of empty input for lower agg,
    // the result is empty, if we merge them, we end up with 1 result with
    // 0, which is wrong.
    if (!isAggregateSupported(bottomCall)
        || (bottomCall.getAggregation() == SqlStdOperatorTable.COUNT
             && hasEmptyGroup)) {
      return;
    }
    SqlSplittableAggFunction splitter = Objects.requireNonNull(
        bottomCall.getAggregation().unwrap(SqlSplittableAggFunction.class));
    AggregateCall finalCall = splitter.merge(topCall, bottomCall);
    // fail to merge the aggregate call, bail out
    if (finalCall == null) {
      return;
    }
    finalCalls.add(finalCall);
  }

  // re-map grouping sets
  ImmutableList<ImmutableBitSet> newGroupingSets = null;
  if (topAgg.getGroupType() != Group.SIMPLE) {
    newGroupingSets =
        ImmutableBitSet.ORDERING.immutableSortedCopy(
            ImmutableBitSet.permute(topAgg.getGroupSets(), map));
  }

  final Aggregate finalAgg =
      topAgg.copy(topAgg.getTraitSet(), bottomAgg.getInput(), topGroupSet,
          newGroupingSets, finalCalls);
  call.transformTo(finalAgg);
}
 

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