Add useRawBytes config to Murmur and Murmur3 partition functions (#17932)

* Add useRawBytes config to Murmur and Murmur3 partition functions

When partitioning on BYTES columns, the partition value is hex-encoded.
With useRawBytes=true in functionConfig, the hex string is decoded back
to raw bytes before hashing, ensuring partition assignment matches the
original byte values rather than treating the hex string as UTF-8 text.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Override getFunctionConfig() and restore backward-compatible constructor

- Murmur and Murmur3 partition functions now persist and return the
  functionConfig via getFunctionConfig(), ensuring partition metadata
  written to ZK retains the config so brokers can reconstruct the
  function correctly.
- Re-introduce MurmurPartitionFunction(int) constructor delegating to
  the new (int, Map) constructor to preserve SPI backward compatibility.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Defensive copy and @nullable annotations for functionConfig

- Store functionConfig as Collections.unmodifiableMap() to prevent
  external mutation from diverging with parsed fields.
- Annotate _functionConfig field and getFunctionConfig() with @nullable
  in both MurmurPartitionFunction and Murmur3PartitionFunction.
- Annotate Murmur3PartitionFunction constructor parameter with @nullable.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Add factory end-to-end test for useRawBytes with Murmur/Murmur2/Murmur3

Verify that PartitionFunctionFactory.getPartitionFunction() correctly
wires functionConfig through to the partition function for all three
aliases, and that getFunctionConfig() roundtrips the config.

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: xiangfu0

pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/Murmur3PartitionFunction.java

@@ -19,8 +19,11 @@
 package org.apache.pinot.segment.spi.partition;
 
 import com.google.common.base.Preconditions;
+import java.util.Collections;
 import java.util.Map;
+import javax.annotation.Nullable;
 import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.utils.BytesUtils;
 import org.apache.pinot.spi.utils.hash.MurmurHashFunctions;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -33,21 +36,27 @@ public class Murmur3PartitionFunction implements PartitionFunction {
   private static final String NAME = "Murmur3";
   private static final String SEED_KEY = "seed";
   private static final String VARIANT_KEY = "variant";
+  private static final String USE_RAW_BYTES_KEY = "useRawBytes";
   private final int _numPartitions;
+  @Nullable
+  private final Map<String, String> _functionConfig;
   private final int _seed;
   private final boolean _useX64;
+  private final boolean _useRawBytes;
 
   /**
    * Constructor for the class.
    * @param numPartitions Number of partitions.
    * @param functionConfig to extract configurations for the partition function.
    */
-  public Murmur3PartitionFunction(int numPartitions, Map<String, String> functionConfig) {
+  public Murmur3PartitionFunction(int numPartitions, @Nullable Map<String, String> functionConfig) {
     Preconditions.checkArgument(numPartitions > 0, "Number of partitions must be > 0");
     _numPartitions = numPartitions;
+    _functionConfig = functionConfig != null ? Collections.unmodifiableMap(functionConfig) : null;
 
     int seed = 0;
     boolean useX64 = false;
+    boolean useRawBytes = false;
     if (functionConfig != null) {
       String seedString = functionConfig.get(SEED_KEY);
       if (StringUtils.isNotEmpty(seedString)) {
@@ -62,13 +71,21 @@ public Murmur3PartitionFunction(int numPartitions, Map<String, String> functionC
               "Murmur3 variant must be either x86_32 or x64_32");
         }
       }
+      useRawBytes = Boolean.parseBoolean(functionConfig.get(USE_RAW_BYTES_KEY));
     }
     _seed = seed;
     _useX64 = useX64;
+    _useRawBytes = useRawBytes;
   }
 
   @Override
   public int getPartition(String value) {
+    if (_useRawBytes) {
+      byte[] bytes = BytesUtils.toBytes(value);
+      int hash = _useX64 ? MurmurHashFunctions.murmurHash3X64Bit32(bytes, _seed)
+          : MurmurHashFunctions.murmurHash3X86Bit32(bytes, _seed);
+      return (hash & Integer.MAX_VALUE) % _numPartitions;
+    }
     int hash = _useX64 ? MurmurHashFunctions.murmurHash3X64Bit32(value, _seed)
         : MurmurHashFunctions.murmurHash3X86Bit32(value.getBytes(UTF_8), _seed);
     return (hash & Integer.MAX_VALUE) % _numPartitions;
@@ -84,6 +101,12 @@ public int getNumPartitions() {
     return _numPartitions;
   }
 
+  @Nullable
+  @Override
+  public Map<String, String> getFunctionConfig() {
+    return _functionConfig;
+  }
+
   // Keep it for backward-compatibility, use getName() instead
   @Override
   public String toString() {

pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/MurmurPartitionFunction.java

@@ -19,6 +19,10 @@
 package org.apache.pinot.segment.spi.partition;
 
 import com.google.common.base.Preconditions;
+import java.util.Collections;
+import java.util.Map;
+import javax.annotation.Nullable;
+import org.apache.pinot.spi.utils.BytesUtils;
 import org.apache.pinot.spi.utils.hash.MurmurHashFunctions;
 
 import static java.nio.charset.StandardCharsets.UTF_8;
@@ -29,20 +33,36 @@
  */
 public class MurmurPartitionFunction implements PartitionFunction {
   private static final String NAME = "Murmur";
+  private static final String USE_RAW_BYTES_KEY = "useRawBytes";
   private final int _numPartitions;
+  @Nullable
+  private final Map<String, String> _functionConfig;
+  private final boolean _useRawBytes;
 
   /**
-   * Constructor for the class.
+   * Constructor for backward compatibility.
    * @param numPartitions Number of partitions.
    */
   public MurmurPartitionFunction(int numPartitions) {
+    this(numPartitions, null);
+  }
+
+  /**
+   * Constructor for the class.
+   * @param numPartitions Number of partitions.
+   * @param functionConfig to extract configurations for the partition function.
+   */
+  public MurmurPartitionFunction(int numPartitions, @Nullable Map<String, String> functionConfig) {
     Preconditions.checkArgument(numPartitions > 0, "Number of partitions must be > 0");
     _numPartitions = numPartitions;
+    _functionConfig = functionConfig != null ? Collections.unmodifiableMap(functionConfig) : null;
+    _useRawBytes = functionConfig != null && Boolean.parseBoolean(functionConfig.get(USE_RAW_BYTES_KEY));
   }
 
   @Override
   public int getPartition(String value) {
-    return (MurmurHashFunctions.murmurHash2(value.getBytes(UTF_8)) & Integer.MAX_VALUE) % _numPartitions;
+    byte[] bytes = _useRawBytes ? BytesUtils.toBytes(value) : value.getBytes(UTF_8);
+    return (MurmurHashFunctions.murmurHash2(bytes) & Integer.MAX_VALUE) % _numPartitions;
   }
 
   @Override
@@ -55,6 +75,12 @@ public int getNumPartitions() {
     return _numPartitions;
   }
 
+  @Nullable
+  @Override
+  public Map<String, String> getFunctionConfig() {
+    return _functionConfig;
+  }
+
   // Keep it for backward-compatibility, use getName() instead
   @Override
   public String toString() {

pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunctionFactory.java

@@ -77,7 +77,7 @@ public static PartitionFunction getPartitionFunction(String functionName, int nu
 
       case Murmur:
       case Murmur2:
-        return new MurmurPartitionFunction(numPartitions);
+        return new MurmurPartitionFunction(numPartitions, functionConfig);
 
       case Murmur3:
         return new Murmur3PartitionFunction(numPartitions, functionConfig);

pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/partition/PartitionFunctionTest.java

@@ -22,6 +22,7 @@
 import java.util.HashMap;
 import java.util.Map;
 import java.util.Random;
+import org.apache.pinot.spi.utils.BytesUtils;
 import org.apache.pinot.spi.utils.JsonUtils;
 import org.apache.pinot.spi.utils.hash.MurmurHashFunctions;
 import org.testng.annotations.Test;
@@ -439,6 +440,100 @@ public void testMurmurPartitionFunctionEquivalence() {
     testPartitionFunction(murmurPartitionFunction, expectedPartitions);
   }
 
+  @Test
+  public void testMurmurPartitionFunctionUseRawBytes() {
+    int numPartitions = 5;
+    Map<String, String> functionConfig = new HashMap<>();
+    functionConfig.put("useRawBytes", "true");
+    MurmurPartitionFunction partitionFunction = new MurmurPartitionFunction(numPartitions, functionConfig);
+
+    // When useRawBytes is true, the hex-encoded value should be decoded back to raw bytes before hashing.
+    // This means getPartition(hexString) should produce the same result as hashing the original raw bytes.
+    byte[] rawBytes = new byte[]{0x01, 0x02, 0x03, 0x04, 0x05};
+    String hexValue = BytesUtils.toHexString(rawBytes);
+
+    int expectedPartition =
+        (MurmurHashFunctions.murmurHash2(rawBytes) & Integer.MAX_VALUE) % numPartitions;
+    assertEquals(partitionFunction.getPartition(hexValue), expectedPartition);
+
+    // Without useRawBytes, the same hex string should be treated as UTF-8 text (different result).
+    MurmurPartitionFunction defaultPartitionFunction = new MurmurPartitionFunction(numPartitions, null);
+    int defaultPartition =
+        (MurmurHashFunctions.murmurHash2(hexValue.getBytes(UTF_8)) & Integer.MAX_VALUE) % numPartitions;
+    assertEquals(defaultPartitionFunction.getPartition(hexValue), defaultPartition);
+  }
+
+  @Test
+  public void testMurmur3PartitionFunctionUseRawBytes() {
+    int numPartitions = 5;
+    byte[] rawBytes = new byte[]{0x01, 0x02, 0x03, 0x04, 0x05};
+    String hexValue = BytesUtils.toHexString(rawBytes);
+
+    // Test x86_32 variant with useRawBytes
+    Map<String, String> functionConfig = new HashMap<>();
+    functionConfig.put("useRawBytes", "true");
+    Murmur3PartitionFunction partitionFunction = new Murmur3PartitionFunction(numPartitions, functionConfig);
+
+    int expectedPartition =
+        (MurmurHashFunctions.murmurHash3X86Bit32(rawBytes, 0) & Integer.MAX_VALUE) % numPartitions;
+    assertEquals(partitionFunction.getPartition(hexValue), expectedPartition);
+
+    // Test x64_32 variant with useRawBytes
+    functionConfig.put("variant", "x64_32");
+    Murmur3PartitionFunction partitionFunctionX64 = new Murmur3PartitionFunction(numPartitions, functionConfig);
+
+    int expectedPartitionX64 =
+        (MurmurHashFunctions.murmurHash3X64Bit32(rawBytes, 0) & Integer.MAX_VALUE) % numPartitions;
+    assertEquals(partitionFunctionX64.getPartition(hexValue), expectedPartitionX64);
+
+    // Test x86_32 variant with useRawBytes and non-zero seed
+    functionConfig.remove("variant");
+    functionConfig.put("seed", "9001");
+    Murmur3PartitionFunction partitionFunctionWithSeed = new Murmur3PartitionFunction(numPartitions, functionConfig);
+
+    int expectedPartitionWithSeed =
+        (MurmurHashFunctions.murmurHash3X86Bit32(rawBytes, 9001) & Integer.MAX_VALUE) % numPartitions;
+    assertEquals(partitionFunctionWithSeed.getPartition(hexValue), expectedPartitionWithSeed);
+
+    // Without useRawBytes, the same hex string should be treated as UTF-8 text.
+    Murmur3PartitionFunction defaultPartitionFunction = new Murmur3PartitionFunction(numPartitions, null);
+    int defaultPartition =
+        (MurmurHashFunctions.murmurHash3X86Bit32(hexValue.getBytes(UTF_8), 0) & Integer.MAX_VALUE) % numPartitions;
+    assertEquals(defaultPartitionFunction.getPartition(hexValue), defaultPartition);
+  }
+
+  @Test
+  public void testUseRawBytesThroughFactory() {
+    int numPartitions = 5;
+    byte[] rawBytes = new byte[]{0x01, 0x02, 0x03, 0x04, 0x05};
+    String hexValue = BytesUtils.toHexString(rawBytes);
+    int expectedPartition =
+        (MurmurHashFunctions.murmurHash2(rawBytes) & Integer.MAX_VALUE) % numPartitions;
+
+    Map<String, String> functionConfig = new HashMap<>();
+    functionConfig.put("useRawBytes", "true");
+
+    // Test Murmur alias through factory
+    PartitionFunction murmurFn =
+        PartitionFunctionFactory.getPartitionFunction("Murmur", numPartitions, functionConfig);
+    assertEquals(murmurFn.getPartition(hexValue), expectedPartition);
+    assertEquals(murmurFn.getFunctionConfig().get("useRawBytes"), "true");
+
+    // Test Murmur2 alias through factory
+    PartitionFunction murmur2Fn =
+        PartitionFunctionFactory.getPartitionFunction("Murmur2", numPartitions, functionConfig);
+    assertEquals(murmur2Fn.getPartition(hexValue), expectedPartition);
+    assertEquals(murmur2Fn.getFunctionConfig().get("useRawBytes"), "true");
+
+    // Test Murmur3 through factory
+    int expectedPartitionMurmur3 =
+        (MurmurHashFunctions.murmurHash3X86Bit32(rawBytes, 0) & Integer.MAX_VALUE) % numPartitions;
+    PartitionFunction murmur3Fn =
+        PartitionFunctionFactory.getPartitionFunction("Murmur3", numPartitions, functionConfig);
+    assertEquals(murmur3Fn.getPartition(hexValue), expectedPartitionMurmur3);
+    assertEquals(murmur3Fn.getFunctionConfig().get("useRawBytes"), "true");
+  }
+
   @Test
   public void testMurmur3PartitionFunctionEquivalence() {