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

/**
 * Most of the aggregations can be retrieved with single
 * <a href="https://www.elastic.co/guide/en/elasticsearch/reference/current/search-aggregations-metrics-stats-aggregation.html">stats</a>
 * function. But currently only one-to-one mapping is supported between sql agg and elastic
 * aggregation.
 */
private static String toElasticAggregate(AggregateCall call) {
  final SqlKind kind = call.getAggregation().getKind();
  switch (kind) {
  case COUNT:
    // approx_count_distinct() vs count()
    return call.isDistinct() && call.isApproximate() ? "cardinality" : "value_count";
  case SUM:
    return "sum";
  case MIN:
    return "min";
  case MAX:
    return "max";
  case AVG:
    return "avg";
  case ANY_VALUE:
    return "terms";
  default:
    throw new IllegalArgumentException("Unknown aggregation kind " + kind + " for " + call);
  }
}
 

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

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

@Override
public void implementOLAP(OLAPImplementor implementor) {
    implementor.fixSharedOlapTableScan(this);
    implementor.visitChild(getInput(), this);

    this.context = implementor.getContext();
    this.columnRowType = buildColumnRowType();
    this.afterAggregate = this.context.afterAggregate;

    // only translate the innermost aggregation
    if (!this.afterAggregate) {
        addToContextGroupBy(this.groups);
        this.context.aggregations.addAll(this.aggregations);
        this.context.aggrOutCols
                .addAll(columnRowType.getAllColumns().subList(groups.size(), columnRowType.getAllColumns().size()));
        this.context.afterAggregate = true;

        if (this.context.afterLimit) {
            this.context.limitPrecedesAggr = true;
        }
    } else {
        this.context.afterOuterAggregate = true;
        for (AggregateCall aggCall : aggCalls) {
            // check if supported by kylin
            if (aggCall.isDistinct()) {
                throw new IllegalStateException("Distinct count is only allowed in innermost sub-query.");
            }
        }
    }
}
 

/**
 * Returns true if and only if this <code>ComplexMetric</code>
 * can be used in the given {@link AggregateCall}.
 * */
public boolean canBeUsed(AggregateCall call) {
  switch (type) {
  case HYPER_UNIQUE:
  case THETA_SKETCH:
    return call != null
          && call.getAggregation().getKind() == SqlKind.COUNT
          && call.isDistinct();
  default:
    return false;
  }
}
 

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(), aggCall.ignoreNulls(), newArgs, -1, aggCall.collation,
            aggCall.getType(), aggCall.getName());
    newAggCalls.set(i, newAggCall);
  }
}
 

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

@Override
public void implementOLAP(OLAPImplementor implementor) {
    implementor.fixSharedOlapTableScan(this);
    implementor.visitChild(getInput(), this);

    this.context = implementor.getContext();
    this.columnRowType = buildColumnRowType();
    this.afterAggregate = this.context.afterAggregate;

    // only translate the innermost aggregation
    if (!this.afterAggregate) {
        addToContextGroupBy(this.groups);
        this.context.aggregations.addAll(this.aggregations);
        this.context.aggrOutCols
                .addAll(columnRowType.getAllColumns().subList(groups.size(), columnRowType.getAllColumns().size()));
        this.context.afterAggregate = true;

        if (this.context.afterLimit) {
            this.context.limitPrecedesAggr = true;
        }
    } else {
        this.context.afterOuterAggregate = true;
        for (AggregateCall aggCall : aggCalls) {
            // check if supported by kylin
            if (aggCall.isDistinct()) {
                throw new IllegalStateException("Distinct count is only allowed in innermost sub-query.");
            }
        }
    }
}
 

static String getSqlFuncName(AggregateCall aggCall) {
    String sqlName = aggCall.getAggregation().getName();
    if (aggCall.isDistinct()) {
        sqlName = sqlName + "_DISTINCT";
    }
    return sqlName;
}
 

static String getSqlFuncName(AggregateCall aggCall) {
    String sqlName = aggCall.getAggregation().getName();
    if (aggCall.isDistinct()) {
        sqlName = sqlName + "_DISTINCT";
    }
    return sqlName;
}
 

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

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

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.collation, aggCall.getType(), aggCall.getName());
        newAggCalls.set(i, newAggCall);
    }
}
 

private Measure toMeasure(AggregateCall aggCall) {
  return new Measure(aggCall.getAggregation(), aggCall.isDistinct(),
      aggCall.name, Lists.transform(aggCall.getArgList(), columns::get));
}
 

@Override
public void onMatch(RelOptRuleCall call) {
  final LogicalAggregate agg = call.rel(0);
  final RelNode input = agg.getInput();

  boolean distinctReplaced = false;
  List<AggregateCall> calls = new ArrayList<>();

  RelMetadataQuery query = null;

  final Boolean[] memo = new Boolean[agg.getInput().getRowType().getFieldCount()];
  for (AggregateCall c : agg.getAggCallList()) {
    final boolean candidate = c.isDistinct() && c.getArgList().size() == 1 && "COUNT".equals(c.getAggregation().getName());

    if(!candidate) {
      calls.add(c);
      continue;
    }

    final int inputOrdinal = c.getArgList().get(0);
    boolean allowed = false;
    if(memo[inputOrdinal] != null) {
      allowed = memo[inputOrdinal];
    } else {
      if(query == null) {
        query = agg.getCluster().getMetadataQuery();
      }

      Set<RelColumnOrigin> origins = query.getColumnOrigins(input, inputOrdinal);

      // see if any column origin allowed a transformation.
      for(RelColumnOrigin o : origins) {
        RelOptTable table = o.getOriginTable();
        NamespaceTable namespaceTable = table.unwrap(NamespaceTable.class);
        if(namespaceTable == null) {
          // unable to decide, no way to transform.
          return;
        }

        if(namespaceTable.isApproximateStatsAllowed()) {
          allowed = true;
        }
      }

      memo[inputOrdinal] = allowed;

    }


    if(allowed) {
      calls.add(AggregateCall.create(HyperLogLog.NDV, false, c.getArgList(), -1, c.getType(), c.getName()));
      distinctReplaced = true;
    } else {
      calls.add(c);
    }

  }

  if(!distinctReplaced) {
    return;
  }

  final RelBuilder builder = relBuilderFactory.create(agg.getCluster(), null);
  builder.push(agg.getInput());
  builder.aggregate(builder.groupKey(agg.getGroupSet().toArray()), calls);
  call.transformTo(builder.build());
}
 

public KylinAggregateCall(AggregateCall aggCall, FunctionDesc func) {
    super(aggCall.getAggregation(), aggCall.isDistinct(), aggCall.getArgList(), aggCall.type, aggCall.name);
    this.func = func;
}
 

private static AggregateCall rollUp(int groupCount, RelBuilder relBuilder,
    AggregateCall aggregateCall, TileKey tileKey) {
  if (aggregateCall.isDistinct()) {
    return null;
  }
  final SqlAggFunction aggregation = aggregateCall.getAggregation();
  final Pair<SqlAggFunction, List<Integer>> seek =
      Pair.of(aggregation, aggregateCall.getArgList());
  final int offset = tileKey.dimensions.cardinality();
  final ImmutableList<Lattice.Measure> measures = tileKey.measures;

  // First, try to satisfy the aggregation by rolling up an aggregate in the
  // materialization.
  final int i = find(measures, seek);
tryRoll:
  if (i >= 0) {
    final SqlAggFunction roll = SubstitutionVisitor.getRollup(aggregation);
    if (roll == null) {
      break tryRoll;
    }
    return AggregateCall.create(roll, false, aggregateCall.isApproximate(),
        aggregateCall.ignoreNulls(), ImmutableList.of(offset + i), -1,
        aggregateCall.collation,
        groupCount, relBuilder.peek(), null, aggregateCall.name);
  }

  // Second, try to satisfy the aggregation based on group set columns.
tryGroup:
  {
    List<Integer> newArgs = new ArrayList<>();
    for (Integer arg : aggregateCall.getArgList()) {
      int z = tileKey.dimensions.indexOf(arg);
      if (z < 0) {
        break tryGroup;
      }
      newArgs.add(z);
    }
    return AggregateCall.create(aggregation, false,
        aggregateCall.isApproximate(), aggregateCall.ignoreNulls(),
        newArgs, -1, aggregateCall.collation,
        groupCount, relBuilder.peek(), null, aggregateCall.name);
  }

  // No roll up possible.
  return null;
}
 

private Measure toMeasure(AggregateCall aggCall) {
  return new Measure(aggCall.getAggregation(), aggCall.isDistinct(),
      aggCall.name, Lists.transform(aggCall.getArgList(), columns::get));
}
 

private String getInputFieldNameForAggCall(String relAlias, AggregateCall aggCall,
    int fieldIndex) {
  final String inputField = getInputFieldName(fieldIndex);
  return aggCall.isDistinct() ? (inputField + DISTINCT_FIELD_SUFFIX)
      : (relAlias + '.' + inputField);
}
 

@Override
protected RelNode createUnion(RelBuilder relBuilder, RexBuilder rexBuilder, RelNode topProject,
        RelNode unionInputQuery, RelNode unionInputView) {
    // Union
    relBuilder.push(unionInputQuery);
    relBuilder.push(unionInputView);
    relBuilder.union(true);
    List<RexNode> exprList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
    List<String> nameList = new ArrayList<>(relBuilder.peek().getRowType().getFieldCount());
    for (int i = 0; i < relBuilder.peek().getRowType().getFieldCount(); i++) {
        // We can take unionInputQuery as it is query based.
        RelDataTypeField field = unionInputQuery.getRowType().getFieldList().get(i);
        exprList.add(
                rexBuilder.ensureType(field.getType(), rexBuilder.makeInputRef(relBuilder.peek(), i), true));
        nameList.add(field.getName());
    }
    relBuilder.project(exprList, nameList);
    // Rollup aggregate
    Aggregate aggregate = (Aggregate) unionInputQuery;
    final ImmutableBitSet groupSet = ImmutableBitSet.range(aggregate.getGroupCount());
    final List<AggCall> aggregateCalls = new ArrayList<>();
    for (int i = 0; i < aggregate.getAggCallList().size(); i++) {
        AggregateCall aggCall = aggregate.getAggCallList().get(i);
        if (aggCall.isDistinct()) {
            // Cannot ROLLUP distinct
            return null;
        }
        SqlAggFunction rollupAgg = getRollup(aggCall.getAggregation());
        if (rollupAgg == null) {
            // Cannot rollup this aggregate, bail out
            return null;
        }
        final RexInputRef operand = rexBuilder.makeInputRef(relBuilder.peek(), aggregate.getGroupCount() + i);
        aggregateCalls.add(
                // TODO: handle aggregate ordering
                relBuilder.aggregateCall(rollupAgg, operand).distinct(aggCall.isDistinct())
                        .approximate(aggCall.isApproximate()).as(aggCall.name));
    }
    RelNode prevNode = relBuilder.peek();
    RelNode result = relBuilder.aggregate(relBuilder.groupKey(groupSet), aggregateCalls).build();
    if (prevNode == result && groupSet.cardinality() != result.getRowType().getFieldCount()) {
        // Aggregate was not inserted but we need to prune columns
        result = relBuilder.push(result).project(relBuilder.fields(groupSet.asList())).build();
    }
    if (topProject != null) {
        // Top project
        return topProject.copy(topProject.getTraitSet(), ImmutableList.of(result));
    }
    // Result
    return result;
}