Search in sources :

Example 1 with ApproximateHistogramFoldingAggregatorFactory

use of org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory in project druid by druid-io.

the class QuantileSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(final PlannerContext plannerContext, final RowSignature rowSignature, final VirtualColumnRegistry virtualColumnRegistry, final RexBuilder rexBuilder, final String name, final AggregateCall aggregateCall, final Project project, final List<Aggregation> existingAggregations, final boolean finalizeAggregations) {
    final DruidExpression input = Aggregations.toDruidExpressionForNumericAggregator(plannerContext, rowSignature, Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(0)));
    if (input == null) {
        return null;
    }
    final AggregatorFactory aggregatorFactory;
    final String histogramName = StringUtils.format("%s:agg", name);
    final RexNode probabilityArg = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(1));
    if (!probabilityArg.isA(SqlKind.LITERAL)) {
        // Probability must be a literal in order to plan.
        return null;
    }
    final float probability = ((Number) RexLiteral.value(probabilityArg)).floatValue();
    final int resolution;
    if (aggregateCall.getArgList().size() >= 3) {
        final RexNode resolutionArg = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(2));
        if (!resolutionArg.isA(SqlKind.LITERAL)) {
            // Resolution must be a literal in order to plan.
            return null;
        }
        resolution = ((Number) RexLiteral.value(resolutionArg)).intValue();
    } else {
        resolution = ApproximateHistogram.DEFAULT_HISTOGRAM_SIZE;
    }
    final int numBuckets = ApproximateHistogram.DEFAULT_BUCKET_SIZE;
    final float lowerLimit = Float.NEGATIVE_INFINITY;
    final float upperLimit = Float.POSITIVE_INFINITY;
    // Look for existing matching aggregatorFactory.
    for (final Aggregation existing : existingAggregations) {
        for (AggregatorFactory factory : existing.getAggregatorFactories()) {
            if (factory instanceof ApproximateHistogramAggregatorFactory) {
                final ApproximateHistogramAggregatorFactory theFactory = (ApproximateHistogramAggregatorFactory) factory;
                // Check input for equivalence.
                final boolean inputMatches;
                final DruidExpression virtualInput = virtualColumnRegistry.findVirtualColumnExpressions(theFactory.requiredFields()).stream().findFirst().orElse(null);
                if (virtualInput == null) {
                    inputMatches = input.isDirectColumnAccess() && input.getDirectColumn().equals(theFactory.getFieldName());
                } else {
                    inputMatches = virtualInput.equals(input);
                }
                final boolean matches = inputMatches && theFactory.getResolution() == resolution && theFactory.getNumBuckets() == numBuckets && theFactory.getLowerLimit() == lowerLimit && theFactory.getUpperLimit() == upperLimit;
                if (matches) {
                    // Found existing one. Use this.
                    return Aggregation.create(ImmutableList.of(), new QuantilePostAggregator(name, factory.getName(), probability));
                }
            }
        }
    }
    // No existing match found. Create a new one.
    if (input.isDirectColumnAccess()) {
        if (rowSignature.getColumnType(input.getDirectColumn()).map(type -> type.is(ValueType.COMPLEX)).orElse(false)) {
            aggregatorFactory = new ApproximateHistogramFoldingAggregatorFactory(histogramName, input.getDirectColumn(), resolution, numBuckets, lowerLimit, upperLimit, false);
        } else {
            aggregatorFactory = new ApproximateHistogramAggregatorFactory(histogramName, input.getDirectColumn(), resolution, numBuckets, lowerLimit, upperLimit, false);
        }
    } else {
        final String virtualColumnName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(input, ColumnType.FLOAT);
        aggregatorFactory = new ApproximateHistogramAggregatorFactory(histogramName, virtualColumnName, resolution, numBuckets, lowerLimit, upperLimit, false);
    }
    return Aggregation.create(ImmutableList.of(aggregatorFactory), new QuantilePostAggregator(name, histogramName, probability));
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) ReturnTypes(org.apache.calcite.sql.type.ReturnTypes) QuantilePostAggregator(org.apache.druid.query.aggregation.histogram.QuantilePostAggregator) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ApproximateHistogram(org.apache.druid.query.aggregation.histogram.ApproximateHistogram) ImmutableList(com.google.common.collect.ImmutableList) RexNode(org.apache.calcite.rex.RexNode) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Nullable(javax.annotation.Nullable) SqlKind(org.apache.calcite.sql.SqlKind) SqlTypeFamily(org.apache.calcite.sql.type.SqlTypeFamily) SqlTypeName(org.apache.calcite.sql.type.SqlTypeName) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) ApproximateHistogramFoldingAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory) StringUtils(org.apache.druid.java.util.common.StringUtils) ValueType(org.apache.druid.segment.column.ValueType) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) ApproximateHistogramAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramAggregatorFactory) List(java.util.List) Aggregations(org.apache.druid.sql.calcite.aggregation.Aggregations) RowSignature(org.apache.druid.segment.column.RowSignature) OperandTypes(org.apache.calcite.sql.type.OperandTypes) ColumnType(org.apache.druid.segment.column.ColumnType) AggregateCall(org.apache.calcite.rel.core.AggregateCall) SqlAggFunction(org.apache.calcite.sql.SqlAggFunction) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ApproximateHistogramFoldingAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory) QuantilePostAggregator(org.apache.druid.query.aggregation.histogram.QuantilePostAggregator) AggregatorFactory(org.apache.druid.query.aggregation.AggregatorFactory) ApproximateHistogramFoldingAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory) ApproximateHistogramAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramAggregatorFactory) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) ApproximateHistogramAggregatorFactory(org.apache.druid.query.aggregation.histogram.ApproximateHistogramAggregatorFactory) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Aggregations

ImmutableList (com.google.common.collect.ImmutableList)1 List (java.util.List)1 Nullable (javax.annotation.Nullable)1 AggregateCall (org.apache.calcite.rel.core.AggregateCall)1 Project (org.apache.calcite.rel.core.Project)1 RexBuilder (org.apache.calcite.rex.RexBuilder)1 RexLiteral (org.apache.calcite.rex.RexLiteral)1 RexNode (org.apache.calcite.rex.RexNode)1 SqlAggFunction (org.apache.calcite.sql.SqlAggFunction)1 SqlFunctionCategory (org.apache.calcite.sql.SqlFunctionCategory)1 SqlKind (org.apache.calcite.sql.SqlKind)1 OperandTypes (org.apache.calcite.sql.type.OperandTypes)1 ReturnTypes (org.apache.calcite.sql.type.ReturnTypes)1 SqlTypeFamily (org.apache.calcite.sql.type.SqlTypeFamily)1 SqlTypeName (org.apache.calcite.sql.type.SqlTypeName)1 StringUtils (org.apache.druid.java.util.common.StringUtils)1 AggregatorFactory (org.apache.druid.query.aggregation.AggregatorFactory)1 ApproximateHistogram (org.apache.druid.query.aggregation.histogram.ApproximateHistogram)1 ApproximateHistogramAggregatorFactory (org.apache.druid.query.aggregation.histogram.ApproximateHistogramAggregatorFactory)1 ApproximateHistogramFoldingAggregatorFactory (org.apache.druid.query.aggregation.histogram.ApproximateHistogramFoldingAggregatorFactory)1