org.apache.calcite.sql.fun.SqlCountAggFunction Java Examples

/**
 * Creates a reference to an aggregate call, checking for repeated calls.
 *
 * <p>Argument types help to optimize for repeated aggregates.
 * For instance count(42) is equivalent to count(*).</p>
 *
 * @param aggCall aggregate call to be added
 * @param groupCount number of groups in the aggregate relation
 * @param indicator Whether the Aggregate has indicator (GROUPING) columns
 * @param aggCalls destination list of aggregate calls
 * @param aggCallMapping the dictionary of already added calls
 * @param aggArgTypes Argument types, not null
 *
 * @return Rex expression for the given aggregate call
 */
public RexNode addAggCall(AggregateCall aggCall, int groupCount,
    boolean indicator, List<AggregateCall> aggCalls,
    Map<AggregateCall, RexNode> aggCallMapping,
    final List<RelDataType> aggArgTypes) {
  if (aggCall.getAggregation() instanceof SqlCountAggFunction
      && !aggCall.isDistinct()) {
    final List<Integer> args = aggCall.getArgList();
    final List<Integer> nullableArgs = nullableArgs(args, aggArgTypes);
    if (!nullableArgs.equals(args)) {
      aggCall = aggCall.copy(nullableArgs, aggCall.filterArg,
          aggCall.collation);
    }
  }
  RexNode rex = aggCallMapping.get(aggCall);
  if (rex == null) {
    int index = aggCalls.size() + groupCount * (indicator ? 2 : 1);
    aggCalls.add(aggCall);
    rex = makeInputRef(aggCall.getType(), index);
    aggCallMapping.put(aggCall, rex);
  }
  return rex;
}
 

/**
 * Creates a reference to an aggregate call, checking for repeated calls.
 *
 * <p>Argument types help to optimize for repeated aggregates.
 * For instance count(42) is equivalent to count(*).</p>
 *
 * @param aggCall aggregate call to be added
 * @param groupCount number of groups in the aggregate relation
 * @param aggCalls destination list of aggregate calls
 * @param aggCallMapping the dictionary of already added calls
 * @param aggArgTypes Argument types, not null
 *
 * @return Rex expression for the given aggregate call
 */
public RexNode addAggCall(AggregateCall aggCall, int groupCount,
    List<AggregateCall> aggCalls,
    Map<AggregateCall, RexNode> aggCallMapping,
    final List<RelDataType> aggArgTypes) {
  if (aggCall.getAggregation() instanceof SqlCountAggFunction
      && !aggCall.isDistinct()) {
    final List<Integer> args = aggCall.getArgList();
    final List<Integer> nullableArgs = nullableArgs(args, aggArgTypes);
    if (!nullableArgs.equals(args)) {
      aggCall = aggCall.copy(nullableArgs, aggCall.filterArg,
          aggCall.collation);
    }
  }
  RexNode rex = aggCallMapping.get(aggCall);
  if (rex == null) {
    int index = aggCalls.size() + groupCount;
    aggCalls.add(aggCall);
    rex = makeInputRef(aggCall.getType(), index);
    aggCallMapping.put(aggCall, rex);
  }
  return rex;
}
 

private void onMatch2(RelOptRuleCall call, LogicalCorrelate correlate, RelNode leftInput,
        LogicalProject aggOutputProject, LogicalAggregate 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().toJoinType();
    // 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()...)
    //
    LogicalCorrelate newCorrelate = LogicalCorrelate.create(leftInput, aggregate, correlate.getCorrelationId(),
            correlate.getRequiredColumns(), correlate.getJoinType());

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

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

private RexNode reduceSum(
    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 arg = oldCall.getArgList().get(0);
  RelDataType argType =
      getFieldType(
          oldAggRel.getInput(),
          arg);
  final RelDataType sumType;
  final SqlAggFunction sumZeroAgg;
  if (isInferenceEnabled) {
    sumType = oldCall.getType();
  } else {
    sumType =
        typeFactory.createTypeWithNullability(
            oldCall.getType(), argType.isNullable());
  }
  sumZeroAgg = new DrillCalciteSqlAggFunctionWrapper(
      new SqlSumEmptyIsZeroAggFunction(), sumType);
  AggregateCall sumZeroCall = AggregateCall.create(sumZeroAgg, 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);
  // 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.constantNull(),
          sumZeroRef);
}
 

private void onMatch2(
    RelOptRuleCall call,
    LogicalCorrelate correlate,
    RelNode leftInput,
    LogicalProject aggOutputProject,
    LogicalAggregate 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()...)
  //
  LogicalCorrelate newCorrelate =
      LogicalCorrelate.create(leftInput, aggregate,
          correlate.getCorrelationId(), correlate.getRequiredColumns(),
          correlate.getJoinType());

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

private void onMatch2(
    RelOptRuleCall call,
    LogicalCorrelate correlate,
    RelNode leftInput,
    LogicalProject aggOutputProject,
    LogicalAggregate 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()...)
  //
  LogicalCorrelate newCorrelate =
      LogicalCorrelate.create(leftInput, aggregate,
          correlate.getCorrelationId(), correlate.getRequiredColumns(),
          correlate.getJoinType());

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

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