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