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

public void onMatch(RelOptRuleCall call) {
	final Aggregate aggregate = call.rel(0);
	final RelNode input = call.rel(1);

	// Distinct is "GROUP BY c1, c2" (where c1, c2 are a set of columns on
	// which the input is unique, i.e. contain a key) and has no aggregate
	// functions or the functions we enumerated. It can be removed.
	final RelNode newInput = convert(input, aggregate.getTraitSet().simplify());

	// If aggregate was projecting a subset of columns, add a project for the
	// same effect.
	final RelBuilder relBuilder = call.builder();
	relBuilder.push(newInput);
	List<Integer> projectIndices = new ArrayList<>(aggregate.getGroupSet().asList());
	for (AggregateCall aggCall : aggregate.getAggCallList()) {
		projectIndices.addAll(aggCall.getArgList());
	}
	relBuilder.project(relBuilder.fields(projectIndices));
	// Create a project if some of the columns have become
	// NOT NULL due to aggregate functions are removed
	relBuilder.convert(aggregate.getRowType(), true);
	call.transformTo(relBuilder.build());
}
 

@Override
public boolean matches(RelOptRuleCall call) {
	final Aggregate aggregate = call.rel(0);
	final RelNode input = call.rel(1);
	if (aggregate.getGroupCount() == 0 || aggregate.indicator ||
			aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
		return false;
	}
	for (AggregateCall aggCall : aggregate.getAggCallList()) {
		SqlKind aggCallKind = aggCall.getAggregation().getKind();
		// TODO supports more AggregateCalls
		boolean isAllowAggCall = aggCallKind == SqlKind.SUM ||
				aggCallKind == SqlKind.MIN ||
				aggCallKind == SqlKind.MAX ||
				aggCall.getAggregation() instanceof SqlAuxiliaryGroupAggFunction;
		if (!isAllowAggCall || aggCall.filterArg >= 0 || aggCall.getArgList().size() != 1) {
			return false;
		}
	}

	final RelMetadataQuery mq = call.getMetadataQuery();
	return SqlFunctions.isTrue(mq.areColumnsUnique(input, aggregate.getGroupSet()));
}
 

public AggregateNode(Compiler compiler, Aggregate rel) {
  super(compiler, rel);
  this.dataContext = compiler.getDataContext();

  ImmutableBitSet union = ImmutableBitSet.of();

  if (rel.getGroupSets() != null) {
    for (ImmutableBitSet group : rel.getGroupSets()) {
      union = union.union(group);
      groups.add(new Grouping(group));
    }
  }

  this.unionGroups = union;
  this.outputRowLength = unionGroups.cardinality()
      + rel.getAggCallList().size();

  ImmutableList.Builder<AccumulatorFactory> builder = ImmutableList.builder();
  for (AggregateCall aggregateCall : rel.getAggCallList()) {
    builder.add(getAccumulator(aggregateCall, false));
  }
  accumulatorFactories = builder.build();
}
 

/**
 * Compute which input fields are used by the aggregate.
 */
private ImmutableBitSet.Builder getInputFieldUsed(Aggregate aggregate, RelNode input) {
	// 1. group fields are always used
	final ImmutableBitSet.Builder inputFieldsUsed =
		aggregate.getGroupSet().rebuild();
	// 2. agg functions
	for (AggregateCall aggCall : aggregate.getAggCallList()) {
		for (int i : aggCall.getArgList()) {
			inputFieldsUsed.set(i);
		}
		if (aggCall.filterArg >= 0) {
			inputFieldsUsed.set(aggCall.filterArg);
		}
	}
	// 3. window time field if the aggregate is a group window aggregate.
	if (aggregate instanceof LogicalWindowAggregate) {
		inputFieldsUsed.set(getWindowTimeFieldIndex(((LogicalWindowAggregate) aggregate).getWindow(), input));
	}
	return inputFieldsUsed;
}
 

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

@Override
public Void visitAggregate(Aggregate aggregate, List<Void> inputStreams) throws Exception {
  beginAggregateStage(aggregate);
  pw.println("        if (_data != null) {");
  pw.println("        List<Object> curGroupValues = getGroupValues(_data);");
  pw.println("        if (!correlatedGroupedValues.containsKey(curGroupValues)) {");
  pw.println("          correlatedGroupedValues.put(curGroupValues, new ArrayList<CorrelatedValues>());");
  pw.println("        }");
  pw.println("        correlatedGroupedValues.get(curGroupValues).add(_data);");
  pw.println("        if (!state.containsKey(curGroupValues)) {");
  pw.println("          state.put(curGroupValues, new HashMap<String, Object>());");
  pw.println("        }");
  pw.println("        Map<String, Object> accumulators = state.get(curGroupValues);");
  for (AggregateCall call : aggregate.getAggCallList()) {
    aggregate(call);
  }
  pw.println("        }");
  endStage();
  return null;
}
 

private boolean hasNestedAggregations(Aggregate rel) {
  if (node instanceof SqlSelect) {
    final SqlNodeList selectList = ((SqlSelect) node).getSelectList();
    if (selectList != null) {
      final Set<Integer> aggregatesArgs = new HashSet<>();
      for (AggregateCall aggregateCall : rel.getAggCallList()) {
        aggregatesArgs.addAll(aggregateCall.getArgList());
      }
      for (int aggregatesArg : aggregatesArgs) {
        if (selectList.get(aggregatesArg) instanceof SqlBasicCall) {
          final SqlBasicCall call =
              (SqlBasicCall) selectList.get(aggregatesArg);
          for (SqlNode operand : call.getOperands()) {
            if (operand instanceof SqlCall
                && ((SqlCall) operand).getOperator() instanceof SqlAggFunction) {
              return true;
            }
          }
        }
      }
    }
  }
  return false;
}
 

private List<RelBuilder.AggCall> getNewAggCallList(
	Aggregate oldAggregate,
	RelBuilder relBuilder,
	Mapping mapping) {

	final List<RelBuilder.AggCall> newAggCallList = new ArrayList<>();

	for (AggregateCall aggCall : oldAggregate.getAggCallList()) {
		final RexNode filterArg = aggCall.filterArg < 0 ? null
			: relBuilder.field(Mappings.apply(mapping, aggCall.filterArg));
		newAggCallList.add(
			relBuilder
				.aggregateCall(
					aggCall.getAggregation(),
					relBuilder.fields(Mappings.apply2(mapping, aggCall.getArgList())))
				.distinct(aggCall.isDistinct())
				.filter(filterArg)
				.approximate(aggCall.isApproximate())
				.sort(relBuilder.fields(aggCall.collation))
				.as(aggCall.name));
	}
	return newAggCallList;
}
 

/**
 * Builds an aggregate query.
 *
 * @param e The Aggregate node
 * @param builder The SQL builder
 * @param selectList The precomputed group list
 * @param groupByList The precomputed select list
 * @return The aggregate query result
 */
protected Result buildAggregate(Aggregate e, Builder builder,
    List<SqlNode> selectList, List<SqlNode> groupByList) {
  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();
}
 

private static boolean isAggregateSupported(Aggregate aggregate, boolean allowFunctions) {
  if (!allowFunctions && !aggregate.getAggCallList().isEmpty()) {
    return false;
  }
  if (aggregate.getGroupType() != Aggregate.Group.SIMPLE) {
    return false;
  }
  // If any aggregate functions do not support splitting, bail out
  // If any aggregate call has a filter or is distinct, bail out
  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    if (aggregateCall.getAggregation().unwrap(SqlSplittableAggFunction.class)
        == null) {
      return false;
    }
    if (aggregateCall.filterArg >= 0 || aggregateCall.isDistinct()) {
      return false;
    }
  }
  return true;
}
 

/**
 * Takes a bitmap representing a set of input references and extracts the
 * ones that reference the group by columns in an aggregate.
 *
 * @param groupKey the original bitmap
 * @param aggRel   the aggregate
 * @param childKey sets bits from groupKey corresponding to group by columns
 */
public static void setAggChildKeys(
    ImmutableBitSet groupKey,
    Aggregate aggRel,
    ImmutableBitSet.Builder childKey) {
  List<AggregateCall> aggCalls = aggRel.getAggCallList();
  for (int bit : groupKey) {
    if (bit < aggRel.getGroupCount()) {
      // group by column
      childKey.set(bit);
    } else {
      // aggregate column -- set a bit for each argument being
      // aggregated
      AggregateCall agg = aggCalls.get(bit - aggRel.getGroupCount());
      for (Integer arg : agg.getArgList()) {
        childKey.set(arg);
      }
    }
  }
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final Values values = call.rel(1);
  Util.discard(values);
  final RelBuilder relBuilder = call.builder();
  final RexBuilder rexBuilder = relBuilder.getRexBuilder();

  final List<RexLiteral> literals = new ArrayList<>();
  for (final AggregateCall aggregateCall : aggregate.getAggCallList()) {
    switch (aggregateCall.getAggregation().getKind()) {
    case COUNT:
    case SUM0:
      literals.add(
          (RexLiteral) rexBuilder.makeLiteral(
              BigDecimal.ZERO, aggregateCall.getType(), false));
      break;

    case MIN:
    case MAX:
    case SUM:
      literals.add(rexBuilder.makeNullLiteral(aggregateCall.getType()));
      break;

    default:
      // Unknown what this aggregate call should do on empty Values. Bail out to be safe.
      return;
    }
  }

  call.transformTo(
      relBuilder.values(ImmutableList.of(literals), aggregate.getRowType())
          .build());

  // New plan is absolutely better than old plan.
  call.getPlanner().prune(aggregate);
}
 

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

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

@Override public boolean matches(final RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final Project project = call.rel(1);

  for (AggregateCall aggregateCall : aggregate.getAggCallList()) {
    final int singleArg = soleArgument(aggregateCall);
    if (singleArg >= 0
        && isThreeArgCase(project.getProjects().get(singleArg))) {
      return true;
    }
  }

  return false;
}
 

@Override public void onMatch(RelOptRuleCall call) {
  final Aggregate aggregate = call.rel(0);
  final Values values = call.rel(1);
  Util.discard(values);
  final RelBuilder relBuilder = call.builder();
  final RexBuilder rexBuilder = relBuilder.getRexBuilder();

  final List<RexLiteral> literals = new ArrayList<>();
  for (final AggregateCall aggregateCall : aggregate.getAggCallList()) {
    switch (aggregateCall.getAggregation().getKind()) {
    case COUNT:
    case SUM0:
      literals.add((RexLiteral) rexBuilder.makeLiteral(
          BigDecimal.ZERO, aggregateCall.getType(), false));
      break;

    case MIN:
    case MAX:
    case SUM:
      literals.add((RexLiteral) rexBuilder.makeCast(
          aggregateCall.getType(), rexBuilder.constantNull()));
      break;

    default:
      // Unknown what this aggregate call should do on empty Values. Bail out to be safe.
      return;
    }
  }

  call.transformTo(
      relBuilder.values(ImmutableList.of(literals), aggregate.getRowType())
          .build());

  // New plan is absolutely better than old plan.
  call.getPlanner().setImportance(aggregate, 0.0);
}
 

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

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

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

private void onMatch2(
    RelOptRuleCall call,
    Correlate correlate,
    RelNode leftInput,
    Project aggOutputProject,
    Aggregate aggregate) {
  if (generatedCorRels.contains(correlate)) {
    // This Correlate was generated by a previous invocation of
    // this rule. No further work to do.
    return;
  }

  setCurrent(call.getPlanner().getRoot(), correlate);

  // check for this pattern
  // The pattern matching could be simplified if rules can be applied
  // during decorrelation,
  //
  // CorrelateRel(left correlation, condition = true)
  //   leftInput
  //   Project-A (a RexNode)
  //     Aggregate (groupby (0), agg0(), agg1()...)

  // check aggOutputProj projects only one expression
  List<RexNode> aggOutputProjExprs = aggOutputProject.getProjects();
  if (aggOutputProjExprs.size() != 1) {
    return;
  }

  JoinRelType joinType = correlate.getJoinType();
  // corRel.getCondition was here, however Correlate was updated so it
  // never includes a join condition. The code was not modified for brevity.
  RexNode joinCond = relBuilder.literal(true);
  if ((joinType != JoinRelType.LEFT)
      || (joinCond != relBuilder.literal(true))) {
    return;
  }

  // check that the agg is on the entire input
  if (!aggregate.getGroupSet().isEmpty()) {
    return;
  }

  List<AggregateCall> aggCalls = aggregate.getAggCallList();
  Set<Integer> isCount = new HashSet<>();

  // remember the count() positions
  int i = -1;
  for (AggregateCall aggCall : aggCalls) {
    ++i;
    if (aggCall.getAggregation() instanceof SqlCountAggFunction) {
      isCount.add(i);
    }
  }

  // now rewrite the plan to
  //
  // Project-A' (all LHS plus transformed original projections,
  //             replacing references to count() with case statement)
  //   Correlate(left correlation, condition = true)
  //     leftInput
  //     Aggregate(groupby (0), agg0(), agg1()...)
  //
  List<RexNode> requiredNodes =
      correlate.getRequiredColumns().asList().stream()
          .map(ord -> relBuilder.getRexBuilder().makeInputRef(correlate, ord))
          .collect(Collectors.toList());
  Correlate newCorrelate = (Correlate) relBuilder.push(leftInput)
      .push(aggregate).correlate(correlate.getJoinType(),
          correlate.getCorrelationId(),
          requiredNodes).build();


  // remember this rel so we don't fire rule on it again
  // REVIEW jhyde 29-Oct-2007: rules should not save state; rule
  // should recognize patterns where it does or does not need to do
  // work
  generatedCorRels.add(newCorrelate);

  // need to update the mapCorToCorRel Update the output position
  // for the corVars: only pass on the corVars that are not used in
  // the join key.
  if (cm.mapCorToCorRel.get(correlate.getCorrelationId()) == correlate) {
    cm.mapCorToCorRel.put(correlate.getCorrelationId(), newCorrelate);
  }

  RelNode newOutput =
      aggregateCorrelatorOutput(newCorrelate, aggOutputProject, isCount);

  call.transformTo(newOutput);
}