Java Code Examples for org.apache.calcite.rel.core.AggregateCall#create()

@Override
public RelNode visit(LogicalAggregate aggregate) {
  RelNode input = aggregate.getInput().accept(this);
  RelDataType incomingRowType = input.getRowType();
  RelDataTypeField modField = incomingRowType.getField(UPDATE_COLUMN, false, false);
  if (modField == null) {
    return aggregate;
  }

  final AggregateCall aggCall = AggregateCall.create(SqlStdOperatorTable.MAX, false, ImmutableList.of(modField.getIndex()), -1, modField.getType(), UPDATE_COLUMN);
  final List<AggregateCall> aggCalls = FluentIterable.from(aggregate.getAggCallList())
    .append(aggCall)
    .toList();
  return aggregate.copy(
    aggregate.getTraitSet(),
    input,
    aggregate.indicator,
    aggregate.getGroupSet(),
    null,
    aggCalls
  );
}
 

private AggregateCall createAggregateCallWithBinding(
    RelDataTypeFactory typeFactory,
    SqlAggFunction aggFunction,
    RelDataType operandType,
    Aggregate oldAggRel,
    AggregateCall oldCall,
    int argOrdinal,
    int filter) {
  final Aggregate.AggCallBinding binding =
      new Aggregate.AggCallBinding(typeFactory, aggFunction,
          ImmutableList.of(operandType), oldAggRel.getGroupCount(),
          filter >= 0);
  return AggregateCall.create(aggFunction,
      oldCall.isDistinct(),
      oldCall.isApproximate(),
      ImmutableIntList.of(argOrdinal),
      filter,
      oldCall.collation,
      aggFunction.inferReturnType(binding),
      null);
}
 

private RexNode getSumAggregatedRexNode(Aggregate oldAggRel,
    AggregateCall oldCall,
    List<AggregateCall> newCalls,
    Map<AggregateCall, RexNode> aggCallMapping,
    RexBuilder rexBuilder,
    int argOrdinal,
    int filterArg) {
  final AggregateCall aggregateCall =
      AggregateCall.create(SqlStdOperatorTable.SUM,
          oldCall.isDistinct(),
          oldCall.isApproximate(),
          ImmutableIntList.of(argOrdinal),
          filterArg,
          oldCall.collation,
          oldAggRel.getGroupCount(),
          oldAggRel.getInput(),
          null,
          null);
  return rexBuilder.addAggCall(aggregateCall,
      oldAggRel.getGroupCount(),
      oldAggRel.indicator,
      newCalls,
      aggCallMapping,
      ImmutableList.of(aggregateCall.getType()));
}
 

private RexNode getRegrCountRexNode(Aggregate oldAggRel,
    AggregateCall oldCall,
    List<AggregateCall> newCalls,
    Map<AggregateCall, RexNode> aggCallMapping,
    ImmutableIntList argOrdinals,
    ImmutableList<RelDataType> operandTypes,
    int filterArg) {
  final AggregateCall countArgAggCall =
      AggregateCall.create(SqlStdOperatorTable.REGR_COUNT,
          oldCall.isDistinct(),
          oldCall.isApproximate(),
          argOrdinals,
          filterArg,
          oldCall.collation,
          oldAggRel.getGroupCount(),
          oldAggRel,
          null,
          null);

  return oldAggRel.getCluster().getRexBuilder().addAggCall(countArgAggCall,
      oldAggRel.getGroupCount(),
      oldAggRel.indicator,
      newCalls,
      aggCallMapping,
      operandTypes);
}
 

private AggregateCall toAggCall(RelInput relInput, Map<String, Object> jsonAggCall) {
  final Map<String, Object> aggMap = (Map) jsonAggCall.get("agg");
  final SqlAggFunction aggregation =
      relJson.toAggregation(aggMap);
  final Boolean distinct = (Boolean) jsonAggCall.get("distinct");
  @SuppressWarnings("unchecked")
  final List<Integer> operands = (List<Integer>) jsonAggCall.get("operands");
  final Integer filterOperand = (Integer) jsonAggCall.get("filter");
  final RelDataType type =
      relJson.toType(cluster.getTypeFactory(), jsonAggCall.get("type"));
  final String name = (String) jsonAggCall.get("name");
  return AggregateCall.create(aggregation, distinct, false, false, operands,
      filterOperand == null ? -1 : filterOperand,
      RelCollations.EMPTY,
      type, name);
}
 

public AggregateCall topSplit(RexBuilder rexBuilder,
    Registry<RexNode> extra, int offset, RelDataType inputRowType,
    AggregateCall aggregateCall, int leftSubTotal, int rightSubTotal) {
  final List<RexNode> merges = new ArrayList<>();
  if (leftSubTotal >= 0) {
    merges.add(
        rexBuilder.makeInputRef(aggregateCall.type, leftSubTotal));
  }
  if (rightSubTotal >= 0) {
    merges.add(
        rexBuilder.makeInputRef(aggregateCall.type, rightSubTotal));
  }
  RexNode node;
  switch (merges.size()) {
  case 1:
    node = merges.get(0);
    break;
  case 2:
    node = rexBuilder.makeCall(SqlStdOperatorTable.MULTIPLY, merges);
    break;
  default:
    throw new AssertionError("unexpected count " + merges);
  }
  int ordinal = extra.register(node);
  return AggregateCall.create(SqlStdOperatorTable.SUM0, false, false,
      ImmutableList.of(ordinal), -1, aggregateCall.collation,
      aggregateCall.type, aggregateCall.name);
}
 

private List<AggregateCall> transformAggCalls(RelNode input, int groupCount,
    List<AggregateCall> origCalls) {
  final List<AggregateCall> newCalls = new ArrayList<>();
  for (Ord<AggregateCall> ord : Ord.zip(origCalls)) {
    final AggregateCall origCall = ord.e;
    if (origCall.isDistinct()
        || !SUPPORTED_AGGREGATES.containsKey(origCall.getAggregation()
            .getClass())) {
      return null;
    }
    final SqlAggFunction aggFun;
    final RelDataType aggType;
    if (origCall.getAggregation() == SqlStdOperatorTable.COUNT) {
      aggFun = SqlStdOperatorTable.SUM0;
      // count(any) is always not null, however nullability of sum might
      // depend on the number of columns in GROUP BY.
      // Here we use SUM0 since we are sure we will not face nullable
      // inputs nor we'll face empty set.
      aggType = null;
    } else {
      aggFun = origCall.getAggregation();
      aggType = origCall.getType();
    }
    AggregateCall newCall =
        AggregateCall.create(aggFun, origCall.isDistinct(),
            origCall.isApproximate(), origCall.ignoreNulls(),
            ImmutableList.of(groupCount + ord.i), -1, origCall.collation,
            groupCount, input, aggType, origCall.getName());
    newCalls.add(newCall);
  }
  return newCalls;
}
 

private AggregateCall copyOf(AggregateCall call) {
  return AggregateCall.create(
    call.getAggregation(), // doesn't look we need to copy this
    call.isDistinct(),
    call.getArgList(),
    call.filterArg,
    copyOf(call.getType()),
    call.getName());
}
 

private static void rewriteAggCalls(
		List<AggregateCall> newAggCalls,
		List<Integer> argList,
		Map<Integer, Integer> sourceOf) {
	// Rewrite the agg calls. Each distinct agg becomes a non-distinct call
	// to the corresponding field from the right; for example,
	// "COUNT(DISTINCT e.sal)" becomes   "COUNT(distinct_e.sal)".
	for (int i = 0; i < newAggCalls.size(); i++) {
		final AggregateCall aggCall = newAggCalls.get(i);

		// Ignore agg calls which are not distinct or have the wrong set
		// arguments. If we're rewriting aggregates whose args are {sal}, we will
		// rewrite COUNT(DISTINCT sal) and SUM(DISTINCT sal) but ignore
		// COUNT(DISTINCT gender) or SUM(sal).
		if (!aggCall.isDistinct()) {
			continue;
		}
		if (!aggCall.getArgList().equals(argList)) {
			continue;
		}

		// Re-map arguments.
		final int argCount = aggCall.getArgList().size();
		final List<Integer> newArgs = new ArrayList<>(argCount);
		for (int j = 0; j < argCount; j++) {
			final Integer arg = aggCall.getArgList().get(j);
			newArgs.add(sourceOf.get(arg));
		}
		final AggregateCall newAggCall =
				AggregateCall.create(aggCall.getAggregation(), false,
						aggCall.isApproximate(), newArgs, -1,
						aggCall.getType(), aggCall.getName());
		newAggCalls.set(i, newAggCall);
	}
}
 

public AggregateCall topSplit(RexBuilder rexBuilder,
    Registry<RexNode> extra, int offset, RelDataType inputRowType,
    AggregateCall aggregateCall, int leftSubTotal, int rightSubTotal) {
  final List<RexNode> merges = new ArrayList<>();
  if (leftSubTotal >= 0) {
    merges.add(
        rexBuilder.makeInputRef(aggregateCall.type, leftSubTotal));
  }
  if (rightSubTotal >= 0) {
    merges.add(
        rexBuilder.makeInputRef(aggregateCall.type, rightSubTotal));
  }
  RexNode node;
  switch (merges.size()) {
  case 1:
    node = merges.get(0);
    break;
  case 2:
    node = rexBuilder.makeCall(SqlStdOperatorTable.MULTIPLY, merges);
    break;
  default:
    throw new AssertionError("unexpected count " + merges);
  }
  int ordinal = extra.register(node);
  return AggregateCall.create(SqlStdOperatorTable.SUM0, false, false,
      false, ImmutableList.of(ordinal), -1, aggregateCall.collation,
      aggregateCall.type, aggregateCall.name);
}
 

private AggregateCall toAggCall(RelInput relInput, Map<String, Object> jsonAggCall) {
  final String aggName = (String) jsonAggCall.get("agg");
  final SqlAggFunction aggregation =
      relJson.toAggregation(relInput, aggName, jsonAggCall);
  final Boolean distinct = (Boolean) jsonAggCall.get("distinct");
  @SuppressWarnings("unchecked")
  final List<Integer> operands = (List<Integer>) jsonAggCall.get("operands");
  final Integer filterOperand = (Integer) jsonAggCall.get("filter");
  final RelDataType type =
      relJson.toType(cluster.getTypeFactory(), jsonAggCall.get("type"));
  return AggregateCall.create(aggregation, distinct, false, operands,
      filterOperand == null ? -1 : filterOperand,
      RelCollations.EMPTY,
      type, null);
}
 

private static void rewriteAggCalls(
		List<AggregateCall> newAggCalls,
		List<Integer> argList,
		Map<Integer, Integer> sourceOf) {
	// Rewrite the agg calls. Each distinct agg becomes a non-distinct call
	// to the corresponding field from the right; for example,
	// "COUNT(DISTINCT e.sal)" becomes   "COUNT(distinct_e.sal)".
	for (int i = 0; i < newAggCalls.size(); i++) {
		final AggregateCall aggCall = newAggCalls.get(i);

		// Ignore agg calls which are not distinct or have the wrong set
		// arguments. If we're rewriting aggregates whose args are {sal}, we will
		// rewrite COUNT(DISTINCT sal) and SUM(DISTINCT sal) but ignore
		// COUNT(DISTINCT gender) or SUM(sal).
		if (!aggCall.isDistinct()
			&& aggCall.getAggregation().getDistinctOptionality() != Optionality.IGNORED) {
			continue;
		}
		if (!aggCall.getArgList().equals(argList)) {
			continue;
		}

		// Re-map arguments.
		final int argCount = aggCall.getArgList().size();
		final List<Integer> newArgs = new ArrayList<>(argCount);
		for (int j = 0; j < argCount; j++) {
			final Integer arg = aggCall.getArgList().get(j);
			newArgs.add(sourceOf.get(arg));
		}
		final AggregateCall newAggCall =
				AggregateCall.create(aggCall.getAggregation(), false,
					aggCall.isApproximate(), false, newArgs, -1,
					RelCollations.EMPTY, aggCall.getType(), aggCall.getName());
		newAggCalls.set(i, newAggCall);
	}
}
 

public AggregateCall merge(AggregateCall top, AggregateCall bottom) {
  if (bottom.getAggregation().getKind() == SqlKind.COUNT
      && top.getAggregation().getKind() == SqlKind.SUM) {
    return AggregateCall.create(bottom.getAggregation(),
        bottom.isDistinct(), bottom.isApproximate(), false,
        bottom.getArgList(), bottom.filterArg, bottom.getCollation(),
        bottom.getType(), top.getName());
  } else {
    return null;
  }
}
 

private List<AggregateCall> transformAggCalls(RelNode input, int groupCount,
    List<AggregateCall> origCalls) {
  final List<AggregateCall> newCalls = new ArrayList<>();
  for (Ord<AggregateCall> ord : Ord.zip(origCalls)) {
    final AggregateCall origCall = ord.e;
    if (origCall.isDistinct()
        || !SUPPORTED_AGGREGATES.containsKey(origCall.getAggregation()
            .getClass())) {
      return null;
    }
    final SqlAggFunction aggFun;
    final RelDataType aggType;
    if (origCall.getAggregation() == SqlStdOperatorTable.COUNT) {
      aggFun = SqlStdOperatorTable.SUM0;
      // count(any) is always not null, however nullability of sum might
      // depend on the number of columns in GROUP BY.
      // Here we use SUM0 since we are sure we will not face nullable
      // inputs nor we'll face empty set.
      aggType = null;
    } else {
      aggFun = origCall.getAggregation();
      aggType = origCall.getType();
    }
    AggregateCall newCall =
        AggregateCall.create(aggFun, origCall.isDistinct(),
            origCall.isApproximate(),
            ImmutableList.of(groupCount + ord.i), -1,
            origCall.collation,
            groupCount,
            input,
            aggType,
            origCall.getName());
    newCalls.add(newCall);
  }
  return newCalls;
}
 

public AggregateCall other(RelDataTypeFactory typeFactory, AggregateCall e) {
  return AggregateCall.create(SqlStdOperatorTable.COUNT, false, false,
      ImmutableIntList.of(), -1,
      RelCollations.EMPTY,
      typeFactory.createSqlType(SqlTypeName.BIGINT), null);
}
 

private RexNode reduceAvg(
    Aggregate oldAggRel,
    AggregateCall oldCall,
    List<AggregateCall> newCalls,
    Map<AggregateCall, RexNode> aggCallMapping) {
  final PlannerSettings plannerSettings = (PlannerSettings) oldAggRel.getCluster().getPlanner().getContext();
  final boolean isInferenceEnabled = plannerSettings.isTypeInferenceEnabled();
  final int nGroups = oldAggRel.getGroupCount();
  RelDataTypeFactory typeFactory =
      oldAggRel.getCluster().getTypeFactory();
  RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
  int iAvgInput = oldCall.getArgList().get(0);
  RelDataType avgInputType =
      getFieldType(
          oldAggRel.getInput(),
          iAvgInput);
  RelDataType sumType =
      TypeInferenceUtils.getDrillSqlReturnTypeInference(SqlKind.SUM.name(),
          ImmutableList.of())
        .inferReturnType(oldCall.createBinding(oldAggRel));
  sumType =
      typeFactory.createTypeWithNullability(
          sumType,
          sumType.isNullable() || nGroups == 0);
  SqlAggFunction sumAgg =
      new DrillCalciteSqlAggFunctionWrapper(new SqlSumEmptyIsZeroAggFunction(), sumType);
  AggregateCall sumCall = AggregateCall.create(sumAgg, oldCall.isDistinct(),
      oldCall.isApproximate(), oldCall.getArgList(), -1, sumType, null);
  final SqlCountAggFunction countAgg = (SqlCountAggFunction) SqlStdOperatorTable.COUNT;
  final RelDataType countType = countAgg.getReturnType(typeFactory);
  AggregateCall countCall = AggregateCall.create(countAgg, oldCall.isDistinct(),
      oldCall.isApproximate(), oldCall.getArgList(), -1, countType, null);

  RexNode tmpsumRef =
      rexBuilder.addAggCall(
          sumCall,
          nGroups,
          oldAggRel.indicator,
          newCalls,
          aggCallMapping,
          ImmutableList.of(avgInputType));

  RexNode tmpcountRef =
      rexBuilder.addAggCall(
          countCall,
          nGroups,
          oldAggRel.indicator,
          newCalls,
          aggCallMapping,
          ImmutableList.of(avgInputType));

  RexNode n = rexBuilder.makeCall(SqlStdOperatorTable.CASE,
      rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
          tmpcountRef, rexBuilder.makeExactLiteral(BigDecimal.ZERO)),
          rexBuilder.constantNull(),
          tmpsumRef);

  // NOTE:  these references are with respect to the output
  // of newAggRel
  /*
  RexNode numeratorRef =
      rexBuilder.makeCall(CastHighOp,
        rexBuilder.addAggCall(
            sumCall,
            nGroups,
            newCalls,
            aggCallMapping,
            ImmutableList.of(avgInputType))
      );
  */
  RexNode numeratorRef = rexBuilder.makeCall(CastHighOp,  n);

  RexNode denominatorRef =
      rexBuilder.addAggCall(
          countCall,
          nGroups,
          oldAggRel.indicator,
          newCalls,
          aggCallMapping,
          ImmutableList.of(avgInputType));
  if (isInferenceEnabled) {
    return rexBuilder.makeCall(
        new DrillSqlOperator(
            "divide",
            2,
            true,
            oldCall.getType(), false),
        numeratorRef,
        denominatorRef);
  } else {
    final RexNode divideRef =
        rexBuilder.makeCall(
            SqlStdOperatorTable.DIVIDE,
            numeratorRef,
            denominatorRef);
    return rexBuilder.makeCast(
        typeFactory.createSqlType(SqlTypeName.ANY), divideRef);
  }
}
 

public AggregateCall other(RelDataTypeFactory typeFactory, AggregateCall e) {
  return AggregateCall.create(SqlStdOperatorTable.COUNT, false, false,
      false, ImmutableIntList.of(), -1, RelCollations.EMPTY,
      typeFactory.createSqlType(SqlTypeName.BIGINT), null);
}
 

private RexNode reduceSum(
    Aggregate oldAggRel,
    AggregateCall oldCall,
    List<AggregateCall> newCalls,
    Map<AggregateCall, RexNode> aggCallMapping) {
  final int nGroups = oldAggRel.getGroupCount();
  RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
  int arg = oldCall.getArgList().get(0);
  RelDataType argType =
      getFieldType(
          oldAggRel.getInput(),
          arg);

  final AggregateCall sumZeroCall =
      AggregateCall.create(SqlStdOperatorTable.SUM0, oldCall.isDistinct(),
          oldCall.isApproximate(), oldCall.getArgList(), oldCall.filterArg,
          oldCall.collation, oldAggRel.getGroupCount(), oldAggRel.getInput(),
          null, oldCall.name);
  final AggregateCall countCall =
      AggregateCall.create(SqlStdOperatorTable.COUNT,
          oldCall.isDistinct(),
          oldCall.isApproximate(),
          oldCall.getArgList(),
          oldCall.filterArg,
          oldCall.collation,
          oldAggRel.getGroupCount(),
          oldAggRel,
          null,
          null);

  // NOTE:  these references are with respect to the output
  // of newAggRel
  RexNode sumZeroRef =
      rexBuilder.addAggCall(sumZeroCall,
          nGroups,
          oldAggRel.indicator,
          newCalls,
          aggCallMapping,
          ImmutableList.of(argType));
  if (!oldCall.getType().isNullable()) {
    // If SUM(x) is not nullable, the validator must have determined that
    // nulls are impossible (because the group is never empty and x is never
    // null). Therefore we translate to SUM0(x).
    return sumZeroRef;
  }
  RexNode countRef =
      rexBuilder.addAggCall(countCall,
          nGroups,
          oldAggRel.indicator,
          newCalls,
          aggCallMapping,
          ImmutableList.of(argType));
  return rexBuilder.makeCall(SqlStdOperatorTable.CASE,
      rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
          countRef, rexBuilder.makeExactLiteral(BigDecimal.ZERO)),
      rexBuilder.makeCast(sumZeroRef.getType(), rexBuilder.constantNull()),
      sumZeroRef);
}
 

private RexNode reduceAvg(
    Aggregate oldAggRel,
    AggregateCall oldCall,
    List<AggregateCall> newCalls,
    Map<AggregateCall, RexNode> aggCallMapping,
    List<RexNode> inputExprs) {
  final int nGroups = oldAggRel.getGroupCount();
  final RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
  final int iAvgInput = oldCall.getArgList().get(0);
  final RelDataType avgInputType =
      getFieldType(
          oldAggRel.getInput(),
          iAvgInput);
  final AggregateCall sumCall =
      AggregateCall.create(SqlStdOperatorTable.SUM,
          oldCall.isDistinct(),
          oldCall.isApproximate(),
          oldCall.getArgList(),
          oldCall.filterArg,
          oldCall.collation,
          oldAggRel.getGroupCount(),
          oldAggRel.getInput(),
          null,
          null);
  final AggregateCall countCall =
      AggregateCall.create(SqlStdOperatorTable.COUNT,
          oldCall.isDistinct(),
          oldCall.isApproximate(),
          oldCall.getArgList(),
          oldCall.filterArg,
          oldCall.collation,
          oldAggRel.getGroupCount(),
          oldAggRel.getInput(),
          null,
          null);

  // NOTE:  these references are with respect to the output
  // of newAggRel
  RexNode numeratorRef =
      rexBuilder.addAggCall(sumCall,
          nGroups,
          oldAggRel.indicator,
          newCalls,
          aggCallMapping,
          ImmutableList.of(avgInputType));
  final RexNode denominatorRef =
      rexBuilder.addAggCall(countCall,
          nGroups,
          oldAggRel.indicator,
          newCalls,
          aggCallMapping,
          ImmutableList.of(avgInputType));

  final RelDataTypeFactory typeFactory = oldAggRel.getCluster().getTypeFactory();
  final RelDataType avgType = typeFactory.createTypeWithNullability(
      oldCall.getType(), numeratorRef.getType().isNullable());
  numeratorRef = rexBuilder.ensureType(avgType, numeratorRef, true);
  final RexNode divideRef =
      rexBuilder.makeCall(SqlStdOperatorTable.DIVIDE, numeratorRef, denominatorRef);
  return rexBuilder.makeCast(oldCall.getType(), divideRef);
}
 

private RexNode reduceSum(
    Aggregate oldAggRel,
    AggregateCall oldCall,
    List<AggregateCall> newCalls,
    Map<AggregateCall, RexNode> aggCallMapping) {
  final int nGroups = oldAggRel.getGroupCount();
  RexBuilder rexBuilder = oldAggRel.getCluster().getRexBuilder();
  int arg = oldCall.getArgList().get(0);
  RelDataType argType =
      getFieldType(
          oldAggRel.getInput(),
          arg);

  final AggregateCall sumZeroCall =
      AggregateCall.create(SqlStdOperatorTable.SUM0, oldCall.isDistinct(),
          oldCall.isApproximate(), oldCall.ignoreNulls(),
          oldCall.getArgList(), oldCall.filterArg,
          oldCall.collation, oldAggRel.getGroupCount(), oldAggRel.getInput(),
          null, oldCall.name);
  final AggregateCall countCall =
      AggregateCall.create(SqlStdOperatorTable.COUNT,
          oldCall.isDistinct(),
          oldCall.isApproximate(),
          oldCall.ignoreNulls(),
          oldCall.getArgList(),
          oldCall.filterArg,
          oldCall.collation,
          oldAggRel.getGroupCount(),
          oldAggRel,
          null,
          null);

  // NOTE:  these references are with respect to the output
  // of newAggRel
  RexNode sumZeroRef =
      rexBuilder.addAggCall(sumZeroCall,
          nGroups,
          newCalls,
          aggCallMapping,
          ImmutableList.of(argType));
  if (!oldCall.getType().isNullable()) {
    // If SUM(x) is not nullable, the validator must have determined that
    // nulls are impossible (because the group is never empty and x is never
    // null). Therefore we translate to SUM0(x).
    return sumZeroRef;
  }
  RexNode countRef =
      rexBuilder.addAggCall(countCall,
          nGroups,
          newCalls,
          aggCallMapping,
          ImmutableList.of(argType));
  return rexBuilder.makeCall(SqlStdOperatorTable.CASE,
      rexBuilder.makeCall(SqlStdOperatorTable.EQUALS,
          countRef, rexBuilder.makeExactLiteral(BigDecimal.ZERO)),
      rexBuilder.makeNullLiteral(sumZeroRef.getType()),
      sumZeroRef);
}