Apache
/
druid
mirror of https://github.com/apache/druid.git
1
0
Fork 0
Code Issues Packages Projects Releases Wiki Activity

array_concat_agg and array_agg support for array inputs (#12226) 

* array_concat_agg and array_agg support for array inputs
changes:
* added array_concat_agg to aggregate arrays into a single array
* added array_agg support for array inputs to make nested array
* added 'shouldAggregateNullInputs' and 'shouldCombineAggregateNullInputs' to fix a correctness issue with STRING_AGG and ARRAY_AGG when merging results, with dual purpose of being an optimization for aggregating

* fix test

* tie capabilities type to legacy mode flag about coercing arrays to strings

* oops

* better javadoc
This commit is contained in:
Clint Wylie 2022-02-07 19:59:30 -08:00 committed by GitHub
parent 090c429c8c
commit ae71e05fc5
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
32 changed files with 1156 additions and 341 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
core/src/main/java/org/apache/druid/math/expr/ExpressionType.java
View File

@ -130,6 +130,7 @@ public class ExpressionType extends BaseTypeSignature<ExprType>
switch (valueType.getElementType().getType()) {
case LONG:
return LONG_ARRAY;
case FLOAT:
case DOUBLE:
return DOUBLE_ARRAY;
case STRING:
@ -166,6 +167,7 @@ public class ExpressionType extends BaseTypeSignature<ExprType>
switch (valueType.getElementType().getType()) {
case LONG:
return LONG_ARRAY;
case FLOAT:
case DOUBLE:
return DOUBLE_ARRAY;
case STRING:

36
core/src/main/java/org/apache/druid/math/expr/Function.java
View File

@ -31,6 +31,7 @@ import org.apache.druid.math.expr.vector.ExprVectorProcessor;
import org.apache.druid.math.expr.vector.VectorMathProcessors;
import org.apache.druid.math.expr.vector.VectorProcessors;
import org.apache.druid.math.expr.vector.VectorStringProcessors;
import org.apache.druid.segment.column.TypeSignature;
import org.apache.druid.segment.column.TypeStrategy;
import org.joda.time.DateTime;
import org.joda.time.DateTimeZone;
@ -44,11 +45,11 @@ import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.Comparator;
import java.util.HashSet;
import java.util.List;
import java.util.Objects;
import java.util.Optional;
import java.util.Set;
import java.util.TreeSet;
import java.util.function.BinaryOperator;
import java.util.function.DoubleBinaryOperator;
import java.util.function.LongBinaryOperator;
@ -502,15 +503,16 @@ public interface Function
@Override
ExprEval doApply(ExprEval arrayExpr, ExprEval scalarExpr)
{
final ExpressionType arrayType = arrayExpr.asArrayType();
if (!scalarExpr.type().equals(arrayExpr.elementType())) {
// try to cast
ExprEval coerced = scalarExpr.castTo(arrayExpr.elementType());
return ExprEval.ofArray(arrayExpr.asArrayType(), add(arrayExpr.asArray(), coerced.value()).toArray());
return ExprEval.ofArray(arrayType, add(arrayType.getElementType(), arrayExpr.asArray(), coerced.value()).toArray());
}
return ExprEval.ofArray(arrayExpr.asArrayType(), add(arrayExpr.asArray(), scalarExpr.value()).toArray());
return ExprEval.ofArray(arrayType, add(arrayType.getElementType(), arrayExpr.asArray(), scalarExpr.value()).toArray());
}
abstract <T> Stream<T> add(T[] array, @Nullable T val);
abstract <T> Stream<T> add(TypeSignature<ExprType> elementType, T[] array, @Nullable T val);
}
/**
@ -551,16 +553,18 @@ public interface Function
return lhsExpr;
}
final ExpressionType arrayType = lhsExpr.asArrayType();
if (!lhsExpr.asArrayType().equals(rhsExpr.asArrayType())) {
// try to cast if they types don't match
ExprEval coerced = rhsExpr.castTo(lhsExpr.asArrayType());
ExprEval.ofArray(lhsExpr.asArrayType(), merge(lhsExpr.asArray(), coerced.asArray()).toArray());
ExprEval coerced = rhsExpr.castTo(arrayType);
ExprEval.ofArray(arrayType, merge(arrayType.getElementType(), lhsExpr.asArray(), coerced.asArray()).toArray());
}
return ExprEval.ofArray(lhsExpr.asArrayType(), merge(lhsExpr.asArray(), rhsExpr.asArray()).toArray());
return ExprEval.ofArray(arrayType, merge(arrayType.getElementType(), lhsExpr.asArray(), rhsExpr.asArray()).toArray());
}
abstract <T> Stream<T> merge(T[] array1, T[] array2);
abstract <T> Stream<T> merge(TypeSignature<ExprType> elementType, T[] array1, T[] array2);
}
abstract class ReduceFunction implements Function
@ -3402,7 +3406,7 @@ public interface Function
}
@Override
<T> Stream<T> add(T[] array, @Nullable T val)
<T> Stream<T> add(TypeSignature<ExprType> elementType, T[] array, @Nullable T val)
{
List<T> l = new ArrayList<>(Arrays.asList(array));
l.add(val);
@ -3431,7 +3435,7 @@ public interface Function
}
@Override
<T> Stream<T> add(T[] array, @Nullable T val)
<T> Stream<T> add(TypeSignature<ExprType> elementType, T[] array, @Nullable T val)
{
List<T> l = new ArrayList<>(Arrays.asList(array));
l.add(0, val);
@ -3448,7 +3452,7 @@ public interface Function
}
@Override
<T> Stream<T> merge(T[] array1, T[] array2)
<T> Stream<T> merge(TypeSignature<ExprType> elementType, T[] array1, T[] array2)
{
List<T> l = new ArrayList<>(Arrays.asList(array1));
l.addAll(Arrays.asList(array2));
@ -3465,9 +3469,10 @@ public interface Function
}
@Override
<T> Stream<T> add(T[] array, @Nullable T val)
<T> Stream<T> add(TypeSignature<ExprType> elementType, T[] array, @Nullable T val)
{
Set<T> l = new HashSet<>(Arrays.asList(array));
Set<T> l = new TreeSet<>(elementType.getNullableStrategy());
l.addAll(Arrays.asList(array));
l.add(val);
return l.stream();
}
@ -3482,9 +3487,10 @@ public interface Function
}
@Override
<T> Stream<T> merge(T[] array1, T[] array2)
<T> Stream<T> merge(TypeSignature<ExprType> elementType, T[] array1, T[] array2)
{
Set<T> l = new HashSet<>(Arrays.asList(array1));
Set<T> l = new TreeSet<>(elementType.getNullableStrategy());
l.addAll(Arrays.asList(array1));
l.addAll(Arrays.asList(array2));
return l.stream();
}

24
core/src/test/java/org/apache/druid/math/expr/ParserTest.java
View File

@ -270,27 +270,27 @@ public class ParserTest extends InitializedNullHandlingTest
public void testLiteralArraysExplicitTypedEmpties()
{
// legacy explicit array format
validateConstantExpression("<STRING>[]", new Object[0]);
validateConstantExpression("<DOUBLE>[]", new Object[0]);
validateConstantExpression("<LONG>[]", new Object[0]);
validateConstantExpression("ARRAY<STRING>[]", new Object[0]);
validateConstantExpression("ARRAY<DOUBLE>[]", new Object[0]);
validateConstantExpression("ARRAY<LONG>[]", new Object[0]);
}
@Test
public void testLiteralArraysExplicitAllNull()
{
// legacy explicit array format
validateConstantExpression("<DOUBLE>[null, null, null]", new Object[]{null, null, null});
validateConstantExpression("<LONG>[null, null, null]", new Object[]{null, null, null});
validateConstantExpression("<STRING>[null, null, null]", new Object[]{null, null, null});
validateConstantExpression("ARRAY<DOUBLE>[null, null, null]", new Object[]{null, null, null});
validateConstantExpression("ARRAY<LONG>[null, null, null]", new Object[]{null, null, null});
validateConstantExpression("ARRAY<STRING>[null, null, null]", new Object[]{null, null, null});
}
@Test
public void testLiteralArraysExplicitTypes()
{
// legacy explicit array format
validateConstantExpression("<DOUBLE>[1.0, null, 2000.0]", new Object[]{1.0, null, 2000.0});
validateConstantExpression("<LONG>[3, null, 4]", new Object[]{3L, null, 4L});
validateConstantExpression("<STRING>['foo', 'bar', 'baz']", new Object[]{"foo", "bar", "baz"});
validateConstantExpression("ARRAY<DOUBLE>[1.0, null, 2000.0]", new Object[]{1.0, null, 2000.0});
validateConstantExpression("ARRAY<LONG>[3, null, 4]", new Object[]{3L, null, 4L});
validateConstantExpression("ARRAY<STRING>['foo', 'bar', 'baz']", new Object[]{"foo", "bar", "baz"});
}
@Test
@ -298,11 +298,11 @@ public class ParserTest extends InitializedNullHandlingTest
{
// legacy explicit array format
// explicit typed numeric arrays mixed numeric types should coerce to the correct explicit type
validateConstantExpression("<DOUBLE>[3, null, 4, 2.345]", new Object[]{3.0, null, 4.0, 2.345});
validateConstantExpression("<LONG>[1.0, null, 2000.0]", new Object[]{1L, null, 2000L});
validateConstantExpression("ARRAY<DOUBLE>[3, null, 4, 2.345]", new Object[]{3.0, null, 4.0, 2.345});
validateConstantExpression("ARRAY<LONG>[1.0, null, 2000.0]", new Object[]{1L, null, 2000L});
// explicit typed string arrays should accept any literal and convert to string
validateConstantExpression("<STRING>['1', null, 2000, 1.1]", new Object[]{"1", null, "2000", "1.1"});
validateConstantExpression("ARRAY<STRING>['1', null, 2000, 1.1]", new Object[]{"1", null, "2000", "1.1"});
}
@Test

6
docs/querying/sql.md
View File

@ -376,9 +376,11 @@ In the aggregation functions supported by Druid, only `COUNT`, `ARRAY_AGG`, and
|`ANY_VALUE(expr, maxBytesPerString)`|Like `ANY_VALUE(expr)`, but for strings. The `maxBytesPerString` parameter determines how much aggregation space to allocate per string. Strings longer than this limit will be truncated. This parameter should be set as low as possible, since high values will lead to wasted memory.|`null` if `druid.generic.useDefaultValueForNull=false`, otherwise `''`|
|`GROUPING(expr, expr...)`|Returns a number to indicate which groupBy dimension is included in a row, when using `GROUPING SETS`. Refer to [additional documentation](aggregations.md#grouping-aggregator) on how to infer this number.|N/A|
|`ARRAY_AGG(expr, [size])`|Collects all values of `expr` into an ARRAY, including null values, with `size` in bytes limit on aggregation size (default of 1024 bytes). If the aggregated array grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `ARRAY_AGG` expression is not currently supported, and the ordering of results within the output array may vary depending on processing order.|`null`|
|`ARRAY_AGG(DISTINCT expr, [size])`|Collects all distinct values of `expr` into an ARRAY, including null values, with `size` in bytes limit on aggregation size (default of 1024 bytes) per aggregate. If the aggregated array grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `ARRAY_AGG` expression is not currently supported, and the ordering of results within the output array may vary depending on processing order.|`null`|
|`ARRAY_AGG(DISTINCT expr, [size])`|Collects all distinct values of `expr` into an ARRAY, including null values, with `size` in bytes limit on aggregation size (default of 1024 bytes) per aggregate. If the aggregated array grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `ARRAY_AGG` expression is not currently supported, and the ordering of results will be based on the default for the element type.|`null`|
|`ARRAY_CONCAT_AGG(expr, [size])`|Concatenates all array `expr` into a single ARRAY, with `size` in bytes limit on aggregation size (default of 1024 bytes). Input `expr` _must_ be an array. Null `expr` will be ignored, but any null values within an `expr` _will_ be included in the resulting array. If the aggregated array grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `ARRAY_CONCAT_AGG` expression is not currently supported, and the ordering of results within the output array may vary depending on processing order.|`null`|
|`ARRAY_CONCAT_AGG(DISTINCT expr, [size])`|Concatenates all distinct values of all array `expr` into a single ARRAY, with `size` in bytes limit on aggregation size (default of 1024 bytes) per aggregate. Input `expr` _must_ be an array. Null `expr` will be ignored, but any null values within an `expr` _will_ be included in the resulting array. If the aggregated array grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `ARRAY_CONCAT_AGG` expression is not currently supported, and the ordering of results will be based on the default for the element type.|`null`|
|`STRING_AGG(expr, separator, [size])`|Collects all values of `expr` into a single STRING, ignoring null values. Each value is joined by the `separator` which must be a literal STRING. An optional `size` in bytes can be supplied to limit aggregation size (default of 1024 bytes). If the aggregated string grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `STRING_AGG` expression is not currently supported, and the ordering of results within the output string may vary depending on processing order.|`null` if `druid.generic.useDefaultValueForNull=false`, otherwise `''`|
|`STRING_AGG(DISTINCT expr, separator, [size])`|Collects all distinct values of `expr` into a single STRING, ignoring null values. Each value is joined by the `separator` which must be a literal STRING. An optional `size` in bytes can be supplied to limit aggregation size (default of 1024 bytes). If the aggregated string grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `STRING_AGG` expression is not currently supported, and the ordering of results within the output string may vary depending on processing order.|`null` if `druid.generic.useDefaultValueForNull=false`, otherwise `''`|
|`STRING_AGG(DISTINCT expr, separator, [size])`|Collects all distinct values of `expr` into a single STRING, ignoring null values. Each value is joined by the `separator` which must be a literal STRING. An optional `size` in bytes can be supplied to limit aggregation size (default of 1024 bytes). If the aggregated string grows larger than the maximum size in bytes, the query will fail. Use of `ORDER BY` within the `STRING_AGG` expression is not currently supported, and the ordering of results will be based on the default `STRING` ordering.|`null` if `druid.generic.useDefaultValueForNull=false`, otherwise `''`|
|`BIT_AND(expr)`|Performs a bitwise AND operation on all input values.|`null` if `druid.generic.useDefaultValueForNull=false`, otherwise `0`|
|`BIT_OR(expr)`|Performs a bitwise OR operation on all input values.|`null` if `druid.generic.useDefaultValueForNull=false`, otherwise `0`|
|`BIT_XOR(expr)`|Performs a bitwise XOR operation on all input values.|`null` if `druid.generic.useDefaultValueForNull=false`, otherwise `0`|

22
processing/src/main/java/org/apache/druid/query/aggregation/ExpressionLambdaAggregator.java
View File

@ -24,30 +24,40 @@ import org.apache.druid.math.expr.Expr;
import org.apache.druid.math.expr.ExprEval;
import javax.annotation.Nullable;
import java.util.List;
public class ExpressionLambdaAggregator implements Aggregator
{
private final Expr lambda;
private final List<String> inputColumns;
private final ExpressionLambdaAggregatorInputBindings bindings;
private final int maxSizeBytes;
private final boolean aggregateNullInputs;
private boolean hasValue;
public ExpressionLambdaAggregator(
final Expr lambda,
final ExpressionLambdaAggregatorInputBindings bindings,
final boolean isNullUnlessAggregated,
final ExpressionLambdaAggregatorFactory.FactorizePlan thePlan,
final int maxSizeBytes
)
{
this.lambda = lambda;
this.bindings = bindings;
this.lambda = thePlan.getExpression();
this.bindings = thePlan.getBindings();
this.hasValue = !thePlan.isNullUnlessAggregated();
this.aggregateNullInputs = thePlan.shouldAggregateNullInputs();
this.inputColumns = thePlan.getInputs();
this.maxSizeBytes = maxSizeBytes;
this.hasValue = !isNullUnlessAggregated;
}
@Override
public void aggregate()
{
if (!aggregateNullInputs) {
for (String column : inputColumns) {
if (bindings.get(column) == null) {
return;
}
}
}
final ExprEval<?> eval = lambda.eval(bindings);
final int estimatedSize = eval.type().getNullableStrategy().estimateSizeBytes(eval.value());
if (estimatedSize > maxSizeBytes) {

84
processing/src/main/java/org/apache/druid/query/aggregation/ExpressionLambdaAggregatorFactory.java
View File

@ -48,6 +48,7 @@ import org.apache.druid.segment.virtual.ExpressionPlanner;
import org.apache.druid.segment.virtual.ExpressionSelectors;
import javax.annotation.Nullable;
import java.util.ArrayList;
import java.util.Collections;
import java.util.Comparator;
import java.util.List;
@ -74,6 +75,8 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
private final String initialValueExpressionString;
private final String initialCombineValueExpressionString;
private final boolean isNullUnlessAggregated;
private final boolean shouldAggregateNullInputs;
private final boolean shouldCombineAggregateNullInputs;
private final String combineExpressionString;
@Nullable
@ -103,6 +106,8 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
@JsonProperty("initialValue") final String initialValue,
@JsonProperty("initialCombineValue") @Nullable final String initialCombineValue,
@JsonProperty("isNullUnlessAggregated") @Nullable final Boolean isNullUnlessAggregated,
@JsonProperty("shouldAggregateNullInputs") @Nullable Boolean shouldAggregateNullInputs,
@JsonProperty("shouldCombineAggregateNullInputs") @Nullable Boolean shouldCombineAggregateNullInputs,
@JsonProperty("fold") final String foldExpression,
@JsonProperty("combine") @Nullable final String combineExpression,
@JsonProperty("compare") @Nullable final String compareExpression,
@ -120,6 +125,12 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
this.initialValueExpressionString = initialValue;
this.initialCombineValueExpressionString = initialCombineValue == null ? initialValue : initialCombineValue;
this.isNullUnlessAggregated = isNullUnlessAggregated == null ? NullHandling.sqlCompatible() : isNullUnlessAggregated;
this.shouldAggregateNullInputs = shouldAggregateNullInputs == null || shouldAggregateNullInputs;
if (shouldCombineAggregateNullInputs == null) {
this.shouldCombineAggregateNullInputs = this.shouldAggregateNullInputs;
} else {
this.shouldCombineAggregateNullInputs = shouldCombineAggregateNullInputs;
}
this.foldExpressionString = foldExpression;
if (combineExpression != null) {
this.combineExpressionString = combineExpression;
@ -223,6 +234,18 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
return isNullUnlessAggregated;
}
@JsonProperty("shouldAggregateNullInputs")
public boolean getShouldAggregateNullInputs()
{
return shouldAggregateNullInputs;
}
@JsonProperty("shouldCombineAggregateNullInputs")
public boolean getShouldCombineAggregateNullInputs()
{
return shouldCombineAggregateNullInputs;
}
@JsonProperty("fold")
public String getFoldExpressionString()
{
@ -262,9 +285,12 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
.appendStrings(fields)
.appendString(initialValueExpressionString)
.appendString(initialCombineValueExpressionString)
.appendBoolean(isNullUnlessAggregated)
.appendBoolean(shouldAggregateNullInputs)
.appendBoolean(shouldCombineAggregateNullInputs)
.appendCacheable(foldExpression.get())
.appendCacheable(combineExpression.get())
.appendCacheable(combineExpression.get())
.appendCacheable(compareExpression.get())
.appendCacheable(finalizeExpression.get())
.appendInt(maxSizeBytes.getBytesInInt())
.build();
@ -275,9 +301,7 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
{
FactorizePlan thePlan = new FactorizePlan(metricFactory);
return new ExpressionLambdaAggregator(
thePlan.getExpression(),
thePlan.getBindings(),
isNullUnlessAggregated,
thePlan,
maxSizeBytes.getBytesInInt()
);
}
@ -287,10 +311,7 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
{
FactorizePlan thePlan = new FactorizePlan(metricFactory);
return new ExpressionLambdaBufferAggregator(
thePlan.getExpression(),
thePlan.getInitialValue(),
thePlan.getBindings(),
isNullUnlessAggregated,
thePlan,
maxSizeBytes.getBytesInInt()
);
}
@ -310,6 +331,13 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
@Override
public Object combine(@Nullable Object lhs, @Nullable Object rhs)
{
if (!shouldCombineAggregateNullInputs) {
if (lhs == null) {
return rhs;
} else if (rhs == null) {
return lhs;
}
}
// arbitrarily assign lhs and rhs to accumulator and aggregator name inputs to re-use combine function
return combineExpression.get().eval(
combineBindings.get().withBinding(accumulatorId, lhs).withBinding(name, rhs)
@ -353,6 +381,8 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
initialValueExpressionString,
initialCombineValueExpressionString,
isNullUnlessAggregated,
shouldAggregateNullInputs,
shouldCombineAggregateNullInputs,
foldExpressionString,
combineExpressionString,
compareExpressionString,
@ -373,6 +403,8 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
initialValueExpressionString,
initialCombineValueExpressionString,
isNullUnlessAggregated,
shouldAggregateNullInputs,
shouldCombineAggregateNullInputs,
foldExpressionString,
combineExpressionString,
compareExpressionString,
@ -433,6 +465,8 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
&& initialValueExpressionString.equals(that.initialValueExpressionString)
&& initialCombineValueExpressionString.equals(that.initialCombineValueExpressionString)
&& isNullUnlessAggregated == that.isNullUnlessAggregated
&& shouldAggregateNullInputs == that.shouldAggregateNullInputs
&& shouldCombineAggregateNullInputs == that.shouldCombineAggregateNullInputs
&& combineExpressionString.equals(that.combineExpressionString)
&& Objects.equals(compareExpressionString, that.compareExpressionString)
&& Objects.equals(finalizeExpressionString, that.finalizeExpressionString);
@ -449,6 +483,8 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
initialValueExpressionString,
initialCombineValueExpressionString,
isNullUnlessAggregated,
shouldAggregateNullInputs,
shouldCombineAggregateNullInputs,
combineExpressionString,
compareExpressionString,
finalizeExpressionString,
@ -466,7 +502,9 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
", foldExpressionString='" + foldExpressionString + '\'' +
", initialValueExpressionString='" + initialValueExpressionString + '\'' +
", initialCombineValueExpressionString='" + initialCombineValueExpressionString + '\'' +
", nullUnlessAggregated='" + isNullUnlessAggregated + '\'' +
", isNullUnlessAggregated='" + isNullUnlessAggregated + '\'' +
", shouldAggregateNullInputs='" + shouldAggregateNullInputs + '\'' +
", shouldCombineAggregateNullInputs='" + shouldCombineAggregateNullInputs + '\'' +
", combineExpressionString='" + combineExpressionString + '\'' +
", compareExpressionString='" + compareExpressionString + '\'' +
", finalizeExpressionString='" + finalizeExpressionString + '\'' +
@ -477,23 +515,29 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
/**
* Determine how to factorize the aggregator
*/
private class FactorizePlan
public class FactorizePlan
{
private final ExpressionPlan plan;
private final ExprEval<?> seed;
private final ExpressionLambdaAggregatorInputBindings bindings;
private final List<String> inputs;
private final boolean aggregateNullInputs;
FactorizePlan(ColumnSelectorFactory metricFactory)
{
this.inputs = new ArrayList<>();
if (fields != null) {
// if fields are set, we are accumulating from raw inputs, use fold expression
plan = ExpressionPlanner.plan(inspectorWithAccumulator(metricFactory), foldExpression.get());
seed = initialValue.get();
aggregateNullInputs = shouldAggregateNullInputs;
} else {
// else we are merging intermediary results, use combine expression
plan = ExpressionPlanner.plan(inspectorWithAccumulator(metricFactory), combineExpression.get());
seed = initialCombineValue.get();
aggregateNullInputs = shouldCombineAggregateNullInputs;
}
bindings = new ExpressionLambdaAggregatorInputBindings(
@ -501,6 +545,11 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
accumulatorId,
seed
);
for (String input : plan.getAnalysis().getRequiredBindingsList()) {
if (!input.equals(accumulatorId)) {
this.inputs.add(input);
}
}
}
public Expr getExpression()
@ -523,6 +572,21 @@ public class ExpressionLambdaAggregatorFactory extends AggregatorFactory
return bindings;
}
public List<String> getInputs()
{
return inputs;
}
public boolean shouldAggregateNullInputs()
{
return aggregateNullInputs;
}
public boolean isNullUnlessAggregated()
{
return isNullUnlessAggregated;
}
private ColumnInspector inspectorWithAccumulator(ColumnInspector inspector)
{
return new ColumnInspector()

25
processing/src/main/java/org/apache/druid/query/aggregation/ExpressionLambdaBufferAggregator.java
View File

@ -25,6 +25,7 @@ import org.apache.druid.math.expr.ExpressionType;
import javax.annotation.Nullable;
import java.nio.ByteBuffer;
import java.util.List;
public class ExpressionLambdaBufferAggregator implements BufferAggregator
{
@ -35,21 +36,22 @@ public class ExpressionLambdaBufferAggregator implements BufferAggregator
private final ExpressionLambdaAggregatorInputBindings bindings;
private final int maxSizeBytes;
private final boolean isNullUnlessAggregated;
private final boolean aggregateNullInputs;
private final List<String> inputColumns;
private final ExpressionType outputType;
public ExpressionLambdaBufferAggregator(
Expr lambda,
ExprEval<?> initialValue,
ExpressionLambdaAggregatorInputBindings bindings,
boolean isNullUnlessAggregated,
final ExpressionLambdaAggregatorFactory.FactorizePlan thePlan,
int maxSizeBytes
)
{
this.lambda = lambda;
this.initialValue = initialValue;
this.lambda = thePlan.getExpression();
this.initialValue = thePlan.getInitialValue();
this.outputType = initialValue.type();
this.bindings = bindings;
this.isNullUnlessAggregated = isNullUnlessAggregated;
this.bindings = thePlan.getBindings();
this.isNullUnlessAggregated = thePlan.isNullUnlessAggregated();
this.aggregateNullInputs = thePlan.shouldAggregateNullInputs();
this.inputColumns = thePlan.getInputs();
this.maxSizeBytes = maxSizeBytes;
}
@ -66,6 +68,13 @@ public class ExpressionLambdaBufferAggregator implements BufferAggregator
@Override
public void aggregate(ByteBuffer buf, int position)
{
if (!aggregateNullInputs) {
for (String column : inputColumns) {
if (bindings.get(column) == null) {
return;
}
}
}
ExprEval<?> acc = ExprEval.deserialize(buf, position, outputType);
bindings.setAccumulator(acc);
ExprEval<?> newAcc = lambda.eval(bindings);

20
processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/GroupByQueryEngineV2.java
View File

@ -295,7 +295,7 @@ public class GroupByQueryEngineV2
}
// We cannot support array-based aggregation on array based grouping as we we donot have all the indexes up front
// to allocate appropriate values
if (dimensions.get(0).getOutputType().equals(ColumnType.STRING_ARRAY)) {
if (dimensions.get(0).getOutputType().isArray()) {
return -1;
}
@ -327,10 +327,11 @@ public class GroupByQueryEngineV2
}
/**
* Checks whether all "dimensions" are either single-valued,
* or STRING_ARRAY, in case we don't want to explode the underline multi value column,
* or if allowed, nonexistent. Since non-existent columnselectors will show up as full of nulls they are effectively
* single valued, however they can also be null during broker merge, for example with an 'inline' datasource subquery.
* Checks whether all "dimensions" are either single-valued, or if the input column or output dimension spec has
* specified a type that {@link ColumnType#isArray()}. Both cases indicate we don't want to explode the under-lying
* multi value column. Since selectors on non-existent columns will show up as full of nulls, they are effectively
* single valued, however capabilites on columns can also be null, for example during broker merge with an 'inline'
* datasource subquery, so we only return true from this method when capabilities are fully known.
*/
public static boolean hasNoExplodingDimensions(
final ColumnInspector inspector,
@ -347,10 +348,13 @@ public class GroupByQueryEngineV2
return false;
}
// Now check column capabilities, which must be present and explicitly not multi-valued
// Now check column capabilities, which must be present and explicitly not multi-valued and not arrays
final ColumnCapabilities columnCapabilities = inspector.getColumnCapabilities(dimension.getDimension());
return dimension.getOutputType().equals(ColumnType.STRING_ARRAY)
|| (columnCapabilities != null && columnCapabilities.hasMultipleValues().isFalse());
return dimension.getOutputType().isArray()
|| (columnCapabilities != null
&& columnCapabilities.hasMultipleValues().isFalse()
&& !columnCapabilities.isArray()
);
});
}

3
processing/src/main/java/org/apache/druid/query/groupby/epinephelinae/vector/VectorGroupByEngine.java
View File

@ -47,7 +47,6 @@ import org.apache.druid.segment.ColumnProcessors;
import org.apache.druid.segment.StorageAdapter;
import org.apache.druid.segment.VirtualColumns;
import org.apache.druid.segment.column.ColumnCapabilities;
import org.apache.druid.segment.column.ValueType;
import org.apache.druid.segment.filter.Filters;
import org.apache.druid.segment.vector.VectorColumnSelectorFactory;
import org.apache.druid.segment.vector.VectorCursor;
@ -106,7 +105,7 @@ public class VectorGroupByEngine
return false;
}
if (dimension.getOutputType().getType().equals(ValueType.ARRAY)) {
if (dimension.getOutputType().isArray()) {
// group by on arrays is not currently supported in the vector processing engine
return false;
}

3
processing/src/main/java/org/apache/druid/segment/DimensionHandlerUtils.java
View File

@ -264,6 +264,9 @@ public final class DimensionHandlerUtils
capabilities.isDictionaryEncoded().isTrue() &&
fn.getExtractionType() == ExtractionFn.ExtractionType.ONE_TO_ONE
)
.setHasMultipleValues(
capabilities.hasMultipleValues().isMaybeTrue() || capabilities.isArray()
)
.setDictionaryValuesSorted(
capabilities.isDictionaryEncoded().isTrue() && fn.preservesOrdering()
);

7
processing/src/main/java/org/apache/druid/segment/column/ColumnCapabilitiesImpl.java
View File

@ -154,7 +154,6 @@ public class ColumnCapabilitiesImpl implements ColumnCapabilities
/**
* Create a no frills, simple column with {@link ValueType} set and everything else false
* @param valueType
*/
public static ColumnCapabilitiesImpl createSimpleNumericColumnCapabilities(TypeSignature<ValueType> valueType)
{
@ -187,14 +186,12 @@ public class ColumnCapabilitiesImpl implements ColumnCapabilities
}
/**
* Similar to {@link #createSimpleNumericColumnCapabilities} except {@link #hasMultipleValues} is explicitly true
* and {@link #hasNulls} is not set
* @param valueType
* Similar to {@link #createSimpleNumericColumnCapabilities} except {@link #hasNulls} is not set
*/
public static ColumnCapabilitiesImpl createSimpleArrayColumnCapabilities(TypeSignature<ValueType> valueType)
{
ColumnCapabilitiesImpl builder = new ColumnCapabilitiesImpl().setType(valueType)
.setHasMultipleValues(true)
.setHasMultipleValues(false)
.setHasBitmapIndexes(false)
.setDictionaryEncoded(false)
.setDictionaryValuesSorted(false)

4
processing/src/main/java/org/apache/druid/segment/virtual/ExpressionPlan.java
View File

@ -275,8 +275,8 @@ public class ExpressionPlan
// the complete set of input types
if (any(Trait.NON_SCALAR_OUTPUT, Trait.NEEDS_APPLIED)) {
// if the hint requested a string, use a string
if (Types.is(outputTypeHint, ValueType.STRING)) {
return ColumnCapabilitiesImpl.createSimpleArrayColumnCapabilities(ColumnType.STRING);
if (Types.is(outputTypeHint, ValueType.STRING) || inferredValueType.is(ValueType.STRING)) {
return ColumnCapabilitiesImpl.createSimpleSingleValueStringColumnCapabilities().setHasMultipleValues(true);
}
// maybe something is looking for a little fun and wants arrays? let whatever it is through
return ColumnCapabilitiesImpl.createSimpleArrayColumnCapabilities(ExpressionType.toColumnType(outputType));

11
processing/src/main/java/org/apache/druid/segment/virtual/ExpressionPlanner.java
View File

@ -107,10 +107,13 @@ public class ExpressionPlanner
// automatic transformation to map across multi-valued inputs (or row by row detection in the worst case)
if (ExpressionPlan.none(traits, ExpressionPlan.Trait.SINGLE_INPUT_SCALAR)) {
final Set<String> definitelyMultiValued = new HashSet<>();
final Set<String> definitelyArray = new HashSet<>();
for (String column : analysis.getRequiredBindings()) {
final ColumnCapabilities capabilities = inspector.getColumnCapabilities(column);
if (capabilities != null) {
if (capabilities.hasMultipleValues().isTrue()) {
if (capabilities.isArray()) {
definitelyArray.add(column);
} else if (capabilities.is(ValueType.STRING) && capabilities.hasMultipleValues().isTrue()) {
definitelyMultiValued.add(column);
} else if (capabilities.is(ValueType.STRING) &&
capabilities.hasMultipleValues().isMaybeTrue() &&
@ -126,7 +129,11 @@ public class ExpressionPlanner
// find any inputs which will need implicitly mapped across multi-valued rows
needsApplied =
columns.stream()
.filter(c -> definitelyMultiValued.contains(c) && !analysis.getArrayBindings().contains(c))
.filter(
c -> !definitelyArray.contains(c)
&& definitelyMultiValued.contains(c)
&& !analysis.getArrayBindings().contains(c)
)
.collect(Collectors.toList());
// if any multi-value inputs, set flag for non-scalar inputs

4
processing/src/main/java/org/apache/druid/segment/virtual/ExpressionVirtualColumn.java
View File

@ -30,6 +30,7 @@ import com.google.common.base.Suppliers;
import org.apache.druid.java.util.common.logger.Logger;
import org.apache.druid.math.expr.Expr;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.math.expr.ExpressionProcessing;
import org.apache.druid.math.expr.Parser;
import org.apache.druid.query.cache.CacheKeyBuilder;
import org.apache.druid.query.dimension.DimensionSpec;
@ -178,8 +179,7 @@ public class ExpressionVirtualColumn implements VirtualColumn
// inputs or because of unimplemented methods on expression implementations themselves, or, because a
// ColumnInspector is not available
// array types must not currently escape from the expression system
if (outputType != null && outputType.isArray()) {
if (ExpressionProcessing.processArraysAsMultiValueStrings() && outputType != null && outputType.isArray()) {
return new ColumnCapabilitiesImpl().setType(ColumnType.STRING).setHasMultipleValues(true);
}

138
processing/src/test/java/org/apache/druid/query/aggregation/ExpressionLambdaAggregatorFactoryTest.java
View File

@ -60,6 +60,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0.0",
"10.0",
true,
true,
false,
"customAccumulator + some_column + some_other_column",
"customAccumulator + expr_agg_name",
"if (o1 > o2, if (o1 == o2, 0, 1), -1)",
@ -105,6 +107,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"x + y",
null,
true,
false,
false,
"__acc + some_column + some_other_column",
"__acc + expr_agg_name",
null,
@ -129,6 +133,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0.0",
"x + y",
true,
false,
false,
"__acc + some_column + some_other_column",
"__acc + expr_agg_name",
null,
@ -150,6 +156,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0",
null,
true,
false,
false,
"__acc + x",
null,
null,
@ -161,6 +169,60 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
Assert.assertEquals(1L, agg.combine(0L, 1L));
}
@Test
public void testCombineExpressionIgnoresNullsIfCombineSkipsNulls()
{
ExpressionLambdaAggregatorFactory agg = new ExpressionLambdaAggregatorFactory(
"expr_agg_name",
ImmutableSet.of("x"),
null,
"ARRAY<STRING>",
"ARRAY<STRING>[]",
true,
false,
false,
"array_append(__acc, x)",
"array_concat(__acc, expr_agg_name)",
null,
null,
null,
TestExprMacroTable.INSTANCE
);
Assert.assertArrayEquals(new Object[]{"hello"}, (Object[]) agg.combine(null, new Object[]{"hello"}));
Assert.assertArrayEquals(
new Object[]{"hello", "world"},
(Object[]) agg.combine(new Object[]{"hello"}, new Object[]{"world"})
);
}
@Test
public void testCombineExpressionDoesntIgnoreNullsIfCombineDoesntSkipsNulls()
{
ExpressionLambdaAggregatorFactory agg = new ExpressionLambdaAggregatorFactory(
"expr_agg_name",
ImmutableSet.of("x"),
null,
"ARRAY<STRING>",
"ARRAY<STRING>[]",
true,
false,
true,
"array_append(__acc, x)",
"array_concat(__acc, expr_agg_name)",
null,
null,
null,
TestExprMacroTable.INSTANCE
);
Assert.assertNull(agg.combine(null, new Object[]{"hello"}));
Assert.assertArrayEquals(
new Object[]{"hello", "world"},
(Object[]) agg.combine(new Object[]{"hello"}, new Object[]{"world"})
);
}
@Test
public void testFinalizeCanDo()
{
@ -171,6 +233,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0",
null,
true,
false,
false,
"__acc + x",
null,
null,
@ -190,8 +254,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("x"),
null,
"0",
"<LONG>[]",
"ARRAY<LONG>[]",
true,
true,
false,
"array_set_add(__acc, x)",
"array_set_add_all(__acc, expr_agg_name)",
null,
@ -214,6 +280,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"''",
"''",
true,
true,
true,
"concat(__acc, some_column, some_other_column)",
"concat(__acc, expr_agg_name)",
null,
@ -237,6 +305,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0",
null,
null,
false,
false,
"__acc + some_column + some_other_column",
"__acc + expr_agg_name",
null,
@ -260,6 +330,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0.0",
null,
null,
false,
false,
"__acc + some_column + some_other_column",
"__acc + expr_agg_name",
null,
@ -281,8 +353,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"''",
"<STRING>[]",
"ARRAY<STRING>[]",
null,
false,
false,
"concat(__acc, some_column, some_other_column)",
"array_set_add(__acc, expr_agg_name)",
null,
@ -304,8 +378,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"''",
"<STRING>[]",
"ARRAY<STRING>[]",
null,
false,
false,
"concat(__acc, some_column, some_other_column)",
"array_set_add(__acc, expr_agg_name)",
null,
@ -327,8 +403,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"0",
"<LONG>[]",
"ARRAY<LONG>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, expr_agg_name)",
null,
@ -350,8 +428,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"0",
"<LONG>[]",
"ARRAY<LONG>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, expr_agg_name)",
null,
@ -373,8 +453,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"0.0",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, expr_agg_name)",
null,
@ -396,8 +478,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"0.0",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, expr_agg_name)",
null,
@ -421,6 +505,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"hyper_unique()",
null,
null,
false,
false,
"hyper_unique_add(some_column, __acc)",
"hyper_unique_add(__acc, expr_agg_name)",
null,
@ -444,6 +530,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"hyper_unique()",
null,
null,
false,
false,
"hyper_unique_add(some_column, __acc)",
"hyper_unique_add(__acc, expr_agg_name)",
null,
@ -473,6 +561,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"''",
"''",
null,
false,
false,
"concat(__acc, some_column, some_other_column)",
"concat(__acc, string_expr)",
null,
@ -487,6 +577,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0.0",
null,
null,
false,
false,
"__acc + some_column + some_other_column",
"__acc + double_expr",
null,
@ -501,6 +593,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"0",
null,
null,
false,
false,
"__acc + some_column + some_other_column",
"__acc + long_expr",
null,
@ -512,9 +606,11 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"string_array_expr",
ImmutableSet.of("some_column", "some_other_column"),
null,
"<STRING>[]",
"<STRING>[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
null,
true,
false,
"array_set_add(__acc, concat(some_column, some_other_column))",
"array_set_add_all(__acc, string_array_expr)",
null,
@ -527,8 +623,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column_expr"),
null,
"0.0",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, double_array)",
null,
@ -541,8 +639,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"0",
"<LONG>[]",
"ARRAY<LONG>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, long_array_expr)",
null,
@ -555,8 +655,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"''",
"<STRING>[]",
"ARRAY<STRING>[]",
null,
false,
false,
"concat(__acc, some_column, some_other_column)",
"array_set_add(__acc, string_array_expr)",
null,
@ -569,8 +671,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column_expr"),
null,
"0.0",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, double_array)",
null,
@ -583,8 +687,10 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
ImmutableSet.of("some_column", "some_other_column"),
null,
"0",
"<LONG>[]",
"ARRAY<LONG>[]",
null,
false,
false,
"__acc + some_column + some_other_column",
"array_set_add(__acc, long_array_expr)",
null,
@ -599,6 +705,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"hyper_unique()",
null,
null,
false,
false,
"hyper_unique_add(some_column, __acc)",
"hyper_unique_add(__acc, expr_agg_name)",
null,
@ -613,6 +721,8 @@ public class ExpressionLambdaAggregatorFactoryTest extends InitializedNullHandli
"hyper_unique()",
null,
null,
false,
false,
"hyper_unique_add(some_column, __acc)",
"hyper_unique_add(__acc, expr_agg_name)",
null,

39
processing/src/test/java/org/apache/druid/query/groupby/GroupByQueryRunnerTest.java
View File

@ -1734,8 +1734,11 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
ExprMacroTable.nil()
))
.setDimensions(
new ExtractionDimensionSpec("v0", "alias", ColumnType.STRING,
new SubstringDimExtractionFn(1, 1)
new ExtractionDimensionSpec(
"v0",
"alias",
ColumnType.STRING,
new SubstringDimExtractionFn(1, 1)
)
)
@ -12031,6 +12034,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"0",
null,
false,
false,
false,
"__acc + 1",
"__acc + rows",
null,
@ -12045,6 +12050,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"0.0",
null,
null,
false,
false,
"__acc + index",
null,
null,
@ -12265,6 +12272,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"hyper_unique()",
null,
null,
false,
false,
"hyper_unique_add(quality, __acc)",
"hyper_unique_add(carExpr, __acc)",
null,
@ -12311,6 +12320,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"hyper_unique()",
null,
null,
false,
false,
"hyper_unique_add(quality, __acc)",
null,
null,
@ -12353,6 +12364,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"0",
null,
false,
false,
false,
"__acc + 1",
"__acc + rows",
null,
@ -12367,6 +12380,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"0.0",
null,
true,
false,
false,
"__acc + index",
null,
null,
@ -12381,6 +12396,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"[]",
null,
null,
true,
false,
"array_set_add(acc, market)",
"array_set_add_all(acc, array_agg_distinct)",
null,
@ -12451,7 +12468,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"idx",
2871.8866900000003d,
"array_agg_distinct",
new String[] {"upfront", "spot", "total_market"}
new String[] {"spot", "total_market", "upfront"}
),
makeRow(
query,
@ -12475,7 +12492,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"idx",
2900.798647d,
"array_agg_distinct",
new String[] {"upfront", "spot", "total_market"}
new String[] {"spot", "total_market", "upfront"}
),
makeRow(
query,
@ -12560,7 +12577,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"idx",
2448.830613d,
"array_agg_distinct",
new String[] {"upfront", "spot", "total_market"}
new String[] {"spot", "total_market", "upfront"}
),
makeRow(
query,
@ -12584,7 +12601,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"idx",
2506.415148d,
"array_agg_distinct",
new String[] {"upfront", "spot", "total_market"}
new String[] {"spot", "total_market", "upfront"}
),
makeRow(
query,
@ -12640,6 +12657,8 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"[]",
null,
null,
true,
false,
"array_set_add(acc, placementish)",
"array_set_add_all(acc, array_agg_distinct)",
null,
@ -12714,7 +12733,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"alias",
"technology",
"array_agg_distinct",
new String[] {"t", "preferred"}
new String[] {"preferred", "t"}
),
makeRow(
query,
@ -12722,7 +12741,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"alias",
"travel",
"array_agg_distinct",
new String[] {"t", "preferred"}
new String[] {"preferred", "t"}
),
makeRow(
@ -12787,7 +12806,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"alias",
"technology",
"array_agg_distinct",
new String[] {"t", "preferred"}
new String[] {"preferred", "t"}
),
makeRow(
query,
@ -12795,7 +12814,7 @@ public class GroupByQueryRunnerTest extends InitializedNullHandlingTest
"alias",
"travel",
"array_agg_distinct",
new String[] {"t", "preferred"}
new String[] {"preferred", "t"}
)
);

2
processing/src/test/java/org/apache/druid/query/groupby/GroupByTimeseriesQueryRunnerTest.java
View File

@ -337,6 +337,8 @@ public class GroupByTimeseriesQueryRunnerTest extends TimeseriesQueryRunnerTest
"[]",
null,
null,
true,
false,
"array_set_add(acc, market)",
"array_set_add_all(acc, array_agg_distinct)",
null,

16
processing/src/test/java/org/apache/druid/query/timeseries/TimeseriesQueryRunnerTest.java
View File

@ -3031,6 +3031,8 @@ public class TimeseriesQueryRunnerTest extends InitializedNullHandlingTest
"0",
null,
false,
false,
false,
"__acc + 1",
"__acc + diy_count",
null,
@ -3045,6 +3047,8 @@ public class TimeseriesQueryRunnerTest extends InitializedNullHandlingTest
"0.0",
null,
null,
false,
false,
"__acc + index",
null,
null,
@ -3057,8 +3061,10 @@ public class TimeseriesQueryRunnerTest extends InitializedNullHandlingTest
ImmutableSet.of("index"),
null,
"0.0",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
null,
false,
false,
"__acc + index",
"array_concat(__acc, diy_decomposed_sum)",
null,
@ -3073,6 +3079,8 @@ public class TimeseriesQueryRunnerTest extends InitializedNullHandlingTest
"[]",
null,
null,
true,
false,
"array_set_add(acc, market)",
"array_set_add_all(acc, array_agg_distinct)",
null,
@ -3094,7 +3102,7 @@ public class TimeseriesQueryRunnerTest extends InitializedNullHandlingTest
"diy_count", 13L,
"diy_sum", 6626.151569,
"diy_decomposed_sum", 6626.151569,
"array_agg_distinct", new String[] {"upfront", "spot", "total_market"}
"array_agg_distinct", new String[] {"spot", "total_market", "upfront"}
)
)
),
@ -3105,7 +3113,7 @@ public class TimeseriesQueryRunnerTest extends InitializedNullHandlingTest
"diy_count", 13L,
"diy_sum", 5833.209718,
"diy_decomposed_sum", 5833.209718,
"array_agg_distinct", new String[] {"upfront", "spot", "total_market"}
"array_agg_distinct", new String[] {"spot", "total_market", "upfront"}
)
)
)
@ -3137,6 +3145,8 @@ public class TimeseriesQueryRunnerTest extends InitializedNullHandlingTest
"[]",
null,
null,
true,
false,
"array_set_add(acc, market)",
"array_set_add_all(acc, array_agg_distinct)",
null,

14
processing/src/test/java/org/apache/druid/query/topn/TopNQueryRunnerTest.java
View File

@ -5988,6 +5988,8 @@ public class TopNQueryRunnerTest extends InitializedNullHandlingTest
"0",
null,
false,
false,
false,
"__acc + 1",
"__acc + diy_count",
null,
@ -6002,6 +6004,8 @@ public class TopNQueryRunnerTest extends InitializedNullHandlingTest
"0.0",
null,
null,
false,
false,
"__acc + index",
null,
null,
@ -6014,8 +6018,10 @@ public class TopNQueryRunnerTest extends InitializedNullHandlingTest
ImmutableSet.of("index"),
null,
"0.0",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
null,
false,
false,
"__acc + index",
"array_concat(__acc, diy_decomposed_sum)",
null,
@ -6030,6 +6036,8 @@ public class TopNQueryRunnerTest extends InitializedNullHandlingTest
"[]",
null,
null,
true,
false,
"array_set_add(acc, quality)",
"array_set_add_all(acc, array_agg_distinct)",
"if(array_length(o1) > array_length(o2), 1, if (array_length(o1) == array_length(o2), 0, -1))",
@ -6052,7 +6060,7 @@ public class TopNQueryRunnerTest extends InitializedNullHandlingTest
.put("diy_count", 837L)
.put("diy_sum", 95606.57232284546D)
.put("diy_decomposed_sum", 95606.57232284546D)
.put("array_agg_distinct", new String[]{"mezzanine", "news", "premium", "business", "entertainment", "health", "technology", "automotive", "travel"})
.put("array_agg_distinct", new String[]{"automotive", "business", "entertainment", "health", "mezzanine", "news", "premium", "technology", "travel"})
.build(),
ImmutableMap.<String, Object>builder()
.put(QueryRunnerTestHelper.MARKET_DIMENSION, "total_market")
@ -6101,6 +6109,8 @@ public class TopNQueryRunnerTest extends InitializedNullHandlingTest
"hyper_unique()",
null,
null,
false,
false,
"hyper_unique_add(quality, __acc)",
"hyper_unique_add(carExpr, __acc)",
null,

6
processing/src/test/java/org/apache/druid/segment/RowBasedColumnSelectorFactoryTest.java
View File

@ -147,7 +147,7 @@ public class RowBasedColumnSelectorFactoryTest extends InitializedNullHandlingTe
Assert.assertFalse(caps.isDictionaryEncoded().isTrue());
Assert.assertFalse(caps.areDictionaryValuesSorted().isTrue());
Assert.assertFalse(caps.areDictionaryValuesUnique().isTrue());
Assert.assertTrue(caps.hasMultipleValues().isTrue());
Assert.assertTrue(caps.hasMultipleValues().isFalse());
Assert.assertFalse(caps.hasSpatialIndexes());
}
@ -161,7 +161,7 @@ public class RowBasedColumnSelectorFactoryTest extends InitializedNullHandlingTe
Assert.assertFalse(caps.isDictionaryEncoded().isTrue());
Assert.assertFalse(caps.areDictionaryValuesSorted().isTrue());
Assert.assertFalse(caps.areDictionaryValuesUnique().isTrue());
Assert.assertTrue(caps.hasMultipleValues().isTrue());
Assert.assertTrue(caps.hasMultipleValues().isFalse());
Assert.assertFalse(caps.hasSpatialIndexes());
}
@ -175,7 +175,7 @@ public class RowBasedColumnSelectorFactoryTest extends InitializedNullHandlingTe
Assert.assertFalse(caps.isDictionaryEncoded().isTrue());
Assert.assertFalse(caps.areDictionaryValuesSorted().isTrue());
Assert.assertFalse(caps.areDictionaryValuesUnique().isTrue());
Assert.assertTrue(caps.hasMultipleValues().isTrue());
Assert.assertTrue(caps.hasMultipleValues().isFalse());
Assert.assertFalse(caps.hasSpatialIndexes());
}

2
processing/src/test/java/org/apache/druid/segment/virtual/ExpressionPlannerTest.java
View File

@ -644,7 +644,7 @@ public class ExpressionPlannerTest extends InitializedNullHandlingTest
Assert.assertFalse(inferred.isDictionaryEncoded().isMaybeTrue());
Assert.assertFalse(inferred.areDictionaryValuesSorted().isMaybeTrue());
Assert.assertFalse(inferred.areDictionaryValuesUnique().isMaybeTrue());
Assert.assertTrue(inferred.hasMultipleValues().isTrue());
Assert.assertFalse(inferred.hasMultipleValues().isMaybeTrue());
Assert.assertFalse(inferred.hasBitmapIndexes());
Assert.assertFalse(inferred.hasSpatialIndexes());

178
sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/ArrayConcatSqlAggregator.java Normal file
View File

@ -0,0 +1,178 @@
/*
* Licensed to the Apache Software Foundation (ASF) under one
* or more contributor license agreements. See the NOTICE file
* distributed with this work for additional information
* regarding copyright ownership. The ASF licenses this file
* to you under the Apache License, Version 2.0 (the
* "License"); you may not use this file except in compliance
* with the License. You may obtain a copy of the License at
*
* http://www.apache.org/licenses/LICENSE-2.0
*
* Unless required by applicable law or agreed to in writing,
* software distributed under the License is distributed on an
* "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
* KIND, either express or implied. See the License for the
* specific language governing permissions and limitations
* under the License.
*/
package org.apache.druid.sql.calcite.aggregation.builtin;
import com.google.common.collect.ImmutableSet;
import org.apache.calcite.rel.core.AggregateCall;
import org.apache.calcite.rel.core.Project;
import org.apache.calcite.rex.RexBuilder;
import org.apache.calcite.rex.RexLiteral;
import org.apache.calcite.rex.RexNode;
import org.apache.calcite.sql.SqlAggFunction;
import org.apache.calcite.sql.SqlFunctionCategory;
import org.apache.calcite.sql.SqlKind;
import org.apache.calcite.sql.type.InferTypes;
import org.apache.calcite.sql.type.OperandTypes;
import org.apache.calcite.sql.type.ReturnTypes;
import org.apache.calcite.util.Optionality;
import org.apache.druid.java.util.common.HumanReadableBytes;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.math.expr.ExpressionType;
import org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory;
import org.apache.druid.segment.VirtualColumn;
import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.column.RowSignature;
import org.apache.druid.sql.calcite.aggregation.Aggregation;
import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
import org.apache.druid.sql.calcite.expression.DruidExpression;
import org.apache.druid.sql.calcite.expression.Expressions;
import org.apache.druid.sql.calcite.planner.Calcites;
import org.apache.druid.sql.calcite.planner.PlannerContext;
import org.apache.druid.sql.calcite.rel.VirtualColumnRegistry;
import javax.annotation.Nullable;
import java.util.List;
import java.util.stream.Collectors;
public class ArrayConcatSqlAggregator implements SqlAggregator
{
private static final String NAME = "ARRAY_CONCAT_AGG";
private static final SqlAggFunction FUNCTION = new ArrayConcatAggFunction();
@Override
public SqlAggFunction calciteFunction()
{
return FUNCTION;
}
@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));
final ExprMacroTable macroTable = plannerContext.getExprMacroTable();
final String fieldName;
final ColumnType druidType = Calcites.getValueTypeForRelDataTypeFull(aggregateCall.getType());
if (druidType == null || !druidType.isArray()) {
// must be an array
return null;
}
final String initialvalue = ExpressionType.fromColumnTypeStrict(druidType).asTypeString() + "[]";
if (arg.isDirectColumnAccess()) {
fieldName = arg.getDirectColumn();
} else {
VirtualColumn vc = virtualColumnRegistry.getOrCreateVirtualColumnForExpression(plannerContext, arg, druidType);
fieldName = vc.getOutputName();
}
if (aggregateCall.isDistinct()) {
return Aggregation.create(
new ExpressionLambdaAggregatorFactory(
name,
ImmutableSet.of(fieldName),
null,
initialvalue,
null,
true,
false,
false,
StringUtils.format("array_set_add_all(\"__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,
false,
false,
StringUtils.format("array_concat(\"__acc\", \"%s\")", fieldName),
StringUtils.format("array_concat(\"__acc\", \"%s\")", name),
null,
null,
maxSizeBytes != null ? new HumanReadableBytes(maxSizeBytes) : null,
macroTable
)
);
}
}
private static class ArrayConcatAggFunction extends SqlAggFunction
{
ArrayConcatAggFunction()
{
super(
NAME,
null,
SqlKind.OTHER_FUNCTION,
ReturnTypes.ARG0,
InferTypes.ANY_NULLABLE,
OperandTypes.or(
OperandTypes.ARRAY,
OperandTypes.sequence(
StringUtils.format("'%s'(expr, maxSizeBytes)", NAME),
OperandTypes.ARRAY,
OperandTypes.POSITIVE_INTEGER_LITERAL
)
),
SqlFunctionCategory.USER_DEFINED_FUNCTION,
false,
false,
Optionality.IGNORED
);
}
}
}

51
sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/ArraySqlAggregator.java
View File

@ -34,11 +34,11 @@ import org.apache.calcite.sql.type.InferTypes;
import org.apache.calcite.sql.type.OperandTypes;
import org.apache.calcite.sql.type.SqlReturnTypeInference;
import org.apache.calcite.sql.type.SqlTypeFamily;
import org.apache.calcite.sql.type.SqlTypeUtil;
import org.apache.calcite.util.Optionality;
import org.apache.druid.java.util.common.HumanReadableBytes;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.math.expr.ExpressionType;
import org.apache.druid.query.aggregation.ExpressionLambdaAggregatorFactory;
import org.apache.druid.segment.VirtualColumn;
import org.apache.druid.segment.column.ColumnType;
@ -55,7 +55,6 @@ import org.apache.druid.sql.calcite.table.RowSignatures;
import javax.annotation.Nullable;
import java.util.List;
import java.util.Objects;
import java.util.stream.Collectors;
public class ArraySqlAggregator implements SqlAggregator
@ -83,32 +82,26 @@ public class ArraySqlAggregator implements SqlAggregator
boolean finalizeAggregations
)
{
final List<DruidExpression> arguments = aggregateCall
final List<RexNode> 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;
}
Integer maxSizeBytes = null;
if (arguments.size() > 1) {
RexNode maxBytes = Expressions.fromFieldAccess(
rowSignature,
project,
aggregateCall.getArgList().get(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 = arguments.get(0);
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;
@ -119,20 +112,8 @@ public class ArraySqlAggregator implements SqlAggregator
initialvalue = "[]";
elementType = ColumnType.STRING;
} else {
initialvalue = ExpressionType.fromColumnTypeStrict(druidType).asTypeString() + "[]";
elementType = (ColumnType) druidType.getElementType();
// elementType should never be null if druidType.isArray is true
assert elementType != null;
switch (elementType.getType()) {
case LONG:
initialvalue = "<LONG>[]";
break;
case DOUBLE:
initialvalue = "<DOUBLE>[]";
break;
default:
initialvalue = "[]";
break;
}
}
if (arg.isDirectColumnAccess()) {
fieldName = arg.getDirectColumn();
@ -150,6 +131,8 @@ public class ArraySqlAggregator implements SqlAggregator
initialvalue,
null,
true,
true,
false,
StringUtils.format("array_set_add(\"__acc\", \"%s\")", fieldName),
StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name),
null,
@ -167,6 +150,8 @@ public class ArraySqlAggregator implements SqlAggregator
initialvalue,
null,
true,
true,
false,
StringUtils.format("array_append(\"__acc\", \"%s\")", fieldName),
StringUtils.format("array_concat(\"__acc\", \"%s\")", name),
null,
@ -184,16 +169,12 @@ public class ArraySqlAggregator implements SqlAggregator
public RelDataType inferReturnType(SqlOperatorBinding sqlOperatorBinding)
{
RelDataType type = sqlOperatorBinding.getOperandType(0);
if (SqlTypeUtil.isArray(type)) {
throw new UnsupportedSQLQueryException("Cannot use ARRAY_AGG on array inputs %s", type);
}
if (type instanceof RowSignatures.ComplexSqlType) {
throw new UnsupportedSQLQueryException("Cannot use ARRAY_AGG on complex inputs %s", type);
}
return Calcites.createSqlArrayTypeWithNullability(
sqlOperatorBinding.getTypeFactory(),
type.getSqlTypeName(),
true
return sqlOperatorBinding.getTypeFactory().createArrayType(
type,
-1
);
}
}

2
sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/BitwiseSqlAggregator.java
View File

@ -162,6 +162,8 @@ public class BitwiseSqlAggregator implements SqlAggregator
"0",
null,
null,
false,
false,
StringUtils.format("%s(\"__acc\", \"%s\")", op.getDruidFunction(), fieldName),
null,
null,

4
sql/src/main/java/org/apache/druid/sql/calcite/aggregation/builtin/StringSqlAggregator.java
View File

@ -152,6 +152,8 @@ public class StringSqlAggregator implements SqlAggregator
initialvalue,
null,
true,
false,
false,
StringUtils.format("array_set_add(\"__acc\", \"%s\")", fieldName),
StringUtils.format("array_set_add_all(\"__acc\", \"%s\")", name),
null,
@ -173,6 +175,8 @@ public class StringSqlAggregator implements SqlAggregator
initialvalue,
null,
true,
false,
false,
StringUtils.format("array_append(\"__acc\", \"%s\")", fieldName),
StringUtils.format("array_concat(\"__acc\", \"%s\")", name),
null,

4
sql/src/main/java/org/apache/druid/sql/calcite/expression/builtin/MultiValueStringToArrayOperatorConversion.java
View File

@ -43,9 +43,7 @@ public class MultiValueStringToArrayOperatorConversion implements SqlOperatorCon
{
private static final SqlFunction SQL_FUNCTION = OperatorConversions
.operatorBuilder("MV_TO_ARRAY")
.operandTypeChecker(
OperandTypes.family(SqlTypeFamily.STRING)
)
.operandTypeChecker(OperandTypes.family(SqlTypeFamily.STRING))
.functionCategory(SqlFunctionCategory.STRING)
.returnTypeNullableArray(SqlTypeName.VARCHAR)
.build();

12
sql/src/main/java/org/apache/druid/sql/calcite/planner/Calcites.java
View File

@ -173,14 +173,11 @@ public class Calcites
}
return ColumnType.UNKNOWN_COMPLEX;
} else if (sqlTypeName == SqlTypeName.ARRAY) {
SqlTypeName componentType = type.getComponentType().getSqlTypeName();
if (isDoubleType(componentType)) {
return ColumnType.DOUBLE_ARRAY;
ColumnType elementType = getValueTypeForRelDataTypeFull(type.getComponentType());
if (elementType != null) {
return ColumnType.ofArray(elementType);
}
if (isLongType(componentType)) {
return ColumnType.LONG_ARRAY;
}
return ColumnType.STRING_ARRAY;
return null;
} else {
return null;
}
@ -265,6 +262,7 @@ public class Calcites
)
{
final RelDataType dataType = typeFactory.createArrayType(
createSqlTypeWithNullability(typeFactory, elementTypeName, nullable),
-1

2
sql/src/main/java/org/apache/druid/sql/calcite/planner/DruidOperatorTable.java
View File

@ -33,6 +33,7 @@ import org.apache.calcite.sql.validate.SqlNameMatcher;
import org.apache.druid.java.util.common.ISE;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.sql.calcite.aggregation.SqlAggregator;
import org.apache.druid.sql.calcite.aggregation.builtin.ArrayConcatSqlAggregator;
import org.apache.druid.sql.calcite.aggregation.builtin.ArraySqlAggregator;
import org.apache.druid.sql.calcite.aggregation.builtin.AvgSqlAggregator;
import org.apache.druid.sql.calcite.aggregation.builtin.BitwiseSqlAggregator;
@ -143,6 +144,7 @@ public class DruidOperatorTable implements SqlOperatorTable
.add(new SumZeroSqlAggregator())
.add(new GroupingSqlAggregator())
.add(new ArraySqlAggregator())
.add(new ArrayConcatSqlAggregator())
.add(new StringSqlAggregator())
.add(new BitwiseSqlAggregator(BitwiseSqlAggregator.Op.AND))
.add(new BitwiseSqlAggregator(BitwiseSqlAggregator.Op.OR))

45
sql/src/main/java/org/apache/druid/sql/calcite/run/NativeQueryMaker.java
View File

@ -332,21 +332,8 @@ public class NativeQueryMaker implements QueryMaker
// the protobuf jdbc handler prefers lists (it actually can't handle java arrays as sql arrays, only java lists)
// the json handler could handle this just fine, but it handles lists as sql arrays as well so just convert
// here if needed
if (value instanceof List) {
coercedValue = value;
} else if (value instanceof String[]) {
coercedValue = Arrays.asList((String[]) value);
} else if (value instanceof Long[]) {
coercedValue = Arrays.asList((Long[]) value);
} else if (value instanceof Double[]) {
coercedValue = Arrays.asList((Double[]) value);
} else if (value instanceof Object[]) {
coercedValue = Arrays.asList((Object[]) value);
} else if (value instanceof ComparableStringArray) {
coercedValue = Arrays.asList(((ComparableStringArray) value).getDelegate());
} else if (value instanceof ComparableList) {
coercedValue = ((ComparableList) value).getDelegate();
} else {
coercedValue = maybeCoerceArrayToList(value, true);
if (coercedValue == null) {
throw new ISE("Cannot coerce[%s] to %s", value.getClass().getName(), sqlType);
}
}
@ -357,6 +344,34 @@ public class NativeQueryMaker implements QueryMaker
return coercedValue;
}
private static Object maybeCoerceArrayToList(Object value, boolean mustCoerce)
{
if (value instanceof List) {
return value;
} else if (value instanceof String[]) {
return Arrays.asList((String[]) value);
} else if (value instanceof Long[]) {
return Arrays.asList((Long[]) value);
} else if (value instanceof Double[]) {
return Arrays.asList((Double[]) value);
} else if (value instanceof Object[]) {
Object[] array = (Object[]) value;
ArrayList<Object> lst = new ArrayList<>(array.length);
for (Object o : array) {
lst.add(maybeCoerceArrayToList(o, false));
}
return lst;
} else if (value instanceof ComparableStringArray) {
return Arrays.asList(((ComparableStringArray) value).getDelegate());
} else if (value instanceof ComparableList) {
return ((ComparableList) value).getDelegate();
} else if (mustCoerce) {
return null;
}
return value;
}
private static DateTime coerceDateTime(Object value, SqlTypeName sqlType)
{
final DateTime dateTime;

186
sql/src/test/java/org/apache/druid/sql/calcite/BaseCalciteQueryTest.java
View File

@ -624,6 +624,7 @@ public class BaseCalciteQueryTest extends CalciteTestBase
}
}
public void testQuery(
final String sql,
final List<Query> expectedQueries,
@ -696,6 +697,25 @@ public class BaseCalciteQueryTest extends CalciteTestBase
);
}
public void testQuery(
final String sql,
final Map<String, Object> context,
final List<Query> expectedQueries,
final ResultsVerifier expectedResultsVerifier
) throws Exception
{
testQuery(
PLANNER_CONFIG_DEFAULT,
context,
DEFAULT_PARAMETERS,
sql,
CalciteTests.REGULAR_USER_AUTH_RESULT,
expectedQueries,
expectedResultsVerifier,
null
);
}
public void testQuery(
final PlannerConfig plannerConfig,
final Map<String, Object> queryContext,
@ -712,31 +732,26 @@ public class BaseCalciteQueryTest extends CalciteTestBase
verifyResults(sql, expectedQueries, expectedResults, plannerResults);
}
/**
* Override not just the outer query context, but also the contexts of all subqueries.
*/
public <T> Query<T> recursivelyOverrideContext(final Query<T> query, final Map<String, Object> context)
public void testQuery(
final PlannerConfig plannerConfig,
final Map<String, Object> queryContext,
final List<SqlParameter> parameters,
final String sql,
final AuthenticationResult authenticationResult,
final List<Query> expectedQueries,
final List<Object[]> expectedResults
) throws Exception
{
return query.withDataSource(recursivelyOverrideContext(query.getDataSource(), context))
.withOverriddenContext(context);
}
/**
* Override the contexts of all subqueries of a particular datasource.
*/
private DataSource recursivelyOverrideContext(final DataSource dataSource, final Map<String, Object> context)
{
if (dataSource instanceof QueryDataSource) {
final Query subquery = ((QueryDataSource) dataSource).getQuery();
return new QueryDataSource(recursivelyOverrideContext(subquery, context));
} else {
return dataSource.withChildren(
dataSource.getChildren()
.stream()
.map(ds -> recursivelyOverrideContext(ds, context))
.collect(Collectors.toList())
);
}
testQuery(
plannerConfig,
queryContext,
parameters,
sql,
authenticationResult,
expectedQueries,
new DefaultResultsVerifier(expectedResults),
null
);
}
public void testQuery(
@ -746,7 +761,8 @@ public class BaseCalciteQueryTest extends CalciteTestBase
final String sql,
final AuthenticationResult authenticationResult,
final List<Query> expectedQueries,
final List<Object[]> expectedResults
final ResultsVerifier expectedResultsVerifier,
@Nullable final Consumer<ExpectedException> expectedExceptionInitializer
) throws Exception
{
log.info("SQL: %s", sql);
@ -778,10 +794,12 @@ public class BaseCalciteQueryTest extends CalciteTestBase
if (cannotVectorize && "force".equals(vectorize)) {
expectedException.expect(RuntimeException.class);
expectedException.expectMessage("Cannot vectorize");
} else if (expectedExceptionInitializer != null) {
expectedExceptionInitializer.accept(expectedException);
}
final List<Object[]> plannerResults = getResults(plannerConfig, theQueryContext, parameters, sql, authenticationResult);
verifyResults(sql, theQueries, expectedResults, plannerResults);
verifyResults(sql, theQueries, plannerResults, expectedResultsVerifier);
}
}
@ -835,13 +853,22 @@ public class BaseCalciteQueryTest extends CalciteTestBase
final List<Object[]> expectedResults,
final List<Object[]> results
)
{
verifyResults(sql, expectedQueries, results, new DefaultResultsVerifier(expectedResults));
}
public void verifyResults(
final String sql,
final List<Query> expectedQueries,
final List<Object[]> results,
final ResultsVerifier expectedResultsVerifier
)
{
for (int i = 0; i < results.size(); i++) {
log.info("row #%d: %s", i, Arrays.toString(results.get(i)));
}
Assert.assertEquals(StringUtils.format("result count: %s", sql), expectedResults.size(), results.size());
assertResultsEquals(sql, expectedResults, results);
expectedResultsVerifier.verify(sql, results);
verifyQueries(sql, expectedQueries);
}
@ -907,67 +934,18 @@ public class BaseCalciteQueryTest extends CalciteTestBase
final Consumer<ExpectedException> expectedExceptionInitializer
) throws Exception
{
testQueryThrows(
testQuery(
PLANNER_CONFIG_DEFAULT,
queryContext,
DEFAULT_PARAMETERS,
sql,
CalciteTests.REGULAR_USER_AUTH_RESULT,
expectedQueries,
(query, results) -> {},
expectedExceptionInitializer
);
}
public void testQueryThrows(
final PlannerConfig plannerConfig,
final Map<String, Object> queryContext,
final List<SqlParameter> parameters,
final String sql,
final AuthenticationResult authenticationResult,
final List<Query> expectedQueries,
final Consumer<ExpectedException> expectedExceptionInitializer
) throws Exception
{
log.info("SQL: %s", sql);
final List<String> vectorizeValues = new ArrayList<>();
vectorizeValues.add("false");
if (!skipVectorize) {
vectorizeValues.add("force");
}
for (final String vectorize : vectorizeValues) {
queryLogHook.clearRecordedQueries();
final Map<String, Object> theQueryContext = new HashMap<>(queryContext);
theQueryContext.put(QueryContexts.VECTORIZE_KEY, vectorize);
theQueryContext.put(QueryContexts.VECTORIZE_VIRTUAL_COLUMNS_KEY, vectorize);
if (!"false".equals(vectorize)) {
theQueryContext.put(QueryContexts.VECTOR_SIZE_KEY, 2); // Small vector size to ensure we use more than one.
}
final List<Query> theQueries = new ArrayList<>();
for (Query query : expectedQueries) {
theQueries.add(recursivelyOverrideContext(query, theQueryContext));
}
if (cannotVectorize && "force".equals(vectorize)) {
expectedException.expect(RuntimeException.class);
expectedException.expectMessage("Cannot vectorize");
} else {
expectedExceptionInitializer.accept(expectedException);
}
// this should validate expectedException
getResults(plannerConfig, theQueryContext, parameters, sql, authenticationResult);
verifyQueries(sql, theQueries);
}
}
public Set<ResourceAction> analyzeResources(
PlannerConfig plannerConfig,
String sql,
@ -1084,6 +1062,33 @@ public class BaseCalciteQueryTest extends CalciteTestBase
);
}
/**
* Override not just the outer query context, but also the contexts of all subqueries.
*/
public static <T> Query<T> recursivelyOverrideContext(final Query<T> query, final Map<String, Object> context)
{
return query.withDataSource(recursivelyOverrideContext(query.getDataSource(), context))
.withOverriddenContext(context);
}
/**
* Override the contexts of all subqueries of a particular datasource.
*/
private static DataSource recursivelyOverrideContext(final DataSource dataSource, final Map<String, Object> context)
{
if (dataSource instanceof QueryDataSource) {
final Query subquery = ((QueryDataSource) dataSource).getQuery();
return new QueryDataSource(recursivelyOverrideContext(subquery, context));
} else {
return dataSource.withChildren(
dataSource.getChildren()
.stream()
.map(ds -> recursivelyOverrideContext(ds, context))
.collect(Collectors.toList())
);
}
}
/**
* This is a provider of query contexts that should be used by join tests.
* It tests various configs that can be passed to join queries. All the configs provided by this provider should
@ -1184,4 +1189,27 @@ public class BaseCalciteQueryTest extends CalciteTestBase
output.put(GroupByQuery.CTX_TIMESTAMP_RESULT_FIELD_INDEX, timestampResultFieldIndex);
return output;
}
@FunctionalInterface
public interface ResultsVerifier
{
void verify(String sql, List<Object[]> results);
}
public class DefaultResultsVerifier implements ResultsVerifier
{
protected final List<Object[]> expectedResults;
public DefaultResultsVerifier(List<Object[]> expectedResults)
{
this.expectedResults = expectedResults;
}
@Override
public void verify(String sql, List<Object[]> results)
{
Assert.assertEquals(StringUtils.format("result count: %s", sql), expectedResults.size(), results.size());
assertResultsEquals(sql, expectedResults, results);
}
}
}

455
sql/src/test/java/org/apache/druid/sql/calcite/CalciteArraysQueryTest.java
View File

@ -24,6 +24,7 @@ import com.google.common.collect.ImmutableSet;
import org.apache.druid.common.config.NullHandling;
import org.apache.druid.java.util.common.HumanReadableBytes;
import org.apache.druid.java.util.common.IAE;
import org.apache.druid.java.util.common.StringUtils;
import org.apache.druid.java.util.common.granularity.Granularities;
import org.apache.druid.math.expr.ExprMacroTable;
import org.apache.druid.math.expr.ExpressionProcessing;
@ -54,6 +55,7 @@ import org.apache.druid.segment.column.ColumnType;
import org.apache.druid.segment.join.JoinType;
import org.apache.druid.sql.calcite.filtration.Filtration;
import org.apache.druid.sql.calcite.util.CalciteTests;
import org.junit.Assert;
import org.junit.Test;
import java.util.Arrays;
@ -163,33 +165,37 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
.build();
ExpressionProcessing.initializeForTests(true);
// if nested arrays are allowed, dim3 is a multi-valued string column, so the automatic translation will turn this
// expression into
//
// `map((dim3) -> array(concat(dim3,'word'),'up'), dim3)`
//
// this works, but we still translate the output into a string since that is the current output type
// in some future this might not auto-convert to a string type (when we support grouping on arrays maybe?)
try {
ExpressionProcessing.initializeForTests(true);
// if nested arrays are allowed, dim3 is a multi-valued string column, so the automatic translation will turn this
// expression into
//
// `map((dim3) -> array(concat(dim3,'word'),'up'), dim3)`
//
// this works, but we still translate the output into a string since that is the current output type
// in some future this might not auto-convert to a string type (when we support grouping on arrays maybe?)
testQuery(
sql,
ImmutableList.of(scanQuery),
ImmutableList.of(
new Object[]{"[[\"aword\",\"up\"],[\"bword\",\"up\"]]", ""},
new Object[]{"[[\"bword\",\"up\"],[\"cword\",\"up\"]]", "10.1"},
new Object[]{"[[\"dword\",\"up\"]]", "2"},
new Object[]{"[[\"word\",\"up\"]]", "1"},
useDefault ? new Object[]{"[[\"word\",\"up\"]]", "def"} : new Object[]{"[[null,\"up\"]]", "def"}
)
);
ExpressionProcessing.initializeForTests(null);
testQuery(
sql,
ImmutableList.of(scanQuery),
ImmutableList.of(
new Object[]{"[[\"aword\",\"up\"],[\"bword\",\"up\"]]", ""},
new Object[]{"[[\"bword\",\"up\"],[\"cword\",\"up\"]]", "10.1"},
new Object[]{"[[\"dword\",\"up\"]]", "2"},
new Object[]{"[[\"word\",\"up\"]]", "1"},
useDefault ? new Object[]{"[[\"word\",\"up\"]]", "def"} : new Object[]{"[[null,\"up\"]]", "def"}
)
);
}
finally {
ExpressionProcessing.initializeForTests(null);
}
// if nested arrays are not enabled, this doesn't work
expectedException.expect(IAE.class);
expectedException.expectMessage(
"Cannot create a nested array type [ARRAY<ARRAY<STRING>>], 'druid.expressions.allowNestedArrays' must be set to true");
"Cannot create a nested array type [ARRAY<ARRAY<STRING>>], 'druid.expressions.allowNestedArrays' must be set to true"
);
testQuery(
sql,
ImmutableList.of(scanQuery),
@ -1044,11 +1050,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
.setVirtualColumns(expressionVirtualColumn(
"v0",
"array(\"f1\")",
ColumnType.DOUBLE_ARRAY
ColumnType.ofArray(ColumnType.FLOAT)
))
.setDimensions(
dimensions(
new DefaultDimensionSpec("v0", "_d0", ColumnType.DOUBLE_ARRAY)
new DefaultDimensionSpec("v0", "_d0", ColumnType.ofArray(ColumnType.FLOAT))
)
)
.setAggregatorSpecs(aggregators(new LongSumAggregatorFactory("a0", "cnt")))
@ -1385,9 +1391,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_append(\"__acc\", \"dim1\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -1399,9 +1407,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a1",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -1414,9 +1424,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a2",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a2\")",
null,
@ -1433,10 +1445,10 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
),
ImmutableList.of(
useDefault
? new Object[]{"[\"10.1\",\"2\",\"1\",\"def\",\"abc\"]", "[\"1\",\"2\",\"abc\",\"def\",\"10.1\"]", null}
? new Object[]{"[\"10.1\",\"2\",\"1\",\"def\",\"abc\"]", "[\"1\",\"10.1\",\"2\",\"abc\",\"def\"]", null}
: new Object[]{
"[\"\",\"10.1\",\"2\",\"1\",\"def\",\"abc\"]",
"[\"\",\"1\",\"2\",\"abc\",\"def\",\"10.1\"]",
"[\"\",\"1\",\"10.1\",\"2\",\"abc\",\"def\"]",
null
}
)
@ -1460,9 +1472,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("dim3"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_append(\"__acc\", \"dim3\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -1474,9 +1488,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a1",
ImmutableSet.of("dim3"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim3\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -1492,7 +1508,7 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
ImmutableList.of(
useDefault
? new Object[]{"[\"a\",\"b\",\"b\",\"c\",\"d\",null,null,null]", "[null,\"a\",\"b\",\"c\",\"d\"]"}
: new Object[]{"[\"a\",\"b\",\"b\",\"c\",\"d\",\"\",null,null]", "[\"\",null,\"a\",\"b\",\"c\",\"d\"]"}
: new Object[]{"[\"a\",\"b\",\"b\",\"c\",\"d\",\"\",null,null]", "[null,\"\",\"a\",\"b\",\"c\",\"d\"]"}
)
);
}
@ -1514,9 +1530,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("l1"),
"__acc",
"<LONG>[]",
"<LONG>[]",
"ARRAY<LONG>[]",
"ARRAY<LONG>[]",
true,
true,
false,
"array_append(\"__acc\", \"l1\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -1528,9 +1546,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a1",
ImmutableSet.of("l1"),
"__acc",
"<LONG>[]",
"<LONG>[]",
"ARRAY<LONG>[]",
"ARRAY<LONG>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"l1\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -1542,9 +1562,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a2",
ImmutableSet.of("d1"),
"__acc",
"<DOUBLE>[]",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
true,
true,
false,
"array_append(\"__acc\", \"d1\")",
"array_concat(\"__acc\", \"a2\")",
null,
@ -1556,9 +1578,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a3",
ImmutableSet.of("d1"),
"__acc",
"<DOUBLE>[]",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"d1\")",
"array_set_add_all(\"__acc\", \"a3\")",
null,
@ -1570,9 +1594,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a4",
ImmutableSet.of("f1"),
"__acc",
"<DOUBLE>[]",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
true,
true,
false,
"array_append(\"__acc\", \"f1\")",
"array_concat(\"__acc\", \"a4\")",
null,
@ -1584,9 +1610,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a5",
ImmutableSet.of("f1"),
"__acc",
"<DOUBLE>[]",
"<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
"ARRAY<DOUBLE>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"f1\")",
"array_set_add_all(\"__acc\", \"a5\")",
null,
@ -1605,22 +1633,246 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"[7,325323,0,0,0,0]",
"[0,7,325323]",
"[1.0,1.7,0.0,0.0,0.0,0.0]",
"[1.0,0.0,1.7]",
"[0.0,1.0,1.7]",
"[1.0,0.10000000149011612,0.0,0.0,0.0,0.0]",
"[1.0,0.10000000149011612,0.0]"
"[0.0,0.10000000149011612,1.0]"
}
: new Object[]{
"[7,325323,0,null,null,null]",
"[0,null,7,325323]",
"[null,0,7,325323]",
"[1.0,1.7,0.0,null,null,null]",
"[1.0,0.0,null,1.7]",
"[null,0.0,1.0,1.7]",
"[1.0,0.10000000149011612,0.0,null,null,null]",
"[1.0,0.10000000149011612,0.0,null]"
"[null,0.0,0.10000000149011612,1.0]"
}
)
);
}
@Test
public void testArrayAggArrays() throws Exception
{
try {
ExpressionProcessing.initializeForTests(true);
cannotVectorize();
testQuery(
"SELECT ARRAY_AGG(ARRAY[l1, l2]), ARRAY_AGG(DISTINCT ARRAY[l1, l2]) FROM numfoo",
QUERY_CONTEXT_NO_STRINGIFY_ARRAY,
ImmutableList.of(
Druids.newTimeseriesQueryBuilder()
.dataSource(CalciteTests.DATASOURCE3)
.intervals(querySegmentSpec(Filtration.eternity()))
.granularity(Granularities.ALL)
.virtualColumns(
expressionVirtualColumn("v0", "array(\"l1\",\"l2\")", ColumnType.LONG_ARRAY)
)
.aggregators(
aggregators(
new ExpressionLambdaAggregatorFactory(
"a0",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<ARRAY<LONG>>[]",
"ARRAY<ARRAY<LONG>>[]",
true,
true,
false,
"array_append(\"__acc\", \"v0\")",
"array_concat(\"__acc\", \"a0\")",
null,
null,
ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES,
TestExprMacroTable.INSTANCE
),
new ExpressionLambdaAggregatorFactory(
"a1",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<ARRAY<LONG>>[]",
"ARRAY<ARRAY<LONG>>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"v0\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
null,
ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES,
TestExprMacroTable.INSTANCE
)
)
)
.context(QUERY_CONTEXT_NO_STRINGIFY_ARRAY)
.build()
),
(sql, results) -> {
// ordering is not stable in array_agg and array_concat_agg
List<Object[]> expected = ImmutableList.of(
useDefault ?
new Object[]{
Arrays.asList(
Arrays.asList(7L, 0L),
Arrays.asList(325323L, 325323L),
Arrays.asList(0L, 0L),
Arrays.asList(0L, 0L),
Arrays.asList(0L, 0L),
Arrays.asList(0L, 0L)
),
Arrays.asList(
Arrays.asList(0L, 0L),
Arrays.asList(7L, 0L),
Arrays.asList(325323L, 325323L)
)
}
:
new Object[]{
Arrays.asList(
Arrays.asList(7L, null),
Arrays.asList(325323L, 325323L),
Arrays.asList(0L, 0L),
Arrays.asList(null, null),
Arrays.asList(null, null),
Arrays.asList(null, null)
),
Arrays.asList(
Arrays.asList(null, null),
Arrays.asList(0L, 0L),
Arrays.asList(7L, null),
Arrays.asList(325323L, 325323L)
)
}
);
assertResultsDeepEquals(sql, expected, results);
}
);
}
finally {
ExpressionProcessing.initializeForTests(null);
}
}
@Test
public void testArrayAggArraysNoNest() throws Exception
{
cannotVectorize();
testQueryThrows(
"SELECT ARRAY_AGG(ARRAY[l1, l2]), ARRAY_AGG(DISTINCT ARRAY[l1, l2]) FROM numfoo",
QUERY_CONTEXT_NO_STRINGIFY_ARRAY,
ImmutableList.of(
Druids.newTimeseriesQueryBuilder()
.dataSource(CalciteTests.DATASOURCE3)
.intervals(querySegmentSpec(Filtration.eternity()))
.granularity(Granularities.ALL)
.virtualColumns(
expressionVirtualColumn("v0", "array(\"l1\",\"l2\")", ColumnType.LONG_ARRAY)
)
.aggregators(
aggregators(
new ExpressionLambdaAggregatorFactory(
"a0",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<ARRAY<LONG>>[]",
"ARRAY<ARRAY<LONG>>[]",
true,
true,
false,
"array_append(\"__acc\", \"v0\")",
"array_concat(\"__acc\", \"a0\")",
null,
null,
ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES,
TestExprMacroTable.INSTANCE
),
new ExpressionLambdaAggregatorFactory(
"a1",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<ARRAY<LONG>>[]",
"ARRAY<ARRAY<LONG>>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"v0\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
null,
ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES,
TestExprMacroTable.INSTANCE
)
)
)
.context(QUERY_CONTEXT_NO_STRINGIFY_ARRAY)
.build()
),
expected -> {
expected.expect(IAE.class);
expected.expectMessage("Cannot create a nested array type [ARRAY<ARRAY<LONG>>], 'druid.expressions.allowNestedArrays' must be set to true");
}
);
}
@Test
public void testArrayConcatAggArrays() throws Exception
{
cannotVectorize();
testQuery(
"SELECT ARRAY_CONCAT_AGG(ARRAY[l1, l2]), ARRAY_CONCAT_AGG(DISTINCT ARRAY[l1, l2]) FROM numfoo",
ImmutableList.of(
Druids.newTimeseriesQueryBuilder()
.dataSource(CalciteTests.DATASOURCE3)
.intervals(querySegmentSpec(Filtration.eternity()))
.granularity(Granularities.ALL)
.virtualColumns(
expressionVirtualColumn("v0", "array(\"l1\",\"l2\")", ColumnType.LONG_ARRAY)
)
.aggregators(
aggregators(
new ExpressionLambdaAggregatorFactory(
"a0",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<LONG>[]",
"ARRAY<LONG>[]",
true,
false,
false,
"array_concat(\"__acc\", \"v0\")",
"array_concat(\"__acc\", \"a0\")",
null,
null,
ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES,
TestExprMacroTable.INSTANCE
),
new ExpressionLambdaAggregatorFactory(
"a1",
ImmutableSet.of("v0"),
"__acc",
"ARRAY<LONG>[]",
"ARRAY<LONG>[]",
true,
false,
false,
"array_set_add_all(\"__acc\", \"v0\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
null,
ExpressionLambdaAggregatorFactory.DEFAULT_MAX_SIZE_BYTES,
TestExprMacroTable.INSTANCE
)
)
)
.context(QUERY_CONTEXT_DEFAULT)
.build()
),
ImmutableList.of(
useDefault
? new Object[]{"[7,0,325323,325323,0,0,0,0,0,0,0,0]", "[0,7,325323]"}
: new Object[]{"[7,null,325323,325323,0,0,null,null,null,null,null,null]", "[null,0,7,325323]"}
)
);
}
@Test
public void testArrayAggToString() throws Exception
{
@ -1639,9 +1891,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a0\")",
null,
@ -1656,7 +1910,7 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
.build()
),
ImmutableList.of(
useDefault ? new Object[]{"1,2,abc,def,10.1"} : new Object[]{",1,2,abc,def,10.1"}
useDefault ? new Object[]{"1,10.1,2,abc,def"} : new Object[]{",1,10.1,2,abc,def"}
)
);
}
@ -1681,9 +1935,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("v0"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"v0\")",
"array_set_add_all(\"__acc\", \"a0\")",
null,
@ -1698,7 +1954,7 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
.build()
),
ImmutableList.of(
useDefault ? new Object[]{"1a,a,2,abc,10.1,defabc"} : new Object[]{"null,1a,a,2,defabc"}
useDefault ? new Object[]{"10.1,1a,2,a,abc,defabc"} : new Object[]{"null,1a,2,a,defabc"}
)
);
}
@ -1720,9 +1976,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("l1"),
"__acc",
"<LONG>[]",
"<LONG>[]",
"ARRAY<LONG>[]",
"ARRAY<LONG>[]",
true,
true,
false,
"array_append(\"__acc\", \"l1\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -1734,9 +1992,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a1",
ImmutableSet.of("l1"),
"__acc",
"<LONG>[]",
"<LONG>[]",
"ARRAY<LONG>[]",
"ARRAY<LONG>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"l1\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -1752,7 +2012,7 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
ImmutableList.of(
useDefault
? new Object[]{"[7,325323,0,0,0,0]", "[0,7,325323]"}
: new Object[]{"[7,325323,0,null,null,null]", "[0,null,7,325323]"}
: new Object[]{"[7,325323,0,null,null,null]", "[null,0,7,325323]"}
)
);
}
@ -1764,18 +2024,18 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
List<Object[]> expectedResults;
if (useDefault) {
expectedResults = ImmutableList.of(
new Object[]{"a", "[\"2\",\"10.1\"]", "2,10.1"},
new Object[]{"a", "[\"2\",\"10.1\"]", "2,10.1"},
new Object[]{"a", "[\"2\",\"10.1\"]", "2,10.1"},
new Object[]{"a", "[\"10.1\",\"2\"]", "10.1,2"},
new Object[]{"a", "[\"10.1\",\"2\"]", "10.1,2"},
new Object[]{"a", "[\"10.1\",\"2\"]", "10.1,2"},
new Object[]{"b", "[\"1\",\"abc\",\"def\"]", "1,abc,def"},
new Object[]{"b", "[\"1\",\"abc\",\"def\"]", "1,abc,def"},
new Object[]{"b", "[\"1\",\"abc\",\"def\"]", "1,abc,def"}
);
} else {
expectedResults = ImmutableList.of(
new Object[]{"a", "[\"\",\"2\",\"10.1\"]", ",2,10.1"},
new Object[]{"a", "[\"\",\"2\",\"10.1\"]", ",2,10.1"},
new Object[]{"a", "[\"\",\"2\",\"10.1\"]", ",2,10.1"},
new Object[]{"a", "[\"\",\"10.1\",\"2\"]", ",10.1,2"},
new Object[]{"a", "[\"\",\"10.1\",\"2\"]", ",10.1,2"},
new Object[]{"a", "[\"\",\"10.1\",\"2\"]", ",10.1,2"},
new Object[]{"b", "[\"1\",\"abc\",\"def\"]", "1,abc,def"},
new Object[]{"b", "[\"1\",\"abc\",\"def\"]", "1,abc,def"},
new Object[]{"b", "[\"1\",\"abc\",\"def\"]", "1,abc,def"}
@ -1801,9 +2061,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a0\")",
null,
@ -1866,9 +2128,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a0\")",
null,
@ -1890,12 +2154,12 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
),
useDefault ?
ImmutableList.of(
new Object[]{"", ImmutableList.of("2", "abc", "10.1"), 1L},
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("abc", "10.1"), 1L},
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}
@ -1945,9 +2209,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a0\")",
null,
@ -2022,9 +2288,11 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
"a0",
ImmutableSet.of("dim1"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a0\")",
null,
@ -2060,5 +2328,30 @@ public class CalciteArraysQueryTest extends BaseCalciteQueryTest
),
expectedResults
);
}
public static void assertResultsDeepEquals(String sql, List<Object[]> expected, List<Object[]> results)
{
for (int row = 0; row < results.size(); row++) {
for (int col = 0; col < results.get(row).length; col++) {
final String rowString = StringUtils.format("result #%d: %s", row + 1, sql);
assertDeepEquals(rowString + " - column: " + col + ":", expected.get(row)[col], results.get(row)[col]);
}
}
}
public static void assertDeepEquals(String path, Object expected, Object actual)
{
if (expected instanceof List && actual instanceof List) {
List expectedList = (List) expected;
List actualList = (List) actual;
Assert.assertEquals(path + " arrays length mismatch", expectedList.size(), actualList.size());
for (int i = 0; i < expectedList.size(); i++) {
assertDeepEquals(path + "[" + i + "]", expectedList.get(i), actualList.get(i));
}
} else {
Assert.assertEquals(path, expected, actual);
}
}
}

90
sql/src/test/java/org/apache/druid/sql/calcite/CalciteQueryTest.java
View File

@ -9205,9 +9205,11 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"a6",
ImmutableSet.of("dim3"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim3\")",
"array_set_add_all(\"__acc\", \"a6\")",
null,
@ -9223,6 +9225,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"dim3\")",
"array_set_add_all(\"__acc\", \"a7\")",
null,
@ -9240,6 +9244,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseAnd(\"__acc\", \"l1\")",
"bitwiseAnd(\"__acc\", \"a8\")",
null,
@ -9257,6 +9263,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseOr(\"__acc\", \"l1\")",
"bitwiseOr(\"__acc\", \"a9\")",
null,
@ -9274,6 +9282,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseXor(\"__acc\", \"l1\")",
"bitwiseXor(\"__acc\", \"a10\")",
null,
@ -9502,9 +9512,11 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"a6",
ImmutableSet.of("dim3"),
"__acc",
"[]",
"[]",
"ARRAY<STRING>[]",
"ARRAY<STRING>[]",
true,
true,
false,
"array_set_add(\"__acc\", \"dim3\")",
"array_set_add_all(\"__acc\", \"a6\")",
null,
@ -9522,6 +9534,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"dim3\")",
"array_set_add_all(\"__acc\", \"a7\")",
null,
@ -9542,6 +9556,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseAnd(\"__acc\", \"l1\")",
"bitwiseAnd(\"__acc\", \"a8\")",
null,
@ -9559,6 +9575,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseOr(\"__acc\", \"l1\")",
"bitwiseOr(\"__acc\", \"a9\")",
null,
@ -9576,6 +9594,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseXor(\"__acc\", \"l1\")",
"bitwiseXor(\"__acc\", \"a10\")",
null,
@ -12852,6 +12872,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseAnd(\"__acc\", \"l1\")",
"bitwiseAnd(\"__acc\", \"a0\")",
null,
@ -12869,6 +12891,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseOr(\"__acc\", \"l1\")",
"bitwiseOr(\"__acc\", \"a1\")",
null,
@ -12886,6 +12910,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseXor(\"__acc\", \"l1\")",
"bitwiseXor(\"__acc\", \"a2\")",
null,
@ -12935,6 +12961,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseAnd(\"__acc\", \"l1\")",
"bitwiseAnd(\"__acc\", \"a0\")",
null,
@ -12952,6 +12980,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseOr(\"__acc\", \"l1\")",
"bitwiseOr(\"__acc\", \"a1\")",
null,
@ -12969,6 +12999,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"0",
"0",
NullHandling.sqlCompatible(),
false,
false,
"bitwiseXor(\"__acc\", \"l1\")",
"bitwiseXor(\"__acc\", \"a2\")",
null,
@ -13033,6 +13065,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_append(\"__acc\", \"dim1\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -13050,6 +13084,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -13067,6 +13103,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"dim1\")",
"array_set_add_all(\"__acc\", \"a2\")",
null,
@ -13086,8 +13124,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
),
ImmutableList.of(
useDefault
? new Object[]{"10.1,2,1,def,abc", "1,2,abc,def,10.1", ""}
: new Object[]{",10.1,2,1,def,abc", ",1,2,abc,def,10.1", null}
? new Object[]{"10.1,2,1,def,abc", "1,10.1,2,abc,def", ""}
: new Object[]{",10.1,2,1,def,abc", ",1,10.1,2,abc,def", null}
)
);
}
@ -13113,6 +13151,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_append(\"__acc\", \"dim3\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -13130,6 +13170,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"dim3\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -13173,6 +13215,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_append(\"__acc\", \"l1\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -13190,6 +13234,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"l1\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -13207,6 +13253,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_append(\"__acc\", \"d1\")",
"array_concat(\"__acc\", \"a2\")",
null,
@ -13224,6 +13272,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"d1\")",
"array_set_add_all(\"__acc\", \"a3\")",
null,
@ -13241,6 +13291,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_append(\"__acc\", \"f1\")",
"array_concat(\"__acc\", \"a4\")",
null,
@ -13258,6 +13310,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"f1\")",
"array_set_add_all(\"__acc\", \"a5\")",
null,
@ -13276,19 +13330,19 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
useDefault
? new Object[]{
"7,325323,0,0,0,0",
"0,7,325323",
"0,325323,7",
"1.0,1.7,0.0,0.0,0.0,0.0",
"0.0,1.0,1.7",
"1.0,0.10000000149011612,0.0,0.0,0.0,0.0",
"0.10000000149011612,0.0,1.0"
"0.0,0.10000000149011612,1.0"
}
: new Object[]{
"7,325323,0",
"0,7,325323",
"0,325323,7",
"1.0,1.7,0.0",
"0.0,1.0,1.7",
"1.0,0.10000000149011612,0.0",
"0.10000000149011612,0.0,1.0"
"0.0,0.10000000149011612,1.0"
}
)
);
@ -13318,6 +13372,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"v0\")",
"array_set_add_all(\"__acc\", \"a0\")",
null,
@ -13335,6 +13391,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"v0\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -13351,8 +13409,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
),
ImmutableList.of(
useDefault
? new Object[]{"1a,a,2,abc,10.1,defabc", "1a||a||2||abc||10.1||defabc"}
: new Object[]{"1a,a,2,defabc", "1a||a||2||defabc"}
? new Object[]{"10.1,1a,2,a,abc,defabc", "10.1||1a||2||a||abc||defabc"}
: new Object[]{"1a,2,a,defabc", "1a||2||a||defabc"}
)
);
}
@ -13388,6 +13446,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_append(\"__acc\", \"l1\")",
"array_concat(\"__acc\", \"a0\")",
null,
@ -13405,6 +13465,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
"[]",
"[]",
true,
false,
false,
"array_set_add(\"__acc\", \"l1\")",
"array_set_add_all(\"__acc\", \"a1\")",
null,
@ -13421,8 +13483,8 @@ public class CalciteQueryTest extends BaseCalciteQueryTest
),
ImmutableList.of(
useDefault
? new Object[]{"7,325323,0,0,0,0", "0,7,325323"}
: new Object[]{"7,325323,0", "0,7,325323"}
? new Object[]{"7,325323,0,0,0,0", "0,325323,7"}
: new Object[]{"7,325323,0", "0,325323,7"}
)
);
}