Search in sources :

Example 26 with ExpressionLambdaAggregatorFactory

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

the class CalciteArraysQueryTest method testArrayAggExpression.

@Test
public void testArrayAggExpression() throws Exception {
    cannotVectorize();
    testQuery("SELECT ARRAY_TO_STRING(ARRAY_AGG(DISTINCT CONCAT(dim1, dim2)), ',') FROM foo", ImmutableList.of(Druids.newTimeseriesQueryBuilder().dataSource(CalciteTests.DATASOURCE1).intervals(querySegmentSpec(Filtration.eternity())).granularity(Granularities.ALL).virtualColumns(expressionVirtualColumn("v0", "concat(\"dim1\",\"dim2\")", ColumnType.STRING)).aggregators(aggregators(new ExpressionLambdaAggregatorFactory("a0", ImmutableSet.of("v0"), "__acc", "ARRAY<STRING>[]", "ARRAY<STRING>[]", true, true, false, "array_set_add(\"__acc\", \"v0\")", "array_set_add_all(\"__acc\", \"a0\")", null, null, ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES, TestExprMacroTable.INSTANCE))).postAggregators(expressionPostAgg("p0", "array_to_string(\"a0\",',')")).context(QUERY_CONTEXT_DEFAULT).build()), ImmutableList.of(useDefault ? new Object[] { "10.1,1a,2,a,abc,defabc" } : new Object[] { "null,1a,2,a,defabc" }));
}
Also used : ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) Test(org.junit.Test)

Example 27 with ExpressionLambdaAggregatorFactory

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

the class CalciteArraysQueryTest method testArrayAggGroupByArrayAggFromSubquery.

@Test
public void testArrayAggGroupByArrayAggFromSubquery() throws Exception {
    cannotVectorize();
    testQuery("SELECT dim2, arr, COUNT(*) FROM (SELECT dim2, ARRAY_AGG(DISTINCT dim1) as arr FROM foo WHERE dim1 is not null GROUP BY 1 LIMIT 5) GROUP BY 1,2", QUERY_CONTEXT_NO_STRINGIFY_ARRAY, ImmutableList.of(GroupByQuery.builder().setDataSource(new TopNQuery(new TableDataSource(CalciteTests.DATASOURCE1), null, new DefaultDimensionSpec("dim2", "d0", ColumnType.STRING), new DimensionTopNMetricSpec(null, StringComparators.LEXICOGRAPHIC), 5, querySegmentSpec(Filtration.eternity()), new NotDimFilter(new SelectorDimFilter("dim1", null, null)), Granularities.ALL, aggregators(new ExpressionLambdaAggregatorFactory("a0", ImmutableSet.of("dim1"), "__acc", "ARRAY<STRING>[]", "ARRAY<STRING>[]", true, true, false, "array_set_add(\"__acc\", \"dim1\")", "array_set_add_all(\"__acc\", \"a0\")", null, null, new HumanReadableBytes(1024), ExprMacroTable.nil())), null, QUERY_CONTEXT_NO_STRINGIFY_ARRAY)).setInterval(querySegmentSpec(Filtration.eternity())).setGranularity(Granularities.ALL).setDimFilter(null).setGranularity(Granularities.ALL).setDimensions(dimensions(new DefaultDimensionSpec("d0", "_d0", ColumnType.STRING), new DefaultDimensionSpec("a0", "_d1", ColumnType.STRING_ARRAY))).setAggregatorSpecs(aggregators(new CountAggregatorFactory("_a0"))).setContext(QUERY_CONTEXT_NO_STRINGIFY_ARRAY).build()), useDefault ? ImmutableList.of(new Object[] { "", ImmutableList.of("10.1", "2", "abc"), 1L }, new Object[] { "a", ImmutableList.of("1"), 1L }, new Object[] { "abc", ImmutableList.of("def"), 1L }) : ImmutableList.of(new Object[] { null, ImmutableList.of("10.1", "abc"), 1L }, new Object[] { "", ImmutableList.of("2"), 1L }, new Object[] { "a", ImmutableList.of("", "1"), 1L }, new Object[] { "abc", ImmutableList.of("def"), 1L }));
}
Also used : DimensionTopNMetricSpec(org.apache.druid.query.topn.DimensionTopNMetricSpec) NotDimFilter(org.apache.druid.query.filter.NotDimFilter) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) TableDataSource(org.apache.druid.query.TableDataSource) CountAggregatorFactory(org.apache.druid.query.aggregation.CountAggregatorFactory) SelectorDimFilter(org.apache.druid.query.filter.SelectorDimFilter) TopNQuery(org.apache.druid.query.topn.TopNQuery) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) DefaultDimensionSpec(org.apache.druid.query.dimension.DefaultDimensionSpec) Test(org.junit.Test)

Example 28 with ExpressionLambdaAggregatorFactory

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

the class StringSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(PlannerContext plannerContext, RowSignature rowSignature, VirtualColumnRegistry virtualColumnRegistry, RexBuilder rexBuilder, String name, AggregateCall aggregateCall, Project project, List<Aggregation> existingAggregations, boolean finalizeAggregations) {
    final List<DruidExpression> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).map(rexNode -> Expressions.toDruidExpression(plannerContext, rowSignature, rexNode)).collect(Collectors.toList());
    if (arguments.stream().anyMatch(Objects::isNull)) {
        return null;
    }
    RexNode separatorNode = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(1));
    if (!separatorNode.isA(SqlKind.LITERAL)) {
        // separator must be a literal
        return null;
    }
    String separator = RexLiteral.stringValue(separatorNode);
    if (separator == null) {
        // separator must not be null
        return null;
    }
    Integer maxSizeBytes = null;
    if (arguments.size() > 2) {
        RexNode maxBytes = Expressions.fromFieldAccess(rowSignature, project, aggregateCall.getArgList().get(2));
        if (!maxBytes.isA(SqlKind.LITERAL)) {
            // maxBytes must be a literal
            return null;
        }
        maxSizeBytes = ((Number) RexLiteral.value(maxBytes)).intValue();
    }
    final DruidExpression arg = arguments.get(0);
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String initialvalue = "[]";
    final ColumnType elementType = ColumnType.STRING;
    final String fieldName;
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        fieldName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(arg, elementType);
    }
    final String finalizer = StringUtils.format("if(array_length(o) == 0, null, array_to_string(o, '%s'))", separator);
    final NotDimFilter dimFilter = new NotDimFilter(new SelectorDimFilter(fieldName, null, null));
    if (aggregateCall.isDistinct()) {
        return Aggregation.create(// string_agg ignores nulls
        new FilteredAggregatorFactory(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_set_add(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name), null, finalizer, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable), dimFilter));
    } else {
        return Aggregation.create(// string_agg ignores nulls
        new FilteredAggregatorFactory(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, false, false, StringUtils.format("array_append(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_concat(\"__acc\", \"%s\")", name), null, finalizer, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable), dimFilter));
    }
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) FilteredAggregatorFactory(org.apache.druid.query.aggregation.FilteredAggregatorFactory) RowSignatures(org.apache.druid.sql.calcite.table.RowSignatures) UnsupportedSQLQueryException(org.apache.druid.sql.calcite.planner.UnsupportedSQLQueryException) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) Optionality(org.apache.calcite.util.Optionality) SelectorDimFilter(org.apache.druid.query.filter.SelectorDimFilter) 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) SqlOperatorBinding(org.apache.calcite.sql.SqlOperatorBinding) RelDataType(org.apache.calcite.rel.type.RelDataType) ImmutableSet(com.google.common.collect.ImmutableSet) SqlKind(org.apache.calcite.sql.SqlKind) SqlTypeFamily(org.apache.calcite.sql.type.SqlTypeFamily) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) NotDimFilter(org.apache.druid.query.filter.NotDimFilter) SqlTypeName(org.apache.calcite.sql.type.SqlTypeName) InferTypes(org.apache.calcite.sql.type.InferTypes) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) StringUtils(org.apache.druid.java.util.common.StringUtils) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Collectors(java.util.stream.Collectors) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) Objects(java.util.Objects) SqlReturnTypeInference(org.apache.calcite.sql.type.SqlReturnTypeInference) List(java.util.List) 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) Calcites(org.apache.druid.sql.calcite.planner.Calcites) Expressions(org.apache.druid.sql.calcite.expression.Expressions) FilteredAggregatorFactory(org.apache.druid.query.aggregation.FilteredAggregatorFactory) ColumnType(org.apache.druid.segment.column.ColumnType) NotDimFilter(org.apache.druid.query.filter.NotDimFilter) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) SelectorDimFilter(org.apache.druid.query.filter.SelectorDimFilter) Objects(java.util.Objects) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Example 29 with ExpressionLambdaAggregatorFactory

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

the class BitwiseSqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(PlannerContext plannerContext, RowSignature rowSignature, VirtualColumnRegistry virtualColumnRegistry, RexBuilder rexBuilder, String name, AggregateCall aggregateCall, Project project, List<Aggregation> existingAggregations, boolean finalizeAggregations) {
    final List<DruidExpression> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).map(rexNode -> Expressions.toDruidExpression(plannerContext, rowSignature, rexNode)).collect(Collectors.toList());
    if (arguments.stream().anyMatch(Objects::isNull)) {
        return null;
    }
    final DruidExpression arg = arguments.get(0);
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        fieldName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(arg, ColumnType.LONG);
    }
    return Aggregation.create(new FilteredAggregatorFactory(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, "0", null, null, false, false, StringUtils.format("%s(\"__acc\", \"%s\")", op.getDruidFunction(), fieldName), null, null, null, null, macroTable), new NotDimFilter(new SelectorDimFilter(fieldName, null, null))));
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) FilteredAggregatorFactory(org.apache.druid.query.aggregation.FilteredAggregatorFactory) ReturnTypes(org.apache.calcite.sql.type.ReturnTypes) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) Optionality(org.apache.calcite.util.Optionality) SelectorDimFilter(org.apache.druid.query.filter.SelectorDimFilter) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Nullable(javax.annotation.Nullable) ImmutableSet(com.google.common.collect.ImmutableSet) SqlKind(org.apache.calcite.sql.SqlKind) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) NotDimFilter(org.apache.druid.query.filter.NotDimFilter) SqlTypeName(org.apache.calcite.sql.type.SqlTypeName) InferTypes(org.apache.calcite.sql.type.InferTypes) RexBuilder(org.apache.calcite.rex.RexBuilder) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) StringUtils(org.apache.druid.java.util.common.StringUtils) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Collectors(java.util.stream.Collectors) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) Objects(java.util.Objects) List(java.util.List) SqlStdOperatorTable(org.apache.calcite.sql.fun.SqlStdOperatorTable) 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) FilteredAggregatorFactory(org.apache.druid.query.aggregation.FilteredAggregatorFactory) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) NotDimFilter(org.apache.druid.query.filter.NotDimFilter) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) SelectorDimFilter(org.apache.druid.query.filter.SelectorDimFilter) Objects(java.util.Objects) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) Nullable(javax.annotation.Nullable)

Example 30 with ExpressionLambdaAggregatorFactory

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

the class ArraySqlAggregator method toDruidAggregation.

@Nullable
@Override
public Aggregation toDruidAggregation(PlannerContext plannerContext, RowSignature rowSignature, VirtualColumnRegistry virtualColumnRegistry, RexBuilder rexBuilder, String name, AggregateCall aggregateCall, Project project, List<Aggregation> existingAggregations, boolean finalizeAggregations) {
    final List<RexNode> arguments = aggregateCall.getArgList().stream().map(i -> Expressions.fromFieldAccess(rowSignature, project, i)).collect(Collectors.toList());
    Integer maxSizeBytes = null;
    if (arguments.size() > 1) {
        RexNode maxBytes = arguments.get(1);
        if (!maxBytes.isA(SqlKind.LITERAL)) {
            // maxBytes must be a literal
            return null;
        }
        maxSizeBytes = ((Number) RexLiteral.value(maxBytes)).intValue();
    }
    final DruidExpression arg = Expressions.toDruidExpression(plannerContext, rowSignature, arguments.get(0));
    if (arg == null) {
        // can't translate argument
        return null;
    }
    final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
    final String fieldName;
    final String initialvalue;
    final ColumnType druidType = Calcites.getValueTypeForRelDataTypeFull(aggregateCall.getType());
    final ColumnType elementType;
    if (druidType == null || !druidType.isArray()) {
        initialvalue = "[]";
        elementType = ColumnType.STRING;
    } else {
        initialvalue = ExpressionType.fromColumnTypeStrict(druidType).asTypeString() + "[]";
        elementType = (ColumnType) druidType.getElementType();
    }
    if (arg.isDirectColumnAccess()) {
        fieldName = arg.getDirectColumn();
    } else {
        fieldName = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(arg, elementType);
    }
    if (aggregateCall.isDistinct()) {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, true, false, StringUtils.format("array_set_add(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    } else {
        return Aggregation.create(new ExpressionLambdaAggregatorFactory(name, ImmutableSet.of(fieldName), null, initialvalue, null, true, true, false, StringUtils.format("array_append(\"__acc\", \"%s\")", fieldName), StringUtils.format("array_concat(\"__acc\", \"%s\")", name), null, null, maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null, macroTable));
    }
}
Also used : Project(org.apache.calcite.rel.core.Project) SqlAggregator(org.apache.druid.sql.calcite.aggregation.SqlAggregator) RowSignatures(org.apache.druid.sql.calcite.table.RowSignatures) UnsupportedSQLQueryException(org.apache.druid.sql.calcite.planner.UnsupportedSQLQueryException) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) Optionality(org.apache.calcite.util.Optionality) RexNode(org.apache.calcite.rex.RexNode) ExpressionType(org.apache.druid.math.expr.ExpressionType) VirtualColumnRegistry(org.apache.druid.sql.calcite.rel.VirtualColumnRegistry) PlannerContext(org.apache.druid.sql.calcite.planner.PlannerContext) Nullable(javax.annotation.Nullable) SqlOperatorBinding(org.apache.calcite.sql.SqlOperatorBinding) RelDataType(org.apache.calcite.rel.type.RelDataType) ImmutableSet(com.google.common.collect.ImmutableSet) SqlKind(org.apache.calcite.sql.SqlKind) SqlTypeFamily(org.apache.calcite.sql.type.SqlTypeFamily) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) InferTypes(org.apache.calcite.sql.type.InferTypes) RexBuilder(org.apache.calcite.rex.RexBuilder) RexLiteral(org.apache.calcite.rex.RexLiteral) SqlFunctionCategory(org.apache.calcite.sql.SqlFunctionCategory) StringUtils(org.apache.druid.java.util.common.StringUtils) Aggregation(org.apache.druid.sql.calcite.aggregation.Aggregation) Collectors(java.util.stream.Collectors) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) SqlReturnTypeInference(org.apache.calcite.sql.type.SqlReturnTypeInference) List(java.util.List) 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) Calcites(org.apache.druid.sql.calcite.planner.Calcites) Expressions(org.apache.druid.sql.calcite.expression.Expressions) ColumnType(org.apache.druid.segment.column.ColumnType) ExpressionLambdaAggregatorFactory(org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory) DruidExpression(org.apache.druid.sql.calcite.expression.DruidExpression) HumanReadableBytes(org.apache.druid.java.util.common.HumanReadableBytes) ExprMacroTable(org.apache.druid.math.expr.ExprMacroTable) RexNode(org.apache.calcite.rex.RexNode) Nullable(javax.annotation.Nullable)

Aggregations

ExpressionLambdaAggregatorFactory (org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory)30 Test (org.junit.Test)25 FilteredAggregatorFactory (org.apache.druid.query.aggregation.FilteredAggregatorFactory)12 DefaultDimensionSpec (org.apache.druid.query.dimension.DefaultDimensionSpec)11 InitializedNullHandlingTest (org.apache.druid.testing.InitializedNullHandlingTest)8 HumanReadableBytes (org.apache.druid.java.util.common.HumanReadableBytes)6 ImmutableSet (com.google.common.collect.ImmutableSet)4 List (java.util.List)4 Collectors (java.util.stream.Collectors)4 Nullable (javax.annotation.Nullable)4 AggregateCall (org.apache.calcite.rel.core.AggregateCall)4 Project (org.apache.calcite.rel.core.Project)4 RexBuilder (org.apache.calcite.rex.RexBuilder)4 SqlAggFunction (org.apache.calcite.sql.SqlAggFunction)4 SqlFunctionCategory (org.apache.calcite.sql.SqlFunctionCategory)4 SqlKind (org.apache.calcite.sql.SqlKind)4 InferTypes (org.apache.calcite.sql.type.InferTypes)4 OperandTypes (org.apache.calcite.sql.type.OperandTypes)4 Optionality (org.apache.calcite.util.Optionality)4 StringUtils (org.apache.druid.java.util.common.StringUtils)4