[UUID 7/8] UUID partitioning

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);
     }

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);
     }

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<String, String> 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;
+  }
+}

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;
 }

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);

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,15 +643,15 @@ 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:
         result = getDistinctCountSmartULLResult(Objects.requireNonNull(dataSource.getDictionary()),
             (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));