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/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-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/query/aggregation/function/AggregationFunctionUtils.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java index 527b7817621c..bf9f716d9933 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/AggregationFunctionUtils.java @@ -69,6 +69,7 @@ import org.apache.pinot.segment.spi.AggregationFunctionType; import org.apache.pinot.segment.spi.SegmentContext; import org.apache.pinot.segment.spi.datasource.DataSource; +import org.apache.pinot.segment.spi.datasource.DataSourceMetadata; import org.apache.pinot.segment.spi.index.reader.Dictionary; import org.apache.pinot.segment.spi.index.startree.AggregationFunctionColumnPair; import org.apache.pinot.spi.data.FieldSpec; @@ -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,10 +800,12 @@ 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) { - if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { - // Treat BYTES value as serialized HyperLogLog + Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); + DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); + if (metadata.getDataType() == FieldSpec.DataType.BYTES) { + // Logical BYTES dictionary entries are serialized HyperLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLog hll = ObjectSerDeUtils.HYPER_LOG_LOG_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -820,10 +823,12 @@ private static HyperLogLog getDistinctCountHLLResult(Dictionary dictionary, } } - private static HyperLogLogPlus getDistinctCountHLLPlusResult(Dictionary dictionary, + private static HyperLogLogPlus getDistinctCountHLLPlusResult(DataSource dataSource, DistinctCountHLLPlusAggregationFunction function, String explainPlanName) { - if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { - // Treat BYTES value as serialized HyperLogLogPlus + Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); + DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); + if (metadata.getDataType() == FieldSpec.DataType.BYTES) { + // Logical BYTES dictionary entries are serialized HyperLogLogPlus objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); HyperLogLogPlus hllplus = ObjectSerDeUtils.HYPER_LOG_LOG_PLUS_SER_DE.deserialize(dictionary.getBytesValue(0)); @@ -861,10 +866,12 @@ private static Object getDistinctCountSmartHLLPlusResult(Dictionary dictionary, } } - private static UltraLogLog getDistinctCountULLResult(Dictionary dictionary, + private static UltraLogLog getDistinctCountULLResult(DataSource dataSource, DistinctCountULLAggregationFunction function, String explainPlanName) { - if (dictionary.getValueType() == FieldSpec.DataType.BYTES) { - // Treat BYTES value as serialized UltraLogLog and merge + Dictionary dictionary = Objects.requireNonNull(dataSource.getDictionary()); + DataSourceMetadata metadata = dataSource.getDataSourceMetadata(); + if (metadata.getDataType() == FieldSpec.DataType.BYTES) { + // Logical BYTES dictionary entries are serialized UltraLogLog objects. try { QueryThreadContext.checkTerminationAndSampleUsage(explainPlanName); UltraLogLog ull = ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(dictionary.getBytesValue(0)); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java index 4fe96b819dbb..e71632673b29 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountBitmapAggregationFunction.java @@ -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; @@ -71,9 +72,9 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized RoaringBitmap - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized RoaringBitmap and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); RoaringBitmap valueBitmap = aggregationResultHolder.getResult(); if (valueBitmap != null) { @@ -90,6 +91,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSV(length, aggregationResultHolder, blockValSet, storedType); } else { @@ -138,6 +141,12 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult valueBitmap.add(stringValues[i].hashCode()); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + valueBitmap.add(Arrays.hashCode(bytesValues[i])); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -198,6 +207,14 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (byte[] value : bytesValues[i]) { + valueBitmap.add(Arrays.hashCode(value)); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -209,9 +226,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized RoaringBitmap - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized RoaringBitmap and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); @@ -226,6 +243,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } else { @@ -277,6 +296,12 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(stringValues[i].hashCode()); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getValueBitmap(groupByResultHolder, groupKeyArray[i]).add(Arrays.hashCode(bytesValues[i])); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -339,6 +364,15 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + RoaringBitmap bitmap = getValueBitmap(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : bytesValues[i]) { + bitmap.add(Arrays.hashCode(value)); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -350,9 +384,9 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized RoaringBitmap - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized RoaringBitmap and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); for (int i = 0; i < length; i++) { RoaringBitmap value = RoaringBitmapUtils.deserialize(bytesValues[i]); @@ -369,6 +403,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } else { @@ -420,6 +456,12 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i].hashCode()); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], Arrays.hashCode(bytesValues[i])); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -494,6 +536,17 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + RoaringBitmap bitmap = getValueBitmap(groupByResultHolder, groupKey); + for (byte[] value : bytesValues[i]) { + bitmap.add(Arrays.hashCode(value)); + } + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_BITMAP aggregation function: " + storedType); @@ -660,6 +713,11 @@ private static RoaringBitmap convertToValueBitmap(DictIdsWrapper dictIdsWrapper) valueBitmap.add(dictionary.getStringValue(iterator.next()).hashCode()); } break; + 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); 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..8194222fb694 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountCPCSketchAggregationFunction.java @@ -135,9 +135,9 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized CPC Sketch - FieldSpec.DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES stores serialized CpcSketch objects in the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(aggregationResultHolder); @@ -153,6 +153,16 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { + aggregateSV(length, aggregationResultHolder, blockValSet, storedType); + } else { + aggregateMV(length, aggregationResultHolder, blockValSet, storedType); + } + } + + protected void aggregateSV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, + DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; if (dictionary != null) { @@ -194,11 +204,44 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde cpcSketch.update(stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + cpcSketch.update(bytesValues[i]); + } + break; + default: + throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); + } + } + + protected void aggregateMV(int length, AggregationResultHolder aggregationResultHolder, BlockValSet blockValSet, + DataType storedType) { + // For dictionary-encoded expression, store dictionary ids into the bitmap + Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; + if (dictionary != null) { + int[][] dictIds = blockValSet.getDictionaryIdsMV(); + RoaringBitmap dictIdBitmap = getDictIdBitmap(aggregationResultHolder, dictionary); + for (int i = 0; i < length; i++) { + dictIdBitmap.add(dictIds[i]); + } + return; + } + + // For non-dictionary-encoded expression, store values into the CpcSketch + switch (storedType) { + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + CpcSketch cpcSketch = getCpcSketch(aggregationResultHolder); + for (int i = 0; i < length; i++) { + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } - CpcSketchAccumulator cpcSketchAccumulator = getAccumulator(aggregationResultHolder); - cpcSketchAccumulator.apply(cpcSketch); } @Override @@ -206,9 +249,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized CPC Sketch - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == FieldSpec.DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES stores serialized CpcSketch objects in the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { CpcSketch[] sketches = deserializeSketches(bytesValues, length); @@ -225,6 +268,16 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { + aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); + } else { + aggregateMVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); + } + } + + protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; if (dictionary != null) { @@ -267,6 +320,40 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getCpcSketch(groupByResultHolder, groupKeyArray[i]).update(bytesValues[i]); + } + break; + default: + throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); + } + } + + protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType storedType) { + // For dictionary-encoded expression, store dictionary ids into the bitmap + Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; + if (dictionary != null) { + int[][] dictIds = blockValSet.getDictionaryIdsMV(); + for (int i = 0; i < length; i++) { + getDictIdBitmap(groupByResultHolder, groupKeyArray[i], dictionary).add(dictIds[i]); + } + return; + } + + // For non-dictionary-encoded expression, store values into the CpcSketch + switch (storedType) { + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } @@ -277,11 +364,9 @@ 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(); - boolean singleValue = blockValSet.isSingleValue(); - - if (singleValue && storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES stores serialized CpcSketch objects in the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { CpcSketch[] sketches = deserializeSketches(bytesValues, length); @@ -298,6 +383,16 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { + aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); + } else { + aggregateMVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); + } + } + + protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType storedType) { // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; if (dictionary != null) { @@ -350,6 +445,47 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + getCpcSketch(groupByResultHolder, groupKey).update(bytesValues[i]); + } + } + break; + default: + throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); + } + } + + protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByResultHolder groupByResultHolder, + BlockValSet blockValSet, DataType storedType) { + // For dictionary-encoded expression, store dictionary ids into the bitmap + Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; + if (dictionary != null) { + int[][] dictIds = blockValSet.getDictionaryIdsMV(); + for (int i = 0; i < length; i++) { + int[] rowDictIds = dictIds[i]; + for (int groupKey : groupKeysArray[i]) { + getDictIdBitmap(groupByResultHolder, groupKey, dictionary).add(rowDictIds); + } + } + return; + } + + // For non-dictionary-encoded expression, store values into the CpcSketch + switch (storedType) { + case BYTES: + byte[][][] bytesValues = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + CpcSketch cpcSketch = getCpcSketch(groupByResultHolder, groupKey); + for (byte[] value : bytesValues[i]) { + cpcSketch.update(value); + } + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_CPC aggregation function: " + storedType); } @@ -524,6 +660,8 @@ private CpcSketch dictionaryToCpcSketch(DictIdsWrapper dictIdsWrapper) { private void addObjectToSketch(Object rawValue, CpcSketch sketch) { if (rawValue instanceof String) { sketch.update((String) rawValue); + } else if (rawValue instanceof byte[]) { + sketch.update((byte[]) rawValue); } else if (rawValue instanceof Integer) { sketch.update((Integer) rawValue); } else if (rawValue instanceof Long) { 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..1d68cde59588 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLAggregationFunction.java @@ -81,9 +81,9 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized HyperLogLog - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized HyperLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { HyperLogLog hyperLogLog = aggregationResultHolder.getResult(); @@ -104,6 +104,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSV(length, aggregationResultHolder, blockValSet, storedType); } else { @@ -159,6 +161,12 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult hyperLogLog.offer(stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + hyperLogLog.offer(bytesValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -222,6 +230,14 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; + case BYTES: + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (byte[] value : bytesValuesArray[i]) { + hyperLogLog.offer(value); + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -232,9 +248,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized HyperLogLog - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized HyperLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -253,6 +269,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } else { @@ -307,6 +325,12 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getHyperLogLog(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -371,6 +395,15 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; + case BYTES: + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + HyperLogLog hyperLogLog = getHyperLogLog(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : bytesValuesArray[i]) { + hyperLogLog.offer(value); + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -381,9 +414,9 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized HyperLogLog - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized HyperLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -405,6 +438,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } else { @@ -459,6 +494,12 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } @@ -541,6 +582,18 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; + case BYTES: + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + byte[][] bytesValues = bytesValuesArray[i]; + for (int groupKey : groupKeysArray[i]) { + HyperLogLog hyperLogLog = getHyperLogLog(groupByResultHolder, groupKey); + for (byte[] value : bytesValues) { + hyperLogLog.offer(value); + } + } + } + break; default: throw new IllegalStateException("Illegal data type for DISTINCT_COUNT_HLL aggregation function: " + storedType); } 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..dfda4ed0bc7e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountHLLPlusAggregationFunction.java @@ -93,9 +93,9 @@ 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(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized HyperLogLogPlus and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { HyperLogLogPlus hyperLogLogPlus = aggregationResultHolder.getResult(); @@ -114,6 +114,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSV(length, aggregationResultHolder, blockValSet, storedType); } else { @@ -164,6 +166,12 @@ protected void aggregateSV(int length, AggregationResultHolder aggregationResult hyperLogLogPlus.offer(stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + hyperLogLogPlus.offer(bytesValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -226,6 +234,14 @@ protected void aggregateMV(int length, AggregationResultHolder aggregationResult } } break; + case BYTES: + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + for (byte[] value : bytesValuesArray[i]) { + hyperLogLogPlus.offer(value); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -237,9 +253,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized HyperLogLogPlus - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized HyperLogLogPlus and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -258,6 +274,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSVGroupBySV(length, groupKeyArray, groupByResultHolder, blockValSet, storedType); } else { @@ -309,6 +327,12 @@ protected void aggregateSVGroupBySV(int length, int[] groupKeyArray, GroupByResu getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]).offer(bytesValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -374,6 +398,15 @@ protected void aggregateMVGroupBySV(int length, int[] groupKeyArray, GroupByResu } } break; + case BYTES: + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(groupByResultHolder, groupKeyArray[i]); + for (byte[] value : bytesValuesArray[i]) { + hyperLogLogPlus.offer(value); + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -385,9 +418,9 @@ 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(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized HyperLogLogPlus and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -409,6 +442,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult return; } + DataType storedType = dataType.getStoredType(); + if (blockValSet.isSingleValue()) { aggregateSVGroupByMV(length, groupKeysArray, groupByResultHolder, blockValSet, storedType); } else { @@ -460,6 +495,12 @@ protected void aggregateSVGroupByMV(int length, int[][] groupKeysArray, GroupByR setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); @@ -542,6 +583,18 @@ protected void aggregateMVGroupByMV(int length, int[][] groupKeysArray, GroupByR } } break; + case BYTES: + byte[][][] bytesValuesArray = blockValSet.getBytesValuesMV(); + for (int i = 0; i < length; i++) { + byte[][] bytesValues = bytesValuesArray[i]; + for (int groupKey : groupKeysArray[i]) { + HyperLogLogPlus hyperLogLogPlus = getHyperLogLogPlus(groupByResultHolder, groupKey); + for (byte[] value : bytesValues) { + hyperLogLogPlus.offer(value); + } + } + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java index a2859fe471ec..9118d766124c 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountThetaSketchAggregationFunction.java @@ -196,7 +196,7 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde if (valueTypes[0] != DataType.BYTES) { List updateSketches = getUpdateSketches(aggregationResultHolder); if (singleValues[0]) { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[] intValues = (int[]) valueArrays[0]; if (_includeDefaultSketch) { @@ -287,13 +287,31 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } break; + case BYTES: + byte[][] bytesValues = (byte[][]) valueArrays[0]; + if (_includeDefaultSketch) { + UpdatableThetaSketch defaultSketch = updateSketches.get(0); + for (int i = 0; i < length; i++) { + defaultSketch.update(bytesValues[i]); + } + } + for (int i = 0; i < numFilters; i++) { + FilterEvaluator filterEvaluator = _filterEvaluators.get(i); + UpdatableThetaSketch updateSketch = updateSketches.get(i + 1); + for (int j = 0; j < length; j++) { + if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { + updateSketch.update(bytesValues[j]); + } + } + } + break; default: throw new IllegalStateException( "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } else { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[][] intValues = (int[][]) valueArrays[0]; if (_includeDefaultSketch) { @@ -404,13 +422,35 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde } } break; + case BYTES: + byte[][][] bytesValues = (byte[][][]) valueArrays[0]; + if (_includeDefaultSketch) { + UpdatableThetaSketch defaultSketch = updateSketches.get(0); + for (int i = 0; i < length; i++) { + for (byte[] value : bytesValues[i]) { + defaultSketch.update(value); + } + } + } + for (int i = 0; i < numFilters; i++) { + FilterEvaluator filterEvaluator = _filterEvaluators.get(i); + UpdatableThetaSketch updateSketch = updateSketches.get(i + 1); + for (int j = 0; j < length; j++) { + if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { + for (byte[] value : bytesValues[j]) { + updateSketch.update(value); + } + } + } + } + break; default: throw new IllegalStateException( "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } } else { - // Serialized sketch + // Logical BYTES stores serialized ThetaSketch objects in the single-value representation. List thetaSketchAccumulators = getUnions(aggregationResultHolder); ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { @@ -444,7 +484,7 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol // Main expression is always index 0 if (valueTypes[0] != DataType.BYTES) { if (singleValues[0]) { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[] intValues = (int[]) valueArrays[0]; for (int i = 0; i < length; i++) { @@ -520,13 +560,29 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } break; + case BYTES: + byte[][] bytesValues = (byte[][]) valueArrays[0]; + for (int i = 0; i < length; i++) { + List updateSketches = + getUpdateSketches(groupByResultHolder, groupKeyArray[i]); + byte[] value = bytesValues[i]; + if (_includeDefaultSketch) { + updateSketches.get(0).update(value); + } + for (int j = 0; j < numFilters; j++) { + if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) { + updateSketches.get(j + 1).update(value); + } + } + } + break; default: throw new IllegalStateException( "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } else { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[][] intValues = (int[][]) valueArrays[0]; for (int i = 0; i < length; i++) { @@ -632,13 +688,35 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol } } break; + case BYTES: + byte[][][] bytesValues = (byte[][][]) valueArrays[0]; + for (int i = 0; i < length; i++) { + List updateSketches = + getUpdateSketches(groupByResultHolder, groupKeyArray[i]); + byte[][] values = bytesValues[i]; + if (_includeDefaultSketch) { + UpdatableThetaSketch defaultSketch = updateSketches.get(0); + for (byte[] value : values) { + defaultSketch.update(value); + } + } + for (int j = 0; j < numFilters; j++) { + if (_filterEvaluators.get(j).evaluate(singleValues, valueTypes, valueArrays, i)) { + UpdatableThetaSketch updateSketch = updateSketches.get(j + 1); + for (byte[] value : values) { + updateSketch.update(value); + } + } + } + } + break; default: throw new IllegalStateException( "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } } else { - // Serialized sketch + // Logical BYTES stores serialized ThetaSketch objects in the single-value representation. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); for (int i = 0; i < length; i++) { List thetaSketchAccumulators = getUnions(groupByResultHolder, groupKeyArray[i]); @@ -668,7 +746,7 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult // Main expression is always index 0 if (valueTypes[0] != DataType.BYTES) { if (singleValues[0]) { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[] intValues = (int[]) valueArrays[0]; if (_includeDefaultSketch) { @@ -769,13 +847,33 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } break; + case BYTES: + byte[][] bytesValues = (byte[][]) valueArrays[0]; + if (_includeDefaultSketch) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + getUpdateSketches(groupByResultHolder, groupKey).get(0).update(bytesValues[i]); + } + } + } + for (int i = 0; i < numFilters; i++) { + FilterEvaluator filterEvaluator = _filterEvaluators.get(i); + for (int j = 0; j < length; j++) { + if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { + for (int groupKey : groupKeysArray[j]) { + getUpdateSketches(groupByResultHolder, groupKey).get(i + 1).update(bytesValues[j]); + } + } + } + } + break; default: throw new IllegalStateException( "Illegal single-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } else { - switch (valueTypes[0]) { + switch (valueTypes[0].getStoredType()) { case INT: int[][] intValues = (int[][]) valueArrays[0]; if (_includeDefaultSketch) { @@ -906,13 +1004,40 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult } } break; + case BYTES: + byte[][][] bytesValues = (byte[][][]) valueArrays[0]; + if (_includeDefaultSketch) { + for (int i = 0; i < length; i++) { + for (int groupKey : groupKeysArray[i]) { + UpdatableThetaSketch defaultSketch = getUpdateSketches(groupByResultHolder, groupKey).get(0); + for (byte[] value : bytesValues[i]) { + defaultSketch.update(value); + } + } + } + } + for (int i = 0; i < numFilters; i++) { + FilterEvaluator filterEvaluator = _filterEvaluators.get(i); + for (int j = 0; j < length; j++) { + if (filterEvaluator.evaluate(singleValues, valueTypes, valueArrays, j)) { + for (int groupKey : groupKeysArray[j]) { + UpdatableThetaSketch updateSketch = + getUpdateSketches(groupByResultHolder, groupKey).get(i + 1); + for (byte[] value : bytesValues[j]) { + updateSketch.update(value); + } + } + } + } + } + break; default: throw new IllegalStateException( "Illegal multi-value data type for DISTINCT_COUNT_THETA_SKETCH aggregation function: " + valueTypes[0]); } } } else { - // Serialized sketch + // Logical BYTES stores serialized ThetaSketch objects in the single-value representation. ThetaSketch[] sketches = deserializeSketches((byte[][]) valueArrays[0], length); if (_includeDefaultSketch) { for (int i = 0; i < length; i++) { @@ -1232,9 +1357,10 @@ private void extractValues(Map blockValSetMap, b 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; - valueTypes[i] = storedType; + valueTypes[i] = dataType; if (singleValue) { switch (storedType) { case INT: @@ -1275,6 +1401,9 @@ private void extractValues(Map blockValSetMap, b case STRING: valueArrays[i] = blockValSet.getStringValuesMV(); break; + case BYTES: + valueArrays[i] = blockValSet.getBytesValuesMV(); + break; default: throw new IllegalStateException(); } @@ -1513,7 +1642,7 @@ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] _predicateEvaluator = PredicateEvaluatorProvider.getPredicateEvaluator(_predicate, null, valueType, null); } if (singleValue) { - switch (valueType) { + switch (valueType.getStoredType()) { case INT: return _predicateEvaluator.applySV(((int[]) valueArray)[index]); case LONG: @@ -1530,7 +1659,7 @@ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] throw new IllegalStateException(); } } else { - switch (valueType) { + switch (valueType.getStoredType()) { case INT: int[] intValues = ((int[][]) valueArray)[index]; return _predicateEvaluator.applyMV(intValues, intValues.length); @@ -1546,6 +1675,9 @@ public boolean evaluate(boolean[] singleValues, DataType[] valueTypes, Object[] case STRING: String[] stringValues = ((String[][]) valueArray)[index]; return _predicateEvaluator.applyMV(stringValues, stringValues.length); + case BYTES: + byte[][] bytesValues = ((byte[][][]) valueArray)[index]; + return _predicateEvaluator.applyMV(bytesValues, bytesValues.length); default: throw new IllegalStateException(); } diff --git a/pinot-core/src/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..7ef3142ddb55 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/aggregation/function/DistinctCountULLAggregationFunction.java @@ -82,9 +82,9 @@ 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(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized UltraLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { UltraLogLog ull = aggregationResultHolder.getResult(); @@ -103,6 +103,8 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde return; } + DataType storedType = dataType.getStoredType(); + // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; if (dictionary != null) { @@ -144,6 +146,12 @@ public void aggregate(int length, AggregationResultHolder aggregationResultHolde UltraLogLogUtils.hashObject(stringValues[i]).ifPresent(ull::add); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(bytesValues[i]).ifPresent(ull::add); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_ULL aggregation function: " + storedType); @@ -155,9 +163,9 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol Map blockValSetMap) { BlockValSet blockValSet = blockValSetMap.get(_expression); - // Treat BYTES value as serialized UltraLogLogs - DataType storedType = blockValSet.getValueType().getStoredType(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized UltraLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -176,6 +184,8 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol return; } + DataType storedType = dataType.getStoredType(); + // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; if (dictionary != null) { @@ -223,6 +233,13 @@ public void aggregateGroupBySV(int length, int[] groupKeyArray, GroupByResultHol .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + UltraLogLogUtils.hashObject(bytesValues[i]) + .ifPresent(getULL(groupByResultHolder, groupKeyArray[i])::add); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_ULL aggregation function: " + storedType); @@ -234,9 +251,9 @@ 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(); - if (storedType == DataType.BYTES) { + DataType dataType = blockValSet.getValueType(); + if (dataType == DataType.BYTES) { + // Logical BYTES is a serialized UltraLogLog and always uses the single-value representation. byte[][] bytesValues = blockValSet.getBytesValuesSV(); try { for (int i = 0; i < length; i++) { @@ -246,7 +263,6 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult if (ull != null) { ull.add(value); } else { - // Create a new HyperLogLogPlus for the group groupByResultHolder.setValueForKey(groupKey, ObjectSerDeUtils.ULTRA_LOG_LOG_OBJECT_SER_DE.deserialize(bytesValues[i])); } @@ -258,6 +274,8 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult return; } + DataType storedType = dataType.getStoredType(); + // For dictionary-encoded expression, store dictionary ids into the bitmap Dictionary dictionary = blockValSet.isDictionaryEncoded() ? blockValSet.getDictionary() : null; if (dictionary != null) { @@ -300,6 +318,12 @@ public void aggregateGroupByMV(int length, int[][] groupKeysArray, GroupByResult setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], stringValues[i]); } break; + case BYTES: + byte[][] bytesValues = blockValSet.getBytesValuesSV(); + for (int i = 0; i < length; i++) { + setValueForGroupKeys(groupByResultHolder, groupKeysArray[i], bytesValues[i]); + } + break; default: throw new IllegalStateException( "Illegal data type for DISTINCT_COUNT_HLL_PLUS aggregation function: " + storedType); diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java index 386648b64f5b..8d227df72bb7 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/distinct/table/BytesDistinctTable.java @@ -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; @@ -290,9 +291,10 @@ private ResultTable toResultTableWithOrderBy() { return new ResultTable(_dataSchema, rows); } - private static void addRows(ByteArray[] values, int length, List rows) { + private void addRows(ByteArray[] values, int length, List rows) { + ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); for (int i = 0; i < length; i++) { - rows.add(new Object[]{values[i].toHexString()}); + rows.add(new Object[]{columnDataType.convertAndFormat(values[i])}); } } @@ -311,9 +313,10 @@ private ResultTable toResultTableWithoutOrderBy() { return new ResultTable(_dataSchema, rows); } - private static void addRows(HashSet values, List rows) { + private void addRows(HashSet values, List rows) { + ColumnDataType columnDataType = _dataSchema.getColumnDataType(0); for (ByteArray value : values) { - rows.add(new Object[]{value.toHexString()}); + rows.add(new Object[]{columnDataType.convertAndFormat(value)}); } } } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java index c1add8ca7082..570cc0704119 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/query/reduce/GroupByDataTableReducer.java @@ -531,6 +531,8 @@ private Object getConvertedKey(DataTable dataTable, ColumnDataType columnDataTyp return dataTable.getString(rowId, colId); case BYTES: return dataTable.getBytes(rowId, colId).getBytes(); + case UUID: + return columnDataType.convert(dataTable.getBytes(rowId, colId)); default: throw new IllegalStateException("Illegal column data type in group key: " + columnDataType); } diff --git a/pinot-core/src/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); } } 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..1b828a499986 --- /dev/null +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/UuidAggregationTest.java @@ -0,0 +1,212 @@ +/** + * 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 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 UUID aggregation coverage over dictionary-encoded and raw SV/MV columns. +@Test(suiteName = "CustomClusterIntegrationTest") +public class UuidAggregationTest extends CustomDataQueryClusterIntegrationTest { + private static final String TABLE_NAME = "UuidAggregationTest"; + private static final String UUID_DICT_SV_COLUMN = "uuidDictSv"; + private static final String UUID_DICT_MV_COLUMN = "uuidDictMv"; + private static final String UUID_RAW_SV_COLUMN = "uuidRawSv"; + private static final String UUID_RAW_MV_COLUMN = "uuidRawMv"; + + private static final String UUID_0 = "550e8400-e29b-41d4-a716-446655440000"; + private static final String UUID_0_HEX = "550e8400e29b41d4a716446655440000"; + private static final String UUID_1 = "550e8400-e29b-41d4-a716-446655440001"; + private static final String UUID_2 = "550e8400-e29b-41d4-a716-446655440002"; + private static final String UUID_3 = "550e8400-e29b-41d4-a716-446655440003"; + + private static final List UUID_SV_VALUES = List.of(UUID_0, UUID_0, UUID_1, UUID_2); + private static final List> UUID_MV_VALUES = + List.of(List.of(UUID_0, UUID_1), List.of(UUID_1, UUID_2), List.of(UUID_0), List.of(UUID_3)); + + @Override + public String getTableName() { + return TABLE_NAME; + } + + @Override + protected long getCountStarResult() { + return UUID_SV_VALUES.size(); + } + + @Override + public int getNumAvroFiles() { + return 1; + } + + @Override + public TableConfig createOfflineTableConfig() { + return new TableConfigBuilder(TableType.OFFLINE).setTableName(getTableName()) + .setNoDictionaryColumns(List.of(UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN)).build(); + } + + @Override + public Schema createSchema() { + return new Schema.SchemaBuilder().setSchemaName(getTableName()) + .addSingleValueDimension(UUID_DICT_SV_COLUMN, DataType.UUID) + .addMultiValueDimension(UUID_DICT_MV_COLUMN, DataType.UUID) + .addSingleValueDimension(UUID_RAW_SV_COLUMN, DataType.UUID) + .addMultiValueDimension(UUID_RAW_MV_COLUMN, DataType.UUID) + .build(); + } + + @Override + public List createAvroFiles() + throws Exception { + org.apache.avro.Schema uuidSchema = org.apache.avro.Schema.create(org.apache.avro.Schema.Type.STRING); + org.apache.avro.Schema avroSchema = org.apache.avro.Schema.createRecord("uuidRecord", null, null, false); + avroSchema.setFields(List.of( + new org.apache.avro.Schema.Field(UUID_DICT_SV_COLUMN, uuidSchema, null, null), + new org.apache.avro.Schema.Field(UUID_DICT_MV_COLUMN, org.apache.avro.Schema.createArray(uuidSchema), null, + null), + new org.apache.avro.Schema.Field(UUID_RAW_SV_COLUMN, uuidSchema, null, null), + new org.apache.avro.Schema.Field(UUID_RAW_MV_COLUMN, org.apache.avro.Schema.createArray(uuidSchema), null, + null))); + + try (AvroFilesAndWriters avroFilesAndWriters = createAvroFilesAndWriters(avroSchema)) { + DataFileWriter writer = avroFilesAndWriters.getWriters().get(0); + for (int i = 0; i < UUID_SV_VALUES.size(); i++) { + GenericData.Record record = new GenericData.Record(avroSchema); + record.put(UUID_DICT_SV_COLUMN, UUID_SV_VALUES.get(i)); + record.put(UUID_DICT_MV_COLUMN, UUID_MV_VALUES.get(i)); + record.put(UUID_RAW_SV_COLUMN, UUID_SV_VALUES.get(i)); + record.put(UUID_RAW_MV_COLUMN, UUID_MV_VALUES.get(i)); + writer.append(record); + } + return avroFilesAndWriters.getAvroFiles(); + } + } + + @Test + public void testGroupByHavingReturningFinalResult() + throws Exception { + setUseMultiStageQueryEngine(false); + JsonNode rows = query(String.format( + "SELECT %1$s, COUNT(*) FROM %2$s GROUP BY %1$s HAVING %1$s = '%3$s' " + + "OPTION(serverReturnFinalResult=true)", + UUID_DICT_SV_COLUMN, getTableName(), UUID_0_HEX)); + + assertEquals(rows.size(), 1, rows.toPrettyString()); + assertEquals(rows.get(0).get(0).asText(), UUID_0, rows.toPrettyString()); + assertEquals(rows.get(0).get(1).asLong(), 2L, rows.toPrettyString()); + } + + @Test(dataProvider = "useBothQueryEngines") + public void testDistinctOnUuidColumn(boolean useMultiStageQueryEngine) + throws Exception { + setUseMultiStageQueryEngine(useMultiStageQueryEngine); + JsonNode rows = query(String.format("SELECT DISTINCT %1$s FROM %2$s ORDER BY %1$s", UUID_DICT_SV_COLUMN, + getTableName())); + + assertEquals(rows.size(), 3, rows.toPrettyString()); + for (int i = 0; i < rows.size(); i++) { + assertEquals(rows.get(i).get(0).asText(), List.of(UUID_0, UUID_1, UUID_2).get(i), rows.toPrettyString()); + } + } + + @Test + public void testDistinctCountOnUuidColumns() + throws Exception { + setUseMultiStageQueryEngine(false); + for (String function : List.of("DISTINCTCOUNT", "DISTINCTCOUNTHLL", "DISTINCTCOUNTHLLPLUS", + "DISTINCTCOUNTBITMAP", "DISTINCTCOUNTTHETASKETCH", "DISTINCTCOUNTCPCSKETCH")) { + JsonNode rows = query(String.format("SELECT %1$s(%2$s), %1$s(%3$s), %1$s(%4$s) FROM %5$s", function, + UUID_DICT_SV_COLUMN, UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN, getTableName())); + assertCounts(rows.get(0), 3L, 3L, 4L); + } + + // DISTINCTCOUNTULL currently supports only single-value inputs. + JsonNode rows = query(String.format("SELECT DISTINCTCOUNTULL(%s), DISTINCTCOUNTULL(%s) FROM %s", + UUID_DICT_SV_COLUMN, UUID_RAW_SV_COLUMN, getTableName())); + assertCounts(rows.get(0), 3L, 3L); + + rows = query(String.format( + "SELECT DISTINCTCOUNTTHETASKETCH(%1$s, '', '%1$s = ''%3$s''', '$1'), " + + "DISTINCTCOUNTTHETASKETCH(%2$s, '', '%2$s = ''%3$s''', '$1') FROM %4$s", + UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN, UUID_0, getTableName())); + assertCounts(rows.get(0), 1L, 2L); + } + + @Test + public void testCpcAndThetaGroupByUuidColumns() + throws Exception { + setUseMultiStageQueryEngine(false); + JsonNode rows = queryGroupBy(UUID_RAW_SV_COLUMN); + assertEquals(rows.size(), 3, rows.toPrettyString()); + assertGroupRow(rows.get(0), UUID_0, 1L, 3L, 1L, 3L); + assertGroupRow(rows.get(1), UUID_1, 1L, 1L, 1L, 1L); + assertGroupRow(rows.get(2), UUID_2, 1L, 1L, 1L, 1L); + + // UUID multi-value group keys require a dictionary, while the aggregate input remains raw. + rows = queryGroupBy(UUID_DICT_MV_COLUMN); + assertEquals(rows.size(), 4, rows.toPrettyString()); + assertGroupRow(rows.get(0), UUID_0, 2L, 2L, 2L, 2L); + assertGroupRow(rows.get(1), UUID_1, 1L, 3L, 1L, 3L); + assertGroupRow(rows.get(2), UUID_2, 1L, 2L, 1L, 2L); + assertGroupRow(rows.get(3), UUID_3, 1L, 1L, 1L, 1L); + } + + private JsonNode queryGroupBy(String groupByColumn) + throws Exception { + return query(String.format( + "SELECT %1$s, DISTINCTCOUNTCPCSKETCH(%2$s), DISTINCTCOUNTCPCSKETCH(%3$s), " + + "DISTINCTCOUNTTHETASKETCH(%2$s), DISTINCTCOUNTTHETASKETCH(%3$s) " + + "FROM %4$s GROUP BY %1$s ORDER BY %1$s", + groupByColumn, UUID_RAW_SV_COLUMN, UUID_RAW_MV_COLUMN, getTableName())); + } + + private JsonNode query(String sql) + throws Exception { + JsonNode response = postQuery(sql); + assertTrue(response.path("exceptions").isEmpty(), sql + " -> " + response.toPrettyString()); + return response.path("resultTable").path("rows"); + } + + private static void assertGroupRow(JsonNode row, String groupKey, long... expectedCounts) { + assertEquals(row.get(0).asText(), groupKey, row.toPrettyString()); + for (int i = 0; i < expectedCounts.length; i++) { + assertEquals(row.get(i + 1).asLong(), expectedCounts[i], row.toPrettyString()); + } + } + + private static void assertCounts(JsonNode row, long... expectedCounts) { + assertEquals(row.size(), expectedCounts.length, row.toPrettyString()); + for (int i = 0; i < expectedCounts.length; i++) { + assertEquals(row.get(i).asLong(), expectedCounts[i], row.toPrettyString()); + } + } +} 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); + } +}