From 7eb393ed067a8e520ad5541c6c1879bd82d207c5 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 7 Jul 2026 14:27:43 -0700 Subject: [PATCH 01/11] [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 31c6b7a55ca9467e7abb6b1ab4d8bc0618de0ab1 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sat, 8 Aug 2026 15:24:01 -0700 Subject: [PATCH 02/11] 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 2d70510d1de5fed45504b8f71b52051f7d5b9948 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sat, 8 Aug 2026 16:40:41 -0700 Subject: [PATCH 03/11] 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 92f322978ffc92dea4786817991d8e22a9a5b7a6 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 9 Aug 2026 02:06:00 -0700 Subject: [PATCH 04/11] 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 4f5576e6b23d7c0f4e01f66c5eedbb5e9dcd6598 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 9 Aug 2026 03:47:45 -0700 Subject: [PATCH 05/11] 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 0476d365f92fc05f9798efe1fd1c303b45b62be3 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Sun, 9 Aug 2026 13:09:33 -0700 Subject: [PATCH 06/11] 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 14afb78c658f5dc3b194f83311653fdd68dc28cd Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 12 Aug 2026 17:57:26 -0700 Subject: [PATCH 07/11] 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 986b5f8920befaae32cb843692e5b889efa4459d Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Wed, 12 Aug 2026 18:14:56 -0700 Subject: [PATCH 08/11] 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 e8957f8c075c3c89c48f4153e5d705572fb97fe8 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 7 Jul 2026 14:27:43 -0700 Subject: [PATCH 09/11] [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 0bb8d39db0dffdb64f1d1da4dd1678e9cf15fc4c Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 7 Jul 2026 14:27:44 -0700 Subject: [PATCH 10/11] [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); } } From cc8213a9a0b45575e3a4c5e03c2385d885bd8249 Mon Sep 17 00:00:00 2001 From: Xiang Fu Date: Tue, 7 Jul 2026 14:27:44 -0700 Subject: [PATCH 11/11] [UUID 8/8] UUID integration tests, benchmarks and docs Part 8/8 of splitting apache/pinot#18140 (logical UUID type). Rebased onto latest master; stacked on uuid-split/07-udfs-partitioning. Downstream references use the UuidKey class merged in #18869. --- README.md | 50 ++ .../pinot/integration/tests/ClusterTest.java | 2 + ...CustomDataQueryClusterIntegrationTest.java | 3 +- .../tests/custom/UuidTypeRealtimeTest.java | 34 ++ .../tests/custom/UuidTypeTest.java | 540 ++++++++++++++++++ .../tests/custom/UuidUpsertRealtimeTest.java | 337 +++++++++++ pinot-perf/pom.xml | 4 + .../perf/BenchmarkUuidGroupingAndLookup.java | 273 +++++++++ .../perf/BenchmarkUuidQueryExecution.java | 360 ++++++++++++ 9 files changed, 1602 insertions(+), 1 deletion(-) create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeRealtimeTest.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeTest.java create mode 100644 pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidUpsertRealtimeTest.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidGroupingAndLookup.java create mode 100644 pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidQueryExecution.java diff --git a/README.md b/README.md index c8c0aa1266ec..1ae4f8f05cb4 100644 --- a/README.md +++ b/README.md @@ -193,5 +193,55 @@ Check out [Pinot documentation](https://docs.pinot.apache.org/) for a complete d - [Pinot Architecture](https://docs.pinot.apache.org/basics/architecture) - [Pinot Query Language](https://docs.pinot.apache.org/users/user-guide-query/pinot-query-language) +### UUID Logical Type + +Pinot supports a logical `UUID` type for both single- and multi-value columns. In v1, Pinot stores `UUID` values +using the existing 16-byte `BYTES` representation, while schema definitions and query results use canonical +lowercase RFC 4122 strings. + +Schema example: +```json +{ + "schemaName": "events", + "dimensionFieldSpecs": [ + { + "name": "eventId", + "dataType": "UUID" + } + ] +} +``` + +Query example: +```sql +SELECT eventId +FROM events +WHERE eventId = CAST('550e8400-e29b-41d4-a716-446655440000' AS UUID) +``` + +UUID conversion helpers: +```sql +SELECT + TO_UUID('550E8400-E29B-41D4-A716-446655440000'), + UUID_TO_STRING(eventId), + UUID_TO_BYTES(eventId), + BYTES_TO_UUID(eventIdBytes), + IS_UUID(eventIdBytes) +FROM events +``` + +Behavior notes: +- Pinot accepts canonical RFC 4122 UUID strings in either upper or lower case on ingest and in functions/casts. +- Pinot always renders `UUID` results as canonical lowercase strings. +- `CAST(... AS UUID)` accepts canonical strings and 16-byte `BYTES` values. + +Migration notes: +- Existing `BYTES` columns keep returning hex strings. Pinot only renders canonical UUID strings for columns declared as `UUID`. +- Pinot does not support changing the data type of an existing column in place. To adopt `UUID` for existing + `STRING` or `BYTES` UUID-shaped data, create a new `UUID` column or a new table/schema and reingest/backfill the + data into it. +- The `UUID` type itself does not require a segment or wire format bump in v1, but migration still requires rebuild or + reingest because schema type mutation is unsupported. + ## License Apache Pinot is under [Apache License, Version 2.0](http://www.apache.org/licenses/LICENSE-2.0) diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java index 4ce8b26ef999..3fb01c6d701a 100644 --- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java +++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/ClusterTest.java @@ -752,6 +752,7 @@ private static JsonNode extractArray(DataSchema.ColumnDataType columnDataType, J case BIG_DECIMAL_ARRAY: case TIMESTAMP_ARRAY: case STRING_ARRAY: + case UUID_ARRAY: case BYTES_ARRAY: array[k] = jsonValue.get(k).textValue(); break; @@ -787,6 +788,7 @@ private static JsonNode extractValue(DataSchema.ColumnDataType columnDataType, J case TIMESTAMP: case STRING: case BYTES: + case UUID: case JSON: object = jsonValue.textValue(); break; diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/CustomDataQueryClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/CustomDataQueryClusterIntegrationTest.java index 21b188837e30..1903ff04984e 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/CustomDataQueryClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/CustomDataQueryClusterIntegrationTest.java @@ -140,7 +140,8 @@ protected void setUpTable() if (isRealtimeTable()) { // In suite mode multiple realtime tests use different topics, so make sure // this class-specific topic exists before the controller validates stream metadata. - _sharedClusterTestSuite.createKafkaTopic(getKafkaTopic()); + // Pass getNumKafkaPartitions() explicitly so subclass overrides are respected. + _sharedClusterTestSuite.createKafkaTopic(getKafkaTopic(), getNumKafkaPartitions()); waitForKafkaTopicMetadataReadyForConsumer(getKafkaTopic(), getNumKafkaPartitions()); // create realtime table diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeRealtimeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeRealtimeTest.java new file mode 100644 index 000000000000..72106bb2e017 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeRealtimeTest.java @@ -0,0 +1,34 @@ +/** + * 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; + + +public class UuidTypeRealtimeTest extends UuidTypeTest { + private static final String REALTIME_TABLE_NAME = "UuidTypeRealtimeTest"; + + @Override + public String getTableName() { + return REALTIME_TABLE_NAME; + } + + @Override + public boolean isRealtimeTable() { + return true; + } +} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeTest.java new file mode 100644 index 000000000000..9cfadcdccd46 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidTypeTest.java @@ -0,0 +1,540 @@ +/** + * 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.nio.ByteBuffer; +import java.util.List; +import org.apache.avro.Schema.Field; +import org.apache.avro.Schema.Type; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.UuidUtils; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.create; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertTrue; + + +@Test(suiteName = "CustomClusterIntegrationTest") +public class UuidTypeTest extends CustomDataQueryClusterIntegrationTest { + private static final String DEFAULT_TABLE_NAME = "UuidTypeTest"; + private static final String ID_COLUMN = "id"; + private static final String UUID_FROM_STRING_COLUMN = "uuidFromString"; + private static final String UUID_FROM_BYTES_COLUMN = "uuidFromBytes"; + private static final String UUID_AS_BYTES_COLUMN = "uuidAsBytes"; + private static final String UUID_ARRAY_FROM_STRING_COLUMN = "uuidArrayFromString"; + private static final String UUID_ARRAY_FROM_BYTES_COLUMN = "uuidArrayFromBytes"; + private static final String TIME_COLUMN = "ts"; + private static final List UUID_INPUT_VALUES = List.of( + "550E8400-E29B-41D4-A716-446655440000", + "550e8400-e29b-41d4-a716-446655440001", + "550e8400-E29B-41D4-A716-446655440000", + "550e8400-e29b-41d4-a716-446655440002", + "550E8400-E29B-41D4-A716-446655440001", + "550e8400-e29b-41d4-a716-446655440003"); + private static final List UUID_VALUES = List.of( + "550e8400-e29b-41d4-a716-446655440000", + "550e8400-e29b-41d4-a716-446655440001", + "550e8400-e29b-41d4-a716-446655440000", + "550e8400-e29b-41d4-a716-446655440002", + "550e8400-e29b-41d4-a716-446655440001", + "550e8400-e29b-41d4-a716-446655440003"); + private static final List DISTINCT_UUID_VALUES = List.of( + "550e8400-e29b-41d4-a716-446655440000", + "550e8400-e29b-41d4-a716-446655440001", + "550e8400-e29b-41d4-a716-446655440002", + "550e8400-e29b-41d4-a716-446655440003"); + private static final List SORTED_UUID_IDS = List.of(0, 2, 1, 4, 3, 5); + private static final long BASE_TIMESTAMP_MILLIS = 1_700_000_000_000L; + + @Override + public String getTableName() { + return DEFAULT_TABLE_NAME; + } + + @Override + public String getTimeColumnName() { + return TIME_COLUMN; + } + + @Override + protected long getCountStarResult() { + return UUID_VALUES.size(); + } + + @Override + public int getNumAvroFiles() { + return 1; + } + + @Override + protected int getRealtimeSegmentFlushSize() { + return 2; + } + + @Override + protected List getInvertedIndexColumns() { + return List.of(UUID_FROM_STRING_COLUMN); + } + + @Override + protected List getNoDictionaryColumns() { + return List.of(UUID_FROM_BYTES_COLUMN, UUID_AS_BYTES_COLUMN, UUID_ARRAY_FROM_BYTES_COLUMN); + } + + @Override + protected List getBloomFilterColumns() { + return List.of(UUID_FROM_STRING_COLUMN, UUID_FROM_BYTES_COLUMN, UUID_AS_BYTES_COLUMN); + } + + @Override + public Schema createSchema() { + return new Schema.SchemaBuilder().setSchemaName(getTableName()) + .addSingleValueDimension(ID_COLUMN, FieldSpec.DataType.INT) + .addSingleValueDimension(UUID_FROM_STRING_COLUMN, FieldSpec.DataType.UUID) + .addSingleValueDimension(UUID_FROM_BYTES_COLUMN, FieldSpec.DataType.UUID) + .addSingleValueDimension(UUID_AS_BYTES_COLUMN, FieldSpec.DataType.BYTES) + .addMultiValueDimension(UUID_ARRAY_FROM_STRING_COLUMN, FieldSpec.DataType.UUID) + .addMultiValueDimension(UUID_ARRAY_FROM_BYTES_COLUMN, FieldSpec.DataType.UUID) + .addDateTimeField(TIME_COLUMN, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .build(); + } + + @Override + public List createAvroFiles() + throws Exception { + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("uuidRecord", null, null, false); + org.apache.avro.Schema uuidStringArraySchema = org.apache.avro.Schema.createArray(create(Type.STRING)); + org.apache.avro.Schema uuidBytesArraySchema = org.apache.avro.Schema.createArray(create(Type.BYTES)); + avroSchema.setFields(List.of( + new Field(ID_COLUMN, create(Type.INT), null, null), + new Field(UUID_FROM_STRING_COLUMN, create(Type.STRING), null, null), + new Field(UUID_FROM_BYTES_COLUMN, create(Type.BYTES), null, null), + new Field(UUID_AS_BYTES_COLUMN, create(Type.BYTES), null, null), + new Field(UUID_ARRAY_FROM_STRING_COLUMN, uuidStringArraySchema, null, null), + new Field(UUID_ARRAY_FROM_BYTES_COLUMN, uuidBytesArraySchema, null, null), + new Field(TIME_COLUMN, create(Type.LONG), null, null))); + + try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) { + List> writers = avroFilesAndWriters.getWriters(); + for (int i = 0; i < UUID_VALUES.size(); i++) { + String uuidInputValue = UUID_INPUT_VALUES.get(i); + byte[] uuidBytes = UuidUtils.toBytes(uuidInputValue); + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(ID_COLUMN, i); + record.put(UUID_FROM_STRING_COLUMN, uuidInputValue); + record.put(UUID_FROM_BYTES_COLUMN, ByteBuffer.wrap(uuidBytes)); + record.put(UUID_AS_BYTES_COLUMN, ByteBuffer.wrap(uuidBytes)); + record.put(UUID_ARRAY_FROM_STRING_COLUMN, uuidStringArray(i, uuidStringArraySchema)); + record.put(UUID_ARRAY_FROM_BYTES_COLUMN, uuidBytesArray(i, uuidBytesArraySchema)); + record.put(TIME_COLUMN, BASE_TIMESTAMP_MILLIS + i); + writers.get(0).append(record); + } + return avroFilesAndWriters.getAvroFiles(); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testSelectAndPredicateQueries(boolean useMultiStageQueryEngine) + throws Exception { + JsonNode rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s, %s, %s, %s FROM %s ORDER BY %s", ID_COLUMN, UUID_FROM_STRING_COLUMN, UUID_FROM_BYTES_COLUMN, + UUID_AS_BYTES_COLUMN, getTableName(), ID_COLUMN)); + + assertEquals(rows.size(), UUID_VALUES.size()); + for (int i = 0; i < UUID_VALUES.size(); i++) { + assertEquals(rows.get(i).get(0).asInt(), i); + assertEquals(rows.get(i).get(1).asText(), UUID_VALUES.get(i)); + assertEquals(rows.get(i).get(2).asText(), UUID_VALUES.get(i)); + assertEquals(rows.get(i).get(3).asText(), uuidHex(UUID_VALUES.get(i))); + } + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE %s = CAST('%s' AS UUID) ORDER BY %s", ID_COLUMN, getTableName(), + UUID_FROM_STRING_COLUMN, UUID_INPUT_VALUES.get(1).toUpperCase(), ID_COLUMN)); + assertIdRows(rows, 1, 4); + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE %s IN (CAST('%s' AS UUID), CAST('%s' AS UUID)) ORDER BY %s", ID_COLUMN, + getTableName(), UUID_FROM_BYTES_COLUMN, DISTINCT_UUID_VALUES.get(0).toUpperCase(), DISTINCT_UUID_VALUES.get(2), + ID_COLUMN)); + assertIdRows(rows, 0, 2, 3); + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE %s = TO_UUID('%s') ORDER BY %s", ID_COLUMN, getTableName(), UUID_FROM_STRING_COLUMN, + DISTINCT_UUID_VALUES.get(3).toUpperCase(), ID_COLUMN)); + assertIdRows(rows, 5); + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE %s = UUID_TO_BYTES(CAST('%s' AS UUID)) ORDER BY %s", ID_COLUMN, getTableName(), + UUID_AS_BYTES_COLUMN, DISTINCT_UUID_VALUES.get(1), ID_COLUMN)); + assertIdRows(rows, 1, 4); + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE CAST(%s AS UUID) = TO_UUID('%s') ORDER BY %s", ID_COLUMN, getTableName(), + UUID_AS_BYTES_COLUMN, DISTINCT_UUID_VALUES.get(2).toUpperCase(), ID_COLUMN)); + assertIdRows(rows, 3); + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE %s NOT IN (CAST('%s' AS UUID), CAST('%s' AS UUID)) ORDER BY %s", ID_COLUMN, + getTableName(), UUID_FROM_STRING_COLUMN, DISTINCT_UUID_VALUES.get(0), DISTINCT_UUID_VALUES.get(3), ID_COLUMN)); + assertIdRows(rows, 1, 3, 4); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testUuidArrayProjection(boolean useMultiStageQueryEngine) + throws Exception { + JsonNode rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s, %s, %s, arrayLength(%s), arrayLength(%s) FROM %s ORDER BY %s", ID_COLUMN, + UUID_ARRAY_FROM_STRING_COLUMN, UUID_ARRAY_FROM_BYTES_COLUMN, UUID_ARRAY_FROM_STRING_COLUMN, + UUID_ARRAY_FROM_BYTES_COLUMN, getTableName(), ID_COLUMN)); + + assertEquals(rows.size(), UUID_VALUES.size()); + for (int i = 0; i < UUID_VALUES.size(); i++) { + assertEquals(rows.get(i).get(0).asInt(), i); + assertUuidArray(rows.get(i).get(1), i); + assertUuidArray(rows.get(i).get(2), i); + assertEquals(rows.get(i).get(3).asInt(), 2); + assertEquals(rows.get(i).get(4).asInt(), 2); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testUuidFunctionsAndCaseQueries(boolean useMultiStageQueryEngine) + throws Exception { + JsonNode rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT UUID_TO_STRING(TO_UUID(UUID_TO_STRING(%s))), UUID_TO_STRING(TO_UUID(%s)), " + + "UUID_TO_STRING(BYTES_TO_UUID(%s)), UUID_TO_BYTES(%s), " + + "IS_UUID(UUID_TO_STRING(%s)), IS_UUID(%s), IS_UUID('not-a-uuid') " + + "FROM %s ORDER BY %s", UUID_FROM_STRING_COLUMN, UUID_AS_BYTES_COLUMN, UUID_FROM_BYTES_COLUMN, + UUID_FROM_BYTES_COLUMN, UUID_FROM_STRING_COLUMN, UUID_AS_BYTES_COLUMN, getTableName(), ID_COLUMN)); + + assertEquals(rows.size(), UUID_VALUES.size()); + for (int i = 0; i < UUID_VALUES.size(); i++) { + assertEquals(rows.get(i).get(0).asText(), UUID_VALUES.get(i)); + assertEquals(rows.get(i).get(1).asText(), UUID_VALUES.get(i)); + assertEquals(rows.get(i).get(2).asText(), UUID_VALUES.get(i)); + assertEquals(rows.get(i).get(3).asText(), uuidHex(UUID_VALUES.get(i))); + assertEquals(rows.get(i).get(4).asBoolean(), true); + assertEquals(rows.get(i).get(5).asBoolean(), true); + assertEquals(rows.get(i).get(6).asBoolean(), false); + } + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT CAST('%s' AS UUID) FROM %s ORDER BY %s LIMIT 1", DISTINCT_UUID_VALUES.get(0).toUpperCase(), + getTableName(), ID_COLUMN)); + assertEquals(rows.size(), 1); + assertEquals(rows.get(0).get(0).asText(), DISTINCT_UUID_VALUES.get(0)); + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT CASE WHEN %s < 2 THEN CAST('%s' AS UUID) ELSE BYTES_TO_UUID(%s) END " + + "FROM %s ORDER BY %s", ID_COLUMN, DISTINCT_UUID_VALUES.get(3).toUpperCase(), UUID_AS_BYTES_COLUMN, + getTableName(), ID_COLUMN)); + assertEquals(rows.size(), UUID_VALUES.size()); + assertEquals(rows.get(0).get(0).asText(), DISTINCT_UUID_VALUES.get(3)); + assertEquals(rows.get(1).get(0).asText(), DISTINCT_UUID_VALUES.get(3)); + for (int i = 2; i < UUID_VALUES.size(); i++) { + assertEquals(rows.get(i).get(0).asText(), UUID_VALUES.get(i)); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testGroupByDistinctOrderByAndParity(boolean useMultiStageQueryEngine) + throws Exception { + JsonNode rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s, COUNT(*) FROM %s GROUP BY %s ORDER BY %s", UUID_FROM_STRING_COLUMN, getTableName(), + UUID_FROM_STRING_COLUMN, UUID_FROM_STRING_COLUMN)); + assertGroupedCounts(rows); + + JsonNode rawRows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s, COUNT(*) FROM %s GROUP BY %s ORDER BY %s", UUID_FROM_BYTES_COLUMN, getTableName(), + UUID_FROM_BYTES_COLUMN, UUID_FROM_BYTES_COLUMN)); + assertEquals(rawRows, rows); + + rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT DISTINCT BYTES_TO_UUID(%s) FROM %s ORDER BY BYTES_TO_UUID(%s)", UUID_AS_BYTES_COLUMN, getTableName(), + UUID_AS_BYTES_COLUMN)); + assertDistinctRows(rows); + + JsonNode dictRows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s, %s FROM %s ORDER BY %s, %s", UUID_FROM_STRING_COLUMN, ID_COLUMN, getTableName(), + UUID_FROM_STRING_COLUMN, ID_COLUMN)); + JsonNode rawOrderRows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s, %s FROM %s ORDER BY %s, %s", UUID_FROM_BYTES_COLUMN, ID_COLUMN, getTableName(), + UUID_FROM_BYTES_COLUMN, ID_COLUMN)); + assertEquals(rawOrderRows, dictRows); + + for (int i = 0; i < SORTED_UUID_IDS.size(); i++) { + int expectedId = SORTED_UUID_IDS.get(i); + assertEquals(dictRows.get(i).get(0).asText(), UUID_VALUES.get(expectedId)); + assertEquals(dictRows.get(i).get(1).asInt(), expectedId); + } + } + + /// Tests range predicates on UUID columns (both dictionary-backed and no-dictionary). + /// + /// UUID byte order determines sort order: the last segment of each test UUID differs only in the final byte + /// (00–03), so the four distinct values have a well-defined byte-ordered range. + /// + /// Only the single-stage engine is tested here; MSQE range pushdown for UUID columns is verified separately + /// through the SSE/MSE parity checks in [#testSseMseParity]. + @Test + public void testRangePredicates() + throws Exception { + // uuidFromString is dictionary-backed; uuidFromBytes is no-dictionary (raw bytes) + // Sorted distinct values: 440000 < 440001 < 440002 < 440003 + + // GT on dictionary-backed UUID column: > 440001 → values 440002 (id=3), 440003 (id=5) + JsonNode rows = postRows(false, String.format( + "SELECT %s FROM %s WHERE %s > '%s' ORDER BY %s", ID_COLUMN, getTableName(), + UUID_FROM_STRING_COLUMN, DISTINCT_UUID_VALUES.get(1), ID_COLUMN)); + assertIdRows(rows, 3, 5); + + // LT on dictionary-backed UUID column: < 440002 → values 440000 (ids=0,2), 440001 (ids=1,4) + rows = postRows(false, String.format( + "SELECT %s FROM %s WHERE %s < '%s' ORDER BY %s", ID_COLUMN, getTableName(), + UUID_FROM_STRING_COLUMN, DISTINCT_UUID_VALUES.get(2), ID_COLUMN)); + assertIdRows(rows, 0, 1, 2, 4); + + // BETWEEN on no-dictionary UUID column: 440001–440002 inclusive → ids 1, 3, 4 + rows = postRows(false, String.format( + "SELECT %s FROM %s WHERE %s BETWEEN '%s' AND '%s' ORDER BY %s", ID_COLUMN, getTableName(), + UUID_FROM_BYTES_COLUMN, DISTINCT_UUID_VALUES.get(1), DISTINCT_UUID_VALUES.get(2), ID_COLUMN)); + assertIdRows(rows, 1, 3, 4); + + // GTE on no-dictionary UUID column: >= 440002 → values 440002 (id=3), 440003 (id=5) + rows = postRows(false, String.format( + "SELECT %s FROM %s WHERE %s >= '%s' ORDER BY %s", ID_COLUMN, getTableName(), + UUID_FROM_BYTES_COLUMN, DISTINCT_UUID_VALUES.get(2), ID_COLUMN)); + assertIdRows(rows, 3, 5); + + // LTE on no-dictionary UUID column: <= 440001 → values 440000 (ids=0,2), 440001 (ids=1,4) + rows = postRows(false, String.format( + "SELECT %s FROM %s WHERE %s <= '%s' ORDER BY %s", ID_COLUMN, getTableName(), + UUID_FROM_BYTES_COLUMN, DISTINCT_UUID_VALUES.get(1), ID_COLUMN)); + assertIdRows(rows, 0, 1, 2, 4); + + // Mixed-case UUID bounds are normalized: upper-case variant of 440001 + rows = postRows(false, String.format( + "SELECT %s FROM %s WHERE %s > '%s' ORDER BY %s", ID_COLUMN, getTableName(), + UUID_FROM_STRING_COLUMN, DISTINCT_UUID_VALUES.get(1).toUpperCase(), ID_COLUMN)); + assertIdRows(rows, 3, 5); + } + + @Test + public void testSseMseParity() + throws Exception { + List queries = List.of( + String.format("SELECT %s, UUID_TO_STRING(BYTES_TO_UUID(%s)) AS uuidText FROM %s ORDER BY %s", + UUID_FROM_STRING_COLUMN, UUID_AS_BYTES_COLUMN, getTableName(), ID_COLUMN), + String.format("SELECT COUNT(*) AS cnt FROM %s WHERE %s = TO_UUID('%s')", getTableName(), + UUID_FROM_STRING_COLUMN, DISTINCT_UUID_VALUES.get(1).toUpperCase()), + String.format("SELECT BYTES_TO_UUID(%s) AS uuidValue, COUNT(*) AS cnt FROM %s " + + "GROUP BY BYTES_TO_UUID(%s) ORDER BY uuidValue", + UUID_AS_BYTES_COLUMN, getTableName(), UUID_AS_BYTES_COLUMN)); + + for (String query : queries) { + JsonNode sseResult = postResultTable(false, query); + JsonNode mseResult = postResultTable(true, query); + assertEquals(mseResult, sseResult, query); + } + } + + /// Regression for the DISTINCTCOUNT-family canonical-string contract: for identifier expressions the + /// aggregation receives a ProjectionBlockValSet whose getStringValuesSV() renders stored BYTES as bare hex, + /// so the UUID branches must fetch raw bytes and canonicalize themselves. Asserts + /// DISTINCTCOUNTX(uuidCol) == DISTINCTCOUNTX(CAST(uuidCol AS STRING)) end-to-end on real segments — + /// a parity that silently breaks if any branch falls back to the projection string path. + @Test(dataProvider = "useBothQueryEngines") + public void testDistinctCountFamilyUuidMatchesCastStringParity(boolean useMultiStageQueryEngine) + throws Exception { + for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTBITMAP", "DISTINCTCOUNTHLL", + "DISTINCTCOUNTHLLPLUS", "DISTINCTCOUNTULL", "DISTINCTCOUNTCPCSKETCH", "DISTINCTCOUNTTHETASKETCH")) { + String onUuid = String.format("SELECT %s(%s) FROM %s", function, UUID_FROM_STRING_COLUMN, getTableName()); + String onCastString = String.format("SELECT %s(CAST(%s AS STRING)) FROM %s", function, + UUID_FROM_STRING_COLUMN, getTableName()); + JsonNode uuidRows = postRows(useMultiStageQueryEngine, onUuid); + JsonNode castRows = postRows(useMultiStageQueryEngine, onCastString); + assertEquals(uuidRows.get(0).get(0).asLong(), castRows.get(0).get(0).asLong(), + function + " on UUID column must match the same function on CAST(uuidCol AS STRING)"); + assertEquals(uuidRows.get(0).get(0).asLong(), DISTINCT_UUID_VALUES.size(), + function + " on UUID column must count the distinct UUID values"); + } + + // Also pin the filterless dictionary-based plan (NonScanBasedAggregationOperator): a WHERE clause forces the + // scan path above on some engines, while the bare aggregates here may be served purely from the dictionary. + // Both shapes must agree for UUID columns. + JsonNode filteredRows = postRows(useMultiStageQueryEngine, String.format( + "SELECT DISTINCTCOUNTHLL(%s) FROM %s WHERE %s >= 0", UUID_FROM_STRING_COLUMN, getTableName(), ID_COLUMN)); + JsonNode unfilteredRows = postRows(useMultiStageQueryEngine, String.format( + "SELECT DISTINCTCOUNTHLL(%s) FROM %s", UUID_FROM_STRING_COLUMN, getTableName())); + assertEquals(unfilteredRows.get(0).get(0).asLong(), filteredRows.get(0).get(0).asLong(), + "Dictionary-based (filterless) and scan-based DISTINCTCOUNTHLL must agree on UUID columns"); + } + + @Test(dataProvider = "useV2QueryEngine") + public void testUuidEqualityJoinWithExpressionKey(boolean useMultiStageQueryEngine) + throws Exception { + JsonNode rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT a.%s, b.%s, a.%s " + + "FROM (SELECT %s, %s FROM %s WHERE %s < 5) a " + + "JOIN (SELECT %s, %s, %s FROM %s WHERE %s > 0) b " + + "ON a.%s = BYTES_TO_UUID(UUID_TO_BYTES(b.%s)) " + + "WHERE a.%s < b.%s " + + "ORDER BY a.%s, b.%s", + ID_COLUMN, ID_COLUMN, UUID_FROM_STRING_COLUMN, ID_COLUMN, UUID_FROM_STRING_COLUMN, getTableName(), ID_COLUMN, + ID_COLUMN, UUID_FROM_BYTES_COLUMN, UUID_AS_BYTES_COLUMN, getTableName(), ID_COLUMN, UUID_FROM_STRING_COLUMN, + UUID_FROM_BYTES_COLUMN, ID_COLUMN, ID_COLUMN, ID_COLUMN, ID_COLUMN)); + + assertEquals(rows.size(), 2); + assertEquals(rows.get(0).get(0).asInt(), 0); + assertEquals(rows.get(0).get(1).asInt(), 2); + assertEquals(rows.get(0).get(2).asText(), DISTINCT_UUID_VALUES.get(0)); + assertEquals(rows.get(1).get(0).asInt(), 1); + assertEquals(rows.get(1).get(1).asInt(), 4); + assertEquals(rows.get(1).get(2).asText(), DISTINCT_UUID_VALUES.get(1)); + } + + /// LEFT JOIN on a UUID key with unmatched left rows: ids 0 and 2 carry the only UUID (…440000) absent from the + /// right side (id >= 3), so they must surface with a NULL right id while every other row matches. Locks in + /// UuidLookupTable/HashJoinOperator key normalization for the non-INNER join paths. + @Test(dataProvider = "useV2QueryEngine") + public void testUuidLeftJoin(boolean useMultiStageQueryEngine) + throws Exception { + JsonNode rows = postRows(useMultiStageQueryEngine, String.format( + "SELECT a.%s, b.%s " + + "FROM (SELECT %s, %s FROM %s) a " + + "LEFT JOIN (SELECT %s, %s FROM %s WHERE %s >= 3) b " + + "ON a.%s = b.%s " + + "ORDER BY a.%s, b.%s", + ID_COLUMN, ID_COLUMN, + ID_COLUMN, UUID_FROM_STRING_COLUMN, getTableName(), + ID_COLUMN, UUID_FROM_STRING_COLUMN, getTableName(), ID_COLUMN, + UUID_FROM_STRING_COLUMN, UUID_FROM_STRING_COLUMN, + ID_COLUMN, ID_COLUMN)); + + // Left ids 0 and 2 (uuid …440000) have no right match; 1 and 4 match right id 4 (…440001); 3 matches 3; + // 5 matches 5. + assertEquals(rows.size(), 6); + assertEquals(rows.get(0).get(0).asInt(), 0); + assertTrue(rows.get(0).get(1).isNull(), "Left row with unmatched UUID must produce NULL right id"); + assertEquals(rows.get(1).get(0).asInt(), 1); + assertEquals(rows.get(1).get(1).asInt(), 4); + assertEquals(rows.get(2).get(0).asInt(), 2); + assertTrue(rows.get(2).get(1).isNull(), "Left row with unmatched UUID must produce NULL right id"); + assertEquals(rows.get(3).get(0).asInt(), 3); + assertEquals(rows.get(3).get(1).asInt(), 3); + assertEquals(rows.get(4).get(0).asInt(), 4); + assertEquals(rows.get(4).get(1).asInt(), 4); + assertEquals(rows.get(5).get(0).asInt(), 5); + assertEquals(rows.get(5).get(1).asInt(), 5); + } + + /// SEMI and ANTI joins on a UUID key (planned from IN / NOT IN subqueries in the multi-stage engine). The right + /// side carries UUIDs of ids >= 3 ({…440001, …440002, …440003}); ids 1/3/4/5 are semi-matched and ids 0/2 + /// (uuid …440000) are anti-matched. + @Test(dataProvider = "useV2QueryEngine") + public void testUuidSemiAndAntiJoin(boolean useMultiStageQueryEngine) + throws Exception { + JsonNode semiRows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE %s IN (SELECT %s FROM %s WHERE %s >= 3) ORDER BY %s", + ID_COLUMN, getTableName(), UUID_FROM_STRING_COLUMN, + UUID_FROM_STRING_COLUMN, getTableName(), ID_COLUMN, ID_COLUMN)); + assertIdRows(semiRows, 1, 3, 4, 5); + + JsonNode antiRows = postRows(useMultiStageQueryEngine, String.format( + "SELECT %s FROM %s WHERE %s NOT IN (SELECT %s FROM %s WHERE %s >= 3) ORDER BY %s", + ID_COLUMN, getTableName(), UUID_FROM_STRING_COLUMN, + UUID_FROM_STRING_COLUMN, getTableName(), ID_COLUMN, ID_COLUMN)); + assertIdRows(antiRows, 0, 2); + } + + private JsonNode postRows(boolean useMultiStageQueryEngine, String query) + throws Exception { + return postResultTable(useMultiStageQueryEngine, query).get("rows"); + } + + private JsonNode postResultTable(boolean useMultiStageQueryEngine, String query) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + JsonNode response = postQuery(query); + assertNoExceptions(response); + return response.get("resultTable"); + } + + private static void assertGroupedCounts(JsonNode rows) { + assertEquals(rows.size(), DISTINCT_UUID_VALUES.size()); + assertEquals(rows.get(0).get(0).asText(), DISTINCT_UUID_VALUES.get(0)); + assertEquals(rows.get(0).get(1).asInt(), 2); + assertEquals(rows.get(1).get(0).asText(), DISTINCT_UUID_VALUES.get(1)); + assertEquals(rows.get(1).get(1).asInt(), 2); + assertEquals(rows.get(2).get(0).asText(), DISTINCT_UUID_VALUES.get(2)); + assertEquals(rows.get(2).get(1).asInt(), 1); + assertEquals(rows.get(3).get(0).asText(), DISTINCT_UUID_VALUES.get(3)); + assertEquals(rows.get(3).get(1).asInt(), 1); + } + + private static void assertDistinctRows(JsonNode rows) { + assertEquals(rows.size(), DISTINCT_UUID_VALUES.size()); + for (int i = 0; i < DISTINCT_UUID_VALUES.size(); i++) { + assertEquals(rows.get(i).get(0).asText(), DISTINCT_UUID_VALUES.get(i)); + } + } + + private static void assertIdRows(JsonNode rows, int... expectedIds) { + assertEquals(rows.size(), expectedIds.length); + for (int i = 0; i < expectedIds.length; i++) { + assertEquals(rows.get(i).get(0).asInt(), expectedIds[i]); + } + } + + private static GenericData.Array uuidStringArray(int rowIndex, org.apache.avro.Schema schema) { + GenericData.Array array = new GenericData.Array<>(2, schema); + array.add(UUID_INPUT_VALUES.get(rowIndex)); + array.add(UUID_INPUT_VALUES.get((rowIndex + 1) % UUID_INPUT_VALUES.size())); + return array; + } + + private static GenericData.Array uuidBytesArray(int rowIndex, org.apache.avro.Schema schema) { + GenericData.Array array = new GenericData.Array<>(2, schema); + array.add(ByteBuffer.wrap(UuidUtils.toBytes(UUID_INPUT_VALUES.get(rowIndex)))); + array.add(ByteBuffer.wrap(UuidUtils.toBytes(UUID_INPUT_VALUES.get((rowIndex + 1) % UUID_INPUT_VALUES.size())))); + return array; + } + + private static void assertUuidArray(JsonNode array, int rowIndex) { + assertEquals(array.size(), 2); + assertEquals(array.get(0).asText(), UUID_VALUES.get(rowIndex)); + assertEquals(array.get(1).asText(), UUID_VALUES.get((rowIndex + 1) % UUID_VALUES.size())); + } + + private static String uuidHex(String uuidValue) { + return BytesUtils.toHexString(UuidUtils.toBytes(uuidValue)); + } + + private static void assertNoExceptions(JsonNode response) { + assertEquals(response.get("exceptions").size(), 0, response.toPrettyString()); + } +} diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidUpsertRealtimeTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidUpsertRealtimeTest.java new file mode 100644 index 000000000000..7a56b77762d1 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidUpsertRealtimeTest.java @@ -0,0 +1,337 @@ +/** + * 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.List; +import java.util.Map; +import java.util.concurrent.atomic.AtomicReference; +import org.apache.avro.Schema.Field; +import org.apache.avro.Schema.Type; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericData; +import org.apache.pinot.integration.tests.ClusterTest.AvroFileSchemaKafkaAvroMessageDecoder; +import org.apache.pinot.spi.config.table.ColumnPartitionConfig; +import org.apache.pinot.spi.config.table.ReplicaGroupStrategyConfig; +import org.apache.pinot.spi.config.table.RoutingConfig; +import org.apache.pinot.spi.config.table.SegmentPartitionConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.config.table.UpsertConfig; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.util.TestUtils; +import org.testng.annotations.Test; + +import static org.apache.avro.Schema.create; +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertNotNull; + +/// Realtime upsert coverage for UUID primary keys. +/// +/// This test uses a single Kafka partition. UUID primary keys must be in canonical lowercase form (enforced at +/// ingestion time by `DataTypeTransformer`) because Kafka partition routing is determined by the raw string value +/// before Pinot normalization; non-canonical values would cause silent dedup failures in multi-partition tables. +@Test(suiteName = "CustomClusterIntegrationTest") +public class UuidUpsertRealtimeTest extends CustomDataQueryClusterIntegrationTest { + private static final long DEDUPLICATED_RECORDS_READY_TIMEOUT_MS = 120_000L; + private static final String TABLE_NAME = "UuidUpsertRealtimeTest"; + private static final String UUID_PK_COLUMN = "uuidPk"; + private static final String PAYLOAD_COLUMN = "payload"; + private static final String TIME_COLUMN = "ts"; + // All primary key values must be canonical lowercase UUIDs. DataTypeTransformer rejects non-canonical + // (uppercase) UUID primary keys for upsert tables to prevent silent dedup failures when the same + // logical UUID is routed to different Kafka partitions before normalization. + private static final List UUID_INPUT_VALUES = List.of( + "550e8400-e29b-41d4-a716-446655440000", + "550e8400-e29b-41d4-a716-446655440001", + "550e8400-e29b-41d4-a716-446655440000", // duplicate of index 0, triggers upsert dedup + "550e8400-e29b-41d4-a716-446655440002", + "550e8400-e29b-41d4-a716-446655440001"); // duplicate of index 1, triggers upsert dedup + private static final List PAYLOAD_VALUES = List.of("alpha-v1", "beta-v1", "alpha-v2", "gamma-v1", + "beta-v2"); + private static final List EXPECTED_UUID_VALUES = List.of( + "550e8400-e29b-41d4-a716-446655440000", + "550e8400-e29b-41d4-a716-446655440001", + "550e8400-e29b-41d4-a716-446655440002"); + private static final List EXPECTED_PAYLOAD_VALUES = List.of("alpha-v2", "beta-v2", "gamma-v1"); + private static final String COUNT_QUERY = String.format("SELECT COUNT(*) FROM %s", TABLE_NAME); + private static final String RAW_COUNT_QUERY = + String.format("SELECT COUNT(*) FROM %s OPTION(skipUpsert=true)", TABLE_NAME); + private static final String ORDERED_ROWS_QUERY = + String.format("SELECT %s, %s FROM %s ORDER BY %s", UUID_PK_COLUMN, PAYLOAD_COLUMN, TABLE_NAME, UUID_PK_COLUMN); + private static final String FILTER_QUERY = String.format("SELECT %s FROM %s WHERE %s = TO_UUID('%s')", + PAYLOAD_COLUMN, TABLE_NAME, UUID_PK_COLUMN, UUID_INPUT_VALUES.get(0)); + private static final int TOTAL_RAW_RECORDS = UUID_INPUT_VALUES.size(); + private static final long BASE_TIMESTAMP_MILLIS = 1_700_100_000_000L; + + @Override + public String getTableName() { + return TABLE_NAME; + } + + @Override + public boolean isRealtimeTable() { + return true; + } + + @Override + public String getTimeColumnName() { + return TIME_COLUMN; + } + + @Override + protected UpsertConfig getUpsertConfig() { + return new UpsertConfig(UpsertConfig.Mode.FULL); + } + + @Override + protected RoutingConfig getRoutingConfig() { + return new RoutingConfig(null, null, RoutingConfig.STRICT_REPLICA_GROUP_INSTANCE_SELECTOR_TYPE, false); + } + + @Override + protected long getCountStarResult() { + return EXPECTED_UUID_VALUES.size(); + } + + @Override + protected int getNumKafkaPartitions() { + return 1; + } + + @Override + public int getNumAvroFiles() { + return 1; + } + + @Override + protected int getRealtimeSegmentFlushSize() { + return 2; + } + + @Override + protected TableConfig createRealtimeTableConfig(File sampleAvroFile) { + AvroFileSchemaKafkaAvroMessageDecoder._avroFile = sampleAvroFile; + SegmentPartitionConfig segmentPartitionConfig = new SegmentPartitionConfig( + Map.of(UUID_PK_COLUMN, new ColumnPartitionConfig("Murmur", getNumKafkaPartitions()))); + return getTableConfigBuilder(TableType.REALTIME) + .setSegmentPartitionConfig(segmentPartitionConfig) + .setReplicaGroupStrategyConfig(new ReplicaGroupStrategyConfig(UUID_PK_COLUMN, 1)) + .build(); + } + + @Override + protected void waitForAllDocsLoaded(long timeoutMs) + throws Exception { + TestUtils.waitForCondition(aVoid -> { + try { + return queryCountStarWithoutUpsert() == TOTAL_RAW_RECORDS; + } catch (Exception e) { + return null; + } + }, 100L, timeoutMs, "Failed to load raw UUID upsert records"); + } + + @Override + public Schema createSchema() { + return new Schema.SchemaBuilder().setSchemaName(getTableName()) + .addSingleValueDimension(UUID_PK_COLUMN, FieldSpec.DataType.UUID) + .addSingleValueDimension(PAYLOAD_COLUMN, FieldSpec.DataType.STRING) + .addDateTimeField(TIME_COLUMN, FieldSpec.DataType.LONG, "1:MILLISECONDS:EPOCH", "1:MILLISECONDS") + .setPrimaryKeyColumns(List.of(UUID_PK_COLUMN)) + .build(); + } + + @Override + public List createAvroFiles() + throws Exception { + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("uuidUpsertRecord", null, null, false); + avroSchema.setFields(List.of( + new Field(UUID_PK_COLUMN, create(Type.STRING), null, null), + new Field(PAYLOAD_COLUMN, create(Type.STRING), null, null), + new Field(TIME_COLUMN, create(Type.LONG), null, null))); + + try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) { + List> writers = avroFilesAndWriters.getWriters(); + for (int i = 0; i < UUID_INPUT_VALUES.size(); i++) { + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(UUID_PK_COLUMN, UUID_INPUT_VALUES.get(i)); + record.put(PAYLOAD_COLUMN, PAYLOAD_VALUES.get(i)); + record.put(TIME_COLUMN, BASE_TIMESTAMP_MILLIS + i); + writers.get(0).append(record); + } + return avroFilesAndWriters.getAvroFiles(); + } + } + + @Test(dataProvider = "useBothQueryEngines") + public void testUuidPrimaryKeyUpsertDedup(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + UuidQueryResults queryResults = waitForExpectedQueryResults(DEDUPLICATED_RECORDS_READY_TIMEOUT_MS); + + assertExpectedCountResponse(queryResults._countResponse, EXPECTED_UUID_VALUES.size()); + assertExpectedCountResponse(queryResults._rawCountResponse, TOTAL_RAW_RECORDS); + assertExpectedOrderedRowsResponse(queryResults._orderedRowsResponse); + assertExpectedFilterResponse(queryResults._filterResponse); + } + + private static void assertNoExceptions(JsonNode response) { + assertEquals(getExceptionCount(response), 0, response.toPrettyString()); + } + + private void assertExpectedCountResponse(JsonNode response, long expectedCount) { + assertNoExceptions(response); + JsonNode firstRowFirstColumn = getFirstRowFirstColumn(response); + assertNotNull(firstRowFirstColumn, response.toPrettyString()); + assertEquals(firstRowFirstColumn.asLong(), expectedCount); + } + + private void assertExpectedOrderedRowsResponse(JsonNode response) { + assertNoExceptions(response); + JsonNode rows = getRows(response); + assertNotNull(rows, response.toPrettyString()); + assertEquals(rows.size(), EXPECTED_UUID_VALUES.size()); + for (int i = 0; i < EXPECTED_UUID_VALUES.size(); i++) { + assertEquals(rows.get(i).get(0).asText(), EXPECTED_UUID_VALUES.get(i)); + assertEquals(rows.get(i).get(1).asText(), EXPECTED_PAYLOAD_VALUES.get(i)); + } + } + + private void assertExpectedFilterResponse(JsonNode response) { + assertNoExceptions(response); + JsonNode firstRowFirstColumn = getFirstRowFirstColumn(response); + assertNotNull(firstRowFirstColumn, response.toPrettyString()); + assertEquals(firstRowFirstColumn.asText(), "alpha-v2"); + } + + private UuidQueryResults waitForExpectedQueryResults(long timeoutMs) + throws Exception { + AtomicReference queryResultsRef = new AtomicReference<>(); + TestUtils.waitForCondition(aVoid -> { + try { + UuidQueryResults queryResults = fetchQueryResults(); + if (matchesExpectedQueryResults(queryResults)) { + queryResultsRef.set(queryResults); + return true; + } + return false; + } catch (Exception e) { + return null; + } + }, 100L, timeoutMs, "Failed to observe expected UUID upsert query results"); + return queryResultsRef.get(); + } + + private UuidQueryResults fetchQueryResults() + throws Exception { + return new UuidQueryResults(postQuery(COUNT_QUERY), postQuery(RAW_COUNT_QUERY), postQuery(ORDERED_ROWS_QUERY), + postQuery(FILTER_QUERY)); + } + + private boolean matchesExpectedQueryResults(UuidQueryResults queryResults) { + return hasExpectedCount(queryResults._countResponse, EXPECTED_UUID_VALUES.size()) + && hasExpectedCount(queryResults._rawCountResponse, TOTAL_RAW_RECORDS) + && hasExpectedOrderedRows(queryResults._orderedRowsResponse) + && hasExpectedFilterValue(queryResults._filterResponse, "alpha-v2"); + } + + private boolean hasExpectedCount(JsonNode response, long expectedCount) { + JsonNode firstRowFirstColumn = getFirstRowFirstColumn(response); + return hasNoExceptions(response) && firstRowFirstColumn != null + && firstRowFirstColumn.asLong(Long.MIN_VALUE) == expectedCount; + } + + private boolean hasExpectedOrderedRows(JsonNode response) { + if (!hasNoExceptions(response)) { + return false; + } + JsonNode rows = getRows(response); + if (rows == null) { + return false; + } + if (rows.size() != EXPECTED_UUID_VALUES.size()) { + return false; + } + for (int i = 0; i < EXPECTED_UUID_VALUES.size(); i++) { + JsonNode row = rows.get(i); + if (row == null || row.size() < 2 || !EXPECTED_UUID_VALUES.get(i).equals(row.get(0).asText()) + || !EXPECTED_PAYLOAD_VALUES.get(i).equals(row.get(1).asText())) { + return false; + } + } + return true; + } + + private boolean hasExpectedFilterValue(JsonNode response, String expectedValue) { + JsonNode firstRowFirstColumn = getFirstRowFirstColumn(response); + return hasNoExceptions(response) && firstRowFirstColumn != null + && expectedValue.equals(firstRowFirstColumn.asText()); + } + + private static boolean hasNoExceptions(JsonNode response) { + return getExceptionCount(response) == 0; + } + + private static int getExceptionCount(JsonNode response) { + JsonNode exceptions = response.get("exceptions"); + return exceptions != null ? exceptions.size() : 0; + } + + private static JsonNode getRows(JsonNode response) { + JsonNode resultTable = response.get("resultTable"); + if (resultTable == null) { + return null; + } + return resultTable.get("rows"); + } + + private static JsonNode getFirstRowFirstColumn(JsonNode response) { + JsonNode rows = getRows(response); + if (rows == null || rows.isEmpty()) { + return null; + } + JsonNode firstRow = rows.get(0); + if (firstRow == null || firstRow.isEmpty()) { + return null; + } + return firstRow.get(0); + } + + private long queryCountStarWithoutUpsert() { + return getPinotConnection().execute(RAW_COUNT_QUERY).getResultSet(0).getLong(0); + } + + private static final class UuidQueryResults { + private final JsonNode _countResponse; + private final JsonNode _rawCountResponse; + private final JsonNode _orderedRowsResponse; + private final JsonNode _filterResponse; + + private UuidQueryResults(JsonNode countResponse, JsonNode rawCountResponse, JsonNode orderedRowsResponse, + JsonNode filterResponse) { + _countResponse = countResponse; + _rawCountResponse = rawCountResponse; + _orderedRowsResponse = orderedRowsResponse; + _filterResponse = filterResponse; + } + } +} diff --git a/pinot-perf/pom.xml b/pinot-perf/pom.xml index 4009800d0c21..11e15b57bf96 100644 --- a/pinot-perf/pom.xml +++ b/pinot-perf/pom.xml @@ -49,6 +49,10 @@ org.apache.pinot pinot-server + + org.apache.pinot + pinot-query-runtime + org.apache.pinot pinot-kafka-3.0 diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidGroupingAndLookup.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidGroupingAndLookup.java new file mode 100644 index 000000000000..0479e4d45d17 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidGroupingAndLookup.java @@ -0,0 +1,273 @@ +/** + * 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.perf; + +import it.unimi.dsi.fastutil.objects.Object2IntOpenHashMap; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import org.apache.pinot.query.runtime.operator.groupby.OneObjectKeyGroupIdGenerator; +import org.apache.pinot.query.runtime.operator.groupby.OneUuidKeyGroupIdGenerator; +import org.apache.pinot.query.runtime.operator.join.ObjectLookupTable; +import org.apache.pinot.query.runtime.operator.join.UuidLookupTable; +import org.apache.pinot.spi.utils.ByteArray; +import org.apache.pinot.spi.utils.UuidKey; +import org.apache.pinot.spi.utils.UuidUtils; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// Benchmarks UUID grouping and lookup hot paths using Pinot's current [ByteArray]-based representation versus a +/// two-long UUID key representation and [UUID]. +/// +/// The benchmark is intentionally engine-adjacent: +/// +/// - V1 grouping uses the same `Object2IntOpenHashMap` shape as no-dictionary UUID group-by. +/// - V2 grouping uses [OneUuidKeyGroupIdGenerator]. +/// - Lookup uses [UuidLookupTable], which is the UUID-specific path used by MSE hash join. +@OutputTimeUnit(TimeUnit.MICROSECONDS) +@Fork(1) +@Warmup(iterations = 3, time = 5) +@Measurement(iterations = 5, time = 5) +@State(Scope.Benchmark) +public class BenchmarkUuidGroupingAndLookup { + private static final long RANDOM_SEED = 677_280_899_123L; + + @Param({"262144"}) + public int _numRows; + + @Param({"65536"}) + public int _cardinality; + + @Param({"8192"}) + public int _numProbes; + + private byte[][] _rowUuidBytes; + private ByteArray[] _rowByteArrays; + private UuidKey[] _rowUuidKeys; + private UUID[] _rowJavaUuids; + private ByteArray[] _distinctByteArrays; + private UuidKey[] _distinctUuidKeys; + private UUID[] _distinctJavaUuids; + private Object[][] _distinctByteArrayRows; + private Object[][] _distinctJavaUuidRows; + private ByteArray[] _probeByteArrays; + private UUID[] _probeJavaUuids; + + @Setup + public void setUp() { + Random random = new Random(RANDOM_SEED); + + byte[][] distinctUuidBytes = new byte[_cardinality][]; + _distinctByteArrays = new ByteArray[_cardinality]; + _distinctUuidKeys = new UuidKey[_cardinality]; + _distinctJavaUuids = new UUID[_cardinality]; + _distinctByteArrayRows = new Object[_cardinality][]; + _distinctJavaUuidRows = new Object[_cardinality][]; + + for (int i = 0; i < _cardinality; i++) { + UUID uuid = new UUID(random.nextLong(), random.nextLong()); + byte[] uuidBytes = UuidUtils.toBytes(uuid); + + distinctUuidBytes[i] = uuidBytes; + _distinctByteArrays[i] = new ByteArray(uuidBytes); + _distinctUuidKeys[i] = UuidKey.fromLongs(uuid.getMostSignificantBits(), uuid.getLeastSignificantBits()); + _distinctJavaUuids[i] = uuid; + _distinctByteArrayRows[i] = new Object[]{i}; + _distinctJavaUuidRows[i] = new Object[]{i}; + } + + _rowUuidBytes = new byte[_numRows][]; + _rowByteArrays = new ByteArray[_numRows]; + _rowUuidKeys = new UuidKey[_numRows]; + _rowJavaUuids = new UUID[_numRows]; + + for (int i = 0; i < _numRows; i++) { + int dictId = random.nextInt(_cardinality); + _rowUuidBytes[i] = distinctUuidBytes[dictId]; + _rowByteArrays[i] = _distinctByteArrays[dictId]; + _rowUuidKeys[i] = _distinctUuidKeys[dictId]; + _rowJavaUuids[i] = _distinctJavaUuids[dictId]; + } + + _probeByteArrays = new ByteArray[_numProbes]; + _probeJavaUuids = new UUID[_numProbes]; + for (int i = 0; i < _numProbes; i++) { + int dictId = random.nextInt(_cardinality); + _probeByteArrays[i] = _distinctByteArrays[dictId]; + _probeJavaUuids[i] = _distinctJavaUuids[dictId]; + } + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int v1CurrentByteArrayGrouping() { + Object2IntOpenHashMap groupIdMap = new Object2IntOpenHashMap<>(_cardinality); + groupIdMap.defaultReturnValue(-1); + + int checksum = 0; + for (int i = 0; i < _numRows; i++) { + checksum += getOrCreateGroupId(groupIdMap, new ByteArray(_rowUuidBytes[i])); + } + return checksum + groupIdMap.size(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int v1UuidKeyGrouping() { + Object2IntOpenHashMap groupIdMap = new Object2IntOpenHashMap<>(_cardinality); + groupIdMap.defaultReturnValue(-1); + + int checksum = 0; + for (int i = 0; i < _numRows; i++) { + checksum += getOrCreateGroupId(groupIdMap, _rowUuidKeys[i]); + } + return checksum + groupIdMap.size(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int v1JavaUuidGrouping() { + Object2IntOpenHashMap groupIdMap = new Object2IntOpenHashMap<>(_cardinality); + groupIdMap.defaultReturnValue(-1); + + int checksum = 0; + for (int i = 0; i < _numRows; i++) { + checksum += getOrCreateGroupId(groupIdMap, _rowJavaUuids[i]); + } + return checksum + groupIdMap.size(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int v2CurrentByteArrayGroupIdGenerator() { + OneObjectKeyGroupIdGenerator groupIdGenerator = new OneObjectKeyGroupIdGenerator(_cardinality, _cardinality); + + int checksum = 0; + for (ByteArray rowByteArray : _rowByteArrays) { + checksum += groupIdGenerator.getGroupId(rowByteArray); + } + return checksum + groupIdGenerator.getNumGroups(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int v2UuidKeyGroupIdGenerator() { + OneUuidKeyGroupIdGenerator groupIdGenerator = new OneUuidKeyGroupIdGenerator(_cardinality, _cardinality); + + int checksum = 0; + for (UuidKey rowUuidKey : _rowUuidKeys) { + checksum += groupIdGenerator.getGroupId(rowUuidKey); + } + return checksum + groupIdGenerator.getNumGroups(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int v2JavaUuidGroupIdGenerator() { + OneObjectKeyGroupIdGenerator groupIdGenerator = new OneObjectKeyGroupIdGenerator(_cardinality, _cardinality); + + int checksum = 0; + for (UUID rowJavaUuid : _rowJavaUuids) { + checksum += groupIdGenerator.getGroupId(rowJavaUuid); + } + return checksum + groupIdGenerator.getNumGroups(); + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int currentObjectLookupTableBuildAndProbe() { + ObjectLookupTable lookupTable = new ObjectLookupTable(); + for (int i = 0; i < _cardinality; i++) { + lookupTable.addRow(_distinctByteArrays[i], _distinctByteArrayRows[i]); + } + lookupTable.finish(); + + int checksum = lookupTable.size(); + for (ByteArray probeByteArray : _probeByteArrays) { + Object[] row = (Object[]) lookupTable.lookup(probeByteArray); + if (row != null) { + checksum += (int) row[0]; + } + } + return checksum; + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int uuidLookupTableBuildAndProbe() { + UuidLookupTable lookupTable = new UuidLookupTable(); + for (int i = 0; i < _cardinality; i++) { + lookupTable.addRow(_distinctByteArrays[i], _distinctByteArrayRows[i]); + } + lookupTable.finish(); + + int checksum = lookupTable.size(); + for (ByteArray probeByteArray : _probeByteArrays) { + Object[] row = (Object[]) lookupTable.lookup(probeByteArray); + if (row != null) { + checksum += (int) row[0]; + } + } + return checksum; + } + + @Benchmark + @BenchmarkMode(Mode.AverageTime) + public int javaUuidObjectLookupTableBuildAndProbe() { + ObjectLookupTable lookupTable = new ObjectLookupTable(); + for (int i = 0; i < _cardinality; i++) { + lookupTable.addRow(_distinctJavaUuids[i], _distinctJavaUuidRows[i]); + } + lookupTable.finish(); + + int checksum = lookupTable.size(); + for (UUID probeJavaUuid : _probeJavaUuids) { + Object[] row = (Object[]) lookupTable.lookup(probeJavaUuid); + if (row != null) { + checksum += (int) row[0]; + } + } + return checksum; + } + + private int getOrCreateGroupId(Object2IntOpenHashMap groupIdMap, T key) { + int numGroups = groupIdMap.size(); + if (numGroups < _cardinality) { + return groupIdMap.computeIfAbsent(key, ignored -> numGroups); + } + return groupIdMap.getInt(key); + } + + public static void main(String[] args) + throws Exception { + new Runner(new OptionsBuilder().include(BenchmarkUuidGroupingAndLookup.class.getSimpleName()).build()).run(); + } +} diff --git a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidQueryExecution.java b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidQueryExecution.java new file mode 100644 index 000000000000..27254a249540 --- /dev/null +++ b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkUuidQueryExecution.java @@ -0,0 +1,360 @@ +/** + * 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.perf; + +import java.io.File; +import java.util.ArrayList; +import java.util.List; +import java.util.Random; +import java.util.UUID; +import java.util.concurrent.TimeUnit; +import org.apache.commons.io.FileUtils; +import org.apache.pinot.common.response.broker.BrokerResponseNative; +import org.apache.pinot.queries.BaseQueriesTest; +import org.apache.pinot.segment.local.indexsegment.immutable.ImmutableSegmentLoader; +import org.apache.pinot.segment.local.segment.creator.impl.SegmentIndexCreationDriverImpl; +import org.apache.pinot.segment.local.segment.index.loader.IndexLoadingConfig; +import org.apache.pinot.segment.local.segment.readers.GenericRowRecordReader; +import org.apache.pinot.segment.spi.IndexSegment; +import org.apache.pinot.segment.spi.creator.SegmentGeneratorConfig; +import org.apache.pinot.spi.config.table.TableConfig; +import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.data.FieldSpec; +import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.data.readers.GenericRow; +import org.apache.pinot.spi.utils.BytesUtils; +import org.apache.pinot.spi.utils.UuidUtils; +import org.apache.pinot.spi.utils.builder.TableConfigBuilder; +import org.openjdk.jmh.annotations.Benchmark; +import org.openjdk.jmh.annotations.BenchmarkMode; +import org.openjdk.jmh.annotations.Fork; +import org.openjdk.jmh.annotations.Measurement; +import org.openjdk.jmh.annotations.Mode; +import org.openjdk.jmh.annotations.OutputTimeUnit; +import org.openjdk.jmh.annotations.Param; +import org.openjdk.jmh.annotations.Scope; +import org.openjdk.jmh.annotations.Setup; +import org.openjdk.jmh.annotations.State; +import org.openjdk.jmh.annotations.TearDown; +import org.openjdk.jmh.annotations.Warmup; +import org.openjdk.jmh.runner.Runner; +import org.openjdk.jmh.runner.options.OptionsBuilder; + + +/// Compares query execution performance for UUID values stored in three representations: +/// +/// - **STRING** — canonical RFC 4122 form, e.g. `"550e8400-e29b-41d4-a716-446655440000"` (36-char string key in +/// group-by). +/// - **BYTES** — raw 16-byte binary; group-by key is a `ByteArray` (heap allocation per row). +/// - **UUID** — native `DataType.UUID`; group-by key is a `UuidKey` (two `long` fields, no heap allocation). +/// +/// Each representation is tested with both raw (no-dictionary) and dictionary-encoded columns. The benchmarked +/// operations are GROUP BY, COUNT(DISTINCT), and equality-filter COUNT(*). +/// +/// Run with: +/// ``` +/// ./mvnw package -DskipTests -pl pinot-perf -am -Ppinot-fastdev +/// java -jar pinot-perf/target/benchmarks.jar BenchmarkUuidQueryExecution +/// ``` +@BenchmarkMode(Mode.AverageTime) +@OutputTimeUnit(TimeUnit.MILLISECONDS) +@Fork(1) +@Warmup(iterations = 3, time = 3) +@Measurement(iterations = 5, time = 3) +@State(Scope.Benchmark) +public class BenchmarkUuidQueryExecution extends BaseQueriesTest { + private static final File INDEX_DIR = + new File(FileUtils.getTempDirectory(), "BenchmarkUuidQueryExecution"); + private static final String TABLE_NAME = "uuidBenchTable"; + private static final String SEGMENT_NAME = "uuidBenchSegment"; + + // STRING columns: UUID in canonical "8-4-4-4-12" dash-separated form + private static final String STR_UUID_RAW = "str_uuid_raw"; + private static final String STR_UUID_DICT = "str_uuid_dict"; + // BYTES columns: UUID as raw 16-byte binary; group-by uses ByteArray + private static final String BYTES_UUID_RAW = "bytes_uuid_raw"; + private static final String BYTES_UUID_DICT = "bytes_uuid_dict"; + // UUID columns: native DataType.UUID; group-by uses UuidKey (two longs) + private static final String UUID_RAW = "uuid_raw"; + private static final String UUID_DICT = "uuid_dict"; + // Two-LONG columns: UUID split into MSB + LSB; GROUP BY both simultaneously + private static final String LONG_MSB_RAW = "long_msb_raw"; + private static final String LONG_LSB_RAW = "long_lsb_raw"; + private static final String LONG_MSB_DICT = "long_msb_dict"; + private static final String LONG_LSB_DICT = "long_lsb_dict"; + + private static final TableConfig TABLE_CONFIG = new TableConfigBuilder(TableType.OFFLINE) + .setTableName(TABLE_NAME) + .setNoDictionaryColumns(List.of(STR_UUID_RAW, BYTES_UUID_RAW, UUID_RAW, LONG_MSB_RAW, LONG_LSB_RAW)) + .build(); + + private static final Schema SCHEMA = new Schema.SchemaBuilder() + .addSingleValueDimension(STR_UUID_RAW, FieldSpec.DataType.STRING) + .addSingleValueDimension(STR_UUID_DICT, FieldSpec.DataType.STRING) + .addSingleValueDimension(BYTES_UUID_RAW, FieldSpec.DataType.BYTES) + .addSingleValueDimension(BYTES_UUID_DICT, FieldSpec.DataType.BYTES) + .addSingleValueDimension(UUID_RAW, FieldSpec.DataType.UUID) + .addSingleValueDimension(UUID_DICT, FieldSpec.DataType.UUID) + .addSingleValueDimension(LONG_MSB_RAW, FieldSpec.DataType.LONG) + .addSingleValueDimension(LONG_LSB_RAW, FieldSpec.DataType.LONG) + .addSingleValueDimension(LONG_MSB_DICT, FieldSpec.DataType.LONG) + .addSingleValueDimension(LONG_LSB_DICT, FieldSpec.DataType.LONG) + .build(); + + /// Total rows in the segment. + @Param("500000") + private int _numRows; + + /// Number of distinct UUID values. Controls group-by cardinality. + @Param("1000") + private int _numUniqueUuids; + + private IndexSegment _indexSegment; + private List _indexSegments; + + /// UUID string used in equality-filter benchmarks (RFC 4122 form). + private String _filterUuidString; + /// Same UUID as hex (no dashes) for BYTES column equality filter. + private String _filterBytesHex; + /// MSB of the filter UUID, for two-long filter benchmarks. + private long _filterMsb; + /// LSB of the filter UUID, for two-long filter benchmarks. + private long _filterLsb; + + @Setup + public void setUp() + throws Exception { + FileUtils.deleteQuietly(INDEX_DIR); + INDEX_DIR.mkdirs(); + + Random random = new Random(42L); + + // Pre-generate the UUID value pool + String[] uuidStrings = new String[_numUniqueUuids]; + byte[][] uuidBytesArr = new byte[_numUniqueUuids][]; + long[] msbArr = new long[_numUniqueUuids]; + long[] lsbArr = new long[_numUniqueUuids]; + for (int i = 0; i < _numUniqueUuids; i++) { + UUID uuid = new UUID(random.nextLong(), random.nextLong()); + uuidStrings[i] = uuid.toString(); + uuidBytesArr[i] = UuidUtils.toBytes(uuid); + msbArr[i] = uuid.getMostSignificantBits(); + lsbArr[i] = uuid.getLeastSignificantBits(); + } + _filterUuidString = uuidStrings[0]; + _filterBytesHex = BytesUtils.toHexString(uuidBytesArr[0]); + _filterMsb = msbArr[0]; + _filterLsb = lsbArr[0]; + + // Build rows with uniform distribution over the UUID pool + List rows = new ArrayList<>(_numRows); + for (int i = 0; i < _numRows; i++) { + int idx = i % _numUniqueUuids; + GenericRow row = new GenericRow(); + row.putValue(STR_UUID_RAW, uuidStrings[idx]); + row.putValue(STR_UUID_DICT, uuidStrings[idx]); + row.putValue(BYTES_UUID_RAW, uuidBytesArr[idx]); + row.putValue(BYTES_UUID_DICT, uuidBytesArr[idx]); + row.putValue(UUID_RAW, uuidBytesArr[idx]); + row.putValue(UUID_DICT, uuidBytesArr[idx]); + row.putValue(LONG_MSB_RAW, msbArr[idx]); + row.putValue(LONG_LSB_RAW, lsbArr[idx]); + row.putValue(LONG_MSB_DICT, msbArr[idx]); + row.putValue(LONG_LSB_DICT, lsbArr[idx]); + rows.add(row); + } + + SegmentGeneratorConfig config = new SegmentGeneratorConfig(TABLE_CONFIG, SCHEMA); + config.setOutDir(INDEX_DIR.getPath()); + config.setTableName(TABLE_NAME); + config.setSegmentName(SEGMENT_NAME); + + SegmentIndexCreationDriverImpl driver = new SegmentIndexCreationDriverImpl(); + driver.init(config, new GenericRowRecordReader(rows)); + driver.build(); + + IndexLoadingConfig loadingConfig = new IndexLoadingConfig(TABLE_CONFIG, SCHEMA); + _indexSegment = ImmutableSegmentLoader.load(new File(INDEX_DIR, SEGMENT_NAME), loadingConfig); + _indexSegments = List.of(_indexSegment); + } + + @TearDown + public void tearDown() { + _indexSegment.destroy(); + FileUtils.deleteQuietly(INDEX_DIR); + } + + // ---- GROUP BY -------------------------------------------------------- + + @Benchmark + public BrokerResponseNative groupByStrUuidRaw() { + return getBrokerResponse( + "SELECT " + STR_UUID_RAW + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + STR_UUID_RAW); + } + + @Benchmark + public BrokerResponseNative groupByStrUuidDict() { + return getBrokerResponse( + "SELECT " + STR_UUID_DICT + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + STR_UUID_DICT); + } + + @Benchmark + public BrokerResponseNative groupByBytesUuidRaw() { + return getBrokerResponse( + "SELECT " + BYTES_UUID_RAW + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + BYTES_UUID_RAW); + } + + @Benchmark + public BrokerResponseNative groupByBytesUuidDict() { + return getBrokerResponse( + "SELECT " + BYTES_UUID_DICT + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + BYTES_UUID_DICT); + } + + @Benchmark + public BrokerResponseNative groupByNativeUuidRaw() { + return getBrokerResponse( + "SELECT " + UUID_RAW + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + UUID_RAW); + } + + @Benchmark + public BrokerResponseNative groupByNativeUuidDict() { + return getBrokerResponse( + "SELECT " + UUID_DICT + ", COUNT(*) FROM " + TABLE_NAME + " GROUP BY " + UUID_DICT); + } + + @Benchmark + public BrokerResponseNative groupByTwoLongRaw() { + return getBrokerResponse( + "SELECT " + LONG_MSB_RAW + ", " + LONG_LSB_RAW + ", COUNT(*) FROM " + TABLE_NAME + + " GROUP BY " + LONG_MSB_RAW + ", " + LONG_LSB_RAW); + } + + @Benchmark + public BrokerResponseNative groupByTwoLongDict() { + return getBrokerResponse( + "SELECT " + LONG_MSB_DICT + ", " + LONG_LSB_DICT + ", COUNT(*) FROM " + TABLE_NAME + + " GROUP BY " + LONG_MSB_DICT + ", " + LONG_LSB_DICT); + } + + // ---- COUNT(DISTINCT) ------------------------------------------------ + + @Benchmark + public BrokerResponseNative countDistinctStrUuidRaw() { + return getBrokerResponse("SELECT COUNT(DISTINCT " + STR_UUID_RAW + ") FROM " + TABLE_NAME); + } + + @Benchmark + public BrokerResponseNative countDistinctStrUuidDict() { + return getBrokerResponse("SELECT COUNT(DISTINCT " + STR_UUID_DICT + ") FROM " + TABLE_NAME); + } + + @Benchmark + public BrokerResponseNative countDistinctBytesUuidRaw() { + return getBrokerResponse("SELECT COUNT(DISTINCT " + BYTES_UUID_RAW + ") FROM " + TABLE_NAME); + } + + @Benchmark + public BrokerResponseNative countDistinctBytesUuidDict() { + return getBrokerResponse("SELECT COUNT(DISTINCT " + BYTES_UUID_DICT + ") FROM " + TABLE_NAME); + } + + @Benchmark + public BrokerResponseNative countDistinctNativeUuidRaw() { + return getBrokerResponse("SELECT COUNT(DISTINCT " + UUID_RAW + ") FROM " + TABLE_NAME); + } + + @Benchmark + public BrokerResponseNative countDistinctNativeUuidDict() { + return getBrokerResponse("SELECT COUNT(DISTINCT " + UUID_DICT + ") FROM " + TABLE_NAME); + } + + // ---- Equality filter (full scan) ------------------------------------ + // For BYTES, the literal is a lowercase hex string (no dashes). + // For STRING and UUID, the literal is the RFC 4122 dash-separated string. + + @Benchmark + public BrokerResponseNative filterStrUuidRaw() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + STR_UUID_RAW + " = '" + _filterUuidString + "'"); + } + + @Benchmark + public BrokerResponseNative filterStrUuidDict() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + STR_UUID_DICT + " = '" + _filterUuidString + "'"); + } + + @Benchmark + public BrokerResponseNative filterBytesUuidRaw() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + BYTES_UUID_RAW + " = '" + _filterBytesHex + "'"); + } + + @Benchmark + public BrokerResponseNative filterBytesUuidDict() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + BYTES_UUID_DICT + " = '" + _filterBytesHex + "'"); + } + + @Benchmark + public BrokerResponseNative filterNativeUuidRaw() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + UUID_RAW + " = '" + _filterUuidString + "'"); + } + + @Benchmark + public BrokerResponseNative filterNativeUuidDict() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + UUID_DICT + " = '" + _filterUuidString + "'"); + } + + @Benchmark + public BrokerResponseNative filterTwoLongRaw() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + LONG_MSB_RAW + " = " + _filterMsb + + " AND " + LONG_LSB_RAW + " = " + _filterLsb); + } + + @Benchmark + public BrokerResponseNative filterTwoLongDict() { + return getBrokerResponse( + "SELECT COUNT(*) FROM " + TABLE_NAME + " WHERE " + LONG_MSB_DICT + " = " + _filterMsb + + " AND " + LONG_LSB_DICT + " = " + _filterLsb); + } + + @Override + protected String getFilter() { + return null; + } + + @Override + protected IndexSegment getIndexSegment() { + return _indexSegment; + } + + @Override + protected List getIndexSegments() { + return _indexSegments; + } + + public static void main(String[] args) + throws Exception { + new Runner(new OptionsBuilder().include(BenchmarkUuidQueryExecution.class.getSimpleName()).build()).run(); + } +}