From de48b97ff9f60c5b559b877cc53e9db299af5ec7 Mon Sep 17 00:00:00 2001 From: praveenc7 Date: Wed, 17 Jun 2026 15:18:35 -0700 Subject: [PATCH] Precompute PartitionFunction key to reduce broker pruning CPU getPartitionFunctionKey() is called once per segment in SinglePartitionColumnSegmentPruner.prune (broker routing hot path). The default and existing overrides rebuilt the key string via concatenation on every call. Precompute it once in the constructor and return the cached field. The produced key is unchanged; this is a pure allocation/CPU optimization. Co-Authored-By: Claude Opus 4.8 (1M context) --- .../BoundedColumnValuePartitionFunction.java | 6 +++-- .../partition/ByteArrayPartitionFunction.java | 7 +++++ .../partition/HashCodePartitionFunction.java | 7 +++++ .../partition/ModuloPartitionFunction.java | 7 +++++ .../partition/Murmur3PartitionFunction.java | 4 ++- .../partition/MurmurPartitionFunction.java | 7 +++++ .../spi/partition/PartitionFunctionTest.java | 26 +++++++++++++++++++ 7 files changed, 61 insertions(+), 3 deletions(-) diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/BoundedColumnValuePartitionFunction.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/BoundedColumnValuePartitionFunction.java index ce52909fe3..ed550974d4 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/BoundedColumnValuePartitionFunction.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/BoundedColumnValuePartitionFunction.java @@ -47,6 +47,7 @@ public class BoundedColumnValuePartitionFunction implements PartitionFunction { private final int _numPartitions; private final Map _functionConfig; private final String[] _values; + private final String _partitionFunctionKey; public BoundedColumnValuePartitionFunction(int numPartitions, Map functionConfig) { Preconditions.checkArgument(functionConfig != null && functionConfig.size() > 0, @@ -59,6 +60,8 @@ public BoundedColumnValuePartitionFunction(int numPartitions, Map 0, "Number of partitions must be > 0, specified", numPartitions); _numPartitions = numPartitions; + _partitionFunctionKey = NAME + "_" + _numPartitions; } @Override @@ -56,6 +58,11 @@ public int getNumPartitions() { return _numPartitions; } + @Override + public String getPartitionFunctionKey() { + return _partitionFunctionKey; + } + // Keep it for backward-compatibility, use getName() instead @Override public String toString() { diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/HashCodePartitionFunction.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/HashCodePartitionFunction.java index 182760cf44..528fb7fa1f 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/HashCodePartitionFunction.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/HashCodePartitionFunction.java @@ -30,10 +30,12 @@ public class HashCodePartitionFunction implements PartitionFunction { private static final String NAME = "HashCode"; private final int _numPartitions; + private final String _partitionFunctionKey; public HashCodePartitionFunction(int numPartitions) { Preconditions.checkArgument(numPartitions > 0, "Number of partitions must be > 0, specified", numPartitions); _numPartitions = numPartitions; + _partitionFunctionKey = NAME + "_" + _numPartitions; } @Override @@ -51,6 +53,11 @@ public int getNumPartitions() { return _numPartitions; } + @Override + public String getPartitionFunctionKey() { + return _partitionFunctionKey; + } + // Keep it for backward-compatibility, use getName() instead @Override public String toString() { diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/ModuloPartitionFunction.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/ModuloPartitionFunction.java index a4b8eb49ab..6afbcb7585 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/ModuloPartitionFunction.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/ModuloPartitionFunction.java @@ -31,6 +31,7 @@ public class ModuloPartitionFunction implements PartitionFunction { private static final String NAME = "Modulo"; private final int _numPartitions; + private final String _partitionFunctionKey; /** * Constructor for the class. @@ -39,6 +40,7 @@ public class ModuloPartitionFunction implements PartitionFunction { public ModuloPartitionFunction(int numPartitions) { Preconditions.checkArgument(numPartitions > 0, "Number of partitions must be > 0, specified", numPartitions); _numPartitions = numPartitions; + _partitionFunctionKey = NAME + "_" + _numPartitions; } /** @@ -63,6 +65,11 @@ public int getNumPartitions() { return _numPartitions; } + @Override + public String getPartitionFunctionKey() { + return _partitionFunctionKey; + } + // Keep it for backward-compatibility, use getName() instead @Override public String toString() { diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/Murmur3PartitionFunction.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/Murmur3PartitionFunction.java index 18b03b5de4..8a82373e7d 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/Murmur3PartitionFunction.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/Murmur3PartitionFunction.java @@ -36,6 +36,7 @@ public class Murmur3PartitionFunction implements PartitionFunction { private final int _numPartitions; private final int _seed; private final boolean _useX64; + private final String _partitionFunctionKey; /** * Constructor for the class. @@ -65,6 +66,7 @@ public Murmur3PartitionFunction(int numPartitions, Map functionC } _seed = seed; _useX64 = useX64; + _partitionFunctionKey = NAME + "_" + _numPartitions + "_" + _seed + "_" + (_useX64 ? "x64" : "x86"); } @Override @@ -92,6 +94,6 @@ public String toString() { @Override public String getPartitionFunctionKey() { - return NAME + "_" + _numPartitions + "_" + _seed + "_" + (_useX64 ? "x64" : "x86"); + return _partitionFunctionKey; } } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/MurmurPartitionFunction.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/MurmurPartitionFunction.java index 8251450c44..3970db9b9c 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/MurmurPartitionFunction.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/MurmurPartitionFunction.java @@ -30,6 +30,7 @@ public class MurmurPartitionFunction implements PartitionFunction { private static final String NAME = "Murmur"; private final int _numPartitions; + private final String _partitionFunctionKey; /** * Constructor for the class. @@ -38,6 +39,7 @@ public class MurmurPartitionFunction implements PartitionFunction { public MurmurPartitionFunction(int numPartitions) { Preconditions.checkArgument(numPartitions > 0, "Number of partitions must be > 0"); _numPartitions = numPartitions; + _partitionFunctionKey = NAME + "_" + _numPartitions; } @Override @@ -55,6 +57,11 @@ public int getNumPartitions() { return _numPartitions; } + @Override + public String getPartitionFunctionKey() { + return _partitionFunctionKey; + } + // Keep it for backward-compatibility, use getName() instead @Override public String toString() { diff --git a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/partition/PartitionFunctionTest.java b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/partition/PartitionFunctionTest.java index c2f5b7d75d..69bb422856 100644 --- a/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/partition/PartitionFunctionTest.java +++ b/pinot-segment-spi/src/test/java/org/apache/pinot/segment/spi/partition/PartitionFunctionTest.java @@ -28,6 +28,7 @@ import static java.nio.charset.StandardCharsets.UTF_8; import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertSame; import static org.testng.Assert.assertTrue; @@ -417,6 +418,31 @@ public void testPartitionFunctionKey() { assertTrue(!murmur3Function.getPartitionFunctionKey().equals(murmur3DifferentSeed.getPartitionFunctionKey())); } + @Test + public void testPartitionFunctionKeyIsPrecomputed() { + // The key is computed once and called per-segment on the broker's hot pruning path. Verify each implementation + // returns the same cached instance on repeated calls rather than rebuilding the string every time. + Map murmur3Config = new HashMap<>(); + murmur3Config.put("seed", "42"); + murmur3Config.put("variant", "x64_32"); + Map boundedConfig = new HashMap<>(); + boundedConfig.put("columnValues", "a|b|c"); + boundedConfig.put("columnValuesDelimiter", "|"); + + PartitionFunction[] functions = { + PartitionFunctionFactory.getPartitionFunction("Modulo", 10, null), + PartitionFunctionFactory.getPartitionFunction("HashCode", 20, null), + PartitionFunctionFactory.getPartitionFunction("Murmur", 15, null), + PartitionFunctionFactory.getPartitionFunction("Murmur3", 25, murmur3Config), + PartitionFunctionFactory.getPartitionFunction("ByteArray", 12, null), + PartitionFunctionFactory.getPartitionFunction("BoundedColumnValue", 4, boundedConfig) + }; + for (PartitionFunction function : functions) { + assertSame(function.getPartitionFunctionKey(), function.getPartitionFunctionKey(), + function.getName() + " should return the same precomputed key instance on every call"); + } + } + private void testBasicProperties(PartitionFunction partitionFunction, String functionName, int numPartitions) { testBasicProperties(partitionFunction, functionName, numPartitions, null); }