You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@pinot.apache.org by ja...@apache.org on 2020/09/29 20:17:22 UTC
[incubator-pinot] branch master updated: Enhance
DistinctCountThetaSketchAggregationFunction (#6004)
This is an automated email from the ASF dual-hosted git repository.
jackie pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/incubator-pinot.git
The following commit(s) were added to refs/heads/master by this push:
new e892cb2 Enhance DistinctCountThetaSketchAggregationFunction (#6004)
e892cb2 is described below
commit e892cb2942f1f41ac7056d977eeca0e5b22897d0
Author: Xiaotian (Jackie) Jiang <17...@users.noreply.github.com>
AuthorDate: Tue Sep 29 13:17:05 2020 -0700
Enhance DistinctCountThetaSketchAggregationFunction (#6004)
## Description
Enhance `DistinctCountThetaSketchAggregationFunction`, and add the following supports:
- Support aggregating on raw values of all data types (both SV and MV)
- Support nested filter (E.g. `A = 1 AND (B = 2 OR C = 3)`)
- Support filter on all data types (both SV and MV)
- Support simple union without filters (E.g. `thetaSketch(col)`)
- Support `$0` as the default sketch (sketch without filter)
Also change the `DistinctCountRawThetaSketchAggregationFunction` to have the same usage as the `DistinctCountThetaSketchAggregationFunction` and returns base64 encoded sketch.
Fix the issue where theta-sketch parameters cannot be correctly handled in `BaseBrokerRequestHandler.updateColumnNames()`
## Release Notes
This aggregation function is still in beta version. This PR involves change on the format of data sent from server to broker, so it works only when both broker and server are upgraded to the new version.
---
.../requesthandler/BaseBrokerRequestHandler.java | 14 +-
.../common/function/AggregationFunctionType.java | 1 -
.../apache/pinot/core/common/ObjectSerDeUtils.java | 81 +-
.../function/AggregationFunctionFactory.java | 2 -
...inctCountRawThetaSketchAggregationFunction.java | 106 +-
...istinctCountThetaSketchAggregationFunction.java | 1691 ++++++++++++++------
.../RawThetaSketchAggregationFunction.java | 563 -------
.../BrokerRequestToQueryContextConverter.java | 23 +-
.../DistinctCountThetaSketchQueriesTest.java | 343 ++++
.../queries/DistinctCountThetaSketchTest.java | 378 -----
.../tests/ThetaSketchIntegrationTest.java | 119 +-
11 files changed, 1757 insertions(+), 1564 deletions(-)
diff --git a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
index 4925781..4c38a11 100644
--- a/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
+++ b/pinot-broker/src/main/java/org/apache/pinot/broker/requesthandler/BaseBrokerRequestHandler.java
@@ -750,10 +750,20 @@ public abstract class BaseBrokerRequestHandler implements BrokerRequestHandler {
}
if (brokerRequest.isSetAggregationsInfo()) {
for (AggregationInfo info : brokerRequest.getAggregationsInfo()) {
- if (!info.getAggregationType().equalsIgnoreCase(AggregationFunctionType.COUNT.getName())) {
+ String functionName = StringUtils.remove(info.getAggregationType(), '_');
+ if (!functionName.equalsIgnoreCase(AggregationFunctionType.COUNT.getName())) {
// Always read from backward compatible api in AggregationFunctionUtils.
List<String> arguments = AggregationFunctionUtils.getArguments(info);
- arguments.replaceAll(e -> fixColumnName(rawTableName, e, columnNameMap));
+
+ if (functionName.equalsIgnoreCase(AggregationFunctionType.DISTINCT.getName())) {
+ // For DISTINCT query, all arguments are expressions
+ arguments.replaceAll(e -> fixColumnName(rawTableName, e, columnNameMap));
+ } else {
+ // For non-DISTINCT query, only the first argument is expression, others are literals
+ // NOTE: We skip fixing the literal arguments because of the legacy behavior of PQL compiler treating string
+ // literal as identifier in the aggregation function.
+ arguments.set(0, fixColumnName(rawTableName, arguments.get(0), columnNameMap));
+ }
info.setExpressions(arguments);
}
}
diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/AggregationFunctionType.java b/pinot-common/src/main/java/org/apache/pinot/common/function/AggregationFunctionType.java
index 20f9761..aeae907 100644
--- a/pinot-common/src/main/java/org/apache/pinot/common/function/AggregationFunctionType.java
+++ b/pinot-common/src/main/java/org/apache/pinot/common/function/AggregationFunctionType.java
@@ -43,7 +43,6 @@ public enum AggregationFunctionType {
PERCENTILE("percentile"),
PERCENTILEEST("percentileEst"),
PERCENTILETDIGEST("percentileTDigest"),
- RAWTHETASKETCH("rawThetaSketch"),
IDSET("idSet"),
// Geo aggregation functions
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java
index 66c1b66..45ed146 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/common/ObjectSerDeUtils.java
@@ -41,8 +41,10 @@ import java.io.ByteArrayOutputStream;
import java.io.DataOutputStream;
import java.io.IOException;
import java.nio.ByteBuffer;
+import java.util.ArrayList;
import java.util.HashMap;
import java.util.Iterator;
+import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.datasketches.memory.Memory;
@@ -91,7 +93,8 @@ public class ObjectSerDeUtils {
DoubleSet(17),
StringSet(18),
BytesSet(19),
- IdSet(20);
+ IdSet(20),
+ List(21);
private final int _value;
@@ -149,6 +152,8 @@ public class ObjectSerDeUtils {
}
} else if (value instanceof IdSet) {
return ObjectType.IdSet;
+ } else if (value instanceof List) {
+ return ObjectType.List;
} else {
throw new IllegalArgumentException("Unsupported type of value: " + value.getClass().getSimpleName());
}
@@ -684,9 +689,12 @@ public class ObjectSerDeUtils {
};
public static final ObjectSerDe<Sketch> DATA_SKETCH_SER_DE = new ObjectSerDe<Sketch>() {
+
@Override
public byte[] serialize(Sketch value) {
- return value.compact().toByteArray();
+ // NOTE: Compact the sketch in unsorted, on-heap fashion for performance concern.
+ // See https://datasketches.apache.org/docs/Theta/ThetaSize.html for more details.
+ return value.compact(false, null).toByteArray();
}
@Override
@@ -703,6 +711,7 @@ public class ObjectSerDeUtils {
};
public static final ObjectSerDe<Geometry> GEOMETRY_SER_DE = new ObjectSerDe<Geometry>() {
+
@Override
public byte[] serialize(Geometry value) {
return GeometrySerializer.serialize(value);
@@ -722,6 +731,7 @@ public class ObjectSerDeUtils {
};
public static final ObjectSerDe<RoaringBitmap> ROARING_BITMAP_SER_DE = new ObjectSerDe<RoaringBitmap>() {
+
@Override
public byte[] serialize(RoaringBitmap bitmap) {
byte[] bytes = new byte[bitmap.serializedSizeInBytes()];
@@ -748,6 +758,7 @@ public class ObjectSerDeUtils {
};
public static final ObjectSerDe<IdSet> ID_SET_SER_DE = new ObjectSerDe<IdSet>() {
+
@Override
public byte[] serialize(IdSet idSet) {
try {
@@ -776,6 +787,69 @@ public class ObjectSerDeUtils {
}
};
+ public static final ObjectSerDe<List<Object>> LIST_SER_DE = new ObjectSerDe<List<Object>>() {
+
+ @Override
+ public byte[] serialize(List<Object> list) {
+ int size = list.size();
+
+ // Directly return the size (0) for empty list
+ if (size == 0) {
+ return new byte[Integer.BYTES];
+ }
+
+ // No need to close these 2 streams (close() is no-op)
+ ByteArrayOutputStream byteArrayOutputStream = new ByteArrayOutputStream();
+ DataOutputStream dataOutputStream = new DataOutputStream(byteArrayOutputStream);
+
+ try {
+ // Write the size of the list
+ dataOutputStream.writeInt(size);
+
+ // Write the value type
+ Object firstValue = list.get(0);
+ int valueType = ObjectType.getObjectType(firstValue).getValue();
+ dataOutputStream.writeInt(valueType);
+
+ // Write the serialized values
+ for (Object value : list) {
+ byte[] bytes = ObjectSerDeUtils.serialize(value, valueType);
+ dataOutputStream.writeInt(bytes.length);
+ dataOutputStream.write(bytes);
+ }
+ } catch (IOException e) {
+ throw new RuntimeException("Caught exception while serializing List", e);
+ }
+
+ return byteArrayOutputStream.toByteArray();
+ }
+
+ @Override
+ public ArrayList<Object> deserialize(byte[] bytes) {
+ return deserialize(ByteBuffer.wrap(bytes));
+ }
+
+ @Override
+ public ArrayList<Object> deserialize(ByteBuffer byteBuffer) {
+ int size = byteBuffer.getInt();
+ ArrayList<Object> list = new ArrayList<>(size);
+
+ // De-serialize the values
+ if (size != 0) {
+ int valueType = byteBuffer.getInt();
+ for (int i = 0; i < size; i++) {
+ int numBytes = byteBuffer.getInt();
+ ByteBuffer slice = byteBuffer.slice();
+ slice.limit(numBytes);
+ list.add(ObjectSerDeUtils.deserialize(slice, valueType));
+ byteBuffer.position(byteBuffer.position() + numBytes);
+ }
+ }
+
+ return list;
+ }
+ };
+
// NOTE: DO NOT change the order, it has to be the same order as the ObjectType
//@formatter:off
private static final ObjectSerDe[] SER_DES = {
@@ -799,7 +873,8 @@ public class ObjectSerDeUtils {
DOUBLE_SET_SER_DE,
STRING_SET_SER_DE,
BYTES_SET_SER_DE,
- ID_SET_SER_DE
+ ID_SET_SER_DE,
+ LIST_SER_DE
};
//@formatter:on
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java
index f0a7d8e..d74db3e 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionFactory.java
@@ -137,8 +137,6 @@ public class AggregationFunctionFactory {
return new DistinctCountThetaSketchAggregationFunction(arguments);
case DISTINCTCOUNTRAWTHETASKETCH:
return new DistinctCountRawThetaSketchAggregationFunction(arguments);
- case RAWTHETASKETCH:
- return new RawThetaSketchAggregationFunction(arguments);
case IDSET:
return new IdSetAggregationFunction(arguments);
case COUNTMV:
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountRawThetaSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountRawThetaSketchAggregationFunction.java
index c4185b0..c78f191 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountRawThetaSketchAggregationFunction.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountRawThetaSketchAggregationFunction.java
@@ -18,35 +18,22 @@
*/
package org.apache.pinot.core.query.aggregation.function;
+import java.util.Base64;
import java.util.List;
-import java.util.Map;
-import org.apache.calcite.sql.parser.SqlParseException;
import org.apache.datasketches.theta.Sketch;
import org.apache.pinot.common.function.AggregationFunctionType;
-import org.apache.pinot.common.utils.DataSchema;
-import org.apache.pinot.core.common.BlockValSet;
-import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
-import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.query.request.context.ExpressionContext;
-import org.apache.pinot.spi.utils.ByteArray;
-
-import static org.apache.pinot.common.utils.DataSchema.ColumnDataType.BYTES;
/**
- * A variation of the {@link DistinctCountThetaSketchAggregationFunction} that returns the serialized bytes
- * of the theta-sketch, as opposed to the actual distinct value.
- *
- * Note: It would have been natural for this class to extend the {@link DistinctCountThetaSketchAggregationFunction},
- * except that the return type for this class is a String, as opposed to Integer for the former, due to which the
- * extension is not possible.
+ * The {@code DistinctCountRawThetaSketchAggregationFunction} shares the same usage as the
+ * {@link DistinctCountThetaSketchAggregationFunction}, and returns the sketch as a base64 encoded string.
*/
-public class DistinctCountRawThetaSketchAggregationFunction implements AggregationFunction<Map<String, Sketch>, ByteArray> {
- private final DistinctCountThetaSketchAggregationFunction _thetaSketchAggregationFunction;
+public class DistinctCountRawThetaSketchAggregationFunction extends DistinctCountThetaSketchAggregationFunction {
- public DistinctCountRawThetaSketchAggregationFunction(List<ExpressionContext> arguments)
- throws SqlParseException {
- _thetaSketchAggregationFunction = new DistinctCountThetaSketchAggregationFunction(arguments);
+ public DistinctCountRawThetaSketchAggregationFunction(List<ExpressionContext> arguments) {
+ super(arguments);
}
@Override
@@ -55,81 +42,16 @@ public class DistinctCountRawThetaSketchAggregationFunction implements Aggregati
}
@Override
- public String getColumnName() {
- return _thetaSketchAggregationFunction.getColumnName();
- }
-
- @Override
- public String getResultColumnName() {
- return _thetaSketchAggregationFunction.getResultColumnName();
- }
-
- @Override
- public List<ExpressionContext> getInputExpressions() {
- return _thetaSketchAggregationFunction.getInputExpressions();
- }
-
- @Override
- public AggregationResultHolder createAggregationResultHolder() {
- return _thetaSketchAggregationFunction.createAggregationResultHolder();
- }
-
- @Override
- public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int maxCapacity) {
- return _thetaSketchAggregationFunction.createGroupByResultHolder(initialCapacity, maxCapacity);
- }
-
- @Override
- public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
- Map<ExpressionContext, BlockValSet> blockValSetMap) {
- _thetaSketchAggregationFunction.aggregate(length, aggregationResultHolder, blockValSetMap);
- }
-
- @Override
- public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
- Map<ExpressionContext, BlockValSet> blockValSetMap) {
- _thetaSketchAggregationFunction.aggregateGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSetMap);
+ public ColumnDataType getFinalResultColumnType() {
+ return ColumnDataType.STRING;
}
@Override
- public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
- Map<ExpressionContext, BlockValSet> blockValSetMap) {
- _thetaSketchAggregationFunction.aggregateGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSetMap);
- }
-
- @Override
- public Map<String, Sketch> extractAggregationResult(AggregationResultHolder aggregationResultHolder) {
- return _thetaSketchAggregationFunction.extractAggregationResult(aggregationResultHolder);
- }
-
- @Override
- public Map<String, Sketch> extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) {
- return _thetaSketchAggregationFunction.extractGroupByResult(groupByResultHolder, groupKey);
- }
-
- @Override
- public Map<String, Sketch> merge(Map<String, Sketch> intermediateResult1, Map<String, Sketch> intermediateResult2) {
- return _thetaSketchAggregationFunction.merge(intermediateResult1, intermediateResult2);
- }
+ public String extractFinalResult(List<Sketch> sketches) {
+ Sketch sketch = evaluatePostAggregationExpression(sketches);
- @Override
- public boolean isIntermediateResultComparable() {
- return _thetaSketchAggregationFunction.isIntermediateResultComparable();
- }
-
- @Override
- public DataSchema.ColumnDataType getIntermediateResultColumnType() {
- return _thetaSketchAggregationFunction.getIntermediateResultColumnType();
- }
-
- @Override
- public DataSchema.ColumnDataType getFinalResultColumnType() {
- return BYTES;
- }
-
- @Override
- public ByteArray extractFinalResult(Map<String, Sketch> intermediateResult) {
- Sketch finalSketch = _thetaSketchAggregationFunction.extractFinalSketch(intermediateResult);
- return new ByteArray(finalSketch.compact().toByteArray());
+ // NOTE: Compact the sketch in unsorted, on-heap fashion for performance concern.
+ // See https://datasketches.apache.org/docs/Theta/ThetaSize.html for more details.
+ return Base64.getEncoder().encodeToString(sketch.compact(false, null).toByteArray());
}
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java
index 049bead..3ef8a28 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java
@@ -19,136 +19,143 @@
package org.apache.pinot.core.query.aggregation.function;
import com.google.common.base.Preconditions;
-import com.google.common.collect.ImmutableList;
import java.util.ArrayList;
import java.util.Collections;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
-import java.util.regex.Matcher;
-import java.util.regex.Pattern;
import org.apache.calcite.sql.parser.SqlParseException;
-import org.apache.commons.collections.MapUtils;
+import org.apache.commons.lang3.StringUtils;
+import org.apache.datasketches.Util;
import org.apache.datasketches.memory.Memory;
import org.apache.datasketches.theta.AnotB;
import org.apache.datasketches.theta.Intersection;
import org.apache.datasketches.theta.SetOperationBuilder;
import org.apache.datasketches.theta.Sketch;
import org.apache.datasketches.theta.Union;
+import org.apache.datasketches.theta.UpdateSketch;
+import org.apache.datasketches.theta.UpdateSketchBuilder;
import org.apache.pinot.common.function.AggregationFunctionType;
-import org.apache.pinot.common.utils.DataSchema;
+import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
import org.apache.pinot.core.common.BlockValSet;
import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluator;
import org.apache.pinot.core.operator.filter.predicate.PredicateEvaluatorProvider;
import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
import org.apache.pinot.core.query.aggregation.ObjectAggregationResultHolder;
-import org.apache.pinot.core.query.aggregation.function.RawThetaSketchAggregationFunction.Parameters;
import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
import org.apache.pinot.core.query.aggregation.groupby.ObjectGroupByResultHolder;
import org.apache.pinot.core.query.request.context.ExpressionContext;
import org.apache.pinot.core.query.request.context.FilterContext;
+import org.apache.pinot.core.query.request.context.FunctionContext;
import org.apache.pinot.core.query.request.context.predicate.Predicate;
import org.apache.pinot.core.query.request.context.utils.QueryContextConverterUtils;
-import org.apache.pinot.spi.data.FieldSpec;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
import org.apache.pinot.sql.parsers.CalciteSqlParser;
/**
- * Implementation of {@link AggregationFunction} to perform the distinct count aggregation using
- * Theta Sketches.
- * <p>TODO: For performance concern, use {@code List<Sketch>} as the intermediate result.
+ * The {@code DistinctCountThetaSketchAggregationFunction} can be used in 2 modes:
+ * <ul>
+ * <li>
+ * Simple union without post-aggregation (1 or 2 arguments): main expression to aggregate on, optional theta-sketch
+ * parameters
+ * <p>E.g. DISTINCT_COUNT_THETA_SKETCH(col)
+ * </li>
+ * <li>
+ * Union with post-aggregation (at least 4 arguments): main expression to aggregate on, theta-sketch parameters,
+ * filter(s), post-aggregation expression
+ * <p>E.g. DISTINCT_COUNT_THETA_SKETCH(col, '', 'dimName=''gender'' AND dimValue=''male''',
+ * 'dimName=''course'' AND dimValue=''math''', 'SET_INTERSECT($1,$2)')
+ * </li>
+ * </ul>
+ * Currently there is only 1 parameter for the function:
+ * <ul>
+ * <li>
+ * nominalEntries: The nominal entries used to create the sketch. (Default 4096)
+ * </li>
+ * </ul>
+ * <p>E.g. DISTINCT_COUNT_THETA_SKETCH(col, 'nominalEntries=8192')
*/
-public class DistinctCountThetaSketchAggregationFunction implements AggregationFunction<Map<String, Sketch>, Long> {
+@SuppressWarnings({"rawtypes", "unchecked"})
+public class DistinctCountThetaSketchAggregationFunction extends BaseSingleInputAggregationFunction<List<Sketch>, Comparable> {
+ private static final String SET_UNION = "SET_UNION";
+ private static final String SET_INTERSECT = "SET_INTERSECT";
+ private static final String SET_DIFF = "SET_DIFF";
+ private static final String DEFAULT_SKETCH_IDENTIFIER = "$0";
+ private static final Sketch EMPTY_SKETCH = new UpdateSketchBuilder().build().compact();
- public enum MergeFunction {
- SET_UNION, SET_INTERSECT, SET_DIFF;
-
- public static final ImmutableList<String> STRING_VALUES =
- ImmutableList.of(SET_UNION.name(), SET_INTERSECT.name(), SET_DIFF.name());
-
- public static final String CSV_VALUES = String.join(",", STRING_VALUES);
-
- public static boolean isValid(String name) {
- return SET_UNION.name().equalsIgnoreCase(name) || SET_INTERSECT.name().equalsIgnoreCase(name) || SET_DIFF.name()
- .equalsIgnoreCase(name);
- }
- }
-
- private static final Pattern ARGUMENT_SUBSTITUTION = Pattern.compile("\\$(\\d+)");
-
- private final ExpressionContext _thetaSketchColumn;
- private final SetOperationBuilder _setOperationBuilder;
private final List<ExpressionContext> _inputExpressions;
+ private final boolean _includeDefaultSketch;
+ private final List<FilterEvaluator> _filterEvaluators;
private final ExpressionContext _postAggregationExpression;
- private final List<Predicate> _predicates;
- private final Map<Predicate, PredicateInfo> _predicateInfoMap;
-
- /**
- * Constructor for the class.
- * @param arguments List of parameters as arguments strings. At least three arguments are expected:
- * <ul>
- * <li> Required: First expression is interpreted as theta sketch column to aggregate on. </li>
- * <li> Required: Second argument is the thetaSketchParams. </li>
- * <li> Optional: Third to penultimate are predicates with LHS and RHS. </li>
- * <li> Required: Last expression is the one that will be evaluated to compute final result. </li>
- * </ul>
- */
- public DistinctCountThetaSketchAggregationFunction(List<ExpressionContext> arguments)
- throws SqlParseException {
- int numArguments = arguments.size();
-
- // NOTE: This function expects at least 4 arguments: theta-sketch column, nominalEntries, predicate(s), post-aggregation expression.
- Preconditions.checkArgument(numArguments > 3,
- "DistinctCountThetaSketch expects at least four arguments (theta-sketch column, parameter(s), post-aggregation expression), got: ",
- numArguments);
-
- // Initialize the theta-sketch column
- _thetaSketchColumn = arguments.get(0);
- Preconditions.checkArgument(_thetaSketchColumn.getType() == ExpressionContext.Type.IDENTIFIER,
- "First argument of DistinctCountThetaSketch must be identifier (theta-sketch column)");
-
- // Initialize the theta-sketch parameters
- ExpressionContext parametersExpression = arguments.get(1);
- Preconditions.checkArgument(parametersExpression.getType() == ExpressionContext.Type.LITERAL,
- "Second argument of DistinctCountThetaSketch must be literal (parameters)");
- Parameters parameters = new Parameters(parametersExpression.getLiteral());
+ private final UpdateSketchBuilder _updateSketchBuilder = new UpdateSketchBuilder();
+ private final SetOperationBuilder _setOperationBuilder = new SetOperationBuilder();
- // Initialize the theta-sketch set operation builder
- _setOperationBuilder = new SetOperationBuilder().setNominalEntries(parameters.getNominalEntries());
+ public DistinctCountThetaSketchAggregationFunction(List<ExpressionContext> arguments) {
+ super(arguments.get(0));
- // Index of the original input predicates
- // This list is zero indexed, whereas argument substitution is 1-indexed: index[0] = $1
- _predicates = new ArrayList<>();
+ // Initialize the UpdateSketchBuilder and SetOperationBuilder with the parameters
+ int numArguments = arguments.size();
+ if (numArguments > 1) {
+ ExpressionContext paramsExpression = arguments.get(1);
+ Preconditions.checkArgument(paramsExpression.getType() == ExpressionContext.Type.LITERAL,
+ "Second argument of DISTINCT_COUNT_THETA_SKETCH aggregation function must be literal (parameters)");
+ Parameters parameters = new Parameters(paramsExpression.getLiteral());
+ int nominalEntries = parameters.getNominalEntries();
+ _updateSketchBuilder.setNominalEntries(nominalEntries);
+ _setOperationBuilder.setNominalEntries(nominalEntries);
+ }
- // Initialize the input expressions
- // NOTE: It is expected to cover the theta-sketch column and the lhs of the predicates.
- _inputExpressions = new ArrayList<>();
- _inputExpressions.add(_thetaSketchColumn);
+ if (numArguments < 4) {
+ // Simple union without post-aggregation
- // Initialize the post-aggregation expression
- // NOTE: It is modeled as a filter
- ExpressionContext postAggregationExpression = arguments.get(numArguments - 1);
- Preconditions.checkArgument(parametersExpression.getType() == ExpressionContext.Type.LITERAL,
- "Last argument of DistinctCountThetaSketch must be literal (post-aggregation expression)");
- _postAggregationExpression = QueryContextConverterUtils
- .getExpression(CalciteSqlParser.compileToExpression(postAggregationExpression.getLiteral()));
+ _inputExpressions = Collections.singletonList(_expression);
+ _includeDefaultSketch = true;
+ _filterEvaluators = Collections.emptyList();
+ _postAggregationExpression = ExpressionContext.forIdentifier(DEFAULT_SKETCH_IDENTIFIER);
+ } else {
+ // Union with post-aggregation
+
+ // Input expressions should include the main expression and the lhs of the predicates in the filters
+ _inputExpressions = new ArrayList<>();
+ _inputExpressions.add(_expression);
+ Map<ExpressionContext, Integer> expressionIndexMap = new HashMap<>();
+ expressionIndexMap.put(_expression, 0);
+
+ // Process the filter expressions
+ _filterEvaluators = new ArrayList<>(numArguments - 3);
+ for (int i = 2; i < numArguments - 1; i++) {
+ ExpressionContext filterExpression = arguments.get(i);
+ Preconditions.checkArgument(filterExpression.getType() == ExpressionContext.Type.LITERAL,
+ "Third to second last argument of DISTINCT_COUNT_THETA_SKETCH aggregation function must be literal (filter expression)");
+ FilterContext filter;
+ try {
+ filter =
+ QueryContextConverterUtils.getFilter(CalciteSqlParser.compileToExpression(filterExpression.getLiteral()));
+ } catch (SqlParseException e) {
+ throw new IllegalArgumentException("Invalid filter expression: " + filterExpression.getLiteral());
+ }
+ // NOTE: Collect expressions before constructing the FilterInfo so that expressionIndexMap always include the
+ // expressions in the filter.
+ collectExpressions(filter, _inputExpressions, expressionIndexMap);
+ _filterEvaluators.add(getFilterEvaluator(filter, expressionIndexMap));
+ }
- // Initialize the predicate map
- _predicateInfoMap = new HashMap<>();
+ // Process the post-aggregation expression
+ ExpressionContext postAggregationExpression = arguments.get(numArguments - 1);
+ Preconditions.checkArgument(postAggregationExpression.getType() == ExpressionContext.Type.LITERAL,
+ "Last argument of DISTINCT_COUNT_THETA_SKETCH aggregation function must be literal (post-aggregation expression)");
+ try {
+ _postAggregationExpression = QueryContextConverterUtils
+ .getExpression(CalciteSqlParser.compileToExpression(postAggregationExpression.getLiteral()));
+ } catch (SqlParseException e) {
+ throw new IllegalArgumentException(
+ "Invalid post-aggregation expression: " + postAggregationExpression.getLiteral());
+ }
- // Predicates are explicitly specified
- for (int i = 2; i < numArguments - 1; i++) {
- ExpressionContext predicateExpression = arguments.get(i);
- Preconditions.checkArgument(predicateExpression.getType() == ExpressionContext.Type.LITERAL,
- "Third to second last argument of DistinctCountThetaSketch must be literal (predicate expression)");
- Predicate predicate = getPredicate(predicateExpression.getLiteral());
- _inputExpressions.add(predicate.getLhs());
- _predicates.add(predicate);
- _predicateInfoMap.put(predicate, new PredicateInfo(predicate));
+ // Validate the post-aggregation expression
+ _includeDefaultSketch = validatePostAggregationExpression(_postAggregationExpression, _filterEvaluators.size());
}
-
- // First expression is the nominal entries parameter
- validatePostAggregationExpression(_postAggregationExpression, _predicates.size());
}
@Override
@@ -157,16 +164,6 @@ public class DistinctCountThetaSketchAggregationFunction implements AggregationF
}
@Override
- public String getColumnName() {
- return AggregationFunctionType.DISTINCTCOUNTTHETASKETCH.getName() + "_" + _thetaSketchColumn;
- }
-
- @Override
- public String getResultColumnName() {
- return AggregationFunctionType.DISTINCTCOUNTTHETASKETCH.getName().toLowerCase() + "(" + _thetaSketchColumn + ")";
- }
-
- @Override
public List<ExpressionContext> getInputExpressions() {
return _inputExpressions;
}
@@ -184,67 +181,248 @@ public class DistinctCountThetaSketchAggregationFunction implements AggregationF
@Override
public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
Map<ExpressionContext, BlockValSet> blockValSetMap) {
- Map<Predicate, Union> unionMap = getUnionMap(aggregationResultHolder);
-
- Sketch[] sketches = deserializeSketches(blockValSetMap.get(_thetaSketchColumn).getBytesValuesSV(), length);
- for (PredicateInfo predicateInfo : _predicateInfoMap.values()) {
- Predicate predicate = predicateInfo.getPredicate();
- BlockValSet blockValSet = blockValSetMap.get(predicate.getLhs());
- FieldSpec.DataType valueType = blockValSet.getValueType();
- PredicateEvaluator predicateEvaluator = predicateInfo.getPredicateEvaluator(valueType);
-
- Union union = unionMap.get(predicate);
- switch (valueType) {
- case INT:
- int[] intValues = blockValSet.getIntValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(intValues[i])) {
- union.update(sketches[i]);
+ int numExpressions = _inputExpressions.size();
+ boolean[] singleValues = new boolean[numExpressions];
+ DataType[] valueTypes = new DataType[numExpressions];
+ Object[] valueArrays = new Object[numExpressions];
+ extractValues(blockValSetMap, singleValues, valueTypes, valueArrays);
+ int numFilters = _filterEvaluators.size();
+
+ // Main expression is always index 0
+ if (valueTypes[0] != DataType.BYTES) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(aggregationResultHolder);
+ if (singleValues[0]) {
+ switch (valueTypes[0]) {
+ case INT:
+ int[] intValues = (int[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ defaultSketch.update(intValues[i]);
+ }
}
- }
- break;
- case LONG:
- long[] longValues = blockValSet.getLongValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(longValues[i])) {
- union.update(sketches[i]);
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ updateSketch.update(intValues[j]);
+ }
+ }
}
- }
- break;
- case FLOAT:
- float[] floatValues = blockValSet.getFloatValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(floatValues[i])) {
- union.update(sketches[i]);
+ break;
+ case LONG:
+ long[] longValues = (long[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ defaultSketch.update(longValues[i]);
+ }
}
- }
- break;
- case DOUBLE:
- double[] doubleValues = blockValSet.getDoubleValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(doubleValues[i])) {
- union.update(sketches[i]);
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ updateSketch.update(longValues[j]);
+ }
+ }
}
- }
- break;
- case STRING:
- String[] stringValues = blockValSet.getStringValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(stringValues[i])) {
- union.update(sketches[i]);
+ break;
+ case FLOAT:
+ float[] floatValues = (float[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ defaultSketch.update(floatValues[i]);
+ }
}
- }
- break;
- case BYTES:
- byte[][] bytesValues = blockValSet.getBytesValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(bytesValues[i])) {
- union.update(sketches[i]);
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ updateSketch.update(floatValues[j]);
+ }
+ }
+ }
+ break;
+ case DOUBLE:
+ double[] doubleValues = (double[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ defaultSketch.update(doubleValues[i]);
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ updateSketch.update(doubleValues[j]);
+ }
+ }
+ }
+ break;
+ case STRING:
+ String[] stringValues = (String[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ defaultSketch.update(stringValues[i]);
+ }
}
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ updateSketch.update(stringValues[j]);
+ }
+ }
+ }
+ break;
+ default:
+ throw new IllegalStateException(
+ "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: "
+ + valueTypes[0]);
+ }
+ } else {
+ switch (valueTypes[0]) {
+ case INT:
+ int[][] intValues = (int[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ for (int value : intValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int value : intValues[j]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case LONG:
+ long[][] longValues = (long[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ for (long value : longValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (long value : longValues[j]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case FLOAT:
+ float[][] floatValues = (float[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ for (float value : floatValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (float value : floatValues[j]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case DOUBLE:
+ double[][] doubleValues = (double[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ for (double value : doubleValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (double value : doubleValues[j]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case STRING:
+ String[][] stringValues = (String[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int i = 0; i < length; i++) {
+ for (String value : stringValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ UpdateSketch updateSketch = updateSketches.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (String value : stringValues[j]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ default:
+ throw new IllegalStateException(
+ "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]);
+ }
+ }
+ } else {
+ // Serialized sketch
+ List<Union> unions = getUnions(aggregationResultHolder);
+ Sketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length);
+ if (_includeDefaultSketch) {
+ Union defaultUnion = unions.get(0);
+ for (Sketch sketch : sketches) {
+ defaultUnion.update(sketch);
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ Union union = unions.get(i + 1);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ union.update(sketches[j]);
}
- break;
- default:
- throw new IllegalStateException();
+ }
}
}
}
@@ -252,64 +430,223 @@ public class DistinctCountThetaSketchAggregationFunction implements AggregationF
@Override
public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
Map<ExpressionContext, BlockValSet> blockValSetMap) {
- Sketch[] sketches = deserializeSketches(blockValSetMap.get(_thetaSketchColumn).getBytesValuesSV(), length);
- for (PredicateInfo predicateInfo : _predicateInfoMap.values()) {
- Predicate predicate = predicateInfo.getPredicate();
- BlockValSet blockValSet = blockValSetMap.get(predicate.getLhs());
- FieldSpec.DataType valueType = blockValSet.getValueType();
- PredicateEvaluator predicateEvaluator = predicateInfo.getPredicateEvaluator(valueType);
-
- switch (valueType) {
- case INT:
- int[] intValues = blockValSet.getIntValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(intValues[i])) {
- getUnionMap(groupByResultHolder, groupKeyArray[i]).get(predicate).update(sketches[i]);
+ int numExpressions = _inputExpressions.size();
+ boolean[] singleValues = new boolean[numExpressions];
+ DataType[] valueTypes = new DataType[numExpressions];
+ Object[] valueArrays = new Object[numExpressions];
+ extractValues(blockValSetMap, singleValues, valueTypes, valueArrays);
+ int numFilters = _filterEvaluators.size();
+
+ // Main expression is always index 0
+ if (valueTypes[0] != DataType.BYTES) {
+ if (singleValues[0]) {
+ switch (valueTypes[0]) {
+ case INT:
+ int[] intValues = (int[]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ int value = intValues[i];
+ if (_includeDefaultSketch) {
+ updateSketches.get(0).update(value);
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ updateSketches.get(j + 1).update(value);
+ }
+ }
}
- }
- break;
- case LONG:
- long[] longValues = blockValSet.getLongValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(longValues[i])) {
- getUnionMap(groupByResultHolder, groupKeyArray[i]).get(predicate).update(sketches[i]);
+ break;
+ case LONG:
+ long[] longValues = (long[]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ long value = longValues[i];
+ if (_includeDefaultSketch) {
+ updateSketches.get(0).update(value);
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ updateSketches.get(j + 1).update(value);
+ }
+ }
}
- }
- break;
- case FLOAT:
- float[] floatValues = blockValSet.getFloatValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(floatValues[i])) {
- getUnionMap(groupByResultHolder, groupKeyArray[i]).get(predicate).update(sketches[i]);
+ break;
+ case FLOAT:
+ float[] floatValues = (float[]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ float value = floatValues[i];
+ if (_includeDefaultSketch) {
+ updateSketches.get(0).update(value);
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ updateSketches.get(j + 1).update(value);
+ }
+ }
}
- }
- break;
- case DOUBLE:
- double[] doubleValues = blockValSet.getDoubleValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(doubleValues[i])) {
- getUnionMap(groupByResultHolder, groupKeyArray[i]).get(predicate).update(sketches[i]);
+ break;
+ case DOUBLE:
+ double[] doubleValues = (double[]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ double value = doubleValues[i];
+ if (_includeDefaultSketch) {
+ updateSketches.get(0).update(value);
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ updateSketches.get(j + 1).update(value);
+ }
+ }
}
- }
- break;
- case STRING:
- String[] stringValues = blockValSet.getStringValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(stringValues[i])) {
- getUnionMap(groupByResultHolder, groupKeyArray[i]).get(predicate).update(sketches[i]);
+ break;
+ case STRING:
+ String[] stringValues = (String[]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ String value = stringValues[i];
+ if (_includeDefaultSketch) {
+ updateSketches.get(0).update(value);
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ updateSketches.get(j + 1).update(value);
+ }
+ }
}
- }
- break;
- case BYTES:
- byte[][] bytesValues = blockValSet.getBytesValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(bytesValues[i])) {
- getUnionMap(groupByResultHolder, groupKeyArray[i]).get(predicate).update(sketches[i]);
+ break;
+ default:
+ throw new IllegalStateException(
+ "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: "
+ + valueTypes[0]);
+ }
+ } else {
+ switch (valueTypes[0]) {
+ case INT:
+ int[][] intValues = (int[][]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ int[] values = intValues[i];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (int value : values) {
+ defaultSketch.update(value);
+ }
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ UpdateSketch updateSketch = updateSketches.get(j + 1);
+ for (int value : values) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case LONG:
+ long[][] longValues = (long[][]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ long[] values = longValues[i];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (long value : values) {
+ defaultSketch.update(value);
+ }
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ UpdateSketch updateSketch = updateSketches.get(j + 1);
+ for (long value : values) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case FLOAT:
+ float[][] floatValues = (float[][]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ float[] values = floatValues[i];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (float value : values) {
+ defaultSketch.update(value);
+ }
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ UpdateSketch updateSketch = updateSketches.get(j + 1);
+ for (float value : values) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case DOUBLE:
+ double[][] doubleValues = (double[][]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ double[] values = doubleValues[i];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (double value : values) {
+ defaultSketch.update(value);
+ }
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ UpdateSketch updateSketch = updateSketches.get(j + 1);
+ for (double value : values) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ break;
+ case STRING:
+ String[][] stringValues = (String[][]) valueArrays[0];
+ for (int i = 0; i < length; i++) {
+ List<UpdateSketch> updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]);
+ String[] values = stringValues[i];
+ if (_includeDefaultSketch) {
+ UpdateSketch defaultSketch = updateSketches.get(0);
+ for (String value : values) {
+ defaultSketch.update(value);
+ }
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ UpdateSketch updateSketch = updateSketches.get(j + 1);
+ for (String value : values) {
+ updateSketch.update(value);
+ }
+ }
+ }
}
+ break;
+ default:
+ throw new IllegalStateException(
+ "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]);
+ }
+ }
+ } else {
+ // Serialized sketch
+ Sketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length);
+ for (int i = 0; i < length; i++) {
+ List<Union> unions = getUnions(groupByResultHolder, groupKeyArray[i]);
+ Sketch sketch = sketches[i];
+ if (_includeDefaultSketch) {
+ unions.get(0).update(sketch);
+ }
+ for (int j = 0; j < numFilters; j++) {
+ if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) {
+ unions.get(j + 1).update(sketch);
}
- break;
- default:
- throw new IllegalStateException();
+ }
}
}
}
@@ -317,213 +654,590 @@ public class DistinctCountThetaSketchAggregationFunction implements AggregationF
@Override
public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
Map<ExpressionContext, BlockValSet> blockValSetMap) {
- Sketch[] sketches = deserializeSketches(blockValSetMap.get(_thetaSketchColumn).getBytesValuesSV(), length);
- for (PredicateInfo predicateInfo : _predicateInfoMap.values()) {
- Predicate predicate = predicateInfo.getPredicate();
- BlockValSet blockValSet = blockValSetMap.get(predicate.getLhs());
- FieldSpec.DataType valueType = blockValSet.getValueType();
- PredicateEvaluator predicateEvaluator = predicateInfo.getPredicateEvaluator(valueType);
-
- switch (valueType) {
- case INT:
- int[] intValues = blockValSet.getIntValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(intValues[i])) {
- for (int groupKey : groupKeysArray[i]) {
- getUnionMap(groupByResultHolder, groupKey).get(predicate).update(sketches[i]);
+ int numExpressions = _inputExpressions.size();
+ boolean[] singleValues = new boolean[numExpressions];
+ DataType[] valueTypes = new DataType[numExpressions];
+ Object[] valueArrays = new Object[numExpressions];
+ extractValues(blockValSetMap, singleValues, valueTypes, valueArrays);
+ int numFilters = _filterEvaluators.size();
+
+ // Main expression is always index 0
+ if (valueTypes[0] != DataType.BYTES) {
+ if (singleValues[0]) {
+ switch (valueTypes[0]) {
+ case INT:
+ int[] intValues = (int[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(0).update(intValues[i]);
+ }
}
}
- }
- break;
- case LONG:
- long[] longValues = blockValSet.getLongValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(longValues[i])) {
- for (int groupKey : groupKeysArray[i]) {
- getUnionMap(groupByResultHolder, groupKey).get(predicate).update(sketches[i]);
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(intValues[j]);
+ }
+ }
}
}
- }
- break;
- case FLOAT:
- float[] floatValues = blockValSet.getFloatValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(floatValues[i])) {
- for (int groupKey : groupKeysArray[i]) {
- getUnionMap(groupByResultHolder, groupKey).get(predicate).update(sketches[i]);
+ break;
+ case LONG:
+ long[] longValues = (long[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(0).update(longValues[i]);
+ }
}
}
- }
- break;
- case DOUBLE:
- double[] doubleValues = blockValSet.getDoubleValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(doubleValues[i])) {
- for (int groupKey : groupKeysArray[i]) {
- getUnionMap(groupByResultHolder, groupKey).get(predicate).update(sketches[i]);
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(longValues[j]);
+ }
+ }
}
}
- }
- break;
- case STRING:
- String[] stringValues = blockValSet.getStringValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(stringValues[i])) {
- for (int groupKey : groupKeysArray[i]) {
- getUnionMap(groupByResultHolder, groupKey).get(predicate).update(sketches[i]);
+ break;
+ case FLOAT:
+ float[] floatValues = (float[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(0).update(floatValues[i]);
+ }
}
}
- }
- break;
- case BYTES:
- byte[][] bytesValues = blockValSet.getBytesValuesSV();
- for (int i = 0; i < length; i++) {
- if (predicateEvaluator.applySV(bytesValues[i])) {
- for (int groupKey : groupKeysArray[i]) {
- getUnionMap(groupByResultHolder, groupKey).get(predicate).update(sketches[i]);
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(floatValues[j]);
+ }
+ }
+ }
+ }
+ break;
+ case DOUBLE:
+ double[] doubleValues = (double[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(0).update(doubleValues[i]);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(doubleValues[j]);
+ }
+ }
+ }
+ }
+ break;
+ case STRING:
+ String[] stringValues = (String[]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(0).update(stringValues[i]);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(stringValues[j]);
+ }
+ }
}
}
+ break;
+ default:
+ throw new IllegalStateException(
+ "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: "
+ + valueTypes[0]);
+ }
+ } else {
+ switch (valueTypes[0]) {
+ case INT:
+ int[][] intValues = (int[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0);
+ for (int value : intValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1);
+ for (int value : intValues[i]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ }
+ break;
+ case LONG:
+ long[][] longValues = (long[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0);
+ for (long value : longValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1);
+ for (long value : longValues[i]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ }
+ break;
+ case FLOAT:
+ float[][] floatValues = (float[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0);
+ for (float value : floatValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1);
+ for (float value : floatValues[i]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ }
+ break;
+ case DOUBLE:
+ double[][] doubleValues = (double[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0);
+ for (double value : doubleValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1);
+ for (double value : doubleValues[i]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ }
+ break;
+ case STRING:
+ String[][] stringValues = (String[][]) valueArrays[0];
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0);
+ for (String value : stringValues[i]) {
+ defaultSketch.update(value);
+ }
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ UpdateSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1);
+ for (String value : stringValues[i]) {
+ updateSketch.update(value);
+ }
+ }
+ }
+ }
+ }
+ break;
+ default:
+ throw new IllegalStateException(
+ "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]);
+ }
+ }
+ } else {
+ // Serialized sketch
+ Sketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length);
+ if (_includeDefaultSketch) {
+ for (int i = 0; i < length; i++) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUnions(groupByResultHolder, groupKey).get(0).update(sketches[i]);
+ }
+ }
+ }
+ for (int i = 0; i < numFilters; i++) {
+ FilterEvaluator filterEvaluator = _filterEvaluators.get(i);
+ for (int j = 0; j < length; j++) {
+ if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) {
+ for (int groupKey : groupKeysArray[i]) {
+ getUnions(groupByResultHolder, groupKey).get(i + 1).update(sketches[i]);
+ }
}
- break;
- default:
- throw new IllegalStateException();
+ }
}
}
}
@Override
- public Map<String, Sketch> extractAggregationResult(AggregationResultHolder aggregationResultHolder) {
- Map<Predicate, Union> unionMap = aggregationResultHolder.getResult();
- if (unionMap == null) {
- return Collections.emptyMap();
+ public List<Sketch> extractAggregationResult(AggregationResultHolder aggregationResultHolder) {
+ List result = aggregationResultHolder.getResult();
+ if (result == null) {
+ int numSketches = _filterEvaluators.size() + 1;
+ List<Sketch> sketches = new ArrayList<>(numSketches);
+ for (int i = 0; i < numSketches; i++) {
+ sketches.add(EMPTY_SKETCH);
+ }
+ return sketches;
}
- Map<String, Sketch> result = new HashMap<>();
- for (PredicateInfo predicateInfo : _predicateInfoMap.values()) {
- Sketch sketch = unionMap.get(predicateInfo.getPredicate()).getResult();
-
- // Skip empty sketches, as they lead to unnecessary unions (and cost performance)
- if (!sketch.isEmpty()) {
- result.put(predicateInfo.getStringPredicate(), sketch);
- }
+ if (result.get(0) instanceof Sketch) {
+ return result;
+ } else {
+ return convertToSketches(result);
}
- return result;
}
@Override
- public Map<String, Sketch> extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) {
- Map<Predicate, Union> unionMap = groupByResultHolder.getResult(groupKey);
- if (unionMap == null) {
- return Collections.emptyMap();
+ public List<Sketch> extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) {
+ List result = groupByResultHolder.getResult(groupKey);
+ if (result.get(0) instanceof Sketch) {
+ return result;
+ } else {
+ return convertToSketches(result);
}
+ }
- Map<String, Sketch> result = new HashMap<>();
- for (PredicateInfo predicateInfo : _predicateInfoMap.values()) {
- Sketch sketch = unionMap.get(predicateInfo.getPredicate()).getResult();
-
- // Skip empty sketches, as they lead to unnecessary unions (and cost performance)
- if (!sketch.isEmpty()) {
- result.put(predicateInfo.getStringPredicate(), sketch);
+ @Override
+ public List<Sketch> merge(List<Sketch> sketches1, List<Sketch> sketches2) {
+ int numSketches = sketches1.size();
+ List<Sketch> mergedSketches = new ArrayList<>(numSketches);
+ for (int i = 0; i < numSketches; i++) {
+ Sketch sketch1 = sketches1.get(i);
+ Sketch sketch2 = sketches2.get(i);
+ if (sketch1.isEmpty()) {
+ mergedSketches.add(sketch2);
+ continue;
}
+ if (sketch2.isEmpty()) {
+ mergedSketches.add(sketch1);
+ continue;
+ }
+ Union union = _setOperationBuilder.buildUnion();
+ union.update(sketch1);
+ union.update(sketch2);
+ // NOTE: Compact the sketch in unsorted, on-heap fashion for performance concern.
+ // See https://datasketches.apache.org/docs/Theta/ThetaSize.html for more details.
+ mergedSketches.add(union.getResult(false, null));
}
- return result;
+ return mergedSketches;
}
@Override
- public Map<String, Sketch> merge(Map<String, Sketch> intermediateResult1, Map<String, Sketch> intermediateResult2) {
- if (MapUtils.isEmpty(intermediateResult1)) {
- return intermediateResult2;
- }
- if (MapUtils.isEmpty(intermediateResult2)) {
- return intermediateResult1;
+ public boolean isIntermediateResultComparable() {
+ return false;
+ }
+
+ @Override
+ public ColumnDataType getIntermediateResultColumnType() {
+ return ColumnDataType.OBJECT;
+ }
+
+ @Override
+ public ColumnDataType getFinalResultColumnType() {
+ return ColumnDataType.LONG;
+ }
+
+ @Override
+ public Comparable extractFinalResult(List<Sketch> sketches) {
+ return Math.round(evaluatePostAggregationExpression(_postAggregationExpression, sketches).getEstimate());
+ }
+
+ /**
+ * Helper method to collect expressions in the filter.
+ */
+ private static void collectExpressions(FilterContext filter, List<ExpressionContext> expressions,
+ Map<ExpressionContext, Integer> expressionIndexMap) {
+ List<FilterContext> children = filter.getChildren();
+ if (children != null) {
+ for (FilterContext child : children) {
+ collectExpressions(child, expressions, expressionIndexMap);
+ }
+ } else {
+ ExpressionContext expression = filter.getPredicate().getLhs();
+ if (expressionIndexMap.putIfAbsent(expression, expressions.size()) == null) {
+ expressions.add(expression);
+ }
}
+ }
- // Add sketches from intermediateResult1, merged with overlapping ones from intermediateResult2
- Map<String, Sketch> mergedResult = new HashMap<>();
- for (Map.Entry<String, Sketch> entry : intermediateResult1.entrySet()) {
- String predicate = entry.getKey();
- Sketch sketch = intermediateResult2.get(predicate);
- if (sketch != null) {
- // Merge the overlapping ones
- Union union = _setOperationBuilder.buildUnion();
- union.update(entry.getValue());
- union.update(sketch);
- mergedResult.put(predicate, union.getResult());
+ /**
+ * Creates a FilterEvaluator for the given filter.
+ */
+ private static FilterEvaluator getFilterEvaluator(FilterContext filter,
+ Map<ExpressionContext, Integer> expressionIndexMap) {
+ List<FilterContext> children = filter.getChildren();
+ if (children != null) {
+ // AND/OR
+ List<FilterEvaluator> childEvaluators = new ArrayList<>(children.size());
+ for (FilterContext child : children) {
+ childEvaluators.add(getFilterEvaluator(child, expressionIndexMap));
+ }
+ if (filter.getType() == FilterContext.Type.AND) {
+ return new AndFilterEvaluator(childEvaluators);
} else {
- // Collect the non-overlapping ones
- mergedResult.put(predicate, entry.getValue());
+ return new OrFilterEvaluator(childEvaluators);
}
+ } else {
+ // Predicate
+ Predicate predicate = filter.getPredicate();
+ int expressionIndex = expressionIndexMap.get(predicate.getLhs());
+ return new PredicateFilterEvaluator(predicate, expressionIndex);
}
+ }
- // Add sketches that are only in intermediateResult2
- for (Map.Entry<String, Sketch> entry : intermediateResult2.entrySet()) {
- // If key already present, it was already merged in the previous iteration.
- mergedResult.putIfAbsent(entry.getKey(), entry.getValue());
+ /**
+ * Validates the post-aggregation expression:
+ * - The sketch id ($0, $1, etc.) does not exceed the number of filters
+ * - Only contains valid set operations (SET_UNION/SET_INTERSECT/SET_DIFF)
+ * - SET_UNION/SET_INTERSECT contains at least 2 arguments
+ * - SET_DIFF contains exactly 2 arguments
+ * Returns whether the post-aggregation expression contains the default sketch ($0).
+ */
+ private static boolean validatePostAggregationExpression(ExpressionContext expression, int numFilters) {
+ Preconditions.checkArgument(expression.getType() != ExpressionContext.Type.LITERAL,
+ "Post-aggregation expression should not contain literal expression: %s", expression.getLiteral());
+
+ if (expression.getType() == ExpressionContext.Type.IDENTIFIER) {
+ int sketchId = extractSketchId(expression.getIdentifier());
+ Preconditions
+ .checkArgument(sketchId <= numFilters, "Sketch id: %s exceeds number of filters: %s", sketchId, numFilters);
+ return sketchId == 0;
}
- return mergedResult;
+ FunctionContext function = expression.getFunction();
+ String functionName = function.getFunctionName();
+ List<ExpressionContext> arguments = function.getArguments();
+ int numArguments = arguments.size();
+ boolean includeDefaultSketch = false;
+ switch (functionName.toUpperCase()) {
+ case SET_UNION:
+ case SET_INTERSECT:
+ Preconditions
+ .checkArgument(numArguments >= 2, "SET_UNION and SET_INTERSECT should have at least 2 arguments, got: %s",
+ numArguments);
+ for (ExpressionContext argument : arguments) {
+ includeDefaultSketch |= validatePostAggregationExpression(argument, numFilters);
+ }
+ break;
+ case SET_DIFF:
+ Preconditions.checkArgument(numArguments == 2, "SET_DIFF should have 2 arguments, got: %s", numArguments);
+ for (ExpressionContext argument : arguments) {
+ includeDefaultSketch |= validatePostAggregationExpression(argument, numFilters);
+ }
+ break;
+ default:
+ throw new IllegalArgumentException("Invalid set operation: " + functionName);
+ }
+ return includeDefaultSketch;
}
- @Override
- public boolean isIntermediateResultComparable() {
- return false;
+ /**
+ * Extracts the sketch id from the identifier (e.g. $0 -> 0, $1 -> 1).
+ */
+ private static int extractSketchId(String identifier) {
+ Preconditions
+ .checkArgument(identifier.charAt(0) == '$', "Invalid identifier: %s, expecting $0, $1, etc.", identifier);
+ int sketchId = Integer.parseInt(identifier.substring(1));
+ Preconditions.checkArgument(sketchId >= 0, "Invalid identifier: %s, expecting $0, $1, etc.", identifier);
+ return sketchId;
}
- @Override
- public DataSchema.ColumnDataType getIntermediateResultColumnType() {
- return DataSchema.ColumnDataType.OBJECT;
+ /**
+ * Extracts values from the BlockValSet map.
+ */
+ private void extractValues(Map<ExpressionContext, BlockValSet> blockValSetMap, boolean[] singleValues,
+ DataType[] valueTypes, Object[] valueArrays) {
+ int numExpressions = _inputExpressions.size();
+ for (int i = 0; i < numExpressions; i++) {
+ BlockValSet blockValSet = blockValSetMap.get(_inputExpressions.get(i));
+ boolean singleValue = blockValSet.isSingleValue();
+ DataType valueType = blockValSet.getValueType();
+ singleValues[i] = singleValue;
+ valueTypes[i] = valueType;
+ if (singleValue) {
+ switch (valueType) {
+ case INT:
+ valueArrays[i] = blockValSet.getIntValuesSV();
+ break;
+ case LONG:
+ valueArrays[i] = blockValSet.getLongValuesSV();
+ break;
+ case FLOAT:
+ valueArrays[i] = blockValSet.getFloatValuesSV();
+ break;
+ case DOUBLE:
+ valueArrays[i] = blockValSet.getDoubleValuesSV();
+ break;
+ case STRING:
+ valueArrays[i] = blockValSet.getStringValuesSV();
+ break;
+ case BYTES:
+ valueArrays[i] = blockValSet.getBytesValuesSV();
+ break;
+ default:
+ throw new IllegalStateException();
+ }
+ } else {
+ switch (valueType) {
+ case INT:
+ valueArrays[i] = blockValSet.getIntValuesMV();
+ break;
+ case LONG:
+ valueArrays[i] = blockValSet.getLongValuesMV();
+ break;
+ case FLOAT:
+ valueArrays[i] = blockValSet.getFloatValuesMV();
+ break;
+ case DOUBLE:
+ valueArrays[i] = blockValSet.getDoubleValuesMV();
+ break;
+ case STRING:
+ valueArrays[i] = blockValSet.getStringValuesMV();
+ break;
+ default:
+ throw new IllegalStateException();
+ }
+ }
+ }
}
- @Override
- public DataSchema.ColumnDataType getFinalResultColumnType() {
- return DataSchema.ColumnDataType.LONG;
+ /**
+ * Returns the UpdateSketch list from the result holder or creates a new one if it does not exist.
+ */
+ private List<UpdateSketch> getUpdateSketches(AggregationResultHolder aggregationResultHolder) {
+ List<UpdateSketch> updateSketches = aggregationResultHolder.getResult();
+ if (updateSketches == null) {
+ updateSketches = buildUpdateSketches();
+ aggregationResultHolder.setValue(updateSketches);
+ }
+ return updateSketches;
}
- @Override
- public Long extractFinalResult(Map<String, Sketch> intermediateResult) {
- Sketch finalSketch = extractFinalSketch(intermediateResult);
- return finalSketch != null ? Math.round(finalSketch.getEstimate()) : 0;
+ /**
+ * Returns the Union list from the result holder or creates a new one if it does not exist.
+ */
+ private List<Union> getUnions(AggregationResultHolder aggregationResultHolder) {
+ List<Union> unions = aggregationResultHolder.getResult();
+ if (unions == null) {
+ unions = buildUnions();
+ aggregationResultHolder.setValue(unions);
+ }
+ return unions;
}
- private Predicate getPredicate(String predicateString) {
- FilterContext filter;
- try {
- filter = QueryContextConverterUtils.getFilter(CalciteSqlParser.compileToExpression(predicateString));
- } catch (SqlParseException e) {
- throw new IllegalArgumentException("Invalid predicate string: " + predicateString);
+ /**
+ * Returns the UpdateSketch list for the given group key or creates a new one if it does not exist.
+ */
+ private List<UpdateSketch> getUpdateSketches(GroupByResultHolder groupByResultHolder, int groupKey) {
+ List<UpdateSketch> updateSketches = groupByResultHolder.getResult(groupKey);
+ if (updateSketches == null) {
+ updateSketches = buildUpdateSketches();
+ groupByResultHolder.setValueForKey(groupKey, updateSketches);
}
- // TODO: Add support for complex predicates with AND/OR.
- Preconditions.checkArgument(filter.getType() == FilterContext.Type.PREDICATE, "Invalid predicate string: %s",
- predicateString);
- return filter.getPredicate();
+ return updateSketches;
}
- private Map<Predicate, Union> getUnionMap(AggregationResultHolder aggregationResultHolder) {
- Map<Predicate, Union> unionMap = aggregationResultHolder.getResult();
- if (unionMap == null) {
- unionMap = getDefaultUnionMap();
- aggregationResultHolder.setValue(unionMap);
+ /**
+ * Returns the Union list for the given group key or creates a new one if it does not exist.
+ */
+ private List<Union> getUnions(GroupByResultHolder groupByResultHolder, int groupKey) {
+ List<Union> unions = groupByResultHolder.getResult(groupKey);
+ if (unions == null) {
+ unions = buildUnions();
+ groupByResultHolder.setValueForKey(groupKey, unions);
}
- return unionMap;
+ return unions;
}
- private Map<Predicate, Union> getUnionMap(GroupByResultHolder groupByResultHolder, int groupKey) {
- Map<Predicate, Union> unionMap = groupByResultHolder.getResult(groupKey);
- if (unionMap == null) {
- unionMap = getDefaultUnionMap();
- groupByResultHolder.setValueForKey(groupKey, unionMap);
+ /**
+ * Builds the UpdateSketch list.
+ */
+ private List<UpdateSketch> buildUpdateSketches() {
+ int numSketches = _filterEvaluators.size() + 1;
+ List<UpdateSketch> updateSketches = new ArrayList<>(numSketches);
+ for (int i = 0; i < numSketches; i++) {
+ updateSketches.add(_updateSketchBuilder.build());
}
- return unionMap;
+ return updateSketches;
}
- private Map<Predicate, Union> getDefaultUnionMap() {
- Map<Predicate, Union> unionMap = new HashMap<>();
- for (Predicate predicate : _predicateInfoMap.keySet()) {
- unionMap.put(predicate, _setOperationBuilder.buildUnion());
+ /**
+ * Builds the Union list.
+ */
+ private List<Union> buildUnions() {
+ int numUnions = _filterEvaluators.size() + 1;
+ List<Union> unions = new ArrayList<>(numUnions);
+ for (int i = 0; i < numUnions; i++) {
+ unions.add(_setOperationBuilder.buildUnion());
}
- return unionMap;
+ return unions;
}
+ /**
+ * Deserializes the sketches from the bytes.
+ */
private Sketch[] deserializeSketches(byte[][] serializedSketches, int length) {
Sketch[] sketches = new Sketch[length];
for (int i = 0; i < length; i++) {
@@ -533,174 +1247,197 @@ public class DistinctCountThetaSketchAggregationFunction implements AggregationF
}
/**
- * Evaluates the theta-sketch post-aggregation expression, which is composed by performing AND/OR on top of the
- * pre-defined predicates. These predicates are evaluated during the aggregation phase, and the cached results are
- * passed to this method to be used when evaluating the expression.
- *
- * @param postAggregationExpression Post-aggregation expression to evaluate (modeled as a filter)
- * @param sketchMap Precomputed sketches for predicates that are part of the expression.
- * @return Overall evaluated sketch for the expression.
+ * Converts the given Unions to Sketches.
*/
- private Sketch evalPostAggregationExpression(ExpressionContext postAggregationExpression,
- Map<Predicate, Sketch> sketchMap) {
- if (postAggregationExpression.getType() == ExpressionContext.Type.LITERAL) {
- throw new IllegalArgumentException("Literal not supported in post-aggregation function");
+ private List<Sketch> convertToSketches(List<Union> unions) {
+ int numUnions = unions.size();
+ List<Sketch> sketches = new ArrayList<>(numUnions);
+ for (Union union : unions) {
+ // NOTE: Compact the sketch in unsorted, on-heap fashion for performance concern.
+ // See https://datasketches.apache.org/docs/Theta/ThetaSize.html for more details.
+ sketches.add(union.getResult(false, null));
}
+ return sketches;
+ }
- if (postAggregationExpression.getType() == ExpressionContext.Type.IDENTIFIER) {
- final Predicate exp = _predicates.get(extractSubstitutionPosition(postAggregationExpression.getLiteral()) - 1);
- return sketchMap.get(exp);
- }
+ /**
+ * Evaluates the post-aggregation expression.
+ */
+ protected Sketch evaluatePostAggregationExpression(List<Sketch> sketches) {
+ return evaluatePostAggregationExpression(_postAggregationExpression, sketches);
+ }
- // shouldn't throw exception because of the validation in the constructor
- MergeFunction func = MergeFunction.valueOf(postAggregationExpression.getFunction().getFunctionName().toUpperCase());
+ /**
+ * Evaluates the post-aggregation expression.
+ */
+ private Sketch evaluatePostAggregationExpression(ExpressionContext expression, List<Sketch> sketches) {
+ if (expression.getType() == ExpressionContext.Type.IDENTIFIER) {
+ return sketches.get(extractSketchId(expression.getIdentifier()));
+ }
- // handle functions recursively
- switch (func) {
+ // NOTE: Compact the sketch in unsorted, on-heap fashion for performance concern.
+ // See https://datasketches.apache.org/docs/Theta/ThetaSize.html for more details.
+ FunctionContext function = expression.getFunction();
+ String functionName = function.getFunctionName();
+ List<ExpressionContext> arguments = function.getArguments();
+ switch (functionName.toUpperCase()) {
case SET_UNION:
Union union = _setOperationBuilder.buildUnion();
- for (ExpressionContext exp : postAggregationExpression.getFunction().getArguments()) {
- union.update(evalPostAggregationExpression(exp, sketchMap));
+ for (ExpressionContext argument : arguments) {
+ union.update(evaluatePostAggregationExpression(argument, sketches));
}
- return union.getResult();
+ return union.getResult(false, null);
case SET_INTERSECT:
Intersection intersection = _setOperationBuilder.buildIntersection();
- for (ExpressionContext exp : postAggregationExpression.getFunction().getArguments()) {
- intersection.update(evalPostAggregationExpression(exp, sketchMap));
+ for (ExpressionContext argument : arguments) {
+ intersection.update(evaluatePostAggregationExpression(argument, sketches));
}
- return intersection.getResult();
+ return intersection.getResult(false, null);
case SET_DIFF:
- List<ExpressionContext> args = postAggregationExpression.getFunction().getArguments();
AnotB diff = _setOperationBuilder.buildANotB();
- Sketch a = evalPostAggregationExpression(args.get(0), sketchMap);
- Sketch b = evalPostAggregationExpression(args.get(1), sketchMap);
- diff.update(a, b);
- return diff.getResult();
+ diff.update(evaluatePostAggregationExpression(arguments.get(0), sketches),
+ evaluatePostAggregationExpression(arguments.get(1), sketches));
+ return diff.getResult(false, null);
default:
- throw new IllegalStateException(String.format("Invalid post-aggregation function: %s",
- postAggregationExpression.getFunction().getFunctionName().toUpperCase()));
+ throw new IllegalStateException();
}
}
/**
- * Extracts the final sketch from the intermediate result by applying the post-aggregation expression on it.
- *
- * @param intermediateResult Intermediate result
- * @return Final Sketch obtained by computing the post-aggregation expression on intermediate result
+ * Helper class to wrap the theta-sketch parameters.
*/
- protected Sketch extractFinalSketch(Map<String, Sketch> intermediateResult) {
- Map<Predicate, Sketch> sketchMap = new HashMap<>();
- for (Map.Entry<String, Sketch> entry : intermediateResult.entrySet()) {
- Predicate predicate = getPredicate(entry.getKey());
- sketchMap.put(predicate, entry.getValue());
+ private static class Parameters {
+ private static final char PARAMETER_DELIMITER = ';';
+ private static final char PARAMETER_KEY_VALUE_SEPARATOR = '=';
+ private static final String NOMINAL_ENTRIES_KEY = "nominalEntries";
+
+ private int _nominalEntries = Util.DEFAULT_NOMINAL_ENTRIES;
+
+ Parameters(String parametersString) {
+ StringUtils.deleteWhitespace(parametersString);
+ String[] keyValuePairs = StringUtils.split(parametersString, PARAMETER_DELIMITER);
+ for (String keyValuePair : keyValuePairs) {
+ String[] keyAndValue = StringUtils.split(keyValuePair, PARAMETER_KEY_VALUE_SEPARATOR);
+ Preconditions.checkArgument(keyAndValue.length == 2, "Invalid parameter: %s", keyValuePair);
+ String key = keyAndValue[0];
+ String value = keyAndValue[1];
+ if (key.equalsIgnoreCase(NOMINAL_ENTRIES_KEY)) {
+ _nominalEntries = Integer.parseInt(value);
+ } else {
+ throw new IllegalArgumentException("Invalid parameter key: " + key);
+ }
+ }
+ }
+
+ int getNominalEntries() {
+ return _nominalEntries;
}
- return evalPostAggregationExpression(_postAggregationExpression, sketchMap);
}
/**
- * Validates that the function context's substitution parameters ($1, $2, etc) does not exceed the number
- * of predicates passed into the post-aggregation function.
- *
- * For example, if the post aggregation function is:
- * INTERSECT($1, $2, $3)
- * But there are only 2 arguments passed into the aggregation function, throw an error
- *
- * SET_DIFF should contain exactly 2 arguments - throw an error otherwise.
- * SET_UNION and SET_INTERSECT should contain 2 or more arguments - throw an error otherwise.
- *
- * @param context The parsed function context that's a tree structure
- * @param numPredicates Max number of predicates available to be substituted
+ * Helper interface to evaluate the filter on the values.
*/
- private static void validatePostAggregationExpression(ExpressionContext context, int numPredicates) {
- if (context.getType() == ExpressionContext.Type.LITERAL) {
- throw new IllegalArgumentException("Invalid post-aggregation function expression syntax.");
- }
+ private interface FilterEvaluator {
- if (context.getType() == ExpressionContext.Type.IDENTIFIER) {
- int id = extractSubstitutionPosition(context.getIdentifier());
- if (id <= 0) {
- throw new IllegalArgumentException("Argument substitution starts at $1");
- }
- if (id > numPredicates) {
- throw new IllegalArgumentException("Argument substitution exceeded number of predicates");
- }
- // if none of the invalid conditions are met above, exit out early
- return;
- }
+ /**
+ * Evaluates the given values with the filter, returns {@code true} if the values pass the filter, {@code false}
+ * otherwise.
+ */
+ boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] valueArrays, int index);
+ }
+
+ private static class AndFilterEvaluator implements FilterEvaluator {
+ final List<FilterEvaluator> _children;
- if (!MergeFunction.isValid(context.getFunction().getFunctionName())) {
- throw new IllegalArgumentException(
- String.format("Invalid Theta Sketch aggregation function. Allowed: [%s]", MergeFunction.CSV_VALUES));
+ private AndFilterEvaluator(List<FilterEvaluator> children) {
+ _children = children;
}
- switch (MergeFunction.valueOf(context.getFunction().getFunctionName().toUpperCase())) {
- case SET_DIFF:
- // set diff can only have 2 arguments
- if (context.getFunction().getArguments().size() != 2) {
- throw new IllegalArgumentException("SET_DIFF function can only have 2 arguments.");
+ @Override
+ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] valueArrays, int index) {
+ for (FilterEvaluator child : _children) {
+ if (!child.evaluate(singleValues, valueTypes, valueArrays, index)) {
+ return false;
}
- validatePostAggregationExpression(context.getFunction().getArguments().get(0), numPredicates);
- validatePostAggregationExpression(context.getFunction().getArguments().get(1), numPredicates);
- break;
- case SET_UNION:
- case SET_INTERSECT:
- if (context.getFunction().getArguments().size() < 2) {
- throw new IllegalArgumentException("SET_UNION and SET_INTERSECT should have at least 2 arguments.");
- }
- for (ExpressionContext arg : context.getFunction().getArguments()) {
- validatePostAggregationExpression(arg, numPredicates);
- }
- break;
- default:
- throw new IllegalStateException("Invalid merge function");
+ }
+ return true;
}
}
- private static int extractSubstitutionPosition(String str) {
- Matcher matcher = ARGUMENT_SUBSTITUTION.matcher(str);
- if (matcher.find() && matcher.groupCount() == 1) {
- return Integer.parseInt(matcher.group(1));
- } else {
- throw new IllegalArgumentException(
- String.format("Invalid argument substitution: [%s]. Use $1, $2, ... (starting from $1)", str));
+ private static class OrFilterEvaluator implements FilterEvaluator {
+ final List<FilterEvaluator> _children;
+
+ private OrFilterEvaluator(List<FilterEvaluator> children) {
+ _children = children;
+ }
+
+ @Override
+ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] valueArrays, int index) {
+ for (FilterEvaluator child : _children) {
+ if (child.evaluate(singleValues, valueTypes, valueArrays, index)) {
+ return true;
+ }
+ }
+ return false;
}
}
- /**
- * Helper class to store predicate related information:
- * <ul>
- * <li>Predicate</li>
- * <li>String representation of the predicate</li>
- * <li>Predicate evaluator</li>
- * </ul>
- */
- private static class PredicateInfo {
+ private static class PredicateFilterEvaluator implements FilterEvaluator {
final Predicate _predicate;
- final String _stringPredicate;
+ final int _expressionIndex;
PredicateEvaluator _predicateEvaluator;
- PredicateInfo(Predicate predicate) {
+ private PredicateFilterEvaluator(Predicate predicate, int expressionIndex) {
_predicate = predicate;
- _stringPredicate = predicate.toString();
- _predicateEvaluator = null; // Initialized lazily
- }
-
- Predicate getPredicate() {
- return _predicate;
+ _expressionIndex = expressionIndex;
}
- String getStringPredicate() {
- return _stringPredicate;
- }
-
- /**
- * Since PredicateEvaluator requires data-type, it is initialized lazily.
- */
- PredicateEvaluator getPredicateEvaluator(FieldSpec.DataType dataType) {
+ @Override
+ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] valueArrays, int index) {
+ boolean singleValue = singleValues[_expressionIndex];
+ DataType valueType = valueTypes[_expressionIndex];
+ Object valueArray = valueArrays[_expressionIndex];
if (_predicateEvaluator == null) {
- _predicateEvaluator = PredicateEvaluatorProvider.getPredicateEvaluator(_predicate, null, dataType);
+ _predicateEvaluator = PredicateEvaluatorProvider.getPredicateEvaluator(_predicate, null, valueType);
+ }
+ if (singleValue) {
+ switch (valueType) {
+ case INT:
+ return _predicateEvaluator.applySV(((int[]) valueArray)[index]);
+ case LONG:
+ return _predicateEvaluator.applySV(((long[]) valueArray)[index]);
+ case FLOAT:
+ return _predicateEvaluator.applySV(((float[]) valueArray)[index]);
+ case DOUBLE:
+ return _predicateEvaluator.applySV(((double[]) valueArray)[index]);
+ case STRING:
+ return _predicateEvaluator.applySV(((String[]) valueArray)[index]);
+ case BYTES:
+ return _predicateEvaluator.applySV(((byte[][]) valueArray)[index]);
+ default:
+ throw new IllegalStateException();
+ }
+ } else {
+ switch (valueType) {
+ case INT:
+ int[] intValues = ((int[][]) valueArray)[index];
+ return _predicateEvaluator.applyMV(intValues, intValues.length);
+ case LONG:
+ long[] longValues = ((long[][]) valueArray)[index];
+ return _predicateEvaluator.applyMV(longValues, longValues.length);
+ case FLOAT:
+ float[] floatValues = ((float[][]) valueArray)[index];
+ return _predicateEvaluator.applyMV(floatValues, floatValues.length);
+ case DOUBLE:
+ double[] doubleValues = ((double[][]) valueArray)[index];
+ return _predicateEvaluator.applyMV(doubleValues, doubleValues.length);
+ case STRING:
+ String[] stringValues = ((String[][]) valueArray)[index];
+ return _predicateEvaluator.applyMV(stringValues, stringValues.length);
+ default:
+ throw new IllegalStateException();
+ }
}
- return _predicateEvaluator;
}
}
}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/RawThetaSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/RawThetaSketchAggregationFunction.java
deleted file mode 100644
index 94c952d..0000000
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/RawThetaSketchAggregationFunction.java
+++ /dev/null
@@ -1,563 +0,0 @@
-/**
- * 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.pinot.core.query.aggregation.function;
-
-import com.google.common.base.Preconditions;
-import java.util.Base64;
-import java.util.List;
-import java.util.Map;
-import org.apache.commons.lang3.StringUtils;
-import org.apache.datasketches.Util;
-import org.apache.datasketches.memory.Memory;
-import org.apache.datasketches.theta.SetOperationBuilder;
-import org.apache.datasketches.theta.Sketch;
-import org.apache.datasketches.theta.Union;
-import org.apache.datasketches.theta.UpdateSketch;
-import org.apache.datasketches.theta.UpdateSketchBuilder;
-import org.apache.pinot.common.function.AggregationFunctionType;
-import org.apache.pinot.common.utils.DataSchema.ColumnDataType;
-import org.apache.pinot.core.common.BlockValSet;
-import org.apache.pinot.core.query.aggregation.AggregationResultHolder;
-import org.apache.pinot.core.query.aggregation.ObjectAggregationResultHolder;
-import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder;
-import org.apache.pinot.core.query.aggregation.groupby.ObjectGroupByResultHolder;
-import org.apache.pinot.core.query.request.context.ExpressionContext;
-import org.apache.pinot.spi.data.FieldSpec;
-
-
-/**
- * The {@code RawThetaSketchAggregationFunction} collects the values for a given expression (can be single-valued or
- * multi-valued) into a {@link Sketch} object, and returns the sketch as a base64 encoded string. It treats BYTES
- * expression as serialized sketches.
- * <p>The function takes an optional second argument as the parameters for the function. Currently there is only 1
- * parameter for the function:
- * <ul>
- * <li>
- * nominalEntries: The nominal entries used to create the sketch. (Default 4096)
- * </li>
- * </ul>
- * <p>Example: RAW_THETA_SKETCH(col, 'nominalEntries=8192')
- */
-public class RawThetaSketchAggregationFunction extends BaseSingleInputAggregationFunction<Sketch, String> {
- private final UpdateSketchBuilder _updateSketchBuilder = new UpdateSketchBuilder();
- private final SetOperationBuilder _setOperationBuilder = new SetOperationBuilder();
-
- public RawThetaSketchAggregationFunction(List<ExpressionContext> arguments) {
- super(arguments.get(0));
-
- // Optional second argument for theta-sketch parameters
- if (arguments.size() > 1) {
- ExpressionContext paramsExpression = arguments.get(1);
- Preconditions.checkArgument(paramsExpression.getType() == ExpressionContext.Type.LITERAL,
- "Second argument of RAW_THETA_SKETCH aggregation function must be literal (parameters)");
- Parameters parameters = new Parameters(paramsExpression.getLiteral());
- int nominalEntries = parameters.getNominalEntries();
- _updateSketchBuilder.setNominalEntries(nominalEntries);
- _setOperationBuilder.setNominalEntries(nominalEntries);
- }
- }
-
- @Override
- public AggregationFunctionType getType() {
- return AggregationFunctionType.RAWTHETASKETCH;
- }
-
- @Override
- public AggregationResultHolder createAggregationResultHolder() {
- return new ObjectAggregationResultHolder();
- }
-
- @Override
- public GroupByResultHolder createGroupByResultHolder(int initialCapacity, int maxCapacity) {
- return new ObjectGroupByResultHolder(initialCapacity, maxCapacity);
- }
-
- @Override
- public void aggregate(int length, AggregationResultHolder aggregationResultHolder,
- Map<ExpressionContext, BlockValSet> blockValSetMap) {
- BlockValSet blockValSet = blockValSetMap.get(_expression);
- FieldSpec.DataType valueType = blockValSet.getValueType();
-
- if (valueType != FieldSpec.DataType.BYTES) {
- UpdateSketch updateSketch = getUpdateSketch(aggregationResultHolder);
- if (blockValSet.isSingleValue()) {
- switch (valueType) {
- case INT:
- int[] intValues = blockValSet.getIntValuesSV();
- for (int i = 0; i < length; i++) {
- updateSketch.update(intValues[i]);
- }
- break;
- case LONG:
- long[] longValues = blockValSet.getLongValuesSV();
- for (int i = 0; i < length; i++) {
- updateSketch.update(longValues[i]);
- }
- break;
- case FLOAT:
- float[] floatValues = blockValSet.getFloatValuesSV();
- for (int i = 0; i < length; i++) {
- updateSketch.update(floatValues[i]);
- }
- break;
- case DOUBLE:
- double[] doubleValues = blockValSet.getDoubleValuesSV();
- for (int i = 0; i < length; i++) {
- updateSketch.update(doubleValues[i]);
- }
- break;
- case STRING:
- String[] stringValues = blockValSet.getStringValuesSV();
- for (int i = 0; i < length; i++) {
- updateSketch.update(stringValues[i]);
- }
- break;
- default:
- throw new IllegalStateException(
- "Illegal single-value data type for RAW_THETA_SKETCH aggregation function: " + valueType);
- }
- } else {
- switch (valueType) {
- case INT:
- int[][] intValues = blockValSet.getIntValuesMV();
- for (int i = 0; i < length; i++) {
- for (int value : intValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case LONG:
- long[][] longValues = blockValSet.getLongValuesMV();
- for (int i = 0; i < length; i++) {
- for (long value : longValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case FLOAT:
- float[][] floatValues = blockValSet.getFloatValuesMV();
- for (int i = 0; i < length; i++) {
- for (float value : floatValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case DOUBLE:
- double[][] doubleValues = blockValSet.getDoubleValuesMV();
- for (int i = 0; i < length; i++) {
- for (double value : doubleValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case STRING:
- String[][] stringValues = blockValSet.getStringValuesMV();
- for (int i = 0; i < length; i++) {
- for (String value : stringValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- default:
- throw new IllegalStateException(
- "Illegal multi-value data type for RAW_THETA_SKETCH aggregation function: " + valueType);
- }
- }
- } else {
- // Serialized sketch
- Union union = getUnion(aggregationResultHolder);
- byte[][] bytesValues = blockValSet.getBytesValuesSV();
- for (int i = 0; i < length; i++) {
- union.update(Sketch.wrap(Memory.wrap(bytesValues[i])));
- }
- }
- }
-
- @Override
- public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder,
- Map<ExpressionContext, BlockValSet> blockValSetMap) {
- BlockValSet blockValSet = blockValSetMap.get(_expression);
- FieldSpec.DataType valueType = blockValSet.getValueType();
-
- if (blockValSet.isSingleValue()) {
- switch (valueType) {
- case INT:
- int[] intValues = blockValSet.getIntValuesSV();
- for (int i = 0; i < length; i++) {
- getUpdateSketch(groupByResultHolder, groupKeyArray[i]).update(intValues[i]);
- }
- break;
- case LONG:
- long[] longValues = blockValSet.getLongValuesSV();
- for (int i = 0; i < length; i++) {
- getUpdateSketch(groupByResultHolder, groupKeyArray[i]).update(longValues[i]);
- }
- break;
- case FLOAT:
- float[] floatValues = blockValSet.getFloatValuesSV();
- for (int i = 0; i < length; i++) {
- getUpdateSketch(groupByResultHolder, groupKeyArray[i]).update(floatValues[i]);
- }
- break;
- case DOUBLE:
- double[] doubleValues = blockValSet.getDoubleValuesSV();
- for (int i = 0; i < length; i++) {
- getUpdateSketch(groupByResultHolder, groupKeyArray[i]).update(doubleValues[i]);
- }
- break;
- case STRING:
- String[] stringValues = blockValSet.getStringValuesSV();
- for (int i = 0; i < length; i++) {
- getUpdateSketch(groupByResultHolder, groupKeyArray[i]).update(stringValues[i]);
- }
- break;
- case BYTES:
- // Serialized sketch
- byte[][] bytesValues = blockValSet.getBytesValuesSV();
- for (int i = 0; i < length; i++) {
- getUnion(groupByResultHolder, groupKeyArray[i]).update(Sketch.wrap(Memory.wrap(bytesValues[i])));
- }
- break;
- default:
- throw new IllegalStateException(
- "Illegal single-value data type for RAW_THETA_SKETCH aggregation function: " + valueType);
- }
- } else {
- switch (valueType) {
- case INT:
- int[][] intValues = blockValSet.getIntValuesMV();
- for (int i = 0; i < length; i++) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKeyArray[i]);
- for (int value : intValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case LONG:
- long[][] longValues = blockValSet.getLongValuesMV();
- for (int i = 0; i < length; i++) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKeyArray[i]);
- for (long value : longValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case FLOAT:
- float[][] floatValues = blockValSet.getFloatValuesMV();
- for (int i = 0; i < length; i++) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKeyArray[i]);
- for (float value : floatValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case DOUBLE:
- double[][] doubleValues = blockValSet.getDoubleValuesMV();
- for (int i = 0; i < length; i++) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKeyArray[i]);
- for (double value : doubleValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- case STRING:
- String[][] stringValues = blockValSet.getStringValuesMV();
- for (int i = 0; i < length; i++) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKeyArray[i]);
- for (String value : stringValues[i]) {
- updateSketch.update(value);
- }
- }
- break;
- default:
- throw new IllegalStateException(
- "Illegal multi-value data type for RAW_THETA_SKETCH aggregation function: " + valueType);
- }
- }
- }
-
- @Override
- public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder,
- Map<ExpressionContext, BlockValSet> blockValSetMap) {
- BlockValSet blockValSet = blockValSetMap.get(_expression);
- FieldSpec.DataType valueType = blockValSet.getValueType();
-
- if (blockValSet.isSingleValue()) {
- switch (valueType) {
- case INT:
- int[] intValues = blockValSet.getIntValuesSV();
- for (int i = 0; i < length; i++) {
- int value = intValues[i];
- for (int groupKey : groupKeysArray[i]) {
- getUpdateSketch(groupByResultHolder, groupKey).update(value);
- }
- }
- break;
- case LONG:
- long[] longValues = blockValSet.getLongValuesSV();
- for (int i = 0; i < length; i++) {
- long value = longValues[i];
- for (int groupKey : groupKeysArray[i]) {
- getUpdateSketch(groupByResultHolder, groupKey).update(value);
- }
- }
- break;
- case FLOAT:
- float[] floatValues = blockValSet.getFloatValuesSV();
- for (int i = 0; i < length; i++) {
- float value = floatValues[i];
- for (int groupKey : groupKeysArray[i]) {
- getUpdateSketch(groupByResultHolder, groupKey).update(value);
- }
- }
- break;
- case DOUBLE:
- double[] doubleValues = blockValSet.getDoubleValuesSV();
- for (int i = 0; i < length; i++) {
- double value = doubleValues[i];
- for (int groupKey : groupKeysArray[i]) {
- getUpdateSketch(groupByResultHolder, groupKey).update(value);
- }
- }
- break;
- case STRING:
- String[] stringValues = blockValSet.getStringValuesSV();
- for (int i = 0; i < length; i++) {
- String value = stringValues[i];
- for (int groupKey : groupKeysArray[i]) {
- getUpdateSketch(groupByResultHolder, groupKey).update(value);
- }
- }
- break;
- case BYTES:
- // Serialized sketch
- byte[][] bytesValues = blockValSet.getBytesValuesSV();
- for (int i = 0; i < length; i++) {
- Sketch sketch = Sketch.wrap(Memory.wrap(bytesValues[i]));
- for (int groupKey : groupKeysArray[i]) {
- getUnion(groupByResultHolder, groupKey).update(sketch);
- }
- }
- break;
- default:
- throw new IllegalStateException(
- "Illegal single-value data type for RAW_THETA_SKETCH aggregation function: " + valueType);
- }
- } else {
- switch (valueType) {
- case INT:
- int[][] intValues = blockValSet.getIntValuesMV();
- for (int i = 0; i < length; i++) {
- int[] values = intValues[i];
- for (int groupKey : groupKeysArray[i]) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKey);
- for (int value : values) {
- updateSketch.update(value);
- }
- }
- }
- break;
- case LONG:
- long[][] longValues = blockValSet.getLongValuesMV();
- for (int i = 0; i < length; i++) {
- long[] values = longValues[i];
- for (int groupKey : groupKeysArray[i]) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKey);
- for (long value : values) {
- updateSketch.update(value);
- }
- }
- }
- break;
- case FLOAT:
- float[][] floatValues = blockValSet.getFloatValuesMV();
- for (int i = 0; i < length; i++) {
- float[] values = floatValues[i];
- for (int groupKey : groupKeysArray[i]) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKey);
- for (float value : values) {
- updateSketch.update(value);
- }
- }
- }
- break;
- case DOUBLE:
- double[][] doubleValues = blockValSet.getDoubleValuesMV();
- for (int i = 0; i < length; i++) {
- double[] values = doubleValues[i];
- for (int groupKey : groupKeysArray[i]) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKey);
- for (double value : values) {
- updateSketch.update(value);
- }
- }
- }
- break;
- case STRING:
- String[][] stringValues = blockValSet.getStringValuesMV();
- for (int i = 0; i < length; i++) {
- String[] values = stringValues[i];
- for (int groupKey : groupKeysArray[i]) {
- UpdateSketch updateSketch = getUpdateSketch(groupByResultHolder, groupKey);
- for (String value : values) {
- updateSketch.update(value);
- }
- }
- }
- break;
- default:
- throw new IllegalStateException(
- "Illegal multi-value data type for RAW_THETA_SKETCH aggregation function: " + valueType);
- }
- }
- }
-
- @Override
- public Sketch extractAggregationResult(AggregationResultHolder aggregationResultHolder) {
- Object result = aggregationResultHolder.getResult();
- if (result == null) {
- return _updateSketchBuilder.build();
- } else {
- if (result instanceof Sketch) {
- return (Sketch) result;
- } else {
- assert result instanceof Union;
- return ((Union) result).getResult();
- }
- }
- }
-
- @Override
- public Sketch extractGroupByResult(GroupByResultHolder groupByResultHolder, int groupKey) {
- Object result = groupByResultHolder.getResult(groupKey);
- if (result instanceof Sketch) {
- return (Sketch) result;
- } else {
- assert result instanceof Union;
- return ((Union) result).getResult();
- }
- }
-
- @Override
- public Sketch merge(Sketch sketch1, Sketch sketch2) {
- Union union = _setOperationBuilder.buildUnion();
- union.update(sketch1);
- union.update(sketch2);
- return union.getResult();
- }
-
- @Override
- public boolean isIntermediateResultComparable() {
- return false;
- }
-
- @Override
- public ColumnDataType getIntermediateResultColumnType() {
- return ColumnDataType.OBJECT;
- }
-
- @Override
- public ColumnDataType getFinalResultColumnType() {
- return ColumnDataType.STRING;
- }
-
- @Override
- public String extractFinalResult(Sketch sketch) {
- return Base64.getEncoder().encodeToString(sketch.compact().toByteArray());
- }
-
- /**
- * Returns the UpdateSketch from the result holder or creates a new one if it does not exist.
- */
- private UpdateSketch getUpdateSketch(AggregationResultHolder aggregationResultHolder) {
- UpdateSketch updateSketch = aggregationResultHolder.getResult();
- if (updateSketch == null) {
- updateSketch = _updateSketchBuilder.build();
- aggregationResultHolder.setValue(updateSketch);
- }
- return updateSketch;
- }
-
- /**
- * Returns the Union from the result holder or creates a new one if it does not exist.
- */
- private Union getUnion(AggregationResultHolder aggregationResultHolder) {
- Union union = aggregationResultHolder.getResult();
- if (union == null) {
- union = _setOperationBuilder.buildUnion();
- aggregationResultHolder.setValue(union);
- }
- return union;
- }
-
- /**
- * Returns the UpdateSketch for the given group key or creates a new one if it does not exist.
- */
- private UpdateSketch getUpdateSketch(GroupByResultHolder groupByResultHolder, int groupKey) {
- UpdateSketch updateSketch = groupByResultHolder.getResult(groupKey);
- if (updateSketch == null) {
- updateSketch = _updateSketchBuilder.build();
- groupByResultHolder.setValueForKey(groupKey, updateSketch);
- }
- return updateSketch;
- }
-
- /**
- * Returns the UpdateSketch for the given group key or creates a new one if it does not exist.
- */
- private Union getUnion(GroupByResultHolder groupByResultHolder, int groupKey) {
- Union union = groupByResultHolder.getResult(groupKey);
- if (union == null) {
- union = _setOperationBuilder.buildUnion();
- groupByResultHolder.setValueForKey(groupKey, union);
- }
- return union;
- }
-
- /**
- * Helper class to wrap the theta-sketch parameters.
- */
- static class Parameters {
- private static final char PARAMETER_DELIMITER = ';';
- private static final char PARAMETER_KEY_VALUE_SEPARATOR = '=';
- private static final String NOMINAL_ENTRIES_KEY = "nominalEntries";
-
- private int _nominalEntries = Util.DEFAULT_NOMINAL_ENTRIES;
-
- Parameters(String parametersString) {
- StringUtils.deleteWhitespace(parametersString);
- String[] keyValuePairs = StringUtils.split(parametersString, PARAMETER_DELIMITER);
- for (String keyValuePair : keyValuePairs) {
- String[] keyAndValue = StringUtils.split(keyValuePair, PARAMETER_KEY_VALUE_SEPARATOR);
- Preconditions.checkArgument(keyAndValue.length == 2, "Invalid parameter: %s", keyValuePair);
- String key = keyAndValue[0];
- String value = keyAndValue[1];
- if (key.equalsIgnoreCase(NOMINAL_ENTRIES_KEY)) {
- _nominalEntries = Integer.parseInt(value);
- } else {
- throw new IllegalArgumentException("Invalid parameter key: " + key);
- }
- }
- }
-
- int getNominalEntries() {
- return _nominalEntries;
- }
- }
-}
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverter.java b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverter.java
index 57beb3c..af57353 100644
--- a/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverter.java
+++ b/pinot-core/src/main/java/org/apache/pinot/core/query/request/context/utils/BrokerRequestToQueryContextConverter.java
@@ -26,6 +26,7 @@ import java.util.List;
import java.util.Map;
import java.util.Set;
import org.apache.commons.collections.CollectionUtils;
+import org.apache.commons.lang3.StringUtils;
import org.apache.pinot.common.function.AggregationFunctionType;
import org.apache.pinot.common.request.AggregationInfo;
import org.apache.pinot.common.request.BrokerRequest;
@@ -114,25 +115,23 @@ public class BrokerRequestToQueryContextConverter {
List<AggregationInfo> aggregationsInfo = brokerRequest.getAggregationsInfo();
selectExpressions = new ArrayList<>(aggregationsInfo.size());
for (AggregationInfo aggregationInfo : aggregationsInfo) {
- String functionName = aggregationInfo.getAggregationType();
+ String functionName = StringUtils.remove(aggregationInfo.getAggregationType(), '_');
List<String> stringExpressions = aggregationInfo.getExpressions();
int numArguments = stringExpressions.size();
List<ExpressionContext> arguments = new ArrayList<>(numArguments);
- if (functionName.equalsIgnoreCase(AggregationFunctionType.DISTINCTCOUNTTHETASKETCH.getName()) || functionName
- .equalsIgnoreCase(AggregationFunctionType.DISTINCTCOUNTRAWTHETASKETCH.getName()) || functionName
- .equalsIgnoreCase(AggregationFunctionType.RAWTHETASKETCH.name()) || functionName
- .equalsIgnoreCase(AggregationFunctionType.IDSET.getName())) {
- // NOTE: For DistinctCountThetaSketch, DistinctCountRawThetaSketch, RawThetaSketch and IdSet, because of the
- // legacy behavior of PQL compiler treating string literal as identifier in aggregation, here we treat
- // all expressions except for the first one as string literal.
+ if (functionName.equalsIgnoreCase(AggregationFunctionType.DISTINCT.getName())) {
+ // For DISTINCT query, all arguments are expressions
+ for (String expression : stringExpressions) {
+ arguments.add(QueryContextConverterUtils.getExpression(expression));
+ }
+ } else {
+ // For non-DISTINCT query, only the first argument is expression, others are literals
+ // NOTE: We directly use the string as the literal value because of the legacy behavior of PQL compiler
+ // treating string literal as identifier in the aggregation function.
arguments.add(QueryContextConverterUtils.getExpression(stringExpressions.get(0)));
for (int i = 1; i < numArguments; i++) {
arguments.add(ExpressionContext.forLiteral(stringExpressions.get(i)));
}
- } else {
- for (String expression : stringExpressions) {
- arguments.add(QueryContextConverterUtils.getExpression(expression));
- }
}
FunctionContext function = new FunctionContext(FunctionContext.Type.AGGREGATION, functionName, arguments);
selectExpressions.add(ExpressionContext.forFunction(function));
diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/DistinctCountThetaSketchQueriesTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/DistinctCountThetaSketchQueriesTest.java
new file mode 100644
index 0000000..6db80d6
--- /dev/null
+++ b/pinot-core/src/test/java/org/apache/pinot/queries/DistinctCountThetaSketchQueriesTest.java
@@ -0,0 +1,343 @@
+/**
+ * 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.pinot.queries;
+
+import java.io.File;
+import java.io.IOException;
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.Base64;
+import java.util.Iterator;
+import java.util.List;
+import org.apache.commons.io.FileUtils;
+import org.apache.datasketches.theta.Sketch;
+import org.apache.datasketches.theta.UpdateSketch;
+import org.apache.datasketches.theta.UpdateSketchBuilder;
+import org.apache.pinot.common.response.broker.AggregationResult;
+import org.apache.pinot.common.response.broker.BrokerResponseNative;
+import org.apache.pinot.common.response.broker.GroupByResult;
+import org.apache.pinot.common.segment.ReadMode;
+import org.apache.pinot.core.common.ObjectSerDeUtils;
+import org.apache.pinot.core.data.readers.GenericRowRecordReader;
+import org.apache.pinot.core.indexsegment.IndexSegment;
+import org.apache.pinot.core.indexsegment.generator.SegmentGeneratorConfig;
+import org.apache.pinot.core.indexsegment.immutable.ImmutableSegment;
+import org.apache.pinot.core.indexsegment.immutable.ImmutableSegmentLoader;
+import org.apache.pinot.core.operator.blocks.IntermediateResultsBlock;
+import org.apache.pinot.core.operator.query.AggregationGroupByOperator;
+import org.apache.pinot.core.operator.query.AggregationOperator;
+import org.apache.pinot.core.query.aggregation.groupby.AggregationGroupByResult;
+import org.apache.pinot.core.query.aggregation.groupby.GroupKeyGenerator;
+import org.apache.pinot.core.query.exception.BadQueryRequestException;
+import org.apache.pinot.core.segment.creator.impl.SegmentIndexCreationDriverImpl;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.data.FieldSpec.DataType;
+import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.data.readers.GenericRow;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.AfterClass;
+import org.testng.annotations.BeforeClass;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotNull;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.fail;
+
+
+/**
+ * Queries test for DISTINCT_COUNT_THETA_SKETCH queries.
+ */
+@SuppressWarnings("unchecked")
+public class DistinctCountThetaSketchQueriesTest extends BaseQueriesTest {
+ private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "DistinctCountThetaSketchQueriesTest");
+ private static final String RAW_TABLE_NAME = "testTable";
+ private static final String SEGMENT_NAME = "testSegment";
+
+ private static final int NUM_RECORDS = 1000;
+
+ private static final String INT_SV_COLUMN = "intSVColumn";
+ private static final String LONG_SV_COLUMN = "longSVColumn";
+ private static final String FLOAT_SV_COLUMN = "floatSVColumn";
+ private static final String DOUBLE_SV_COLUMN = "doubleSVColumn";
+ private static final String STRING_SV_COLUMN = "stringSVColumn";
+ private static final String INT_MV_COLUMN = "intMVColumn";
+ private static final String LONG_MV_COLUMN = "longMVColumn";
+ private static final String FLOAT_MV_COLUMN = "floatMVColumn";
+ private static final String DOUBLE_MV_COLUMN = "doubleMVColumn";
+ private static final String STRING_MV_COLUMN = "stringMVColumn";
+ private static final String BYTES_COLUMN = "bytesColumn";
+ private static final Schema SCHEMA = new Schema.SchemaBuilder().addSingleValueDimension(INT_SV_COLUMN, DataType.INT)
+ .addSingleValueDimension(LONG_SV_COLUMN, DataType.LONG).addSingleValueDimension(FLOAT_SV_COLUMN, DataType.FLOAT)
+ .addSingleValueDimension(DOUBLE_SV_COLUMN, DataType.DOUBLE)
+ .addSingleValueDimension(STRING_SV_COLUMN, DataType.STRING).addMultiValueDimension(INT_MV_COLUMN, DataType.INT)
+ .addMultiValueDimension(LONG_MV_COLUMN, DataType.LONG).addMultiValueDimension(FLOAT_MV_COLUMN, DataType.FLOAT)
+ .addMultiValueDimension(DOUBLE_MV_COLUMN, DataType.DOUBLE)
+ .addMultiValueDimension(STRING_MV_COLUMN, DataType.STRING).addMetric(BYTES_COLUMN, DataType.BYTES).build();
+ private static final TableConfig TABLE_CONFIG =
+ new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).build();
+
+ private IndexSegment _indexSegment;
+ private List<IndexSegment> _indexSegments;
+
+ @Override
+ protected String getFilter() {
+ return "";
+ }
+
+ @Override
+ protected IndexSegment getIndexSegment() {
+ return _indexSegment;
+ }
+
+ @Override
+ protected List<IndexSegment> getIndexSegments() {
+ return _indexSegments;
+ }
+
+ @BeforeClass
+ public void setUp()
+ throws Exception {
+ FileUtils.deleteDirectory(INDEX_DIR);
+ UpdateSketchBuilder sketchBuilder = new UpdateSketchBuilder();
+
+ List<GenericRow> records = new ArrayList<>(NUM_RECORDS);
+ for (int i = 0; i < NUM_RECORDS; i++) {
+ GenericRow record = new GenericRow();
+ record.putValue(INT_SV_COLUMN, i);
+ record.putValue(LONG_SV_COLUMN, i);
+ record.putValue(FLOAT_SV_COLUMN, i);
+ record.putValue(DOUBLE_SV_COLUMN, i);
+ record.putValue(STRING_SV_COLUMN, i);
+ Integer[] mvEntry = new Integer[]{i, i + NUM_RECORDS, i + 2 * NUM_RECORDS};
+ record.putValue(INT_MV_COLUMN, mvEntry);
+ record.putValue(LONG_MV_COLUMN, mvEntry);
+ record.putValue(FLOAT_MV_COLUMN, mvEntry);
+ record.putValue(DOUBLE_MV_COLUMN, mvEntry);
+ record.putValue(STRING_MV_COLUMN, mvEntry);
+ // Store serialized sketches in the BYTES column
+ UpdateSketch sketch = sketchBuilder.build();
+ sketch.update(i);
+ sketch.update(i + NUM_RECORDS);
+ sketch.update(i + 2 * NUM_RECORDS);
+ record.putValue(BYTES_COLUMN, sketch.compact().toByteArray());
+ records.add(record);
+ }
+
+ SegmentGeneratorConfig segmentGeneratorConfig = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA);
+ segmentGeneratorConfig.setTableName(RAW_TABLE_NAME);
+ segmentGeneratorConfig.setSegmentName(SEGMENT_NAME);
+ segmentGeneratorConfig.setOutDir(INDEX_DIR.getPath());
+
+ SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
+ driver.init(segmentGeneratorConfig, new GenericRowRecordReader(records));
+ driver.build();
+
+ ImmutableSegment immutableSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), ReadMode.mmap);
+ _indexSegment = immutableSegment;
+ _indexSegments = Arrays.asList(immutableSegment, immutableSegment);
+ }
+
+ @Test
+ public void testAggregationOnly() {
+ String query = "SELECT DISTINCT_COUNT_THETA_SKETCH(intSVColumn), DISTINCT_COUNT_THETA_SKETCH(longSVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(floatSVColumn), DISTINCT_COUNT_THETA_SKETCH(doubleSVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(stringSVColumn), DISTINCT_COUNT_THETA_SKETCH(intMVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(longMVColumn), DISTINCT_COUNT_THETA_SKETCH(floatMVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(doubleMVColumn), DISTINCT_COUNT_THETA_SKETCH(stringMVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(bytesColumn) FROM testTable";
+
+ // Inner segment
+ AggregationOperator aggregationOperator = getOperatorForPqlQuery(query);
+ IntermediateResultsBlock resultsBlock = aggregationOperator.nextBlock();
+ QueriesTestUtils.testInnerSegmentExecutionStatistics(aggregationOperator.getExecutionStatistics(), NUM_RECORDS, 0,
+ 11 * NUM_RECORDS, NUM_RECORDS);
+ List<Object> aggregationResult = resultsBlock.getAggregationResult();
+ assertNotNull(aggregationResult);
+ assertEquals(aggregationResult.size(), 11);
+ for (int i = 0; i < 11; i++) {
+ List<Sketch> sketches = (List<Sketch>) aggregationResult.get(i);
+ assertEquals(sketches.size(), 1);
+ Sketch sketch = sketches.get(0);
+ if (i < 5) {
+ assertEquals(Math.round(sketch.getEstimate()), NUM_RECORDS);
+ } else {
+ assertEquals(Math.round(sketch.getEstimate()), 3 * NUM_RECORDS);
+ }
+ }
+
+ // Inter segments
+ String[] expectedResults = new String[11];
+ for (int i = 0; i < 11; i++) {
+ if (i < 5) {
+ expectedResults[i] = Integer.toString(NUM_RECORDS);
+ } else {
+ expectedResults[i] = Integer.toString(3 * NUM_RECORDS);
+ }
+ }
+ BrokerResponseNative brokerResponse = getBrokerResponseForPqlQuery(query);
+ QueriesTestUtils
+ .testInterSegmentAggregationResult(brokerResponse, 4 * NUM_RECORDS, 0, 4 * 11 * NUM_RECORDS, 4 * NUM_RECORDS,
+ expectedResults);
+ }
+
+ @Test
+ public void testAggregationGroupBy() {
+ String baseQuery = "SELECT DISTINCT_COUNT_THETA_SKETCH(intSVColumn), DISTINCT_COUNT_THETA_SKETCH(longSVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(floatSVColumn), DISTINCT_COUNT_THETA_SKETCH(doubleSVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(stringSVColumn), DISTINCT_COUNT_THETA_SKETCH(intMVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(longMVColumn), DISTINCT_COUNT_THETA_SKETCH(floatMVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(doubleMVColumn), DISTINCT_COUNT_THETA_SKETCH(stringMVColumn), "
+ + "DISTINCT_COUNT_THETA_SKETCH(bytesColumn) FROM testTable GROUP BY ";
+ for (boolean groupBySV : new boolean[]{true, false}) {
+ String query = baseQuery + (groupBySV ? "intSVColumn" : "intMVColumn");
+
+ // Inner segment
+ AggregationGroupByOperator aggregationGroupByOperator = getOperatorForPqlQuery(query);
+ IntermediateResultsBlock resultsBlock = aggregationGroupByOperator.nextBlock();
+ QueriesTestUtils
+ .testInnerSegmentExecutionStatistics(aggregationGroupByOperator.getExecutionStatistics(), NUM_RECORDS, 0,
+ 11 * NUM_RECORDS, NUM_RECORDS);
+ AggregationGroupByResult aggregationGroupByResult = resultsBlock.getAggregationGroupByResult();
+ assertNotNull(aggregationGroupByResult);
+ int numGroups = 0;
+ Iterator<GroupKeyGenerator.GroupKey> groupKeyIterator = aggregationGroupByResult.getGroupKeyIterator();
+ while (groupKeyIterator.hasNext()) {
+ numGroups++;
+ GroupKeyGenerator.GroupKey groupKey = groupKeyIterator.next();
+ for (int i = 0; i < 6; i++) {
+ List<Sketch> sketches = (List<Sketch>) aggregationGroupByResult.getResultForKey(groupKey, i);
+ assertEquals(sketches.size(), 1);
+ Sketch sketch = sketches.get(0);
+ if (i < 5) {
+ assertEquals(Math.round(sketch.getEstimate()), 1);
+ } else {
+ assertEquals(Math.round(sketch.getEstimate()), 3);
+ }
+ }
+ }
+ if (groupBySV) {
+ assertEquals(numGroups, NUM_RECORDS);
+ } else {
+ assertEquals(numGroups, 3 * NUM_RECORDS);
+ }
+
+ // Inter segments
+ BrokerResponseNative brokerResponse = getBrokerResponseForPqlQuery(query);
+ assertEquals(brokerResponse.getNumDocsScanned(), 4 * NUM_RECORDS);
+ assertEquals(brokerResponse.getNumEntriesScannedInFilter(), 0);
+ assertEquals(brokerResponse.getNumEntriesScannedPostFilter(), 4 * 11 * NUM_RECORDS);
+ assertEquals(brokerResponse.getTotalDocs(), 4 * NUM_RECORDS);
+ List<AggregationResult> aggregationResults = brokerResponse.getAggregationResults();
+ assertEquals(aggregationResults.size(), 11);
+ for (int i = 0; i < 11; i++) {
+ AggregationResult aggregationResult = aggregationResults.get(i);
+ List<GroupByResult> groupByResults = aggregationResult.getGroupByResult();
+ assertNotNull(groupByResults);
+ for (GroupByResult groupByResult : groupByResults) {
+ List<String> group = groupByResult.getGroup();
+ assertEquals(group.size(), 1);
+ if (i < 5) {
+ assertEquals(groupByResult.getValue(), Integer.toString(1));
+ } else {
+ assertEquals(groupByResult.getValue(), Integer.toString(3));
+ }
+ }
+ }
+ }
+ }
+
+ @Test
+ public void testPostAggregation() {
+ String query = "SELECT DISTINCT_COUNT_THETA_SKETCH(intSVColumn, '', "
+ // [300, 500), [800, 900)
+ + "'longSVColumn >= 300 AND (floatSVColumn < 500 OR doubleSVColumn BETWEEN 800 AND 899)', "
+ // [400, 850)
+ + "'intMVColumn >= 2400 AND longMVColumn < 850', "
+ // [825, 1000)
+ + "'floatMVColumn >= 2825', "
+ // [0, 100)
+ + "'doubleMVColumn < 100', "
+ // Expected: [0, 100), [400, 500), [800, 825)
+ + "'SET_UNION($4,SET_DIFF(SET_INTERSECT($1,$2),$3))') FROM testTable";
+
+ // Inner segment
+ AggregationOperator aggregationOperator = getOperatorForPqlQuery(query);
+ IntermediateResultsBlock resultsBlock = aggregationOperator.nextBlock();
+ QueriesTestUtils.testInnerSegmentExecutionStatistics(aggregationOperator.getExecutionStatistics(), NUM_RECORDS, 0,
+ 8 * NUM_RECORDS, NUM_RECORDS);
+ List<Object> aggregationResult = resultsBlock.getAggregationResult();
+ assertNotNull(aggregationResult);
+ assertEquals(aggregationResult.size(), 1);
+ List<Sketch> sketches = (List<Sketch>) aggregationResult.get(0);
+ assertEquals(sketches.size(), 5);
+ assertTrue(sketches.get(0).isEmpty());
+ assertEquals(Math.round(sketches.get(1).getEstimate()), 300);
+ assertEquals(Math.round(sketches.get(2).getEstimate()), 450);
+ assertEquals(Math.round(sketches.get(3).getEstimate()), 175);
+ assertEquals(Math.round(sketches.get(4).getEstimate()), 100);
+
+ // Inter segments
+ String[] expectedResults = new String[]{Integer.toString(225)};
+ BrokerResponseNative brokerResponse = getBrokerResponseForPqlQuery(query);
+ QueriesTestUtils
+ .testInterSegmentAggregationResult(brokerResponse, 4 * NUM_RECORDS, 0, 4 * 8 * NUM_RECORDS, 4 * NUM_RECORDS,
+ expectedResults);
+ }
+
+ @Test
+ public void testDistinctCountRawThetaSketch() {
+ String query = "SELECT DISTINCT_COUNT_RAW_THETA_SKETCH(intSVColumn) FROM testTable";
+ BrokerResponseNative brokerResponse = getBrokerResponseForPqlQuery(query);
+ String serializedSketch = (String) brokerResponse.getAggregationResults().get(0).getValue();
+ Sketch sketch = ObjectSerDeUtils.DATA_SKETCH_SER_DE.deserialize(Base64.getDecoder().decode(serializedSketch));
+ assertEquals(Math.round(sketch.getEstimate()), NUM_RECORDS);
+ }
+
+ @Test
+ public void testInvalidQueries() {
+ testInvalidQuery("select DISTINCT_COUNT_THETA_SKETCH(intSVColumn, '', 'longSVColumn < 100', '$2') from testTable");
+ testInvalidQuery("select DISTINCT_COUNT_THETA_SKETCH(intSVColumn, '', 'longSVColumn < 100', 'foo') from testTable");
+ testInvalidQuery(
+ "select DISTINCT_COUNT_THETA_SKETCH(intSVColumn, '', 'longSVColumn < 100', 'SET_UNION($1)') from testTable");
+ testInvalidQuery(
+ "select DISTINCT_COUNT_THETA_SKETCH(intSVColumn, '', 'longSVColumn < 100', 'SET_INTERSECT($1)') from testTable");
+ testInvalidQuery(
+ "select DISTINCT_COUNT_THETA_SKETCH(intSVColumn, '', 'longSVColumn < 100', 'SET_DIFF($1)') from testTable");
+ testInvalidQuery(
+ "select DISTINCT_COUNT_THETA_SKETCH(intSVColumn, '', 'longSVColumn < 100', 'floatSVColumn > 500', 'SET_DIFF($0,$1,$2)') from testTable");
+ }
+
+ private void testInvalidQuery(String query) {
+ try {
+ getBrokerResponseForSqlQuery(query);
+ fail();
+ } catch (BadQueryRequestException e) {
+ // Expected
+ }
+ }
+
+ @AfterClass
+ public void tearDown()
+ throws IOException {
+ _indexSegment.destroy();
+ FileUtils.deleteDirectory(INDEX_DIR);
+ }
+}
diff --git a/pinot-core/src/test/java/org/apache/pinot/queries/DistinctCountThetaSketchTest.java b/pinot-core/src/test/java/org/apache/pinot/queries/DistinctCountThetaSketchTest.java
deleted file mode 100644
index 2627b78..0000000
--- a/pinot-core/src/test/java/org/apache/pinot/queries/DistinctCountThetaSketchTest.java
+++ /dev/null
@@ -1,378 +0,0 @@
-/**
- * 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.pinot.queries;
-
-import java.io.File;
-import java.util.ArrayList;
-import java.util.Arrays;
-import java.util.Base64;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.List;
-import java.util.Random;
-import org.apache.commons.io.FileUtils;
-import org.apache.commons.lang3.StringUtils;
-import org.apache.datasketches.memory.Memory;
-import org.apache.datasketches.theta.Sketch;
-import org.apache.datasketches.theta.UpdateSketch;
-import org.apache.datasketches.theta.UpdateSketchBuilder;
-import org.apache.pinot.common.function.AggregationFunctionType;
-import org.apache.pinot.common.response.broker.BrokerResponseNative;
-import org.apache.pinot.common.response.broker.GroupByResult;
-import org.apache.pinot.common.segment.ReadMode;
-import org.apache.pinot.core.data.readers.GenericRowRecordReader;
-import org.apache.pinot.core.indexsegment.IndexSegment;
-import org.apache.pinot.core.indexsegment.generator.SegmentGeneratorConfig;
-import org.apache.pinot.core.indexsegment.immutable.ImmutableSegment;
-import org.apache.pinot.core.indexsegment.immutable.ImmutableSegmentLoader;
-import org.apache.pinot.core.query.exception.BadQueryRequestException;
-import org.apache.pinot.core.segment.creator.impl.SegmentIndexCreationDriverImpl;
-import org.apache.pinot.spi.config.table.TableType;
-import org.apache.pinot.spi.data.DimensionFieldSpec;
-import org.apache.pinot.spi.data.FieldSpec;
-import org.apache.pinot.spi.data.MetricFieldSpec;
-import org.apache.pinot.spi.data.Schema;
-import org.apache.pinot.spi.data.readers.GenericRow;
-import org.apache.pinot.spi.data.readers.RecordReader;
-import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
-import org.testng.annotations.AfterClass;
-import org.testng.annotations.BeforeClass;
-import org.testng.annotations.DataProvider;
-import org.testng.annotations.Test;
-
-import static org.testng.Assert.assertEquals;
-
-
-/**
- * Unit tests for {@link org.apache.pinot.core.query.aggregation.function.DistinctCountThetaSketchAggregationFunction}.
- * <ul>
- * <li> Generates a segment with 3 dimension columns with low cardinality (to increase distinct count). </li>
- * <li> Runs various queries and compares result of distinctCountThetaSketch against distinctCount function. </li>
- * </ul>
- */
-public class DistinctCountThetaSketchTest extends BaseQueriesTest {
- private static final File INDEX_DIR = new File(FileUtils.getTempDirectory(), "DistinctCountThetaSketchTest");
- private static final String TABLE_NAME = "testTable";
- private static final String SEGMENT_NAME = "testSegment";
- private static final String THETA_SKETCH_COLUMN = "colTS";
- private static final String DISTINCT_COLUMN = "distinctColumn";
-
- private static final int NUM_ROWS = 1001;
- private static final int MAX_CARDINALITY = 5; // 3 columns will lead to at most 125 groups
-
- private static final long RANDOM_SEED = System.nanoTime();
- private static final Random RANDOM = new Random(RANDOM_SEED);
- private static final String ERROR_MESSAGE = "Random seed: " + RANDOM_SEED;
-
- private IndexSegment _indexSegment;
- private List<IndexSegment> _indexSegments;
-
- @BeforeClass
- public void setup()
- throws Exception {
- FileUtils.deleteQuietly(INDEX_DIR);
-
- File segmentFile = buildSegment(buildSchema());
- ImmutableSegment immutableSegment = ImmutableSegmentLoader.load(segmentFile, ReadMode.mmap);
- _indexSegment = immutableSegment;
- _indexSegments = Arrays.asList(immutableSegment, immutableSegment);
- }
-
- @AfterClass
- public void tearDown() {
- _indexSegment.destroy();
- FileUtils.deleteQuietly(INDEX_DIR);
- }
-
- @Test
- public void testAggregationPql() {
- testThetaSketches(false, false);
- }
-
- @Test
- public void testAggregationSql() {
- testThetaSketches(false, true);
- }
-
- @Test
- public void testGroupByPql() {
- testThetaSketches(true, false);
- }
-
- @Test
- public void testGroupBySql() {
- testThetaSketches(true, true);
- }
-
- @Test(expectedExceptions = BadQueryRequestException.class, dataProvider = "badQueries")
- public void testInvalidNoPredicates(final String query) {
- getBrokerResponseForSqlQuery(query);
- }
-
- @DataProvider(name = "badQueries")
- public Object[][] badQueries() {
- return new Object[][]{
- // need at least 4 arguments in agg func
- {"select distinctCountThetaSketch(colTS, 'nominalEntries=123', '$0') from testTable"},
- // substitution arguments should start at $1
- {"select distinctCountThetaSketch(colTS, 'nominalEntries=123', 'colA = 1', '$0') from testTable"},
- // substituting variable has numeric value higher than the number of predicates provided
- {"select distinctCountThetaSketch(colTS, 'nominalEntries=123', 'colA = 1', '$5') from testTable"},
- // SET_DIFF requires exactly 2 arguments
- {"select distinctCountThetaSketch(colTS, 'nominalEntries=123', 'colA = 1', 'SET_DIFF($1)') from testTable"},
- // invalid merging function
- {"select distinctCountThetaSketch(colTS, 'nominalEntries=123', 'colA = 1', 'asdf') from testTable"},
- // union with < 2 arguments
- {"select distinctCountThetaSketch(colTS, 'nominalEntries=123', 'colA = 1', 'SET_UNION($1)')"},
- // intersect with < 2 arguments
- {"select distinctCountThetaSketch(colTS, 'nominalEntries=123', 'colA = 1', 'SET_INTERSECT($1)')"}};
- }
-
- private void testThetaSketches(boolean groupBy, boolean sql) {
- String parameters = "nominalEntries=1001";
-
- List<String> predicates = Collections.singletonList("colA = 1");
- String postAggregationExpression = "$1";
- String filter = "colA = 1";
- testThetaSketch(parameters, predicates, postAggregationExpression, filter, groupBy, sql);
-
- // Test Intersection (AND)
- predicates = Arrays.asList("colA = 1", "colB >= 2.0", "colC <> 'colC_1'");
- postAggregationExpression = "SET_INTERSECT($1, $2, $3)";
- filter = StringUtils.join(predicates, " and ");
- testThetaSketch(parameters, predicates, postAggregationExpression, filter, groupBy, sql);
-
- // Test Union (OR)
- predicates = Arrays.asList("colA = 1", "colB = 1.9");
- postAggregationExpression = "SET_UNION($1, $2)";
- filter = StringUtils.join(predicates, " or ");
- testThetaSketch(parameters, predicates, postAggregationExpression, filter, groupBy, sql);
-
- // Test complex predicates
- predicates = Arrays.asList("colA in (1, 2)", "colB not in (3.0)", "colC between 'colC_1' and 'colC_5'");
- postAggregationExpression = "SET_UNION(SET_INTERSECT($1, $2), SET_INTERSECT($1, $3))";
- filter = '(' + predicates.get(0) + " and " + predicates.get(1) + ") or (" + predicates.get(0) + " and " + predicates
- .get(2) + ')';
- testThetaSketch(parameters, predicates, postAggregationExpression, filter, groupBy, sql);
- }
-
- private void testThetaSketch(String parameters, List<String> predicates, String postAggregationExpression,
- String filter, boolean groupBy, boolean sql) {
- String distinctCountThetaSketchQuery =
- buildDistinctCountThetaSketchQuery(parameters, predicates, postAggregationExpression, filter, groupBy);
- String rawThetaSketchQuery1 = buildRawThetaSketchQuery(THETA_SKETCH_COLUMN, parameters, filter, groupBy);
- String rawThetaSketchQuery2 = buildRawThetaSketchQuery(DISTINCT_COLUMN, parameters, filter, groupBy);
- String distinctCountQuery = buildDistinctCountQuery(filter, groupBy);
- testQuery(distinctCountThetaSketchQuery, distinctCountQuery, groupBy, sql, false);
- testQuery(rawThetaSketchQuery1, distinctCountQuery, groupBy, sql, true);
- testQuery(rawThetaSketchQuery2, distinctCountQuery, groupBy, sql, true);
- }
-
- private void testQuery(String tsQuery, String distinctQuery, boolean groupBy, boolean sql, boolean raw) {
- BrokerResponseNative actualResponse =
- (sql) ? getBrokerResponseForSqlQuery(tsQuery) : getBrokerResponseForPqlQuery(tsQuery);
-
- BrokerResponseNative expectedResponse =
- (sql) ? getBrokerResponseForSqlQuery(distinctQuery) : getBrokerResponseForPqlQuery(distinctQuery);
-
- if (groupBy) {
- compareGroupBy(actualResponse, expectedResponse, sql, raw);
- } else {
- compareAggregation(actualResponse, expectedResponse, sql, raw);
- }
- }
-
- private void compareAggregation(BrokerResponseNative actualResponse, BrokerResponseNative expectedResponse,
- boolean sql, boolean raw) {
- if (sql) {
- compareSql(actualResponse, expectedResponse, raw);
- } else {
- compareAggregationPql(actualResponse, expectedResponse, raw);
- }
- }
-
- private void compareGroupBy(BrokerResponseNative actualResponse, BrokerResponseNative expectedResponse, boolean sql,
- boolean raw) {
- if (sql) {
- compareSql(actualResponse, expectedResponse, raw);
- } else {
- compareGroupByPql(actualResponse, expectedResponse, raw);
- }
- }
-
- private void compareSql(BrokerResponseNative actualResponse, BrokerResponseNative expectedResponse, boolean raw) {
- List<Object[]> actualRows = actualResponse.getResultTable().getRows();
- List<Object[]> expectedRows = expectedResponse.getResultTable().getRows();
- int numRows = actualRows.size();
- assertEquals(numRows, expectedRows.size(), ERROR_MESSAGE);
- for (int i = 0; i < numRows; i++) {
- int actual = getSketchValue(actualRows.get(i)[0].toString(), raw);
- int expected = (int) expectedRows.get(i)[0];
- assertEquals(actual, expected, ERROR_MESSAGE);
- }
- }
-
- private void compareAggregationPql(BrokerResponseNative actualResponse, BrokerResponseNative expectedResponse,
- boolean raw) {
- int actual = getSketchValue((String) actualResponse.getAggregationResults().get(0).getValue(), raw);
- int expected = Integer.parseInt((String) expectedResponse.getAggregationResults().get(0).getValue());
- assertEquals(actual, expected, ERROR_MESSAGE);
- }
-
- private void compareGroupByPql(BrokerResponseNative actualResponse, BrokerResponseNative expectedResponse,
- boolean raw) {
- List<GroupByResult> actualResults = actualResponse.getAggregationResults().get(0).getGroupByResult();
- List<GroupByResult> expectedResults = expectedResponse.getAggregationResults().get(0).getGroupByResult();
- int numGroups = actualResults.size();
- assertEquals(numGroups, expectedResults.size(), ERROR_MESSAGE);
- for (int i = 0; i < numGroups; i++) {
- int actual = getSketchValue((String) actualResults.get(i).getValue(), raw);
- int expected = Integer.parseInt((String) expectedResults.get(i).getValue());
- assertEquals(actual, expected, ERROR_MESSAGE);
- }
- }
-
- private int getSketchValue(String value, boolean raw) {
- if (!raw) {
- return Integer.parseInt(value);
- } else {
- byte[] bytes = Base64.getDecoder().decode(value);
- return (int) Math.round(Sketch.wrap(Memory.wrap(bytes)).getEstimate());
- }
- }
-
- private String buildDistinctCountThetaSketchQuery(String parameters, List<String> predicates,
- String postAggregationExpression, String filter, boolean groupBy) {
- StringBuilder stringBuilder =
- new StringBuilder("select ").append(AggregationFunctionType.DISTINCTCOUNTTHETASKETCH.getName()).append('(')
- .append(THETA_SKETCH_COLUMN).append(",'").append(parameters).append("','");
- for (String predicate : predicates) {
- stringBuilder.append(predicate.replace("'", "''")).append("','");
- }
- stringBuilder.append(postAggregationExpression.replace("'", "''")).append("')");
-
- stringBuilder.append(" from ").append(TABLE_NAME).append(" where ").append(filter);
- if (groupBy) {
- stringBuilder.append(" group by colA, colB");
- }
- return stringBuilder.toString();
- }
-
- private String buildRawThetaSketchQuery(String column, String parameters, String filter, boolean groupBy) {
- StringBuilder stringBuilder =
- new StringBuilder("select ").append(AggregationFunctionType.RAWTHETASKETCH.getName()).append('(').append(column)
- .append(",'").append(parameters).append("')");
-
- stringBuilder.append(" from ").append(TABLE_NAME).append(" where ").append(filter);
- if (groupBy) {
- stringBuilder.append(" group by colA, colB");
- }
- return stringBuilder.toString();
- }
-
- private String buildDistinctCountQuery(String filter, boolean groupBy) {
- StringBuilder stringBuilder =
- new StringBuilder("select ").append(AggregationFunctionType.DISTINCTCOUNT.getName()).append('(')
- .append(DISTINCT_COLUMN).append(')');
-
- stringBuilder.append(" from ").append(TABLE_NAME).append(" where ").append(filter);
- if (groupBy) {
- stringBuilder.append(" group by colA, colB");
- }
- return stringBuilder.toString();
- }
-
- @Override
- protected String getFilter() {
- return ""; // No filters required for this test.
- }
-
- @Override
- protected IndexSegment getIndexSegment() {
- return _indexSegment;
- }
-
- @Override
- protected List<IndexSegment> getIndexSegments() {
- return _indexSegments;
- }
-
- protected File buildSegment(Schema schema)
- throws Exception {
-
- StringBuilder stringBuilder = new StringBuilder();
- List<GenericRow> rows = new ArrayList<>(NUM_ROWS);
-
- for (int i = 0; i < NUM_ROWS; i++) {
- stringBuilder.setLength(0);
- HashMap<String, Object> valueMap = new HashMap<>();
-
- int value = (i % (1 + RANDOM.nextInt(MAX_CARDINALITY)));
- valueMap.put("colA", value);
- stringBuilder.append(value);
-
- value = (i % (1 + RANDOM.nextInt(MAX_CARDINALITY)));
- valueMap.put("colB", (i % (1 + RANDOM.nextInt(MAX_CARDINALITY))));
- stringBuilder.append(value);
-
- String sValue = "colC" + "_" + (i % (1 + RANDOM.nextInt(MAX_CARDINALITY)));
- valueMap.put("colC", sValue);
- stringBuilder.append(sValue);
-
- String distinctValue = stringBuilder.toString();
- valueMap.put(DISTINCT_COLUMN, distinctValue);
-
- UpdateSketch sketch = new UpdateSketchBuilder().build();
- sketch.update(distinctValue);
- valueMap.put(THETA_SKETCH_COLUMN, sketch.compact().toByteArray());
-
- GenericRow genericRow = new GenericRow();
- genericRow.init(valueMap);
- rows.add(genericRow);
- }
-
- SegmentGeneratorConfig config =
- new SegmentGeneratorConfig(new TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE_NAME).build(), schema);
- config.setOutDir(INDEX_DIR.getPath());
- config.setTableName(TABLE_NAME);
- config.setSegmentName(SEGMENT_NAME);
-
- config.setRawIndexCreationColumns(Collections.singletonList(THETA_SKETCH_COLUMN));
- SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl();
-
- try (RecordReader recordReader = new GenericRowRecordReader(rows)) {
- driver.init(config, recordReader);
- driver.build();
- }
-
- return driver.getOutputDirectory();
- }
-
- private Schema buildSchema() {
- Schema schema = new Schema();
-
- schema.addField(new DimensionFieldSpec("colA", FieldSpec.DataType.INT, true));
- schema.addField(new DimensionFieldSpec("colB", FieldSpec.DataType.DOUBLE, true));
- schema.addField(new DimensionFieldSpec("colC", FieldSpec.DataType.STRING, true));
-
- schema.addField(new DimensionFieldSpec(DISTINCT_COLUMN, FieldSpec.DataType.STRING, true));
- schema.addField(new MetricFieldSpec(THETA_SKETCH_COLUMN, FieldSpec.DataType.BYTES));
- return schema;
- }
-}
diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ThetaSketchIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ThetaSketchIntegrationTest.java
index 4bc3a24..efc03b3 100644
--- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ThetaSketchIntegrationTest.java
+++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/ThetaSketchIntegrationTest.java
@@ -52,7 +52,6 @@ import static org.testng.Assert.assertEquals;
public class ThetaSketchIntegrationTest extends BaseClusterIntegrationTest {
-
private static final String DIM_NAME = "dimName";
private static final String DIM_VALUE = "dimValue";
private static final String SHARD_ID = "shardId";
@@ -125,45 +124,97 @@ public class ThetaSketchIntegrationTest extends BaseClusterIntegrationTest {
*/
// gender = female
- String query = "select distinctCountThetaSketch(thetaSketchCol, '', "
- + "'dimName = ''gender''', 'dimValue = ''Female''', 'SET_INTERSECT($1, $2)') from "
- + DEFAULT_TABLE_NAME + " where dimName = 'gender' AND dimValue = 'Female'";
- runAndAssert(query, 50 + 60 + 70 + 110 + 120 + 130);
+ {
+ String query = "select distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME
+ + " where dimName = 'gender' and dimValue = 'Female'";
+ int expected = 50 + 60 + 70 + 110 + 120 + 130;
+ runAndAssert(query, expected);
+
+ query =
+ "select distinctCountThetaSketch(thetaSketchCol, '', 'dimName = ''gender'' and dimValue = ''Female''', '$1') from "
+ + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+
+ query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''gender''', 'dimValue = ''Female''', 'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+ }
// gender = male
- query = "select distinctCountThetaSketch(thetaSketchCol, '', "
- + "'dimName = ''gender''', 'dimValue = ''Male''', 'SET_INTERSECT($1, $2)') from "
- + DEFAULT_TABLE_NAME + " where dimName = 'gender' AND dimValue = 'Male'";
- runAndAssert(query, 80 + 90 + 100 + 140 + 150 + 160);
+ {
+ String query = "select distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME
+ + " where dimName = 'gender' and dimValue = 'Male'";
+ int expected = 80 + 90 + 100 + 140 + 150 + 160;
+ runAndAssert(query, expected);
+
+ query =
+ "select distinctCountThetaSketch(thetaSketchCol, '', 'dimName = ''gender'' and dimValue = ''Male''', '$1') from "
+ + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+
+ query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''gender''', 'dimValue = ''Male''', 'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+ }
// course = math
- query = "select distinctCountThetaSketch(thetaSketchCol, '', "
- + "'dimName = ''course''', 'dimValue = ''Math''', 'SET_INTERSECT($1, $2)') from "
- + DEFAULT_TABLE_NAME + " where dimName = 'course' AND dimValue = 'Math'";
- runAndAssert(query, 50 + 80 + 110 + 140);
-
- // gender = female AND course = math
- query = "select distinctCountThetaSketch(thetaSketchCol, '', "
- + "'dimName = ''gender''', 'dimValue = ''Female''', 'dimName = ''course''', 'dimValue = ''Math''', "
- + "'SET_INTERSECT(SET_INTERSECT($1, $2), SET_INTERSECT($3, $4))') from "
- + DEFAULT_TABLE_NAME
- + " where (dimName = 'gender' AND dimValue = 'Female') OR (dimName = 'course' AND dimValue = 'Math')";
- runAndAssert(query, 50 + 110);
-
- // gender = male OR course = biology
- query = "select distinctCountThetaSketch(thetaSketchCol, '', "
- + "'dimName = ''gender''', 'dimValue = ''Male''', 'dimName = ''course''', 'dimValue = ''Biology''', "
- + "'SET_UNION(SET_INTERSECT($1, $2), SET_INTERSECT($3, $4))') from "
- + DEFAULT_TABLE_NAME
- + " where (dimName = 'gender' AND dimValue = 'Male') OR (dimName = 'course' AND dimValue = 'Biology')";
- runAndAssert(query, 70 + 80 + 90 + 100 + 130 + 140 + 150 + 160);
+ {
+ String query = "select distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME
+ + " where dimName = 'course' AND dimValue = 'Math'";
+ int expected = 50 + 80 + 110 + 140;
+ runAndAssert(query, expected);
+
+ query =
+ "select distinctCountThetaSketch(thetaSketchCol, '', 'dimName = ''course'' and dimValue = ''Math''', '$1') from "
+ + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+
+ query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''course''', 'dimValue = ''Math''', 'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+ }
+
+ // gender = female INTERSECT course = math
+ {
+ String query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''gender'' and dimValue = ''Female''', 'dimName = ''course'' and dimValue = ''Math''', "
+ + "'SET_INTERSECT($1, $2)') from " + DEFAULT_TABLE_NAME;
+ int expected = 50 + 110;
+ runAndAssert(query, expected);
+
+ query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''gender''', 'dimValue = ''Female''', 'dimName = ''course''', 'dimValue = ''Math''', "
+ + "'SET_INTERSECT($1, $2, $3, $4)') from " + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+
+ query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''gender''', 'dimValue = ''Female''', 'dimName = ''course''', 'dimValue = ''Math''', "
+ + "'SET_INTERSECT(SET_INTERSECT($1, $2), SET_INTERSECT($3, $4))') from " + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+ }
+
+ // gender = male UNION course = biology
+ {
+ String query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''gender'' and dimValue = ''Male''', 'dimName = ''course'' and dimValue = ''Biology''', "
+ + "'SET_UNION($1, $2)') from " + DEFAULT_TABLE_NAME;
+ int expected = 70 + 80 + 90 + 100 + 130 + 140 + 150 + 160;
+ runAndAssert(query, expected);
+
+ query = "select distinctCountThetaSketch(thetaSketchCol, '', "
+ + "'dimName = ''gender''', 'dimValue = ''Male''', 'dimName = ''course''', 'dimValue = ''Biology''', "
+ + "'SET_UNION(SET_INTERSECT($1, $2), SET_INTERSECT($3, $4))') from " + DEFAULT_TABLE_NAME;
+ runAndAssert(query, expected);
+ }
// group by gender
- query = "select dimValue, distinctCountThetaSketch(thetaSketchCol, '', 'dimName = ''gender''', '$1') from "
- + DEFAULT_TABLE_NAME + " where dimName = 'gender' AND (dimValue = 'Female' OR dimValue = 'Male') "
- + "group by dimValue";
- runAndAssert(query,
- ImmutableMap.of("Female", 50 + 60 + 70 + 110 + 120 + 130, "Male", 80 + 90 + 100 + 140 + 150 + 160));
+ {
+ String query = "select dimValue, distinctCountThetaSketch(thetaSketchCol) from " + DEFAULT_TABLE_NAME
+ + " where dimName = 'gender' group by dimValue";
+ ImmutableMap<String, Integer> expected =
+ ImmutableMap.of("Female", 50 + 60 + 70 + 110 + 120 + 130, "Male", 80 + 90 + 100 + 140 + 150 + 160);
+ runAndAssert(query, expected);
+ }
}
private void runAndAssert(String query, int expected)
---------------------------------------------------------------------
To unsubscribe, e-mail: commits-unsubscribe@pinot.apache.org
For additional commands, e-mail: commits-help@pinot.apache.org