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

public RelNode convert(RelNode rel) {
  final Aggregate agg = (Aggregate) rel;
  if (agg.getGroupSets().size() != 1) {
    // GROUPING SETS not supported; see
    // [CALCITE-734] Push GROUPING SETS to underlying SQL via JDBC adapter
    return null;
  }
  final RelTraitSet traitSet =
      agg.getTraitSet().replace(out);
  try {
    return new JdbcAggregate(rel.getCluster(), traitSet,
        convert(agg.getInput(), out), false, agg.getGroupSet(),
        agg.getGroupSets(), agg.getAggCallList());
  } catch (InvalidRelException e) {
    LOGGER.debug(e.toString());
    return null;
  }
}
 

protected void generateGroupBy(DremioRelToSqlConverter.Builder builder, Aggregate e) {
  List<SqlNode> groupByList = Expressions.list();
  final List<SqlNode> selectList = new ArrayList<>();
  for (int group : e.getGroupSet()) {
    final SqlNode field = builder.context.field(group);
    addSelect(selectList, field, e.getRowType());
    addGroupBy(groupByList, field);
  }

  for (AggregateCall aggCall : e.getAggCallList()) {
    SqlNode aggCallSqlNode = ((DremioContext) builder.context).toSql(aggCall, e);

    if (aggCall.getAggregation() instanceof SqlSingleValueAggFunction) {
      aggCallSqlNode = dialect.
        rewriteSingleValueExpr(aggCallSqlNode);
    }
    addSelect(selectList, aggCallSqlNode, e.getRowType());
  }
  builder.setSelect(new SqlNodeList(selectList, POS));
  if (!groupByList.isEmpty() || e.getAggCallList().isEmpty()) {
    // Some databases don't support "GROUP BY ()". We can omit it as long
    // as there is at least one aggregate function.
    builder.setGroupBy(new SqlNodeList(groupByList, POS));
  }
}
 

private String getGroupByIndices(Aggregate n) {
  StringBuilder res = new StringBuilder();
  int count = 0;
  for (int i : n.getGroupSet()) {
    if (++count > 1) {
      res.append(", ");
    }
    res.append(i);
  }
  return res.toString();
}
 

/**
 * Infers predicates for an Aggregate.
 *
 * <p>Pulls up predicates that only contains references to columns in the
 * GroupSet. For e.g.
 *
 * <blockquote><pre>
 * inputPullUpExprs : { a &gt; 7, b + c &lt; 10, a + e = 9}
 * groupSet         : { a, b}
 * pulledUpExprs    : { a &gt; 7}
 * </pre></blockquote>
 */
public RelOptPredicateList getPredicates(Aggregate agg, RelMetadataQuery mq) {
  final RelNode input = agg.getInput();
  final RexBuilder rexBuilder = agg.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> aggPullUpPredicates = new ArrayList<>();

  ImmutableBitSet groupKeys = agg.getGroupSet();
  if (groupKeys.isEmpty()) {
    // "GROUP BY ()" can convert an empty relation to a non-empty relation, so
    // it is not valid to pull up predicates. In particular, consider the
    // predicate "false": it is valid on all input rows (trivially - there are
    // no rows!) but not on the output (there is one row).
    return RelOptPredicateList.EMPTY;
  }
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(), agg.getRowType().getFieldCount());

  int i = 0;
  for (int j : groupKeys) {
    m.set(j, i++);
  }

  for (RexNode r : inputInfo.pulledUpPredicates) {
    ImmutableBitSet rCols = RelOptUtil.InputFinder.bits(r);
    if (groupKeys.contains(rCols)) {
      r = r.accept(new RexPermuteInputsShuttle(m, input));
      aggPullUpPredicates.add(r);
    }
  }
  return RelOptPredicateList.of(rexBuilder, aggPullUpPredicates);
}
 

public Double getRowCount(Aggregate rel, RelMetadataQuery mq) {
  ImmutableBitSet groupKey = rel.getGroupSet();

  // rowCount is the cardinality of the group by columns
  Double distinctRowCount =
      mq.getDistinctRowCount(rel.getInput(), groupKey, null);
  if (distinctRowCount == null) {
    distinctRowCount = mq.getRowCount(rel.getInput()) / 10;
  }

  // Grouping sets multiply
  distinctRowCount *= rel.getGroupSets().size();

  return distinctRowCount;
}
 

public List<Double> averageColumnSizes(Aggregate rel, RelMetadataQuery mq) {
  final List<Double> inputColumnSizes =
      mq.getAverageColumnSizesNotNull(rel.getInput());
  final ImmutableList.Builder<Double> list = ImmutableList.builder();
  for (int key : rel.getGroupSet()) {
    list.add(inputColumnSizes.get(key));
  }
  for (AggregateCall aggregateCall : rel.getAggCallList()) {
    list.add(averageTypeValueSize(aggregateCall.type));
  }
  return list.build();
}
 

public Result visitAggregate(Aggregate e) {
  // "select a, b, sum(x) from ( ... ) group by a, b"
  final Result x = visitChild(0, e.getInput());
  final Builder builder =
      x.builder(e, Clause.GROUP_BY);
  List<SqlNode> groupByList = Expressions.list();
  final List<SqlNode> selectList = new ArrayList<>();
  for (int group : e.getGroupSet()) {
    final SqlNode field = builder.context.field(group);
    addSelect(selectList, field, e.getRowType());
    groupByList.add(field);
  }
  for (AggregateCall aggCall : e.getAggCallList()) {
    SqlNode aggCallSqlNode = builder.context.toSql(aggCall);
    if (aggCall.getAggregation() instanceof SqlSingleValueAggFunction) {
      aggCallSqlNode =
          rewriteSingleValueExpr(aggCallSqlNode, dialect);
    }
    addSelect(selectList, aggCallSqlNode, e.getRowType());
  }
  builder.setSelect(new SqlNodeList(selectList, POS));
  if (!groupByList.isEmpty() || e.getAggCallList().isEmpty()) {
    // Some databases don't support "GROUP BY ()". We can omit it as long
    // as there is at least one aggregate function.
    builder.setGroupBy(new SqlNodeList(groupByList, POS));
  }
  return builder.result();
}
 

/**
 * 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;
}
 

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

public List<Double> averageColumnSizes(Aggregate rel, RelMetadataQuery mq) {
  final List<Double> inputColumnSizes =
      mq.getAverageColumnSizesNotNull(rel.getInput());
  final ImmutableList.Builder<Double> list = ImmutableList.builder();
  for (int key : rel.getGroupSet()) {
    list.add(inputColumnSizes.get(key));
  }
  for (AggregateCall aggregateCall : rel.getAggCallList()) {
    list.add(averageTypeValueSize(aggregateCall.type));
  }
  return list.build();
}
 

/**
 * @see #dispatch
 */
public Result visit(Aggregate e) {
  // "select a, b, sum(x) from ( ... ) group by a, b"
  final Result x = visitChild(0, e.getInput());
  final Builder builder;
  if (e.getInput() instanceof Project) {
    builder = x.builder(e);
    builder.clauses.add(Clause.GROUP_BY);
  } else {
    builder = x.builder(e, Clause.GROUP_BY);
  }
  List<SqlNode> groupByList = Expressions.list();
  final List<SqlNode> selectList = new ArrayList<>();
  for (int group : e.getGroupSet()) {
    final SqlNode field = builder.context.field(group);
    addSelect(selectList, field, e.getRowType());
    groupByList.add(field);
  }
  for (AggregateCall aggCall : e.getAggCallList()) {
    SqlNode aggCallSqlNode = builder.context.toSql(aggCall);
    if (aggCall.getAggregation() instanceof SqlSingleValueAggFunction) {
      aggCallSqlNode = dialect.
        rewriteSingleValueExpr(aggCallSqlNode);
    }
    addSelect(selectList, aggCallSqlNode, e.getRowType());
  }
  builder.setSelect(new SqlNodeList(selectList, POS));
  if (!groupByList.isEmpty() || e.getAggCallList().isEmpty()) {
    // Some databases don't support "GROUP BY ()". We can omit it as long
    // as there is at least one aggregate function.
    builder.setGroupBy(new SqlNodeList(groupByList, POS));
  }
  return builder.result();
}
 

/**
 * Builds the group list for an Aggregate node.
 *
 * @param e The Aggregate node
 * @param builder The SQL builder
 * @param groupByList output group list
 * @param selectList output select list
 */
protected void buildAggGroupList(Aggregate e, Builder builder,
    List<SqlNode> groupByList, List<SqlNode> selectList) {
  for (int group : e.getGroupSet()) {
    final SqlNode field = builder.context.field(group);
    addSelect(selectList, field, e.getRowType());
    groupByList.add(field);
  }
}
 

/** @see #dispatch */
public Result visit(Aggregate e) {
  // "select a, b, sum(x) from ( ... ) group by a, b"
  final Result x = visitChild(0, e.getInput());
  final Builder builder;
  if (e.getInput() instanceof Project) {
    builder = x.builder(e);
    builder.clauses.add(Clause.GROUP_BY);
  } else {
    builder = x.builder(e, Clause.GROUP_BY);
  }
  List<SqlNode> groupByList = FluentListUtils.list();
  final List<SqlNode> selectList = new ArrayList<>();
  for (int group : e.getGroupSet()) {
    final SqlNode field = builder.context.field(group);
    addSelect(selectList, field, e.getRowType());
    groupByList.add(field);
  }
  for (AggregateCall aggCall : e.getAggCallList()) {
    SqlNode aggCallSqlNode = builder.context.toSql(aggCall);
    if (aggCall.getAggregation() instanceof SqlSingleValueAggFunction) {
      aggCallSqlNode = dialect.
          rewriteSingleValueExpr(aggCallSqlNode);
    }
    addSelect(selectList, aggCallSqlNode, e.getRowType());
  }
  builder.setSelect(new SqlNodeList(selectList, POS));
  if (!groupByList.isEmpty() || e.getAggCallList().isEmpty()) {
    // Some databases don't support "GROUP BY ()". We can omit it as long
    // as there is at least one aggregate function.
    builder.setGroupBy(new SqlNodeList(groupByList, POS));
  }
  return builder.result();
}
 

/**
 * Infers predicates for an Aggregate.
 *
 * <p>Pulls up predicates that only contains references to columns in the
 * GroupSet. For e.g.
 *
 * <blockquote><pre>
 * inputPullUpExprs : { a &gt; 7, b + c &lt; 10, a + e = 9}
 * groupSet         : { a, b}
 * pulledUpExprs    : { a &gt; 7}
 * </pre></blockquote>
 */
public RelOptPredicateList getPredicates(Aggregate agg, RelMetadataQuery mq) {
  final RelNode input = agg.getInput();
  final RexBuilder rexBuilder = agg.getCluster().getRexBuilder();
  final RelOptPredicateList inputInfo = mq.getPulledUpPredicates(input);
  final List<RexNode> aggPullUpPredicates = new ArrayList<>();

  ImmutableBitSet groupKeys = agg.getGroupSet();
  if (groupKeys.isEmpty()) {
    // "GROUP BY ()" can convert an empty relation to a non-empty relation, so
    // it is not valid to pull up predicates. In particular, consider the
    // predicate "false": it is valid on all input rows (trivially - there are
    // no rows!) but not on the output (there is one row).
    return RelOptPredicateList.EMPTY;
  }
  Mapping m = Mappings.create(MappingType.PARTIAL_FUNCTION,
      input.getRowType().getFieldCount(), agg.getRowType().getFieldCount());

  int i = 0;
  for (int j : groupKeys) {
    m.set(j, i++);
  }

  for (RexNode r : inputInfo.pulledUpPredicates) {
    ImmutableBitSet rCols = RelOptUtil.InputFinder.bits(r);
    if (groupKeys.contains(rCols)) {
      r = r.accept(new RexPermuteInputsShuttle(m, input));
      aggPullUpPredicates.add(r);
    }
  }
  return RelOptPredicateList.of(rexBuilder, aggPullUpPredicates);
}
 

/**
 * Given an {@link org.apache.calcite.rel.core.Aggregate}
 * and the ordinals of the arguments to a
 * particular call to an aggregate function, creates a 'select distinct'
 * relational expression which projects the group columns and those
 * arguments but nothing else.
 *
 * <p>For example, given
 *
 * <blockquote>
 * <pre>select f0, count(distinct f1), count(distinct f2)
 * from t group by f0</pre>
 * </blockquote>
 *
 * <p>and the argument list
 *
 * <blockquote>{2}</blockquote>
 *
 * <p>returns
 *
 * <blockquote>
 * <pre>select distinct f0, f2 from t</pre>
 * </blockquote>
 *
 * <p>The <code>sourceOf</code> map is populated with the source of each
 * column; in this case sourceOf.get(0) = 0, and sourceOf.get(1) = 2.
 *
 * @param relBuilder Relational expression builder
 * @param aggregate Aggregate relational expression
 * @param argList   Ordinals of columns to make distinct
 * @param filterArg Ordinal of column to filter on, or -1
 * @param sourceOf  Out parameter, is populated with a map of where each
 *                  output field came from
 * @return Aggregate relational expression which projects the required
 * columns
 */
private RelBuilder createSelectDistinct(RelBuilder relBuilder,
		Aggregate aggregate, List<Integer> argList, int filterArg,
		Map<Integer, Integer> sourceOf) {
	relBuilder.push(aggregate.getInput());
	final List<Pair<RexNode, String>> projects = new ArrayList<>();
	final List<RelDataTypeField> childFields =
			relBuilder.peek().getRowType().getFieldList();
	for (int i : aggregate.getGroupSet()) {
		sourceOf.put(i, projects.size());
		projects.add(RexInputRef.of2(i, childFields));
	}
	if (filterArg >= 0) {
		sourceOf.put(filterArg, projects.size());
		projects.add(RexInputRef.of2(filterArg, childFields));
	}
	for (Integer arg : argList) {
		if (filterArg >= 0) {
			// Implement
			//   agg(DISTINCT arg) FILTER $f
			// by generating
			//   SELECT DISTINCT ... CASE WHEN $f THEN arg ELSE NULL END AS arg
			// and then applying
			//   agg(arg)
			// as usual.
			//
			// It works except for (rare) agg functions that need to see null
			// values.
			final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
			final RexInputRef filterRef = RexInputRef.of(filterArg, childFields);
			final Pair<RexNode, String> argRef = RexInputRef.of2(arg, childFields);
			RexNode condition =
					rexBuilder.makeCall(SqlStdOperatorTable.CASE, filterRef,
							argRef.left,
							rexBuilder.ensureType(argRef.left.getType(),
									rexBuilder.makeCast(argRef.left.getType(),
											rexBuilder.constantNull()),
									true));
			sourceOf.put(arg, projects.size());
			projects.add(Pair.of(condition, "i$" + argRef.right));
			continue;
		}
		if (sourceOf.get(arg) != null) {
			continue;
		}
		sourceOf.put(arg, projects.size());
		projects.add(RexInputRef.of2(arg, childFields));
	}
	relBuilder.project(Pair.left(projects), Pair.right(projects));

	// Get the distinct values of the GROUP BY fields and the arguments
	// to the agg functions.
	relBuilder.push(
			aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), false,
					ImmutableBitSet.range(projects.size()),
					null, com.google.common.collect.ImmutableList.<AggregateCall>of()));
	return relBuilder;
}
 

/**
 * Given an {@link org.apache.calcite.rel.core.Aggregate}
 * and the ordinals of the arguments to a
 * particular call to an aggregate function, creates a 'select distinct'
 * relational expression which projects the group columns and those
 * arguments but nothing else.
 *
 * <p>For example, given
 *
 * <blockquote>
 * <pre>select f0, count(distinct f1), count(distinct f2)
 * from t group by f0</pre>
 * </blockquote>
 *
 * <p>and the argument list
 *
 * <blockquote>{2}</blockquote>
 *
 * <p>returns
 *
 * <blockquote>
 * <pre>select distinct f0, f2 from t</pre>
 * </blockquote>
 *
 * <p>The <code>sourceOf</code> map is populated with the source of each
 * column; in this case sourceOf.get(0) = 0, and sourceOf.get(1) = 2.
 *
 * @param relBuilder Relational expression builder
 * @param aggregate Aggregate relational expression
 * @param argList   Ordinals of columns to make distinct
 * @param filterArg Ordinal of column to filter on, or -1
 * @param sourceOf  Out parameter, is populated with a map of where each
 *                  output field came from
 * @return Aggregate relational expression which projects the required
 * columns
 */
private RelBuilder createSelectDistinct(RelBuilder relBuilder,
		Aggregate aggregate, List<Integer> argList, int filterArg,
		Map<Integer, Integer> sourceOf) {
	relBuilder.push(aggregate.getInput());
	final List<Pair<RexNode, String>> projects = new ArrayList<>();
	final List<RelDataTypeField> childFields =
			relBuilder.peek().getRowType().getFieldList();
	for (int i : aggregate.getGroupSet()) {
		sourceOf.put(i, projects.size());
		projects.add(RexInputRef.of2(i, childFields));
	}
	if (filterArg >= 0) {
		sourceOf.put(filterArg, projects.size());
		projects.add(RexInputRef.of2(filterArg, childFields));
	}
	for (Integer arg : argList) {
		if (filterArg >= 0) {
			// Implement
			//   agg(DISTINCT arg) FILTER $f
			// by generating
			//   SELECT DISTINCT ... CASE WHEN $f THEN arg ELSE NULL END AS arg
			// and then applying
			//   agg(arg)
			// as usual.
			//
			// It works except for (rare) agg functions that need to see null
			// values.
			final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
			final RexInputRef filterRef = RexInputRef.of(filterArg, childFields);
			final Pair<RexNode, String> argRef = RexInputRef.of2(arg, childFields);
			RexNode condition =
					rexBuilder.makeCall(SqlStdOperatorTable.CASE, filterRef,
							argRef.left,
							rexBuilder.makeNullLiteral(argRef.left.getType()));
			sourceOf.put(arg, projects.size());
			projects.add(Pair.of(condition, "i$" + argRef.right));
			continue;
		}
		if (sourceOf.get(arg) != null) {
			continue;
		}
		sourceOf.put(arg, projects.size());
		projects.add(RexInputRef.of2(arg, childFields));
	}
	relBuilder.project(Pair.left(projects), Pair.right(projects));

	// Get the distinct values of the GROUP BY fields and the arguments
	// to the agg functions.
	relBuilder.push(
			aggregate.copy(aggregate.getTraitSet(), relBuilder.build(),
					ImmutableBitSet.range(projects.size()),
					null, com.google.common.collect.ImmutableList.<AggregateCall>of()));
	return relBuilder;
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final Project project = call.rel(1);
  final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
  final List<AggregateCall> newCalls =
      new ArrayList<>(aggregate.getAggCallList().size());
  final List<RexNode> newProjects = new ArrayList<>(project.getProjects());
  final List<RexNode> newCasts = new ArrayList<>();

  for (int fieldNumber : aggregate.getGroupSet()) {
    newCasts.add(
        rexBuilder.makeInputRef(
            project.getProjects().get(fieldNumber).getType(), fieldNumber));
  }

  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    AggregateCall newCall =
        transform(aggregateCall, project, newProjects);

    // Possibly CAST the new aggregator to an appropriate type.
    final int i = newCasts.size();
    final RelDataType oldType =
        aggregate.getRowType().getFieldList().get(i).getType();
    if (newCall == null) {
      newCalls.add(aggregateCall);
      newCasts.add(rexBuilder.makeInputRef(oldType, i));
    } else {
      newCalls.add(newCall);
      newCasts.add(
          rexBuilder.makeCast(oldType,
              rexBuilder.makeInputRef(newCall.getType(), i)));
    }
  }

  if (newCalls.equals(aggregate.getAggCallList())) {
    return;
  }

  final RelBuilder relBuilder = call.builder()
      .push(project.getInput())
      .project(newProjects);

  final RelBuilder.GroupKey groupKey =
      relBuilder.groupKey(aggregate.getGroupSet(),
          (Iterable<ImmutableBitSet>) aggregate.getGroupSets());

  relBuilder.aggregate(groupKey, newCalls)
      .convert(aggregate.getRowType(), false);

  call.transformTo(relBuilder.build());
  call.getPlanner().prune(aggregate);
}
 

/**
 * Given an {@link org.apache.calcite.rel.core.Aggregate}
 * and the ordinals of the arguments to a
 * particular call to an aggregate function, creates a 'select distinct'
 * relational expression which projects the group columns and those
 * arguments but nothing else.
 *
 * <p>For example, given
 *
 * <blockquote>
 * <pre>select f0, count(distinct f1), count(distinct f2)
 * from t group by f0</pre>
 * </blockquote>
 *
 * <p>and the argument list
 *
 * <blockquote>{2}</blockquote>
 *
 * <p>returns
 *
 * <blockquote>
 * <pre>select distinct f0, f2 from t</pre>
 * </blockquote>
 *
 * <p>The <code>sourceOf</code> map is populated with the source of each
 * column; in this case sourceOf.get(0) = 0, and sourceOf.get(1) = 2.
 *
 * @param relBuilder Relational expression builder
 * @param aggregate Aggregate relational expression
 * @param argList   Ordinals of columns to make distinct
 * @param filterArg Ordinal of column to filter on, or -1
 * @param sourceOf  Out parameter, is populated with a map of where each
 *                  output field came from
 * @return Aggregate relational expression which projects the required
 * columns
 */
private RelBuilder createSelectDistinct(RelBuilder relBuilder, Aggregate aggregate, List<Integer> argList,
        int filterArg, Map<Integer, Integer> sourceOf) {
    relBuilder.push(aggregate.getInput());
    final List<Pair<RexNode, String>> projects = new ArrayList<>();
    final List<RelDataTypeField> childFields = relBuilder.peek().getRowType().getFieldList();
    for (int i : aggregate.getGroupSet()) {
        sourceOf.put(i, projects.size());
        projects.add(RexInputRef.of2(i, childFields));
    }
    for (Integer arg : argList) {
        if (filterArg >= 0) {
            // Implement
            // agg(DISTINCT arg) FILTER $f
            // by generating
            // SELECT DISTINCT ... CASE WHEN $f THEN arg ELSE NULL END AS arg
            // and then applying
            // agg(arg)
            // as usual.
            //
            // It works except for (rare) agg functions that need to see null
            // values.
            final RexBuilder rexBuilder = aggregate.getCluster().getRexBuilder();
            final RexInputRef filterRef = RexInputRef.of(filterArg, childFields);
            final Pair<RexNode, String> argRef = RexInputRef.of2(arg, childFields);
            RexNode condition = rexBuilder.makeCall(SqlStdOperatorTable.CASE, filterRef, argRef.left,
                    rexBuilder.ensureType(argRef.left.getType(),
                            rexBuilder.makeCast(argRef.left.getType(), rexBuilder.constantNull()), true));
            sourceOf.put(arg, projects.size());
            projects.add(Pair.of(condition, "i$" + argRef.right));
            continue;
        }
        if (sourceOf.get(arg) != null) {
            continue;
        }
        sourceOf.put(arg, projects.size());
        projects.add(RexInputRef.of2(arg, childFields));
    }
    relBuilder.project(Pair.left(projects), Pair.right(projects));

    // Get the distinct values of the GROUP BY fields and the arguments
    // to the agg functions.
    relBuilder.push(aggregate.copy(aggregate.getTraitSet(), relBuilder.build(), false,
            ImmutableBitSet.range(projects.size()), null, ImmutableList.of()));
    return relBuilder;
}
 

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

/**
 * Do a shallow clone of oldAggRel and update aggCalls. Could be refactored
 * into Aggregate and subclasses - but it's only needed for some
 * subclasses.
 *
 * @param oldAggRel AggregateRel to clone.
 * @param inputRel  Input relational expression
 * @param newCalls  New list of AggregateCalls
 * @return shallow clone with new list of AggregateCalls.
 */
protected Aggregate newAggregateRel(
    Aggregate oldAggRel,
    RelNode inputRel,
    List<AggregateCall> newCalls) {
  RelOptCluster cluster = inputRel.getCluster();
  return new LogicalAggregate(cluster, cluster.traitSetOf(Convention.NONE),
      inputRel, oldAggRel.indicator, oldAggRel.getGroupSet(), oldAggRel.getGroupSets(), newCalls);
}