From 2b98f4ceadb1a91420a76c60794655c663387d28 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 7 Jul 2026 14:27:43 -0700 Subject: [PATCH 01/19] [UUID 5/8] UUID aggregation, group-by and distinct Part 5/8 of splitting apache/pinot#18140 (logical UUID type). Rebased onto latest master; stacked on uuid-split/04-sse-predicates-cast. Downstream references use the UuidKey class merged in #18869. --- .../function/AggregationFunctionUtils.java | 55 +++++++++-- .../function/AnyValueAggregationFunction.java | 7 ++ ...istinctCountBitmapAggregationFunction.java | 47 ++++++++- ...inctCountCPCSketchAggregationFunction.java | 53 ++++++++-- .../DistinctCountHLLAggregationFunction.java | 48 +++++++++- ...stinctCountHLLPlusAggregationFunction.java | 45 ++++++++- ...ctCountThetaSketchAggregationFunction.java | 42 +++++++- .../DistinctCountULLAggregationFunction.java | 51 +++++++++- ...IntegerTupleSketchAggregationFunction.java | 25 ++++- ...ictionaryMultiColumnGroupKeyGenerator.java | 68 ++++++++++--- ...ctionarySingleColumnGroupKeyGenerator.java | 91 +++++++++++++++--- .../groupby/utils/UuidToIdMap.java | 59 ++++++++++++ .../groupby/utils/ValueToIdMapFactory.java | 2 + .../distinct/table/BytesDistinctTable.java | 16 +++- .../query/reduce/GroupByDataTableReducer.java | 1 + ...stinctCountHLLAggregationFunctionTest.java | 96 +++++++++++++++++++ .../NoDictionaryGroupKeyGeneratorTest.java | 48 ++++++++-- .../table/BytesDistinctTableTest.java | 70 ++++++++++++++ 18 files changed, 747 insertions(+), 77 deletions(-) create mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java create mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index 527b7817621c..c5edbd429485 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -74,6 +74,7 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.query.QueryThreadContext; import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; /// The `AggregationFunctionUtils` class provides utility methods for aggregation function. @@ -610,23 +611,23 @@ public static Object getAggregationResult(AggregationFunction aggregationFunctio break; case DISTINCTCOUNTHLL: case DISTINCTCOUNTHLLMV: - result = getDistinctCountHLLResult(Objects.requireNonNull(dataSource.getDictionary()), + result = getDistinctCountHLLResult(dataSource, (DistinctCountHLLAggregationFunction) aggregationFunction, explainPlanName); break; case DISTINCTCOUNTRAWHLL: case DISTINCTCOUNTRAWHLLMV: - result = getDistinctCountHLLResult(Objects.requireNonNull(dataSource.getDictionary()), + result = getDistinctCountHLLResult(dataSource, ((DistinctCountRawHLLAggregationFunction) aggregationFunction).getDistinctCountHLLAggregationFunction(), explainPlanName); break; case DISTINCTCOUNTHLLPLUS: case DISTINCTCOUNTHLLPLUSMV: - result = getDistinctCountHLLPlusResult(Objects.requireNonNull(dataSource.getDictionary()), + result = getDistinctCountHLLPlusResult(dataSource, (DistinctCountHLLPlusAggregationFunction) aggregationFunction, explainPlanName); break; case DISTINCTCOUNTRAWHLLPLUS: case DISTINCTCOUNTRAWHLLPLUSMV: - result = getDistinctCountHLLPlusResult(Objects.requireNonNull(dataSource.getDictionary()), + result = getDistinctCountHLLPlusResult(dataSource, ((DistinctCountRawHLLPlusAggregationFunction) aggregationFunction) .getDistinctCountHLLPlusAggregationFunction(), explainPlanName); break; @@ -642,7 +643,7 @@ public static Object getAggregationResult(AggregationFunction aggregationFunctio (DistinctCountSmartHLLPlusAggregationFunction) aggregationFunction, explainPlanName); break; case DISTINCTCOUNTULL: - result = getDistinctCountULLResult(Objects.requireNonNull(dataSource.getDictionary()), + result = getDistinctCountULLResult(dataSource, (DistinctCountULLAggregationFunction) aggregationFunction, explainPlanName); break; case DISTINCTCOUNTSMARTULL: @@ -650,7 +651,7 @@ public static Object getAggregationResult(AggregationFunction aggregationFunctio (DistinctCountSmartULLAggregationFunction) aggregationFunction, explainPlanName); break; case DISTINCTCOUNTRAWULL: - result = getDistinctCountULLResult(Objects.requireNonNull(dataSource.getDictionary()), + result = getDistinctCountULLResult(dataSource, (DistinctCountULLAggregationFunction) aggregationFunction, explainPlanName); break; default: @@ -799,8 +800,20 @@ private static HyperLogLogPlus getDistinctValueHLLPlus(Dictionary dictionary, in return hllPlus; } - private static HyperLogLog getDistinctCountHLLResult(Dictionary dictionary, + private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, DistinctCountHLLAggregationFunction function, String explainPlanName) { + Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); + // UUID dictionary entries are logical scalar values, not serialized HyperLogLogs. Offer their canonical string + // representation to match the scan-based path and DISTINCTCOUNTHLL(CAST(uuidColumn AS STRING)). + if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.UUID) { + HyperLogLog hll = new HyperLogLog(function.getLog2m()); + int length = dictionary.length(); + for (int i = 0; i < length; i++) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(i, explainPlanName); + hll.offer(UuidUtils.toString(dictionary.getBytesValue(i))); + } + return hll; + } if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { // Treat BYTES value as serialized HyperLogLog try { @@ -820,8 +833,20 @@ private static HyperLogLog getDistinctCountHLLResult(Dictionary dictionary, } } - private static HyperLogLogPlus getDistinctCountHLLPlusResult(Dictionary dictionary, + private static HyperLogLogPlus getDistinctCountHLLPlusResult(DataSource dataSource, DistinctCountHLLPlusAggregationFunction function, String explainPlanName) { + Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); + // UUID dictionary entries are logical scalar values, not serialized HyperLogLogPluses. Offer their canonical + // string representation to match the scan-based path and DISTINCTCOUNTHLLPLUS(CAST(uuidColumn AS STRING)). + if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.UUID) { + HyperLogLogPlus hllPlus = new HyperLogLogPlus(function.getP(), function.getSp()); + int length = dictionary.length(); + for (int i = 0; i < length; i++) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(i, explainPlanName); + hllPlus.offer(UuidUtils.toString(dictionary.getBytesValue(i))); + } + return hllPlus; + } if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { // Treat BYTES value as serialized HyperLogLogPlus try { @@ -861,8 +886,20 @@ private static Object getDistinctCountSmartHLLPlusResult(Dictionary dictionary, } } - private static UltraLogLog getDistinctCountULLResult(Dictionary dictionary, + private static UltraLogLog getDistinctCountULLResult(DataSource dataSource, DistinctCountULLAggregationFunction function, String explainPlanName) { + Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); + // UUID dictionary entries are logical scalar values, not serialized UltraLogLogs. Hash their canonical string + // representation to match the scan-based path and DISTINCTCOUNTULL(CAST(uuidColumn AS STRING)). + if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.UUID) { + UltraLogLog ull = UltraLogLog.create(function.getP()); + int length = dictionary.length(); + for (int i = 0; i < length; i++) { + QueryThreadContext.checkTerminationAndSampleUsagePeriodically(i, explainPlanName); + UltraLogLogUtils.hashObject(UuidUtils.toString(dictionary.getBytesValue(i))).ifPresent(ull::add); + } + return ull; + } if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { // Treat BYTES value as serialized UltraLogLog and merge try { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java index a649fe6d7a17..7a3ba85337ca 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java @@ -327,6 +327,13 @@ private void ensureResultType(BlockValSet bvs) { if (_resultType != null) { return; } + // Inspect the logical type first so a UUID column reports ColumnDataType.UUID (and the broker renders canonical + // RFC-4122 strings) rather than collapsing to BYTES (which would render hex). All other dispatch keys off the + // stored type, matching the BYTES/STRING storage convention. + if (bvs.getValueType() == FieldSpec.DataType.UUID) { + _resultType = ColumnDataType.UUID; + return; + } switch (bvs.getValueType().getStoredType()) { case INT: _resultType = ColumnDataType.INT; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index 4fe96b819dbb..04a3bc0d9b4d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -34,6 +34,7 @@ import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -71,8 +72,22 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized RoaringBitmap state. Add the + // hashCode of the canonical UUID string so DISTINCTCOUNTBITMAP(uuidCol) matches + // DISTINCTCOUNTBITMAP(CAST(uuidCol AS STRING)). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + RoaringBitmap bitmap = getValueBitmap(aggregationResultHolder); + for (int i = 0; i < length; i++) { + bitmap.add(UuidUtils.toString(uuidBytesValues[i]).hashCode()); + } + return; + } + // Treat BYTES value as serialized RoaringBitmap - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); RoaringBitmap valueBitmap = aggregationResultHolder.getResult(); @@ -209,8 +224,20 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: add hashCode of canonical UUID string (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getValueBitmap(groupByResultHolder, groupKeyArray[i]) + .add(UuidUtils.toString(uuidBytesValues[i]).hashCode()); + } + return; + } + // Treat BYTES value as serialized RoaringBitmap - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { @@ -350,8 +377,22 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: add hashCode of canonical UUID string (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + int hash = UuidUtils.toString(uuidBytesValues[i]).hashCode(); + for (int groupKey : groupKeysArray[i]) { + getValueBitmap(groupByResultHolder, groupKey).add(hash); + } + } + return; + } + // Treat BYTES value as serialized RoaringBitmap - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index ec273c066ab1..7111f211e509 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -41,6 +41,7 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -135,8 +136,24 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + FieldSpec.DataType dataType = blockValSet.getValueType(); + FieldSpec.DataType storedType = dataType.getStoredType(); + + // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized CPC Sketch state. Update + // the sketch with the canonical UUID string so DISTINCTCOUNTCPC(uuidCol) matches + // DISTINCTCOUNTCPC(CAST(uuidCol AS STRING)). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + // Leave the updated CpcSketch in the holder; extractAggregationResult converts it to an accumulator. + // Calling getAccumulator here would read the holder slot already occupied by the sketch and fail. + CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); + for (int i = 0; i < length; i++) { + cpcSketch.update(UuidUtils.toString(uuidBytesValues[i])); + } + return; + } + // Treat BYTES value as serialized CPC Sketch - FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -197,8 +214,9 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } - CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(aggregationResultHolder); - cpcSketchAccumulator.apply(cpcSketch); + // The updated CpcSketch already lives in the holder (getCpcSketch stored it); extractAggregationResult + // converts it to a CpcSketchAccumulator. Reading the holder as an accumulator here would + // ClassCastException — the holder slot contains the sketch, not an accumulator. } @Override @@ -206,8 +224,19 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: update with canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(UuidUtils.toString(uuidBytesValues[i])); + } + return; + } + // Treat BYTES value as serialized CPC Sketch - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == FieldSpec.DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -277,10 +306,22 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized CPC Sketch - DataType storedType = blockValSet.getValueType().getStoredType(); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); boolean singleValue = blockValSet.isSingleValue(); + // UUID columns: update with canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID && singleValue) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + String canonical = UuidUtils.toString(uuidBytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + getCpcSketch(groupByResultHolder, groupKey).update(canonical); + } + } + return; + } + if (singleValue && storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java index 07bfb6bd41af..8a74f972f7f3 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -38,6 +38,7 @@ import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.RoaringBitmap; @@ -81,8 +82,24 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns are stored as 16-byte BYTES, but a UUID value is a logical scalar — not a serialized + // HyperLogLog. Offer the canonical UUID string so the result matches DISTINCTCOUNTHLL on a STRING column + // holding the same logical UUIDs. NOTE: fetch raw bytes and convert explicitly — for identifier expressions + // the BlockValSet is a ProjectionBlockValSet whose getStringValuesSV() renders stored BYTES as bare hex, + // not the canonical RFC-4122 form. + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + HyperLogLog hyperLogLog = getHyperLogLog(aggregationResultHolder); + for (int i = 0; i < length; i++) { + hyperLogLog.offer(UuidUtils.toString(uuidBytesValues[i])); + } + return; + } + // Treat BYTES value as serialized HyperLogLog - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -232,8 +249,19 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(UuidUtils.toString(uuidBytesValues[i])); + } + return; + } + // Treat BYTES value as serialized HyperLogLog - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -381,8 +409,22 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + String canonical = UuidUtils.toString(uuidBytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + getHyperLogLog(groupByResultHolder, groupKey).offer(canonical); + } + } + return; + } + // Treat BYTES value as serialized HyperLogLog - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java index fd337f433a34..1c2762670c61 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -37,6 +37,7 @@ import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -93,8 +94,21 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized HyperLogLogPlus state. Offer the + // canonical UUID string so DISTINCTCOUNTHLLPLUS(uuidCol) matches DISTINCTCOUNTHLLPLUS(CAST(uuidCol AS STRING)). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(aggregationResultHolder); + for (int i = 0; i < length; i++) { + hyperLogLogPlus.offer(UuidUtils.toString(uuidBytesValues[i])); + } + return; + } + // Treat BYTES value as serialized HyperLogLogPlus - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -237,8 +251,19 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(UuidUtils.toString(uuidBytesValues[i])); + } + return; + } + // Treat BYTES value as serialized HyperLogLogPlus - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -385,8 +410,22 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + String canonical = UuidUtils.toString(uuidBytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + getHyperLogLogPlus(groupByResultHolder, groupKey).offer(canonical); + } + } + return; + } + // Treat BYTES value as serialized HyperLogLogPlus - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { 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 a2859fe471ec..67e9c2996ef9 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 @@ -56,6 +56,7 @@ import org.apache.pinot.segment.spi.Constants; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.sql.parsers.CalciteSqlParser; @@ -189,7 +190,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde boolean[] singleValues = new boolean[numExpressions]; DataType[] valueTypes = new DataType[numExpressions]; Object[] valueArrays = new Object[numExpressions]; - extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); + extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); // Main expression is always index 0 @@ -438,7 +439,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol boolean[] singleValues = new boolean[numExpressions]; DataType[] valueTypes = new DataType[numExpressions]; Object[] valueArrays = new Object[numExpressions]; - extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); + extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); // Main expression is always index 0 @@ -662,7 +663,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult boolean[] singleValues = new boolean[numExpressions]; DataType[] valueTypes = new DataType[numExpressions]; Object[] valueArrays = new Object[numExpressions]; - extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); + extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); // Main expression is always index 0 @@ -1226,14 +1227,45 @@ private static int extractSketchId(String identifier) { } /// Extracts values from the BlockValSet map. - private void extractValues(Map blockValSetMap, boolean[] singleValues, + private void extractValues(int length, Map 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 storedType = blockValSet.getValueType().getStoredType(); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); singleValues[i] = singleValue; + // UUID columns are stored as 16-byte BYTES but a UUID value is a logical scalar, not a pre-serialized + // theta sketch. Surface UUID as STRING (canonical UUID form) so the downstream update-sketch path + // matches DISTINCTCOUNTTHETASKETCH(CAST(uuidCol AS STRING)). Without this branch, the function would + // take the serialized-sketch path below and Sketch.wrap would fail on raw 16-byte UUID content. + // NOTE: fetch raw bytes and convert explicitly — for identifier expressions the BlockValSet is a + // ProjectionBlockValSet whose getStringValuesSV() renders stored BYTES as bare hex, not canonical form. + if (dataType == DataType.UUID) { + valueTypes[i] = DataType.STRING; + if (singleValue) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + String[] canonicalValues = new String[length]; + for (int j = 0; j < length; j++) { + canonicalValues[j] = UuidUtils.toString(uuidBytesValues[j]); + } + valueArrays[i] = canonicalValues; + } else { + byte[][][] uuidBytesValuesMV = blockValSet.getBytesValuesMV(); + String[][] canonicalValuesMV = new String[length][]; + for (int j = 0; j < length; j++) { + byte[][] row = uuidBytesValuesMV[j]; + String[] canonicalRow = new String[row.length]; + for (int k = 0; k < row.length; k++) { + canonicalRow[k] = UuidUtils.toString(row[k]); + } + canonicalValuesMV[j] = canonicalRow; + } + valueArrays[i] = canonicalValuesMV; + } + continue; + } valueTypes[i] = storedType; if (singleValue) { switch (storedType) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java index 5309a9713af3..cad535aa2aea 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -38,6 +38,7 @@ import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -82,8 +83,21 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized HyperLogLogPlus - DataType storedType = blockValSet.getValueType().getStoredType(); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized UltraLogLog state. Hash the + // canonical UUID string so DISTINCTCOUNTULL(uuidCol) matches DISTINCTCOUNTULL(CAST(uuidCol AS STRING)). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + UltraLogLog ull = getULL(aggregationResultHolder); + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(UuidUtils.toString(uuidBytesValues[i])).ifPresent(ull::add); + } + return; + } + + // Treat BYTES value as serialized UltraLogLog if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -155,8 +169,20 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: hash canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); + UltraLogLogUtils.hashObject(UuidUtils.toString(uuidBytesValues[i])).ifPresent(ull::add); + } + return; + } + // Treat BYTES value as serialized UltraLogLogs - DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -234,8 +260,23 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized HyperLogLogPlus - DataType storedType = blockValSet.getValueType().getStoredType(); + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + + // UUID columns: hash canonical UUID strings converted from raw bytes (see aggregate() for rationale). + if (dataType == DataType.UUID) { + byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + String canonical = UuidUtils.toString(uuidBytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + UltraLogLog ull = getULL(groupByResultHolder, groupKey); + UltraLogLogUtils.hashObject(canonical).ifPresent(ull::add); + } + } + return; + } + + // Treat BYTES value as serialized UltraLogLogs if (storedType == DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java index 7591272c0d5b..58b4f7c68607 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java @@ -152,8 +152,15 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); + FieldSpec.DataType dataType = blockValSet.getValueType(); + FieldSpec.DataType storedType = dataType.getStoredType(); + // UUID columns are stored as BYTES but contain raw 16-byte UUID values, not serialized tuple sketches. + // Surface a clear error rather than letting the deserialize step fail with a confusing sketch-format message. + if (dataType == FieldSpec.DataType.UUID) { + throw new IllegalStateException(getType() + " does not accept raw UUID values; pre-serialize the column as " + + "Integer Tuple Sketches first"); + } // Treat BYTES value as serialized Integer Tuple Sketch - FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == FieldSpec.DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -184,8 +191,12 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized Integer Tuple Sketch - FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); + FieldSpec.DataType dataType = blockValSet.getValueType(); + FieldSpec.DataType storedType = dataType.getStoredType(); + if (dataType == FieldSpec.DataType.UUID) { + throw new IllegalStateException(getType() + " does not accept raw UUID values; pre-serialize the column as " + + "Integer Tuple Sketches first"); + } if (storedType == FieldSpec.DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); @@ -209,8 +220,12 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized Integer Tuple Sketch - FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); + FieldSpec.DataType dataType = blockValSet.getValueType(); + FieldSpec.DataType storedType = dataType.getStoredType(); + if (dataType == FieldSpec.DataType.UUID) { + throw new IllegalStateException(getType() + " does not accept raw UUID values; pre-serialize the column as " + + "Integer Tuple Sketches first"); + } boolean singleValue = blockValSet.isSingleValue(); if (singleValue && storedType == FieldSpec.DataType.BYTES) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java index a10de9f47a58..f6d3784ebf96 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java @@ -36,6 +36,7 @@ import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.FixedIntArray; +import org.apache.pinot.spi.utils.UuidKey; import org.roaringbitmap.RoaringBitmap; @@ -51,7 +52,11 @@ public class NoDictionaryMultiColumnGroupKeyGenerator implements GroupKeyGenerat private final ExpressionContext[] _groupByExpressions; private final int _numGroupByExpressions; - private final DataType[] _storedTypes; + /// Per-column group-key dispatch type: stored type of each column, except UUID is preserved as + /// [DataType#UUID] so the on-the-fly dictionary keys on [org.apache.pinot.spi.utils.UuidKey] + /// instead of [org.apache.pinot.spi.utils.ByteArray]. For every other logical type this equals + /// `logicalType.getStoredType()`. + private final DataType[] _dataTypes; private final Dictionary[] _dictionaries; private final ValueToIdMap[] _onTheFlyDictionaries; private final Object2IntOpenHashMap _groupKeyMap; @@ -65,7 +70,7 @@ public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOp Map groupByExpressionSizesFromPredicates) { _groupByExpressions = groupByExpressions; _numGroupByExpressions = groupByExpressions.length; - _storedTypes = new DataType[_numGroupByExpressions]; + _dataTypes = new DataType[_numGroupByExpressions]; _dictionaries = new Dictionary[_numGroupByExpressions]; _onTheFlyDictionaries = new ValueToIdMap[_numGroupByExpressions]; _isSingleValueExpressions = new boolean[_numGroupByExpressions]; @@ -75,7 +80,9 @@ public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOp for (int i = 0; i < _numGroupByExpressions; i++) { ExpressionContext groupByExpression = groupByExpressions[i]; ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpression); - _storedTypes[i] = columnContext.getDataType().getStoredType(); + DataType logicalType = columnContext.getDataType(); + // Normalize to stored type, but preserve UUID so it uses UuidKey instead of ByteArray. + _dataTypes[i] = logicalType == DataType.UUID ? DataType.UUID : logicalType.getStoredType(); // Take the dict-id path only when the forward index is dict-encoded. A column with EncodingType.RAW + // dictionaryIndex exposes a Dictionary but BlockValSet#getDictionaryIdsSV throws on its RAW forward // index — fall back to an on-the-fly dictionary on raw values for that case. @@ -84,7 +91,7 @@ public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOp if (dictionary != null) { _dictionaries[i] = dictionary; } else { - _onTheFlyDictionaries[i] = ValueToIdMapFactory.get(_storedTypes[i]); + _onTheFlyDictionaries[i] = ValueToIdMapFactory.get(_dataTypes[i]); } if (canOptimizeGroupByUpperBound) { Integer size = groupByExpressionSizesFromPredicates.get(groupByExpression); @@ -121,7 +128,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { if (_dictionaries[i] != null) { values[i] = blockValSet.getDictionaryIdsSV(); } else { - switch (_storedTypes[i]) { + switch (_dataTypes[i]) { case INT: values[i] = blockValSet.getIntValuesSV(); break; @@ -141,10 +148,11 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { values[i] = blockValSet.getStringValuesSV(); break; case BYTES: + case UUID: values[i] = blockValSet.getBytesValuesSV(); break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedTypes[i]); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataTypes[i]); } } } @@ -176,7 +184,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.put(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = onTheFlyDictionary.put(new ByteArray(((byte[][]) columnValues)[row])); + keyValue = putBytesValue(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); } else { keyValue = onTheFlyDictionary.put(((Object[]) columnValues)[row]); } @@ -200,7 +208,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.getId(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = onTheFlyDictionary.getId(new ByteArray(((byte[][]) columnValues)[row])); + keyValue = getBytesValueId(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); } else { keyValue = onTheFlyDictionary.getId(((Object[]) columnValues)[row]); } @@ -242,7 +250,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.put(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = onTheFlyDictionary.put(new ByteArray(((byte[][]) columnValues)[row])); + keyValue = putBytesValue(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); } else { keyValue = onTheFlyDictionary.put(((Object[]) columnValues)[row]); } @@ -263,7 +271,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.getId(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = onTheFlyDictionary.getId(new ByteArray(((byte[][]) columnValues)[row])); + keyValue = getBytesValueId(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); } else { keyValue = onTheFlyDictionary.getId(((Object[]) columnValues)[row]); } @@ -310,7 +318,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { } else { ValueToIdMap onTheFlyDictionary = _onTheFlyDictionaries[i]; if (_isSingleValueExpressions[i]) { - switch (_storedTypes[i]) { + switch (_dataTypes[i]) { case INT: int[] intValues = blockValSet.getIntValuesSV(); for (int j = 0; j < numDocs; j++) { @@ -341,6 +349,12 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { keys[j][i] = new int[]{onTheFlyDictionary.put(stringValues[j])}; } break; + case UUID: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int j = 0; j < numDocs; j++) { + keys[j][i] = new int[]{onTheFlyDictionary.put(UuidKey.fromBytes(uuidValues[j]))}; + } + break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int j = 0; j < numDocs; j++) { @@ -349,10 +363,10 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { break; default: throw new IllegalArgumentException( - "Illegal data type for no-dictionary key generator: " + _storedTypes[i]); + "Illegal data type for no-dictionary key generator: " + _dataTypes[i]); } } else { - switch (_storedTypes[i]) { + switch (_dataTypes[i]) { case INT: int[][] intValues = blockValSet.getIntValuesMV(); for (int j = 0; j < numDocs; j++) { @@ -408,9 +422,20 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { keys[j][i] = mvKeys; } break; + case UUID: + byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + for (int j = 0; j < numDocs; j++) { + int mvSize = uuidValues[j].length; + int[] mvKeys = new int[mvSize]; + for (int k = 0; k < mvSize; k++) { + mvKeys[k] = onTheFlyDictionary.put(UuidKey.fromBytes(uuidValues[j][k])); + } + keys[j][i] = mvKeys; + } + break; default: throw new IllegalArgumentException( - "Illegal data type for no-dictionary key generator: " + _storedTypes[i]); + "Illegal data type for no-dictionary key generator: " + _dataTypes[i]); } } } @@ -516,4 +541,19 @@ private Object[] buildKeysFromIds(FixedIntArray keyList) { } return keys; } + + private static int putBytesValue(ValueToIdMap onTheFlyDictionary, byte[][] columnValues, int row, DataType dataType) { + if (dataType == DataType.UUID) { + return onTheFlyDictionary.put(UuidKey.fromBytes(columnValues[row])); + } + return onTheFlyDictionary.put(new ByteArray(columnValues[row])); + } + + private static int getBytesValueId(ValueToIdMap onTheFlyDictionary, byte[][] columnValues, int row, + DataType dataType) { + if (dataType == DataType.UUID) { + return onTheFlyDictionary.getId(UuidKey.fromBytes(columnValues[row])); + } + return onTheFlyDictionary.getId(new ByteArray(columnValues[row])); + } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java index d6b0bfcce935..a0cc335c3d93 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java @@ -41,6 +41,7 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidKey; import org.roaringbitmap.RoaringBitmap; @@ -49,7 +50,11 @@ @SuppressWarnings({"rawtypes", "unchecked"}) public class NoDictionarySingleColumnGroupKeyGenerator implements GroupKeyGenerator { private final ExpressionContext _groupByExpression; - private final DataType _storedType; + /// Group-key dispatch type: stored type of the column, except UUID is preserved as [DataType#UUID] so the + /// group-key map keys on [org.apache.pinot.spi.utils.UuidKey] (two primitive longs) instead of + /// [org.apache.pinot.spi.utils.ByteArray]. For every other logical type this equals + /// `logicalType.getStoredType()`. + private final DataType _dataType; private final Map _groupKeyMap; private final int _globalGroupIdUpperBound; // TODO(nhejazi): Most of the logic between _nullHandlingEnabled=true/false is not sharable, so consider making a @@ -66,8 +71,10 @@ public NoDictionarySingleColumnGroupKeyGenerator(BaseProjectOperator projectO @Nullable Map groupByExpressionSizesFromPredicates) { _groupByExpression = groupByExpression; ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpression); - _storedType = columnContext.getDataType().getStoredType(); - _groupKeyMap = createGroupKeyMap(_storedType); + DataType logicalType = columnContext.getDataType(); + // Normalize to stored type, but preserve UUID so it uses UuidKey instead of ByteArray + _dataType = logicalType == DataType.UUID ? DataType.UUID : logicalType.getStoredType(); + _groupKeyMap = createGroupKeyMap(_dataType); if (groupByExpressionSizesFromPredicates != null) { Integer size = groupByExpressionSizesFromPredicates.get(groupByExpression); _globalGroupIdUpperBound = size != null ? Math.min(size, numGroupsLimit) : numGroupsLimit; @@ -95,7 +102,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } int numDocs = valueBlock.getNumDocs(); - switch (_storedType) { + switch (_dataType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -132,6 +139,12 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { groupKeys[i] = getKeyForValue(stringValues[i]); } break; + case UUID: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < numDocs; i++) { + groupKeys[i] = getKeyForValue(UuidKey.fromBytes(uuidValues[i])); + } + break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -139,7 +152,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); } } @@ -149,7 +162,7 @@ public void generateKeysForBlockNullHandlingEnabled(ValueBlock valueBlock, int[] BlockValSet blockValSet = valueBlock.getBlockValueSet(_groupByExpression); int numDocs = valueBlock.getNumDocs(); - switch (_storedType) { + switch (_dataType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); if (nullBitmap.getCardinality() < numDocs) { @@ -210,6 +223,16 @@ public void generateKeysForBlockNullHandlingEnabled(ValueBlock valueBlock, int[] Arrays.fill(groupKeys, 0, numDocs, getKeyForValue((String) null)); } break; + case UUID: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + if (nullBitmap.getCardinality() < numDocs) { + for (int i = 0; i < numDocs; i++) { + groupKeys[i] = getKeyForValue(nullBitmap.contains(i) ? null : UuidKey.fromBytes(uuidValues[i])); + } + } else if (numDocs > 0) { + Arrays.fill(groupKeys, 0, numDocs, getKeyForValue((UuidKey) null)); + } + break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (nullBitmap.getCardinality() < numDocs) { @@ -221,7 +244,7 @@ public void generateKeysForBlockNullHandlingEnabled(ValueBlock valueBlock, int[] } break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); } } @@ -256,6 +279,10 @@ private Map createGroupKeyMap(DataType keyType) { Object2IntOpenHashMap stringMap = new Object2IntOpenHashMap<>(); stringMap.defaultReturnValue(INVALID_ID); return stringMap; + case UUID: + Object2IntOpenHashMap uuidMap = new Object2IntOpenHashMap<>(); + uuidMap.defaultReturnValue(INVALID_ID); + return uuidMap; case BYTES: Object2IntOpenHashMap bytesMap = new Object2IntOpenHashMap<>(); bytesMap.defaultReturnValue(INVALID_ID); @@ -271,7 +298,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { BlockValSet blockValSet = valueBlock.getBlockValueSet(_groupByExpression); if (_isSingleValueExpression) { - switch (_storedType) { + switch (_dataType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -302,6 +329,12 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { groupKeys[i] = new int[]{getKeyForValue(stringValues[i])}; } break; + case UUID: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < numDocs; i++) { + groupKeys[i] = new int[]{getKeyForValue(UuidKey.fromBytes(uuidValues[i]))}; + } + break; case BYTES: byte[][] byteValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -309,10 +342,10 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { } break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); } } else { - switch (_storedType) { + switch (_dataType) { case INT: int[][] intValues = blockValSet.getIntValuesMV(); for (int i = 0; i < numDocs; i++) { @@ -368,8 +401,19 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { groupKeys[i] = mvKeys; } break; + case UUID: + byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < numDocs; i++) { + int mvSize = uuidValues[i].length; + int[] mvKeys = new int[mvSize]; + for (int j = 0; j < mvSize; j++) { + mvKeys[j] = getKeyForValue(UuidKey.fromBytes(uuidValues[i][j])); + } + groupKeys[i] = mvKeys; + } + break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); } } } @@ -381,7 +425,7 @@ public int getCurrentGroupKeyUpperBound() { @Override public Iterator getGroupKeys() { - switch (_storedType) { + switch (_dataType) { case INT: return new IntGroupKeyIterator((Int2IntOpenHashMap) _groupKeyMap, _groupIdForNullValue); case LONG: @@ -393,7 +437,8 @@ public Iterator getGroupKeys() { case BIG_DECIMAL: case STRING: case BYTES: - return new ObjectGroupKeyIterator((Object2IntOpenHashMap) _groupKeyMap); + case UUID: + return new ObjectGroupKeyIterator((Object2IntOpenHashMap) _groupKeyMap, _dataType); default: throw new IllegalStateException(); } @@ -485,6 +530,16 @@ private int getKeyForValue(ByteArray value) { return groupId; } + private int getKeyForValue(UuidKey value) { + Object2IntMap map = (Object2IntMap) _groupKeyMap; + int groupId = map.getInt(value); + if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { + groupId = _numGroups++; + map.put(value, groupId); + } + return groupId; + } + private static class IntGroupKeyIterator implements Iterator { final Iterator _iterator; final GroupKey _groupKey; @@ -633,10 +688,12 @@ public void remove() { private static class ObjectGroupKeyIterator implements Iterator { final ObjectIterator _iterator; final GroupKey _groupKey; + final DataType _dataType; - ObjectGroupKeyIterator(Object2IntOpenHashMap objectMap) { + ObjectGroupKeyIterator(Object2IntOpenHashMap objectMap, DataType dataType) { _iterator = objectMap.object2IntEntrySet().fastIterator(); _groupKey = new GroupKey(); + _dataType = dataType; } @Override @@ -648,7 +705,11 @@ public boolean hasNext() { public GroupKey next() { Object2IntMap.Entry entry = _iterator.next(); _groupKey._groupId = entry.getIntValue(); - _groupKey._keys = new Object[]{entry.getKey()}; + Object key = entry.getKey(); + if (_dataType == DataType.UUID && key != null) { + key = ((UuidKey) key).toByteArray(); + } + _groupKey._keys = new Object[]{key}; return _groupKey; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java new file mode 100644 index 000000000000..845ad21c37a9 --- /dev/null +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java @@ -0,0 +1,59 @@ +/** + * 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.groupby.utils; + +import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; +import java.util.ArrayList; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidKey; + + +/// Implementation of [ValueToIdMap] for Pinot's logical UUID type. +public class UuidToIdMap implements ValueToIdMap { + private final Object2IntOpenHashMap _valueToIdMap; + private final ArrayList _idToValueMap; + + public UuidToIdMap() { + _valueToIdMap = new Object2IntOpenHashMap<>(); + _valueToIdMap.defaultReturnValue(INVALID_KEY); + _idToValueMap = new ArrayList<>(); + } + + @Override + public int put(Object value) { + UuidKey uuidKey = UuidKey.fromObject(value); + int id = _valueToIdMap.getInt(uuidKey); + if (id == INVALID_KEY) { + id = _valueToIdMap.size(); + _valueToIdMap.put(uuidKey, id); + _idToValueMap.add(uuidKey.toByteArray()); + } + return id; + } + + @Override + public int getId(Object value) { + return _valueToIdMap.getInt(UuidKey.fromObject(value)); + } + + @Override + public Object get(int id) { + return _idToValueMap.get(id); + } +} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java index 4ce47caa196b..a15f9bf9a263 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java @@ -36,6 +36,8 @@ public static ValueToIdMap get(DataType dataType) { return new FloatToIdMap(); case DOUBLE: return new DoubleToIdMap(); + case UUID: + return new UuidToIdMap(); default: return new ObjectToIdMap(); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java index 386648b64f5b..42649c9bd28a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java @@ -287,12 +287,13 @@ private ResultTable toResultTableWithOrderBy() { rows = new ArrayList<>(numValues); addRows(sortedValues, numValues, rows); } + formatRows(rows); return new ResultTable(_dataSchema, rows); } private static void addRows(ByteArray[] values, int length, List rows) { for (int i = 0; i < length; i++) { - rows.add(new Object[]{values[i].toHexString()}); + rows.add(new Object[]{values[i]}); } } @@ -308,12 +309,23 @@ private ResultTable toResultTableWithoutOrderBy() { rows = new ArrayList<>(numValues); addRows(_valueSet, rows); } + formatRows(rows); return new ResultTable(_dataSchema, rows); } private static void addRows(HashSet values, List rows) { for (ByteArray value : values) { - rows.add(new Object[]{value.toHexString()}); + rows.add(new Object[]{value}); + } + } + + private void formatRows(List rows) { + DataSchema.ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); + for (Object[] row : rows) { + Object value = row[0]; + if (value != null) { + row[0] = columnDataType.convertAndFormat(value); + } } } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java index c1add8ca7082..b0abf23afb08 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java @@ -530,6 +530,7 @@ private Object getConvertedKey(DataTable dataTable, ColumnDataType columnDataTyp case JSON: return dataTable.getString(rowId, colId); case BYTES: + case UUID: return dataTable.getBytes(rowId, colId).getBytes(); default: throw new IllegalStateException("Illegal column data type in group key: " + columnDataType); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java index 2520115affa0..71feab4aa979 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java @@ -20,13 +20,19 @@ import com.clearspring.analytics.stream.cardinality.HyperLogLog; import java.util.BitSet; +import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.RequestContextUtils; +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.segment.spi.Constants; import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; @@ -55,6 +61,96 @@ public void testCanUseStarTreeDefaultLog2m() { Assert.assertFalse(function.canUseStarTree(Map.of(Constants.HLL_LOG2M_KEY, "16"))); } + /// Regression: UUID columns have storedType=BYTES, but a UUID value is a logical scalar, not a serialized + /// HyperLogLog. The aggregator must route UUID columns through the same content-hash path as STRING (offering + /// canonical UUID strings) instead of trying to deserialize each 16-byte value as an HLL. + @Test + public void testAggregateOnUuidColumnOffersCanonicalStringsAndProducesExactDistinctCount() { + ExpressionContext expression = RequestContextUtils.getExpression("uuidCol"); + DistinctCountHLLAggregationFunction function = new DistinctCountHLLAggregationFunction(List.of(expression)); + + // Three distinct UUIDs across six rows; the same UUID repeats twice on rows 0/3, 1/4, 2/5. + String[] uuidStrings = new String[]{ + "550e8400-e29b-41d4-a716-446655440000", + "12345678-1234-1234-1234-1234567890ab", + "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12", + "550e8400-e29b-41d4-a716-446655440000", + "12345678-1234-1234-1234-1234567890ab", + "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12" + }; + + // Stub the BYTES fetch (raw 16-byte values) — the production path converts bytes to canonical strings + // itself because ProjectionBlockValSet.getStringValuesSV() would render stored BYTES as bare hex. + byte[][] uuidBytes = new byte[uuidStrings.length][]; + for (int i = 0; i < uuidStrings.length; i++) { + uuidBytes[i] = UuidUtils.toBytes(uuidStrings[i]); + } + BlockValSet uuidBlockValSet = mock(BlockValSet.class); + when(uuidBlockValSet.getValueType()).thenReturn(DataType.UUID); + when(uuidBlockValSet.getBytesValuesSV()).thenReturn(uuidBytes); + when(uuidBlockValSet.isSingleValue()).thenReturn(true); + when(uuidBlockValSet.getDictionary()).thenReturn(null); + + Map blockValSetMap = new HashMap<>(); + blockValSetMap.put(expression, uuidBlockValSet); + + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(uuidStrings.length, resultHolder, blockValSetMap); + + Object intermediate = function.extractAggregationResult(resultHolder); + Assert.assertTrue(intermediate instanceof HyperLogLog, + "Intermediate result must be a HyperLogLog, not a dictionary bitmap"); + long cardinality = ((HyperLogLog) intermediate).cardinality(); + Assert.assertEquals(cardinality, 3L, + "HLL cardinality must equal the 3 distinct UUIDs; got " + cardinality); + } + + /// Cross-type consistency: DISTINCTCOUNTHLL(uuidCol) must produce the same HLL cardinality as + /// DISTINCTCOUNTHLL(stringRepresentationOfSameUuids). Locks in the design contract that UUID columns are + /// hashed as canonical UUID strings. + @Test + public void testUuidDistinctCountHllMatchesStringDistinctCountHll() { + String[] uuidStrings = new String[]{ + "550e8400-e29b-41d4-a716-446655440000", + "12345678-1234-1234-1234-1234567890ab", + "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12" + }; + + long uuidHllCardinality = computeHllCardinality(uuidStrings, DataType.UUID); + long stringHllCardinality = computeHllCardinality(uuidStrings, DataType.STRING); + + Assert.assertEquals(uuidHllCardinality, stringHllCardinality, + "DISTINCTCOUNTHLL(uuidCol) must match DISTINCTCOUNTHLL(CAST(uuidCol AS STRING))"); + } + + private long computeHllCardinality(String[] values, DataType valueType) { + ExpressionContext expression = RequestContextUtils.getExpression("col"); + DistinctCountHLLAggregationFunction function = new DistinctCountHLLAggregationFunction(List.of(expression)); + + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(valueType); + if (valueType == DataType.UUID) { + // UUID path fetches raw bytes and converts to canonical form itself (projection string fetch returns hex) + byte[][] uuidBytes = new byte[values.length][]; + for (int i = 0; i < values.length; i++) { + uuidBytes[i] = UuidUtils.toBytes(values[i]); + } + when(blockValSet.getBytesValuesSV()).thenReturn(uuidBytes); + } else { + when(blockValSet.getStringValuesSV()).thenReturn(values); + } + when(blockValSet.isSingleValue()).thenReturn(true); + when(blockValSet.getDictionary()).thenReturn(null); + + Map blockValSetMap = new HashMap<>(); + blockValSetMap.put(expression, blockValSet); + + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(values.length, resultHolder, blockValSetMap); + Object intermediate = function.extractAggregationResult(resultHolder); + return ((HyperLogLog) intermediate).cardinality(); + } + @Test public void testCanUseStarTreeCustomLog2m() { DistinctCountHLLAggregationFunction function = new DistinctCountHLLAggregationFunction( diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java index 2ed7bff83f1b..2192848feb3d 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java @@ -27,6 +27,7 @@ import java.util.List; import java.util.Random; import java.util.Set; +import java.util.UUID; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; import org.apache.commons.lang3.StringUtils; @@ -51,6 +52,7 @@ import org.apache.pinot.spi.utils.BytesUtils; import org.apache.pinot.spi.utils.CommonConstants.Server; import org.apache.pinot.spi.utils.ReadMode; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; @@ -75,12 +77,19 @@ public class NoDictionaryGroupKeyGeneratorTest { private static final String STRING_COLUMN = "stringColumn"; private static final String BYTES_COLUMN = "bytesColumn"; private static final String BYTES_DICT_COLUMN = "bytesDictColumn"; + private static final String UUID_COLUMN = "uuidColumn"; + private static final String BOOLEAN_COLUMN = "booleanColumn"; + private static final String TIMESTAMP_COLUMN = "timestampColumn"; + private static final String UUID_DICT_COLUMN = "uuidDictColumn"; private static final List COLUMNS = Arrays.asList(INT_COLUMN, LONG_COLUMN, FLOAT_COLUMN, DOUBLE_COLUMN, STRING_COLUMN, BYTES_COLUMN, - BYTES_DICT_COLUMN); + BYTES_DICT_COLUMN, UUID_COLUMN, BOOLEAN_COLUMN, TIMESTAMP_COLUMN, UUID_DICT_COLUMN); private static final int NUM_COLUMNS = COLUMNS.size(); + private static final Set UUID_COLUMNS = Set.of(UUID_COLUMN, UUID_DICT_COLUMN); private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) - .setNoDictionaryColumns(COLUMNS.subList(0, NUM_COLUMNS - 1)).build(); + .setNoDictionaryColumns( + List.of(INT_COLUMN, LONG_COLUMN, FLOAT_COLUMN, DOUBLE_COLUMN, STRING_COLUMN, BYTES_COLUMN, UUID_COLUMN, + BOOLEAN_COLUMN, TIMESTAMP_COLUMN)).build(); private static final Schema SCHEMA = new Schema.SchemaBuilder().addSingleValueDimension(INT_COLUMN, FieldSpec.DataType.INT) .addSingleValueDimension(LONG_COLUMN, FieldSpec.DataType.LONG) @@ -88,7 +97,11 @@ public class NoDictionaryGroupKeyGeneratorTest { .addSingleValueDimension(DOUBLE_COLUMN, FieldSpec.DataType.DOUBLE) .addSingleValueDimension(STRING_COLUMN, FieldSpec.DataType.STRING) .addSingleValueDimension(BYTES_COLUMN, FieldSpec.DataType.BYTES) - .addSingleValueDimension(BYTES_DICT_COLUMN, FieldSpec.DataType.BYTES).build(); + .addSingleValueDimension(BYTES_DICT_COLUMN, FieldSpec.DataType.BYTES) + .addSingleValueDimension(UUID_COLUMN, FieldSpec.DataType.UUID) + .addSingleValueDimension(BOOLEAN_COLUMN, FieldSpec.DataType.BOOLEAN) + .addSingleValueDimension(TIMESTAMP_COLUMN, FieldSpec.DataType.TIMESTAMP) + .addSingleValueDimension(UUID_DICT_COLUMN, FieldSpec.DataType.UUID).build(); private static final int NUM_RECORDS = 1000; private static final int NUM_UNIQUE_RECORDS = 100; @@ -129,6 +142,19 @@ public void setUp() record.putValue(BYTES_DICT_COLUMN, bytesValue); values[5] = BytesUtils.toHexString(bytesValue); values[6] = values[5]; + byte[] uuidBytes = UuidUtils.toBytes(new UUID(RANDOM.nextLong(), RANDOM.nextLong())); + record.putValue(UUID_COLUMN, uuidBytes); + values[7] = UuidUtils.toString(uuidBytes); + // BOOLEAN stored as INT (0/1) — exercises the logical→stored-type normalization fix + int boolIntValue = RANDOM.nextBoolean() ? 1 : 0; + record.putValue(BOOLEAN_COLUMN, boolIntValue); + values[8] = Integer.toString(boolIntValue); + // TIMESTAMP stored as LONG — exercises the logical→stored-type normalization fix + long timestampValue = Math.abs(RANDOM.nextLong()); + record.putValue(TIMESTAMP_COLUMN, timestampValue); + values[9] = Long.toString(timestampValue); + record.putValue(UUID_DICT_COLUMN, uuidBytes); + values[10] = values[7]; for (int j = 0; j < NUM_RECORDS / NUM_UNIQUE_RECORDS; j++) { records.add(record); } @@ -173,9 +199,12 @@ public void testMultiColumnGroupKeyGenerator() { testGroupKeyGenerator(new int[]{0, 1}); testGroupKeyGenerator(new int[]{2, 3}); testGroupKeyGenerator(new int[]{4, 5}); + testGroupKeyGenerator(new int[]{7, 10}); + testGroupKeyGenerator(new int[]{8, 9}); testGroupKeyGenerator(new int[]{1, 2, 3}); testGroupKeyGenerator(new int[]{4, 5, 0}); - testGroupKeyGenerator(new int[]{5, 4, 3, 2, 1, 0}); + testGroupKeyGenerator(new int[]{7, 5, 4}); + testGroupKeyGenerator(new int[]{7, 5, 4, 3, 2, 1, 0}); } /// Tests multi-column group key generator when at least one column as dictionary, and others don't. @@ -212,7 +241,7 @@ private void testGroupKeyGenerator(int[] groupByColumnIndexes) { Iterator groupKeys = groupKeyGenerator.getGroupKeys(); while (groupKeys.hasNext()) { GroupKeyGenerator.GroupKey groupKey = groupKeys.next(); - assertTrue(expectedGroupKeys.contains(getActualGroupKey(groupKey._keys))); + assertTrue(expectedGroupKeys.contains(getActualGroupKey(groupKey._keys, groupByColumnIndexes))); } } @@ -234,13 +263,18 @@ private Set getExpectedGroupKeys(int[] groupByColumnIndexes) { return groupKeys; } - private String getActualGroupKey(Object[] groupKeys) { + private String getActualGroupKey(Object[] groupKeys, int[] groupByColumnIndexes) { StringBuilder stringBuilder = new StringBuilder(); for (int i = 0; i < groupKeys.length; i++) { if (i > 0) { stringBuilder.append(GroupKeyGenerator.DELIMITER); } - stringBuilder.append(groupKeys[i]); + int columnIndex = groupByColumnIndexes[i]; + if (UUID_COLUMNS.contains(COLUMNS.get(columnIndex))) { + stringBuilder.append(UuidUtils.toString(((org.apache.pinot.spi.utils.ByteArray) groupKeys[i]).getBytes())); + } else { + stringBuilder.append(groupKeys[i]); + } } return stringBuilder.toString(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java new file mode 100644 index 000000000000..6ca1c2298cb3 --- /dev/null +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java @@ -0,0 +1,70 @@ +/** + * 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.distinct.table; + +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.common.request.context.OrderByExpressionContext; +import org.apache.pinot.common.response.broker.ResultTable; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.UuidUtils; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +/// Tests for [BytesDistinctTable]. +public class BytesDistinctTableTest { + private static final String UUID_COLUMN = "uuidCol"; + private static final String UUID_VALUE_1 = "550e8400-e29b-41d4-a716-446655440000"; + private static final String UUID_VALUE_2 = "550e8400-e29b-41d4-a716-446655440001"; + + @Test + public void testToResultTableFormatsUuidAndBytesWithoutOrderBy() { + BytesDistinctTable uuidTable = new BytesDistinctTable( + new DataSchema(new String[]{UUID_COLUMN}, new ColumnDataType[]{ColumnDataType.UUID}), 10, false, null); + uuidTable.addUnbounded(new ByteArray(UuidUtils.toBytes(UUID_VALUE_1))); + + ResultTable uuidResultTable = uuidTable.toResultTable(); + assertEquals(uuidResultTable.getRows().get(0)[0], UUID_VALUE_1); + + byte[] bytesValue = new byte[]{0x01, 0x23, 0x45}; + BytesDistinctTable bytesTable = new BytesDistinctTable( + new DataSchema(new String[]{"bytesCol"}, new ColumnDataType[]{ColumnDataType.BYTES}), 10, false, null); + bytesTable.addUnbounded(new ByteArray(bytesValue)); + + ResultTable bytesResultTable = bytesTable.toResultTable(); + assertEquals(bytesResultTable.getRows().get(0)[0], BytesUtils.toHexString(bytesValue)); + } + + @Test + public void testToResultTableFormatsUuidWithOrderBy() { + BytesDistinctTable uuidTable = new BytesDistinctTable( + new DataSchema(new String[]{UUID_COLUMN}, new ColumnDataType[]{ColumnDataType.UUID}), 10, false, + new OrderByExpressionContext(ExpressionContext.forIdentifier(UUID_COLUMN), true)); + uuidTable.addUnbounded(new ByteArray(UuidUtils.toBytes(UUID_VALUE_2))); + uuidTable.addUnbounded(new ByteArray(UuidUtils.toBytes(UUID_VALUE_1))); + + ResultTable resultTable = uuidTable.toResultTable(); + assertEquals(resultTable.getRows().get(0)[0], UUID_VALUE_1); + assertEquals(resultTable.getRows().get(1)[0], UUID_VALUE_2); + } +} From 70b28aa862dd04b6b89ac16f2146c796845c8304 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sat, 8 Aug 2026 15:24:01 -0700 Subject: [PATCH 02/19] Return the converted UUID form for group keys in the DataTable reduce path getConvertedKey had `case UUID` falling through to BYTES, returning the raw byte[]. The other reduce path converts group keys via ColumnDataType#convert, and UUID is the one type whose converted form is not its stored bytes -- it yields a java.util.UUID. PredicateRowMatcher casts that directly (see #18872), so the byte[] made GROUP BY ... HAVING over a UUID column fail with "ClassCastException: class [B cannot be cast to class java.util.UUID". Delegating to columnDataType.convert(...) keeps the two paths identical by construction rather than by duplicated knowledge. Covered by a new UuidAggregationTest integration test rather than a unit test. The unit-level BaseQueriesTest harness cannot reach this code, which is why it was uncovered; a query-level test goes through the real broker reduce. Verified by reverting the fix: testGroupByUuidColumnWithHaving and testGroupByUuidColumnWithHavingReturningFinalResult both fail with the ClassCastException and both pass with it. The test also covers GROUP BY key rendering, DISTINCT (BytesDistinctTable no longer hard-codes hex) and DISTINCTCOUNT / DISTINCTCOUNTHLL / DISTINCTCOUNTBITMAP over a UUID column. --- .../query/reduce/GroupByDataTableReducer.java | 8 +- .../tests/custom/UuidAggregationTest.java | 209 ++++++++++++++++++ 2 files changed, 216 insertions(+), 1 deletion(-) create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java index b0abf23afb08..be757ecd9770 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java @@ -530,8 +530,14 @@ private Object getConvertedKey(DataTable dataTable, ColumnDataType columnDataTyp case JSON: return dataTable.getString(rowId, colId); case BYTES: - case UUID: return dataTable.getBytes(rowId, colId).getBytes(); + case UUID: + // Deliberately delegated to ColumnDataType#convert rather than falling through to BYTES. The other reduce + // path (reduceWithIndexedTable) converts group keys with exactly that method, and UUID is the one type + // whose converted form is not its stored bytes -- it yields a java.util.UUID. PredicateRowMatcher casts + // directly on that, so returning the raw byte[] here makes GROUP BY ... HAVING on a UUID column throw + // ClassCastException. Delegating keeps the two paths identical by construction. + return columnDataType.convert(dataTable.getBytes(rowId, colId)); default: throw new IllegalStateException("Illegal column data type in group key: " + columnDataType); } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java new file mode 100644 index 000000000000..016ed27e3606 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java @@ -0,0 +1,209 @@ +/** + * 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.integration.tests.custom; + +import com.fasterxml.jackson.databind.JsonNode; +import java.io.File; +import java.util.ArrayList; +import java.util.List; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +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.utils.builder.TableConfigBuilder; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +/// End-to-end coverage for aggregating, grouping and de-duplicating a UUID column. +/// +/// These run through a real broker reduce, which is the point: the group-key conversion in +/// `GroupByDataTableReducer#getConvertedKey` is only reachable when the broker reduces a *single* data table, and the +/// unit-level `BaseQueriesTest` harness always reduces two. A `case UUID` there that returns the stored `byte[]` +/// instead of the converted `java.util.UUID` makes `GROUP BY ... HAVING` over a UUID column fail with +/// `ClassCastException: class [B cannot be cast to class java.util.UUID`, and only a query-level test catches it. +@Test(suiteName = "CustomClusterIntegrationTest") +public class UuidAggregationTest extends CustomDataQueryClusterIntegrationTest { + private static final String TABLE_NAME = "UuidAggregationTest"; + private static final String UUID_COLUMN = "uuidColumn"; + private static final String UUID_0 = "550e8400-e29b-41d4-a716-446655440000"; + private static final String UUID_0_HEX = "550e8400e29b41d4a716446655440000"; + private static final String UUID_1 = "550e8400-e29b-41d4-a716-446655440001"; + private static final String UUID_2 = "550e8400-e29b-41d4-a716-446655440002"; + + /// `UUID_0` appears twice so grouping and distinct are distinguishable from a plain row count. + private static final List ROWS = List.of(UUID_0, UUID_0, UUID_1, UUID_2); + private static final int NUM_DISTINCT = 3; + + @Override + public String getTableName() { + return TABLE_NAME; + } + + @Override + protected long getCountStarResult() { + return ROWS.size(); + } + + @Override + public int getNumAvroFiles() { + return 1; + } + + @Override + public TableConfig createOfflineTableConfig() { + return new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName()).build(); + } + + @Override + public Schema createSchema() { + return new Schema.SchemaBuilder().setSchemaName(getTableName()) + .addSingleValueDimension(UUID_COLUMN, DataType.UUID) + .build(); + } + + @Override + public List createAvroFiles() + throws Exception { + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("uuidRecord", null, null, false); + avroSchema.setFields(List.of(new org.apache.avro.Schema.Field(UUID_COLUMN, + org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING), null, null))); + + try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) { + DataFileWriter writer = avroFilesAndWriters.getWriters().get(0); + for (String uuid : ROWS) { + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(UUID_COLUMN, uuid); + writer.append(record); + } + return avroFilesAndWriters.getAvroFiles(); + } + } + + @Test + public void testGroupByUuidColumn() + throws Exception { + setUseMultiStageQueryEngine(false); + JsonNode rows = query( + String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s ORDER BY %s", UUID_COLUMN, getTableName(), UUID_COLUMN, + UUID_COLUMN)); + assertEquals(rows.size(), NUM_DISTINCT, rows.toPrettyString()); + + // Group keys must come back as canonical UUIDs, not hex and not a byte-array rendering. + List keys = new ArrayList<>(); + for (JsonNode row : rows) { + keys.add(row.get(0).asText()); + } + assertEquals(keys, List.of(UUID_0, UUID_1, UUID_2), rows.toPrettyString()); + assertEquals(rows.get(0).get(1).asLong(), 2, rows.toPrettyString()); + assertEquals(rows.get(1).get(1).asLong(), 1, rows.toPrettyString()); + } + + /// The regression that motivated this class: `GROUP BY` a UUID column with a `HAVING` predicate on that same + /// column runs the group key through `getConvertedKey` and then straight into `PredicateRowMatcher`, which casts + /// to `java.util.UUID`. + @Test + public void testGroupByUuidColumnWithHaving() + throws Exception { + setUseMultiStageQueryEngine(false); + JsonNode rows = query( + String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s HAVING %s = '%s'", UUID_COLUMN, getTableName(), + UUID_COLUMN, UUID_COLUMN, UUID_0_HEX)); + assertEquals(rows.size(), 1, rows.toPrettyString()); + assertEquals(rows.get(0).get(0).asText(), UUID_0, rows.toPrettyString()); + assertEquals(rows.get(0).get(1).asLong(), 2, rows.toPrettyString()); + + // Same thing via an explicit CAST of the canonical form. + rows = query(String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s HAVING %s = CAST('%s' AS UUID)", UUID_COLUMN, + getTableName(), UUID_COLUMN, UUID_COLUMN, UUID_1)); + assertEquals(rows.size(), 1, rows.toPrettyString()); + assertEquals(rows.get(0).get(0).asText(), UUID_1, rows.toPrettyString()); + assertEquals(rows.get(0).get(1).asLong(), 1, rows.toPrettyString()); + } + + /// The `serverReturnFinalResult` variant. This is the option that gates the single-data-table branch in + /// `GroupByDataTableReducer#reduceAndSetResults` (`isServerReturnFinalResult() && dataTables.size() == 1`), so it + /// covers a different reduce branch from the test above. Verified by reverting the `getConvertedKey` fix: both + /// this and the plain `GROUP BY ... HAVING` test fail with + /// `ClassCastException: class [B cannot be cast to class java.util.UUID`, and both pass with it. + @Test + public void testGroupByUuidColumnWithHavingReturningFinalResult() + throws Exception { + setUseMultiStageQueryEngine(false); + JsonNode rows = query( + String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s HAVING %s = '%s' OPTION(serverReturnFinalResult=true)", + UUID_COLUMN, getTableName(), UUID_COLUMN, UUID_COLUMN, UUID_0_HEX)); + assertEquals(rows.size(), 1, rows.toPrettyString()); + assertEquals(rows.get(0).get(0).asText(), UUID_0, rows.toPrettyString()); + assertEquals(rows.get(0).get(1).asLong(), 2, rows.toPrettyString()); + } + + /// Group keys must also render canonically on the `getConvertedKey` path, not as hex. + @Test + public void testGroupByUuidColumnReturningFinalResult() + throws Exception { + setUseMultiStageQueryEngine(false); + JsonNode rows = query( + String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s ORDER BY %s OPTION(serverReturnFinalResult=true)", + UUID_COLUMN, getTableName(), UUID_COLUMN, UUID_COLUMN)); + assertEquals(rows.size(), NUM_DISTINCT, rows.toPrettyString()); + List keys = new ArrayList<>(); + for (JsonNode row : rows) { + keys.add(row.get(0).asText()); + } + assertEquals(keys, List.of(UUID_0, UUID_1, UUID_2), rows.toPrettyString()); + } + + @Test + public void testDistinctOnUuidColumn() + throws Exception { + setUseMultiStageQueryEngine(false); + JsonNode rows = + query(String.format("SELECT DISTINCT %s FROM %s ORDER BY %s", UUID_COLUMN, getTableName(), UUID_COLUMN)); + assertEquals(rows.size(), NUM_DISTINCT, rows.toPrettyString()); + + // BytesDistinctTable used to hard-code toHexString(); a UUID column must render canonically. + List values = new ArrayList<>(); + for (JsonNode row : rows) { + values.add(row.get(0).asText()); + } + assertEquals(values, List.of(UUID_0, UUID_1, UUID_2), rows.toPrettyString()); + } + + @Test + public void testDistinctCountOnUuidColumn() + throws Exception { + setUseMultiStageQueryEngine(false); + for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTBITMAP")) { + JsonNode rows = query(String.format("SELECT %s(%s) FROM %s", function, UUID_COLUMN, getTableName())); + assertEquals(rows.get(0).get(0).asLong(), NUM_DISTINCT, function + ": " + rows.toPrettyString()); + } + } + + private JsonNode query(String sql) + throws Exception { + JsonNode response = postQuery(sql); + assertTrue(response.path("exceptions").isEmpty(), sql + " -> " + response.toPrettyString()); + return response.path("resultTable").path("rows"); + } +} From 798bf216e97077d31ddb3adbfad6b28a1ee095c7 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sat, 8 Aug 2026 16:40:41 -0700 Subject: [PATCH 03/19] Cast directly to UuidKey in UuidToIdMap Both callers key on UuidKey already (NoDictionary{Single,Multi}ColumnGroupKey Generator, via UuidKey.fromBytes), so UuidKey.fromObject accepted five input types where exactly one is ever passed, and ran an instanceof chain per row in the group-by loop. Casting directly matches the sibling maps -- DoubleToIdMap casts to double -- and keeps the input type deterministic, which is what was asked for on the equivalent PredicateRowMatcher branch in #18872. --- .../query/aggregation/groupby/utils/UuidToIdMap.java | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java index 845ad21c37a9..0ee338e97013 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java @@ -35,9 +35,14 @@ public UuidToIdMap() { _idToValueMap = new ArrayList<>(); } + /// Both callers -- [org.apache.pinot.core.query.aggregation.groupby.NoDictionaryMultiColumnGroupKeyGenerator] and + /// [org.apache.pinot.core.query.aggregation.groupby.NoDictionarySingleColumnGroupKeyGenerator] -- key on + /// [UuidKey] already, so this casts directly rather than going through `UuidKey#fromObject`. That matches the + /// sibling maps (e.g. [DoubleToIdMap] casts to `double`) and keeps the per-row `instanceof` chain out of the + /// group-by loop. @Override public int put(Object value) { - UuidKey uuidKey = UuidKey.fromObject(value); + UuidKey uuidKey = (UuidKey) value; int id = _valueToIdMap.getInt(uuidKey); if (id == INVALID_KEY) { id = _valueToIdMap.size(); @@ -49,7 +54,7 @@ public int put(Object value) { @Override public int getId(Object value) { - return _valueToIdMap.getInt(UuidKey.fromObject(value)); + return _valueToIdMap.getInt((UuidKey) value); } @Override From 70ba71b1f89d7607d591f549244cac7285616af6 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 9 Aug 2026 02:06:00 -0700 Subject: [PATCH 04/19] Hash UUID's stored bytes in distinct-count aggregations, not a canonical string The previous version rendered each UUID as its 36-char canonical string so that DISTINCTCOUNTHLL(uuidCol) would equal DISTINCTCOUNTHLL(CAST(uuidCol AS STRING)). No other logical type provides that guarantee: the scan path switches on the stored type, so TIMESTAMP offers its raw millis and BOOLEAN its int, and neither matches a CAST to STRING. The UUID rendering was inventing a cross-type equivalence at the cost of a String allocation per row in the aggregation loop. UUID now hashes its stored 16 bytes. Verified byte[] is content-hashed rather than identity-hashed by both HyperLogLog (clearspring MurmurHash) and UltraLogLogUtils.OBJECT_FUNNEL (putBytes). A minimal guard is still needed at each site, because unlike LONG or INT the stored BYTES type is not a scalar case in this family: the scan path has no `case BYTES`, and the dictionary path reads BYTES as serialized sketch state. - AggregationFunctionUtils: the three UUID blocks are gone; the BYTES guard now excludes UUID so it falls through to the scalar path, which offers dictionary.get(i) -- the stored byte[] -- exactly as the scan path does. - DistinctCountBitmap hashes Arrays.hashCode(bytes). - DistinctCountThetaSketch cannot take scalar bytes (BYTES there means "serialized sketch"), so it surfaces the stored hex rendering instead. Replaces testUuidDistinctCountHllMatchesStringDistinctCountHll, which asserted the invariant being dropped, with one pinning the new behaviour. --- .../function/AggregationFunctionUtils.java | 52 ++++++------------- ...istinctCountBitmapAggregationFunction.java | 8 +-- ...inctCountCPCSketchAggregationFunction.java | 7 ++- .../DistinctCountHLLAggregationFunction.java | 7 ++- ...stinctCountHLLPlusAggregationFunction.java | 7 ++- ...ctCountThetaSketchAggregationFunction.java | 28 +++++----- .../DistinctCountULLAggregationFunction.java | 7 ++- ...stinctCountHLLAggregationFunctionTest.java | 30 ++++++++--- 8 files changed, 67 insertions(+), 79 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index c5edbd429485..6908ef1f19d8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -74,7 +74,6 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.query.QueryThreadContext; import org.apache.pinot.spi.utils.ByteArray; -import org.apache.pinot.spi.utils.UuidUtils; /// The `AggregationFunctionUtils` class provides utility methods for aggregation function. @@ -803,18 +802,11 @@ private static HyperLogLogPlus getDistinctValueHLLPlus(Dictionary dictionary, in private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, DistinctCountHLLAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - // UUID dictionary entries are logical scalar values, not serialized HyperLogLogs. Offer their canonical string - // representation to match the scan-based path and DISTINCTCOUNTHLL(CAST(uuidColumn AS STRING)). - if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.UUID) { - HyperLogLog hll = new HyperLogLog(function.getLog2m()); - int length = dictionary.length(); - for (int i = 0; i < length; i++) { - QueryThreadContext.checkTerminationAndSampleUsagePeriodically(i, explainPlanName); - hll.offer(UuidUtils.toString(dictionary.getBytesValue(i))); - } - return hll; - } - if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { + // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary), but its entries are logical + // scalars, not serialized sketch state. Excluding it here lets it fall through to the scalar path + // below, which offers dictionary.get(i) -- the stored byte[] -- exactly as the scan path does. + if (dataSource.getDataSourceMetadata().getDataType() != FieldSpec.DataType.UUID + && dictionary.getValueType() == FieldSpec.DataType.BYTES) { // Treat BYTES value as serialized HyperLogLog try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); @@ -836,18 +828,11 @@ private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, private static HyperLogLogPlus getDistinctCountHLLPlusResult(DataSource dataSource, DistinctCountHLLPlusAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - // UUID dictionary entries are logical scalar values, not serialized HyperLogLogPluses. Offer their canonical - // string representation to match the scan-based path and DISTINCTCOUNTHLLPLUS(CAST(uuidColumn AS STRING)). - if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.UUID) { - HyperLogLogPlus hllPlus = new HyperLogLogPlus(function.getP(), function.getSp()); - int length = dictionary.length(); - for (int i = 0; i < length; i++) { - QueryThreadContext.checkTerminationAndSampleUsagePeriodically(i, explainPlanName); - hllPlus.offer(UuidUtils.toString(dictionary.getBytesValue(i))); - } - return hllPlus; - } - if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { + // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary), but its entries are logical + // scalars, not serialized sketch state. Excluding it here lets it fall through to the scalar path + // below, which offers dictionary.get(i) -- the stored byte[] -- exactly as the scan path does. + if (dataSource.getDataSourceMetadata().getDataType() != FieldSpec.DataType.UUID + && dictionary.getValueType() == FieldSpec.DataType.BYTES) { // Treat BYTES value as serialized HyperLogLogPlus try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); @@ -889,18 +874,11 @@ private static Object getDistinctCountSmartHLLPlusResult(Dictionary dictionary, private static UltraLogLog getDistinctCountULLResult(DataSource dataSource, DistinctCountULLAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - // UUID dictionary entries are logical scalar values, not serialized UltraLogLogs. Hash their canonical string - // representation to match the scan-based path and DISTINCTCOUNTULL(CAST(uuidColumn AS STRING)). - if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.UUID) { - UltraLogLog ull = UltraLogLog.create(function.getP()); - int length = dictionary.length(); - for (int i = 0; i < length; i++) { - QueryThreadContext.checkTerminationAndSampleUsagePeriodically(i, explainPlanName); - UltraLogLogUtils.hashObject(UuidUtils.toString(dictionary.getBytesValue(i))).ifPresent(ull::add); - } - return ull; - } - if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { + // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary), but its entries are logical + // scalars, not serialized sketch state. Excluding it here lets it fall through to the scalar path + // below, which offers dictionary.get(i) -- the stored byte[] -- exactly as the scan path does. + if (dataSource.getDataSourceMetadata().getDataType() != FieldSpec.DataType.UUID + && dictionary.getValueType() == FieldSpec.DataType.BYTES) { // Treat BYTES value as serialized UltraLogLog and merge try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index 04a3bc0d9b4d..03c10883f36a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.util.Arrays; import java.util.List; import java.util.Map; import javax.annotation.Nullable; @@ -34,7 +35,6 @@ import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -82,7 +82,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); RoaringBitmap bitmap = getValueBitmap(aggregationResultHolder); for (int i = 0; i < length; i++) { - bitmap.add(UuidUtils.toString(uuidBytesValues[i]).hashCode()); + bitmap.add(Arrays.hashCode(uuidBytesValues[i])); } return; } @@ -232,7 +232,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { getValueBitmap(groupByResultHolder, groupKeyArray[i]) - .add(UuidUtils.toString(uuidBytesValues[i]).hashCode()); + .add(Arrays.hashCode(uuidBytesValues[i])); } return; } @@ -384,7 +384,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult if (dataType == DataType.UUID) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - int hash = UuidUtils.toString(uuidBytesValues[i]).hashCode(); + int hash = Arrays.hashCode(uuidBytesValues[i]); for (int groupKey : groupKeysArray[i]) { getValueBitmap(groupByResultHolder, groupKey).add(hash); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index 7111f211e509..96e286aae1e9 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -41,7 +41,6 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; -import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -148,7 +147,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde // Calling getAccumulator here would read the holder slot already occupied by the sketch and fail. CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); for (int i = 0; i < length; i++) { - cpcSketch.update(UuidUtils.toString(uuidBytesValues[i])); + cpcSketch.update(uuidBytesValues[i]); } return; } @@ -231,7 +230,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol if (dataType == DataType.UUID) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(UuidUtils.toString(uuidBytesValues[i])); + getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(uuidBytesValues[i]); } return; } @@ -314,7 +313,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult if (dataType == DataType.UUID && singleValue) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - String canonical = UuidUtils.toString(uuidBytesValues[i]); + byte[] canonical = uuidBytesValues[i]; for (int groupKey : groupKeysArray[i]) { getCpcSketch(groupByResultHolder, groupKey).update(canonical); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java index 8a74f972f7f3..8a837029e977 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -38,7 +38,6 @@ import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; -import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.RoaringBitmap; @@ -94,7 +93,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); HyperLogLog hyperLogLog = getHyperLogLog(aggregationResultHolder); for (int i = 0; i < length; i++) { - hyperLogLog.offer(UuidUtils.toString(uuidBytesValues[i])); + hyperLogLog.offer(uuidBytesValues[i]); } return; } @@ -256,7 +255,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol if (dataType == DataType.UUID) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(UuidUtils.toString(uuidBytesValues[i])); + getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(uuidBytesValues[i]); } return; } @@ -416,7 +415,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult if (dataType == DataType.UUID) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - String canonical = UuidUtils.toString(uuidBytesValues[i]); + byte[] canonical = uuidBytesValues[i]; for (int groupKey : groupKeysArray[i]) { getHyperLogLog(groupByResultHolder, groupKey).offer(canonical); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java index 1c2762670c61..452f33412099 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -37,7 +37,6 @@ import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; -import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -103,7 +102,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(aggregationResultHolder); for (int i = 0; i < length; i++) { - hyperLogLogPlus.offer(UuidUtils.toString(uuidBytesValues[i])); + hyperLogLogPlus.offer(uuidBytesValues[i]); } return; } @@ -258,7 +257,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol if (dataType == DataType.UUID) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(UuidUtils.toString(uuidBytesValues[i])); + getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(uuidBytesValues[i]); } return; } @@ -417,7 +416,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult if (dataType == DataType.UUID) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - String canonical = UuidUtils.toString(uuidBytesValues[i]); + byte[] canonical = uuidBytesValues[i]; for (int groupKey : groupKeysArray[i]) { getHyperLogLogPlus(groupByResultHolder, groupKey).offer(canonical); } 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 67e9c2996ef9..263cf8af0091 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 @@ -55,8 +55,8 @@ import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.Constants; import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.BytesUtils; import org.apache.pinot.spi.utils.CommonConstants; -import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.sql.parsers.CalciteSqlParser; @@ -1237,32 +1237,32 @@ private void extractValues(int length, Map block DataType storedType = dataType.getStoredType(); singleValues[i] = singleValue; // UUID columns are stored as 16-byte BYTES but a UUID value is a logical scalar, not a pre-serialized - // theta sketch. Surface UUID as STRING (canonical UUID form) so the downstream update-sketch path - // matches DISTINCTCOUNTTHETASKETCH(CAST(uuidCol AS STRING)). Without this branch, the function would - // take the serialized-sketch path below and Sketch.wrap would fail on raw 16-byte UUID content. - // NOTE: fetch raw bytes and convert explicitly — for identifier expressions the BlockValSet is a - // ProjectionBlockValSet whose getStringValuesSV() renders stored BYTES as bare hex, not canonical form. + // theta sketch: without this branch the function takes the serialized-sketch path below and Sketch.wrap + // fails on raw 16-byte UUID content. Unlike the other distinct-count functions this one cannot consume the + // stored bytes directly -- DataType.BYTES here means "serialized sketch", with no scalar-bytes mode -- so + // the stored value is surfaced as its hex rendering, the same form used at every other String-typed UUID + // boundary (see PredicateUtils#getStoredValue and the Bloom filter key). if (dataType == DataType.UUID) { valueTypes[i] = DataType.STRING; if (singleValue) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - String[] canonicalValues = new String[length]; + String[] hexValues = new String[length]; for (int j = 0; j < length; j++) { - canonicalValues[j] = UuidUtils.toString(uuidBytesValues[j]); + hexValues[j] = BytesUtils.toHexString(uuidBytesValues[j]); } - valueArrays[i] = canonicalValues; + valueArrays[i] = hexValues; } else { byte[][][] uuidBytesValuesMV = blockValSet.getBytesValuesMV(); - String[][] canonicalValuesMV = new String[length][]; + String[][] hexValuesMV = new String[length][]; for (int j = 0; j < length; j++) { byte[][] row = uuidBytesValuesMV[j]; - String[] canonicalRow = new String[row.length]; + String[] hexRow = new String[row.length]; for (int k = 0; k < row.length; k++) { - canonicalRow[k] = UuidUtils.toString(row[k]); + hexRow[k] = BytesUtils.toHexString(row[k]); } - canonicalValuesMV[j] = canonicalRow; + hexValuesMV[j] = hexRow; } - valueArrays[i] = canonicalValuesMV; + valueArrays[i] = hexValuesMV; } continue; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java index cad535aa2aea..3a9319e00c7c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -38,7 +38,6 @@ import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.CommonConstants; -import org.apache.pinot.spi.utils.UuidUtils; import org.roaringbitmap.PeekableIntIterator; import org.roaringbitmap.RoaringBitmap; @@ -92,7 +91,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); UltraLogLog ull = getULL(aggregationResultHolder); for (int i = 0; i < length; i++) { - UltraLogLogUtils.hashObject(UuidUtils.toString(uuidBytesValues[i])).ifPresent(ull::add); + UltraLogLogUtils.hashObject(uuidBytesValues[i]).ifPresent(ull::add); } return; } @@ -177,7 +176,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); - UltraLogLogUtils.hashObject(UuidUtils.toString(uuidBytesValues[i])).ifPresent(ull::add); + UltraLogLogUtils.hashObject(uuidBytesValues[i]).ifPresent(ull::add); } return; } @@ -267,7 +266,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult if (dataType == DataType.UUID) { byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - String canonical = UuidUtils.toString(uuidBytesValues[i]); + byte[] canonical = uuidBytesValues[i]; for (int groupKey : groupKeysArray[i]) { UltraLogLog ull = getULL(groupByResultHolder, groupKey); UltraLogLogUtils.hashObject(canonical).ifPresent(ull::add); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java index 71feab4aa979..3049fae37358 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java @@ -19,6 +19,7 @@ package org.apache.pinot.core.query.aggregation.function; import com.clearspring.analytics.stream.cardinality.HyperLogLog; +import java.util.Arrays; import java.util.BitSet; import java.util.HashMap; import java.util.List; @@ -32,6 +33,7 @@ import org.apache.pinot.segment.spi.Constants; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; @@ -105,22 +107,34 @@ public void testAggregateOnUuidColumnOffersCanonicalStringsAndProducesExactDisti "HLL cardinality must equal the 3 distinct UUIDs; got " + cardinality); } - /// Cross-type consistency: DISTINCTCOUNTHLL(uuidCol) must produce the same HLL cardinality as - /// DISTINCTCOUNTHLL(stringRepresentationOfSameUuids). Locks in the design contract that UUID columns are - /// hashed as canonical UUID strings. + /// UUID columns hash their **stored bytes**, exactly as TIMESTAMP hashes its stored millis rather than a + /// formatted string. Consequence: DISTINCTCOUNTHLL(uuidCol) does NOT equal + /// DISTINCTCOUNTHLL(CAST(uuidCol AS STRING)) -- and neither does it for TIMESTAMP, so this is the consistent + /// behaviour for a logical type, not a gap. Pinned here so nobody "fixes" it back into a canonical-string + /// rendering, which would reintroduce a per-row String allocation in the aggregation loop. @Test - public void testUuidDistinctCountHllMatchesStringDistinctCountHll() { + public void testUuidDistinctCountHllHashesStoredBytesNotCanonicalString() + throws java.io.IOException { String[] uuidStrings = new String[]{ "550e8400-e29b-41d4-a716-446655440000", "12345678-1234-1234-1234-1234567890ab", "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12" }; - long uuidHllCardinality = computeHllCardinality(uuidStrings, DataType.UUID); - long stringHllCardinality = computeHllCardinality(uuidStrings, DataType.STRING); + // Cardinality is still exact for a small distinct set... + Assert.assertEquals(computeHllCardinality(uuidStrings, DataType.UUID), 3L); - Assert.assertEquals(uuidHllCardinality, stringHllCardinality, - "DISTINCTCOUNTHLL(uuidCol) must match DISTINCTCOUNTHLL(CAST(uuidCol AS STRING))"); + // ...but the sketch is built over the 16 stored bytes, so a HyperLogLog fed the canonical strings differs. + HyperLogLog fromCanonicalStrings = new HyperLogLog(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M); + for (String uuid : uuidStrings) { + fromCanonicalStrings.offer(uuid); + } + HyperLogLog fromStoredBytes = new HyperLogLog(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M); + for (String uuid : uuidStrings) { + fromStoredBytes.offer(UuidUtils.toBytes(uuid)); + } + Assert.assertFalse(Arrays.equals(fromCanonicalStrings.getBytes(), fromStoredBytes.getBytes()), + "stored-bytes and canonical-string sketches are expected to differ"); } private long computeHllCardinality(String[] values, DataType valueType) { From dd51284f9f6830975e2db91d2177eccb76186209 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 9 Aug 2026 03:47:45 -0700 Subject: [PATCH 05/19] Drop UUID special-casing that the stored type already covers Applying the rule that if TIMESTAMP and BIG_DECIMAL need no special handling at a site, UUID does not either -- their stored types carry them, and UUID's should too. Removed: - NoDictionary{Single,Multi}ColumnGroupKeyGenerator, UuidToIdMap and its ValueToIdMapFactory entry. `case BYTES` is already supported there and is structurally identical to the UUID branch: same getBytesValuesSV(), same loop, same map, differing only in wrapping UuidKey vs ByteArray. Both yield ByteArray downstream, which is what ColumnDataType.UUID#convert consumes, so this was a pure key-representation optimization -- two primitive longs instead of a byte[] wrapper -- with no benchmark to justify 5 files and +237/-36. UuidAggregationTest#testGroupByUuidColumn still passes, confirming group keys render identically without it. - AnyValueAggregationFunction. Its switch is on getStoredType(), so TIMESTAMP already collapses to LONG (raw millis) and BOOLEAN to INT. UUID collapsing to BYTES is the consistent behaviour; the branch made it the odd one out. - IntegerTupleSketchAggregationFunction. The UUID branch only produced a friendlier error message; TIMESTAMP gets no such treatment when it is equally unusable there. Kept, because the same rule shows they are needed: - BytesDistinctTable: LongDistinctTable and BigDecimalDistinctTable already render by ColumnDataType, and MultiColumnDistinctTable already calls convertAndFormat. BytesDistinctTable hard-coding toHexString() was the outlier; this brings it in line rather than special-casing UUID. - The distinct-count guards and GroupByDataTableReducer, both verified necessary by probe -- BYTES means "serialized sketch" in one and the group key must convert in the other. 19 files / +958 -80 -> 12 files / +694 -39. --- .../function/AnyValueAggregationFunction.java | 7 -- ...IntegerTupleSketchAggregationFunction.java | 25 +---- ...ictionaryMultiColumnGroupKeyGenerator.java | 68 +++----------- ...ctionarySingleColumnGroupKeyGenerator.java | 91 +++---------------- .../groupby/utils/UuidToIdMap.java | 64 ------------- .../groupby/utils/ValueToIdMapFactory.java | 2 - .../NoDictionaryGroupKeyGeneratorTest.java | 48 ++-------- 7 files changed, 41 insertions(+), 264 deletions(-) delete mode 100644 pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java index 7a3ba85337ca..a649fe6d7a17 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AnyValueAggregationFunction.java @@ -327,13 +327,6 @@ private void ensureResultType(BlockValSet bvs) { if (_resultType != null) { return; } - // Inspect the logical type first so a UUID column reports ColumnDataType.UUID (and the broker renders canonical - // RFC-4122 strings) rather than collapsing to BYTES (which would render hex). All other dispatch keys off the - // stored type, matching the BYTES/STRING storage convention. - if (bvs.getValueType() == FieldSpec.DataType.UUID) { - _resultType = ColumnDataType.UUID; - return; - } switch (bvs.getValueType().getStoredType()) { case INT: _resultType = ColumnDataType.INT; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java index 58b4f7c68607..7591272c0d5b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/IntegerTupleSketchAggregationFunction.java @@ -152,15 +152,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - FieldSpec.DataType dataType = blockValSet.getValueType(); - FieldSpec.DataType storedType = dataType.getStoredType(); - // UUID columns are stored as BYTES but contain raw 16-byte UUID values, not serialized tuple sketches. - // Surface a clear error rather than letting the deserialize step fail with a confusing sketch-format message. - if (dataType == FieldSpec.DataType.UUID) { - throw new IllegalStateException(getType() + " does not accept raw UUID values; pre-serialize the column as " - + "Integer Tuple Sketches first"); - } // Treat BYTES value as serialized Integer Tuple Sketch + FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == FieldSpec.DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { @@ -191,12 +184,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol BlockValSet blockValSet = blockValSetMap.get(_expression); - FieldSpec.DataType dataType = blockValSet.getValueType(); - FieldSpec.DataType storedType = dataType.getStoredType(); - if (dataType == FieldSpec.DataType.UUID) { - throw new IllegalStateException(getType() + " does not accept raw UUID values; pre-serialize the column as " - + "Integer Tuple Sketches first"); - } + // Treat BYTES value as serialized Integer Tuple Sketch + FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); if (storedType == FieldSpec.DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); @@ -220,12 +209,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - FieldSpec.DataType dataType = blockValSet.getValueType(); - FieldSpec.DataType storedType = dataType.getStoredType(); - if (dataType == FieldSpec.DataType.UUID) { - throw new IllegalStateException(getType() + " does not accept raw UUID values; pre-serialize the column as " - + "Integer Tuple Sketches first"); - } + // Treat BYTES value as serialized Integer Tuple Sketch + FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); boolean singleValue = blockValSet.isSingleValue(); if (singleValue && storedType == FieldSpec.DataType.BYTES) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java index f6d3784ebf96..a10de9f47a58 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryMultiColumnGroupKeyGenerator.java @@ -36,7 +36,6 @@ import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.FixedIntArray; -import org.apache.pinot.spi.utils.UuidKey; import org.roaringbitmap.RoaringBitmap; @@ -52,11 +51,7 @@ public class NoDictionaryMultiColumnGroupKeyGenerator implements GroupKeyGenerat private final ExpressionContext[] _groupByExpressions; private final int _numGroupByExpressions; - /// Per-column group-key dispatch type: stored type of each column, except UUID is preserved as - /// [DataType#UUID] so the on-the-fly dictionary keys on [org.apache.pinot.spi.utils.UuidKey] - /// instead of [org.apache.pinot.spi.utils.ByteArray]. For every other logical type this equals - /// `logicalType.getStoredType()`. - private final DataType[] _dataTypes; + private final DataType[] _storedTypes; private final Dictionary[] _dictionaries; private final ValueToIdMap[] _onTheFlyDictionaries; private final Object2IntOpenHashMap _groupKeyMap; @@ -70,7 +65,7 @@ public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOp Map groupByExpressionSizesFromPredicates) { _groupByExpressions = groupByExpressions; _numGroupByExpressions = groupByExpressions.length; - _dataTypes = new DataType[_numGroupByExpressions]; + _storedTypes = new DataType[_numGroupByExpressions]; _dictionaries = new Dictionary[_numGroupByExpressions]; _onTheFlyDictionaries = new ValueToIdMap[_numGroupByExpressions]; _isSingleValueExpressions = new boolean[_numGroupByExpressions]; @@ -80,9 +75,7 @@ public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOp for (int i = 0; i < _numGroupByExpressions; i++) { ExpressionContext groupByExpression = groupByExpressions[i]; ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpression); - DataType logicalType = columnContext.getDataType(); - // Normalize to stored type, but preserve UUID so it uses UuidKey instead of ByteArray. - _dataTypes[i] = logicalType == DataType.UUID ? DataType.UUID : logicalType.getStoredType(); + _storedTypes[i] = columnContext.getDataType().getStoredType(); // Take the dict-id path only when the forward index is dict-encoded. A column with EncodingType.RAW + // dictionaryIndex exposes a Dictionary but BlockValSet#getDictionaryIdsSV throws on its RAW forward // index — fall back to an on-the-fly dictionary on raw values for that case. @@ -91,7 +84,7 @@ public NoDictionaryMultiColumnGroupKeyGenerator(BaseProjectOperator projectOp if (dictionary != null) { _dictionaries[i] = dictionary; } else { - _onTheFlyDictionaries[i] = ValueToIdMapFactory.get(_dataTypes[i]); + _onTheFlyDictionaries[i] = ValueToIdMapFactory.get(_storedTypes[i]); } if (canOptimizeGroupByUpperBound) { Integer size = groupByExpressionSizesFromPredicates.get(groupByExpression); @@ -128,7 +121,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { if (_dictionaries[i] != null) { values[i] = blockValSet.getDictionaryIdsSV(); } else { - switch (_dataTypes[i]) { + switch (_storedTypes[i]) { case INT: values[i] = blockValSet.getIntValuesSV(); break; @@ -148,11 +141,10 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { values[i] = blockValSet.getStringValuesSV(); break; case BYTES: - case UUID: values[i] = blockValSet.getBytesValuesSV(); break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataTypes[i]); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedTypes[i]); } } } @@ -184,7 +176,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.put(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = putBytesValue(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); + keyValue = onTheFlyDictionary.put(new ByteArray(((byte[][]) columnValues)[row])); } else { keyValue = onTheFlyDictionary.put(((Object[]) columnValues)[row]); } @@ -208,7 +200,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.getId(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = getBytesValueId(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); + keyValue = onTheFlyDictionary.getId(new ByteArray(((byte[][]) columnValues)[row])); } else { keyValue = onTheFlyDictionary.getId(((Object[]) columnValues)[row]); } @@ -250,7 +242,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.put(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = putBytesValue(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); + keyValue = onTheFlyDictionary.put(new ByteArray(((byte[][]) columnValues)[row])); } else { keyValue = onTheFlyDictionary.put(((Object[]) columnValues)[row]); } @@ -271,7 +263,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } else if (columnValues instanceof double[]) { keyValue = onTheFlyDictionary.getId(((double[]) columnValues)[row]); } else if (columnValues instanceof byte[][]) { - keyValue = getBytesValueId(onTheFlyDictionary, (byte[][]) columnValues, row, _dataTypes[col]); + keyValue = onTheFlyDictionary.getId(new ByteArray(((byte[][]) columnValues)[row])); } else { keyValue = onTheFlyDictionary.getId(((Object[]) columnValues)[row]); } @@ -318,7 +310,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { } else { ValueToIdMap onTheFlyDictionary = _onTheFlyDictionaries[i]; if (_isSingleValueExpressions[i]) { - switch (_dataTypes[i]) { + switch (_storedTypes[i]) { case INT: int[] intValues = blockValSet.getIntValuesSV(); for (int j = 0; j < numDocs; j++) { @@ -349,12 +341,6 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { keys[j][i] = new int[]{onTheFlyDictionary.put(stringValues[j])}; } break; - case UUID: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int j = 0; j < numDocs; j++) { - keys[j][i] = new int[]{onTheFlyDictionary.put(UuidKey.fromBytes(uuidValues[j]))}; - } - break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int j = 0; j < numDocs; j++) { @@ -363,10 +349,10 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { break; default: throw new IllegalArgumentException( - "Illegal data type for no-dictionary key generator: " + _dataTypes[i]); + "Illegal data type for no-dictionary key generator: " + _storedTypes[i]); } } else { - switch (_dataTypes[i]) { + switch (_storedTypes[i]) { case INT: int[][] intValues = blockValSet.getIntValuesMV(); for (int j = 0; j < numDocs; j++) { @@ -422,20 +408,9 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { keys[j][i] = mvKeys; } break; - case UUID: - byte[][][] uuidValues = blockValSet.getBytesValuesMV(); - for (int j = 0; j < numDocs; j++) { - int mvSize = uuidValues[j].length; - int[] mvKeys = new int[mvSize]; - for (int k = 0; k < mvSize; k++) { - mvKeys[k] = onTheFlyDictionary.put(UuidKey.fromBytes(uuidValues[j][k])); - } - keys[j][i] = mvKeys; - } - break; default: throw new IllegalArgumentException( - "Illegal data type for no-dictionary key generator: " + _dataTypes[i]); + "Illegal data type for no-dictionary key generator: " + _storedTypes[i]); } } } @@ -541,19 +516,4 @@ private Object[] buildKeysFromIds(FixedIntArray keyList) { } return keys; } - - private static int putBytesValue(ValueToIdMap onTheFlyDictionary, byte[][] columnValues, int row, DataType dataType) { - if (dataType == DataType.UUID) { - return onTheFlyDictionary.put(UuidKey.fromBytes(columnValues[row])); - } - return onTheFlyDictionary.put(new ByteArray(columnValues[row])); - } - - private static int getBytesValueId(ValueToIdMap onTheFlyDictionary, byte[][] columnValues, int row, - DataType dataType) { - if (dataType == DataType.UUID) { - return onTheFlyDictionary.getId(UuidKey.fromBytes(columnValues[row])); - } - return onTheFlyDictionary.getId(new ByteArray(columnValues[row])); - } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java index a0cc335c3d93..d6b0bfcce935 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionarySingleColumnGroupKeyGenerator.java @@ -41,7 +41,6 @@ import org.apache.pinot.core.operator.blocks.ValueBlock; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.utils.ByteArray; -import org.apache.pinot.spi.utils.UuidKey; import org.roaringbitmap.RoaringBitmap; @@ -50,11 +49,7 @@ @SuppressWarnings({"rawtypes", "unchecked"}) public class NoDictionarySingleColumnGroupKeyGenerator implements GroupKeyGenerator { private final ExpressionContext _groupByExpression; - /// Group-key dispatch type: stored type of the column, except UUID is preserved as [DataType#UUID] so the - /// group-key map keys on [org.apache.pinot.spi.utils.UuidKey] (two primitive longs) instead of - /// [org.apache.pinot.spi.utils.ByteArray]. For every other logical type this equals - /// `logicalType.getStoredType()`. - private final DataType _dataType; + private final DataType _storedType; private final Map _groupKeyMap; private final int _globalGroupIdUpperBound; // TODO(nhejazi): Most of the logic between _nullHandlingEnabled=true/false is not sharable, so consider making a @@ -71,10 +66,8 @@ public NoDictionarySingleColumnGroupKeyGenerator(BaseProjectOperator projectO @Nullable Map groupByExpressionSizesFromPredicates) { _groupByExpression = groupByExpression; ColumnContext columnContext = projectOperator.getResultColumnContext(groupByExpression); - DataType logicalType = columnContext.getDataType(); - // Normalize to stored type, but preserve UUID so it uses UuidKey instead of ByteArray - _dataType = logicalType == DataType.UUID ? DataType.UUID : logicalType.getStoredType(); - _groupKeyMap = createGroupKeyMap(_dataType); + _storedType = columnContext.getDataType().getStoredType(); + _groupKeyMap = createGroupKeyMap(_storedType); if (groupByExpressionSizesFromPredicates != null) { Integer size = groupByExpressionSizesFromPredicates.get(groupByExpression); _globalGroupIdUpperBound = size != null ? Math.min(size, numGroupsLimit) : numGroupsLimit; @@ -102,7 +95,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } int numDocs = valueBlock.getNumDocs(); - switch (_dataType) { + switch (_storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -139,12 +132,6 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { groupKeys[i] = getKeyForValue(stringValues[i]); } break; - case UUID: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < numDocs; i++) { - groupKeys[i] = getKeyForValue(UuidKey.fromBytes(uuidValues[i])); - } - break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -152,7 +139,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[] groupKeys) { } break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); } } @@ -162,7 +149,7 @@ public void generateKeysForBlockNullHandlingEnabled(ValueBlock valueBlock, int[] BlockValSet blockValSet = valueBlock.getBlockValueSet(_groupByExpression); int numDocs = valueBlock.getNumDocs(); - switch (_dataType) { + switch (_storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); if (nullBitmap.getCardinality() < numDocs) { @@ -223,16 +210,6 @@ public void generateKeysForBlockNullHandlingEnabled(ValueBlock valueBlock, int[] Arrays.fill(groupKeys, 0, numDocs, getKeyForValue((String) null)); } break; - case UUID: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - if (nullBitmap.getCardinality() < numDocs) { - for (int i = 0; i < numDocs; i++) { - groupKeys[i] = getKeyForValue(nullBitmap.contains(i) ? null : UuidKey.fromBytes(uuidValues[i])); - } - } else if (numDocs > 0) { - Arrays.fill(groupKeys, 0, numDocs, getKeyForValue((UuidKey) null)); - } - break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (nullBitmap.getCardinality() < numDocs) { @@ -244,7 +221,7 @@ public void generateKeysForBlockNullHandlingEnabled(ValueBlock valueBlock, int[] } break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); } } @@ -279,10 +256,6 @@ private Map createGroupKeyMap(DataType keyType) { Object2IntOpenHashMap stringMap = new Object2IntOpenHashMap<>(); stringMap.defaultReturnValue(INVALID_ID); return stringMap; - case UUID: - Object2IntOpenHashMap uuidMap = new Object2IntOpenHashMap<>(); - uuidMap.defaultReturnValue(INVALID_ID); - return uuidMap; case BYTES: Object2IntOpenHashMap bytesMap = new Object2IntOpenHashMap<>(); bytesMap.defaultReturnValue(INVALID_ID); @@ -298,7 +271,7 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { BlockValSet blockValSet = valueBlock.getBlockValueSet(_groupByExpression); if (_isSingleValueExpression) { - switch (_dataType) { + switch (_storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -329,12 +302,6 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { groupKeys[i] = new int[]{getKeyForValue(stringValues[i])}; } break; - case UUID: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < numDocs; i++) { - groupKeys[i] = new int[]{getKeyForValue(UuidKey.fromBytes(uuidValues[i]))}; - } - break; case BYTES: byte[][] byteValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < numDocs; i++) { @@ -342,10 +309,10 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { } break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); } } else { - switch (_dataType) { + switch (_storedType) { case INT: int[][] intValues = blockValSet.getIntValuesMV(); for (int i = 0; i < numDocs; i++) { @@ -401,19 +368,8 @@ public void generateKeysForBlock(ValueBlock valueBlock, int[][] groupKeys) { groupKeys[i] = mvKeys; } break; - case UUID: - byte[][][] uuidValues = blockValSet.getBytesValuesMV(); - for (int i = 0; i < numDocs; i++) { - int mvSize = uuidValues[i].length; - int[] mvKeys = new int[mvSize]; - for (int j = 0; j < mvSize; j++) { - mvKeys[j] = getKeyForValue(UuidKey.fromBytes(uuidValues[i][j])); - } - groupKeys[i] = mvKeys; - } - break; default: - throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _dataType); + throw new IllegalArgumentException("Illegal data type for no-dictionary key generator: " + _storedType); } } } @@ -425,7 +381,7 @@ public int getCurrentGroupKeyUpperBound() { @Override public Iterator getGroupKeys() { - switch (_dataType) { + switch (_storedType) { case INT: return new IntGroupKeyIterator((Int2IntOpenHashMap) _groupKeyMap, _groupIdForNullValue); case LONG: @@ -437,8 +393,7 @@ public Iterator getGroupKeys() { case BIG_DECIMAL: case STRING: case BYTES: - case UUID: - return new ObjectGroupKeyIterator((Object2IntOpenHashMap) _groupKeyMap, _dataType); + return new ObjectGroupKeyIterator((Object2IntOpenHashMap) _groupKeyMap); default: throw new IllegalStateException(); } @@ -530,16 +485,6 @@ private int getKeyForValue(ByteArray value) { return groupId; } - private int getKeyForValue(UuidKey value) { - Object2IntMap map = (Object2IntMap) _groupKeyMap; - int groupId = map.getInt(value); - if (groupId == INVALID_ID && _numGroups < _globalGroupIdUpperBound) { - groupId = _numGroups++; - map.put(value, groupId); - } - return groupId; - } - private static class IntGroupKeyIterator implements Iterator { final Iterator _iterator; final GroupKey _groupKey; @@ -688,12 +633,10 @@ public void remove() { private static class ObjectGroupKeyIterator implements Iterator { final ObjectIterator _iterator; final GroupKey _groupKey; - final DataType _dataType; - ObjectGroupKeyIterator(Object2IntOpenHashMap objectMap, DataType dataType) { + ObjectGroupKeyIterator(Object2IntOpenHashMap objectMap) { _iterator = objectMap.object2IntEntrySet().fastIterator(); _groupKey = new GroupKey(); - _dataType = dataType; } @Override @@ -705,11 +648,7 @@ public boolean hasNext() { public GroupKey next() { Object2IntMap.Entry entry = _iterator.next(); _groupKey._groupId = entry.getIntValue(); - Object key = entry.getKey(); - if (_dataType == DataType.UUID && key != null) { - key = ((UuidKey) key).toByteArray(); - } - _groupKey._keys = new Object[]{key}; + _groupKey._keys = new Object[]{entry.getKey()}; return _groupKey; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java deleted file mode 100644 index 0ee338e97013..000000000000 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/UuidToIdMap.java +++ /dev/null @@ -1,64 +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.groupby.utils; - -import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; -import java.util.ArrayList; -import org.apache.pinot.spi.utils.ByteArray; -import org.apache.pinot.spi.utils.UuidKey; - - -/// Implementation of [ValueToIdMap] for Pinot's logical UUID type. -public class UuidToIdMap implements ValueToIdMap { - private final Object2IntOpenHashMap _valueToIdMap; - private final ArrayList _idToValueMap; - - public UuidToIdMap() { - _valueToIdMap = new Object2IntOpenHashMap<>(); - _valueToIdMap.defaultReturnValue(INVALID_KEY); - _idToValueMap = new ArrayList<>(); - } - - /// Both callers -- [org.apache.pinot.core.query.aggregation.groupby.NoDictionaryMultiColumnGroupKeyGenerator] and - /// [org.apache.pinot.core.query.aggregation.groupby.NoDictionarySingleColumnGroupKeyGenerator] -- key on - /// [UuidKey] already, so this casts directly rather than going through `UuidKey#fromObject`. That matches the - /// sibling maps (e.g. [DoubleToIdMap] casts to `double`) and keeps the per-row `instanceof` chain out of the - /// group-by loop. - @Override - public int put(Object value) { - UuidKey uuidKey = (UuidKey) value; - int id = _valueToIdMap.getInt(uuidKey); - if (id == INVALID_KEY) { - id = _valueToIdMap.size(); - _valueToIdMap.put(uuidKey, id); - _idToValueMap.add(uuidKey.toByteArray()); - } - return id; - } - - @Override - public int getId(Object value) { - return _valueToIdMap.getInt((UuidKey) value); - } - - @Override - public Object get(int id) { - return _idToValueMap.get(id); - } -} diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java index a15f9bf9a263..4ce47caa196b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/groupby/utils/ValueToIdMapFactory.java @@ -36,8 +36,6 @@ public static ValueToIdMap get(DataType dataType) { return new FloatToIdMap(); case DOUBLE: return new DoubleToIdMap(); - case UUID: - return new UuidToIdMap(); default: return new ObjectToIdMap(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java index 2192848feb3d..2ed7bff83f1b 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/groupby/NoDictionaryGroupKeyGeneratorTest.java @@ -27,7 +27,6 @@ import java.util.List; import java.util.Random; import java.util.Set; -import java.util.UUID; import org.apache.commons.io.FileUtils; import org.apache.commons.lang3.RandomStringUtils; import org.apache.commons.lang3.StringUtils; @@ -52,7 +51,6 @@ import org.apache.pinot.spi.utils.BytesUtils; import org.apache.pinot.spi.utils.CommonConstants.Server; import org.apache.pinot.spi.utils.ReadMode; -import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.testng.annotations.AfterClass; import org.testng.annotations.BeforeClass; @@ -77,19 +75,12 @@ public class NoDictionaryGroupKeyGeneratorTest { private static final String STRING_COLUMN = "stringColumn"; private static final String BYTES_COLUMN = "bytesColumn"; private static final String BYTES_DICT_COLUMN = "bytesDictColumn"; - private static final String UUID_COLUMN = "uuidColumn"; - private static final String BOOLEAN_COLUMN = "booleanColumn"; - private static final String TIMESTAMP_COLUMN = "timestampColumn"; - private static final String UUID_DICT_COLUMN = "uuidDictColumn"; private static final List COLUMNS = Arrays.asList(INT_COLUMN, LONG_COLUMN, FLOAT_COLUMN, DOUBLE_COLUMN, STRING_COLUMN, BYTES_COLUMN, - BYTES_DICT_COLUMN, UUID_COLUMN, BOOLEAN_COLUMN, TIMESTAMP_COLUMN, UUID_DICT_COLUMN); + BYTES_DICT_COLUMN); private static final int NUM_COLUMNS = COLUMNS.size(); - private static final Set UUID_COLUMNS = Set.of(UUID_COLUMN, UUID_DICT_COLUMN); private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME) - .setNoDictionaryColumns( - List.of(INT_COLUMN, LONG_COLUMN, FLOAT_COLUMN, DOUBLE_COLUMN, STRING_COLUMN, BYTES_COLUMN, UUID_COLUMN, - BOOLEAN_COLUMN, TIMESTAMP_COLUMN)).build(); + .setNoDictionaryColumns(COLUMNS.subList(0, NUM_COLUMNS - 1)).build(); private static final Schema SCHEMA = new Schema.SchemaBuilder().addSingleValueDimension(INT_COLUMN, FieldSpec.DataType.INT) .addSingleValueDimension(LONG_COLUMN, FieldSpec.DataType.LONG) @@ -97,11 +88,7 @@ public class NoDictionaryGroupKeyGeneratorTest { .addSingleValueDimension(DOUBLE_COLUMN, FieldSpec.DataType.DOUBLE) .addSingleValueDimension(STRING_COLUMN, FieldSpec.DataType.STRING) .addSingleValueDimension(BYTES_COLUMN, FieldSpec.DataType.BYTES) - .addSingleValueDimension(BYTES_DICT_COLUMN, FieldSpec.DataType.BYTES) - .addSingleValueDimension(UUID_COLUMN, FieldSpec.DataType.UUID) - .addSingleValueDimension(BOOLEAN_COLUMN, FieldSpec.DataType.BOOLEAN) - .addSingleValueDimension(TIMESTAMP_COLUMN, FieldSpec.DataType.TIMESTAMP) - .addSingleValueDimension(UUID_DICT_COLUMN, FieldSpec.DataType.UUID).build(); + .addSingleValueDimension(BYTES_DICT_COLUMN, FieldSpec.DataType.BYTES).build(); private static final int NUM_RECORDS = 1000; private static final int NUM_UNIQUE_RECORDS = 100; @@ -142,19 +129,6 @@ public void setUp() record.putValue(BYTES_DICT_COLUMN, bytesValue); values[5] = BytesUtils.toHexString(bytesValue); values[6] = values[5]; - byte[] uuidBytes = UuidUtils.toBytes(new UUID(RANDOM.nextLong(), RANDOM.nextLong())); - record.putValue(UUID_COLUMN, uuidBytes); - values[7] = UuidUtils.toString(uuidBytes); - // BOOLEAN stored as INT (0/1) — exercises the logical→stored-type normalization fix - int boolIntValue = RANDOM.nextBoolean() ? 1 : 0; - record.putValue(BOOLEAN_COLUMN, boolIntValue); - values[8] = Integer.toString(boolIntValue); - // TIMESTAMP stored as LONG — exercises the logical→stored-type normalization fix - long timestampValue = Math.abs(RANDOM.nextLong()); - record.putValue(TIMESTAMP_COLUMN, timestampValue); - values[9] = Long.toString(timestampValue); - record.putValue(UUID_DICT_COLUMN, uuidBytes); - values[10] = values[7]; for (int j = 0; j < NUM_RECORDS / NUM_UNIQUE_RECORDS; j++) { records.add(record); } @@ -199,12 +173,9 @@ public void testMultiColumnGroupKeyGenerator() { testGroupKeyGenerator(new int[]{0, 1}); testGroupKeyGenerator(new int[]{2, 3}); testGroupKeyGenerator(new int[]{4, 5}); - testGroupKeyGenerator(new int[]{7, 10}); - testGroupKeyGenerator(new int[]{8, 9}); testGroupKeyGenerator(new int[]{1, 2, 3}); testGroupKeyGenerator(new int[]{4, 5, 0}); - testGroupKeyGenerator(new int[]{7, 5, 4}); - testGroupKeyGenerator(new int[]{7, 5, 4, 3, 2, 1, 0}); + testGroupKeyGenerator(new int[]{5, 4, 3, 2, 1, 0}); } /// Tests multi-column group key generator when at least one column as dictionary, and others don't. @@ -241,7 +212,7 @@ private void testGroupKeyGenerator(int[] groupByColumnIndexes) { Iterator groupKeys = groupKeyGenerator.getGroupKeys(); while (groupKeys.hasNext()) { GroupKeyGenerator.GroupKey groupKey = groupKeys.next(); - assertTrue(expectedGroupKeys.contains(getActualGroupKey(groupKey._keys, groupByColumnIndexes))); + assertTrue(expectedGroupKeys.contains(getActualGroupKey(groupKey._keys))); } } @@ -263,18 +234,13 @@ private Set getExpectedGroupKeys(int[] groupByColumnIndexes) { return groupKeys; } - private String getActualGroupKey(Object[] groupKeys, int[] groupByColumnIndexes) { + private String getActualGroupKey(Object[] groupKeys) { StringBuilder stringBuilder = new StringBuilder(); for (int i = 0; i < groupKeys.length; i++) { if (i > 0) { stringBuilder.append(GroupKeyGenerator.DELIMITER); } - int columnIndex = groupByColumnIndexes[i]; - if (UUID_COLUMNS.contains(COLUMNS.get(columnIndex))) { - stringBuilder.append(UuidUtils.toString(((org.apache.pinot.spi.utils.ByteArray) groupKeys[i]).getBytes())); - } else { - stringBuilder.append(groupKeys[i]); - } + stringBuilder.append(groupKeys[i]); } return stringBuilder.toString(); } From 03602dd1cd5d6d0ff4d150cf4993267e4fdba0ff Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 9 Aug 2026 13:09:33 -0700 Subject: [PATCH 06/19] Format DISTINCT rows in one pass in BytesDistinctTable The previous version added every row holding a ByteArray, then walked the whole list again in formatRows() to rewrite row[0]. Two traversals and two writes per row where master did one. The formatting now happens inline in addRows, with the ColumnDataType passed in -- the same shape LongDistinctTable already uses for its TIMESTAMP handling. Not an extra String allocation either way: ColumnDataType#convertAndFormat for BYTES is ((ByteArray) value).toHexString(), the exact call master made, and ByteArray#getBytes() returns the internal array without copying. --- .../distinct/table/BytesDistinctTable.java | 37 +++++++------------ 1 file changed, 14 insertions(+), 23 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java index 42649c9bd28a..f41794fb7420 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java @@ -32,6 +32,7 @@ import org.apache.pinot.common.request.context.OrderByExpressionContext; import org.apache.pinot.common.response.broker.ResultTable; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.core.common.datatable.DataTableBuilder; import org.apache.pinot.core.common.datatable.DataTableBuilderFactory; import org.apache.pinot.spi.query.QueryThreadContext; @@ -263,69 +264,59 @@ private ResultTable toResultTableWithOrderBy() { } int numValues = sortedValues.length; assert numValues <= _limit; + ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); List rows; if (_hasNull) { if (numValues == _limit) { rows = new ArrayList<>(_limit); if (_orderByExpression.isNullsLast()) { - addRows(sortedValues, numValues, rows); + addRows(columnDataType, sortedValues, numValues, rows); } else { rows.add(new Object[]{null}); - addRows(sortedValues, numValues - 1, rows); + addRows(columnDataType, sortedValues, numValues - 1, rows); } } else { rows = new ArrayList<>(numValues + 1); if (_orderByExpression.isNullsLast()) { - addRows(sortedValues, numValues, rows); + addRows(columnDataType, sortedValues, numValues, rows); rows.add(new Object[]{null}); } else { rows.add(new Object[]{null}); - addRows(sortedValues, numValues, rows); + addRows(columnDataType, sortedValues, numValues, rows); } } } else { rows = new ArrayList<>(numValues); - addRows(sortedValues, numValues, rows); + addRows(columnDataType, sortedValues, numValues, rows); } - formatRows(rows); return new ResultTable(_dataSchema, rows); } - private static void addRows(ByteArray[] values, int length, List rows) { + private static void addRows(ColumnDataType columnDataType, ByteArray[] values, int length, List rows) { for (int i = 0; i < length; i++) { - rows.add(new Object[]{values[i]}); + rows.add(new Object[]{columnDataType.convertAndFormat(values[i])}); } } private ResultTable toResultTableWithoutOrderBy() { int numValues = _valueSet.size(); assert numValues <= _limit; + ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); List rows; if (_hasNull && numValues < _limit) { rows = new ArrayList<>(numValues + 1); - addRows(_valueSet, rows); + addRows(columnDataType, _valueSet, rows); rows.add(new Object[]{null}); } else { rows = new ArrayList<>(numValues); - addRows(_valueSet, rows); + addRows(columnDataType, _valueSet, rows); } - formatRows(rows); return new ResultTable(_dataSchema, rows); } - private static void addRows(HashSet values, List rows) { + private static void addRows(ColumnDataType columnDataType, HashSet values, List rows) { for (ByteArray value : values) { - rows.add(new Object[]{value}); - } - } - - private void formatRows(List rows) { - DataSchema.ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); - for (Object[] row : rows) { - Object value = row[0]; - if (value != null) { - row[0] = columnDataType.convertAndFormat(value); - } + rows.add(new Object[]{columnDataType.convertAndFormat(value)}); } } } From 628fb5391b5d8d900d4ab050c22aa36ae36e20f7 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 12 Aug 2026 17:57:26 -0700 Subject: [PATCH 07/19] Route UUID through the normal dispatch in DistinctCountBitmap Replaces the three `if (dataType == DataType.UUID)` early-returns with a narrowed guard on the serialized-bitmap branch: if (storedType == DataType.BYTES && dataType != DataType.UUID) UUID now falls through to aggregateSV/aggregateMV and the group-by variants like any other type, handled by a `case BYTES` in each switch. That fixes MV UUID, which the early-returns silently broke -- they called getBytesValuesSV() above the isSingleValue() dispatch, so an MV column took the SV accessor. Seven switches needed the case, not six: the dictionary path in convertToValueBitmap also throws on BYTES, which the integration test caught -- DISTINCTCOUNTBITMAP(uuidColumn) failed with "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: BYTES" until it was added. DistinctCountHLL, DistinctCountHLLPlus, DistinctCountCPCSketch and DistinctCountULL still use the early-return shape and need the same treatment. --- ...istinctCountBitmapAggregationFunction.java | 99 ++++++++++++------- 1 file changed, 62 insertions(+), 37 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index 03c10883f36a..20cbc7999e68 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -75,20 +75,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized RoaringBitmap state. Add the - // hashCode of the canonical UUID string so DISTINCTCOUNTBITMAP(uuidCol) matches - // DISTINCTCOUNTBITMAP(CAST(uuidCol AS STRING)). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - RoaringBitmap bitmap = getValueBitmap(aggregationResultHolder); - for (int i = 0; i < length; i++) { - bitmap.add(Arrays.hashCode(uuidBytesValues[i])); - } - return; - } - // Treat BYTES value as serialized RoaringBitmap - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); RoaringBitmap valueBitmap = aggregationResultHolder.getResult(); if (valueBitmap != null) { @@ -153,6 +141,13 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult valueBitmap.add(stringValues[i].hashCode()); } break; + // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + valueBitmap.add(Arrays.hashCode(uuidValues[i])); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -213,6 +208,15 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; + // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. + case BYTES: + byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (byte[] value : uuidValues[i]) { + valueBitmap.add(Arrays.hashCode(value)); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -227,18 +231,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: add hashCode of canonical UUID string (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getValueBitmap(groupByResultHolder, groupKeyArray[i]) - .add(Arrays.hashCode(uuidBytesValues[i])); - } - return; - } - // Treat BYTES value as serialized RoaringBitmap - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); @@ -304,6 +298,13 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(stringValues[i].hashCode()); } break; + // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(Arrays.hashCode(uuidValues[i])); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -366,6 +367,16 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; + // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. + case BYTES: + byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + RoaringBitmap bitmap = getValueBitmap(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : uuidValues[i]) { + bitmap.add(Arrays.hashCode(value)); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -380,20 +391,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: add hashCode of canonical UUID string (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - int hash = Arrays.hashCode(uuidBytesValues[i]); - for (int groupKey : groupKeysArray[i]) { - getValueBitmap(groupByResultHolder, groupKey).add(hash); - } - } - return; - } - // Treat BYTES value as serialized RoaringBitmap - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); @@ -461,6 +460,13 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i].hashCode()); } break; + // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], Arrays.hashCode(uuidValues[i])); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -535,6 +541,18 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; + // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. + case BYTES: + byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + RoaringBitmap bitmap = getValueBitmap(groupByResultHolder, groupKey); + for (byte[] value : uuidValues[i]) { + bitmap.add(Arrays.hashCode(value)); + } + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -701,6 +719,13 @@ private static RoaringBitmap convertToValueBitmap(DictIdsWrapper dictIdsWrapper) valueBitmap.add(dictionary.getStringValue(iterator.next()).hashCode()); } break; + // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary); hash the stored bytes to match + // the scan path above. A real BYTES column holds serialized bitmaps and never reaches the dictionary path. + case BYTES: + while (iterator.hasNext()) { + valueBitmap.add(Arrays.hashCode(dictionary.getBytesValue(iterator.next()))); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); From 96e464a77be6a6c3708cb325effaa581bd3b66a4 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 12 Aug 2026 18:14:56 -0700 Subject: [PATCH 08/19] Route UUID through the normal dispatch in DistinctCountHLL Same shape as DistinctCountBitmap: the three `if (dataType == DataType.UUID)` early-returns are gone, the serialized-HLL guard is narrowed to if (storedType == DataType.BYTES && dataType != DataType.UUID) and UUID falls through to aggregateSV/aggregateMV and the group-by variants, handled by a `case BYTES` in each of the six switches. This fixes MV UUID, which the early-returns broke by calling getBytesValuesSV() above the isSingleValue() dispatch. No seventh switch here: unlike DistinctCountBitmap, HLL's dictionary path lives in AggregationFunctionUtils#getDistinctCountHLLResult, which already excludes UUID from the serialized branch so it falls through to the scalar path. --- .../DistinctCountHLLAggregationFunction.java | 94 +++++++++++-------- 1 file changed, 56 insertions(+), 38 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java index 8a837029e977..666903857e9b 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -84,22 +84,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns are stored as 16-byte BYTES, but a UUID value is a logical scalar — not a serialized - // HyperLogLog. Offer the canonical UUID string so the result matches DISTINCTCOUNTHLL on a STRING column - // holding the same logical UUIDs. NOTE: fetch raw bytes and convert explicitly — for identifier expressions - // the BlockValSet is a ProjectionBlockValSet whose getStringValuesSV() renders stored BYTES as bare hex, - // not the canonical RFC-4122 form. - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - HyperLogLog hyperLogLog = getHyperLogLog(aggregationResultHolder); - for (int i = 0; i < length; i++) { - hyperLogLog.offer(uuidBytesValues[i]); - } - return; - } - // Treat BYTES value as serialized HyperLogLog - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); @@ -175,6 +161,13 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult hyperLogLog.offer(stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + hyperLogLog.offer(uuidValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -238,6 +231,15 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; + // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. + case BYTES: + byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (byte[] value : uuidValuesArray[i]) { + hyperLogLog.offer(value); + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -251,17 +253,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(uuidBytesValues[i]); - } - return; - } - // Treat BYTES value as serialized HyperLogLog - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -334,6 +327,13 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(uuidValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -398,6 +398,16 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; + // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. + case BYTES: + byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + HyperLogLog hyperLogLog = getHyperLogLog(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : uuidValuesArray[i]) { + hyperLogLog.offer(value); + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -411,20 +421,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - byte[] canonical = uuidBytesValues[i]; - for (int groupKey : groupKeysArray[i]) { - getHyperLogLog(groupByResultHolder, groupKey).offer(canonical); - } - } - return; - } - // Treat BYTES value as serialized HyperLogLog - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -500,6 +498,13 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], uuidValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -582,6 +587,19 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; + // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. + case BYTES: + byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + byte[][] uuidValues = uuidValuesArray[i]; + for (int groupKey : groupKeysArray[i]) { + HyperLogLog hyperLogLog = getHyperLogLog(groupByResultHolder, groupKey); + for (byte[] value : uuidValues) { + hyperLogLog.offer(value); + } + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } From 89007bbc1090e2c834289ea03e3df5927d93da65 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Thu, 13 Aug 2026 15:28:18 -0700 Subject: [PATCH 09/19] Route UUID through the normal dispatch in DistinctCountHLLPlus and ULL Same shape as Bitmap and HLL: the three `if (dataType == DataType.UUID)` early-returns are removed, each serialized-sketch guard is narrowed to if (storedType == DataType.BYTES && dataType != DataType.UUID) and UUID reaches the ordinary dispatch via a `case BYTES` in each switch. HLLPlus has 6 switches (SV/MV x aggregate, groupBySV, groupByMV). ULL has 3: it is single-value only, so all of its switches read getStringValuesSV and the BYTES cases mirror that -- no MV variants to add. Both dictionary paths live in AggregationFunctionUtils, which already excludes UUID from the serialized branch, so neither needed a seventh switch. --- ...stinctCountHLLPlusAggregationFunction.java | 91 ++++++++++++------- .../DistinctCountULLAggregationFunction.java | 62 +++++-------- 2 files changed, 81 insertions(+), 72 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java index 452f33412099..dce1e55ceaf5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -96,19 +96,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized HyperLogLogPlus state. Offer the - // canonical UUID string so DISTINCTCOUNTHLLPLUS(uuidCol) matches DISTINCTCOUNTHLLPLUS(CAST(uuidCol AS STRING)). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(aggregationResultHolder); - for (int i = 0; i < length; i++) { - hyperLogLogPlus.offer(uuidBytesValues[i]); - } - return; - } - // Treat BYTES value as serialized HyperLogLogPlus - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); @@ -177,6 +166,13 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult hyperLogLogPlus.offer(stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + hyperLogLogPlus.offer(uuidValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -239,6 +235,15 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; + // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. + case BYTES: + byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (byte[] value : uuidValuesArray[i]) { + hyperLogLogPlus.offer(value); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -253,17 +258,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(uuidBytesValues[i]); - } - return; - } - // Treat BYTES value as serialized HyperLogLogPlus - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -333,6 +329,13 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(uuidValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -398,6 +401,16 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; + // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. + case BYTES: + byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : uuidValuesArray[i]) { + hyperLogLogPlus.offer(value); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -412,20 +425,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: offer canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - byte[] canonical = uuidBytesValues[i]; - for (int groupKey : groupKeysArray[i]) { - getHyperLogLogPlus(groupByResultHolder, groupKey).offer(canonical); - } - } - return; - } - // Treat BYTES value as serialized HyperLogLogPlus - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -498,6 +499,13 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], uuidValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -580,6 +588,19 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; + // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. + case BYTES: + byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + byte[][] uuidValues = uuidValuesArray[i]; + for (int groupKey : groupKeysArray[i]) { + HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(groupByResultHolder, groupKey); + for (byte[] value : uuidValues) { + hyperLogLogPlus.offer(value); + } + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java index 3a9319e00c7c..9dcf2f9e6f04 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -85,19 +85,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized UltraLogLog state. Hash the - // canonical UUID string so DISTINCTCOUNTULL(uuidCol) matches DISTINCTCOUNTULL(CAST(uuidCol AS STRING)). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - UltraLogLog ull = getULL(aggregationResultHolder); - for (int i = 0; i < length; i++) { - UltraLogLogUtils.hashObject(uuidBytesValues[i]).ifPresent(ull::add); - } - return; - } - // Treat BYTES value as serialized UltraLogLog - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { UltraLogLog ull = aggregationResultHolder.getResult(); @@ -157,6 +146,13 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde UltraLogLogUtils.hashObject(stringValues[i]).ifPresent(ull::add); } break; + // Reached only by UUID: a real BYTES column is serialized ULL state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(uuidValues[i]).ifPresent(ull::add); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_ULL aggregation function: " + storedType); @@ -171,18 +167,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: hash canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - UltraLogLog ull = getULL(groupByResultHolder, groupKeyArray[i]); - UltraLogLogUtils.hashObject(uuidBytesValues[i]).ifPresent(ull::add); - } - return; - } - // Treat BYTES value as serialized UltraLogLogs - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -248,6 +234,14 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); } break; + // Reached only by UUID: a real BYTES column is serialized ULL state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(uuidValues[i]) + .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_ULL aggregation function: " + storedType); @@ -262,21 +256,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: hash canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - byte[] canonical = uuidBytesValues[i]; - for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = getULL(groupByResultHolder, groupKey); - UltraLogLogUtils.hashObject(canonical).ifPresent(ull::add); - } - } - return; - } - // Treat BYTES value as serialized UltraLogLogs - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -340,6 +321,13 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is serialized ULL state and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], uuidValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); From fbe5cefda6278ab44e31273894ebee34b2f79247 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Thu, 13 Aug 2026 21:38:03 -0700 Subject: [PATCH 10/19] Read UUID as bytes in DistinctCountThetaSketch extractValues no longer special-cases UUID: the block that materialized a String[] / String[][] of hex renderings is gone, so a UUID column now reports its stored type (BYTES) and carries the raw byte[][] like any bytes column. All the handling moved into the aggregateXXX methods. Each of the three dispatch sites now reads the logical type from the block val set it already has and narrows its guard, so UUID takes the scalar path rather than the serialized-sketch path, and a "case BYTES" in each of the six scalar switches feeds the stored bytes straight to UpdatableThetaSketch#update(byte[]) -- no String allocated anywhere. Same shape as Bitmap, HLL, HLLPlus and ULL. This also fixes MV UUID here for the same reason: the work now happens inside the SV and MV branches instead of above them. --- ...ctCountThetaSketchAggregationFunction.java | 177 ++++++++++++++---- 1 file changed, 140 insertions(+), 37 deletions(-) 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 263cf8af0091..a104bfd4ea8a 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 @@ -55,7 +55,6 @@ import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.Constants; import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.apache.pinot.spi.utils.BytesUtils; import org.apache.pinot.spi.utils.CommonConstants; import org.apache.pinot.sql.parsers.CalciteSqlParser; @@ -193,8 +192,12 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0 - if (valueTypes[0] != DataType.BYTES) { + // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but + // its values are logical scalars rather than serialized sketches, so it takes the scalar path below and is + // handled by `case BYTES`. + boolean uuidMainExpression = + blockValSetMap.get(_inputExpressions.get(0)).getValueType() == DataType.UUID; + if (valueTypes[0] != DataType.BYTES || uuidMainExpression) { List updateSketches = getUpdateSketches(aggregationResultHolder); if (singleValues[0]) { switch (valueTypes[0]) { @@ -288,6 +291,24 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } break; + case BYTES: + byte[][] uuidValues = (byte[][]) valueArrays[0]; + if (_includeDefaultSketch) { + UpdatableThetaSketch defaultSketch = updateSketches.get(0); + for (int i = 0; i < length; i++) { + defaultSketch.update(uuidValues[i]); + } + } + for (int i = 0; i < numFilters; i++) { + FilterEvaluator filterEvaluator = _filterEvaluators.get(i); + UpdatableThetaSketch updateSketch = updateSketches.get(i + 1); + for (int j = 0; j < length; j++) { + if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { + updateSketch.update(uuidValues[j]); + } + } + } + break; default: throw new IllegalStateException( "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " @@ -405,6 +426,28 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } break; + case BYTES: + byte[][][] uuidValues = (byte[][][]) valueArrays[0]; + if (_includeDefaultSketch) { + UpdatableThetaSketch defaultSketch = updateSketches.get(0); + for (int i = 0; i < length; i++) { + for (byte[] value : uuidValues[i]) { + defaultSketch.update(value); + } + } + } + for (int i = 0; i < numFilters; i++) { + FilterEvaluator filterEvaluator = _filterEvaluators.get(i); + UpdatableThetaSketch updateSketch = updateSketches.get(i + 1); + for (int j = 0; j < length; j++) { + if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { + for (byte[] value : uuidValues[j]) { + updateSketch.update(value); + } + } + } + } + break; default: throw new IllegalStateException( "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); @@ -442,8 +485,12 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0 - if (valueTypes[0] != DataType.BYTES) { + // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but + // its values are logical scalars rather than serialized sketches, so it takes the scalar path below and is + // handled by `case BYTES`. + boolean uuidMainExpression = + blockValSetMap.get(_inputExpressions.get(0)).getValueType() == DataType.UUID; + if (valueTypes[0] != DataType.BYTES || uuidMainExpression) { if (singleValues[0]) { switch (valueTypes[0]) { case INT: @@ -521,6 +568,21 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } break; + case BYTES: + byte[][] uuidValues = (byte[][]) valueArrays[0]; + for (int i = 0; i < length; i++) { + List updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]); + byte[] value = uuidValues[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; default: throw new IllegalStateException( "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " @@ -633,6 +695,27 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } break; + case BYTES: + byte[][][] uuidValues = (byte[][][]) valueArrays[0]; + for (int i = 0; i < length; i++) { + List updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]); + byte[][] values = uuidValues[i]; + if (_includeDefaultSketch) { + UpdatableThetaSketch defaultSketch = updateSketches.get(0); + for (byte[] value : values) { + defaultSketch.update(value); + } + } + for (int j = 0; j < numFilters; j++) { + if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) { + UpdatableThetaSketch updateSketch = updateSketches.get(j + 1); + for (byte[] value : values) { + updateSketch.update(value); + } + } + } + } + break; default: throw new IllegalStateException( "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); @@ -666,8 +749,12 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0 - if (valueTypes[0] != DataType.BYTES) { + // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but + // its values are logical scalars rather than serialized sketches, so it takes the scalar path below and is + // handled by `case BYTES`. + boolean uuidMainExpression = + blockValSetMap.get(_inputExpressions.get(0)).getValueType() == DataType.UUID; + if (valueTypes[0] != DataType.BYTES || uuidMainExpression) { if (singleValues[0]) { switch (valueTypes[0]) { case INT: @@ -770,6 +857,26 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } break; + case BYTES: + byte[][] uuidValues = (byte[][]) valueArrays[0]; + if (_includeDefaultSketch) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + getUpdateSketches(groupByResultHolder, groupKey).get(0).update(uuidValues[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(uuidValues[j]); + } + } + } + } + break; default: throw new IllegalStateException( "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " @@ -907,6 +1014,32 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } break; + case BYTES: + byte[][][] uuidValues = (byte[][][]) valueArrays[0]; + if (_includeDefaultSketch) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + UpdatableThetaSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0); + for (byte[] value : uuidValues[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]) { + UpdatableThetaSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1); + for (byte[] value : uuidValues[i]) { + updateSketch.update(value); + } + } + } + } + } + break; default: throw new IllegalStateException( "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); @@ -1236,36 +1369,6 @@ private void extractValues(int length, Map block DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); singleValues[i] = singleValue; - // UUID columns are stored as 16-byte BYTES but a UUID value is a logical scalar, not a pre-serialized - // theta sketch: without this branch the function takes the serialized-sketch path below and Sketch.wrap - // fails on raw 16-byte UUID content. Unlike the other distinct-count functions this one cannot consume the - // stored bytes directly -- DataType.BYTES here means "serialized sketch", with no scalar-bytes mode -- so - // the stored value is surfaced as its hex rendering, the same form used at every other String-typed UUID - // boundary (see PredicateUtils#getStoredValue and the Bloom filter key). - if (dataType == DataType.UUID) { - valueTypes[i] = DataType.STRING; - if (singleValue) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - String[] hexValues = new String[length]; - for (int j = 0; j < length; j++) { - hexValues[j] = BytesUtils.toHexString(uuidBytesValues[j]); - } - valueArrays[i] = hexValues; - } else { - byte[][][] uuidBytesValuesMV = blockValSet.getBytesValuesMV(); - String[][] hexValuesMV = new String[length][]; - for (int j = 0; j < length; j++) { - byte[][] row = uuidBytesValuesMV[j]; - String[] hexRow = new String[row.length]; - for (int k = 0; k < row.length; k++) { - hexRow[k] = BytesUtils.toHexString(row[k]); - } - hexValuesMV[j] = hexRow; - } - valueArrays[i] = hexValuesMV; - } - continue; - } valueTypes[i] = storedType; if (singleValue) { switch (storedType) { From 3a694ff0d3762644b9210e0f171a18c0b1f8a6b7 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Thu, 13 Aug 2026 21:49:11 -0700 Subject: [PATCH 11/19] Leave extractValues untouched in DistinctCountThetaSketch The length parameter was only needed to size the hex String arrays that the UUID block used to build. With that block gone, length was unused in the body, so the signature change and its three call sites were dead weight. extractValues is now byte-identical to master: all UUID handling lives in the aggregateXXX methods. --- .../DistinctCountThetaSketchAggregationFunction.java | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) 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 a104bfd4ea8a..89b70290567b 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 @@ -189,7 +189,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde boolean[] singleValues = new boolean[numExpressions]; DataType[] valueTypes = new DataType[numExpressions]; Object[] valueArrays = new Object[numExpressions]; - extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); + extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but @@ -482,7 +482,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol boolean[] singleValues = new boolean[numExpressions]; DataType[] valueTypes = new DataType[numExpressions]; Object[] valueArrays = new Object[numExpressions]; - extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); + extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but @@ -746,7 +746,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult boolean[] singleValues = new boolean[numExpressions]; DataType[] valueTypes = new DataType[numExpressions]; Object[] valueArrays = new Object[numExpressions]; - extractValues(length, blockValSetMap, singleValues, valueTypes, valueArrays); + extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but @@ -1360,7 +1360,7 @@ private static int extractSketchId(String identifier) { } /// Extracts values from the BlockValSet map. - private void extractValues(int length, Map blockValSetMap, boolean[] singleValues, + private void extractValues(Map blockValSetMap, boolean[] singleValues, DataType[] valueTypes, Object[] valueArrays) { int numExpressions = _inputExpressions.size(); for (int i = 0; i < numExpressions; i++) { From b2e7f9e045f18f4b17807d8c26540dc8dd7d9b97 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Thu, 13 Aug 2026 22:09:38 -0700 Subject: [PATCH 12/19] Route UUID through the normal dispatch in DistinctCountCPCSketch Last of the six. Same shape as the others: the UUID early-returns are gone, the serialized-sketch guard is narrowed with `&& dataType != DataType.UUID`, and a `case BYTES` in each switch feeds the stored bytes to CpcSketch#update. This one had 2 UUID blocks and 1 BYTES guard rather than 3 and 3 -- the counts differ per function, so they were verified rather than assumed. No `if (dataType == DataType.UUID)` early-return remains anywhere in the aggregation function package. --- ...inctCountCPCSketchAggregationFunction.java | 48 +++++++++---------- 1 file changed, 24 insertions(+), 24 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index 96e286aae1e9..2be253256157 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -138,22 +138,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde FieldSpec.DataType dataType = blockValSet.getValueType(); FieldSpec.DataType storedType = dataType.getStoredType(); - // UUID values are logical scalars (stored as 16-byte BYTES) — not serialized CPC Sketch state. Update - // the sketch with the canonical UUID string so DISTINCTCOUNTCPC(uuidCol) matches - // DISTINCTCOUNTCPC(CAST(uuidCol AS STRING)). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - // Leave the updated CpcSketch in the holder; extractAggregationResult converts it to an accumulator. - // Calling getAccumulator here would read the holder slot already occupied by the sketch and fail. - CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); - for (int i = 0; i < length; i++) { - cpcSketch.update(uuidBytesValues[i]); - } - return; - } - // Treat BYTES value as serialized CPC Sketch - if (storedType == DataType.BYTES) { + if (storedType == DataType.BYTES && dataType != DataType.UUID) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(aggregationResultHolder); @@ -210,6 +196,13 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde cpcSketch.update(stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is a serialized CPC sketch and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + cpcSketch.update(uuidValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } @@ -226,15 +219,6 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // UUID columns: update with canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(uuidBytesValues[i]); - } - return; - } - // Treat BYTES value as serialized CPC Sketch if (storedType == FieldSpec.DataType.BYTES) { byte[][] bytesValues = blockValSet.getBytesValuesSV(); @@ -295,6 +279,13 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(stringValues[i]); } break; + // Reached only by UUID: a real BYTES column is a serialized CPC sketch and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(uuidValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } @@ -390,6 +381,15 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } break; + // Reached only by UUID: a real BYTES column is a serialized CPC sketch and is handled above. + case BYTES: + byte[][] uuidValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + getCpcSketch(groupByResultHolder, groupKey).update(uuidValues[i]); + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } From 04650c80de54dd4e0285a7632b249a8ce714f512 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 14 Aug 2026 10:57:12 -0700 Subject: [PATCH 13/19] Fix UUID aggregation dispatch for CPC and Theta sketches Route UUID values through the single-value and multi-value switch branches for aggregate and group-by paths while preserving serialized BYTES sketch handling. Add unit and ingestion/query coverage for UUID and multi-value serialized sketches. --- ...inctCountCPCSketchAggregationFunction.java | 440 ++++++++++++++---- ...ctCountThetaSketchAggregationFunction.java | 178 +++++-- ...CountCPCSketchAggregationFunctionTest.java | 163 +++++++ ...untThetaSketchAggregationFunctionTest.java | 155 ++++++ .../tests/custom/UuidAggregationTest.java | 3 +- 5 files changed, 808 insertions(+), 131 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index 2be253256157..ee53f89064eb 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -135,80 +135,156 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - FieldSpec.DataType dataType = blockValSet.getValueType(); - FieldSpec.DataType storedType = dataType.getStoredType(); - - // Treat BYTES value as serialized CPC Sketch - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(aggregationResultHolder); - CpcSketch[] sketches = deserializeSketches(bytesValues, length); - for (CpcSketch sketch : sketches) { - if (sketch != null) { - cpcSketchAccumulator.apply(sketch); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging CPC sketches", e); - } - return; + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { + aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); + } else { + aggregateMV(length, aggregationResultHolder, blockValSet, dataType, storedType); } + } + protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, + DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; } // For non-dictionary-encoded expression, store values into the CpcSketch - CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); switch (storedType) { case INT: + CpcSketch intSketch = getCpcSketch(aggregationResultHolder); int[] intValues = blockValSet.getIntValuesSV(); for (int i = 0; i < length; i++) { - cpcSketch.update(intValues[i]); + intSketch.update(intValues[i]); } break; case LONG: + CpcSketch longSketch = getCpcSketch(aggregationResultHolder); long[] longValues = blockValSet.getLongValuesSV(); for (int i = 0; i < length; i++) { - cpcSketch.update(longValues[i]); + longSketch.update(longValues[i]); } break; case FLOAT: + CpcSketch floatSketch = getCpcSketch(aggregationResultHolder); float[] floatValues = blockValSet.getFloatValuesSV(); for (int i = 0; i < length; i++) { - cpcSketch.update(floatValues[i]); + floatSketch.update(floatValues[i]); } break; case DOUBLE: + CpcSketch doubleSketch = getCpcSketch(aggregationResultHolder); double[] doubleValues = blockValSet.getDoubleValuesSV(); for (int i = 0; i < length; i++) { - cpcSketch.update(doubleValues[i]); + doubleSketch.update(doubleValues[i]); } break; case STRING: + CpcSketch stringSketch = getCpcSketch(aggregationResultHolder); String[] stringValues = blockValSet.getStringValuesSV(); for (int i = 0; i < length; i++) { - cpcSketch.update(stringValues[i]); + stringSketch.update(stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is a serialized CPC sketch and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.UUID) { + CpcSketch uuidSketch = getCpcSketch(aggregationResultHolder); + for (int i = 0; i < length; i++) { + uuidSketch.update(bytesValues[i]); + } + } else { + mergeSerializedSketches(aggregationResultHolder, bytesValues, length); + } + break; + default: + throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); + } + } + + protected void aggregateMV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, + DataType dataType, DataType storedType) { + // For dictionary-encoded expression, store dictionary ids into the bitmap + Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; + if (dictionary != null && dataType != DataType.BYTES) { + int[][] dictIds = blockValSet.getDictionaryIdsMV(); + RoaringBitmap dictIdBitmap = getDictIdBitmap(aggregationResultHolder, dictionary); + for (int i = 0; i < length; i++) { + dictIdBitmap.add(dictIds[i]); + } + return; + } + + // For non-dictionary-encoded expression, store values into the CpcSketch + switch (storedType) { + case INT: + CpcSketch intSketch = getCpcSketch(aggregationResultHolder); + int[][] intValues = blockValSet.getIntValuesMV(); + for (int i = 0; i < length; i++) { + for (int value : intValues[i]) { + intSketch.update(value); + } + } + break; + case LONG: + CpcSketch longSketch = getCpcSketch(aggregationResultHolder); + long[][] longValues = blockValSet.getLongValuesMV(); for (int i = 0; i < length; i++) { - cpcSketch.update(uuidValues[i]); + for (long value : longValues[i]) { + longSketch.update(value); + } + } + break; + case FLOAT: + CpcSketch floatSketch = getCpcSketch(aggregationResultHolder); + float[][] floatValues = blockValSet.getFloatValuesMV(); + for (int i = 0; i < length; i++) { + for (float value : floatValues[i]) { + floatSketch.update(value); + } + } + break; + case DOUBLE: + CpcSketch doubleSketch = getCpcSketch(aggregationResultHolder); + double[][] doubleValues = blockValSet.getDoubleValuesMV(); + for (int i = 0; i < length; i++) { + for (double value : doubleValues[i]) { + doubleSketch.update(value); + } + } + break; + case STRING: + CpcSketch stringSketch = getCpcSketch(aggregationResultHolder); + String[][] stringValues = blockValSet.getStringValuesMV(); + for (int i = 0; i < length; i++) { + for (String value : stringValues[i]) { + stringSketch.update(value); + } + } + break; + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + if (dataType == DataType.UUID) { + CpcSketch uuidSketch = getCpcSketch(aggregationResultHolder); + for (int i = 0; i < length; i++) { + for (byte[] value : bytesValues[i]) { + uuidSketch.update(value); + } + } + } else { + for (int i = 0; i < length; i++) { + mergeSerializedSketches(aggregationResultHolder, bytesValues[i], bytesValues[i].length); + } } break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } - // The updated CpcSketch already lives in the holder (getCpcSketch stored it); extractAggregationResult - // converts it to a CpcSketchAccumulator. Reading the holder as an accumulator here would - // ClassCastException — the holder slot contains the sketch, not an accumulator. } @Override @@ -218,28 +294,18 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - - // Treat BYTES value as serialized CPC Sketch - if (storedType == FieldSpec.DataType.BYTES) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - CpcSketch[] sketches = deserializeSketches(bytesValues, length); - for (int i = 0; i < length; i++) { - CpcSketch sketch = sketches[i]; - if (sketch != null) { - CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(groupByResultHolder, groupKeyArray[i]); - cpcSketchAccumulator.apply(sketch); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while aggregating CPC Sketches", e); - } - return; + if (blockValSet.isSingleValue()) { + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); + } else { + aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); } + } + protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -279,11 +345,20 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is a serialized CPC sketch and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.UUID) { + for (int i = 0; i < length; i++) { + getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(bytesValues[i]); + } + } else { + CpcSketch[] sketches = deserializeSketches(bytesValues, length); + for (int i = 0; i < length; i++) { + CpcSketch sketch = sketches[i]; + if (sketch != null) { + getAccumulator(groupByResultHolder, groupKeyArray[i]).apply(sketch); + } + } } break; default: @@ -291,47 +366,110 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } - @Override - public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - Map blockValSetMap) { - BlockValSet blockValSet = blockValSetMap.get(_expression); - - DataType dataType = blockValSet.getValueType(); - DataType storedType = dataType.getStoredType(); - boolean singleValue = blockValSet.isSingleValue(); - - // UUID columns: update with canonical UUID strings converted from raw bytes (see aggregate() for rationale). - if (dataType == DataType.UUID && singleValue) { - byte[][] uuidBytesValues = blockValSet.getBytesValuesSV(); + protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType dataType, DataType storedType) { + // For dictionary-encoded expression, store dictionary ids into the bitmap + Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; + if (dictionary != null && dataType != DataType.BYTES) { + int[][] dictIds = blockValSet.getDictionaryIdsMV(); for (int i = 0; i < length; i++) { - byte[] canonical = uuidBytesValues[i]; - for (int groupKey : groupKeysArray[i]) { - getCpcSketch(groupByResultHolder, groupKey).update(canonical); - } + getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); } return; } - if (singleValue && storedType == DataType.BYTES) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - CpcSketch[] sketches = deserializeSketches(bytesValues, length); + // For non-dictionary-encoded expression, store values into the CpcSketch + switch (storedType) { + case INT: + int[][] intValues = blockValSet.getIntValuesMV(); for (int i = 0; i < length; i++) { - if (sketches[i] != null) { - for (int groupKey : groupKeysArray[i]) { - getAccumulator(groupByResultHolder, groupKey).apply(sketches[i]); + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (int value : intValues[i]) { + cpcSketch.update(value); + } + } + break; + case LONG: + long[][] longValues = blockValSet.getLongValuesMV(); + for (int i = 0; i < length; i++) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (long value : longValues[i]) { + cpcSketch.update(value); + } + } + break; + case FLOAT: + float[][] floatValues = blockValSet.getFloatValuesMV(); + for (int i = 0; i < length; i++) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (float value : floatValues[i]) { + cpcSketch.update(value); + } + } + break; + case DOUBLE: + double[][] doubleValues = blockValSet.getDoubleValuesMV(); + for (int i = 0; i < length; i++) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (double value : doubleValues[i]) { + cpcSketch.update(value); + } + } + break; + case STRING: + String[][] stringValues = blockValSet.getStringValuesMV(); + for (int i = 0; i < length; i++) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (String value : stringValues[i]) { + cpcSketch.update(value); + } + } + break; + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + if (dataType == DataType.UUID) { + for (int i = 0; i < length; i++) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); + } + } + } else { + for (int i = 0; i < length; i++) { + CpcSketch[] sketches = deserializeSketches(bytesValues[i], bytesValues[i].length); + CpcSketchAccumulator accumulator = getAccumulator(groupByResultHolder, groupKeyArray[i]); + for (CpcSketch sketch : sketches) { + if (sketch != null) { + accumulator.apply(sketch); + } } } } - } catch (Exception e) { - throw new RuntimeException("Caught exception while aggregating CPC sketches", e); - } - return; + break; + default: + throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); + } + } + + @Override + public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, + Map blockValSetMap) { + BlockValSet blockValSet = blockValSetMap.get(_expression); + + DataType dataType = blockValSet.getValueType(); + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); + } else { + aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); } + } + protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -381,12 +519,125 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } break; - // Reached only by UUID: a real BYTES column is a serialized CPC sketch and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.UUID) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + getCpcSketch(groupByResultHolder, groupKey).update(bytesValues[i]); + } + } + } else { + CpcSketch[] sketches = deserializeSketches(bytesValues, length); + for (int i = 0; i < length; i++) { + CpcSketch sketch = sketches[i]; + if (sketch != null) { + for (int groupKey : groupKeysArray[i]) { + getAccumulator(groupByResultHolder, groupKey).apply(sketch); + } + } + } + } + break; + default: + throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); + } + } + + protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType dataType, DataType storedType) { + // For dictionary-encoded expression, store dictionary ids into the bitmap + Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; + if (dictionary != null && dataType != DataType.BYTES) { + int[][] dictIds = blockValSet.getDictionaryIdsMV(); + for (int i = 0; i < length; i++) { + int[] rowDictIds = dictIds[i]; + for (int groupKey : groupKeysArray[i]) { + getDictIdBitmap(groupByResultHolder, groupKey, dictionary).add(rowDictIds); + } + } + return; + } + + // For non-dictionary-encoded expression, store values into the CpcSketch + switch (storedType) { + case INT: + int[][] intValues = blockValSet.getIntValuesMV(); for (int i = 0; i < length; i++) { for (int groupKey : groupKeysArray[i]) { - getCpcSketch(groupByResultHolder, groupKey).update(uuidValues[i]); + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (int value : intValues[i]) { + cpcSketch.update(value); + } + } + } + break; + case LONG: + long[][] longValues = blockValSet.getLongValuesMV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (long value : longValues[i]) { + cpcSketch.update(value); + } + } + } + break; + case FLOAT: + float[][] floatValues = blockValSet.getFloatValuesMV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (float value : floatValues[i]) { + cpcSketch.update(value); + } + } + } + break; + case DOUBLE: + double[][] doubleValues = blockValSet.getDoubleValuesMV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (double value : doubleValues[i]) { + cpcSketch.update(value); + } + } + } + break; + case STRING: + String[][] stringValues = blockValSet.getStringValuesMV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (String value : stringValues[i]) { + cpcSketch.update(value); + } + } + } + break; + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + if (dataType == DataType.UUID) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); + } + } + } + } else { + for (int i = 0; i < length; i++) { + CpcSketch[] sketches = deserializeSketches(bytesValues[i], bytesValues[i].length); + for (int groupKey : groupKeysArray[i]) { + CpcSketchAccumulator accumulator = getAccumulator(groupByResultHolder, groupKey); + for (CpcSketch sketch : sketches) { + if (sketch != null) { + accumulator.apply(sketch); + } + } + } } } break; @@ -564,6 +815,8 @@ private CpcSketch dictionaryToCpcSketch(DictIdsWrapper dictIdsWrapper) { private void addObjectToSketch(Object rawValue, CpcSketch sketch) { if (rawValue instanceof String) { sketch.update((String) rawValue); + } else if (rawValue instanceof byte[]) { + sketch.update((byte[]) rawValue); } else if (rawValue instanceof Integer) { sketch.update((Integer) rawValue); } else if (rawValue instanceof Long) { @@ -585,6 +838,10 @@ private void addObjectsToSketch(Object[] rawValues, CpcSketch sketch) { for (String s : (String[]) rawValues) { sketch.update(s); } + } else if (rawValues instanceof byte[][]) { + for (byte[] bytes : (byte[][]) rawValues) { + sketch.update(bytes); + } } else if (rawValues instanceof Integer[]) { for (Integer i : (Integer[]) rawValues) { sketch.update(i); @@ -627,6 +884,21 @@ private CpcSketchAccumulator getAccumulator(GroupByResultHolder groupByResultHol return accumulator; } + private void mergeSerializedSketches(AggregationResultHolder aggregationResultHolder, byte[][] bytesValues, + int length) { + try { + CpcSketchAccumulator accumulator = getAccumulator(aggregationResultHolder); + CpcSketch[] sketches = deserializeSketches(bytesValues, length); + for (CpcSketch sketch : sketches) { + if (sketch != null) { + accumulator.apply(sketch); + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging CPC sketches", e); + } + } + /// Deserializes the sketches from the bytes. Returns null for empty byte arrays which represent /// the default null value for BYTES columns in Pinot. Callers must handle null entries. @SuppressWarnings({"unchecked"}) 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 89b70290567b..408bcf943bb5 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 @@ -91,6 +91,7 @@ public class DistinctCountThetaSketchAggregationFunction private final List _filterEvaluators; private final ExpressionContext _postAggregationExpression; private final UpdatableThetaSketchBuilder _updateSketchBuilder = new UpdatableThetaSketchBuilder(); + private final ThetaSketch _emptySketch; private int _nominalEntries = ThetaUtil.DEFAULT_NOMINAL_ENTRIES; protected final ThetaSetOperationBuilder _setOperationBuilder = new ThetaSetOperationBuilder(); protected int _accumulatorThreshold = DEFAULT_ACCUMULATOR_THRESHOLD; @@ -116,6 +117,7 @@ public DistinctCountThetaSketchAggregationFunction(List argum _setOperationBuilder.setP(p); _updateSketchBuilder.setP(p); } + _emptySketch = _updateSketchBuilder.build().compact(); if (numArguments < 4) { // Simple union without post-aggregation @@ -192,15 +194,13 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but - // its values are logical scalars rather than serialized sketches, so it takes the scalar path below and is - // handled by `case BYTES`. - boolean uuidMainExpression = - blockValSetMap.get(_inputExpressions.get(0)).getValueType() == DataType.UUID; - if (valueTypes[0] != DataType.BYTES || uuidMainExpression) { + // Main expression is always index 0. Logical BYTES values contain serialized sketches; all other logical types, + // including UUID, use the raw-value path and are handled by their stored-type switch branch. + DataType dataType = valueTypes[0]; + if (dataType != DataType.BYTES) { List updateSketches = getUpdateSketches(aggregationResultHolder); if (singleValues[0]) { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[] intValues = (int[]) valueArrays[0]; if (_includeDefaultSketch) { @@ -315,7 +315,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde + valueTypes[0]); } } else { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[][] intValues = (int[][]) valueArrays[0]; if (_includeDefaultSketch) { @@ -450,11 +450,11 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde break; default: throw new IllegalStateException( - "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); + "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + + valueTypes[0]); } } - } else { - // Serialized sketch + } else if (singleValues[0]) { List thetaSketchAccumulators = getUnions(aggregationResultHolder); ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { @@ -472,6 +472,28 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } } + } else { + List thetaSketchAccumulators = getUnions(aggregationResultHolder); + ThetaSketch[][] sketches = deserializeSketches((byte[][][]) valueArrays[0], length); + if (_includeDefaultSketch) { + ThetaSketchAccumulator defaultThetaAccumulator = thetaSketchAccumulators.get(0); + for (int i = 0; i < length; i++) { + for (ThetaSketch sketch : sketches[i]) { + defaultThetaAccumulator.apply(sketch); + } + } + } + for (int i = 0; i < numFilters; i++) { + FilterEvaluator filterEvaluator = _filterEvaluators.get(i); + ThetaSketchAccumulator thetaSketchAccumulator = thetaSketchAccumulators.get(i + 1); + for (int j = 0; j < length; j++) { + if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { + for (ThetaSketch sketch : sketches[j]) { + thetaSketchAccumulator.apply(sketch); + } + } + } + } } } @@ -485,14 +507,12 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but - // its values are logical scalars rather than serialized sketches, so it takes the scalar path below and is - // handled by `case BYTES`. - boolean uuidMainExpression = - blockValSetMap.get(_inputExpressions.get(0)).getValueType() == DataType.UUID; - if (valueTypes[0] != DataType.BYTES || uuidMainExpression) { + // Main expression is always index 0. Logical BYTES values contain serialized sketches; all other logical types, + // including UUID, use the raw-value path and are handled by their stored-type switch branch. + DataType dataType = valueTypes[0]; + if (dataType != DataType.BYTES) { if (singleValues[0]) { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[] intValues = (int[]) valueArrays[0]; for (int i = 0; i < length; i++) { @@ -571,7 +591,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol case BYTES: byte[][] uuidValues = (byte[][]) valueArrays[0]; for (int i = 0; i < length; i++) { - List updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]); + List updateSketches = + getUpdateSketches(groupByResultHolder, groupKeyArray[i]); byte[] value = uuidValues[i]; if (_includeDefaultSketch) { updateSketches.get(0).update(value); @@ -589,7 +610,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol + valueTypes[0]); } } else { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[][] intValues = (int[][]) valueArrays[0]; for (int i = 0; i < length; i++) { @@ -698,7 +719,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol case BYTES: byte[][][] uuidValues = (byte[][][]) valueArrays[0]; for (int i = 0; i < length; i++) { - List updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]); + List updateSketches = + getUpdateSketches(groupByResultHolder, groupKeyArray[i]); byte[][] values = uuidValues[i]; if (_includeDefaultSketch) { UpdatableThetaSketch defaultSketch = updateSketches.get(0); @@ -718,14 +740,15 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol break; default: throw new IllegalStateException( - "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); + "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + + valueTypes[0]); } } - } else { - // Serialized sketch + } else if (singleValues[0]) { ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); for (int i = 0; i < length; i++) { - List thetaSketchAccumulators = getUnions(groupByResultHolder, groupKeyArray[i]); + List thetaSketchAccumulators = + getUnions(groupByResultHolder, groupKeyArray[i]); ThetaSketch sketch = sketches[i]; if (_includeDefaultSketch) { thetaSketchAccumulators.get(0).apply(sketch); @@ -736,6 +759,26 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } } + } else { + ThetaSketch[][] sketches = deserializeSketches((byte[][][]) valueArrays[0], length); + for (int i = 0; i < length; i++) { + List thetaSketchAccumulators = + getUnions(groupByResultHolder, groupKeyArray[i]); + if (_includeDefaultSketch) { + ThetaSketchAccumulator defaultThetaAccumulator = thetaSketchAccumulators.get(0); + for (ThetaSketch sketch : sketches[i]) { + defaultThetaAccumulator.apply(sketch); + } + } + for (int j = 0; j < numFilters; j++) { + if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) { + ThetaSketchAccumulator thetaSketchAccumulator = thetaSketchAccumulators.get(j + 1); + for (ThetaSketch sketch : sketches[i]) { + thetaSketchAccumulator.apply(sketch); + } + } + } + } } } @@ -749,14 +792,12 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0. A UUID column reports BYTES here because that is its stored type, but - // its values are logical scalars rather than serialized sketches, so it takes the scalar path below and is - // handled by `case BYTES`. - boolean uuidMainExpression = - blockValSetMap.get(_inputExpressions.get(0)).getValueType() == DataType.UUID; - if (valueTypes[0] != DataType.BYTES || uuidMainExpression) { + // Main expression is always index 0. Logical BYTES values contain serialized sketches; all other logical types, + // including UUID, use the raw-value path and are handled by their stored-type switch branch. + DataType dataType = valueTypes[0]; + if (dataType != DataType.BYTES) { if (singleValues[0]) { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[] intValues = (int[]) valueArrays[0]; if (_includeDefaultSketch) { @@ -870,7 +911,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult FilterEvaluator filterEvaluator = _filterEvaluators.get(i); for (int j = 0; j < length; j++) { if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { - for (int groupKey : groupKeysArray[i]) { + for (int groupKey : groupKeysArray[j]) { getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(uuidValues[j]); } } @@ -883,7 +924,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult + valueTypes[0]); } } else { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[][] intValues = (int[][]) valueArrays[0]; if (_includeDefaultSketch) { @@ -1030,9 +1071,10 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult FilterEvaluator filterEvaluator = _filterEvaluators.get(i); for (int j = 0; j < length; j++) { if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { - for (int groupKey : groupKeysArray[i]) { - UpdatableThetaSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1); - for (byte[] value : uuidValues[i]) { + for (int groupKey : groupKeysArray[j]) { + UpdatableThetaSketch updateSketch = + getUpdateSketches(groupByResultHolder, groupKey).get(i + 1); + for (byte[] value : uuidValues[j]) { updateSketch.update(value); } } @@ -1042,11 +1084,11 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult break; default: throw new IllegalStateException( - "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); + "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + + valueTypes[0]); } } - } else { - // Serialized sketch + } else if (singleValues[0]) { ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { @@ -1059,8 +1101,33 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult 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).apply(sketches[i]); + for (int groupKey : groupKeysArray[j]) { + getUnions(groupByResultHolder, groupKey).get(i + 1).apply(sketches[j]); + } + } + } + } + } else { + ThetaSketch[][] sketches = deserializeSketches((byte[][][]) valueArrays[0], length); + if (_includeDefaultSketch) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + ThetaSketchAccumulator defaultThetaAccumulator = getUnions(groupByResultHolder, groupKey).get(0); + for (ThetaSketch sketch : sketches[i]) { + defaultThetaAccumulator.apply(sketch); + } + } + } + } + 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[j]) { + ThetaSketchAccumulator thetaSketchAccumulator = getUnions(groupByResultHolder, groupKey).get(i + 1); + for (ThetaSketch sketch : sketches[j]) { + thetaSketchAccumulator.apply(sketch); + } } } } @@ -1369,7 +1436,7 @@ private void extractValues(Map blockValSetMap, b DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); singleValues[i] = singleValue; - valueTypes[i] = storedType; + valueTypes[i] = dataType; if (singleValue) { switch (storedType) { case INT: @@ -1410,6 +1477,9 @@ private void extractValues(Map blockValSetMap, b case STRING: valueArrays[i] = blockValSet.getStringValuesMV(); break; + case BYTES: + valueArrays[i] = blockValSet.getBytesValuesMV(); + break; default: throw new IllegalStateException(); } @@ -1483,11 +1553,24 @@ private List buildUnions() { private ThetaSketch[] deserializeSketches(byte[][] serializedSketches, int length) { ThetaSketch[] sketches = new ThetaSketch[length]; for (int i = 0; i < length; i++) { - sketches[i] = ThetaSketch.wrap(MemorySegment.ofArray(serializedSketches[i]).asReadOnly()); + sketches[i] = deserializeSketch(serializedSketches[i]); } return sketches; } + private ThetaSketch[][] deserializeSketches(byte[][][] serializedSketches, int length) { + ThetaSketch[][] sketches = new ThetaSketch[length][]; + for (int i = 0; i < length; i++) { + sketches[i] = deserializeSketches(serializedSketches[i], serializedSketches[i].length); + } + return sketches; + } + + private ThetaSketch deserializeSketch(byte[] serializedSketch) { + return serializedSketch.length == 0 ? _emptySketch + : ThetaSketch.wrap(MemorySegment.ofArray(serializedSketch).asReadOnly()); + } + /// Evaluates the post-aggregation expression. protected ThetaSketch evaluatePostAggregationExpression(List sketches) { return evaluatePostAggregationExpression(_postAggregationExpression, sketches); @@ -1648,7 +1731,7 @@ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] _predicateEvaluator = PredicateEvaluatorProvider.getPredicateEvaluator(_predicate, null, valueType, null); } if (singleValue) { - switch (valueType) { + switch (valueType.getStoredType()) { case INT: return _predicateEvaluator.applySV(((int[]) valueArray)[index]); case LONG: @@ -1665,7 +1748,7 @@ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] throw new IllegalStateException(); } } else { - switch (valueType) { + switch (valueType.getStoredType()) { case INT: int[] intValues = ((int[][]) valueArray)[index]; return _predicateEvaluator.applyMV(intValues, intValues.length); @@ -1681,6 +1764,9 @@ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] case STRING: String[] stringValues = ((String[][]) valueArray)[index]; return _predicateEvaluator.applyMV(stringValues, stringValues.length); + case BYTES: + byte[][] bytesValues = ((byte[][][]) valueArray)[index]; + return _predicateEvaluator.applyMV(bytesValues, bytesValues.length); default: throw new IllegalStateException(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java index a3bcda6204da..80faf2555253 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java @@ -26,15 +26,39 @@ import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; 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.groupby.GroupByResultHolder; import org.apache.pinot.segment.local.customobject.CpcSketchAccumulator; import org.apache.pinot.segment.local.customobject.SerializedCPCSketch; import org.apache.pinot.segment.spi.Constants; +import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + public class DistinctCountCPCSketchAggregationFunctionTest { + private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); + private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); + private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); + private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); + private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); + private static final byte[][] UUID_VALUES_SV = {UUID_0, UUID_1, UUID_0, UUID_2}; + private static final byte[][][] UUID_VALUES_MV = {{UUID_0, UUID_1}, {UUID_1, UUID_2}, {UUID_0}, {UUID_3}}; + + @DataProvider(name = "uuidValueModes") + public static Object[][] uuidValueModes() { + return new Object[][]{{true}, {false}}; + } @Test public void testCanUseStarTreeDefaultLgK() { @@ -138,4 +162,143 @@ public void testMergeWithEmptyAccumulators() { result = function.merge(new CpcSketchAccumulator(12, 2), new CpcSketchAccumulator(12, 2)); Assert.assertTrue(result.isEmpty()); } + + @Test(dataProvider = "uuidValueModes") + public void testAggregateUuid(boolean singleValue) { + DistinctCountCPCSketchAggregationFunction function = + new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); + BlockValSet blockValSet = mockUuidBlockValSet(singleValue); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + + function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder), singleValue ? 3L : 4L); + verifyBytesAccessor(blockValSet, singleValue); + } + + @Test(dataProvider = "uuidValueModes") + public void testAggregateUuidGroupBySV(boolean singleValue) { + DistinctCountCPCSketchAggregationFunction function = + new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); + BlockValSet blockValSet = mockUuidBlockValSet(singleValue); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); + + function.aggregateGroupBySV(UUID_VALUES_SV.length, new int[]{0, 0, 1, 1}, resultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 2L : 3L); + Assert.assertEquals(extractFinalResult(function, resultHolder, 1), 2L); + verifyBytesAccessor(blockValSet, singleValue); + } + + @Test(dataProvider = "uuidValueModes") + public void testAggregateUuidGroupByMV(boolean singleValue) { + DistinctCountCPCSketchAggregationFunction function = + new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); + BlockValSet blockValSet = mockUuidBlockValSet(singleValue); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); + + function.aggregateGroupByMV(UUID_VALUES_SV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, resultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 1L : 2L); + Assert.assertEquals(extractFinalResult(function, resultHolder, 1), singleValue ? 3L : 4L); + verifyBytesAccessor(blockValSet, singleValue); + } + + @Test + public void testAggregateDictionaryEncodedUuid() { + DistinctCountCPCSketchAggregationFunction function = + new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); + Dictionary dictionary = mock(Dictionary.class); + when(dictionary.get(0)).thenReturn(UUID_0); + when(dictionary.get(1)).thenReturn(UUID_1); + when(dictionary.get(2)).thenReturn(UUID_2); + + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); + when(blockValSet.isSingleValue()).thenReturn(true); + when(blockValSet.isDictionaryEncoded()).thenReturn(true); + when(blockValSet.getDictionary()).thenReturn(dictionary); + when(blockValSet.getDictionaryIdsSV()).thenReturn(new int[]{0, 1, 0, 2}); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + + function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder), 3L); + } + + @Test + public void testAggregateMultiValueSerializedSketches() { + byte[][][] serializedSketches = { + {serializedSketch("a"), serializedSketch("b")}, + {serializedSketch("b"), serializedSketch("c")}, + {new byte[0]}, + {serializedSketch("d")} + }; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + when(blockValSet.getBytesValuesMV()).thenReturn(serializedSketches); + DistinctCountCPCSketchAggregationFunction function = + new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 1L); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + + private static BlockValSet mockUuidBlockValSet(boolean singleValue) { + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); + when(blockValSet.isSingleValue()).thenReturn(singleValue); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + if (singleValue) { + when(blockValSet.getBytesValuesSV()).thenReturn(UUID_VALUES_SV); + } else { + when(blockValSet.getBytesValuesMV()).thenReturn(UUID_VALUES_MV); + } + return blockValSet; + } + + private static void verifyBytesAccessor(BlockValSet blockValSet, boolean singleValue) { + if (singleValue) { + verify(blockValSet, atLeastOnce()).getBytesValuesSV(); + verify(blockValSet, never()).getBytesValuesMV(); + } else { + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + } + + private static byte[] serializedSketch(String value) { + CpcSketch sketch = new CpcSketch(); + sketch.update(value); + return sketch.toByteArray(); + } + + private static long extractFinalResult(DistinctCountCPCSketchAggregationFunction function, + AggregationResultHolder resultHolder) { + return ((Number) function.extractFinalResult(function.extractAggregationResult(resultHolder))).longValue(); + } + + private static long extractFinalResult(DistinctCountCPCSketchAggregationFunction function, + GroupByResultHolder resultHolder, int groupKey) { + return ((Number) function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey))).longValue(); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java index 85eac01d28b7..f26b69caca78 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java @@ -20,14 +20,40 @@ import java.util.List; import java.util.Map; +import org.apache.datasketches.theta.UpdatableThetaSketch; +import org.apache.datasketches.theta.UpdatableThetaSketchBuilder; import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; +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.segment.spi.Constants; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; +import org.testng.annotations.DataProvider; import org.testng.annotations.Test; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + public class DistinctCountThetaSketchAggregationFunctionTest { + private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); + private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); + private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); + private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); + private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); + private static final byte[][] UUID_VALUES_SV = {UUID_0, UUID_1, UUID_0, UUID_2}; + private static final byte[][][] UUID_VALUES_MV = {{UUID_0, UUID_1}, {UUID_1, UUID_2}, {UUID_0}, {UUID_3}}; + + @DataProvider(name = "uuidValueModes") + public static Object[][] uuidValueModes() { + return new Object[][]{{true}, {false}}; + } @Test public void testCanUseStarTreeDefaultK() { @@ -54,4 +80,133 @@ public void testCanUseCustomK() { Assert.assertTrue(function.canUseStarTree(Map.of(Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES, 32768))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES, "32768"))); } + + @Test(dataProvider = "uuidValueModes") + public void testAggregateUuid(boolean singleValue) { + DistinctCountThetaSketchAggregationFunction function = + new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); + BlockValSet blockValSet = mockUuidBlockValSet(singleValue); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + + function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder), singleValue ? 3L : 4L); + verifyBytesAccessor(blockValSet, singleValue); + } + + @Test(dataProvider = "uuidValueModes") + public void testAggregateUuidGroupBySV(boolean singleValue) { + DistinctCountThetaSketchAggregationFunction function = + new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); + BlockValSet blockValSet = mockUuidBlockValSet(singleValue); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); + + function.aggregateGroupBySV(UUID_VALUES_SV.length, new int[]{0, 0, 1, 1}, resultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 2L : 3L); + Assert.assertEquals(extractFinalResult(function, resultHolder, 1), 2L); + verifyBytesAccessor(blockValSet, singleValue); + } + + @Test(dataProvider = "uuidValueModes") + public void testAggregateUuidGroupByMV(boolean singleValue) { + DistinctCountThetaSketchAggregationFunction function = + new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); + BlockValSet blockValSet = mockUuidBlockValSet(singleValue); + GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); + + function.aggregateGroupByMV(UUID_VALUES_SV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, resultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 1L : 2L); + Assert.assertEquals(extractFinalResult(function, resultHolder, 1), singleValue ? 3L : 4L); + verifyBytesAccessor(blockValSet, singleValue); + } + + @Test + public void testUuidPredicateUsesLogicalType() { + DistinctCountThetaSketchAggregationFunction function = new DistinctCountThetaSketchAggregationFunction( + List.of(UUID_EXPRESSION, ExpressionContext.forLiteral(Literal.stringValue("")), + ExpressionContext.forLiteral(Literal.stringValue("uuidCol = '550e8400-e29b-41d4-a716-446655440000'")), + ExpressionContext.forLiteral(Literal.stringValue("$1")))); + BlockValSet blockValSet = mockUuidBlockValSet(true); + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + + function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder), 1L); + } + + @Test + public void testAggregateMultiValueSerializedSketches() { + byte[][][] serializedSketches = { + {serializedSketch("a"), serializedSketch("b")}, + {serializedSketch("b"), serializedSketch("c")}, + {new byte[0]}, + {serializedSketch("d")} + }; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.getBytesValuesMV()).thenReturn(serializedSketches); + DistinctCountThetaSketchAggregationFunction function = + new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 1L); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + + private static BlockValSet mockUuidBlockValSet(boolean singleValue) { + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); + when(blockValSet.isSingleValue()).thenReturn(singleValue); + if (singleValue) { + when(blockValSet.getBytesValuesSV()).thenReturn(UUID_VALUES_SV); + } else { + when(blockValSet.getBytesValuesMV()).thenReturn(UUID_VALUES_MV); + } + return blockValSet; + } + + private static void verifyBytesAccessor(BlockValSet blockValSet, boolean singleValue) { + if (singleValue) { + verify(blockValSet, atLeastOnce()).getBytesValuesSV(); + verify(blockValSet, never()).getBytesValuesMV(); + } else { + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + } + + private static byte[] serializedSketch(String value) { + UpdatableThetaSketch sketch = new UpdatableThetaSketchBuilder().build(); + sketch.update(value); + return sketch.compact().toByteArray(); + } + + private static long extractFinalResult(DistinctCountThetaSketchAggregationFunction function, + AggregationResultHolder resultHolder) { + return ((Number) function.extractFinalResult(function.extractAggregationResult(resultHolder))).longValue(); + } + + private static long extractFinalResult(DistinctCountThetaSketchAggregationFunction function, + GroupByResultHolder resultHolder, int groupKey) { + return ((Number) function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey))).longValue(); + } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java index 016ed27e3606..b0101f6bb43d 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java @@ -194,7 +194,8 @@ public void testDistinctOnUuidColumn() public void testDistinctCountOnUuidColumn() throws Exception { setUseMultiStageQueryEngine(false); - for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTBITMAP")) { + for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTBITMAP", + "DISTINCTCOUNTTHETASKETCH", "DISTINCTCOUNTCPCSKETCH")) { JsonNode rows = query(String.format("SELECT %s(%s) FROM %s", function, UUID_COLUMN, getTableName())); assertEquals(rows.get(0).get(0).asLong(), NUM_DISTINCT, function + ": " + rows.toPrettyString()); } From 81667cadfce9f1915c86854629ae1538622b6a1b Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 14 Aug 2026 15:22:27 -0700 Subject: [PATCH 14/19] Address UUID aggregation review feedback --- .../function/AggregationFunctionUtils.java | 21 +-- ...istinctCountBitmapAggregationFunction.java | 130 +++++--------- ...inctCountCPCSketchAggregationFunction.java | 110 +++++------- .../DistinctCountHLLAggregationFunction.java | 163 +++++++----------- ...stinctCountHLLPlusAggregationFunction.java | 163 ++++++++---------- ...ctCountThetaSketchAggregationFunction.java | 128 +++----------- .../DistinctCountULLAggregationFunction.java | 147 +++++++--------- .../AggregationFunctionUtilsTest.java | 84 +++++++++ ...tCountBitmapMVAggregationFunctionTest.java | 61 +++++++ ...CountCPCSketchAggregationFunctionTest.java | 95 ++++++++-- ...stinctCountHLLAggregationFunctionTest.java | 11 +- ...inctCountHLLMVAggregationFunctionTest.java | 61 +++++++ ...CountHLLPlusMVAggregationFunctionTest.java | 61 +++++++ ...untThetaSketchAggregationFunctionTest.java | 68 ++++++-- ...stinctCountULLAggregationFunctionTest.java | 112 ++++++++++++ .../tests/custom/UuidAggregationTest.java | 4 +- 16 files changed, 837 insertions(+), 582 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index 6908ef1f19d8..8cb23df58471 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -802,11 +802,8 @@ private static HyperLogLogPlus getDistinctValueHLLPlus(Dictionary dictionary, in private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, DistinctCountHLLAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary), but its entries are logical - // scalars, not serialized sketch state. Excluding it here lets it fall through to the scalar path - // below, which offers dictionary.get(i) -- the stored byte[] -- exactly as the scan path does. - if (dataSource.getDataSourceMetadata().getDataType() != FieldSpec.DataType.UUID - && dictionary.getValueType() == FieldSpec.DataType.BYTES) { + if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES + && dataSource.getDataSourceMetadata().isSingleValue()) { // Treat BYTES value as serialized HyperLogLog try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); @@ -828,11 +825,8 @@ private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, private static HyperLogLogPlus getDistinctCountHLLPlusResult(DataSource dataSource, DistinctCountHLLPlusAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary), but its entries are logical - // scalars, not serialized sketch state. Excluding it here lets it fall through to the scalar path - // below, which offers dictionary.get(i) -- the stored byte[] -- exactly as the scan path does. - if (dataSource.getDataSourceMetadata().getDataType() != FieldSpec.DataType.UUID - && dictionary.getValueType() == FieldSpec.DataType.BYTES) { + if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES + && dataSource.getDataSourceMetadata().isSingleValue()) { // Treat BYTES value as serialized HyperLogLogPlus try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); @@ -874,11 +868,8 @@ private static Object getDistinctCountSmartHLLPlusResult(Dictionary dictionary, private static UltraLogLog getDistinctCountULLResult(DataSource dataSource, DistinctCountULLAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary), but its entries are logical - // scalars, not serialized sketch state. Excluding it here lets it fall through to the scalar path - // below, which offers dictionary.get(i) -- the stored byte[] -- exactly as the scan path does. - if (dataSource.getDataSourceMetadata().getDataType() != FieldSpec.DataType.UUID - && dictionary.getValueType() == FieldSpec.DataType.BYTES) { + if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES + && dataSource.getDataSourceMetadata().isSingleValue()) { // Treat BYTES value as serialized UltraLogLog and merge try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index 20cbc7999e68..55e9feb1420a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -75,36 +75,18 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized RoaringBitmap - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - RoaringBitmap valueBitmap = aggregationResultHolder.getResult(); - if (valueBitmap != null) { - for (int i = 0; i < length; i++) { - valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); - } - } else { - valueBitmap = RoaringBitmapUtils.deserialize(bytesValues[0]); - aggregationResultHolder.setValue(valueBitmap); - for (int i = 1; i < length; i++) { - valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); - } - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSV(length, aggregationResultHolder, blockValSet, storedType); + aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); } else { aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType storedType) { + DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; @@ -141,11 +123,16 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult valueBitmap.add(stringValues[i].hashCode()); } break; - // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - valueBitmap.add(Arrays.hashCode(uuidValues[i])); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); + } + } else { + for (int i = 0; i < length; i++) { + valueBitmap.add(Arrays.hashCode(bytesValues[i])); + } } break; default: @@ -208,11 +195,10 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; - // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. case BYTES: - byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { - for (byte[] value : uuidValues[i]) { + for (byte[] value : bytesValues[i]) { valueBitmap.add(Arrays.hashCode(value)); } } @@ -231,34 +217,18 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized RoaringBitmap - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); - int groupKey = groupKeyArray[i]; - RoaringBitmap valueBitmap = groupByResultHolder.getResult(groupKey); - if (valueBitmap != null) { - valueBitmap.or(value); - } else { - groupByResultHolder.setValueForKey(groupKey, value); - } - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); } else { aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType storedType) { + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -298,11 +268,17 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(stringValues[i].hashCode()); } break; - // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(Arrays.hashCode(uuidValues[i])); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + getValueBitmap(groupByResultHolder, groupKeyArray[i]) + .or(RoaringBitmapUtils.deserialize(bytesValues[i])); + } + } else { + for (int i = 0; i < length; i++) { + getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(Arrays.hashCode(bytesValues[i])); + } } break; default: @@ -367,12 +343,11 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; - // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. case BYTES: - byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { RoaringBitmap bitmap = getValueBitmap(groupByResultHolder, groupKeyArray[i]); - for (byte[] value : uuidValues[i]) { + for (byte[] value : bytesValues[i]) { bitmap.add(Arrays.hashCode(value)); } } @@ -391,36 +366,18 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized RoaringBitmap - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); - for (int groupKey : groupKeysArray[i]) { - RoaringBitmap bitmap = groupByResultHolder.getResult(groupKey); - if (bitmap != null) { - bitmap.or(value); - } else { - // Clone a bitmap for the group - groupByResultHolder.setValueForKey(groupKey, value.clone()); - } - } - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); } else { aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType storedType) { + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -460,11 +417,19 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i].hashCode()); } break; - // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], Arrays.hashCode(uuidValues[i])); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + getValueBitmap(groupByResultHolder, groupKey).or(value); + } + } + } else { + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], Arrays.hashCode(bytesValues[i])); + } } break; default: @@ -541,13 +506,12 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; - // Reached only by UUID: a real BYTES column is serialized sketch state and is handled above. case BYTES: - byte[][][] uuidValues = blockValSet.getBytesValuesMV(); + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { for (int groupKey : groupKeysArray[i]) { RoaringBitmap bitmap = getValueBitmap(groupByResultHolder, groupKey); - for (byte[] value : uuidValues[i]) { + for (byte[] value : bytesValues[i]) { bitmap.add(Arrays.hashCode(value)); } } @@ -719,8 +683,6 @@ private static RoaringBitmap convertToValueBitmap(DictIdsWrapper dictIdsWrapper) valueBitmap.add(dictionary.getStringValue(iterator.next()).hashCode()); } break; - // A UUID column's dictionary reports BYTES (it is a plain BytesDictionary); hash the stored bytes to match - // the scan path above. A real BYTES column holds serialized bitmaps and never reaches the dictionary path. case BYTES: while (iterator.hasNext()) { valueBitmap.add(Arrays.hashCode(dictionary.getBytesValue(iterator.next()))); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index ee53f89064eb..5407ccfe3e36 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -140,7 +140,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde if (blockValSet.isSingleValue()) { aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); } else { - aggregateMV(length, aggregationResultHolder, blockValSet, dataType, storedType); + aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } @@ -193,13 +193,13 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.UUID) { - CpcSketch uuidSketch = getCpcSketch(aggregationResultHolder); + if (dataType == DataType.BYTES) { + mergeSerializedSketches(aggregationResultHolder, bytesValues, length); + } else { + CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); for (int i = 0; i < length; i++) { - uuidSketch.update(bytesValues[i]); + cpcSketch.update(bytesValues[i]); } - } else { - mergeSerializedSketches(aggregationResultHolder, bytesValues, length); } break; default: @@ -208,10 +208,10 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult } protected void aggregateMV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType dataType, DataType storedType) { + DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[][] dictIds = blockValSet.getDictionaryIdsMV(); RoaringBitmap dictIdBitmap = getDictIdBitmap(aggregationResultHolder, dictionary); for (int i = 0; i < length; i++) { @@ -269,16 +269,10 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult break; case BYTES: byte[][][] bytesValues = blockValSet.getBytesValuesMV(); - if (dataType == DataType.UUID) { - CpcSketch uuidSketch = getCpcSketch(aggregationResultHolder); - for (int i = 0; i < length; i++) { - for (byte[] value : bytesValues[i]) { - uuidSketch.update(value); - } - } - } else { - for (int i = 0; i < length; i++) { - mergeSerializedSketches(aggregationResultHolder, bytesValues[i], bytesValues[i].length); + CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); + for (int i = 0; i < length; i++) { + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); } } break; @@ -297,7 +291,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol if (blockValSet.isSingleValue()) { aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); } else { - aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } @@ -347,11 +341,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.UUID) { - for (int i = 0; i < length; i++) { - getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(bytesValues[i]); - } - } else { + if (dataType == DataType.BYTES) { CpcSketch[] sketches = deserializeSketches(bytesValues, length); for (int i = 0; i < length; i++) { CpcSketch sketch = sketches[i]; @@ -359,6 +349,10 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getAccumulator(groupByResultHolder, groupKeyArray[i]).apply(sketch); } } + } else { + for (int i = 0; i < length; i++) { + getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(bytesValues[i]); + } } break; default: @@ -367,10 +361,10 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu } protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[][] dictIds = blockValSet.getDictionaryIdsMV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -427,22 +421,10 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu break; case BYTES: byte[][][] bytesValues = blockValSet.getBytesValuesMV(); - if (dataType == DataType.UUID) { - for (int i = 0; i < length; i++) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); - for (byte[] value : bytesValues[i]) { - cpcSketch.update(value); - } - } - } else { - for (int i = 0; i < length; i++) { - CpcSketch[] sketches = deserializeSketches(bytesValues[i], bytesValues[i].length); - CpcSketchAccumulator accumulator = getAccumulator(groupByResultHolder, groupKeyArray[i]); - for (CpcSketch sketch : sketches) { - if (sketch != null) { - accumulator.apply(sketch); - } - } + for (int i = 0; i < length; i++) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); } } break; @@ -461,7 +443,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult if (blockValSet.isSingleValue()) { aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); } else { - aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } @@ -521,13 +503,7 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.UUID) { - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - getCpcSketch(groupByResultHolder, groupKey).update(bytesValues[i]); - } - } - } else { + if (dataType == DataType.BYTES) { CpcSketch[] sketches = deserializeSketches(bytesValues, length); for (int i = 0; i < length; i++) { CpcSketch sketch = sketches[i]; @@ -537,6 +513,12 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } } + } else { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + getCpcSketch(groupByResultHolder, groupKey).update(bytesValues[i]); + } + } } break; default: @@ -545,10 +527,10 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR } protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[][] dictIds = blockValSet.getDictionaryIdsMV(); for (int i = 0; i < length; i++) { int[] rowDictIds = dictIds[i]; @@ -618,25 +600,11 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR break; case BYTES: byte[][][] bytesValues = blockValSet.getBytesValuesMV(); - if (dataType == DataType.UUID) { - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); - for (byte[] value : bytesValues[i]) { - cpcSketch.update(value); - } - } - } - } else { - for (int i = 0; i < length; i++) { - CpcSketch[] sketches = deserializeSketches(bytesValues[i], bytesValues[i].length); - for (int groupKey : groupKeysArray[i]) { - CpcSketchAccumulator accumulator = getAccumulator(groupByResultHolder, groupKey); - for (CpcSketch sketch : sketches) { - if (sketch != null) { - accumulator.apply(sketch); - } - } + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); } } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java index 666903857e9b..98d767b77367 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -84,42 +84,20 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized HyperLogLog - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); - if (hyperLogLog != null) { - for (int i = 0; i < length; i++) { - hyperLogLog.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i])); - } - } else { - hyperLogLog = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[0]); - aggregationResultHolder.setValue(hyperLogLog); - for (int i = 1; i < length; i++) { - hyperLogLog.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i])); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogs", e); - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSV(length, aggregationResultHolder, blockValSet, storedType); + aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); } else { aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType storedType) { + DataType dataType, DataType storedType) { // For dictionary-encoded expression, collect dictionary ids into a BitSet for deduplication. // BitSet gives O(1) insertion with no container-switching overhead (unlike RoaringBitmap), and uses // dictSize/8 bytes of memory (e.g. 128 KB for a 1M-entry dictionary). Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); BitSet bitSet = getDictIdBitSet(aggregationResultHolder, dictionary); for (int i = 0; i < length; i++) { @@ -129,7 +107,7 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult } // For non-dictionary-encoded expression, store values into the HyperLogLog - HyperLogLog hyperLogLog = getHyperLogLog(aggregationResultHolder); + HyperLogLog hyperLogLog = dataType == DataType.BYTES ? null : getHyperLogLog(aggregationResultHolder); switch (storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); @@ -161,11 +139,16 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult hyperLogLog.offer(stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - hyperLogLog.offer(uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLog(aggregationResultHolder, bytesValues[i]); + } + } else { + for (int i = 0; i < length; i++) { + hyperLogLog.offer(bytesValues[i]); + } } break; default: @@ -231,11 +214,10 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; - // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. case BYTES: - byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { - for (byte[] value : uuidValuesArray[i]) { + for (byte[] value : bytesValuesArray[i]) { hyperLogLog.offer(value); } } @@ -253,41 +235,21 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized HyperLogLog - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - for (int i = 0; i < length; i++) { - HyperLogLog value = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i]); - int groupKey = groupKeyArray[i]; - HyperLogLog hyperLogLog = groupByResultHolder.getResult(groupKey); - if (hyperLogLog != null) { - hyperLogLog.addAll(value); - } else { - groupByResultHolder.setValueForKey(groupKey, value); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogs", e); - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); } else { aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType storedType) { + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, collect dictionary ids into a RoaringBitmap for deduplication. // RoaringBitmap is used (not BitSet) because it is sparse: memory scales with the number of distinct dict IDs // seen per group, not with the full dictionary size. This avoids OOM when many groups each see few distinct values // (contrast with the non-group-by path, which uses a single BitSet across the entire dictionary). Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -327,11 +289,16 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLog(groupByResultHolder, groupKeyArray[i], bytesValues[i]); + } + } else { + for (int i = 0; i < length; i++) { + getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); + } } break; default: @@ -398,12 +365,11 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; - // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. case BYTES: - byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { HyperLogLog hyperLogLog = getHyperLogLog(groupByResultHolder, groupKeyArray[i]); - for (byte[] value : uuidValuesArray[i]) { + for (byte[] value : bytesValuesArray[i]) { hyperLogLog.offer(value); } } @@ -421,41 +387,18 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized HyperLogLog - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - for (int i = 0; i < length; i++) { - HyperLogLog value = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i]); - for (int groupKey : groupKeysArray[i]) { - HyperLogLog hyperLogLog = groupByResultHolder.getResult(groupKey); - if (hyperLogLog != null) { - hyperLogLog.addAll(value); - } else { - // Create a new HyperLogLog for the group - groupByResultHolder.setValueForKey(groupKey, - ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i])); - } - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogs", e); - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); } else { aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType storedType) { + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, collect dictionary ids into a RoaringBitmap (see aggregateSVGroupBySV). Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { int dictId = dictIds[i]; @@ -498,11 +441,18 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + mergeSerializedHyperLogLog(groupByResultHolder, groupKey, bytesValues[i]); + } + } + } else { + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); + } } break; default: @@ -587,14 +537,13 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; - // Reached only by UUID: a real BYTES column is serialized HLL state and is handled above. case BYTES: - byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { - byte[][] uuidValues = uuidValuesArray[i]; + byte[][] bytesValues = bytesValuesArray[i]; for (int groupKey : groupKeysArray[i]) { HyperLogLog hyperLogLog = getHyperLogLog(groupByResultHolder, groupKey); - for (byte[] value : uuidValues) { + for (byte[] value : bytesValues) { hyperLogLog.offer(value); } } @@ -723,6 +672,26 @@ protected static BitSet getDictIdBitSet(AggregationResultHolder aggregationResul return dictIdsWrapper._bitSet; } + private void mergeSerializedHyperLogLog(AggregationResultHolder aggregationResultHolder, byte[] bytes) { + HyperLogLog value = deserializeHyperLogLog(bytes); + HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); + aggregationResultHolder.setValue(hyperLogLog == null ? value : merge(hyperLogLog, value)); + } + + private void mergeSerializedHyperLogLog(GroupByResultHolder groupByResultHolder, int groupKey, byte[] bytes) { + HyperLogLog value = deserializeHyperLogLog(bytes); + HyperLogLog hyperLogLog = groupByResultHolder.getResult(groupKey); + groupByResultHolder.setValueForKey(groupKey, hyperLogLog == null ? value : merge(hyperLogLog, value)); + } + + private static HyperLogLog deserializeHyperLogLog(byte[] bytes) { + try { + return ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytes); + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogs", e); + } + } + /// Returns the HyperLogLog from the result holder or creates a new one if it does not exist. protected HyperLogLog getHyperLogLog(AggregationResultHolder aggregationResultHolder) { HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java index dce1e55ceaf5..8570ae195812 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -96,45 +96,26 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized HyperLogLogPlus - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); - if (hyperLogLogPlus == null) { - hyperLogLogPlus = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[0]); - aggregationResultHolder.setValue(hyperLogLogPlus); - } else { - hyperLogLogPlus.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[0])); - } - for (int i = 1; i < length; i++) { - hyperLogLogPlus.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i])); - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSV(length, aggregationResultHolder, blockValSet, storedType); + aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); } else { aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType storedType) { + DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; } // For non-dictionary-encoded expression, store values into the HyperLogLogPlus - HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(aggregationResultHolder); + HyperLogLogPlus hyperLogLogPlus = + dataType == DataType.BYTES ? null : getHyperLogLogPlus(aggregationResultHolder); switch (storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); @@ -166,11 +147,16 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult hyperLogLogPlus.offer(stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - hyperLogLogPlus.offer(uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLogPlus(aggregationResultHolder, bytesValues[i]); + } + } else { + for (int i = 0; i < length; i++) { + hyperLogLogPlus.offer(bytesValues[i]); + } } break; default: @@ -235,11 +221,10 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; - // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. case BYTES: - byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { - for (byte[] value : uuidValuesArray[i]) { + for (byte[] value : bytesValuesArray[i]) { hyperLogLogPlus.offer(value); } } @@ -258,38 +243,18 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized HyperLogLogPlus - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - for (int i = 0; i < length; i++) { - HyperLogLogPlus value = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i]); - int groupKey = groupKeyArray[i]; - HyperLogLogPlus hyperLogLogPlus = groupByResultHolder.getResult(groupKey); - if (hyperLogLogPlus != null) { - hyperLogLogPlus.addAll(value); - } else { - groupByResultHolder.setValueForKey(groupKey, value); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); } else { aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType storedType) { + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -329,11 +294,16 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKeyArray[i], bytesValues[i]); + } + } else { + for (int i = 0; i < length; i++) { + getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); + } } break; default: @@ -401,12 +371,11 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; - // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. case BYTES: - byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]); - for (byte[] value : uuidValuesArray[i]) { + for (byte[] value : bytesValuesArray[i]) { hyperLogLogPlus.offer(value); } } @@ -425,41 +394,18 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized HyperLogLogPlus - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - for (int i = 0; i < length; i++) { - HyperLogLogPlus value = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i]); - for (int groupKey : groupKeysArray[i]) { - HyperLogLogPlus hyperLogLogPlus = groupByResultHolder.getResult(groupKey); - if (hyperLogLogPlus != null) { - hyperLogLogPlus.addAll(value); - } else { - // Create a new HyperLogLogPlus for the group - groupByResultHolder.setValueForKey(groupKey, - ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i])); - } - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); - } - return; - } - if (blockValSet.isSingleValue()) { - aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); } else { aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType storedType) { + BlockValSet blockValSet, DataType dataType, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -499,11 +445,18 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKey, bytesValues[i]); + } + } + } else { + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); + } } break; default: @@ -588,14 +541,13 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; - // Reached only by UUID: a real BYTES column is serialized HLL+ state and is handled above. case BYTES: - byte[][][] uuidValuesArray = blockValSet.getBytesValuesMV(); + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { - byte[][] uuidValues = uuidValuesArray[i]; + byte[][] bytesValues = bytesValuesArray[i]; for (int groupKey : groupKeysArray[i]) { HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(groupByResultHolder, groupKey); - for (byte[] value : uuidValues) { + for (byte[] value : bytesValues) { hyperLogLogPlus.offer(value); } } @@ -714,6 +666,27 @@ protected static RoaringBitmap getDictIdBitmap(AggregationResultHolder aggregati return dictIdsWrapper._dictIdBitmap; } + private void mergeSerializedHyperLogLogPlus(AggregationResultHolder aggregationResultHolder, byte[] bytes) { + HyperLogLogPlus value = deserializeHyperLogLogPlus(bytes); + HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); + aggregationResultHolder.setValue(hyperLogLogPlus == null ? value : merge(hyperLogLogPlus, value)); + } + + private void mergeSerializedHyperLogLogPlus(GroupByResultHolder groupByResultHolder, int groupKey, byte[] bytes) { + HyperLogLogPlus value = deserializeHyperLogLogPlus(bytes); + HyperLogLogPlus hyperLogLogPlus = groupByResultHolder.getResult(groupKey); + groupByResultHolder.setValueForKey(groupKey, + hyperLogLogPlus == null ? value : merge(hyperLogLogPlus, value)); + } + + private static HyperLogLogPlus deserializeHyperLogLogPlus(byte[] bytes) { + try { + return ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytes); + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); + } + } + /// Returns the HyperLogLogPlus from the result holder or creates a new one if it does not exist. protected HyperLogLogPlus getHyperLogLogPlus(AggregationResultHolder aggregationResultHolder) { HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); 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 408bcf943bb5..da483e6b8b8f 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 @@ -194,10 +194,9 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0. Logical BYTES values contain serialized sketches; all other logical types, - // including UUID, use the raw-value path and are handled by their stored-type switch branch. + // Main expression is always index 0 DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES) { + if (dataType != DataType.BYTES || !singleValues[0]) { List updateSketches = getUpdateSketches(aggregationResultHolder); if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { @@ -292,11 +291,11 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } break; case BYTES: - byte[][] uuidValues = (byte[][]) valueArrays[0]; + byte[][] bytesValues = (byte[][]) valueArrays[0]; if (_includeDefaultSketch) { UpdatableThetaSketch defaultSketch = updateSketches.get(0); for (int i = 0; i < length; i++) { - defaultSketch.update(uuidValues[i]); + defaultSketch.update(bytesValues[i]); } } for (int i = 0; i < numFilters; i++) { @@ -304,7 +303,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde UpdatableThetaSketch updateSketch = updateSketches.get(i + 1); for (int j = 0; j < length; j++) { if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { - updateSketch.update(uuidValues[j]); + updateSketch.update(bytesValues[j]); } } } @@ -427,11 +426,11 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } break; case BYTES: - byte[][][] uuidValues = (byte[][][]) valueArrays[0]; + byte[][][] bytesValues = (byte[][][]) valueArrays[0]; if (_includeDefaultSketch) { UpdatableThetaSketch defaultSketch = updateSketches.get(0); for (int i = 0; i < length; i++) { - for (byte[] value : uuidValues[i]) { + for (byte[] value : bytesValues[i]) { defaultSketch.update(value); } } @@ -441,7 +440,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde UpdatableThetaSketch updateSketch = updateSketches.get(i + 1); for (int j = 0; j < length; j++) { if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { - for (byte[] value : uuidValues[j]) { + for (byte[] value : bytesValues[j]) { updateSketch.update(value); } } @@ -454,7 +453,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde + valueTypes[0]); } } - } else if (singleValues[0]) { + } else { List thetaSketchAccumulators = getUnions(aggregationResultHolder); ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { @@ -472,28 +471,6 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } } - } else { - List thetaSketchAccumulators = getUnions(aggregationResultHolder); - ThetaSketch[][] sketches = deserializeSketches((byte[][][]) valueArrays[0], length); - if (_includeDefaultSketch) { - ThetaSketchAccumulator defaultThetaAccumulator = thetaSketchAccumulators.get(0); - for (int i = 0; i < length; i++) { - for (ThetaSketch sketch : sketches[i]) { - defaultThetaAccumulator.apply(sketch); - } - } - } - for (int i = 0; i < numFilters; i++) { - FilterEvaluator filterEvaluator = _filterEvaluators.get(i); - ThetaSketchAccumulator thetaSketchAccumulator = thetaSketchAccumulators.get(i + 1); - for (int j = 0; j < length; j++) { - if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { - for (ThetaSketch sketch : sketches[j]) { - thetaSketchAccumulator.apply(sketch); - } - } - } - } } } @@ -507,10 +484,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0. Logical BYTES values contain serialized sketches; all other logical types, - // including UUID, use the raw-value path and are handled by their stored-type switch branch. + // Main expression is always index 0 DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES) { + if (dataType != DataType.BYTES || !singleValues[0]) { if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { case INT: @@ -589,11 +565,11 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } break; case BYTES: - byte[][] uuidValues = (byte[][]) valueArrays[0]; + byte[][] bytesValues = (byte[][]) valueArrays[0]; for (int i = 0; i < length; i++) { List updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]); - byte[] value = uuidValues[i]; + byte[] value = bytesValues[i]; if (_includeDefaultSketch) { updateSketches.get(0).update(value); } @@ -717,11 +693,11 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } break; case BYTES: - byte[][][] uuidValues = (byte[][][]) valueArrays[0]; + byte[][][] bytesValues = (byte[][][]) valueArrays[0]; for (int i = 0; i < length; i++) { List updateSketches = getUpdateSketches(groupByResultHolder, groupKeyArray[i]); - byte[][] values = uuidValues[i]; + byte[][] values = bytesValues[i]; if (_includeDefaultSketch) { UpdatableThetaSketch defaultSketch = updateSketches.get(0); for (byte[] value : values) { @@ -744,7 +720,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol + valueTypes[0]); } } - } else if (singleValues[0]) { + } else { ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); for (int i = 0; i < length; i++) { List thetaSketchAccumulators = @@ -759,26 +735,6 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } } - } else { - ThetaSketch[][] sketches = deserializeSketches((byte[][][]) valueArrays[0], length); - for (int i = 0; i < length; i++) { - List thetaSketchAccumulators = - getUnions(groupByResultHolder, groupKeyArray[i]); - if (_includeDefaultSketch) { - ThetaSketchAccumulator defaultThetaAccumulator = thetaSketchAccumulators.get(0); - for (ThetaSketch sketch : sketches[i]) { - defaultThetaAccumulator.apply(sketch); - } - } - for (int j = 0; j < numFilters; j++) { - if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) { - ThetaSketchAccumulator thetaSketchAccumulator = thetaSketchAccumulators.get(j + 1); - for (ThetaSketch sketch : sketches[i]) { - thetaSketchAccumulator.apply(sketch); - } - } - } - } } } @@ -792,10 +748,9 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult extractValues(blockValSetMap, singleValues, valueTypes, valueArrays); int numFilters = _filterEvaluators.size(); - // Main expression is always index 0. Logical BYTES values contain serialized sketches; all other logical types, - // including UUID, use the raw-value path and are handled by their stored-type switch branch. + // Main expression is always index 0 DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES) { + if (dataType != DataType.BYTES || !singleValues[0]) { if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { case INT: @@ -899,11 +854,11 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } break; case BYTES: - byte[][] uuidValues = (byte[][]) valueArrays[0]; + byte[][] bytesValues = (byte[][]) valueArrays[0]; if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { for (int groupKey : groupKeysArray[i]) { - getUpdateSketches(groupByResultHolder, groupKey).get(0).update(uuidValues[i]); + getUpdateSketches(groupByResultHolder, groupKey).get(0).update(bytesValues[i]); } } } @@ -912,7 +867,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult for (int j = 0; j < length; j++) { if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { for (int groupKey : groupKeysArray[j]) { - getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(uuidValues[j]); + getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(bytesValues[j]); } } } @@ -1056,12 +1011,12 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } break; case BYTES: - byte[][][] uuidValues = (byte[][][]) valueArrays[0]; + byte[][][] bytesValues = (byte[][][]) valueArrays[0]; if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { for (int groupKey : groupKeysArray[i]) { UpdatableThetaSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0); - for (byte[] value : uuidValues[i]) { + for (byte[] value : bytesValues[i]) { defaultSketch.update(value); } } @@ -1074,7 +1029,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult for (int groupKey : groupKeysArray[j]) { UpdatableThetaSketch updateSketch = getUpdateSketches(groupByResultHolder, groupKey).get(i + 1); - for (byte[] value : uuidValues[j]) { + for (byte[] value : bytesValues[j]) { updateSketch.update(value); } } @@ -1088,7 +1043,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult + valueTypes[0]); } } - } else if (singleValues[0]) { + } else { ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { @@ -1107,31 +1062,6 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } } - } else { - ThetaSketch[][] sketches = deserializeSketches((byte[][][]) valueArrays[0], length); - if (_includeDefaultSketch) { - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - ThetaSketchAccumulator defaultThetaAccumulator = getUnions(groupByResultHolder, groupKey).get(0); - for (ThetaSketch sketch : sketches[i]) { - defaultThetaAccumulator.apply(sketch); - } - } - } - } - 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[j]) { - ThetaSketchAccumulator thetaSketchAccumulator = getUnions(groupByResultHolder, groupKey).get(i + 1); - for (ThetaSketch sketch : sketches[j]) { - thetaSketchAccumulator.apply(sketch); - } - } - } - } - } } } @@ -1558,14 +1488,6 @@ private ThetaSketch[] deserializeSketches(byte[][] serializedSketches, int lengt return sketches; } - private ThetaSketch[][] deserializeSketches(byte[][][] serializedSketches, int length) { - ThetaSketch[][] sketches = new ThetaSketch[length][]; - for (int i = 0; i < length; i++) { - sketches[i] = deserializeSketches(serializedSketches[i], serializedSketches[i].length); - } - return sketches; - } - private ThetaSketch deserializeSketch(byte[] serializedSketch) { return serializedSketch.length == 0 ? _emptySketch : ThetaSketch.wrap(MemorySegment.ofArray(serializedSketch).asReadOnly()); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java index 9dcf2f9e6f04..d6f174f86257 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -85,36 +85,16 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized UltraLogLog - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - UltraLogLog ull = aggregationResultHolder.getResult(); - if (ull == null) { - ull = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0]); - aggregationResultHolder.setValue(ull); - } else { - ull.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0])); - } - for (int i = 1; i < length; i++) { - ull.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging UltraLogLogs", e); - } - return; - } - // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; } // For non-dictionary-encoded expression, store values into the UltraLogLog - UltraLogLog ull = getULL(aggregationResultHolder); + UltraLogLog ull = dataType == DataType.BYTES ? null : getULL(aggregationResultHolder); switch (storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); @@ -146,11 +126,27 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde UltraLogLogUtils.hashObject(stringValues[i]).ifPresent(ull::add); } break; - // Reached only by UUID: a real BYTES column is serialized ULL state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - UltraLogLogUtils.hashObject(uuidValues[i]).ifPresent(ull::add); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + try { + UltraLogLog serializedUll = aggregationResultHolder.getResult(); + if (serializedUll == null) { + serializedUll = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0]); + aggregationResultHolder.setValue(serializedUll); + } else { + serializedUll.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0])); + } + for (int i = 1; i < length; i++) { + serializedUll.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging UltraLogLogs", e); + } + } else { + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(bytesValues[i]).ifPresent(ull::add); + } } break; default: @@ -167,29 +163,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized UltraLogLogs - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - for (int i = 0; i < length; i++) { - UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); - int groupKey = groupKeyArray[i]; - UltraLogLog ull = groupByResultHolder.getResult(groupKey); - if (ull != null) { - ull.add(value); - } else { - groupByResultHolder.setValueForKey(groupKey, value); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging UltraLogLog", e); - } - return; - } - // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -234,12 +210,28 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); } break; - // Reached only by UUID: a real BYTES column is serialized ULL state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - UltraLogLogUtils.hashObject(uuidValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + try { + for (int i = 0; i < length; i++) { + UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); + int groupKey = groupKeyArray[i]; + UltraLogLog ull = groupByResultHolder.getResult(groupKey); + if (ull != null) { + ull.add(value); + } else { + groupByResultHolder.setValueForKey(groupKey, value); + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging UltraLogLog", e); + } + } else { + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(bytesValues[i]) + .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + } } break; default: @@ -256,32 +248,9 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); DataType storedType = dataType.getStoredType(); - // Treat BYTES value as serialized UltraLogLogs - if (storedType == DataType.BYTES && dataType != DataType.UUID) { - byte[][] bytesValues = blockValSet.getBytesValuesSV(); - try { - for (int i = 0; i < length; i++) { - UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); - for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = groupByResultHolder.getResult(groupKey); - if (ull != null) { - ull.add(value); - } else { - // Create a new HyperLogLogPlus for the group - groupByResultHolder.setValueForKey(groupKey, - ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); - } - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging UltraLogLog", e); - } - return; - } - // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null) { + if (dictionary != null && dataType != DataType.BYTES) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -321,11 +290,29 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; - // Reached only by UUID: a real BYTES column is serialized ULL state and is handled above. case BYTES: - byte[][] uuidValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], uuidValues[i]); + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + if (dataType == DataType.BYTES) { + try { + for (int i = 0; i < length; i++) { + UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + UltraLogLog ull = groupByResultHolder.getResult(groupKey); + if (ull != null) { + ull.add(value); + } else { + groupByResultHolder.setValueForKey(groupKey, + ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); + } + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging UltraLogLog", e); + } + } else { + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); + } } break; default: diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java index 29041933af2f..788098efa3af 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java @@ -18,21 +18,34 @@ */ package org.apache.pinot.core.query.aggregation.function; +import com.clearspring.analytics.stream.cardinality.HyperLogLog; +import java.io.IOException; +import java.util.Arrays; +import java.util.List; +import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.common.ObjectSerDeUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.reader.Dictionary; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.testng.annotations.Test; +import static org.mockito.Mockito.anyInt; import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertThrows; +import static org.testng.Assert.assertTrue; /// Unit test for {@link AggregationFunctionUtils#getAggregationResult}, the metadata/dictionary based aggregation /// result resolver used by the non-scan based and partial metadata based aggregation paths. @SuppressWarnings("rawtypes") public class AggregationFunctionUtilsTest { + private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); private static AggregationFunction mockFunction(AggregationFunctionType type) { AggregationFunction aggregationFunction = mock(AggregationFunction.class); @@ -81,4 +94,75 @@ public void testNonCountWithNullDataSourceThrows() { () -> AggregationFunctionUtils.getAggregationResult(mockFunction(AggregationFunctionType.MIN), null, 100, "TEST")); } + + @Test + public void testDistinctCountHllMvOffersDictionaryBytesAsRawValues() + throws IOException { + DistinctCountHLLMVAggregationFunction function = + new DistinctCountHLLMVAggregationFunction(List.of(BYTES_EXPRESSION)); + byte[][] bytesValues = {{1}, {2}, {3}}; + Dictionary dictionary = mock(Dictionary.class); + when(dictionary.length()).thenReturn(bytesValues.length); + for (int i = 0; i < bytesValues.length; i++) { + when(dictionary.get(i)).thenReturn(bytesValues[i]); + } + DataSource dataSource = mockBytesDataSource(dictionary, false); + + HyperLogLog result = (HyperLogLog) AggregationFunctionUtils.getAggregationResult(function, dataSource, 3, "TEST"); + + HyperLogLog expected = new HyperLogLog(function.getLog2m()); + for (byte[] value : bytesValues) { + expected.offer(value); + } + assertEquals(result.cardinality(), 3L); + assertTrue(Arrays.equals(result.getBytes(), expected.getBytes())); + for (int i = 0; i < bytesValues.length; i++) { + verify(dictionary).get(i); + } + verify(dictionary, never()).getBytesValue(anyInt()); + } + + @Test + public void testDistinctCountHllMergesSingleValueSerializedBytes() + throws IOException { + DistinctCountHLLAggregationFunction function = + new DistinctCountHLLAggregationFunction(List.of(BYTES_EXPRESSION)); + byte[] firstSketch = serializedHll(function, "a", "b"); + byte[] secondSketch = serializedHll(function, "b", "c"); + Dictionary dictionary = mock(Dictionary.class); + when(dictionary.length()).thenReturn(2); + when(dictionary.getBytesValue(0)).thenReturn(firstSketch); + when(dictionary.getBytesValue(1)).thenReturn(secondSketch); + DataSource dataSource = mockBytesDataSource(dictionary, true); + + HyperLogLog result = (HyperLogLog) AggregationFunctionUtils.getAggregationResult(function, dataSource, 2, "TEST"); + + HyperLogLog expected = new HyperLogLog(function.getLog2m()); + expected.offer("a"); + expected.offer("b"); + expected.offer("c"); + assertEquals(result.cardinality(), 3L); + assertTrue(Arrays.equals(result.getBytes(), expected.getBytes())); + verify(dictionary).getBytesValue(0); + verify(dictionary).getBytesValue(1); + verify(dictionary, never()).get(anyInt()); + } + + private static DataSource mockBytesDataSource(Dictionary dictionary, boolean singleValue) { + DataSourceMetadata metadata = mock(DataSourceMetadata.class); + when(metadata.getDataType()).thenReturn(DataType.BYTES); + when(metadata.isSingleValue()).thenReturn(singleValue); + DataSource dataSource = mock(DataSource.class); + when(dataSource.getDictionary()).thenReturn(dictionary); + when(dataSource.getDataSourceMetadata()).thenReturn(metadata); + return dataSource; + } + + private static byte[] serializedHll(DistinctCountHLLAggregationFunction function, String... values) { + HyperLogLog hll = new HyperLogLog(function.getLog2m()); + for (String value : values) { + hll.offer(value); + } + return ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.serialize(hll); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java index ec4fccf785f5..1105d53bda10 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java @@ -18,13 +18,28 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.request.context.ExpressionContext; +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.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; +import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + public class DistinctCountBitmapMVAggregationFunctionTest extends AbstractAggregationFunctionTest { + private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); @Test public void testAggregationMV() { @@ -92,4 +107,50 @@ public void testAggregationMVGroupByMV() { "tag1 | 3", // distinct: 1, 2, 3 "tag2 | 3"); // distinct: 1, 2, 3 } + + @Test + public void testMultiValueBytesUseMultiValueAccessor() { + byte[][][] bytesValues = { + {{1}, {2}}, + {{2}, {3}}, + {{1}}, + {{4}} + }; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); + DistinctCountBitmapAggregationFunction function = + new DistinctCountBitmapAggregationFunction(List.of(BYTES_EXPRESSION)); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4); + + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + + private static int extractFinalResult(DistinctCountBitmapAggregationFunction function, + AggregationResultHolder resultHolder) { + return function.extractFinalResult(function.extractAggregationResult(resultHolder)); + } + + private static int extractFinalResult(DistinctCountBitmapAggregationFunction function, + GroupByResultHolder resultHolder, int groupKey) { + return function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey)); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java index 80faf2555253..417c9f9b17ce 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java @@ -48,12 +48,21 @@ public class DistinctCountCPCSketchAggregationFunctionTest { private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); + private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); private static final byte[][] UUID_VALUES_SV = {UUID_0, UUID_1, UUID_0, UUID_2}; private static final byte[][][] UUID_VALUES_MV = {{UUID_0, UUID_1}, {UUID_1, UUID_2}, {UUID_0}, {UUID_3}}; + // These payloads are intentionally too short to be serialized CPC sketches. MV BYTES are raw values. + private static final byte[] RAW_BYTES_0 = {1}; + private static final byte[] RAW_BYTES_1 = {2}; + private static final byte[] RAW_BYTES_2 = {3}; + private static final byte[] RAW_BYTES_3 = {4}; + private static final byte[][][] RAW_BYTES_VALUES_MV = { + {RAW_BYTES_0, RAW_BYTES_1}, {RAW_BYTES_1, RAW_BYTES_2}, {RAW_BYTES_0}, {RAW_BYTES_3} + }; @DataProvider(name = "uuidValueModes") public static Object[][] uuidValueModes() { @@ -229,40 +238,96 @@ public void testAggregateDictionaryEncodedUuid() { } @Test - public void testAggregateMultiValueSerializedSketches() { - byte[][][] serializedSketches = { - {serializedSketch("a"), serializedSketch("b")}, - {serializedSketch("b"), serializedSketch("c")}, - {new byte[0]}, - {serializedSketch("d")} + public void testAggregateSingleValueSerializedSketches() { + byte[][] serializedSketches = { + serializedSketch("a", "b"), + serializedSketch("b", "c"), + new byte[0], + serializedSketch("d") }; BlockValSet blockValSet = mock(BlockValSet.class); when(blockValSet.getValueType()).thenReturn(DataType.BYTES); - when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.isSingleValue()).thenReturn(true); when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(serializedSketches); + when(blockValSet.getBytesValuesSV()).thenReturn(serializedSketches); DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); + new DistinctCountCPCSketchAggregationFunction(List.of(BYTES_EXPRESSION)); AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); + Map.of(BYTES_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 1L); GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); + verify(blockValSet, atLeastOnce()).getBytesValuesSV(); + verify(blockValSet, never()).getBytesValuesMV(); + } + + @Test + public void testAggregateMultiValueRawBytes() { + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + when(blockValSet.getBytesValuesMV()).thenReturn(RAW_BYTES_VALUES_MV); + DistinctCountCPCSketchAggregationFunction function = + new DistinctCountCPCSketchAggregationFunction(List.of(BYTES_EXPRESSION)); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(RAW_BYTES_VALUES_MV.length, aggregationResultHolder, + Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(RAW_BYTES_VALUES_MV.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(RAW_BYTES_VALUES_MV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); verify(blockValSet, atLeastOnce()).getBytesValuesMV(); verify(blockValSet, never()).getBytesValuesSV(); } + @Test + public void testAggregateDictionaryEncodedMultiValueRawBytes() { + Dictionary dictionary = mock(Dictionary.class); + when(dictionary.get(0)).thenReturn(RAW_BYTES_0); + when(dictionary.get(1)).thenReturn(RAW_BYTES_1); + when(dictionary.get(2)).thenReturn(RAW_BYTES_2); + when(dictionary.get(3)).thenReturn(RAW_BYTES_3); + + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.isDictionaryEncoded()).thenReturn(true); + when(blockValSet.getDictionary()).thenReturn(dictionary); + when(blockValSet.getDictionaryIdsMV()).thenReturn(new int[][]{{0, 1}, {1, 2}, {0}, {3}}); + DistinctCountCPCSketchAggregationFunction function = + new DistinctCountCPCSketchAggregationFunction(List.of(BYTES_EXPRESSION)); + + AggregationResultHolder resultHolder = function.createAggregationResultHolder(); + function.aggregate(RAW_BYTES_VALUES_MV.length, resultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + + Assert.assertEquals(extractFinalResult(function, resultHolder), 4L); + verify(blockValSet, atLeastOnce()).getDictionaryIdsMV(); + verify(blockValSet, never()).getBytesValuesSV(); + verify(blockValSet, never()).getBytesValuesMV(); + } + private static BlockValSet mockUuidBlockValSet(boolean singleValue) { BlockValSet blockValSet = mock(BlockValSet.class); when(blockValSet.getValueType()).thenReturn(DataType.UUID); @@ -286,9 +351,11 @@ private static void verifyBytesAccessor(BlockValSet blockValSet, boolean singleV } } - private static byte[] serializedSketch(String value) { + private static byte[] serializedSketch(String... values) { CpcSketch sketch = new CpcSketch(); - sketch.update(value); + for (String value : values) { + sketch.update(value); + } return sketch.toByteArray(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java index 3049fae37358..8dc9ddc0699c 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java @@ -64,10 +64,10 @@ public void testCanUseStarTreeDefaultLog2m() { } /// Regression: UUID columns have storedType=BYTES, but a UUID value is a logical scalar, not a serialized - /// HyperLogLog. The aggregator must route UUID columns through the same content-hash path as STRING (offering - /// canonical UUID strings) instead of trying to deserialize each 16-byte value as an HLL. + /// HyperLogLog. The aggregator must offer the stored bytes as values instead of trying to deserialize each + /// 16-byte value as an HLL. @Test - public void testAggregateOnUuidColumnOffersCanonicalStringsAndProducesExactDistinctCount() { + public void testAggregateOnUuidColumnOffersStoredBytesAndProducesExactDistinctCount() { ExpressionContext expression = RequestContextUtils.getExpression("uuidCol"); DistinctCountHLLAggregationFunction function = new DistinctCountHLLAggregationFunction(List.of(expression)); @@ -81,8 +81,7 @@ public void testAggregateOnUuidColumnOffersCanonicalStringsAndProducesExactDisti "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12" }; - // Stub the BYTES fetch (raw 16-byte values) — the production path converts bytes to canonical strings - // itself because ProjectionBlockValSet.getStringValuesSV() would render stored BYTES as bare hex. + // Stub the BYTES fetch with the raw 16-byte stored values offered by the production path. byte[][] uuidBytes = new byte[uuidStrings.length][]; for (int i = 0; i < uuidStrings.length; i++) { uuidBytes[i] = UuidUtils.toBytes(uuidStrings[i]); @@ -144,7 +143,7 @@ private long computeHllCardinality(String[] values, DataType valueType) { BlockValSet blockValSet = mock(BlockValSet.class); when(blockValSet.getValueType()).thenReturn(valueType); if (valueType == DataType.UUID) { - // UUID path fetches raw bytes and converts to canonical form itself (projection string fetch returns hex) + // UUID path fetches and offers the raw stored bytes. byte[][] uuidBytes = new byte[values.length][]; for (int i = 0; i < values.length; i++) { uuidBytes[i] = UuidUtils.toBytes(values[i]); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java index fe24e95b946e..4a2c5eac5571 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java @@ -18,13 +18,28 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.request.context.ExpressionContext; +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.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; +import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + public class DistinctCountHLLMVAggregationFunctionTest extends AbstractAggregationFunctionTest { + private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); @Test public void testAggregationMV() { @@ -92,4 +107,50 @@ public void testAggregationMVGroupByMV() { "tag1 | 3", // distinct: 1, 2, 3 "tag2 | 3"); // distinct: 1, 2, 3 } + + @Test + public void testMultiValueBytesUseMultiValueAccessor() { + DistinctCountHLLAggregationFunction function = + new DistinctCountHLLAggregationFunction(List.of(BYTES_EXPRESSION)); + byte[][][] bytesValues = { + {{1}, {2}}, + {{2}, {3}}, + {{1}}, + {{4}} + }; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); + + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + + private static long extractFinalResult(DistinctCountHLLAggregationFunction function, + AggregationResultHolder resultHolder) { + return function.extractFinalResult(function.extractAggregationResult(resultHolder)); + } + + private static long extractFinalResult(DistinctCountHLLAggregationFunction function, + GroupByResultHolder resultHolder, int groupKey) { + return function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey)); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java index 9a2ca5f86e79..f680143bbc6d 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java @@ -18,13 +18,28 @@ */ package org.apache.pinot.core.query.aggregation.function; +import java.util.List; +import java.util.Map; +import org.apache.pinot.common.request.context.ExpressionContext; +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.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; +import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + public class DistinctCountHLLPlusMVAggregationFunctionTest extends AbstractAggregationFunctionTest { + private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); @Test public void testAggregationMV() { @@ -92,4 +107,50 @@ public void testAggregationMVGroupByMV() { "tag1 | 3", // distinct: 1, 2, 3 "tag2 | 3"); // distinct: 1, 2, 3 } + + @Test + public void testMultiValueBytesUseMultiValueAccessor() { + DistinctCountHLLPlusAggregationFunction function = + new DistinctCountHLLPlusAggregationFunction(List.of(BYTES_EXPRESSION)); + byte[][][] bytesValues = { + {{1}, {2}}, + {{2}, {3}}, + {{1}}, + {{4}} + }; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(false); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); + + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + + private static long extractFinalResult(DistinctCountHLLPlusAggregationFunction function, + AggregationResultHolder resultHolder) { + return function.extractFinalResult(function.extractAggregationResult(resultHolder)); + } + + private static long extractFinalResult(DistinctCountHLLPlusAggregationFunction function, + GroupByResultHolder resultHolder, int groupKey) { + return function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey)); + } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java index f26b69caca78..906e7fafb897 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java @@ -43,6 +43,7 @@ public class DistinctCountThetaSketchAggregationFunctionTest { private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); + private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); @@ -139,37 +140,72 @@ public void testUuidPredicateUsesLogicalType() { } @Test - public void testAggregateMultiValueSerializedSketches() { - byte[][][] serializedSketches = { - {serializedSketch("a"), serializedSketch("b")}, - {serializedSketch("b"), serializedSketch("c")}, - {new byte[0]}, - {serializedSketch("d")} + public void testAggregateMultiValueBytesAsRawValues() { + // These one-byte values are not serialized sketches, so the test also verifies that MV BYTES are not deserialized. + byte[][][] bytesValues = { + {{1}, {2}}, + {{2}, {3}}, + {{1}}, + {{4}} }; BlockValSet blockValSet = mock(BlockValSet.class); when(blockValSet.getValueType()).thenReturn(DataType.BYTES); when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(serializedSketches); + when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); DistinctCountThetaSketchAggregationFunction function = - new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); + new DistinctCountThetaSketchAggregationFunction(List.of(BYTES_EXPRESSION)); AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, groupByMVResultHolder, + Map.of(BYTES_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); + verify(blockValSet, atLeastOnce()).getBytesValuesMV(); + verify(blockValSet, never()).getBytesValuesSV(); + } + + @Test + public void testAggregateSingleValueSerializedSketches() { + byte[][] serializedSketches = { + serializedSketch("a", "b"), + serializedSketch("b", "c"), + new byte[0], + serializedSketch("d") + }; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isSingleValue()).thenReturn(true); + when(blockValSet.getBytesValuesSV()).thenReturn(serializedSketches); + DistinctCountThetaSketchAggregationFunction function = + new DistinctCountThetaSketchAggregationFunction(List.of(BYTES_EXPRESSION)); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); + Map.of(BYTES_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 1L); GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); + groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); + verify(blockValSet, atLeastOnce()).getBytesValuesSV(); + verify(blockValSet, never()).getBytesValuesMV(); } private static BlockValSet mockUuidBlockValSet(boolean singleValue) { @@ -194,9 +230,11 @@ private static void verifyBytesAccessor(BlockValSet blockValSet, boolean singleV } } - private static byte[] serializedSketch(String value) { + private static byte[] serializedSketch(String... values) { UpdatableThetaSketch sketch = new UpdatableThetaSketchBuilder().build(); - sketch.update(value); + for (String value : values) { + sketch.update(value); + } return sketch.compact().toByteArray(); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java index 06c05b366bdc..e6a6fb0e5717 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java @@ -18,16 +18,34 @@ */ package org.apache.pinot.core.query.aggregation.function; +import com.dynatrace.hash4j.distinctcount.UltraLogLog; import java.util.List; import java.util.Map; import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; +import org.apache.pinot.core.common.BlockValSet; +import org.apache.pinot.core.common.ObjectSerDeUtils; +import org.apache.pinot.core.query.aggregation.AggregationResultHolder; +import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; +import org.apache.pinot.segment.local.utils.UltraLogLogUtils; import org.apache.pinot.segment.spi.Constants; +import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; +import static org.mockito.Mockito.atLeastOnce; +import static org.mockito.Mockito.mock; +import static org.mockito.Mockito.never; +import static org.mockito.Mockito.verify; +import static org.mockito.Mockito.when; + public class DistinctCountULLAggregationFunctionTest { + private static final ExpressionContext INPUT_EXPRESSION = ExpressionContext.forIdentifier("inputCol"); + private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); + private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); + private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); @Test public void testCanUseStarTreeDefaultP() { @@ -58,4 +76,98 @@ public void testCanUseStarTreeCustomP() { Assert.assertTrue(function.canUseStarTree(Map.of(Constants.HLLPLUS_ULL_P_KEY, 16))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.HLLPLUS_ULL_P_KEY, "16"))); } + + @Test + public void testRawUuidUsesStoredBytes() { + DistinctCountULLAggregationFunction function = + new DistinctCountULLAggregationFunction(List.of(INPUT_EXPRESSION)); + byte[][] uuidValues = {UUID_0, UUID_1, UUID_0, UUID_2}; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + when(blockValSet.getBytesValuesSV()).thenReturn(uuidValues); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(INPUT_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 3L); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(INPUT_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 2L); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, groupByMVResultHolder, + Map.of(INPUT_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 1L); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); + + verify(blockValSet, atLeastOnce()).getBytesValuesSV(); + verify(blockValSet, never()).getBytesValuesMV(); + } + + @Test + public void testSerializedBytesUsesStoredSketches() { + DistinctCountULLAggregationFunction function = + new DistinctCountULLAggregationFunction(List.of(INPUT_EXPRESSION)); + byte[][] serializedSketches = { + serializedSketch(function, "a", "b"), + serializedSketch(function, "b", "c"), + serializedSketch(function, "a"), + serializedSketch(function, "d") + }; + BlockValSet blockValSet = mock(BlockValSet.class); + when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.isDictionaryEncoded()).thenReturn(false); + when(blockValSet.getBytesValuesSV()).thenReturn(serializedSketches); + + AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); + function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(INPUT_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), + referenceEstimate(function, "a", "b", "c", "d")); + + GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(INPUT_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), + referenceEstimate(function, "a", "b", "c")); + Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), referenceEstimate(function, "a", "d")); + + GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); + function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(INPUT_EXPRESSION, blockValSet)); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), referenceEstimate(function, "a", "b")); + Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), + referenceEstimate(function, "a", "b", "c", "d")); + + verify(blockValSet, atLeastOnce()).getBytesValuesSV(); + verify(blockValSet, never()).getBytesValuesMV(); + } + + private static byte[] serializedSketch(DistinctCountULLAggregationFunction function, String... values) { + UltraLogLog sketch = UltraLogLog.create(function.getP()); + for (String value : values) { + UltraLogLogUtils.hashObject(value).ifPresent(sketch::add); + } + return ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.serialize(sketch); + } + + private static long referenceEstimate(DistinctCountULLAggregationFunction function, String... values) { + UltraLogLog reference = UltraLogLog.create(function.getP()); + for (String value : values) { + UltraLogLogUtils.hashObject(value).ifPresent(reference::add); + } + return Math.round(reference.getDistinctCountEstimate()); + } + + private static long extractFinalResult(DistinctCountULLAggregationFunction function, + AggregationResultHolder resultHolder) { + return ((Number) function.extractFinalResult(function.extractAggregationResult(resultHolder))).longValue(); + } + + private static long extractFinalResult(DistinctCountULLAggregationFunction function, + GroupByResultHolder resultHolder, int groupKey) { + return ((Number) function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey))).longValue(); + } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java index b0101f6bb43d..f52623e48a0b 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java @@ -194,8 +194,8 @@ public void testDistinctOnUuidColumn() public void testDistinctCountOnUuidColumn() throws Exception { setUseMultiStageQueryEngine(false); - for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTBITMAP", - "DISTINCTCOUNTTHETASKETCH", "DISTINCTCOUNTCPCSKETCH")) { + for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTHLLPLUS", "DISTINCTCOUNTULL", + "DISTINCTCOUNTBITMAP", "DISTINCTCOUNTTHETASKETCH", "DISTINCTCOUNTCPCSKETCH")) { JsonNode rows = query(String.format("SELECT %s(%s) FROM %s", function, UUID_COLUMN, getTableName())); assertEquals(rows.get(0).get(0).asLong(), NUM_DISTINCT, function + ": " + rows.toPrettyString()); } From 677b02f652a747dd02966da45e3cb6f8c375a592 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 14 Aug 2026 16:03:14 -0700 Subject: [PATCH 15/19] Clarify serialized aggregation BYTES handling --- .../aggregation/function/AggregationFunctionUtils.java | 6 +++--- .../function/DistinctCountBitmapAggregationFunction.java | 3 +++ .../function/DistinctCountCPCSketchAggregationFunction.java | 3 +++ .../function/DistinctCountHLLAggregationFunction.java | 3 +++ .../function/DistinctCountHLLPlusAggregationFunction.java | 3 +++ .../DistinctCountThetaSketchAggregationFunction.java | 3 +++ .../function/DistinctCountULLAggregationFunction.java | 3 +++ 7 files changed, 21 insertions(+), 3 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index 8cb23df58471..e650897c8dda 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -804,7 +804,7 @@ private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES && dataSource.getDataSourceMetadata().isSingleValue()) { - // Treat BYTES value as serialized HyperLogLog + // SV logical BYTES dictionary values contain serialized HyperLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLog hll = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -827,7 +827,7 @@ private static HyperLogLogPlus getDistinctCountHLLPlusResult(DataSource dataSour Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES && dataSource.getDataSourceMetadata().isSingleValue()) { - // Treat BYTES value as serialized HyperLogLogPlus + // SV logical BYTES dictionary values contain serialized HyperLogLogPlus objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLogPlus hllplus = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -870,7 +870,7 @@ private static UltraLogLog getDistinctCountULLResult(DataSource dataSource, Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES && dataSource.getDataSourceMetadata().isSingleValue()) { - // Treat BYTES value as serialized UltraLogLog and merge + // SV logical BYTES dictionary values contain serialized UltraLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); UltraLogLog ull = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(dictionary.getBytesValue(0)); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index 55e9feb1420a..badf4e22e82f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -126,6 +126,7 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized RoaringBitmap objects. for (int i = 0; i < length; i++) { valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); } @@ -271,6 +272,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized RoaringBitmap objects. for (int i = 0; i < length; i++) { getValueBitmap(groupByResultHolder, groupKeyArray[i]) .or(RoaringBitmapUtils.deserialize(bytesValues[i])); @@ -420,6 +422,7 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized RoaringBitmap objects. for (int i = 0; i < length; i++) { RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); for (int groupKey : groupKeysArray[i]) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index 5407ccfe3e36..59d6dcd9a70f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -194,6 +194,7 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized CpcSketch objects. mergeSerializedSketches(aggregationResultHolder, bytesValues, length); } else { CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); @@ -342,6 +343,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized CpcSketch objects. CpcSketch[] sketches = deserializeSketches(bytesValues, length); for (int i = 0; i < length; i++) { CpcSketch sketch = sketches[i]; @@ -504,6 +506,7 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized CpcSketch objects. CpcSketch[] sketches = deserializeSketches(bytesValues, length); for (int i = 0; i < length; i++) { CpcSketch sketch = sketches[i]; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java index 98d767b77367..1c190d9d9bd5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -142,6 +142,7 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized HyperLogLog objects. for (int i = 0; i < length; i++) { mergeSerializedHyperLogLog(aggregationResultHolder, bytesValues[i]); } @@ -292,6 +293,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized HyperLogLog objects. for (int i = 0; i < length; i++) { mergeSerializedHyperLogLog(groupByResultHolder, groupKeyArray[i], bytesValues[i]); } @@ -444,6 +446,7 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized HyperLogLog objects. for (int i = 0; i < length; i++) { for (int groupKey : groupKeysArray[i]) { mergeSerializedHyperLogLog(groupByResultHolder, groupKey, bytesValues[i]); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java index 8570ae195812..da0f4163048e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -150,6 +150,7 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized HyperLogLogPlus objects. for (int i = 0; i < length; i++) { mergeSerializedHyperLogLogPlus(aggregationResultHolder, bytesValues[i]); } @@ -297,6 +298,7 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized HyperLogLogPlus objects. for (int i = 0; i < length; i++) { mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKeyArray[i], bytesValues[i]); } @@ -448,6 +450,7 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized HyperLogLogPlus objects. for (int i = 0; i < length; i++) { for (int groupKey : groupKeysArray[i]) { mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKey, bytesValues[i]); 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 da483e6b8b8f..1dbb76df7ffc 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 @@ -454,6 +454,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } } else { + // SV logical BYTES values contain serialized ThetaSketch objects. List thetaSketchAccumulators = getUnions(aggregationResultHolder); ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { @@ -721,6 +722,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } } else { + // SV logical BYTES values contain serialized ThetaSketch objects. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); for (int i = 0; i < length; i++) { List thetaSketchAccumulators = @@ -1044,6 +1046,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } } else { + // SV logical BYTES values contain serialized ThetaSketch objects. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java index d6f174f86257..96d6ff825ea6 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -129,6 +129,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized UltraLogLog objects. try { UltraLogLog serializedUll = aggregationResultHolder.getResult(); if (serializedUll == null) { @@ -213,6 +214,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized UltraLogLog objects. try { for (int i = 0; i < length; i++) { UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); @@ -293,6 +295,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); if (dataType == DataType.BYTES) { + // SV logical BYTES values contain serialized UltraLogLog objects. try { for (int i = 0; i < length; i++) { UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); From 50ff5eb093d95cfde18e6638f6ddaaa905cb6484 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 14 Aug 2026 16:36:34 -0700 Subject: [PATCH 16/19] Restore legacy serialized BYTES dispatch --- .../function/AggregationFunctionUtils.java | 19 +-- ...istinctCountBitmapAggregationFunction.java | 89 ++++++----- ...inctCountCPCSketchAggregationFunction.java | 98 ++++++------ .../DistinctCountHLLAggregationFunction.java | 87 ++++++----- ...stinctCountHLLPlusAggregationFunction.java | 88 ++++++----- ...ctCountThetaSketchAggregationFunction.java | 15 +- .../DistinctCountULLAggregationFunction.java | 143 ++++++++++-------- .../AggregationFunctionUtilsTest.java | 26 ++-- ...tCountBitmapMVAggregationFunctionTest.java | 35 +++-- ...CountCPCSketchAggregationFunctionTest.java | 66 +++----- ...inctCountHLLMVAggregationFunctionTest.java | 35 +++-- ...CountHLLPlusMVAggregationFunctionTest.java | 35 +++-- ...untThetaSketchAggregationFunctionTest.java | 35 ----- 13 files changed, 387 insertions(+), 384 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index e650897c8dda..7c7ac22b2088 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -69,6 +69,7 @@ import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.SegmentContext; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; import org.apache.pinot.spi.data.FieldSpec; @@ -802,9 +803,9 @@ private static HyperLogLogPlus getDistinctValueHLLPlus(Dictionary dictionary, in private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, DistinctCountHLLAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES - && dataSource.getDataSourceMetadata().isSingleValue()) { - // SV logical BYTES dictionary values contain serialized HyperLogLog objects. + DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); + if (metadata.getDataType() == FieldSpec.DataType.BYTES) { + // Logical BYTES dictionary values contain serialized HyperLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLog hll = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -825,9 +826,9 @@ private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, private static HyperLogLogPlus getDistinctCountHLLPlusResult(DataSource dataSource, DistinctCountHLLPlusAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES - && dataSource.getDataSourceMetadata().isSingleValue()) { - // SV logical BYTES dictionary values contain serialized HyperLogLogPlus objects. + DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); + if (metadata.getDataType() == FieldSpec.DataType.BYTES) { + // Logical BYTES dictionary values contain serialized HyperLogLogPlus objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLogPlus hllplus = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -868,9 +869,9 @@ private static Object getDistinctCountSmartHLLPlusResult(Dictionary dictionary, private static UltraLogLog getDistinctCountULLResult(DataSource dataSource, DistinctCountULLAggregationFunction function, String explainPlanName) { Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); - if (dataSource.getDataSourceMetadata().getDataType() == FieldSpec.DataType.BYTES - && dataSource.getDataSourceMetadata().isSingleValue()) { - // SV logical BYTES dictionary values contain serialized UltraLogLog objects. + DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); + if (metadata.getDataType() == FieldSpec.DataType.BYTES) { + // Logical BYTES dictionary values contain serialized UltraLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); UltraLogLog ull = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(dictionary.getBytesValue(0)); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index badf4e22e82f..593ce48fc37a 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -73,20 +73,31 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized RoaringBitmap objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + RoaringBitmap valueBitmap = getValueBitmap(aggregationResultHolder); + for (int i = 0; i < length; i++) { + valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); + aggregateSV(length, aggregationResultHolder, blockValSet, storedType); } else { aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType dataType, DataType storedType) { + DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; @@ -125,15 +136,8 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized RoaringBitmap objects. - for (int i = 0; i < length; i++) { - valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); - } - } else { - for (int i = 0; i < length; i++) { - valueBitmap.add(Arrays.hashCode(bytesValues[i])); - } + for (int i = 0; i < length; i++) { + valueBitmap.add(Arrays.hashCode(bytesValues[i])); } break; default: @@ -216,20 +220,30 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized RoaringBitmap objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getValueBitmap(groupByResultHolder, groupKeyArray[i]).or(RoaringBitmapUtils.deserialize(bytesValues[i])); + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -271,16 +285,8 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized RoaringBitmap objects. - for (int i = 0; i < length; i++) { - getValueBitmap(groupByResultHolder, groupKeyArray[i]) - .or(RoaringBitmapUtils.deserialize(bytesValues[i])); - } - } else { - for (int i = 0; i < length; i++) { - getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(Arrays.hashCode(bytesValues[i])); - } + for (int i = 0; i < length; i++) { + getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(Arrays.hashCode(bytesValues[i])); } break; default: @@ -366,20 +372,33 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized RoaringBitmap objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + getValueBitmap(groupByResultHolder, groupKey).or(value); + } + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -421,18 +440,8 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized RoaringBitmap objects. - for (int i = 0; i < length; i++) { - RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); - for (int groupKey : groupKeysArray[i]) { - getValueBitmap(groupByResultHolder, groupKey).or(value); - } - } - } else { - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], Arrays.hashCode(bytesValues[i])); - } + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], Arrays.hashCode(bytesValues[i])); } break; default: diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index 59d6dcd9a70f..96857e84e03c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -136,19 +136,26 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized CpcSketch objects. + // They always use the single-value representation. + mergeSerializedSketches(aggregationResultHolder, blockValSet.getBytesValuesSV(), length); + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); + aggregateSV(length, aggregationResultHolder, blockValSet, storedType); } else { aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType dataType, DataType storedType) { + DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; @@ -193,14 +200,9 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized CpcSketch objects. - mergeSerializedSketches(aggregationResultHolder, bytesValues, length); - } else { - CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); - for (int i = 0; i < length; i++) { - cpcSketch.update(bytesValues[i]); - } + CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); + for (int i = 0; i < length; i++) { + cpcSketch.update(bytesValues[i]); } break; default: @@ -288,19 +290,32 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized CpcSketch objects. + // They always use the single-value representation. + CpcSketch[] sketches = deserializeSketches(blockValSet.getBytesValuesSV(), length); + for (int i = 0; i < length; i++) { + CpcSketch sketch = sketches[i]; + if (sketch != null) { + getAccumulator(groupByResultHolder, groupKeyArray[i]).apply(sketch); + } + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -342,19 +357,8 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized CpcSketch objects. - CpcSketch[] sketches = deserializeSketches(bytesValues, length); - for (int i = 0; i < length; i++) { - CpcSketch sketch = sketches[i]; - if (sketch != null) { - getAccumulator(groupByResultHolder, groupKeyArray[i]).apply(sketch); - } - } - } else { - for (int i = 0; i < length; i++) { - getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(bytesValues[i]); - } + for (int i = 0; i < length; i++) { + getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(bytesValues[i]); } break; default: @@ -441,19 +445,34 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized CpcSketch objects. + // They always use the single-value representation. + CpcSketch[] sketches = deserializeSketches(blockValSet.getBytesValuesSV(), length); + for (int i = 0; i < length; i++) { + CpcSketch sketch = sketches[i]; + if (sketch != null) { + for (int groupKey : groupKeysArray[i]) { + getAccumulator(groupByResultHolder, groupKey).apply(sketch); + } + } + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -505,22 +524,9 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized CpcSketch objects. - CpcSketch[] sketches = deserializeSketches(bytesValues, length); - for (int i = 0; i < length; i++) { - CpcSketch sketch = sketches[i]; - if (sketch != null) { - for (int groupKey : groupKeysArray[i]) { - getAccumulator(groupByResultHolder, groupKey).apply(sketch); - } - } - } - } else { - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - getCpcSketch(groupByResultHolder, groupKey).update(bytesValues[i]); - } + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + getCpcSketch(groupByResultHolder, groupKey).update(bytesValues[i]); } } break; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java index 1c190d9d9bd5..f4dcced17b70 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -82,22 +82,32 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized HyperLogLog objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLog(aggregationResultHolder, bytesValues[i]); + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); + aggregateSV(length, aggregationResultHolder, blockValSet, storedType); } else { aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType dataType, DataType storedType) { + DataType storedType) { // For dictionary-encoded expression, collect dictionary ids into a BitSet for deduplication. // BitSet gives O(1) insertion with no container-switching overhead (unlike RoaringBitmap), and uses // dictSize/8 bytes of memory (e.g. 128 KB for a 1M-entry dictionary). Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); BitSet bitSet = getDictIdBitSet(aggregationResultHolder, dictionary); for (int i = 0; i < length; i++) { @@ -107,7 +117,7 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult } // For non-dictionary-encoded expression, store values into the HyperLogLog - HyperLogLog hyperLogLog = dataType == DataType.BYTES ? null : getHyperLogLog(aggregationResultHolder); + HyperLogLog hyperLogLog = getHyperLogLog(aggregationResultHolder); switch (storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); @@ -141,15 +151,8 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized HyperLogLog objects. - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLog(aggregationResultHolder, bytesValues[i]); - } - } else { - for (int i = 0; i < length; i++) { - hyperLogLog.offer(bytesValues[i]); - } + for (int i = 0; i < length; i++) { + hyperLogLog.offer(bytesValues[i]); } break; default: @@ -234,23 +237,33 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized HyperLogLog objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLog(groupByResultHolder, groupKeyArray[i], bytesValues[i]); + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, collect dictionary ids into a RoaringBitmap for deduplication. // RoaringBitmap is used (not BitSet) because it is sparse: memory scales with the number of distinct dict IDs // seen per group, not with the full dictionary size. This avoids OOM when many groups each see few distinct values // (contrast with the non-group-by path, which uses a single BitSet across the entire dictionary). Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -292,15 +305,8 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized HyperLogLog objects. - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLog(groupByResultHolder, groupKeyArray[i], bytesValues[i]); - } - } else { - for (int i = 0; i < length; i++) { - getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); - } + for (int i = 0; i < length; i++) { + getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); } break; default: @@ -387,20 +393,32 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized HyperLogLog objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + mergeSerializedHyperLogLog(groupByResultHolder, groupKey, bytesValues[i]); + } + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, collect dictionary ids into a RoaringBitmap (see aggregateSVGroupBySV). Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { int dictId = dictIds[i]; @@ -445,17 +463,8 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized HyperLogLog objects. - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - mergeSerializedHyperLogLog(groupByResultHolder, groupKey, bytesValues[i]); - } - } - } else { - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); - } + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); } break; default: diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java index da0f4163048e..d8282b2bde18 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -94,28 +94,37 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized HyperLogLogPlus objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLogPlus(aggregationResultHolder, bytesValues[i]); + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSV(length, aggregationResultHolder, blockValSet, dataType, storedType); + aggregateSV(length, aggregationResultHolder, blockValSet, storedType); } else { aggregateMV(length, aggregationResultHolder, blockValSet, storedType); } } protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, - DataType dataType, DataType storedType) { + DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; } // For non-dictionary-encoded expression, store values into the HyperLogLogPlus - HyperLogLogPlus hyperLogLogPlus = - dataType == DataType.BYTES ? null : getHyperLogLogPlus(aggregationResultHolder); + HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(aggregationResultHolder); switch (storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); @@ -149,15 +158,8 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized HyperLogLogPlus objects. - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLogPlus(aggregationResultHolder, bytesValues[i]); - } - } else { - for (int i = 0; i < length; i++) { - hyperLogLogPlus.offer(bytesValues[i]); - } + for (int i = 0; i < length; i++) { + hyperLogLogPlus.offer(bytesValues[i]); } break; default: @@ -242,20 +244,30 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized HyperLogLogPlus objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKeyArray[i], bytesValues[i]); + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -297,15 +309,8 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized HyperLogLogPlus objects. - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKeyArray[i], bytesValues[i]); - } - } else { - for (int i = 0; i < length; i++) { - getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); - } + for (int i = 0; i < length; i++) { + getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); } break; default: @@ -394,20 +399,32 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized HyperLogLogPlus objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKey, bytesValues[i]); + } + } + return; + } + DataType storedType = dataType.getStoredType(); if (blockValSet.isSingleValue()) { - aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, dataType, storedType); + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } else { aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } } protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, - BlockValSet blockValSet, DataType dataType, DataType storedType) { + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -449,17 +466,8 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized HyperLogLogPlus objects. - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKey, bytesValues[i]); - } - } - } else { - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); - } + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); } break; default: 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 1dbb76df7ffc..38b27495285a 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 @@ -196,7 +196,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde // Main expression is always index 0 DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES || !singleValues[0]) { + if (dataType != DataType.BYTES) { List updateSketches = getUpdateSketches(aggregationResultHolder); if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { @@ -454,7 +454,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } } else { - // SV logical BYTES values contain serialized ThetaSketch objects. + // Logical BYTES values contain serialized ThetaSketch objects. + // They always use the single-value representation. List thetaSketchAccumulators = getUnions(aggregationResultHolder); ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { @@ -487,7 +488,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol // Main expression is always index 0 DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES || !singleValues[0]) { + if (dataType != DataType.BYTES) { if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { case INT: @@ -722,7 +723,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } } else { - // SV logical BYTES values contain serialized ThetaSketch objects. + // Logical BYTES values contain serialized ThetaSketch objects. + // They always use the single-value representation. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); for (int i = 0; i < length; i++) { List thetaSketchAccumulators = @@ -752,7 +754,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult // Main expression is always index 0 DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES || !singleValues[0]) { + if (dataType != DataType.BYTES) { if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { case INT: @@ -1046,7 +1048,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } } else { - // SV logical BYTES values contain serialized ThetaSketch objects. + // Logical BYTES values contain serialized ThetaSketch objects. + // They always use the single-value representation. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java index 96d6ff825ea6..4358dac24aa6 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -83,18 +83,39 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized UltraLogLog objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + try { + UltraLogLog ull = aggregationResultHolder.getResult(); + if (ull == null) { + ull = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0]); + aggregationResultHolder.setValue(ull); + } else { + ull.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0])); + } + for (int i = 1; i < length; i++) { + ull.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging UltraLogLogs", e); + } + return; + } + DataType storedType = dataType.getStoredType(); // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); getDictIdBitmap(aggregationResultHolder, dictionary).addN(dictIds, 0, length); return; } // For non-dictionary-encoded expression, store values into the UltraLogLog - UltraLogLog ull = dataType == DataType.BYTES ? null : getULL(aggregationResultHolder); + UltraLogLog ull = getULL(aggregationResultHolder); switch (storedType) { case INT: int[] intValues = blockValSet.getIntValuesSV(); @@ -128,26 +149,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized UltraLogLog objects. - try { - UltraLogLog serializedUll = aggregationResultHolder.getResult(); - if (serializedUll == null) { - serializedUll = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0]); - aggregationResultHolder.setValue(serializedUll); - } else { - serializedUll.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[0])); - } - for (int i = 1; i < length; i++) { - serializedUll.add(ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging UltraLogLogs", e); - } - } else { - for (int i = 0; i < length; i++) { - UltraLogLogUtils.hashObject(bytesValues[i]).ifPresent(ull::add); - } + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(bytesValues[i]).ifPresent(ull::add); } break; default: @@ -162,11 +165,32 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized UltraLogLog objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + try { + for (int i = 0; i < length; i++) { + UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); + int groupKey = groupKeyArray[i]; + UltraLogLog ull = groupByResultHolder.getResult(groupKey); + if (ull != null) { + ull.add(value); + } else { + groupByResultHolder.setValueForKey(groupKey, value); + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging UltraLogLog", e); + } + return; + } + DataType storedType = dataType.getStoredType(); // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); @@ -213,27 +237,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized UltraLogLog objects. - try { - for (int i = 0; i < length; i++) { - UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); - int groupKey = groupKeyArray[i]; - UltraLogLog ull = groupByResultHolder.getResult(groupKey); - if (ull != null) { - ull.add(value); - } else { - groupByResultHolder.setValueForKey(groupKey, value); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging UltraLogLog", e); - } - } else { - for (int i = 0; i < length; i++) { - UltraLogLogUtils.hashObject(bytesValues[i]) - .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); - } + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(bytesValues[i]) + .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); } break; default: @@ -248,11 +254,34 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult BlockValSet blockValSet = blockValSetMap.get(_expression); DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES values contain serialized UltraLogLog objects. + // They always use the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + try { + for (int i = 0; i < length; i++) { + UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + UltraLogLog ull = groupByResultHolder.getResult(groupKey); + if (ull != null) { + ull.add(value); + } else { + groupByResultHolder.setValueForKey(groupKey, + ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); + } + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging UltraLogLog", e); + } + return; + } + DataType storedType = dataType.getStoredType(); // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; - if (dictionary != null && dataType != DataType.BYTES) { + if (dictionary != null) { int[] dictIds = blockValSet.getDictionaryIdsSV(); for (int i = 0; i < length; i++) { setDictIdForGroupKeys(groupByResultHolder, groupKeysArray[i], dictionary, dictIds[i]); @@ -294,28 +323,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - if (dataType == DataType.BYTES) { - // SV logical BYTES values contain serialized UltraLogLog objects. - try { - for (int i = 0; i < length; i++) { - UltraLogLog value = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i]); - for (int groupKey : groupKeysArray[i]) { - UltraLogLog ull = groupByResultHolder.getResult(groupKey); - if (ull != null) { - ull.add(value); - } else { - groupByResultHolder.setValueForKey(groupKey, - ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); - } - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging UltraLogLog", e); - } - } else { - for (int i = 0; i < length; i++) { - setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); - } + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); } break; default: diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java index 788098efa3af..20f93011e422 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java @@ -29,6 +29,7 @@ import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.spi.data.FieldSpec.DataType; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.annotations.Test; import static org.mockito.Mockito.anyInt; @@ -45,7 +46,7 @@ /// result resolver used by the non-scan based and partial metadata based aggregation paths. @SuppressWarnings("rawtypes") public class AggregationFunctionUtilsTest { - private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); + private static final ExpressionContext INPUT_EXPRESSION = ExpressionContext.forIdentifier("inputCol"); private static AggregationFunction mockFunction(AggregationFunctionType type) { AggregationFunction aggregationFunction = mock(AggregationFunction.class); @@ -96,17 +97,21 @@ public void testNonCountWithNullDataSourceThrows() { } @Test - public void testDistinctCountHllMvOffersDictionaryBytesAsRawValues() + public void testDistinctCountHllMvOffersUuidDictionaryBytesAsRawValues() throws IOException { DistinctCountHLLMVAggregationFunction function = - new DistinctCountHLLMVAggregationFunction(List.of(BYTES_EXPRESSION)); - byte[][] bytesValues = {{1}, {2}, {3}}; + new DistinctCountHLLMVAggregationFunction(List.of(INPUT_EXPRESSION)); + byte[][] bytesValues = { + UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"), + UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"), + UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002") + }; Dictionary dictionary = mock(Dictionary.class); when(dictionary.length()).thenReturn(bytesValues.length); for (int i = 0; i < bytesValues.length; i++) { when(dictionary.get(i)).thenReturn(bytesValues[i]); } - DataSource dataSource = mockBytesDataSource(dictionary, false); + DataSource dataSource = mockDataSource(dictionary, DataType.UUID, false); HyperLogLog result = (HyperLogLog) AggregationFunctionUtils.getAggregationResult(function, dataSource, 3, "TEST"); @@ -123,17 +128,18 @@ public void testDistinctCountHllMvOffersDictionaryBytesAsRawValues() } @Test - public void testDistinctCountHllMergesSingleValueSerializedBytes() + public void testDistinctCountHllMergesLogicalBytesAsSerializedState() throws IOException { DistinctCountHLLAggregationFunction function = - new DistinctCountHLLAggregationFunction(List.of(BYTES_EXPRESSION)); + new DistinctCountHLLAggregationFunction(List.of(INPUT_EXPRESSION)); byte[] firstSketch = serializedHll(function, "a", "b"); byte[] secondSketch = serializedHll(function, "b", "c"); Dictionary dictionary = mock(Dictionary.class); when(dictionary.length()).thenReturn(2); when(dictionary.getBytesValue(0)).thenReturn(firstSketch); when(dictionary.getBytesValue(1)).thenReturn(secondSketch); - DataSource dataSource = mockBytesDataSource(dictionary, true); + // Logical BYTES determines serialized-state handling independently of the cardinality metadata. + DataSource dataSource = mockDataSource(dictionary, DataType.BYTES, false); HyperLogLog result = (HyperLogLog) AggregationFunctionUtils.getAggregationResult(function, dataSource, 2, "TEST"); @@ -148,9 +154,9 @@ public void testDistinctCountHllMergesSingleValueSerializedBytes() verify(dictionary, never()).get(anyInt()); } - private static DataSource mockBytesDataSource(Dictionary dictionary, boolean singleValue) { + private static DataSource mockDataSource(Dictionary dictionary, DataType dataType, boolean singleValue) { DataSourceMetadata metadata = mock(DataSourceMetadata.class); - when(metadata.getDataType()).thenReturn(DataType.BYTES); + when(metadata.getDataType()).thenReturn(dataType); when(metadata.isSingleValue()).thenReturn(singleValue); DataSource dataSource = mock(DataSource.class); when(dataSource.getDictionary()).thenReturn(dictionary); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java index 1105d53bda10..9c1f4fda78d7 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java @@ -28,6 +28,7 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; @@ -39,7 +40,11 @@ public class DistinctCountBitmapMVAggregationFunctionTest extends AbstractAggregationFunctionTest { - private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); + private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); + private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); + private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); + private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); + private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); @Test public void testAggregationMV() { @@ -109,34 +114,34 @@ public void testAggregationMVGroupByMV() { } @Test - public void testMultiValueBytesUseMultiValueAccessor() { - byte[][][] bytesValues = { - {{1}, {2}}, - {{2}, {3}}, - {{1}}, - {{4}} + public void testMultiValueUuidUsesMultiValueAccessor() { + byte[][][] uuidValues = { + {UUID_0, UUID_1}, + {UUID_1, UUID_2}, + {UUID_0}, + {UUID_3} }; BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); when(blockValSet.isSingleValue()).thenReturn(false); when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); + when(blockValSet.getBytesValuesMV()).thenReturn(uuidValues); DistinctCountBitmapAggregationFunction function = - new DistinctCountBitmapAggregationFunction(List.of(BYTES_EXPRESSION)); + new DistinctCountBitmapAggregationFunction(List.of(UUID_EXPRESSION)); AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4); GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2); GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java index 417c9f9b17ce..4d25ed21341f 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java @@ -55,14 +55,6 @@ public class DistinctCountCPCSketchAggregationFunctionTest { private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); private static final byte[][] UUID_VALUES_SV = {UUID_0, UUID_1, UUID_0, UUID_2}; private static final byte[][][] UUID_VALUES_MV = {{UUID_0, UUID_1}, {UUID_1, UUID_2}, {UUID_0}, {UUID_3}}; - // These payloads are intentionally too short to be serialized CPC sketches. MV BYTES are raw values. - private static final byte[] RAW_BYTES_0 = {1}; - private static final byte[] RAW_BYTES_1 = {2}; - private static final byte[] RAW_BYTES_2 = {3}; - private static final byte[] RAW_BYTES_3 = {4}; - private static final byte[][][] RAW_BYTES_VALUES_MV = { - {RAW_BYTES_0, RAW_BYTES_1}, {RAW_BYTES_1, RAW_BYTES_2}, {RAW_BYTES_0}, {RAW_BYTES_3} - }; @DataProvider(name = "uuidValueModes") public static Object[][] uuidValueModes() { @@ -238,7 +230,7 @@ public void testAggregateDictionaryEncodedUuid() { } @Test - public void testAggregateSingleValueSerializedSketches() { + public void testLogicalBytesUsesSerializedSingleValueAccessorBeforeDispatch() { byte[][] serializedSketches = { serializedSketch("a", "b"), serializedSketch("b", "c"), @@ -247,8 +239,6 @@ public void testAggregateSingleValueSerializedSketches() { }; BlockValSet blockValSet = mock(BlockValSet.class); when(blockValSet.getValueType()).thenReturn(DataType.BYTES); - when(blockValSet.isSingleValue()).thenReturn(true); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); when(blockValSet.getBytesValuesSV()).thenReturn(serializedSketches); DistinctCountCPCSketchAggregationFunction function = new DistinctCountCPCSketchAggregationFunction(List.of(BYTES_EXPRESSION)); @@ -270,59 +260,41 @@ public void testAggregateSingleValueSerializedSketches() { Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); verify(blockValSet, atLeastOnce()).getBytesValuesSV(); verify(blockValSet, never()).getBytesValuesMV(); + verify(blockValSet, never()).isSingleValue(); } @Test - public void testAggregateMultiValueRawBytes() { + public void testAggregateDictionaryEncodedMultiValueUuid() { + Dictionary dictionary = mock(Dictionary.class); + when(dictionary.get(0)).thenReturn(UUID_0); + when(dictionary.get(1)).thenReturn(UUID_1); + when(dictionary.get(2)).thenReturn(UUID_2); + when(dictionary.get(3)).thenReturn(UUID_3); + BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(RAW_BYTES_VALUES_MV); + when(blockValSet.isDictionaryEncoded()).thenReturn(true); + when(blockValSet.getDictionary()).thenReturn(dictionary); + when(blockValSet.getDictionaryIdsMV()).thenReturn(new int[][]{{0, 1}, {1, 2}, {0}, {3}}); DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(BYTES_EXPRESSION)); + new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(RAW_BYTES_VALUES_MV.length, aggregationResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregate(UUID_VALUES_MV.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(RAW_BYTES_VALUES_MV.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupBySV(UUID_VALUES_MV.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(RAW_BYTES_VALUES_MV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupByMV(UUID_VALUES_MV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); - } - - @Test - public void testAggregateDictionaryEncodedMultiValueRawBytes() { - Dictionary dictionary = mock(Dictionary.class); - when(dictionary.get(0)).thenReturn(RAW_BYTES_0); - when(dictionary.get(1)).thenReturn(RAW_BYTES_1); - when(dictionary.get(2)).thenReturn(RAW_BYTES_2); - when(dictionary.get(3)).thenReturn(RAW_BYTES_3); - - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); - when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.isDictionaryEncoded()).thenReturn(true); - when(blockValSet.getDictionary()).thenReturn(dictionary); - when(blockValSet.getDictionaryIdsMV()).thenReturn(new int[][]{{0, 1}, {1, 2}, {0}, {3}}); - DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(BYTES_EXPRESSION)); - - AggregationResultHolder resultHolder = function.createAggregationResultHolder(); - function.aggregate(RAW_BYTES_VALUES_MV.length, resultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder), 4L); verify(blockValSet, atLeastOnce()).getDictionaryIdsMV(); verify(blockValSet, never()).getBytesValuesSV(); verify(blockValSet, never()).getBytesValuesMV(); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java index 4a2c5eac5571..2ec9960741d2 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java @@ -28,6 +28,7 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; @@ -39,7 +40,11 @@ public class DistinctCountHLLMVAggregationFunctionTest extends AbstractAggregationFunctionTest { - private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); + private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); + private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); + private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); + private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); + private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); @Test public void testAggregationMV() { @@ -109,34 +114,34 @@ public void testAggregationMVGroupByMV() { } @Test - public void testMultiValueBytesUseMultiValueAccessor() { + public void testMultiValueUuidUsesMultiValueAccessor() { DistinctCountHLLAggregationFunction function = - new DistinctCountHLLAggregationFunction(List.of(BYTES_EXPRESSION)); - byte[][][] bytesValues = { - {{1}, {2}}, - {{2}, {3}}, - {{1}}, - {{4}} + new DistinctCountHLLAggregationFunction(List.of(UUID_EXPRESSION)); + byte[][][] uuidValues = { + {UUID_0, UUID_1}, + {UUID_1, UUID_2}, + {UUID_0}, + {UUID_3} }; BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); when(blockValSet.isSingleValue()).thenReturn(false); when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); + when(blockValSet.getBytesValuesMV()).thenReturn(uuidValues); AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java index f680143bbc6d..648451f57e51 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java @@ -28,6 +28,7 @@ import org.apache.pinot.spi.data.FieldSpec; import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; @@ -39,7 +40,11 @@ public class DistinctCountHLLPlusMVAggregationFunctionTest extends AbstractAggregationFunctionTest { - private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); + private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); + private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); + private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); + private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); + private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); @Test public void testAggregationMV() { @@ -109,34 +114,34 @@ public void testAggregationMVGroupByMV() { } @Test - public void testMultiValueBytesUseMultiValueAccessor() { + public void testMultiValueUuidUsesMultiValueAccessor() { DistinctCountHLLPlusAggregationFunction function = - new DistinctCountHLLPlusAggregationFunction(List.of(BYTES_EXPRESSION)); - byte[][][] bytesValues = { - {{1}, {2}}, - {{2}, {3}}, - {{1}}, - {{4}} + new DistinctCountHLLPlusAggregationFunction(List.of(UUID_EXPRESSION)); + byte[][][] uuidValues = { + {UUID_0, UUID_1}, + {UUID_1, UUID_2}, + {UUID_0}, + {UUID_3} }; BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); + when(blockValSet.getValueType()).thenReturn(DataType.UUID); when(blockValSet.isSingleValue()).thenReturn(false); when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); + when(blockValSet.getBytesValuesMV()).thenReturn(uuidValues); AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, + Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); + function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, + groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java index 906e7fafb897..2ced8f05ff4f 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java @@ -139,41 +139,6 @@ public void testUuidPredicateUsesLogicalType() { Assert.assertEquals(extractFinalResult(function, resultHolder), 1L); } - @Test - public void testAggregateMultiValueBytesAsRawValues() { - // These one-byte values are not serialized sketches, so the test also verifies that MV BYTES are not deserialized. - byte[][][] bytesValues = { - {{1}, {2}}, - {{2}, {3}}, - {{1}}, - {{4}} - }; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); - when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(bytesValues); - DistinctCountThetaSketchAggregationFunction function = - new DistinctCountThetaSketchAggregationFunction(List.of(BYTES_EXPRESSION)); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(bytesValues.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(bytesValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(bytesValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, groupByMVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); - } - @Test public void testAggregateSingleValueSerializedSketches() { byte[][] serializedSketches = { From 49cc6ddca8ff84a1ce367fcadac51fa09564a948 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Fri, 14 Aug 2026 17:34:31 -0700 Subject: [PATCH 17/19] Simplify UUID aggregation coverage --- .../function/AggregationFunctionUtils.java | 6 +- ...istinctCountBitmapAggregationFunction.java | 40 ++- ...inctCountCPCSketchAggregationFunction.java | 235 ++++-------------- .../DistinctCountHLLAggregationFunction.java | 79 +++--- ...stinctCountHLLPlusAggregationFunction.java | 78 +++--- ...ctCountThetaSketchAggregationFunction.java | 43 +--- .../DistinctCountULLAggregationFunction.java | 9 +- .../distinct/table/BytesDistinctTable.java | 22 +- .../query/reduce/GroupByDataTableReducer.java | 5 - .../AggregationFunctionUtilsTest.java | 90 ------- ...tCountBitmapMVAggregationFunctionTest.java | 66 ----- ...CountCPCSketchAggregationFunctionTest.java | 202 --------------- ...stinctCountHLLAggregationFunctionTest.java | 109 -------- ...inctCountHLLMVAggregationFunctionTest.java | 66 ----- ...CountHLLPlusMVAggregationFunctionTest.java | 66 ----- ...untThetaSketchAggregationFunctionTest.java | 158 ------------ ...stinctCountULLAggregationFunctionTest.java | 112 --------- .../table/BytesDistinctTableTest.java | 70 ------ .../tests/custom/UuidAggregationTest.java | 192 +++++++------- 19 files changed, 288 insertions(+), 1360 deletions(-) delete mode 100644 pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index 7c7ac22b2088..bf9f716d9933 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -805,7 +805,7 @@ private static HyperLogLog getDistinctCountHLLResult(DataSource dataSource, Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); if (metadata.getDataType() == FieldSpec.DataType.BYTES) { - // Logical BYTES dictionary values contain serialized HyperLogLog objects. + // Logical BYTES dictionary entries are serialized HyperLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLog hll = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -828,7 +828,7 @@ private static HyperLogLogPlus getDistinctCountHLLPlusResult(DataSource dataSour Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); if (metadata.getDataType() == FieldSpec.DataType.BYTES) { - // Logical BYTES dictionary values contain serialized HyperLogLogPlus objects. + // Logical BYTES dictionary entries are serialized HyperLogLogPlus objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLogPlus hllplus = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -871,7 +871,7 @@ private static UltraLogLog getDistinctCountULLResult(DataSource dataSource, Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); if (metadata.getDataType() == FieldSpec.DataType.BYTES) { - // Logical BYTES dictionary values contain serialized UltraLogLog objects. + // Logical BYTES dictionary entries are serialized UltraLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); UltraLogLog ull = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(dictionary.getBytesValue(0)); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index 593ce48fc37a..e71632673b29 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -74,12 +74,19 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized RoaringBitmap objects. - // They always use the single-value representation. + // Logical BYTES is a serialized RoaringBitmap and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); - RoaringBitmap valueBitmap = getValueBitmap(aggregationResultHolder); - for (int i = 0; i < length; i++) { - valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); + RoaringBitmap valueBitmap = aggregationResultHolder.getResult(); + if (valueBitmap != null) { + for (int i = 0; i < length; i++) { + valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); + } + } else { + valueBitmap = RoaringBitmapUtils.deserialize(bytesValues[0]); + aggregationResultHolder.setValue(valueBitmap); + for (int i = 1; i < length; i++) { + valueBitmap.or(RoaringBitmapUtils.deserialize(bytesValues[i])); + } } return; } @@ -221,11 +228,17 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized RoaringBitmap objects. - // They always use the single-value representation. + // Logical BYTES is a serialized RoaringBitmap and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { - getValueBitmap(groupByResultHolder, groupKeyArray[i]).or(RoaringBitmapUtils.deserialize(bytesValues[i])); + RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); + int groupKey = groupKeyArray[i]; + RoaringBitmap valueBitmap = groupByResultHolder.getResult(groupKey); + if (valueBitmap != null) { + valueBitmap.or(value); + } else { + groupByResultHolder.setValueForKey(groupKey, value); + } } return; } @@ -373,13 +386,18 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized RoaringBitmap objects. - // They always use the single-value representation. + // Logical BYTES is a serialized RoaringBitmap and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); for (int groupKey : groupKeysArray[i]) { - getValueBitmap(groupByResultHolder, groupKey).or(value); + RoaringBitmap bitmap = groupByResultHolder.getResult(groupKey); + if (bitmap != null) { + bitmap.or(value); + } else { + // Clone a bitmap for the group + groupByResultHolder.setValueForKey(groupKey, value.clone()); + } } } return; diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java index 96857e84e03c..8194222fb694 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -137,9 +137,19 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized CpcSketch objects. - // They always use the single-value representation. - mergeSerializedSketches(aggregationResultHolder, blockValSet.getBytesValuesSV(), length); + // Logical BYTES stores serialized CpcSketch objects in the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + try { + CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(aggregationResultHolder); + CpcSketch[] sketches = deserializeSketches(bytesValues, length); + for (CpcSketch sketch : sketches) { + if (sketch != null) { + cpcSketchAccumulator.apply(sketch); + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging CPC sketches", e); + } return; } @@ -162,45 +172,40 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult } // For non-dictionary-encoded expression, store values into the CpcSketch + CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); switch (storedType) { case INT: - CpcSketch intSketch = getCpcSketch(aggregationResultHolder); int[] intValues = blockValSet.getIntValuesSV(); for (int i = 0; i < length; i++) { - intSketch.update(intValues[i]); + cpcSketch.update(intValues[i]); } break; case LONG: - CpcSketch longSketch = getCpcSketch(aggregationResultHolder); long[] longValues = blockValSet.getLongValuesSV(); for (int i = 0; i < length; i++) { - longSketch.update(longValues[i]); + cpcSketch.update(longValues[i]); } break; case FLOAT: - CpcSketch floatSketch = getCpcSketch(aggregationResultHolder); float[] floatValues = blockValSet.getFloatValuesSV(); for (int i = 0; i < length; i++) { - floatSketch.update(floatValues[i]); + cpcSketch.update(floatValues[i]); } break; case DOUBLE: - CpcSketch doubleSketch = getCpcSketch(aggregationResultHolder); double[] doubleValues = blockValSet.getDoubleValuesSV(); for (int i = 0; i < length; i++) { - doubleSketch.update(doubleValues[i]); + cpcSketch.update(doubleValues[i]); } break; case STRING: - CpcSketch stringSketch = getCpcSketch(aggregationResultHolder); String[] stringValues = blockValSet.getStringValuesSV(); for (int i = 0; i < length; i++) { - stringSketch.update(stringValues[i]); + cpcSketch.update(stringValues[i]); } break; case BYTES: byte[][] bytesValues = blockValSet.getBytesValuesSV(); - CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); for (int i = 0; i < length; i++) { cpcSketch.update(bytesValues[i]); } @@ -225,51 +230,6 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult // For non-dictionary-encoded expression, store values into the CpcSketch switch (storedType) { - case INT: - CpcSketch intSketch = getCpcSketch(aggregationResultHolder); - int[][] intValues = blockValSet.getIntValuesMV(); - for (int i = 0; i < length; i++) { - for (int value : intValues[i]) { - intSketch.update(value); - } - } - break; - case LONG: - CpcSketch longSketch = getCpcSketch(aggregationResultHolder); - long[][] longValues = blockValSet.getLongValuesMV(); - for (int i = 0; i < length; i++) { - for (long value : longValues[i]) { - longSketch.update(value); - } - } - break; - case FLOAT: - CpcSketch floatSketch = getCpcSketch(aggregationResultHolder); - float[][] floatValues = blockValSet.getFloatValuesMV(); - for (int i = 0; i < length; i++) { - for (float value : floatValues[i]) { - floatSketch.update(value); - } - } - break; - case DOUBLE: - CpcSketch doubleSketch = getCpcSketch(aggregationResultHolder); - double[][] doubleValues = blockValSet.getDoubleValuesMV(); - for (int i = 0; i < length; i++) { - for (double value : doubleValues[i]) { - doubleSketch.update(value); - } - } - break; - case STRING: - CpcSketch stringSketch = getCpcSketch(aggregationResultHolder); - String[][] stringValues = blockValSet.getStringValuesMV(); - for (int i = 0; i < length; i++) { - for (String value : stringValues[i]) { - stringSketch.update(value); - } - } - break; case BYTES: byte[][][] bytesValues = blockValSet.getBytesValuesMV(); CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); @@ -291,14 +251,19 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized CpcSketch objects. - // They always use the single-value representation. - CpcSketch[] sketches = deserializeSketches(blockValSet.getBytesValuesSV(), length); - for (int i = 0; i < length; i++) { - CpcSketch sketch = sketches[i]; - if (sketch != null) { - getAccumulator(groupByResultHolder, groupKeyArray[i]).apply(sketch); + // Logical BYTES stores serialized CpcSketch objects in the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + try { + CpcSketch[] sketches = deserializeSketches(bytesValues, length); + for (int i = 0; i < length; i++) { + CpcSketch sketch = sketches[i]; + if (sketch != null) { + CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(groupByResultHolder, groupKeyArray[i]); + cpcSketchAccumulator.apply(sketch); + } } + } catch (Exception e) { + throw new RuntimeException("Caught exception while aggregating CPC Sketches", e); } return; } @@ -380,51 +345,6 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu // For non-dictionary-encoded expression, store values into the CpcSketch switch (storedType) { - case INT: - int[][] intValues = blockValSet.getIntValuesMV(); - for (int i = 0; i < length; i++) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); - for (int value : intValues[i]) { - cpcSketch.update(value); - } - } - break; - case LONG: - long[][] longValues = blockValSet.getLongValuesMV(); - for (int i = 0; i < length; i++) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); - for (long value : longValues[i]) { - cpcSketch.update(value); - } - } - break; - case FLOAT: - float[][] floatValues = blockValSet.getFloatValuesMV(); - for (int i = 0; i < length; i++) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); - for (float value : floatValues[i]) { - cpcSketch.update(value); - } - } - break; - case DOUBLE: - double[][] doubleValues = blockValSet.getDoubleValuesMV(); - for (int i = 0; i < length; i++) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); - for (double value : doubleValues[i]) { - cpcSketch.update(value); - } - } - break; - case STRING: - String[][] stringValues = blockValSet.getStringValuesMV(); - for (int i = 0; i < length; i++) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); - for (String value : stringValues[i]) { - cpcSketch.update(value); - } - } - break; case BYTES: byte[][][] bytesValues = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { @@ -446,16 +366,19 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized CpcSketch objects. - // They always use the single-value representation. - CpcSketch[] sketches = deserializeSketches(blockValSet.getBytesValuesSV(), length); - for (int i = 0; i < length; i++) { - CpcSketch sketch = sketches[i]; - if (sketch != null) { - for (int groupKey : groupKeysArray[i]) { - getAccumulator(groupByResultHolder, groupKey).apply(sketch); + // Logical BYTES stores serialized CpcSketch objects in the single-value representation. + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + try { + CpcSketch[] sketches = deserializeSketches(bytesValues, length); + for (int i = 0; i < length; i++) { + if (sketches[i] != null) { + for (int groupKey : groupKeysArray[i]) { + getAccumulator(groupByResultHolder, groupKey).apply(sketches[i]); + } } } + } catch (Exception e) { + throw new RuntimeException("Caught exception while aggregating CPC sketches", e); } return; } @@ -552,61 +475,6 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR // For non-dictionary-encoded expression, store values into the CpcSketch switch (storedType) { - case INT: - int[][] intValues = blockValSet.getIntValuesMV(); - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); - for (int value : intValues[i]) { - cpcSketch.update(value); - } - } - } - break; - case LONG: - long[][] longValues = blockValSet.getLongValuesMV(); - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); - for (long value : longValues[i]) { - cpcSketch.update(value); - } - } - } - break; - case FLOAT: - float[][] floatValues = blockValSet.getFloatValuesMV(); - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); - for (float value : floatValues[i]) { - cpcSketch.update(value); - } - } - } - break; - case DOUBLE: - double[][] doubleValues = blockValSet.getDoubleValuesMV(); - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); - for (double value : doubleValues[i]) { - cpcSketch.update(value); - } - } - } - break; - case STRING: - String[][] stringValues = blockValSet.getStringValuesMV(); - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); - for (String value : stringValues[i]) { - cpcSketch.update(value); - } - } - } - break; case BYTES: byte[][][] bytesValues = blockValSet.getBytesValuesMV(); for (int i = 0; i < length; i++) { @@ -815,10 +683,6 @@ private void addObjectsToSketch(Object[] rawValues, CpcSketch sketch) { for (String s : (String[]) rawValues) { sketch.update(s); } - } else if (rawValues instanceof byte[][]) { - for (byte[] bytes : (byte[][]) rawValues) { - sketch.update(bytes); - } } else if (rawValues instanceof Integer[]) { for (Integer i : (Integer[]) rawValues) { sketch.update(i); @@ -861,21 +725,6 @@ private CpcSketchAccumulator getAccumulator(GroupByResultHolder groupByResultHol return accumulator; } - private void mergeSerializedSketches(AggregationResultHolder aggregationResultHolder, byte[][] bytesValues, - int length) { - try { - CpcSketchAccumulator accumulator = getAccumulator(aggregationResultHolder); - CpcSketch[] sketches = deserializeSketches(bytesValues, length); - for (CpcSketch sketch : sketches) { - if (sketch != null) { - accumulator.apply(sketch); - } - } - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging CPC sketches", e); - } - } - /// Deserializes the sketches from the bytes. Returns null for empty byte arrays which represent /// the default null value for BYTES columns in Pinot. Callers must handle null entries. @SuppressWarnings({"unchecked"}) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java index f4dcced17b70..1d68cde59588 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -83,11 +83,23 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized HyperLogLog objects. - // They always use the single-value representation. + // Logical BYTES is a serialized HyperLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLog(aggregationResultHolder, bytesValues[i]); + try { + HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); + if (hyperLogLog != null) { + for (int i = 0; i < length; i++) { + hyperLogLog.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i])); + } + } else { + hyperLogLog = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[0]); + aggregationResultHolder.setValue(hyperLogLog); + for (int i = 1; i < length; i++) { + hyperLogLog.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i])); + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogs", e); } return; } @@ -238,11 +250,21 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized HyperLogLog objects. - // They always use the single-value representation. + // Logical BYTES is a serialized HyperLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLog(groupByResultHolder, groupKeyArray[i], bytesValues[i]); + try { + for (int i = 0; i < length; i++) { + HyperLogLog value = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i]); + int groupKey = groupKeyArray[i]; + HyperLogLog hyperLogLog = groupByResultHolder.getResult(groupKey); + if (hyperLogLog != null) { + hyperLogLog.addAll(value); + } else { + groupByResultHolder.setValueForKey(groupKey, value); + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogs", e); } return; } @@ -394,13 +416,24 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized HyperLogLog objects. - // They always use the single-value representation. + // Logical BYTES is a serialized HyperLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - mergeSerializedHyperLogLog(groupByResultHolder, groupKey, bytesValues[i]); + try { + for (int i = 0; i < length; i++) { + HyperLogLog value = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + HyperLogLog hyperLogLog = groupByResultHolder.getResult(groupKey); + if (hyperLogLog != null) { + hyperLogLog.addAll(value); + } else { + // Create a new HyperLogLog for the group + groupByResultHolder.setValueForKey(groupKey, + ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytesValues[i])); + } + } } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogs", e); } return; } @@ -684,26 +717,6 @@ protected static BitSet getDictIdBitSet(AggregationResultHolder aggregationResul return dictIdsWrapper._bitSet; } - private void mergeSerializedHyperLogLog(AggregationResultHolder aggregationResultHolder, byte[] bytes) { - HyperLogLog value = deserializeHyperLogLog(bytes); - HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); - aggregationResultHolder.setValue(hyperLogLog == null ? value : merge(hyperLogLog, value)); - } - - private void mergeSerializedHyperLogLog(GroupByResultHolder groupByResultHolder, int groupKey, byte[] bytes) { - HyperLogLog value = deserializeHyperLogLog(bytes); - HyperLogLog hyperLogLog = groupByResultHolder.getResult(groupKey); - groupByResultHolder.setValueForKey(groupKey, hyperLogLog == null ? value : merge(hyperLogLog, value)); - } - - private static HyperLogLog deserializeHyperLogLog(byte[] bytes) { - try { - return ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(bytes); - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogs", e); - } - } - /// Returns the HyperLogLog from the result holder or creates a new one if it does not exist. protected HyperLogLog getHyperLogLog(AggregationResultHolder aggregationResultHolder) { HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java index d8282b2bde18..dfda4ed0bc7e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -95,11 +95,21 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized HyperLogLogPlus objects. - // They always use the single-value representation. + // Logical BYTES is a serialized HyperLogLogPlus and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLogPlus(aggregationResultHolder, bytesValues[i]); + try { + HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); + if (hyperLogLogPlus == null) { + hyperLogLogPlus = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[0]); + aggregationResultHolder.setValue(hyperLogLogPlus); + } else { + hyperLogLogPlus.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[0])); + } + for (int i = 1; i < length; i++) { + hyperLogLogPlus.addAll(ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i])); + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); } return; } @@ -245,11 +255,21 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized HyperLogLogPlus objects. - // They always use the single-value representation. + // Logical BYTES is a serialized HyperLogLogPlus and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKeyArray[i], bytesValues[i]); + try { + for (int i = 0; i < length; i++) { + HyperLogLogPlus value = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i]); + int groupKey = groupKeyArray[i]; + HyperLogLogPlus hyperLogLogPlus = groupByResultHolder.getResult(groupKey); + if (hyperLogLogPlus != null) { + hyperLogLogPlus.addAll(value); + } else { + groupByResultHolder.setValueForKey(groupKey, value); + } + } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); } return; } @@ -400,13 +420,24 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized HyperLogLogPlus objects. - // They always use the single-value representation. + // Logical BYTES is a serialized HyperLogLogPlus and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); - for (int i = 0; i < length; i++) { - for (int groupKey : groupKeysArray[i]) { - mergeSerializedHyperLogLogPlus(groupByResultHolder, groupKey, bytesValues[i]); + try { + for (int i = 0; i < length; i++) { + HyperLogLogPlus value = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i]); + for (int groupKey : groupKeysArray[i]) { + HyperLogLogPlus hyperLogLogPlus = groupByResultHolder.getResult(groupKey); + if (hyperLogLogPlus != null) { + hyperLogLogPlus.addAll(value); + } else { + // Create a new HyperLogLogPlus for the group + groupByResultHolder.setValueForKey(groupKey, + ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytesValues[i])); + } + } } + } catch (Exception e) { + throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); } return; } @@ -677,27 +708,6 @@ protected static RoaringBitmap getDictIdBitmap(AggregationResultHolder aggregati return dictIdsWrapper._dictIdBitmap; } - private void mergeSerializedHyperLogLogPlus(AggregationResultHolder aggregationResultHolder, byte[] bytes) { - HyperLogLogPlus value = deserializeHyperLogLogPlus(bytes); - HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); - aggregationResultHolder.setValue(hyperLogLogPlus == null ? value : merge(hyperLogLogPlus, value)); - } - - private void mergeSerializedHyperLogLogPlus(GroupByResultHolder groupByResultHolder, int groupKey, byte[] bytes) { - HyperLogLogPlus value = deserializeHyperLogLogPlus(bytes); - HyperLogLogPlus hyperLogLogPlus = groupByResultHolder.getResult(groupKey); - groupByResultHolder.setValueForKey(groupKey, - hyperLogLogPlus == null ? value : merge(hyperLogLogPlus, value)); - } - - private static HyperLogLogPlus deserializeHyperLogLogPlus(byte[] bytes) { - try { - return ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(bytes); - } catch (Exception e) { - throw new RuntimeException("Caught exception while merging HyperLogLogPlus", e); - } - } - /// Returns the HyperLogLogPlus from the result holder or creates a new one if it does not exist. protected HyperLogLogPlus getHyperLogLogPlus(AggregationResultHolder aggregationResultHolder) { HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); 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 38b27495285a..9118d766124c 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 @@ -91,7 +91,6 @@ public class DistinctCountThetaSketchAggregationFunction private final List _filterEvaluators; private final ExpressionContext _postAggregationExpression; private final UpdatableThetaSketchBuilder _updateSketchBuilder = new UpdatableThetaSketchBuilder(); - private final ThetaSketch _emptySketch; private int _nominalEntries = ThetaUtil.DEFAULT_NOMINAL_ENTRIES; protected final ThetaSetOperationBuilder _setOperationBuilder = new ThetaSetOperationBuilder(); protected int _accumulatorThreshold = DEFAULT_ACCUMULATOR_THRESHOLD; @@ -117,7 +116,6 @@ public DistinctCountThetaSketchAggregationFunction(List argum _setOperationBuilder.setP(p); _updateSketchBuilder.setP(p); } - _emptySketch = _updateSketchBuilder.build().compact(); if (numArguments < 4) { // Simple union without post-aggregation @@ -195,8 +193,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde int numFilters = _filterEvaluators.size(); // Main expression is always index 0 - DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES) { + if (valueTypes[0] != DataType.BYTES) { List updateSketches = getUpdateSketches(aggregationResultHolder); if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { @@ -449,13 +446,11 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde break; default: throw new IllegalStateException( - "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " - + valueTypes[0]); + "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } } else { - // Logical BYTES values contain serialized ThetaSketch objects. - // They always use the single-value representation. + // Logical BYTES stores serialized ThetaSketch objects in the single-value representation. List thetaSketchAccumulators = getUnions(aggregationResultHolder); ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { @@ -487,8 +482,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol int numFilters = _filterEvaluators.size(); // Main expression is always index 0 - DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES) { + if (valueTypes[0] != DataType.BYTES) { if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { case INT: @@ -718,17 +712,14 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol break; default: throw new IllegalStateException( - "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " - + valueTypes[0]); + "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } } else { - // Logical BYTES values contain serialized ThetaSketch objects. - // They always use the single-value representation. + // Logical BYTES stores serialized ThetaSketch objects in the single-value representation. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); for (int i = 0; i < length; i++) { - List thetaSketchAccumulators = - getUnions(groupByResultHolder, groupKeyArray[i]); + List thetaSketchAccumulators = getUnions(groupByResultHolder, groupKeyArray[i]); ThetaSketch sketch = sketches[i]; if (_includeDefaultSketch) { thetaSketchAccumulators.get(0).apply(sketch); @@ -753,8 +744,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult int numFilters = _filterEvaluators.size(); // Main expression is always index 0 - DataType dataType = valueTypes[0]; - if (dataType != DataType.BYTES) { + if (valueTypes[0] != DataType.BYTES) { if (singleValues[0]) { switch (valueTypes[0].getStoredType()) { case INT: @@ -1043,13 +1033,11 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult break; default: throw new IllegalStateException( - "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " - + valueTypes[0]); + "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } } else { - // Logical BYTES values contain serialized ThetaSketch objects. - // They always use the single-value representation. + // Logical BYTES stores serialized ThetaSketch objects in the single-value representation. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { @@ -1062,8 +1050,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult FilterEvaluator filterEvaluator = _filterEvaluators.get(i); for (int j = 0; j < length; j++) { if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { - for (int groupKey : groupKeysArray[j]) { - getUnions(groupByResultHolder, groupKey).get(i + 1).apply(sketches[j]); + for (int groupKey : groupKeysArray[i]) { + getUnions(groupByResultHolder, groupKey).get(i + 1).apply(sketches[i]); } } } @@ -1489,16 +1477,11 @@ private List buildUnions() { private ThetaSketch[] deserializeSketches(byte[][] serializedSketches, int length) { ThetaSketch[] sketches = new ThetaSketch[length]; for (int i = 0; i < length; i++) { - sketches[i] = deserializeSketch(serializedSketches[i]); + sketches[i] = ThetaSketch.wrap(MemorySegment.ofArray(serializedSketches[i]).asReadOnly()); } return sketches; } - private ThetaSketch deserializeSketch(byte[] serializedSketch) { - return serializedSketch.length == 0 ? _emptySketch - : ThetaSketch.wrap(MemorySegment.ofArray(serializedSketch).asReadOnly()); - } - /// Evaluates the post-aggregation expression. protected ThetaSketch evaluatePostAggregationExpression(List sketches) { return evaluatePostAggregationExpression(_postAggregationExpression, sketches); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java index 4358dac24aa6..7ef3142ddb55 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -84,8 +84,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized UltraLogLog objects. - // They always use the single-value representation. + // Logical BYTES is a serialized UltraLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { UltraLogLog ull = aggregationResultHolder.getResult(); @@ -166,8 +165,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized UltraLogLog objects. - // They always use the single-value representation. + // Logical BYTES is a serialized UltraLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -255,8 +253,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult DataType dataType = blockValSet.getValueType(); if (dataType == DataType.BYTES) { - // Logical BYTES values contain serialized UltraLogLog objects. - // They always use the single-value representation. + // Logical BYTES is a serialized UltraLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java index f41794fb7420..8d227df72bb7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java @@ -264,35 +264,35 @@ private ResultTable toResultTableWithOrderBy() { } int numValues = sortedValues.length; assert numValues <= _limit; - ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); List rows; if (_hasNull) { if (numValues == _limit) { rows = new ArrayList<>(_limit); if (_orderByExpression.isNullsLast()) { - addRows(columnDataType, sortedValues, numValues, rows); + addRows(sortedValues, numValues, rows); } else { rows.add(new Object[]{null}); - addRows(columnDataType, sortedValues, numValues - 1, rows); + addRows(sortedValues, numValues - 1, rows); } } else { rows = new ArrayList<>(numValues + 1); if (_orderByExpression.isNullsLast()) { - addRows(columnDataType, sortedValues, numValues, rows); + addRows(sortedValues, numValues, rows); rows.add(new Object[]{null}); } else { rows.add(new Object[]{null}); - addRows(columnDataType, sortedValues, numValues, rows); + addRows(sortedValues, numValues, rows); } } } else { rows = new ArrayList<>(numValues); - addRows(columnDataType, sortedValues, numValues, rows); + addRows(sortedValues, numValues, rows); } return new ResultTable(_dataSchema, rows); } - private static void addRows(ColumnDataType columnDataType, ByteArray[] values, int length, List rows) { + private void addRows(ByteArray[] values, int length, List rows) { + ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); for (int i = 0; i < length; i++) { rows.add(new Object[]{columnDataType.convertAndFormat(values[i])}); } @@ -301,20 +301,20 @@ private static void addRows(ColumnDataType columnDataType, ByteArray[] values, i private ResultTable toResultTableWithoutOrderBy() { int numValues = _valueSet.size(); assert numValues <= _limit; - ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); List rows; if (_hasNull && numValues < _limit) { rows = new ArrayList<>(numValues + 1); - addRows(columnDataType, _valueSet, rows); + addRows(_valueSet, rows); rows.add(new Object[]{null}); } else { rows = new ArrayList<>(numValues); - addRows(columnDataType, _valueSet, rows); + addRows(_valueSet, rows); } return new ResultTable(_dataSchema, rows); } - private static void addRows(ColumnDataType columnDataType, HashSet values, List rows) { + private void addRows(HashSet values, List rows) { + ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); for (ByteArray value : values) { rows.add(new Object[]{columnDataType.convertAndFormat(value)}); } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java index be757ecd9770..570cc0704119 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java @@ -532,11 +532,6 @@ private Object getConvertedKey(DataTable dataTable, ColumnDataType columnDataTyp case BYTES: return dataTable.getBytes(rowId, colId).getBytes(); case UUID: - // Deliberately delegated to ColumnDataType#convert rather than falling through to BYTES. The other reduce - // path (reduceWithIndexedTable) converts group keys with exactly that method, and UUID is the one type - // whose converted form is not its stored bytes -- it yields a java.util.UUID. PredicateRowMatcher casts - // directly on that, so returning the raw byte[] here makes GROUP BY ... HAVING on a UUID column throw - // ClassCastException. Delegating keeps the two paths identical by construction. return columnDataType.convert(dataTable.getBytes(rowId, colId)); default: throw new IllegalStateException("Illegal column data type in group key: " + columnDataType); diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java index 20f93011e422..29041933af2f 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtilsTest.java @@ -18,35 +18,21 @@ */ package org.apache.pinot.core.query.aggregation.function; -import com.clearspring.analytics.stream.cardinality.HyperLogLog; -import java.io.IOException; -import java.util.Arrays; -import java.util.List; -import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.core.common.ObjectSerDeUtils; import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.datasource.DataSource; -import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.reader.Dictionary; -import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.apache.pinot.spi.utils.UuidUtils; import org.testng.annotations.Test; -import static org.mockito.Mockito.anyInt; import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; import static org.mockito.Mockito.when; import static org.testng.Assert.assertEquals; import static org.testng.Assert.assertThrows; -import static org.testng.Assert.assertTrue; /// Unit test for {@link AggregationFunctionUtils#getAggregationResult}, the metadata/dictionary based aggregation /// result resolver used by the non-scan based and partial metadata based aggregation paths. @SuppressWarnings("rawtypes") public class AggregationFunctionUtilsTest { - private static final ExpressionContext INPUT_EXPRESSION = ExpressionContext.forIdentifier("inputCol"); private static AggregationFunction mockFunction(AggregationFunctionType type) { AggregationFunction aggregationFunction = mock(AggregationFunction.class); @@ -95,80 +81,4 @@ public void testNonCountWithNullDataSourceThrows() { () -> AggregationFunctionUtils.getAggregationResult(mockFunction(AggregationFunctionType.MIN), null, 100, "TEST")); } - - @Test - public void testDistinctCountHllMvOffersUuidDictionaryBytesAsRawValues() - throws IOException { - DistinctCountHLLMVAggregationFunction function = - new DistinctCountHLLMVAggregationFunction(List.of(INPUT_EXPRESSION)); - byte[][] bytesValues = { - UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"), - UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"), - UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002") - }; - Dictionary dictionary = mock(Dictionary.class); - when(dictionary.length()).thenReturn(bytesValues.length); - for (int i = 0; i < bytesValues.length; i++) { - when(dictionary.get(i)).thenReturn(bytesValues[i]); - } - DataSource dataSource = mockDataSource(dictionary, DataType.UUID, false); - - HyperLogLog result = (HyperLogLog) AggregationFunctionUtils.getAggregationResult(function, dataSource, 3, "TEST"); - - HyperLogLog expected = new HyperLogLog(function.getLog2m()); - for (byte[] value : bytesValues) { - expected.offer(value); - } - assertEquals(result.cardinality(), 3L); - assertTrue(Arrays.equals(result.getBytes(), expected.getBytes())); - for (int i = 0; i < bytesValues.length; i++) { - verify(dictionary).get(i); - } - verify(dictionary, never()).getBytesValue(anyInt()); - } - - @Test - public void testDistinctCountHllMergesLogicalBytesAsSerializedState() - throws IOException { - DistinctCountHLLAggregationFunction function = - new DistinctCountHLLAggregationFunction(List.of(INPUT_EXPRESSION)); - byte[] firstSketch = serializedHll(function, "a", "b"); - byte[] secondSketch = serializedHll(function, "b", "c"); - Dictionary dictionary = mock(Dictionary.class); - when(dictionary.length()).thenReturn(2); - when(dictionary.getBytesValue(0)).thenReturn(firstSketch); - when(dictionary.getBytesValue(1)).thenReturn(secondSketch); - // Logical BYTES determines serialized-state handling independently of the cardinality metadata. - DataSource dataSource = mockDataSource(dictionary, DataType.BYTES, false); - - HyperLogLog result = (HyperLogLog) AggregationFunctionUtils.getAggregationResult(function, dataSource, 2, "TEST"); - - HyperLogLog expected = new HyperLogLog(function.getLog2m()); - expected.offer("a"); - expected.offer("b"); - expected.offer("c"); - assertEquals(result.cardinality(), 3L); - assertTrue(Arrays.equals(result.getBytes(), expected.getBytes())); - verify(dictionary).getBytesValue(0); - verify(dictionary).getBytesValue(1); - verify(dictionary, never()).get(anyInt()); - } - - private static DataSource mockDataSource(Dictionary dictionary, DataType dataType, boolean singleValue) { - DataSourceMetadata metadata = mock(DataSourceMetadata.class); - when(metadata.getDataType()).thenReturn(dataType); - when(metadata.isSingleValue()).thenReturn(singleValue); - DataSource dataSource = mock(DataSource.class); - when(dataSource.getDictionary()).thenReturn(dictionary); - when(dataSource.getDataSourceMetadata()).thenReturn(metadata); - return dataSource; - } - - private static byte[] serializedHll(DistinctCountHLLAggregationFunction function, String... values) { - HyperLogLog hll = new HyperLogLog(function.getLog2m()); - for (String value : values) { - hll.offer(value); - } - return ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.serialize(hll); - } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java index 9c1f4fda78d7..ec4fccf785f5 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapMVAggregationFunctionTest.java @@ -18,33 +18,13 @@ */ package org.apache.pinot.core.query.aggregation.function; -import java.util.List; -import java.util.Map; -import org.apache.pinot.common.request.context.ExpressionContext; -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.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; -import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; -import org.apache.pinot.spi.utils.UuidUtils; -import org.testng.Assert; import org.testng.annotations.Test; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - public class DistinctCountBitmapMVAggregationFunctionTest extends AbstractAggregationFunctionTest { - private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); - private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); - private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); - private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); - private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); @Test public void testAggregationMV() { @@ -112,50 +92,4 @@ public void testAggregationMVGroupByMV() { "tag1 | 3", // distinct: 1, 2, 3 "tag2 | 3"); // distinct: 1, 2, 3 } - - @Test - public void testMultiValueUuidUsesMultiValueAccessor() { - byte[][][] uuidValues = { - {UUID_0, UUID_1}, - {UUID_1, UUID_2}, - {UUID_0}, - {UUID_3} - }; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(uuidValues); - DistinctCountBitmapAggregationFunction function = - new DistinctCountBitmapAggregationFunction(List.of(UUID_EXPRESSION)); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4); - - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); - } - - private static int extractFinalResult(DistinctCountBitmapAggregationFunction function, - AggregationResultHolder resultHolder) { - return function.extractFinalResult(function.extractAggregationResult(resultHolder)); - } - - private static int extractFinalResult(DistinctCountBitmapAggregationFunction function, - GroupByResultHolder resultHolder, int groupKey) { - return function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey)); - } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java index 4d25ed21341f..a3bcda6204da 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunctionTest.java @@ -26,40 +26,15 @@ import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; 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.groupby.GroupByResultHolder; import org.apache.pinot.segment.local.customobject.CpcSketchAccumulator; import org.apache.pinot.segment.local.customobject.SerializedCPCSketch; import org.apache.pinot.segment.spi.Constants; -import org.apache.pinot.segment.spi.index.reader.Dictionary; -import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; -import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - public class DistinctCountCPCSketchAggregationFunctionTest { - private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); - private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); - private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); - private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); - private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); - private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); - private static final byte[][] UUID_VALUES_SV = {UUID_0, UUID_1, UUID_0, UUID_2}; - private static final byte[][][] UUID_VALUES_MV = {{UUID_0, UUID_1}, {UUID_1, UUID_2}, {UUID_0}, {UUID_3}}; - - @DataProvider(name = "uuidValueModes") - public static Object[][] uuidValueModes() { - return new Object[][]{{true}, {false}}; - } @Test public void testCanUseStarTreeDefaultLgK() { @@ -163,181 +138,4 @@ public void testMergeWithEmptyAccumulators() { result = function.merge(new CpcSketchAccumulator(12, 2), new CpcSketchAccumulator(12, 2)); Assert.assertTrue(result.isEmpty()); } - - @Test(dataProvider = "uuidValueModes") - public void testAggregateUuid(boolean singleValue) { - DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); - BlockValSet blockValSet = mockUuidBlockValSet(singleValue); - AggregationResultHolder resultHolder = function.createAggregationResultHolder(); - - function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder), singleValue ? 3L : 4L); - verifyBytesAccessor(blockValSet, singleValue); - } - - @Test(dataProvider = "uuidValueModes") - public void testAggregateUuidGroupBySV(boolean singleValue) { - DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); - BlockValSet blockValSet = mockUuidBlockValSet(singleValue); - GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); - - function.aggregateGroupBySV(UUID_VALUES_SV.length, new int[]{0, 0, 1, 1}, resultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 2L : 3L); - Assert.assertEquals(extractFinalResult(function, resultHolder, 1), 2L); - verifyBytesAccessor(blockValSet, singleValue); - } - - @Test(dataProvider = "uuidValueModes") - public void testAggregateUuidGroupByMV(boolean singleValue) { - DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); - BlockValSet blockValSet = mockUuidBlockValSet(singleValue); - GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); - - function.aggregateGroupByMV(UUID_VALUES_SV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, resultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 1L : 2L); - Assert.assertEquals(extractFinalResult(function, resultHolder, 1), singleValue ? 3L : 4L); - verifyBytesAccessor(blockValSet, singleValue); - } - - @Test - public void testAggregateDictionaryEncodedUuid() { - DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); - Dictionary dictionary = mock(Dictionary.class); - when(dictionary.get(0)).thenReturn(UUID_0); - when(dictionary.get(1)).thenReturn(UUID_1); - when(dictionary.get(2)).thenReturn(UUID_2); - - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isSingleValue()).thenReturn(true); - when(blockValSet.isDictionaryEncoded()).thenReturn(true); - when(blockValSet.getDictionary()).thenReturn(dictionary); - when(blockValSet.getDictionaryIdsSV()).thenReturn(new int[]{0, 1, 0, 2}); - AggregationResultHolder resultHolder = function.createAggregationResultHolder(); - - function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder), 3L); - } - - @Test - public void testLogicalBytesUsesSerializedSingleValueAccessorBeforeDispatch() { - byte[][] serializedSketches = { - serializedSketch("a", "b"), - serializedSketch("b", "c"), - new byte[0], - serializedSketch("d") - }; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); - when(blockValSet.getBytesValuesSV()).thenReturn(serializedSketches); - DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(BYTES_EXPRESSION)); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 1L); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); - verify(blockValSet, atLeastOnce()).getBytesValuesSV(); - verify(blockValSet, never()).getBytesValuesMV(); - verify(blockValSet, never()).isSingleValue(); - } - - @Test - public void testAggregateDictionaryEncodedMultiValueUuid() { - Dictionary dictionary = mock(Dictionary.class); - when(dictionary.get(0)).thenReturn(UUID_0); - when(dictionary.get(1)).thenReturn(UUID_1); - when(dictionary.get(2)).thenReturn(UUID_2); - when(dictionary.get(3)).thenReturn(UUID_3); - - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.isDictionaryEncoded()).thenReturn(true); - when(blockValSet.getDictionary()).thenReturn(dictionary); - when(blockValSet.getDictionaryIdsMV()).thenReturn(new int[][]{{0, 1}, {1, 2}, {0}, {3}}); - DistinctCountCPCSketchAggregationFunction function = - new DistinctCountCPCSketchAggregationFunction(List.of(UUID_EXPRESSION)); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(UUID_VALUES_MV.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(UUID_VALUES_MV.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(UUID_VALUES_MV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); - verify(blockValSet, atLeastOnce()).getDictionaryIdsMV(); - verify(blockValSet, never()).getBytesValuesSV(); - verify(blockValSet, never()).getBytesValuesMV(); - } - - private static BlockValSet mockUuidBlockValSet(boolean singleValue) { - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isSingleValue()).thenReturn(singleValue); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); - if (singleValue) { - when(blockValSet.getBytesValuesSV()).thenReturn(UUID_VALUES_SV); - } else { - when(blockValSet.getBytesValuesMV()).thenReturn(UUID_VALUES_MV); - } - return blockValSet; - } - - private static void verifyBytesAccessor(BlockValSet blockValSet, boolean singleValue) { - if (singleValue) { - verify(blockValSet, atLeastOnce()).getBytesValuesSV(); - verify(blockValSet, never()).getBytesValuesMV(); - } else { - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); - } - } - - private static byte[] serializedSketch(String... values) { - CpcSketch sketch = new CpcSketch(); - for (String value : values) { - sketch.update(value); - } - return sketch.toByteArray(); - } - - private static long extractFinalResult(DistinctCountCPCSketchAggregationFunction function, - AggregationResultHolder resultHolder) { - return ((Number) function.extractFinalResult(function.extractAggregationResult(resultHolder))).longValue(); - } - - private static long extractFinalResult(DistinctCountCPCSketchAggregationFunction function, - GroupByResultHolder resultHolder, int groupKey) { - return ((Number) function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey))).longValue(); - } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java index 8dc9ddc0699c..2520115affa0 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunctionTest.java @@ -19,22 +19,14 @@ package org.apache.pinot.core.query.aggregation.function; import com.clearspring.analytics.stream.cardinality.HyperLogLog; -import java.util.Arrays; import java.util.BitSet; -import java.util.HashMap; import java.util.List; import java.util.Map; import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.request.context.RequestContextUtils; -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.segment.spi.Constants; import org.apache.pinot.segment.spi.index.reader.Dictionary; -import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.apache.pinot.spi.utils.CommonConstants; -import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; @@ -63,107 +55,6 @@ public void testCanUseStarTreeDefaultLog2m() { Assert.assertFalse(function.canUseStarTree(Map.of(Constants.HLL_LOG2M_KEY, "16"))); } - /// Regression: UUID columns have storedType=BYTES, but a UUID value is a logical scalar, not a serialized - /// HyperLogLog. The aggregator must offer the stored bytes as values instead of trying to deserialize each - /// 16-byte value as an HLL. - @Test - public void testAggregateOnUuidColumnOffersStoredBytesAndProducesExactDistinctCount() { - ExpressionContext expression = RequestContextUtils.getExpression("uuidCol"); - DistinctCountHLLAggregationFunction function = new DistinctCountHLLAggregationFunction(List.of(expression)); - - // Three distinct UUIDs across six rows; the same UUID repeats twice on rows 0/3, 1/4, 2/5. - String[] uuidStrings = new String[]{ - "550e8400-e29b-41d4-a716-446655440000", - "12345678-1234-1234-1234-1234567890ab", - "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12", - "550e8400-e29b-41d4-a716-446655440000", - "12345678-1234-1234-1234-1234567890ab", - "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12" - }; - - // Stub the BYTES fetch with the raw 16-byte stored values offered by the production path. - byte[][] uuidBytes = new byte[uuidStrings.length][]; - for (int i = 0; i < uuidStrings.length; i++) { - uuidBytes[i] = UuidUtils.toBytes(uuidStrings[i]); - } - BlockValSet uuidBlockValSet = mock(BlockValSet.class); - when(uuidBlockValSet.getValueType()).thenReturn(DataType.UUID); - when(uuidBlockValSet.getBytesValuesSV()).thenReturn(uuidBytes); - when(uuidBlockValSet.isSingleValue()).thenReturn(true); - when(uuidBlockValSet.getDictionary()).thenReturn(null); - - Map blockValSetMap = new HashMap<>(); - blockValSetMap.put(expression, uuidBlockValSet); - - AggregationResultHolder resultHolder = function.createAggregationResultHolder(); - function.aggregate(uuidStrings.length, resultHolder, blockValSetMap); - - Object intermediate = function.extractAggregationResult(resultHolder); - Assert.assertTrue(intermediate instanceof HyperLogLog, - "Intermediate result must be a HyperLogLog, not a dictionary bitmap"); - long cardinality = ((HyperLogLog) intermediate).cardinality(); - Assert.assertEquals(cardinality, 3L, - "HLL cardinality must equal the 3 distinct UUIDs; got " + cardinality); - } - - /// UUID columns hash their **stored bytes**, exactly as TIMESTAMP hashes its stored millis rather than a - /// formatted string. Consequence: DISTINCTCOUNTHLL(uuidCol) does NOT equal - /// DISTINCTCOUNTHLL(CAST(uuidCol AS STRING)) -- and neither does it for TIMESTAMP, so this is the consistent - /// behaviour for a logical type, not a gap. Pinned here so nobody "fixes" it back into a canonical-string - /// rendering, which would reintroduce a per-row String allocation in the aggregation loop. - @Test - public void testUuidDistinctCountHllHashesStoredBytesNotCanonicalString() - throws java.io.IOException { - String[] uuidStrings = new String[]{ - "550e8400-e29b-41d4-a716-446655440000", - "12345678-1234-1234-1234-1234567890ab", - "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12" - }; - - // Cardinality is still exact for a small distinct set... - Assert.assertEquals(computeHllCardinality(uuidStrings, DataType.UUID), 3L); - - // ...but the sketch is built over the 16 stored bytes, so a HyperLogLog fed the canonical strings differs. - HyperLogLog fromCanonicalStrings = new HyperLogLog(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M); - for (String uuid : uuidStrings) { - fromCanonicalStrings.offer(uuid); - } - HyperLogLog fromStoredBytes = new HyperLogLog(CommonConstants.Helix.DEFAULT_HYPERLOGLOG_LOG2M); - for (String uuid : uuidStrings) { - fromStoredBytes.offer(UuidUtils.toBytes(uuid)); - } - Assert.assertFalse(Arrays.equals(fromCanonicalStrings.getBytes(), fromStoredBytes.getBytes()), - "stored-bytes and canonical-string sketches are expected to differ"); - } - - private long computeHllCardinality(String[] values, DataType valueType) { - ExpressionContext expression = RequestContextUtils.getExpression("col"); - DistinctCountHLLAggregationFunction function = new DistinctCountHLLAggregationFunction(List.of(expression)); - - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(valueType); - if (valueType == DataType.UUID) { - // UUID path fetches and offers the raw stored bytes. - byte[][] uuidBytes = new byte[values.length][]; - for (int i = 0; i < values.length; i++) { - uuidBytes[i] = UuidUtils.toBytes(values[i]); - } - when(blockValSet.getBytesValuesSV()).thenReturn(uuidBytes); - } else { - when(blockValSet.getStringValuesSV()).thenReturn(values); - } - when(blockValSet.isSingleValue()).thenReturn(true); - when(blockValSet.getDictionary()).thenReturn(null); - - Map blockValSetMap = new HashMap<>(); - blockValSetMap.put(expression, blockValSet); - - AggregationResultHolder resultHolder = function.createAggregationResultHolder(); - function.aggregate(values.length, resultHolder, blockValSetMap); - Object intermediate = function.extractAggregationResult(resultHolder); - return ((HyperLogLog) intermediate).cardinality(); - } - @Test public void testCanUseStarTreeCustomLog2m() { DistinctCountHLLAggregationFunction function = new DistinctCountHLLAggregationFunction( diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java index 2ec9960741d2..fe24e95b946e 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLMVAggregationFunctionTest.java @@ -18,33 +18,13 @@ */ package org.apache.pinot.core.query.aggregation.function; -import java.util.List; -import java.util.Map; -import org.apache.pinot.common.request.context.ExpressionContext; -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.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; -import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; -import org.apache.pinot.spi.utils.UuidUtils; -import org.testng.Assert; import org.testng.annotations.Test; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - public class DistinctCountHLLMVAggregationFunctionTest extends AbstractAggregationFunctionTest { - private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); - private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); - private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); - private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); - private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); @Test public void testAggregationMV() { @@ -112,50 +92,4 @@ public void testAggregationMVGroupByMV() { "tag1 | 3", // distinct: 1, 2, 3 "tag2 | 3"); // distinct: 1, 2, 3 } - - @Test - public void testMultiValueUuidUsesMultiValueAccessor() { - DistinctCountHLLAggregationFunction function = - new DistinctCountHLLAggregationFunction(List.of(UUID_EXPRESSION)); - byte[][][] uuidValues = { - {UUID_0, UUID_1}, - {UUID_1, UUID_2}, - {UUID_0}, - {UUID_3} - }; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(uuidValues); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); - - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); - } - - private static long extractFinalResult(DistinctCountHLLAggregationFunction function, - AggregationResultHolder resultHolder) { - return function.extractFinalResult(function.extractAggregationResult(resultHolder)); - } - - private static long extractFinalResult(DistinctCountHLLAggregationFunction function, - GroupByResultHolder resultHolder, int groupKey) { - return function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey)); - } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java index 648451f57e51..9a2ca5f86e79 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusMVAggregationFunctionTest.java @@ -18,33 +18,13 @@ */ package org.apache.pinot.core.query.aggregation.function; -import java.util.List; -import java.util.Map; -import org.apache.pinot.common.request.context.ExpressionContext; -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.queries.FluentQueryTest; import org.apache.pinot.spi.data.FieldSpec; -import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.Schema; -import org.apache.pinot.spi.utils.UuidUtils; -import org.testng.Assert; import org.testng.annotations.Test; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - public class DistinctCountHLLPlusMVAggregationFunctionTest extends AbstractAggregationFunctionTest { - private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); - private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); - private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); - private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); - private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); @Test public void testAggregationMV() { @@ -112,50 +92,4 @@ public void testAggregationMVGroupByMV() { "tag1 | 3", // distinct: 1, 2, 3 "tag2 | 3"); // distinct: 1, 2, 3 } - - @Test - public void testMultiValueUuidUsesMultiValueAccessor() { - DistinctCountHLLPlusAggregationFunction function = - new DistinctCountHLLPlusAggregationFunction(List.of(UUID_EXPRESSION)); - byte[][][] uuidValues = { - {UUID_0, UUID_1}, - {UUID_1, UUID_2}, - {UUID_0}, - {UUID_3} - }; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isSingleValue()).thenReturn(false); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesMV()).thenReturn(uuidValues); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 4L); - - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); - } - - private static long extractFinalResult(DistinctCountHLLPlusAggregationFunction function, - AggregationResultHolder resultHolder) { - return function.extractFinalResult(function.extractAggregationResult(resultHolder)); - } - - private static long extractFinalResult(DistinctCountHLLPlusAggregationFunction function, - GroupByResultHolder resultHolder, int groupKey) { - return function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey)); - } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java index 2ced8f05ff4f..85eac01d28b7 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunctionTest.java @@ -20,41 +20,14 @@ import java.util.List; import java.util.Map; -import org.apache.datasketches.theta.UpdatableThetaSketch; -import org.apache.datasketches.theta.UpdatableThetaSketchBuilder; import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; -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.segment.spi.Constants; -import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; -import org.testng.annotations.DataProvider; import org.testng.annotations.Test; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - public class DistinctCountThetaSketchAggregationFunctionTest { - private static final ExpressionContext UUID_EXPRESSION = ExpressionContext.forIdentifier("uuidCol"); - private static final ExpressionContext BYTES_EXPRESSION = ExpressionContext.forIdentifier("bytesCol"); - private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); - private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); - private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); - private static final byte[] UUID_3 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440003"); - private static final byte[][] UUID_VALUES_SV = {UUID_0, UUID_1, UUID_0, UUID_2}; - private static final byte[][][] UUID_VALUES_MV = {{UUID_0, UUID_1}, {UUID_1, UUID_2}, {UUID_0}, {UUID_3}}; - - @DataProvider(name = "uuidValueModes") - public static Object[][] uuidValueModes() { - return new Object[][]{{true}, {false}}; - } @Test public void testCanUseStarTreeDefaultK() { @@ -81,135 +54,4 @@ public void testCanUseCustomK() { Assert.assertTrue(function.canUseStarTree(Map.of(Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES, 32768))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.THETA_TUPLE_SKETCH_NOMINAL_ENTRIES, "32768"))); } - - @Test(dataProvider = "uuidValueModes") - public void testAggregateUuid(boolean singleValue) { - DistinctCountThetaSketchAggregationFunction function = - new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); - BlockValSet blockValSet = mockUuidBlockValSet(singleValue); - AggregationResultHolder resultHolder = function.createAggregationResultHolder(); - - function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder), singleValue ? 3L : 4L); - verifyBytesAccessor(blockValSet, singleValue); - } - - @Test(dataProvider = "uuidValueModes") - public void testAggregateUuidGroupBySV(boolean singleValue) { - DistinctCountThetaSketchAggregationFunction function = - new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); - BlockValSet blockValSet = mockUuidBlockValSet(singleValue); - GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); - - function.aggregateGroupBySV(UUID_VALUES_SV.length, new int[]{0, 0, 1, 1}, resultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 2L : 3L); - Assert.assertEquals(extractFinalResult(function, resultHolder, 1), 2L); - verifyBytesAccessor(blockValSet, singleValue); - } - - @Test(dataProvider = "uuidValueModes") - public void testAggregateUuidGroupByMV(boolean singleValue) { - DistinctCountThetaSketchAggregationFunction function = - new DistinctCountThetaSketchAggregationFunction(List.of(UUID_EXPRESSION)); - BlockValSet blockValSet = mockUuidBlockValSet(singleValue); - GroupByResultHolder resultHolder = function.createGroupByResultHolder(2, 2); - - function.aggregateGroupByMV(UUID_VALUES_SV.length, new int[][]{{0}, {1}, {0, 1}, {1}}, resultHolder, - Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder, 0), singleValue ? 1L : 2L); - Assert.assertEquals(extractFinalResult(function, resultHolder, 1), singleValue ? 3L : 4L); - verifyBytesAccessor(blockValSet, singleValue); - } - - @Test - public void testUuidPredicateUsesLogicalType() { - DistinctCountThetaSketchAggregationFunction function = new DistinctCountThetaSketchAggregationFunction( - List.of(UUID_EXPRESSION, ExpressionContext.forLiteral(Literal.stringValue("")), - ExpressionContext.forLiteral(Literal.stringValue("uuidCol = '550e8400-e29b-41d4-a716-446655440000'")), - ExpressionContext.forLiteral(Literal.stringValue("$1")))); - BlockValSet blockValSet = mockUuidBlockValSet(true); - AggregationResultHolder resultHolder = function.createAggregationResultHolder(); - - function.aggregate(UUID_VALUES_SV.length, resultHolder, Map.of(UUID_EXPRESSION, blockValSet)); - - Assert.assertEquals(extractFinalResult(function, resultHolder), 1L); - } - - @Test - public void testAggregateSingleValueSerializedSketches() { - byte[][] serializedSketches = { - serializedSketch("a", "b"), - serializedSketch("b", "c"), - new byte[0], - serializedSketch("d") - }; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); - when(blockValSet.isSingleValue()).thenReturn(true); - when(blockValSet.getBytesValuesSV()).thenReturn(serializedSketches); - DistinctCountThetaSketchAggregationFunction function = - new DistinctCountThetaSketchAggregationFunction(List.of(BYTES_EXPRESSION)); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 4L); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 3L); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 1L); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(BYTES_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 2L); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); - verify(blockValSet, atLeastOnce()).getBytesValuesSV(); - verify(blockValSet, never()).getBytesValuesMV(); - } - - private static BlockValSet mockUuidBlockValSet(boolean singleValue) { - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isSingleValue()).thenReturn(singleValue); - if (singleValue) { - when(blockValSet.getBytesValuesSV()).thenReturn(UUID_VALUES_SV); - } else { - when(blockValSet.getBytesValuesMV()).thenReturn(UUID_VALUES_MV); - } - return blockValSet; - } - - private static void verifyBytesAccessor(BlockValSet blockValSet, boolean singleValue) { - if (singleValue) { - verify(blockValSet, atLeastOnce()).getBytesValuesSV(); - verify(blockValSet, never()).getBytesValuesMV(); - } else { - verify(blockValSet, atLeastOnce()).getBytesValuesMV(); - verify(blockValSet, never()).getBytesValuesSV(); - } - } - - private static byte[] serializedSketch(String... values) { - UpdatableThetaSketch sketch = new UpdatableThetaSketchBuilder().build(); - for (String value : values) { - sketch.update(value); - } - return sketch.compact().toByteArray(); - } - - private static long extractFinalResult(DistinctCountThetaSketchAggregationFunction function, - AggregationResultHolder resultHolder) { - return ((Number) function.extractFinalResult(function.extractAggregationResult(resultHolder))).longValue(); - } - - private static long extractFinalResult(DistinctCountThetaSketchAggregationFunction function, - GroupByResultHolder resultHolder, int groupKey) { - return ((Number) function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey))).longValue(); - } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java index e6a6fb0e5717..06c05b366bdc 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunctionTest.java @@ -18,34 +18,16 @@ */ package org.apache.pinot.core.query.aggregation.function; -import com.dynatrace.hash4j.distinctcount.UltraLogLog; import java.util.List; import java.util.Map; import org.apache.pinot.common.request.Literal; import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.core.common.BlockValSet; -import org.apache.pinot.core.common.ObjectSerDeUtils; -import org.apache.pinot.core.query.aggregation.AggregationResultHolder; -import org.apache.pinot.core.query.aggregation.groupby.GroupByResultHolder; -import org.apache.pinot.segment.local.utils.UltraLogLogUtils; import org.apache.pinot.segment.spi.Constants; -import org.apache.pinot.spi.data.FieldSpec.DataType; -import org.apache.pinot.spi.utils.UuidUtils; import org.testng.Assert; import org.testng.annotations.Test; -import static org.mockito.Mockito.atLeastOnce; -import static org.mockito.Mockito.mock; -import static org.mockito.Mockito.never; -import static org.mockito.Mockito.verify; -import static org.mockito.Mockito.when; - public class DistinctCountULLAggregationFunctionTest { - private static final ExpressionContext INPUT_EXPRESSION = ExpressionContext.forIdentifier("inputCol"); - private static final byte[] UUID_0 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000"); - private static final byte[] UUID_1 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440001"); - private static final byte[] UUID_2 = UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440002"); @Test public void testCanUseStarTreeDefaultP() { @@ -76,98 +58,4 @@ public void testCanUseStarTreeCustomP() { Assert.assertTrue(function.canUseStarTree(Map.of(Constants.HLLPLUS_ULL_P_KEY, 16))); Assert.assertTrue(function.canUseStarTree(Map.of(Constants.HLLPLUS_ULL_P_KEY, "16"))); } - - @Test - public void testRawUuidUsesStoredBytes() { - DistinctCountULLAggregationFunction function = - new DistinctCountULLAggregationFunction(List.of(INPUT_EXPRESSION)); - byte[][] uuidValues = {UUID_0, UUID_1, UUID_0, UUID_2}; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.UUID); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesSV()).thenReturn(uuidValues); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(uuidValues.length, aggregationResultHolder, Map.of(INPUT_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), 3L); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(uuidValues.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(INPUT_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), 2L); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), 2L); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(uuidValues.length, new int[][]{{0}, {1}, {0, 1}, {1}}, groupByMVResultHolder, - Map.of(INPUT_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), 1L); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), 3L); - - verify(blockValSet, atLeastOnce()).getBytesValuesSV(); - verify(blockValSet, never()).getBytesValuesMV(); - } - - @Test - public void testSerializedBytesUsesStoredSketches() { - DistinctCountULLAggregationFunction function = - new DistinctCountULLAggregationFunction(List.of(INPUT_EXPRESSION)); - byte[][] serializedSketches = { - serializedSketch(function, "a", "b"), - serializedSketch(function, "b", "c"), - serializedSketch(function, "a"), - serializedSketch(function, "d") - }; - BlockValSet blockValSet = mock(BlockValSet.class); - when(blockValSet.getValueType()).thenReturn(DataType.BYTES); - when(blockValSet.isDictionaryEncoded()).thenReturn(false); - when(blockValSet.getBytesValuesSV()).thenReturn(serializedSketches); - - AggregationResultHolder aggregationResultHolder = function.createAggregationResultHolder(); - function.aggregate(serializedSketches.length, aggregationResultHolder, Map.of(INPUT_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, aggregationResultHolder), - referenceEstimate(function, "a", "b", "c", "d")); - - GroupByResultHolder groupBySVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupBySV(serializedSketches.length, new int[]{0, 0, 1, 1}, groupBySVResultHolder, - Map.of(INPUT_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 0), - referenceEstimate(function, "a", "b", "c")); - Assert.assertEquals(extractFinalResult(function, groupBySVResultHolder, 1), referenceEstimate(function, "a", "d")); - - GroupByResultHolder groupByMVResultHolder = function.createGroupByResultHolder(2, 2); - function.aggregateGroupByMV(serializedSketches.length, new int[][]{{0}, {1}, {0, 1}, {1}}, - groupByMVResultHolder, Map.of(INPUT_EXPRESSION, blockValSet)); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 0), referenceEstimate(function, "a", "b")); - Assert.assertEquals(extractFinalResult(function, groupByMVResultHolder, 1), - referenceEstimate(function, "a", "b", "c", "d")); - - verify(blockValSet, atLeastOnce()).getBytesValuesSV(); - verify(blockValSet, never()).getBytesValuesMV(); - } - - private static byte[] serializedSketch(DistinctCountULLAggregationFunction function, String... values) { - UltraLogLog sketch = UltraLogLog.create(function.getP()); - for (String value : values) { - UltraLogLogUtils.hashObject(value).ifPresent(sketch::add); - } - return ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.serialize(sketch); - } - - private static long referenceEstimate(DistinctCountULLAggregationFunction function, String... values) { - UltraLogLog reference = UltraLogLog.create(function.getP()); - for (String value : values) { - UltraLogLogUtils.hashObject(value).ifPresent(reference::add); - } - return Math.round(reference.getDistinctCountEstimate()); - } - - private static long extractFinalResult(DistinctCountULLAggregationFunction function, - AggregationResultHolder resultHolder) { - return ((Number) function.extractFinalResult(function.extractAggregationResult(resultHolder))).longValue(); - } - - private static long extractFinalResult(DistinctCountULLAggregationFunction function, - GroupByResultHolder resultHolder, int groupKey) { - return ((Number) function.extractFinalResult(function.extractGroupByResult(resultHolder, groupKey))).longValue(); - } } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java b/pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java deleted file mode 100644 index 6ca1c2298cb3..000000000000 --- a/pinot-core/src/test/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTableTest.java +++ /dev/null @@ -1,70 +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.distinct.table; - -import org.apache.pinot.common.request.context.ExpressionContext; -import org.apache.pinot.common.request.context.OrderByExpressionContext; -import org.apache.pinot.common.response.broker.ResultTable; -import org.apache.pinot.common.utils.DataSchema; -import org.apache.pinot.common.utils.DataSchema.ColumnDataType; -import org.apache.pinot.spi.utils.ByteArray; -import org.apache.pinot.spi.utils.BytesUtils; -import org.apache.pinot.spi.utils.UuidUtils; -import org.testng.annotations.Test; - -import static org.testng.Assert.assertEquals; - - -/// Tests for [BytesDistinctTable]. -public class BytesDistinctTableTest { - private static final String UUID_COLUMN = "uuidCol"; - private static final String UUID_VALUE_1 = "550e8400-e29b-41d4-a716-446655440000"; - private static final String UUID_VALUE_2 = "550e8400-e29b-41d4-a716-446655440001"; - - @Test - public void testToResultTableFormatsUuidAndBytesWithoutOrderBy() { - BytesDistinctTable uuidTable = new BytesDistinctTable( - new DataSchema(new String[]{UUID_COLUMN}, new ColumnDataType[]{ColumnDataType.UUID}), 10, false, null); - uuidTable.addUnbounded(new ByteArray(UuidUtils.toBytes(UUID_VALUE_1))); - - ResultTable uuidResultTable = uuidTable.toResultTable(); - assertEquals(uuidResultTable.getRows().get(0)[0], UUID_VALUE_1); - - byte[] bytesValue = new byte[]{0x01, 0x23, 0x45}; - BytesDistinctTable bytesTable = new BytesDistinctTable( - new DataSchema(new String[]{"bytesCol"}, new ColumnDataType[]{ColumnDataType.BYTES}), 10, false, null); - bytesTable.addUnbounded(new ByteArray(bytesValue)); - - ResultTable bytesResultTable = bytesTable.toResultTable(); - assertEquals(bytesResultTable.getRows().get(0)[0], BytesUtils.toHexString(bytesValue)); - } - - @Test - public void testToResultTableFormatsUuidWithOrderBy() { - BytesDistinctTable uuidTable = new BytesDistinctTable( - new DataSchema(new String[]{UUID_COLUMN}, new ColumnDataType[]{ColumnDataType.UUID}), 10, false, - new OrderByExpressionContext(ExpressionContext.forIdentifier(UUID_COLUMN), true)); - uuidTable.addUnbounded(new ByteArray(UuidUtils.toBytes(UUID_VALUE_2))); - uuidTable.addUnbounded(new ByteArray(UuidUtils.toBytes(UUID_VALUE_1))); - - ResultTable resultTable = uuidTable.toResultTable(); - assertEquals(resultTable.getRows().get(0)[0], UUID_VALUE_1); - assertEquals(resultTable.getRows().get(1)[0], UUID_VALUE_2); - } -} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java index f52623e48a0b..1b828a499986 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java @@ -20,7 +20,6 @@ import com.fasterxml.jackson.databind.JsonNode; import java.io.File; -import java.util.ArrayList; import java.util.List; import org.apache.avro.file.DataFileWriter; import org.apache.avro.generic.GenericData; @@ -35,25 +34,24 @@ import static org.testng.Assert.assertTrue; -/// End-to-end coverage for aggregating, grouping and de-duplicating a UUID column. -/// -/// These run through a real broker reduce, which is the point: the group-key conversion in -/// `GroupByDataTableReducer#getConvertedKey` is only reachable when the broker reduces a *single* data table, and the -/// unit-level `BaseQueriesTest` harness always reduces two. A `case UUID` there that returns the stored `byte[]` -/// instead of the converted `java.util.UUID` makes `GROUP BY ... HAVING` over a UUID column fail with -/// `ClassCastException: class [B cannot be cast to class java.util.UUID`, and only a query-level test catches it. +/// End-to-end UUID aggregation coverage over dictionary-encoded and raw SV/MV columns. @Test(suiteName = "CustomClusterIntegrationTest") public class UuidAggregationTest extends CustomDataQueryClusterIntegrationTest { private static final String TABLE_NAME = "UuidAggregationTest"; - private static final String UUID_COLUMN = "uuidColumn"; + private static final String UUID_DICT_SV_COLUMN = "uuidDictSv"; + private static final String UUID_DICT_MV_COLUMN = "uuidDictMv"; + private static final String UUID_RAW_SV_COLUMN = "uuidRawSv"; + private static final String UUID_RAW_MV_COLUMN = "uuidRawMv"; + private static final String UUID_0 = "550e8400-e29b-41d4-a716-446655440000"; private static final String UUID_0_HEX = "550e8400e29b41d4a716446655440000"; private static final String UUID_1 = "550e8400-e29b-41d4-a716-446655440001"; private static final String UUID_2 = "550e8400-e29b-41d4-a716-446655440002"; + private static final String UUID_3 = "550e8400-e29b-41d4-a716-446655440003"; - /// `UUID_0` appears twice so grouping and distinct are distinguishable from a plain row count. - private static final List ROWS = List.of(UUID_0, UUID_0, UUID_1, UUID_2); - private static final int NUM_DISTINCT = 3; + private static final List UUID_SV_VALUES = List.of(UUID_0, UUID_0, UUID_1, UUID_2); + private static final List> UUID_MV_VALUES = + List.of(List.of(UUID_0, UUID_1), List.of(UUID_1, UUID_2), List.of(UUID_0), List.of(UUID_3)); @Override public String getTableName() { @@ -62,7 +60,7 @@ public String getTableName() { @Override protected long getCountStarResult() { - return ROWS.size(); + return UUID_SV_VALUES.size(); } @Override @@ -72,28 +70,41 @@ public int getNumAvroFiles() { @Override public TableConfig createOfflineTableConfig() { - return new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName()).build(); + return new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName()) + .setNoDictionaryColumns(List.of(UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN)).build(); } @Override public Schema createSchema() { return new Schema.SchemaBuilder().setSchemaName(getTableName()) - .addSingleValueDimension(UUID_COLUMN, DataType.UUID) + .addSingleValueDimension(UUID_DICT_SV_COLUMN, DataType.UUID) + .addMultiValueDimension(UUID_DICT_MV_COLUMN, DataType.UUID) + .addSingleValueDimension(UUID_RAW_SV_COLUMN, DataType.UUID) + .addMultiValueDimension(UUID_RAW_MV_COLUMN, DataType.UUID) .build(); } @Override public List createAvroFiles() throws Exception { + org.apache.avro.Schema uuidSchema = org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING); org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("uuidRecord", null, null, false); - avroSchema.setFields(List.of(new org.apache.avro.Schema.Field(UUID_COLUMN, - org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING), null, null))); + avroSchema.setFields(List.of( + new org.apache.avro.Schema.Field(UUID_DICT_SV_COLUMN, uuidSchema, null, null), + new org.apache.avro.Schema.Field(UUID_DICT_MV_COLUMN, org.apache.avro.Schema.createArray(uuidSchema), null, + null), + new org.apache.avro.Schema.Field(UUID_RAW_SV_COLUMN, uuidSchema, null, null), + new org.apache.avro.Schema.Field(UUID_RAW_MV_COLUMN, org.apache.avro.Schema.createArray(uuidSchema), null, + null))); try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) { DataFileWriter writer = avroFilesAndWriters.getWriters().get(0); - for (String uuid : ROWS) { + for (int i = 0; i < UUID_SV_VALUES.size(); i++) { GenericData.Record record = new GenericData.Record(avroSchema); - record.put(UUID_COLUMN, uuid); + record.put(UUID_DICT_SV_COLUMN, UUID_SV_VALUES.get(i)); + record.put(UUID_DICT_MV_COLUMN, UUID_MV_VALUES.get(i)); + record.put(UUID_RAW_SV_COLUMN, UUID_SV_VALUES.get(i)); + record.put(UUID_RAW_MV_COLUMN, UUID_MV_VALUES.get(i)); writer.append(record); } return avroFilesAndWriters.getAvroFiles(); @@ -101,104 +112,81 @@ public List createAvroFiles() } @Test - public void testGroupByUuidColumn() + public void testGroupByHavingReturningFinalResult() throws Exception { setUseMultiStageQueryEngine(false); - JsonNode rows = query( - String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s ORDER BY %s", UUID_COLUMN, getTableName(), UUID_COLUMN, - UUID_COLUMN)); - assertEquals(rows.size(), NUM_DISTINCT, rows.toPrettyString()); - - // Group keys must come back as canonical UUIDs, not hex and not a byte-array rendering. - List keys = new ArrayList<>(); - for (JsonNode row : rows) { - keys.add(row.get(0).asText()); - } - assertEquals(keys, List.of(UUID_0, UUID_1, UUID_2), rows.toPrettyString()); - assertEquals(rows.get(0).get(1).asLong(), 2, rows.toPrettyString()); - assertEquals(rows.get(1).get(1).asLong(), 1, rows.toPrettyString()); - } + JsonNode rows = query(String.format( + "SELECT %1$s, COUNT(*) FROM %2$s GROUP BY %1$s HAVING %1$s = '%3$s' " + + "OPTION(serverReturnFinalResult=true)", + UUID_DICT_SV_COLUMN, getTableName(), UUID_0_HEX)); - /// The regression that motivated this class: `GROUP BY` a UUID column with a `HAVING` predicate on that same - /// column runs the group key through `getConvertedKey` and then straight into `PredicateRowMatcher`, which casts - /// to `java.util.UUID`. - @Test - public void testGroupByUuidColumnWithHaving() - throws Exception { - setUseMultiStageQueryEngine(false); - JsonNode rows = query( - String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s HAVING %s = '%s'", UUID_COLUMN, getTableName(), - UUID_COLUMN, UUID_COLUMN, UUID_0_HEX)); assertEquals(rows.size(), 1, rows.toPrettyString()); assertEquals(rows.get(0).get(0).asText(), UUID_0, rows.toPrettyString()); - assertEquals(rows.get(0).get(1).asLong(), 2, rows.toPrettyString()); - - // Same thing via an explicit CAST of the canonical form. - rows = query(String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s HAVING %s = CAST('%s' AS UUID)", UUID_COLUMN, - getTableName(), UUID_COLUMN, UUID_COLUMN, UUID_1)); - assertEquals(rows.size(), 1, rows.toPrettyString()); - assertEquals(rows.get(0).get(0).asText(), UUID_1, rows.toPrettyString()); - assertEquals(rows.get(0).get(1).asLong(), 1, rows.toPrettyString()); + assertEquals(rows.get(0).get(1).asLong(), 2L, rows.toPrettyString()); } - /// The `serverReturnFinalResult` variant. This is the option that gates the single-data-table branch in - /// `GroupByDataTableReducer#reduceAndSetResults` (`isServerReturnFinalResult() && dataTables.size() == 1`), so it - /// covers a different reduce branch from the test above. Verified by reverting the `getConvertedKey` fix: both - /// this and the plain `GROUP BY ... HAVING` test fail with - /// `ClassCastException: class [B cannot be cast to class java.util.UUID`, and both pass with it. - @Test - public void testGroupByUuidColumnWithHavingReturningFinalResult() + @Test(dataProvider = "useBothQueryEngines") + public void testDistinctOnUuidColumn(boolean useMultiStageQueryEngine) throws Exception { - setUseMultiStageQueryEngine(false); - JsonNode rows = query( - String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s HAVING %s = '%s' OPTION(serverReturnFinalResult=true)", - UUID_COLUMN, getTableName(), UUID_COLUMN, UUID_COLUMN, UUID_0_HEX)); - assertEquals(rows.size(), 1, rows.toPrettyString()); - assertEquals(rows.get(0).get(0).asText(), UUID_0, rows.toPrettyString()); - assertEquals(rows.get(0).get(1).asLong(), 2, rows.toPrettyString()); + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + JsonNode rows = query(String.format("SELECT DISTINCT %1$s FROM %2$s ORDER BY %1$s", UUID_DICT_SV_COLUMN, + getTableName())); + + assertEquals(rows.size(), 3, rows.toPrettyString()); + for (int i = 0; i < rows.size(); i++) { + assertEquals(rows.get(i).get(0).asText(), List.of(UUID_0, UUID_1, UUID_2).get(i), rows.toPrettyString()); + } } - /// Group keys must also render canonically on the `getConvertedKey` path, not as hex. @Test - public void testGroupByUuidColumnReturningFinalResult() + public void testDistinctCountOnUuidColumns() throws Exception { setUseMultiStageQueryEngine(false); - JsonNode rows = query( - String.format("SELECT %s, COUNT(*) FROM %s GROUP BY %s ORDER BY %s OPTION(serverReturnFinalResult=true)", - UUID_COLUMN, getTableName(), UUID_COLUMN, UUID_COLUMN)); - assertEquals(rows.size(), NUM_DISTINCT, rows.toPrettyString()); - List keys = new ArrayList<>(); - for (JsonNode row : rows) { - keys.add(row.get(0).asText()); + for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTHLLPLUS", + "DISTINCTCOUNTBITMAP", "DISTINCTCOUNTTHETASKETCH", "DISTINCTCOUNTCPCSKETCH")) { + JsonNode rows = query(String.format("SELECT %1$s(%2$s), %1$s(%3$s), %1$s(%4$s) FROM %5$s", function, + UUID_DICT_SV_COLUMN, UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN, getTableName())); + assertCounts(rows.get(0), 3L, 3L, 4L); } - assertEquals(keys, List.of(UUID_0, UUID_1, UUID_2), rows.toPrettyString()); + + // DISTINCTCOUNTULL currently supports only single-value inputs. + JsonNode rows = query(String.format("SELECT DISTINCTCOUNTULL(%s), DISTINCTCOUNTULL(%s) FROM %s", + UUID_DICT_SV_COLUMN, UUID_RAW_SV_COLUMN, getTableName())); + assertCounts(rows.get(0), 3L, 3L); + + rows = query(String.format( + "SELECT DISTINCTCOUNTTHETASKETCH(%1$s, '', '%1$s = ''%3$s''', '$1'), " + + "DISTINCTCOUNTTHETASKETCH(%2$s, '', '%2$s = ''%3$s''', '$1') FROM %4$s", + UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN, UUID_0, getTableName())); + assertCounts(rows.get(0), 1L, 2L); } @Test - public void testDistinctOnUuidColumn() + public void testCpcAndThetaGroupByUuidColumns() throws Exception { setUseMultiStageQueryEngine(false); - JsonNode rows = - query(String.format("SELECT DISTINCT %s FROM %s ORDER BY %s", UUID_COLUMN, getTableName(), UUID_COLUMN)); - assertEquals(rows.size(), NUM_DISTINCT, rows.toPrettyString()); - - // BytesDistinctTable used to hard-code toHexString(); a UUID column must render canonically. - List values = new ArrayList<>(); - for (JsonNode row : rows) { - values.add(row.get(0).asText()); - } - assertEquals(values, List.of(UUID_0, UUID_1, UUID_2), rows.toPrettyString()); + JsonNode rows = queryGroupBy(UUID_RAW_SV_COLUMN); + assertEquals(rows.size(), 3, rows.toPrettyString()); + assertGroupRow(rows.get(0), UUID_0, 1L, 3L, 1L, 3L); + assertGroupRow(rows.get(1), UUID_1, 1L, 1L, 1L, 1L); + assertGroupRow(rows.get(2), UUID_2, 1L, 1L, 1L, 1L); + + // UUID multi-value group keys require a dictionary, while the aggregate input remains raw. + rows = queryGroupBy(UUID_DICT_MV_COLUMN); + assertEquals(rows.size(), 4, rows.toPrettyString()); + assertGroupRow(rows.get(0), UUID_0, 2L, 2L, 2L, 2L); + assertGroupRow(rows.get(1), UUID_1, 1L, 3L, 1L, 3L); + assertGroupRow(rows.get(2), UUID_2, 1L, 2L, 1L, 2L); + assertGroupRow(rows.get(3), UUID_3, 1L, 1L, 1L, 1L); } - @Test - public void testDistinctCountOnUuidColumn() + private JsonNode queryGroupBy(String groupByColumn) throws Exception { - setUseMultiStageQueryEngine(false); - for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTHLLPLUS", "DISTINCTCOUNTULL", - "DISTINCTCOUNTBITMAP", "DISTINCTCOUNTTHETASKETCH", "DISTINCTCOUNTCPCSKETCH")) { - JsonNode rows = query(String.format("SELECT %s(%s) FROM %s", function, UUID_COLUMN, getTableName())); - assertEquals(rows.get(0).get(0).asLong(), NUM_DISTINCT, function + ": " + rows.toPrettyString()); - } + return query(String.format( + "SELECT %1$s, DISTINCTCOUNTCPCSKETCH(%2$s), DISTINCTCOUNTCPCSKETCH(%3$s), " + + "DISTINCTCOUNTTHETASKETCH(%2$s), DISTINCTCOUNTTHETASKETCH(%3$s) " + + "FROM %4$s GROUP BY %1$s ORDER BY %1$s", + groupByColumn, UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN, getTableName())); } private JsonNode query(String sql) @@ -207,4 +195,18 @@ private JsonNode query(String sql) assertTrue(response.path("exceptions").isEmpty(), sql + " -> " + response.toPrettyString()); return response.path("resultTable").path("rows"); } + + private static void assertGroupRow(JsonNode row, String groupKey, long... expectedCounts) { + assertEquals(row.get(0).asText(), groupKey, row.toPrettyString()); + for (int i = 0; i < expectedCounts.length; i++) { + assertEquals(row.get(i + 1).asLong(), expectedCounts[i], row.toPrettyString()); + } + } + + private static void assertCounts(JsonNode row, long... expectedCounts) { + assertEquals(row.size(), expectedCounts.length, row.toPrettyString()); + for (int i = 0; i < expectedCounts.length; i++) { + assertEquals(row.get(i).asLong(), expectedCounts[i], row.toPrettyString()); + } + } } From bc1940b0324c25ffc64c63ff274d58d821467c1f Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 7 Jul 2026 14:27:43 -0700 Subject: [PATCH 18/19] [UUID 6/8] UUID multi-stage engine (planner + runtime) Part 6/8 of splitting apache/pinot#18140 (logical UUID type). Rebased onto latest master; stacked on uuid-split/05-agg-groupby-distinct. Downstream references use the UuidKey class merged in #18869. --- pinot-common/src/main/proto/expressions.proto | 4 + .../parser/CalciteRexExpressionParser.java | 11 ++- .../logical/RelToPlanNodeConverter.java | 4 +- .../planner/logical/RexExpressionUtils.java | 18 ++++ .../physical/v2/PRelToPlanNodeConverter.java | 4 +- .../serde/ProtoExpressionToRexExpression.java | 5 + .../serde/RexExpressionToProtoExpression.java | 1 + .../CalciteRexExpressionParserTest.java | 53 +++++++++++ .../logical/RelToPlanNodeConverterTest.java | 8 ++ .../planner/serde/RexExpressionSerDeTest.java | 9 +- .../runtime/operator/HashJoinOperator.java | 22 +++-- .../groupby/GroupIdGeneratorFactory.java | 2 + .../groupby/OneUuidKeyGroupIdGenerator.java | 81 ++++++++++++++++ .../runtime/operator/join/LookupTable.java | 6 ++ .../operator/join/UuidLookupTable.java | 84 ++++++++++++++++ .../plan/server/ServerPlanRequestUtils.java | 7 +- .../operator/HashJoinOperatorTest.java | 39 ++++++++ .../server/ServerPlanRequestUtilsTest.java | 95 +++++++++++++++++++ 18 files changed, 436 insertions(+), 17 deletions(-) create mode 100644 pinot-query-planner/src/test/java/org/apache/pinot/query/parser/CalciteRexExpressionParserTest.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/OneUuidKeyGroupIdGenerator.java create mode 100644 pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/UuidLookupTable.java create mode 100644 pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtilsTest.java diff --git a/pinot-common/src/main/proto/expressions.proto b/pinot-common/src/main/proto/expressions.proto index cd185eba0843..74d318ba15ea 100644 --- a/pinot-common/src/main/proto/expressions.proto +++ b/pinot-common/src/main/proto/expressions.proto @@ -44,6 +44,10 @@ enum ColumnDataType { UNKNOWN = 19; MAP = 20; BIG_DECIMAL_ARRAY = 21; + // Rolling-upgrade limitation for UUID columns: in a mixed-version multi-stage query, an older broker/server that + // does not know UUID = 22 / UUID_ARRAY = 23 will fail planning with UnknownEnumValueException when receiving a plan + // that includes a UUID literal. Avoid issuing UUID queries until all brokers and servers are upgraded. See the + // matching note on DataSchema.toBytes and ProtoExpressionToRexExpression#convertColumnDataType. UUID = 22; UUID_ARRAY = 23; } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java index 432ea445a517..f34575435e87 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/parser/CalciteRexExpressionParser.java @@ -31,6 +31,7 @@ import org.apache.pinot.query.planner.logical.RexExpression; import org.apache.pinot.spi.utils.BooleanUtils; import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.sql.parsers.ParserUtils; @@ -122,7 +123,13 @@ public static Expression toExpression(RexExpression rexNode, List se if (rexNode instanceof RexExpression.InputRef) { return inputRefToIdentifier((RexExpression.InputRef) rexNode, selectList); } else if (rexNode instanceof RexExpression.Literal) { - return RequestUtils.getLiteralExpression(toLiteral((RexExpression.Literal) rexNode)); + RexExpression.Literal literal = (RexExpression.Literal) rexNode; + if (literal.getDataType() == ColumnDataType.UUID) { + return RequestUtils.getFunctionExpression("cast", + RequestUtils.getLiteralExpression(UuidUtils.toString((ByteArray) literal.getValue())), + RequestUtils.getLiteralExpression("UUID")); + } + return RequestUtils.getLiteralExpression(toLiteral(literal)); } else { assert rexNode instanceof RexExpression.FunctionCall; return compileFunctionExpression((RexExpression.FunctionCall) rexNode, selectList); @@ -144,6 +151,8 @@ public static Literal toLiteral(RexExpression.Literal literal) { ColumnDataType dataType = literal.getDataType(); if (dataType == ColumnDataType.BOOLEAN) { value = BooleanUtils.isTrueInternalValue(value); + } else if (dataType == ColumnDataType.UUID) { + value = UuidUtils.toString((ByteArray) value); } else if (dataType == ColumnDataType.BYTES) { value = ((ByteArray) value).getBytes(); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java index 6198f1cbcb0f..65d6ec1aceef 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverter.java @@ -1080,11 +1080,11 @@ public static ColumnDataType convertToColumnDataType(RelDataType relDataType) { case CHAR: case VARCHAR: return isArray ? ColumnDataType.STRING_ARRAY : ColumnDataType.STRING; + case UUID: + return isArray ? ColumnDataType.UUID_ARRAY : ColumnDataType.UUID; case BINARY: case VARBINARY: return isArray ? ColumnDataType.BYTES_ARRAY : ColumnDataType.BYTES; - case UUID: - return isArray ? ColumnDataType.UUID_ARRAY : ColumnDataType.UUID; case MAP: return ColumnDataType.MAP; case OTHER: diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java index 01ba0c1d8dc7..7f91be9c1843 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/logical/RexExpressionUtils.java @@ -26,6 +26,7 @@ import java.util.Calendar; import java.util.List; import java.util.Set; +import java.util.UUID; import javax.annotation.Nullable; import org.apache.calcite.avatica.util.ByteString; import org.apache.calcite.plan.RelOptCluster; @@ -53,6 +54,7 @@ import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.spi.utils.BooleanUtils; import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -149,6 +151,9 @@ public static RexLiteral toRexLiteral(RelBuilder builder, RexExpression.Literal ByteString byteString = new ByteString(bytes); return rexBuilder.makeBinaryLiteral(byteString); } + case UUID: + assert value != null; + return rexBuilder.makeUuidLiteral(UuidUtils.toUUID((ByteArray) value)); default: throw new IllegalStateException("Unsupported ColumnDataType: " + literal.getDataType()); } @@ -264,6 +269,19 @@ private static RexExpression.Literal fromRexLiteralValue(ColumnDataType dataType case BYTES: value = new ByteArray(((ByteString) value).getBytes()); break; + case UUID: + if (value instanceof UUID) { + value = new ByteArray(UuidUtils.toBytes((UUID) value)); + } else if (value instanceof ByteString) { + value = new ByteArray(UuidUtils.toBytes(((ByteString) value).getBytes())); + } else if (value instanceof NlsString) { + value = new ByteArray(UuidUtils.toBytes(((NlsString) value).getValue())); + } else if (value instanceof String) { + value = new ByteArray(UuidUtils.toBytes((String) value)); + } else { + throw new IllegalStateException("Unsupported value type for UUID: " + value.getClass().getName()); + } + break; default: throw new IllegalStateException("Unsupported ColumnDataType: " + dataType); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java index cc9d44165932..76b576368a16 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/PRelToPlanNodeConverter.java @@ -340,11 +340,11 @@ public static ColumnDataType convertToColumnDataType(RelDataType relDataType) { case CHAR: case VARCHAR: return isArray ? ColumnDataType.STRING_ARRAY : ColumnDataType.STRING; + case UUID: + return isArray ? ColumnDataType.UUID_ARRAY : ColumnDataType.UUID; case BINARY: case VARBINARY: return isArray ? ColumnDataType.BYTES_ARRAY : ColumnDataType.BYTES; - case UUID: - return isArray ? ColumnDataType.UUID_ARRAY : ColumnDataType.UUID; case MAP: return ColumnDataType.MAP; case OTHER: diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java index 060ca8447214..bb12758825f3 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/ProtoExpressionToRexExpression.java @@ -144,6 +144,7 @@ public static RexExpression.Literal convertLiteral(Expressions.Literal literal) } return new RexExpression.Literal(dataType, values); } + // NOTE: UUID_ARRAY's stored type is BYTES_ARRAY, so this case handles both. case BYTES_ARRAY: { Expressions.BytesArray bytesArray = literal.getBytesArray(); int numValues = bytesArray.getValuesCount(); @@ -209,6 +210,10 @@ public static ColumnDataType convertColumnDataType(Expressions.ColumnDataType da case UNKNOWN: return ColumnDataType.UNKNOWN; default: + // Rolling-upgrade limitation for UUID columns: an older broker/server that does not know UUID = 22 / + // UUID_ARRAY = 23 from expressions.proto will land here with UNRECOGNIZED and throw. Avoid issuing UUID + // queries until all brokers and servers are upgraded. See the matching note on expressions.proto and + // DataSchema.toBytes. throw new IllegalStateException("Unsupported proto ColumnDataType: " + dataType); } } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java index 7b0c79ec5d81..54575592963c 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/serde/RexExpressionToProtoExpression.java @@ -130,6 +130,7 @@ public static Expressions.Literal convertLiteral(RexExpression.Literal literal) literalBuilder.setStringArray( Expressions.StringArray.newBuilder().addAllValues(Arrays.asList((String[]) value)).build()); break; + // NOTE: UUID_ARRAY's stored type is BYTES_ARRAY, so this case handles both. case BYTES_ARRAY: { ByteArray[] bytesArray = (ByteArray[]) value; Expressions.BytesArray.Builder builder = Expressions.BytesArray.newBuilder(); diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/parser/CalciteRexExpressionParserTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/parser/CalciteRexExpressionParserTest.java new file mode 100644 index 000000000000..8678ae80d42e --- /dev/null +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/parser/CalciteRexExpressionParserTest.java @@ -0,0 +1,53 @@ +/** + * 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.query.parser; + +import java.util.List; +import org.apache.pinot.common.request.Expression; +import org.apache.pinot.common.request.Function; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.query.planner.logical.RexExpression; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertNull; + + +public class CalciteRexExpressionParserTest { + private static final String UUID_VALUE = "550e8400-e29b-41d4-a716-446655440000"; + + @Test + public void testToExpressionPreservesUuidLiteralAsCast() { + RexExpression.Literal uuidLiteral = + new RexExpression.Literal(ColumnDataType.UUID, new ByteArray(UuidUtils.toBytes(UUID_VALUE))); + + Expression expression = CalciteRexExpressionParser.toExpression(uuidLiteral, List.of()); + + assertNull(expression.getLiteral()); + Function function = expression.getFunctionCall(); + assertNotNull(function); + assertEquals(function.getOperator(), "cast"); + assertEquals(function.getOperandsSize(), 2); + assertEquals(function.getOperands().get(0).getLiteral().getStringValue(), UUID_VALUE); + assertEquals(function.getOperands().get(1).getLiteral().getStringValue(), "UUID"); + } +} diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java index 69d6acca0cda..c99eecc2a430 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/logical/RelToPlanNodeConverterTest.java @@ -45,6 +45,7 @@ import org.apache.calcite.sql.type.SqlTypeName; import org.apache.calcite.util.ImmutableBitSet; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.query.planner.physical.v2.PRelToPlanNodeConverter; import org.apache.pinot.query.planner.plannode.FilterNode; import org.apache.pinot.query.planner.plannode.PlanNode; import org.apache.pinot.query.planner.plannode.UnnestNode; @@ -136,6 +137,13 @@ public void testBigDecimal() { DataSchema.ColumnDataType.BIG_DECIMAL); } + @Test + public void testConvertToColumnDataTypeForUUID() { + RelDataType uuidType = new BasicSqlType(RelDataTypeSystem.DEFAULT, SqlTypeName.UUID); + Assert.assertEquals(RelToPlanNodeConverter.convertToColumnDataType(uuidType), DataSchema.ColumnDataType.UUID); + Assert.assertEquals(PRelToPlanNodeConverter.convertToColumnDataType(uuidType), DataSchema.ColumnDataType.UUID); + } + @Test public void testConvertToColumnDataTypeForArray() { Assert.assertEquals(RelToPlanNodeConverter.convertToColumnDataType( diff --git a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java index ffd8a62e76c8..232497b6156f 100644 --- a/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java +++ b/pinot-query-planner/src/test/java/org/apache/pinot/query/planner/serde/RexExpressionSerDeTest.java @@ -36,11 +36,13 @@ public class RexExpressionSerDeTest { private static final List SUPPORTED_DATE_TYPES = List.of(ColumnDataType.INT, ColumnDataType.LONG, ColumnDataType.FLOAT, ColumnDataType.DOUBLE, ColumnDataType.BIG_DECIMAL, ColumnDataType.BOOLEAN, ColumnDataType.TIMESTAMP, ColumnDataType.STRING, - ColumnDataType.BYTES, ColumnDataType.UUID, ColumnDataType.INT_ARRAY, ColumnDataType.LONG_ARRAY, + ColumnDataType.UUID, ColumnDataType.BYTES, ColumnDataType.INT_ARRAY, ColumnDataType.LONG_ARRAY, ColumnDataType.FLOAT_ARRAY, ColumnDataType.DOUBLE_ARRAY, ColumnDataType.BOOLEAN_ARRAY, - ColumnDataType.TIMESTAMP_ARRAY, ColumnDataType.STRING_ARRAY, ColumnDataType.UUID_ARRAY, + ColumnDataType.TIMESTAMP_ARRAY, + ColumnDataType.STRING_ARRAY, ColumnDataType.BYTES_ARRAY, ColumnDataType.UUID_ARRAY, ColumnDataType.UNKNOWN); private static final Random RANDOM = new Random(); + private static final String UUID_VALUE = "550e8400-e29b-41d4-a716-446655440000"; @Test public void testNullLiteral() { @@ -100,8 +102,7 @@ public void testBytesLiteral() { @Test public void testUuidLiteral() { - verifyLiteralSerDe(new RexExpression.Literal(ColumnDataType.UUID, - new ByteArray(UuidUtils.toBytes("550e8400-e29b-41d4-a716-446655440000")))); + verifyLiteralSerDe(new RexExpression.Literal(ColumnDataType.UUID, new ByteArray(UuidUtils.toBytes(UUID_VALUE)))); } @Test diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java index 1a6ea42c2f72..bb1ccc2563a3 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/HashJoinOperator.java @@ -27,6 +27,7 @@ import javax.annotation.Nullable; import org.apache.calcite.rel.core.JoinRelType; import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; import org.apache.pinot.core.data.table.Key; import org.apache.pinot.query.planner.partitioning.KeySelector; import org.apache.pinot.query.planner.partitioning.KeySelectorFactory; @@ -38,6 +39,7 @@ import org.apache.pinot.query.runtime.operator.join.LongLookupTable; import org.apache.pinot.query.runtime.operator.join.LookupTable; import org.apache.pinot.query.runtime.operator.join.ObjectLookupTable; +import org.apache.pinot.query.runtime.operator.join.UuidLookupTable; import org.apache.pinot.query.runtime.plan.OpChainExecutionContext; @@ -95,7 +97,11 @@ private static LookupTable createLookupTable(List joinKeys, DataSchema if (joinKeys.size() > 1) { return new ObjectLookupTable(); } - switch (schema.getColumnDataType(joinKeys.get(0)).getStoredType()) { + ColumnDataType columnDataType = schema.getColumnDataType(joinKeys.get(0)); + if (columnDataType == ColumnDataType.UUID) { + return new UuidLookupTable(); + } + switch (columnDataType.getStoredType()) { case INT: return new IntLookupTable(); case LONG: @@ -203,7 +209,8 @@ private List buildJoinedDataBlockUniqueKeys(MseBlock.Data leftBlock) { if (handleNullKey(key, leftRow, rows)) { continue; } - Object[] rightRow = (Object[]) _rightTable.lookup(key); + Object normalizedKey = _rightTable.normalizeKey(key); + Object[] rightRow = (Object[]) _rightTable.lookup(normalizedKey); if (rightRow == null) { handleUnmatchedLeftRow(leftRow, rows); } else { @@ -216,7 +223,7 @@ private List buildJoinedDataBlockUniqueKeys(MseBlock.Data leftBlock) { checkTerminationAndSampleUsagePeriodically(rows.size(), BUILD_JOINED_ROWS_SCOPE); rows.add(resultRowView.toArray()); if (_matchedRightRows != null) { - _matchedRightRows.put(key, BIT_SET_PLACEHOLDER); + _matchedRightRows.put(normalizedKey, BIT_SET_PLACEHOLDER); } } else { handleUnmatchedLeftRow(leftRow, rows); @@ -238,7 +245,8 @@ private List buildJoinedDataBlockDuplicateKeys(MseBlock.Data leftBlock if (handleNullKey(key, leftRow, rows)) { continue; } - List rightRows = (List) _rightTable.lookup(key); + Object normalizedKey = _rightTable.normalizeKey(key); + List rightRows = (List) _rightTable.lookup(normalizedKey); if (rightRows == null) { handleUnmatchedLeftRow(leftRow, rows); } else { @@ -256,7 +264,7 @@ private List buildJoinedDataBlockDuplicateKeys(MseBlock.Data leftBlock rows.add(resultRowView.toArray()); hasMatchForLeftRow = true; if (_matchedRightRows != null) { - _matchedRightRows.computeIfAbsent(key, k -> new BitSet(numRightRows)).set(i); + _matchedRightRows.computeIfAbsent(normalizedKey, k -> new BitSet(numRightRows)).set(i); } } } @@ -289,7 +297,7 @@ private List buildJoinedDataBlockSemi(MseBlock.Data leftBlock) { for (Object[] leftRow : leftRows) { Object key = _leftKeySelector.getKey(leftRow); - if (_rightTable.containsKey(key)) { + if (_rightTable.containsKey(_rightTable.normalizeKey(key))) { checkTerminationAndSampleUsagePeriodically(rows.size(), BUILD_JOINED_ROWS_SCOPE); rows.add(leftRow); } @@ -305,7 +313,7 @@ private List buildJoinedDataBlockAnti(MseBlock.Data leftBlock) { for (Object[] leftRow : leftRows) { Object key = _leftKeySelector.getKey(leftRow); - if (!_rightTable.containsKey(key)) { + if (!_rightTable.containsKey(_rightTable.normalizeKey(key))) { checkTerminationAndSampleUsagePeriodically(rows.size(), BUILD_JOINED_ROWS_SCOPE); rows.add(leftRow); } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGeneratorFactory.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGeneratorFactory.java index a254c37bf812..c18ac8534416 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGeneratorFactory.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/GroupIdGeneratorFactory.java @@ -38,6 +38,8 @@ public static GroupIdGenerator getGroupIdGenerator(ColumnDataType[] keyTypes, in return new OneFloatKeyGroupIdGenerator(numGroupsLimit, initialCapacity); case DOUBLE: return new OneDoubleKeyGroupIdGenerator(numGroupsLimit, initialCapacity); + case UUID: + return new OneUuidKeyGroupIdGenerator(numGroupsLimit, initialCapacity); default: return new OneObjectKeyGroupIdGenerator(numGroupsLimit, initialCapacity); } diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/OneUuidKeyGroupIdGenerator.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/OneUuidKeyGroupIdGenerator.java new file mode 100644 index 000000000000..9dea1f763b47 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/groupby/OneUuidKeyGroupIdGenerator.java @@ -0,0 +1,81 @@ +/** + * 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.query.runtime.operator.groupby; + +import it.unimi.dsi.fastutil.objects.Object2IntMap; +import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; +import it.unimi.dsi.fastutil.objects.ObjectIterator; +import java.util.Iterator; +import java.util.function.ToIntFunction; +import org.apache.pinot.spi.utils.UuidKey; + + +/// Group-id generator for a single UUID group-by key in the multi-stage engine. Normalizes every incoming key +/// (`byte[]`, `ByteArray`, `String` or `UuidKey`) to [UuidKey] — two primitive longs — so map probing avoids byte-array +/// hashing/equality on the hot path. Not thread-safe; each instance is owned by a single operator thread, matching the +/// other [GroupIdGenerator] implementations. +public class OneUuidKeyGroupIdGenerator implements GroupIdGenerator { + private final Object2IntOpenHashMap _groupIdMap; + private final int _numGroupsLimit; + private final ToIntFunction _groupIdGenerator; + + public OneUuidKeyGroupIdGenerator(int numGroupsLimit, int initialCapacity) { + _groupIdMap = new Object2IntOpenHashMap<>(initialCapacity); + _groupIdMap.defaultReturnValue(INVALID_ID); + _numGroupsLimit = numGroupsLimit; + _groupIdGenerator = ignored -> _groupIdMap.size(); + } + + @Override + public int getGroupId(Object key) { + Object normalizedKey = key != null ? UuidKey.fromObject(key) : null; + if (_groupIdMap.size() < _numGroupsLimit) { + return _groupIdMap.computeIfAbsent(normalizedKey, _groupIdGenerator); + } else { + return _groupIdMap.getInt(normalizedKey); + } + } + + @Override + public int getNumGroups() { + return _groupIdMap.size(); + } + + @Override + public Iterator getGroupKeyIterator(int numColumns) { + return new Iterator() { + final ObjectIterator> _entryIterator = + _groupIdMap.object2IntEntrySet().fastIterator(); + + @Override + public boolean hasNext() { + return _entryIterator.hasNext(); + } + + @Override + public GroupKey next() { + Object2IntMap.Entry entry = _entryIterator.next(); + Object[] row = new Object[numColumns]; + Object key = entry.getKey(); + row[0] = key != null ? ((UuidKey) key).toByteArray() : null; + return new GroupKey(entry.getIntValue(), row); + } + }; + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java index 3d0664dfc79d..640167568d1f 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/LookupTable.java @@ -56,6 +56,12 @@ protected Object computeNewValue(Object[] row, @Nullable Object currentValue) { /// table, and before looking up rows. public abstract void finish(); + /// Normalizes a join key into the internal lookup-table key shape. + @Nullable + public Object normalizeKey(@Nullable Object key) { + return key; + } + protected static void convertValueToList(Map.Entry entry) { Object value = entry.getValue(); if (value instanceof Object[]) { diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/UuidLookupTable.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/UuidLookupTable.java new file mode 100644 index 000000000000..7fc7e69c91f1 --- /dev/null +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/operator/join/UuidLookupTable.java @@ -0,0 +1,84 @@ +/** + * 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.query.runtime.operator.join; + +import com.google.common.collect.Maps; +import java.util.Map; +import java.util.Set; +import javax.annotation.Nullable; +import org.apache.pinot.spi.utils.UuidKey; + + +/// Lookup table optimized for Pinot's logical UUID type — stores entries keyed by [UuidKey] (two primitive longs) so +/// the hot path avoids `ByteArray` wrapping/equals. +/// +/// **Contract:** [#addRow] normalizes the supplied key implicitly via [#normalizeKey]; [#containsKey] and [#lookup] +/// do NOT — callers must pass an already-normalized key (the join operators do this via +/// `_rightTable.normalizeKey(...)`). Passing a raw `byte[]`, `String`, or `ByteArray` to [#containsKey] / [#lookup] +/// will silently miss because the table is keyed on `UuidKey`, whose equality is by primitive longs. +@SuppressWarnings("unchecked") +public class UuidLookupTable extends LookupTable { + private final Map _lookupTable = Maps.newHashMapWithExpectedSize(INITIAL_CAPACITY); + + @Override + public void addRow(@Nullable Object key, Object[] row) { + Object normalizedKey = normalizeKey(key); + if (normalizedKey == null) { + return; + } + _lookupTable.compute(normalizedKey, (k, v) -> computeNewValue(row, v)); + } + + @Override + public void finish() { + if (!_keysUnique) { + for (Map.Entry entry : _lookupTable.entrySet()) { + convertValueToList(entry); + } + } + } + + @Nullable + @Override + public Object normalizeKey(@Nullable Object key) { + return key != null ? UuidKey.fromObject(key) : null; + } + + @Override + public boolean containsKey(@Nullable Object key) { + return key != null && _lookupTable.containsKey(key); + } + + @Nullable + @Override + public Object lookup(@Nullable Object key) { + return key != null ? _lookupTable.get(key) : null; + } + + @SuppressWarnings("rawtypes") + @Override + public Set> entrySet() { + return _lookupTable.entrySet(); + } + + @Override + public int size() { + return _lookupTable.size(); + } +} diff --git a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java index 4cdd84594bf5..6a1bba3997d0 100644 --- a/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java +++ b/pinot-query-runtime/src/main/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtils.java @@ -60,6 +60,7 @@ import org.apache.pinot.spi.data.Schema; import org.apache.pinot.spi.utils.ByteArray; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.apache.pinot.sql.FilterKind; import org.apache.pinot.sql.parsers.rewriter.NonAggregationGroupByToDistinctQueryRewriter; @@ -398,7 +399,11 @@ private static List computeInOperands(List dataContainer, } Arrays.sort(arrBytes); for (int rowIdx = 0; rowIdx < numRows; rowIdx++) { - expressions.add(RequestUtils.getLiteralExpression(arrBytes[rowIdx].getBytes())); + if (columnDataType == DataSchema.ColumnDataType.UUID) { + expressions.add(RequestUtils.getLiteralExpression(UuidUtils.toString(arrBytes[rowIdx]))); + } else { + expressions.add(RequestUtils.getLiteralExpression(arrBytes[rowIdx].getBytes())); + } } break; default: diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java index dcd9c04ce05a..604c0ed4b21a 100644 --- a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/operator/HashJoinOperatorTest.java @@ -32,6 +32,8 @@ import org.apache.pinot.query.runtime.blocks.ErrorMseBlock; import org.apache.pinot.query.runtime.blocks.MseBlock; import org.apache.pinot.spi.exception.QueryErrorCode; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; import org.mockito.Mock; import org.mockito.Mockito; import org.testng.annotations.AfterMethod; @@ -47,6 +49,9 @@ public class HashJoinOperatorTest { + private static final ByteArray UUID_A = uuid("550e8400-e29b-41d4-a716-446655440000"); + private static final ByteArray UUID_B = uuid("550e8400-e29b-41d4-a716-446655440001"); + private static final ByteArray UUID_C = uuid("550e8400-e29b-41d4-a716-446655440002"); private AutoCloseable _mocks; private MultiStageOperator _leftInput; private MultiStageOperator _rightInput; @@ -55,6 +60,8 @@ public class HashJoinOperatorTest { private static final DataSchema DEFAULT_CHILD_SCHEMA = new DataSchema(new String[]{"int_col", "string_col"}, new ColumnDataType[] {ColumnDataType.INT, ColumnDataType.STRING}); + private static final DataSchema UUID_CHILD_SCHEMA = new DataSchema(new String[]{"uuid_col", "int_col"}, + new ColumnDataType[] {ColumnDataType.UUID, ColumnDataType.INT}); @BeforeMethod public void setUp() { _mocks = openMocks(this); @@ -114,6 +121,34 @@ public void shouldHandleInnerJoinOnInt() { "Max rows in join should equal right table size"); } + @Test + public void shouldHandleRightJoinOnUuid() { + _leftInput = new BlockListMultiStageOperator.Builder(UUID_CHILD_SCHEMA) + .addRow(UUID_A, 1) + .addRow(UUID_B, 2) + .buildWithEos(); + _rightInput = new BlockListMultiStageOperator.Builder(UUID_CHILD_SCHEMA) + .addRow(UUID_B, 20) + .addRow(UUID_B, 21) + .addRow(UUID_C, 30) + .buildWithEos(); + DataSchema resultSchema = new DataSchema(new String[]{"uuid_col1", "int_col1", "uuid_col2", "int_col2"}, + new ColumnDataType[]{ColumnDataType.UUID, ColumnDataType.INT, ColumnDataType.UUID, ColumnDataType.INT}); + + HashJoinOperator operator = getOperator(UUID_CHILD_SCHEMA, resultSchema, JoinRelType.RIGHT, List.of(0), List.of(0), + List.of(), PlanNode.NodeHint.EMPTY); + + List resultRows1 = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(resultRows1.size(), 2); + assertTrue(containsRow(resultRows1, new Object[]{UUID_B, 2, UUID_B, 20})); + assertTrue(containsRow(resultRows1, new Object[]{UUID_B, 2, UUID_B, 21})); + + List resultRows2 = ((MseBlock.Data) operator.nextBlock()).asRowHeap().getRows(); + assertEquals(resultRows2.size(), 1); + assertTrue(containsRow(resultRows2, new Object[]{null, null, UUID_C, 30})); + assertTrue(operator.nextBlock().isSuccess()); + } + @Test public void shouldHandleLeftJoin() { _leftInput = new BlockListMultiStageOperator.Builder(DEFAULT_CHILD_SCHEMA) @@ -570,6 +605,10 @@ private boolean containsRow(List rows, Object[] expectedRow) { return false; } + private static ByteArray uuid(String value) { + return new ByteArray(UuidUtils.toBytes(value)); + } + @Test public void shouldHandleSemiJoinWithNulls() { diff --git a/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtilsTest.java b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtilsTest.java new file mode 100644 index 000000000000..3c73797db17f --- /dev/null +++ b/pinot-query-runtime/src/test/java/org/apache/pinot/query/runtime/plan/server/ServerPlanRequestUtilsTest.java @@ -0,0 +1,95 @@ +/** + * 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.query.runtime.plan.server; + +import java.lang.reflect.Method; +import java.util.ArrayList; +import java.util.List; +import org.apache.pinot.common.request.Expression; +import org.apache.pinot.common.utils.DataSchema; +import org.apache.pinot.common.utils.DataSchema.ColumnDataType; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidUtils; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; +import static org.testng.Assert.assertTrue; + + +/// Unit tests for [ServerPlanRequestUtils]. +public class ServerPlanRequestUtilsTest { + private static final String UUID_VALUE = "550e8400-e29b-41d4-a716-446655440000"; + + /// Regression test for the UUID IN-predicate literal fix. + /// + /// Before the fix, UUID ByteArray values were passed as raw byte[] literals into the dynamic filter. The server-side + /// predicate evaluator expected a canonical UUID string, so the filter never matched and UUID join queries silently + /// returned no rows. + /// + /// After the fix, UUID values are emitted as canonical lowercase string literals. + @Test + public void testComputeInOperandsUuidEmitsStringLiterals() + throws Exception { + DataSchema schema = new DataSchema(new String[]{"uuidCol"}, new ColumnDataType[]{ColumnDataType.UUID}); + + List dataContainer = new ArrayList<>(); + dataContainer.add(new Object[]{new ByteArray(UuidUtils.toBytes(UUID_VALUE))}); + + List expressions = invokeComputeInOperands(dataContainer, schema, 0); + + assertEquals(expressions.size(), 1); + Expression expr = expressions.get(0); + assertNotNull(expr.getLiteral(), "UUID operand must be a literal expression"); + // Must be a string literal containing the canonical UUID, not a byte array literal + assertTrue(expr.getLiteral().isSetStringValue(), + "UUID literal must be a string, not bytes. Got: " + expr.getLiteral()); + assertEquals(expr.getLiteral().getStringValue(), UUID_VALUE, + "UUID literal value must be canonical lowercase RFC 4122 string"); + } + + /// Verifies that raw BYTES columns still emit byte-array literals (unchanged behavior). + @Test + public void testComputeInOperandsBytesEmitsByteArrayLiterals() + throws Exception { + DataSchema schema = new DataSchema(new String[]{"bytesCol"}, new ColumnDataType[]{ColumnDataType.BYTES}); + byte[] rawBytes = {0x01, 0x02, 0x03}; + + List dataContainer = new ArrayList<>(); + dataContainer.add(new Object[]{new ByteArray(rawBytes)}); + + List expressions = invokeComputeInOperands(dataContainer, schema, 0); + + assertEquals(expressions.size(), 1); + Expression expr = expressions.get(0); + assertNotNull(expr.getLiteral(), "BYTES operand must be a literal expression"); + assertTrue(expr.getLiteral().isSetBinaryValue(), + "BYTES literal must be binary, not string. Got: " + expr.getLiteral()); + } + + @SuppressWarnings("unchecked") + private static List invokeComputeInOperands(List dataContainer, DataSchema dataSchema, + int colIdx) + throws Exception { + Method method = ServerPlanRequestUtils.class.getDeclaredMethod("computeInOperands", List.class, DataSchema.class, + int.class); + method.setAccessible(true); + return (List) method.invoke(null, dataContainer, dataSchema, colIdx); + } +} From ce012cebfdf485652e168c7276d7119f91bd4013 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 7 Jul 2026 14:27:44 -0700 Subject: [PATCH 19/19] [UUID 7/8] UUID partitioning Part 7/8 of splitting apache/pinot#18140 (logical UUID type). - UuidPartitionFunction: hashes the 16-byte UUID form via Murmur2, matching what an external producer keyed on raw UUID bytes computes - PartitionerFactory / TableConfigPartitioner: thread the column's logical DataType through so UUID columns render canonically instead of as bare hex - UUID_ARRAY entries for the array scalar functions The UUID scalar functions and multi-stage UDF wrappers that were previously part of this layer now live in their own PR (#19091) so they can be reviewed and merged in parallel. --- .../array/ArrayLengthScalarFunction.java | 3 + .../array/ArraysOverlapScalarFunction.java | 3 + .../function/UuidPartitionFunction.java | 78 +++++++++++++++++++ .../function/PartitionFunctionTest.java | 54 +++++++++++++ .../processing/mapper/SegmentMapper.java | 2 +- .../partitioner/PartitionerFactory.java | 28 ++++++- .../partitioner/TableConfigPartitioner.java | 22 +++++- 7 files changed, 185 insertions(+), 5 deletions(-) create mode 100644 pinot-common/src/main/java/org/apache/pinot/common/partition/function/UuidPartitionFunction.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArrayLengthScalarFunction.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArrayLengthScalarFunction.java index 60fab79baa3a..31c91dd6c634 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArrayLengthScalarFunction.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArrayLengthScalarFunction.java @@ -59,6 +59,9 @@ public class ArrayLengthScalarFunction implements PinotScalarFunction { TYPE_FUNCTION_INFO_MAP.put(DataSchema.ColumnDataType.BYTES_ARRAY, new FunctionInfo(ArrayLengthScalarFunction.class.getMethod("arrayLength", byte[][].class), ArrayLengthScalarFunction.class, false)); + TYPE_FUNCTION_INFO_MAP.put(DataSchema.ColumnDataType.UUID_ARRAY, + new FunctionInfo(ArrayLengthScalarFunction.class.getMethod("arrayLength", byte[][].class), + ArrayLengthScalarFunction.class, false)); } catch (NoSuchMethodException e) { throw new RuntimeException(e); } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArraysOverlapScalarFunction.java b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArraysOverlapScalarFunction.java index 029d0b3f45a6..fc6c7523df8c 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArraysOverlapScalarFunction.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/function/scalar/array/ArraysOverlapScalarFunction.java @@ -70,6 +70,9 @@ public class ArraysOverlapScalarFunction implements PinotScalarFunction { TYPE_FUNCTION_INFO_MAP.put(DataSchema.ColumnDataType.BYTES_ARRAY, new FunctionInfo(ArraysOverlapScalarFunction.class.getMethod("arraysOverlap", byte[][].class, byte[][].class), ArraysOverlapScalarFunction.class, false)); + TYPE_FUNCTION_INFO_MAP.put(DataSchema.ColumnDataType.UUID_ARRAY, + new FunctionInfo(ArraysOverlapScalarFunction.class.getMethod("arraysOverlap", byte[][].class, byte[][].class), + ArraysOverlapScalarFunction.class, false)); } catch (NoSuchMethodException e) { throw new RuntimeException(e); } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/UuidPartitionFunction.java b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/UuidPartitionFunction.java new file mode 100644 index 000000000000..55cf151a9c33 --- /dev/null +++ b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/UuidPartitionFunction.java @@ -0,0 +1,78 @@ +/** + * 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.common.partition.function; + +import com.google.common.base.Preconditions; +import java.util.Map; +import javax.annotation.Nullable; +import org.apache.pinot.segment.spi.partition.PartitionFunction; +import org.apache.pinot.segment.spi.partition.PartitionIdNormalizer; +import org.apache.pinot.spi.utils.UuidUtils; +import org.apache.pinot.spi.utils.hash.MurmurHashFunctions; + + +/// Partition function for Pinot's logical UUID type. Parses the canonical RFC 4122 UUID string into its +/// 16-byte binary form, hashes those bytes via Murmur2, and runs the configured [PartitionIdNormalizer] +/// (default [PartitionIdNormalizer#MASK]) to derive the partition id. +/// +/// This matches what an external producer that hashes the 16-byte UUID via Murmur2 would compute (the +/// most common convention for UUID-keyed messages in Kafka/Pulsar/Kinesis). +/// +/// Hashing the binary form avoids two pitfalls of hashing the canonical UUID string directly: the dashes +/// do not contribute entropy, and producers that emit raw UUID bytes on the wire would otherwise need a +/// Pinot-only canonical-format step in their partitioning code. +public class UuidPartitionFunction implements PartitionFunction { + private static final String NAME = "Uuid"; + private static final PartitionIdNormalizer DEFAULT_NORMALIZER = PartitionIdNormalizer.MASK; + private final int _numPartitions; + private final PartitionIdNormalizer _normalizer; + + public UuidPartitionFunction(int numPartitions, @Nullable Map functionConfig) { + Preconditions.checkArgument(numPartitions > 0, "Number of partitions must be > 0, was: %s", numPartitions); + _numPartitions = numPartitions; + _normalizer = PartitionFunctionConfigs.normalizer(functionConfig, DEFAULT_NORMALIZER); + } + + @Override + public int getPartition(String value) { + byte[] uuidBytes = UuidUtils.toBytes(value); + return _normalizer.getPartitionId(MurmurHashFunctions.murmurHash2(uuidBytes), _numPartitions); + } + + @Override + public String getName() { + return NAME; + } + + @Override + public int getNumPartitions() { + return _numPartitions; + } + + @Override + public PartitionIdNormalizer getPartitionIdNormalizer() { + return _normalizer; + } + + // Keep it for backward-compatibility, use getName() instead + @Override + public String toString() { + return NAME; + } +} diff --git a/pinot-common/src/test/java/org/apache/pinot/common/partition/function/PartitionFunctionTest.java b/pinot-common/src/test/java/org/apache/pinot/common/partition/function/PartitionFunctionTest.java index 6e55464c4729..16880a9c0884 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/partition/function/PartitionFunctionTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/partition/function/PartitionFunctionTest.java @@ -26,6 +26,7 @@ import org.apache.pinot.segment.spi.partition.PartitionFunctionFactory; import org.apache.pinot.spi.utils.BytesUtils; import org.apache.pinot.spi.utils.JsonUtils; +import org.apache.pinot.spi.utils.UuidUtils; import org.apache.pinot.spi.utils.hash.FnvHashFunctions; import org.apache.pinot.spi.utils.hash.MurmurHashFunctions; import org.testng.annotations.Test; @@ -755,6 +756,59 @@ public void testByteArrayPartitionFunctionEquivalence() { testPartitionFunction(byteArrayPartitionFunction, expectedPartitions); } + /// Unit test for [UuidPartitionFunction]. + /// + /// - Verifies factory registration. + /// - Verifies partition values are deterministic and in [0, numPartitions). + /// - Verifies the function hashes the 16-byte UUID form (canonical-string format does not affect the hash). + /// - Verifies an invalid UUID string throws. + @Test + public void testUuidPartitioner() { + int numPartitions = 64; + + // Factory registration (case-insensitive name lookup) should produce a UuidPartitionFunction. + PartitionFunction viaFactory = PartitionFunctionFactory.getPartitionFunction("uUiD", numPartitions, null); + assertEquals(viaFactory.getName(), "Uuid"); + assertEquals(viaFactory.getNumPartitions(), numPartitions); + assertTrue(viaFactory instanceof UuidPartitionFunction); + + UuidPartitionFunction direct = new UuidPartitionFunction(numPartitions, null); + testBasicProperties(direct, "Uuid", numPartitions); + + // Determinism + range. + String[] uuids = new String[]{ + "00000000-0000-0000-0000-000000000000", + "ffffffff-ffff-ffff-ffff-ffffffffffff", + "12345678-1234-1234-1234-1234567890ab", + "9c5e1f24-0b8e-4c9d-87f1-0aa64a3b9d12", + "550e8400-e29b-41d4-a716-446655440000" + }; + for (String uuid : uuids) { + int partition = direct.getPartition(uuid); + assertTrue(partition >= 0 && partition < numPartitions, "partition " + partition + " out of range"); + assertEquals(direct.getPartition(uuid), partition, "non-deterministic partition for " + uuid); + + // Hash matches the documented contract: murmurHash2 of the 16-byte canonical form, masked, modulo numPartitions. + byte[] uuidBytes = UuidUtils.toBytes(uuid); + int expected = (MurmurHashFunctions.murmurHash2(uuidBytes) & Integer.MAX_VALUE) % numPartitions; + assertEquals(partition, expected); + } + + // Different UUIDs should not all collapse to the same partition (basic spread sanity check). + int firstPartition = direct.getPartition(uuids[0]); + boolean spread = false; + for (int i = 1; i < uuids.length; i++) { + if (direct.getPartition(uuids[i]) != firstPartition) { + spread = true; + break; + } + } + assertTrue(spread, "UuidPartitionFunction collapsed all sample UUIDs to the same partition"); + + // Invalid UUID must surface as an exception (not a silent zero partition). + expectThrows(IllegalArgumentException.class, () -> direct.getPartition("not-a-uuid")); + } + private void testPartitionInExpectedRange(PartitionFunction partitionFunction, Object value, int numPartitions) { int partition = partitionFunction.getPartition(value.toString()); assertTrue(partition >= 0 && partition < numPartitions); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java index c60abbd12aa9..8545679dc3b6 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/mapper/SegmentMapper.java @@ -117,7 +117,7 @@ public SegmentMapper(List recordReaderFileConfigs, Trans schema.isEnableColumnBasedNullHandling() || tableConfig.getIndexingConfig().isNullHandlingEnabled(); _transformPipeline = transformPipeline; _timeHandler = TimeHandlerFactory.getTimeHandler(processorConfig); - _partitioners = PartitionerFactory.getPartitioners(processorConfig.getPartitionerConfigs()); + _partitioners = PartitionerFactory.getPartitioners(processorConfig.getPartitionerConfigs(), schema); // Time partition + partition from partitioners _partitionsBuffer = new String[_partitioners.length + 1]; _throttledLogger = new ThrottledLogger(LOGGER, tableConfig.getIngestionConfig()); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/PartitionerFactory.java b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/PartitionerFactory.java index 8e5f54540eed..a9d805abe8a5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/PartitionerFactory.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/PartitionerFactory.java @@ -20,6 +20,9 @@ import com.google.common.base.Preconditions; import java.util.List; +import javax.annotation.Nullable; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; /// Factory for Partitioner and PartitionFilter @@ -41,6 +44,15 @@ public enum PartitionerType { /// Construct a Partitioner using the PartitioningConfig public static Partitioner getPartitioner(PartitionerConfig config) { + return getPartitioner(config, null); + } + + /// Construct a Partitioner using the PartitioningConfig. When `schema` is non-null and the partitioner + /// is column-aware (e.g. [PartitionerType#TABLE_PARTITION_CONFIG]), the column's logical + /// [FieldSpec.DataType] is threaded through so values are rendered via + /// [FieldSpec.DataType#toString(Object)] (canonical UUID strings for UUID columns) instead of the bare + /// hex from [FieldSpec#getStringValue]. + public static Partitioner getPartitioner(PartitionerConfig config, @Nullable Schema schema) { Partitioner partitioner = null; switch (config.getPartitionerType()) { @@ -67,7 +79,14 @@ public static Partitioner getPartitioner(PartitionerConfig config) { "Must provide columnName for TABLE_PARTITION_CONFIG Partitioner"); Preconditions.checkState(config.getColumnPartitionConfig() != null, "Must provide columnPartitionConfig for TABLE_PARTITION_CONFIG Partitioner"); - partitioner = new TableConfigPartitioner(config.getColumnName(), config.getColumnPartitionConfig()); + FieldSpec.DataType dataType = null; + if (schema != null) { + FieldSpec fieldSpec = schema.getFieldSpecFor(config.getColumnName()); + if (fieldSpec != null) { + dataType = fieldSpec.getDataType(); + } + } + partitioner = new TableConfigPartitioner(config.getColumnName(), config.getColumnPartitionConfig(), dataType); break; default: break; @@ -79,10 +98,15 @@ public static Partitioner getPartitioner(PartitionerConfig config) { /// /// @return Array of partitioners public static Partitioner[] getPartitioners(List partitionerConfigs) { + return getPartitioners(partitionerConfigs, null); + } + + /// Create partitioner array from configuration, optionally type-aware via the provided [Schema]. + public static Partitioner[] getPartitioners(List partitionerConfigs, @Nullable Schema schema) { int numPartitioners = partitionerConfigs.size(); Partitioner[] partitioners = new Partitioner[numPartitioners]; for (int i = 0; i < numPartitioners; i++) { - partitioners[i] = PartitionerFactory.getPartitioner(partitionerConfigs.get(i)); + partitioners[i] = PartitionerFactory.getPartitioner(partitionerConfigs.get(i), schema); } return partitioners; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/TableConfigPartitioner.java b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/TableConfigPartitioner.java index 36ece750ae98..5f2f151ebe5f 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/TableConfigPartitioner.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/segment/processing/partitioner/TableConfigPartitioner.java @@ -18,10 +18,12 @@ */ package org.apache.pinot.core.segment.processing.partitioner; +import javax.annotation.Nullable; import org.apache.pinot.segment.spi.partition.PartitionFunction; import org.apache.pinot.segment.spi.partition.PartitionFunctionFactory; import org.apache.pinot.spi.config.table.ColumnPartitionConfig; import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.FieldSpec.DataType; import org.apache.pinot.spi.data.readers.GenericRow; @@ -29,15 +31,27 @@ public class TableConfigPartitioner implements Partitioner { private final String _column; private final PartitionFunction _partitionFunction; + /// Non-null when the column's logical type is known. Used to render values via + /// {@link DataType#toString(Object)} so UUID columns produce the canonical RFC 4122 form + /// (matching MutableSegmentImpl's runtime partition path and the {@code Uuid} partition function's + /// expectation) instead of the bare-hex string that {@link FieldSpec#getStringValue} would emit. + @Nullable + private final DataType _dataType; public TableConfigPartitioner(String columnName, ColumnPartitionConfig columnPartitionConfig) { + this(columnName, columnPartitionConfig, null); + } + + public TableConfigPartitioner(String columnName, ColumnPartitionConfig columnPartitionConfig, + @Nullable DataType dataType) { _column = columnName; _partitionFunction = PartitionFunctionFactory.getPartitionFunction(columnPartitionConfig); + _dataType = dataType; } @Override public String getPartition(GenericRow genericRow) { - return String.valueOf(_partitionFunction.getPartition(FieldSpec.getStringValue(genericRow.getValue(_column)))); + return String.valueOf(_partitionFunction.getPartition(toPartitionString(genericRow.getValue(_column)))); } @Override @@ -51,6 +65,10 @@ public String getPartitionFromColumns(Object[] columnValues) { throw new IllegalArgumentException( "TableConfigPartitioner expects exactly 1 column value, got " + columnValues.length); } - return String.valueOf(_partitionFunction.getPartition(FieldSpec.getStringValue(columnValues[0]))); + return String.valueOf(_partitionFunction.getPartition(toPartitionString(columnValues[0]))); + } + + private String toPartitionString(Object value) { + return _dataType != null ? _dataType.toString(value) : FieldSpec.getStringValue(value); } }