This is an automated email from the ASF dual-hosted git repository.
xiangfu0 pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/pinot.git
The following commit(s) were added to refs/heads/master by this push:
new 9c3f0e6d208 Prepare partition pruning predicates once per query
(#19539)
9c3f0e6d208 is described below
commit 9c3f0e6d208f48627414fd3b7cb4e6b32e9db3f0
Author: Xiang Fu <[email protected]>
AuthorDate: Tue Oct 6 21:31:58 2026 -0700
Prepare partition pruning predicates once per query (#19539)
* Reuse partition predicate hashes across broker segments
* Simplify partition function config matching
* Remove partition pruning benchmark
* Fix unused test import
* Address partition pruning review feedback
* Correct partition pruning review fixes and canonicalize cache keys
* Simplify partition pruning cache state
* Limit partition pruning optimization to query-local ID caching
* Limit partition predicate preparation to large segment sets
* Delegate partition ID reuse compatibility to partition functions
* Clarify partition reuse contract and add query-level cache fallback
* Use primitive partition ID sets and configurable pruning cutoff
Keep scalar IDs for equality and single-ID IN prefixes, and retain lazy
evaluation for remaining literals. Use a 32-candidate default based on bounded
JMH comparisons.
Configure pinot.broker.partition.pruning.cache.min.segments in cluster
config, or append .events_OFFLINE for a live table override. For example, use a
cluster default of 256 and an events_OFFLINE override of 32 to canary the lower
cutoff. Removing the table key restores the cluster default;
enablePartitionPruningCache=false still forces the original evaluator.
* Avoid overlap with concurrent routing-manager tests
* Correct partition pruning cache cutoff documentation
* Support config-aware partition ID reuse
* Cache partition function identities
* Avoid per-query partition threshold lookup
* Use partition function equality instead of interned identities
* Avoid copying partition function configuration
* Address partition pruning review feedback
Use numeric preparation thresholds at cluster, table, and query levels.
Reuse a small per-query list of function-specific prepared predicates.
* Use function equality for query-local partition pruning
Read the live preparation threshold once per query so cluster changes do
not rebuild routing entries.
* Avoid hashing functions in common partition-pruning queries
Use a small query-local list and promote to a map after eight distinct
functions. Keep lazy predicate evaluation and cover the promotion case.
* Avoid preparing non-reusable partition functions
* Reuse OR partition IDs after a miss
Once every direct OR value has been visited on a nonmatching segment, reuse
their union for later segment checks. Keep first-hit short circuiting and
multi-partition AND semantics unchanged.
How to reproduce: Compare 256, 512, and 1024 segment pruning for a nested
OR, EQUALS, and first-hit IN filter. Whole-predicate eager preparation slows
the first-hit IN control while the on-miss OR union speeds up the nested case.
* Keep OR preparation off simple predicate hot paths
Reuse existing prepared-predicate fields for OR state and move its union
logic out of matches(). This keeps the EQUALS and first-hit IN paths lean while
retaining the nested-filter gain.
How to reproduce: Run the full-prune JMH fixture with 256, 512, and 1024
independently constructed equal Modulo functions, comparing the pre-change lazy
class to this revision for EQUALS, nested AND/OR, and first-hit IN(128).
* Prepare partition pruning ID sets eagerly
Compute the whole predicate once per equivalent partition function and
query, intersecting AND candidates and unioning OR candidates. Retain the
singleton-ID fast path and use the original evaluator when preparation throws,
preserving its short-circuit error behavior.
How to reproduce: For a single-valued Modulo partition column, a segment
listing partitions 1 and 2 was retained for id=1 AND id=2 even though no row
could match. The prepared intersection now prunes it. Tests also cover nested
set operations and invalid literals after an empty intersection.
* Make partition preparation opt-in and short-circuit unrestricted OR
---
.../broker/broker/helix/BaseBrokerStarter.java | 1 +
.../helix/MultiClusterHelixBrokerStarter.java | 1 +
.../routing/manager/BaseBrokerRoutingManager.java | 29 +-
.../segmentpruner/SegmentPrunerFactory.java | 27 +-
.../SinglePartitionColumnSegmentPruner.java | 186 ++++++
.../routing/manager/BrokerRoutingManagerTest.java | 33 ++
.../SinglePartitionColumnSegmentPrunerTest.java | 632 +++++++++++++++++++++
.../function/ByteArrayPartitionFunction.java | 24 +
.../partition/function/FnvPartitionFunction.java | 28 +
.../function/HashCodePartitionFunction.java | 24 +
.../function/ModuloPartitionFunction.java | 24 +
.../function/Murmur3PartitionFunction.java | 29 +
.../function/MurmurPartitionFunction.java | 26 +
.../common/utils/config/QueryOptionsUtils.java | 8 +
.../partition/function/PartitionFunctionTest.java | 123 +++-
.../common/utils/config/QueryOptionsUtilsTest.java | 10 +
.../segment/spi/partition/PartitionFunction.java | 11 +
.../pinot/spi/config/table/RoutingConfig.java | 19 +-
.../apache/pinot/spi/utils/CommonConstants.java | 8 +
19 files changed, 1231 insertions(+), 12 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
index 18cac8d2825..79ad054c15b 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/BaseBrokerStarter.java
@@ -440,6 +440,7 @@ public abstract class BaseBrokerStarter implements
ServiceStartable {
_serverRoutingStatsManager = new ServerRoutingStatsManager(_brokerConf,
_brokerMetrics);
_serverRoutingStatsManager.init();
initRoutingManager();
+
_clusterConfigChangeHandler.registerClusterConfigChangeListener(_routingManager);
final PinotConfiguration factoryConf =
_brokerConf.subset(Broker.ACCESS_CONTROL_CONFIG_PREFIX);
// Adding cluster name to the config so that it can be used by the
AccessControlFactory
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
index 661c640b520..a1e05842c37 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/broker/helix/MultiClusterHelixBrokerStarter.java
@@ -230,6 +230,7 @@ public class MultiClusterHelixBrokerStarter extends
BaseBrokerStarter {
RemoteClusterBrokerRoutingManager routingManager =
new RemoteClusterBrokerRoutingManager(clusterName, _brokerMetrics,
_serverRoutingStatsManager, _brokerConf);
routingManager.init(entry.getValue());
+
_clusterConfigChangeHandler.registerClusterConfigChangeListener(routingManager);
_remoteRoutingManagers.put(clusterName, routingManager);
} catch (Exception e) {
LOGGER.error("[multi-cluster] Failed to initialize routing manager for
cluster '{}'", clusterName, e);
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
index 25e5b784b7d..829e20c8e82 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/manager/BaseBrokerRoutingManager.java
@@ -90,6 +90,7 @@ import
org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategy;
import org.apache.pinot.core.routing.timeboundary.TimeBoundaryStrategyService;
import org.apache.pinot.core.transport.ServerInstance;
import
org.apache.pinot.core.transport.server.routing.stats.ServerRoutingStatsManager;
+import org.apache.pinot.spi.config.provider.PinotClusterConfigChangeListener;
import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
import org.apache.pinot.spi.config.table.QueryConfig;
import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
@@ -127,7 +128,8 @@ import org.slf4j.LoggerFactory;
///
/// TODO: Expose RoutingEntry class to get a consistent view in the broker
request handler and save the redundant map
/// lookups.
-public abstract class BaseBrokerRoutingManager implements RoutingManager,
ClusterChangeHandler {
+public abstract class BaseBrokerRoutingManager
+ implements RoutingManager, ClusterChangeHandler,
PinotClusterConfigChangeListener {
private static final Logger LOGGER =
LoggerFactory.getLogger(BaseBrokerRoutingManager.class);
protected final BrokerMetrics _brokerMetrics;
@@ -170,6 +172,9 @@ public abstract class BaseBrokerRoutingManager implements
RoutingManager, Cluste
private String _instanceConfigsPath;
protected ZkHelixPropertyStore<ZNRecord> _propertyStore;
+ private volatile int _partitionPruningPreparationThreshold =
+ CommonConstants.Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD;
+
/// Snapshot of `_enabledServerInstanceMap` restricted to
enabled-minus-excluded servers. Replaced atomically under
/// `_globalLock.writeLock()` whenever routing membership changes. Serves as
the single source of truth for routable
/// servers: its `keySet()` is passed to `InstanceSelector` for per-table
selection, and callers that pick workers
@@ -209,6 +214,25 @@ public abstract class BaseBrokerRoutingManager implements
RoutingManager, Cluste
_propertyStore = helixManager.getHelixPropertyStore();
}
+ @Override
+ public void onChange(Set<String> changedConfigs, Map<String, String>
clusterConfigs) {
+ String key =
CommonConstants.Broker.CONFIG_OF_PARTITION_PRUNING_PREPARATION_THRESHOLD;
+ int updatedThreshold =
CommonConstants.Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD;
+ String value = clusterConfigs.get(key);
+ if (value != null) {
+ try {
+ updatedThreshold = Integer.parseInt(value);
+ } catch (NumberFormatException e) {
+ LOGGER.warn("Ignoring invalid partition pruning threshold: {}={}",
key, value);
+ }
+ }
+ _partitionPruningPreparationThreshold = updatedThreshold;
+ }
+
+ int getPartitionPruningPreparationThreshold() {
+ return _partitionPruningPreparationThreshold;
+ }
+
/// Sets a callback to be invoked when a server is re-enabled after being
excluded.
/// This is useful for resetting gRPC channel state to avoid exponential
backoff delays.
public void setServerReenableCallback(Consumer<ServerInstance> callback) {
@@ -812,7 +836,8 @@ public abstract class BaseBrokerRoutingManager implements
RoutingManager, Cluste
segmentSelector.init(idealState, externalView,
preSelectedOnlineSegments);
// Register segment pruners and initialize segment zk metadata fetcher.
- List<SegmentPruner> segmentPruners =
SegmentPrunerFactory.getSegmentPruners(tableConfig, _propertyStore);
+ List<SegmentPruner> segmentPruners =
SegmentPrunerFactory.getSegmentPruners(tableConfig, _propertyStore,
+ this::getPartitionPruningPreparationThreshold);
AdaptiveServerSelector adaptiveServerSelector =
AdaptiveServerSelectorFactory.getAdaptiveServerSelector(_serverRoutingStatsManager,
_pinotConfig);
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerFactory.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerFactory.java
index 6fc8a150c1a..1a084733226 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerFactory.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerFactory.java
@@ -22,6 +22,7 @@ import java.util.ArrayList;
import java.util.List;
import java.util.Map;
import java.util.Set;
+import java.util.function.IntSupplier;
import javax.annotation.Nullable;
import org.apache.commons.collections4.MapUtils;
import org.apache.helix.store.zk.ZkHelixPropertyStore;
@@ -36,6 +37,7 @@ import org.apache.pinot.spi.config.table.TableConfig;
import org.apache.pinot.spi.config.table.TableType;
import org.apache.pinot.spi.data.DateTimeFieldSpec;
import org.apache.pinot.spi.data.Schema;
+import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -51,6 +53,16 @@ public class SegmentPrunerFactory {
public static List<SegmentPruner> getSegmentPruners(TableConfig tableConfig,
ZkHelixPropertyStore<ZNRecord> propertyStore) {
+ return getSegmentPruners(tableConfig, propertyStore,
Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD);
+ }
+
+ public static List<SegmentPruner> getSegmentPruners(TableConfig tableConfig,
+ ZkHelixPropertyStore<ZNRecord> propertyStore, int
defaultPartitionPruningPreparationThreshold) {
+ return getSegmentPruners(tableConfig, propertyStore, () ->
defaultPartitionPruningPreparationThreshold);
+ }
+
+ public static List<SegmentPruner> getSegmentPruners(TableConfig tableConfig,
+ ZkHelixPropertyStore<ZNRecord> propertyStore, IntSupplier
defaultPreparationThreshold) {
List<SegmentPruner> segmentPruners = new ArrayList<>();
boolean needsEmptySegment =
TableConfigUtils.needsEmptySegmentPruner(tableConfig);
if (needsEmptySegment) {
@@ -65,7 +77,8 @@ public class SegmentPrunerFactory {
List<SegmentPruner> configuredSegmentPruners = new
ArrayList<>(segmentPrunerTypes.size());
for (String segmentPrunerType : segmentPrunerTypes) {
if
(RoutingConfig.PARTITION_SEGMENT_PRUNER_TYPE.equalsIgnoreCase(segmentPrunerType))
{
- SegmentPruner partitionSegmentPruner =
getPartitionSegmentPruner(tableConfig);
+ SegmentPruner partitionSegmentPruner =
+ getPartitionSegmentPruner(tableConfig,
defaultPreparationThreshold);
if (partitionSegmentPruner != null) {
configuredSegmentPruners.add(partitionSegmentPruner);
}
@@ -88,7 +101,8 @@ public class SegmentPrunerFactory {
if ((tableType == TableType.OFFLINE &&
LEGACY_PARTITION_AWARE_OFFLINE_ROUTING.equalsIgnoreCase(
routingTableBuilderName)) || (tableType == TableType.REALTIME
&&
LEGACY_PARTITION_AWARE_REALTIME_ROUTING.equalsIgnoreCase(routingTableBuilderName)))
{
- SegmentPruner partitionSegmentPruner =
getPartitionSegmentPruner(tableConfig);
+ SegmentPruner partitionSegmentPruner =
+ getPartitionSegmentPruner(tableConfig,
defaultPreparationThreshold);
if (partitionSegmentPruner != null) {
segmentPruners.add(partitionSegmentPruner);
}
@@ -99,7 +113,8 @@ public class SegmentPrunerFactory {
}
@Nullable
- private static SegmentPruner getPartitionSegmentPruner(TableConfig
tableConfig) {
+ private static SegmentPruner getPartitionSegmentPruner(TableConfig
tableConfig,
+ IntSupplier defaultPreparationThreshold) {
String tableNameWithType = tableConfig.getTableName();
SegmentPartitionConfig segmentPartitionConfig =
tableConfig.getIndexingConfig().getSegmentPartitionConfig();
if (segmentPartitionConfig == null) {
@@ -114,8 +129,12 @@ public class SegmentPrunerFactory {
Set<String> partitionColumns = columnPartitionMap.keySet();
LOGGER.info("Using PartitionSegmentPruner on partition columns: {} for
table: {}", partitionColumns,
tableNameWithType);
+ RoutingConfig routingConfig = tableConfig.getRoutingConfig();
+ Integer tableMinSegments = routingConfig != null ?
routingConfig.getPartitionPruningPreparationThreshold() : null;
+ IntSupplier preparationThreshold =
+ tableMinSegments != null ? () -> tableMinSegments :
defaultPreparationThreshold;
return partitionColumns.size() == 1 ? new
SinglePartitionColumnSegmentPruner(tableNameWithType,
- partitionColumns.iterator().next())
+ partitionColumns.iterator().next(), preparationThreshold)
: new MultiPartitionColumnsSegmentPruner(tableNameWithType,
partitionColumns);
}
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java
index 3e672b88686..3bb79cf4a10 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java
@@ -18,11 +18,18 @@
*/
package org.apache.pinot.broker.routing.segmentpruner;
+import it.unimi.dsi.fastutil.ints.IntIterator;
+import it.unimi.dsi.fastutil.ints.IntOpenHashSet;
+import it.unimi.dsi.fastutil.ints.IntSet;
+import it.unimi.dsi.fastutil.ints.IntSets;
+import java.util.ArrayList;
+import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;
import java.util.concurrent.ConcurrentHashMap;
+import java.util.function.IntSupplier;
import javax.annotation.Nullable;
import org.apache.helix.model.ExternalView;
import org.apache.helix.model.IdealState;
@@ -34,19 +41,35 @@ import org.apache.pinot.common.request.Expression;
import org.apache.pinot.common.request.Function;
import org.apache.pinot.common.request.Identifier;
import org.apache.pinot.common.request.context.RequestContextUtils;
+import org.apache.pinot.common.utils.config.QueryOptionsUtils;
+import org.apache.pinot.segment.spi.partition.PartitionFunction;
+import org.apache.pinot.spi.utils.CommonConstants.Broker;
import org.apache.pinot.sql.FilterKind;
/// The `SinglePartitionColumnSegmentPruner` prunes segments based on their
partition metadata stored in ZK. The
/// pruner supports queries with filter (or nested filter) of EQUALITY and IN
predicates.
public class SinglePartitionColumnSegmentPruner implements SegmentPruner {
+ private static final int MAX_LINEAR_FUNCTIONS = 8;
private final String _tableNameWithType;
private final String _partitionColumn;
+ private final IntSupplier _preparationThreshold;
private final Map<String, SegmentPartitionInfo> _partitionInfoMap = new
ConcurrentHashMap<>();
public SinglePartitionColumnSegmentPruner(String tableNameWithType, String
partitionColumn) {
+ this(tableNameWithType, partitionColumn,
Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD);
+ }
+
+ public SinglePartitionColumnSegmentPruner(String tableNameWithType, String
partitionColumn,
+ int minSegmentsForPreparation) {
+ this(tableNameWithType, partitionColumn, () -> minSegmentsForPreparation);
+ }
+
+ public SinglePartitionColumnSegmentPruner(String tableNameWithType, String
partitionColumn,
+ IntSupplier preparationThreshold) {
_tableNameWithType = tableNameWithType;
_partitionColumn = partitionColumn;
+ _preparationThreshold = preparationThreshold;
}
@Override
@@ -94,6 +117,12 @@ public class SinglePartitionColumnSegmentPruner implements
SegmentPruner {
if (filterExpression == null) {
return segments;
}
+ Integer queryMinSegments =
+
QueryOptionsUtils.getPartitionPruningPreparationThreshold(brokerRequest.getPinotQuery().getQueryOptions());
+ int minSegmentsForPreparation = queryMinSegments != null ?
queryMinSegments : _preparationThreshold.getAsInt();
+ if (minSegmentsForPreparation >= 0 && segments.size() >=
minSegmentsForPreparation) {
+ return pruneWithPreparedPredicate(filterExpression, segments);
+ }
Set<String> selectedSegments = new HashSet<>();
for (String segment : segments) {
SegmentPartitionInfo partitionInfo = _partitionInfoMap.get(segment);
@@ -105,6 +134,62 @@ public class SinglePartitionColumnSegmentPruner implements
SegmentPruner {
return selectedSegments;
}
+ private Set<String> pruneWithPreparedPredicate(Expression filterExpression,
Set<String> segments) {
+ Set<String> selectedSegments = new HashSet<>();
+ List<PreparedPartitionIds> predicates = new ArrayList<>(2);
+ Map<PartitionFunction, PreparedPartitionIds> predicateMap = null;
+ for (String segment : segments) {
+ SegmentPartitionInfo partitionInfo = _partitionInfoMap.get(segment);
+ if (partitionInfo == null || partitionInfo ==
SegmentPartitionUtils.INVALID_PARTITION_INFO) {
+ selectedSegments.add(segment);
+ continue;
+ }
+ PartitionFunction function = partitionInfo.getPartitionFunction();
+ if (!function.supportsPartitionIdPreparation()) {
+ if (isPartitionMatch(filterExpression, partitionInfo)) {
+ selectedSegments.add(segment);
+ }
+ continue;
+ }
+ PreparedPartitionIds predicate = null;
+ if (predicateMap == null) {
+ for (PreparedPartitionIds candidate : predicates) {
+ if (function.equals(candidate._partitionFunction)) {
+ predicate = candidate;
+ break;
+ }
+ }
+ } else {
+ predicate = predicateMap.get(function);
+ }
+ if (predicate == null) {
+ try {
+ predicate = new PreparedPartitionIds(function,
preparePartitionIds(filterExpression, function), false);
+ } catch (RuntimeException e) {
+ // Preparation can reach a value or branch that per-segment
short-circuiting would skip. Defer errors
+ // to the original evaluator, and remember the fallback so this
function is not prepared again.
+ predicate = new PreparedPartitionIds(function, null, true);
+ }
+ if (predicateMap == null && predicates.size() < MAX_LINEAR_FUNCTIONS) {
+ predicates.add(predicate);
+ } else {
+ if (predicateMap == null) {
+ predicateMap = new HashMap<>();
+ for (PreparedPartitionIds candidate : predicates) {
+ predicateMap.put(candidate._partitionFunction, candidate);
+ }
+ }
+ predicateMap.put(function, predicate);
+ }
+ }
+ if (predicate._preparationFailed ? isPartitionMatch(filterExpression,
partitionInfo)
+ : predicate.matches(partitionInfo.getPartitions())) {
+ selectedSegments.add(segment);
+ }
+ }
+ return selectedSegments;
+ }
+
private boolean isPartitionMatch(Expression filterExpression,
SegmentPartitionInfo partitionInfo) {
Function function = filterExpression.getFunctionCall();
FilterKind filterKind = FilterKind.valueOf(function.getOperator());
@@ -152,4 +237,105 @@ public class SinglePartitionColumnSegmentPruner
implements SegmentPruner {
return true;
}
}
+
+ /// Computes conservative candidate IDs for the whole predicate. Null
represents every partition.
+ /// Partition columns are single-valued, so a row satisfying AND must belong
to every child's candidate set.
+ @Nullable
+ private IntSet preparePartitionIds(Expression expression, PartitionFunction
partitionFunction) {
+ Function function = expression.getFunctionCall();
+ FilterKind kind = FilterKind.valueOf(function.getOperator());
+ List<Expression> operands = function.getOperands();
+ switch (kind) {
+ case AND: {
+ IntSet ids = null;
+ for (Expression child : operands) {
+ IntSet childIds = preparePartitionIds(child, partitionFunction);
+ if (childIds != null) {
+ if (ids == null) {
+ ids = childIds;
+ } else {
+ // Singleton leaves are immutable; composite and IN sets can be
intersected in place.
+ if (!(ids instanceof IntOpenHashSet)) {
+ ids = new IntOpenHashSet(ids);
+ }
+ ids.retainAll(childIds);
+ }
+ }
+ }
+ return ids;
+ }
+ case OR: {
+ IntSet ids = new IntOpenHashSet();
+ for (Expression child : operands) {
+ IntSet childIds = preparePartitionIds(child, partitionFunction);
+ if (childIds == null) {
+ return null;
+ }
+ ids.addAll(childIds);
+ }
+ return ids;
+ }
+ case EQUALS:
+ case IN: {
+ Identifier identifier = operands.get(0).getIdentifier();
+ if (identifier == null ||
!identifier.getName().equals(_partitionColumn)) {
+ return null;
+ }
+ int numValues = kind == FilterKind.EQUALS ? 1 : operands.size() - 1;
+ if (numValues == 1) {
+ return
IntSets.singleton(partitionFunction.getPartition(RequestContextUtils.getStringValue(operands.get(1))));
+ }
+ IntSet ids = new IntOpenHashSet(Math.min(numValues,
partitionFunction.getNumPartitions()));
+ for (int i = 1; i <= numValues; i++) {
+
ids.add(partitionFunction.getPartition(RequestContextUtils.getStringValue(operands.get(i))));
+ }
+ return ids;
+ }
+ default:
+ return null;
+ }
+ }
+
+ /// Prepared IDs belong to one prune call. Their sets are read-only after
construction and never shared by queries.
+ private static final class PreparedPartitionIds {
+ private final PartitionFunction _partitionFunction;
+ @Nullable
+ private final IntSet _partitionIds;
+ @Nullable
+ private final Integer _singlePartitionId;
+ private final boolean _preparationFailed;
+
+ private PreparedPartitionIds(PartitionFunction partitionFunction,
@Nullable IntSet partitionIds,
+ boolean preparationFailed) {
+ _partitionFunction = partitionFunction;
+ _partitionIds = partitionIds;
+ _singlePartitionId = partitionIds != null && partitionIds.size() == 1 ?
partitionIds.iterator().nextInt() : null;
+ _preparationFailed = preparationFailed;
+ }
+
+ private boolean matches(Set<Integer> partitions) {
+ if (_partitionIds == null) {
+ return true;
+ }
+ if (_singlePartitionId != null) {
+ return partitions.contains(_singlePartitionId);
+ }
+ // Probe the smaller set; a segment can itself contain many partitions.
+ if (_partitionIds.size() <= partitions.size()) {
+ IntIterator iterator = _partitionIds.iterator();
+ while (iterator.hasNext()) {
+ if (partitions.contains(iterator.nextInt())) {
+ return true;
+ }
+ }
+ } else {
+ for (int partition : partitions) {
+ if (_partitionIds.contains(partition)) {
+ return true;
+ }
+ }
+ }
+ return false;
+ }
+ }
}
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java
index 82668fd1ce1..57592fd309b 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/manager/BrokerRoutingManagerTest.java
@@ -31,6 +31,7 @@ import org.apache.helix.HelixConstants.ChangeType;
import org.apache.helix.HelixDataAccessor;
import org.apache.helix.HelixManager;
import org.apache.helix.PropertyKey;
+import org.apache.helix.model.ClusterConfig;
import org.apache.helix.model.ExternalView;
import org.apache.helix.model.IdealState;
import org.apache.helix.model.InstanceConfig;
@@ -46,6 +47,7 @@ import
org.apache.pinot.broker.routing.segmentpruner.SegmentPruner;
import org.apache.pinot.broker.routing.segmentselector.SegmentSelector;
import org.apache.pinot.broker.routing.tablesampler.TableSampler;
import org.apache.pinot.broker.routing.timeboundary.TimeBoundaryManager;
+import org.apache.pinot.common.config.DefaultClusterConfigChangeHandler;
import org.apache.pinot.common.metrics.BrokerGauge;
import org.apache.pinot.common.metrics.BrokerMeter;
import org.apache.pinot.common.metrics.BrokerMetrics;
@@ -73,6 +75,8 @@ import org.testng.annotations.AfterMethod;
import org.testng.annotations.BeforeMethod;
import org.testng.annotations.Test;
+import static
org.apache.pinot.spi.utils.CommonConstants.Broker.CONFIG_OF_PARTITION_PRUNING_PREPARATION_THRESHOLD;
+import static
org.apache.pinot.spi.utils.CommonConstants.Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD;
import static org.mockito.ArgumentMatchers.any;
import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyLong;
@@ -806,6 +810,35 @@ public class BrokerRoutingManagerTest {
assertTrue(_routingManager.getRoutableServerInstanceMap().containsKey(SERVER_INSTANCE_ID));
}
+ @Test
+ public void testPartitionPruningThresholdUpdates()
+ throws Exception {
+ String key = CONFIG_OF_PARTITION_PRUNING_PREPARATION_THRESHOLD;
+ _routingManager = spy(_routingManager);
+ putRoutingEntry(TEST_TABLE, createRoutingEntry(TEST_TABLE, null, null,
Map.of()));
+ DefaultClusterConfigChangeHandler handler = new
DefaultClusterConfigChangeHandler();
+ ClusterConfig config = new ClusterConfig("testCluster");
+ config.getRecord().setSimpleField(key, "128");
+ handler.onClusterConfigChange(config, null);
+ handler.registerClusterConfigChangeListener(_routingManager);
+ assertEquals(_routingManager.getPartitionPruningPreparationThreshold(),
128);
+
+ config.getRecord().setSimpleField(key, "-1");
+ handler.onClusterConfigChange(config, null);
+ assertEquals(_routingManager.getPartitionPruningPreparationThreshold(),
-1);
+
+ config.getRecord().setSimpleField(key, "invalid");
+ handler.onClusterConfigChange(config, null);
+ assertEquals(_routingManager.getPartitionPruningPreparationThreshold(),
+ DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD);
+
+ config.getRecord().getSimpleFields().clear();
+ handler.onClusterConfigChange(config, null);
+ assertEquals(_routingManager.getPartitionPruningPreparationThreshold(),
+ DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD);
+ verify(_routingManager, never()).buildRouting(TEST_TABLE);
+ }
+
/// Creates a ZNRecord representing an enabled server instance.
private ZNRecord createEnabledServerZNRecord(String instanceId) {
ZNRecord record = new ZNRecord(instanceId);
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPrunerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPrunerTest.java
new file mode 100644
index 00000000000..3d37d2ffd8a
--- /dev/null
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPrunerTest.java
@@ -0,0 +1,632 @@
+/**
+ * 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.broker.routing.segmentpruner;
+
+import java.util.ArrayList;
+import java.util.Arrays;
+import java.util.HashSet;
+import java.util.LinkedHashMap;
+import java.util.List;
+import java.util.Map;
+import java.util.Set;
+import java.util.concurrent.Callable;
+import java.util.concurrent.ExecutorService;
+import java.util.concurrent.Executors;
+import java.util.concurrent.Future;
+import java.util.concurrent.atomic.AtomicInteger;
+import javax.annotation.Nullable;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.common.metadata.segment.SegmentPartitionMetadata;
+import org.apache.pinot.common.request.BrokerRequest;
+import org.apache.pinot.common.request.Expression;
+import org.apache.pinot.common.request.PinotQuery;
+import org.apache.pinot.common.utils.request.RequestUtils;
+import org.apache.pinot.segment.spi.partition.PartitionFunction;
+import org.apache.pinot.segment.spi.partition.PartitionFunctionFactory;
+import org.apache.pinot.segment.spi.partition.PartitionIdNormalizer;
+import org.apache.pinot.segment.spi.partition.metadata.ColumnPartitionMetadata;
+import org.apache.pinot.spi.config.table.ColumnPartitionConfig;
+import org.apache.pinot.spi.config.table.RoutingConfig;
+import org.apache.pinot.spi.config.table.SegmentPartitionConfig;
+import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.config.table.TableType;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.CommonConstants.Broker;
+import
org.apache.pinot.spi.utils.CommonConstants.Broker.Request.QueryOptionKey;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertTrue;
+import static org.testng.Assert.expectThrows;
+
+
+/// Exercises query-local partition predicate preparation with real metadata
initialization and refresh, without ZK.
+public class SinglePartitionColumnSegmentPrunerTest {
+ private static final String COLUMN = "memberId";
+ private static final String TABLE = "testTable_OFFLINE";
+
+ @Test
+ public void testThresholdThroughFactory() throws Exception {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.OFFLINE).setTableName(TABLE).build();
+ tableConfig.getIndexingConfig().setSegmentPartitionConfig(
+ new SegmentPartitionConfig(Map.of(COLUMN, new
ColumnPartitionConfig("PrunerCounting", 8))));
+ AtomicInteger clusterThreshold = new AtomicInteger(100);
+ tableConfig.setRoutingConfig(new RoutingConfig(null, List.of("partition"),
null, null));
+ SegmentPruner livePruner =
SegmentPrunerFactory.getSegmentPruners(tableConfig, null,
clusterThreshold::get).get(0);
+ List<String> segments = List.of("first", "second");
+ List<ZNRecord> metadata = List.of(metadata("first", "PrunerCounting", 8,
Set.of(1), null),
+ metadata("second", "PrunerCounting", 8, Set.of(2), null));
+ livePruner.init(null, null, segments, metadata);
+ BrokerRequest request = request(predicate("EQUALS", "1"));
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(livePruner.prune(request, Set.of("first", "second")),
Set.of("first"));
+ assertEquals(CountingPartitionFunction.CALLS.get(), 2);
+ clusterThreshold.set(2);
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(livePruner.prune(request, Set.of("first", "second")),
Set.of("first"));
+ assertEquals(CountingPartitionFunction.CALLS.get(), 1);
+
+ tableConfig.setRoutingConfig(new RoutingConfig(null, List.of("partition"),
null, null, 2));
+ SegmentPruner pruner = SegmentPrunerFactory.getSegmentPruners(tableConfig,
null, 100).get(0);
+ pruner.init(null, null, segments, metadata);
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request, Set.of("first", "second")),
Set.of("first"));
+ assertEquals(CountingPartitionFunction.CALLS.get(), 1);
+
request.getPinotQuery().setQueryOptions(Map.of(QueryOptionKey.PARTITION_PRUNING_PREPARATION_THRESHOLD,
"-1"));
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request, Set.of("first", "second")),
Set.of("first"));
+ assertEquals(CountingPartitionFunction.CALLS.get(), 2);
+ }
+
+ @Test
+ public void
testComputesPartitionIdOnceAcrossDistinctMetadataInstancesPerQuery() throws
Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ Set<String> expected = new HashSet<>();
+ for (int i = 0; i < 256; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "PrunerCounting", 8, Set.of(i %
8), i % 2 == 0 ? null : Map.of()));
+ if (i % 8 == 3) {
+ expected.add(segment);
+ }
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ CountingPartitionFunction.CALLS.set(0);
+ BrokerRequest request = request(predicate("EQUALS", "3"));
+ assertEquals(pruner.prune(request, records.keySet()), expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 1);
+ assertEquals(pruner.prune(request, records.keySet()), expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 2, "Prepared IDs must
not survive a prune call");
+
request.getPinotQuery().setQueryOptions(Map.of(QueryOptionKey.PARTITION_PRUNING_PREPARATION_THRESHOLD,
"-1"));
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request, records.keySet()), expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 256,
+ "Disabling predicate preparation must use per-segment evaluation");
+ }
+
+ @Test
+ public void testInterleavedFunctionConfigurationsAndPartitionCounts() throws
Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ records.put("a", metadata("a", "PrunerCounting", 8, Set.of(3), null));
+ records.put("sameFunction", metadata("sameFunction", "PrunerCounting", 8,
Set.of(2), null));
+ records.put("b", metadata("b", "PrunerCounting", 8, Set.of(4),
Map.of("offset", "1")));
+ records.put("c", metadata("c", "PrunerCounting", 8, Set.of(2), null));
+ records.put("d", metadata("d", "PrunerCounting", 8, Set.of(3),
Map.of("offset", "1")));
+ records.put("e", metadata("e", "PrunerCounting", 16, Set.of(11), null));
+ Set<String> expected = new HashSet<>(Set.of("a", "b", "e"));
+ for (int i = 0; i < 250; i++) {
+ String segment = "tail_" + i;
+ records.put(segment, metadata(segment, "PrunerCounting", 8, Set.of(3),
null));
+ expected.add(segment);
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request(predicate("EQUALS", "11")),
records.keySet()), expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 3,
+ "Each distinct function should prepare its partition ID once");
+ records.put("b", metadata("b", "PrunerCounting", 16, Set.of(11), null));
+ records.put("d", metadata("d", "PrunerCounting", 8, Set.of(3), null));
+ expected.add("d");
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner(records).prune(request(predicate("EQUALS", "11")),
records.keySet()), expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 2,
+ "Different partition counts must not reuse partition IDs");
+ }
+
+ @Test
+ public void testMoreThanEightDistinctPartitionFunctions() throws Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ Set<String> expected = new HashSet<>();
+ for (int offset = 0; offset < 9; offset++) {
+ String segment = "matching_" + offset;
+ records.put(segment, metadata(segment, "PrunerCounting", 16,
Set.of(offset),
+ Map.of("offset", Integer.toString(offset))));
+ expected.add(segment);
+ }
+ for (int offset = 0; offset < 9; offset++) {
+ String segment = "nonmatching_" + offset;
+ records.put(segment, metadata(segment, "PrunerCounting", 16,
Set.of(offset + 1),
+ Map.of("offset", Integer.toString(offset))));
+ }
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner(records, 1).prune(request(predicate("EQUALS", "0")),
records.keySet()), expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 9);
+ }
+
+ @Test
+ public void testEqualConfigurationsReusePartitionIds() throws Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ for (int i = 0; i <
Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "PrunerCounting", 8, Set.of(4),
Map.of("offset", "1")));
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request(predicate("EQUALS", "3")),
records.keySet()), records.keySet());
+ assertEquals(CountingPartitionFunction.CALLS.get(), 1, "Equal nonempty
configs should share computed IDs");
+ pruner.refreshSegment("segment_1", metadata("segment_1", "PrunerCounting",
8, Set.of(4), Map.of("offset", "2")));
+ Set<String> expected = new HashSet<>(records.keySet());
+ expected.remove("segment_1");
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request(predicate("EQUALS", "3")),
records.keySet()), expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 2, "A changed config
must not reuse the first function's ID");
+ }
+
+ @Test
+ public void testLegacyFunctionKeepsFirstHitShortCircuit() throws Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ for (int i = 0; i <
Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "PrunerLegacy", 8, Set.of(0),
null));
+ }
+ PartitionFunction first =
PartitionFunctionFactory.getPartitionFunction("PrunerLegacy", 8, null);
+ assertFalse(first.supportsPartitionIdPreparation());
+ assertNotEquals(first,
PartitionFunctionFactory.getPartitionFunction("PrunerLegacy", 8, null));
+ String[] values = new String[128];
+ Arrays.setAll(values, i -> Integer.toString(i));
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request(predicate("IN", values)),
records.keySet()), records.keySet());
+ assertEquals(CountingPartitionFunction.CALLS.get(), records.size(),
+ "Identity-only plugins should evaluate only the first matching value
per segment");
+ }
+
+ @Test
+ public void testUnrestrictedOrSkipsLaterPartitionValues() throws Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ for (int i = 0; i <
Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "PrunerCounting", 8, Set.of(0),
null));
+ }
+ String[] values = new String[128];
+ Arrays.setAll(values, i -> Integer.toString(i));
+ values[127] = "invalid-number";
+ Expression filter = function("OR",
+ function("EQUALS", RequestUtils.getIdentifierExpression("other"),
RequestUtils.getLiteralExpression("x")),
+ predicate("IN", values));
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner.prune(request(filter), records.keySet()),
records.keySet());
+ assertEquals(CountingPartitionFunction.CALLS.get(), 0,
+ "An unrestricted OR child should skip all later values, including
invalid ones");
+ }
+
+ @Test
+ public void testDuplicateInPartitionsArePreparedOnce() throws Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ records.put("first", metadata("first", "PrunerCounting", 8, Set.of(1),
null));
+ records.put("configured", metadata("configured", "PrunerCounting", 8,
Set.of(3), Map.of("offset", "1")));
+ records.put("second", metadata("second", "PrunerCounting", 8, Set.of(2),
null));
+ records.put("miss", metadata("miss", "PrunerCounting", 8, Set.of(3, 4),
null));
+ records.put("repeat", metadata("repeat", "PrunerCounting", 8, Set.of(1,
2), null));
+ Set<String> expected = new HashSet<>(Set.of("first", "configured",
"second", "repeat"));
+ for (int i = 0; i < 252; i++) {
+ String segment = "repeat_" + i;
+ records.put(segment, metadata(segment, "PrunerCounting", 8, Set.of(1,
2), null));
+ expected.add(segment);
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ CountingPartitionFunction.CALLS.set(0);
+ // Every literal is prepared once for each function, even when several
literals map to the same partition.
+ assertEquals(pruner.prune(request(predicate("IN", "1", "9", "17", "2")),
records.keySet()),
+ expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 8);
+ Set<String> firstHits = Set.of("first", "repeat");
+ CountingPartitionFunction.CALLS.set(0);
+ assertEquals(pruner(records, 1).prune(request(predicate("IN", "1", "9",
"17", "2")), firstHits), firstHits);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 4, "Preparation must
include values after the first hit");
+ // A configured first segment must not seed IDs for later default-config
segments.
+ records.put("first", metadata("first", "PrunerCounting", 8, Set.of(2),
Map.of("offset", "1")));
+ assertEquals(pruner(records).prune(request(predicate("IN", "1", "9", "17",
"2")), records.keySet()), expected);
+ }
+
+ @Test(dataProvider = "candidateCounts")
+ public void testLargeConfigurationsAndUnrelatedPredicates(int numSegments,
@Nullable Integer queryMinSegments)
+ throws Exception {
+ String values = "first|" + "x".repeat(100_000);
+ Map<String, String> config = Map.of("columnValues", values,
"columnValuesDelimiter", "|");
+ ZNRecord first = metadata("first", "BoundedColumnValue", 3, Set.of(1),
config);
+ ZNRecord second = metadata("second", "BoundedColumnValue", 3, Set.of(2),
config);
+ Map<String, ZNRecord> records = new LinkedHashMap<>(Map.of("first", first,
"second", second));
+ for (int i = 2; i < numSegments; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "BoundedColumnValue", 3,
Set.of(2), config));
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ BrokerRequest request = request(predicate("EQUALS", "first"));
+ if (queryMinSegments != null) {
+ request.getPinotQuery().setQueryOptions(
+ Map.of(QueryOptionKey.PARTITION_PRUNING_PREPARATION_THRESHOLD,
queryMinSegments.toString()));
+ }
+ assertEquals(pruner.prune(request, records.keySet()), Set.of("first"));
+ request.getPinotQuery().setFilterExpression(function("EQUALS",
RequestUtils.getIdentifierExpression("other"),
+ RequestUtils.getLiteralExpression("value")));
+ assertEquals(pruner.prune(request, records.keySet()), records.keySet());
+ pruner.refreshSegment("first", metadata("first", "BoundedColumnValue", 3,
Set.of(2), config));
+ request.getPinotQuery().setFilterExpression(predicate("EQUALS", "first"));
+ assertEquals(pruner.prune(request, records.keySet()), Set.of());
+ }
+
+ @Test
+ public void testMixedFunctionsNormalizersAndFunctionConfig() throws
Exception {
+ Map<String, ZNRecord> moduloRecords = new LinkedHashMap<>();
+ moduloRecords.put("positive", metadata("positive", "Modulo", 8, Set.of(7),
null));
+ moduloRecords.put("abs", metadata("abs", "Modulo", 8, Set.of(1),
Map.of("partitionIdNormalizer", "ABS")));
+ moduloRecords.put("wrongAbs", metadata("wrongAbs", "Modulo", 8, Set.of(7),
Map.of("partitionIdNormalizer", "ABS")));
+ Set<String> expected = new HashSet<>(Set.of("positive", "abs"));
+ for (int i = 0; i < 253; i++) {
+ String segment = "positive_" + i;
+ moduloRecords.put(segment, metadata(segment, "Modulo", 8, Set.of(7),
null));
+ expected.add(segment);
+ }
+ assertEquals(pruner(moduloRecords).prune(request(predicate("EQUALS",
"-1")), moduloRecords.keySet()),
+ expected);
+
+ String value = "80ff0102";
+ Map<String, String> rawConfig = Map.of("useRawBytes", "true");
+ int textPartition =
PartitionFunctionFactory.getPartitionFunction("Murmur", 97,
null).getPartition(value);
+ int rawPartition = PartitionFunctionFactory.getPartitionFunction("Murmur",
97, rawConfig).getPartition(value);
+ assertNotEquals(textPartition, rawPartition, "Fixture must distinguish
raw-byte and string hashing");
+ Map<String, ZNRecord> murmurRecords = new LinkedHashMap<>();
+ murmurRecords.put("text", metadata("text", "Murmur", 97,
Set.of(textPartition), null));
+ murmurRecords.put("raw", metadata("raw", "Murmur", 97,
Set.of(rawPartition), rawConfig));
+ murmurRecords.put("wrongRaw", metadata("wrongRaw", "Murmur", 97,
Set.of(textPartition), rawConfig));
+ expected = new HashSet<>(Set.of("text", "raw"));
+ for (int i = 0; i < 253; i++) {
+ String segment = "text_" + i;
+ murmurRecords.put(segment, metadata(segment, "Murmur", 97,
Set.of(textPartition), null));
+ expected.add(segment);
+ }
+ assertEquals(pruner(murmurRecords).prune(request(predicate("EQUALS",
value)), murmurRecords.keySet()),
+ expected);
+
+ Map<String, ZNRecord> lookupRecords = new LinkedHashMap<>();
+ lookupRecords.put("first", metadata("first", "BoundedColumnValue", 3,
Set.of(1),
+ Map.of("columnValues", "11|12", "columnValuesDelimiter", "|")));
+ lookupRecords.put("second", metadata("second", "BoundedColumnValue", 3,
Set.of(2),
+ Map.of("columnValues", "12|11", "columnValuesDelimiter", "|")));
+ lookupRecords.put("modulo", metadata("modulo", "Modulo", 3, Set.of(2),
null));
+ assertEquals(pruner(lookupRecords).prune(request(predicate("EQUALS",
"11")), lookupRecords.keySet()),
+ lookupRecords.keySet());
+ }
+
+ @Test
+ public void testPreparedAndOrWithMultipleSegmentPartitions() throws
Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ records.put("miss", metadata("miss", "Modulo", 8, Set.of(7), null));
+ records.put("multiple", metadata("multiple", "Modulo", 8, Set.of(1, 2),
null));
+ for (int i = 2; i <
Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "Modulo", 8, Set.of(i % 8),
null));
+ }
+ Expression filter = function("AND", function("OR", predicate("EQUALS",
"1"), predicate("IN", "2", "3")),
+ function("EQUALS", RequestUtils.getIdentifierExpression("other"),
+ RequestUtils.getLiteralExpression("value")));
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ BrokerRequest prepared = request(filter);
+ BrokerRequest legacy = request(filter);
+
legacy.getPinotQuery().setQueryOptions(Map.of(QueryOptionKey.PARTITION_PRUNING_PREPARATION_THRESHOLD,
"-1"));
+ Set<String> expected = pruner.prune(legacy, records.keySet());
+ assertTrue(expected.contains("multiple"));
+ assertEquals(pruner.prune(prepared, records.keySet()), expected);
+ assertEquals(pruner.prune(request(function("AND", predicate("EQUALS",
"1"), predicate("EQUALS", "2"))),
+ records.keySet()), Set.of(), "Different rows cannot satisfy an
impossible single-valued conjunction");
+ Expression intersection = function("AND", predicate("IN", "1", "2"),
predicate("IN", "2", "3"));
+ assertEquals(pruner.prune(request(intersection), records.keySet()),
+ pruner.prune(request(predicate("EQUALS", "2")), records.keySet()));
+ assertEquals(pruner.prune(request(function("OR", intersection,
predicate("EQUALS", "4"))), records.keySet()),
+ pruner.prune(request(predicate("IN", "2", "4")), records.keySet()));
+ expectThrows(NumberFormatException.class, () ->
pruner.prune(request(function("AND", predicate("EQUALS", "1"),
+ predicate("EQUALS", "2"), predicate("EQUALS", "invalid-number"))),
records.keySet()));
+ }
+
+ @DataProvider
+ public Object[][] candidateCounts() {
+ int threshold = Broker.DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD;
+ return new Object[][]{{1, null}, {2, null}, {threshold - 1, null},
{threshold, null}, {threshold, -1}};
+ }
+
+ @Test(dataProvider = "candidateCounts")
+ public void testAndOrUnsupportedPredicatesAndShortCircuitErrors(int
numSegments, @Nullable Integer queryMinSegments)
+ throws Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ for (int i = 0; i < numSegments; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "Modulo", 8, Set.of(1), null));
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records,
queryMinSegments);
+ Set<String> segments = records.keySet();
+ Expression invalidValue = predicate("EQUALS", "invalid-number");
+ assertEquals(pruner.prune(request(predicate("IN", "1", "invalid-number")),
segments), segments);
+ expectThrows(NumberFormatException.class,
+ () -> pruner.prune(request(predicate("IN", "2", "invalid-number")),
segments));
+ assertEquals(pruner.prune(request(function("AND", predicate("EQUALS",
"2"), invalidValue)), segments), Set.of());
+ assertEquals(pruner.prune(request(function("OR", predicate("EQUALS", "1"),
invalidValue)), segments), segments);
+ Expression invalidOperator = function("INVALID_OPERATOR");
+ assertEquals(pruner.prune(request(function("OR", predicate("EQUALS", "1"),
invalidOperator)), segments), segments);
+ assertEquals(pruner.prune(request(function("AND", predicate("EQUALS",
"2"), invalidOperator)), segments), Set.of());
+ expectThrows(IllegalArgumentException.class,
+ () -> pruner.prune(request(function("OR", predicate("EQUALS", "2"),
invalidOperator)), segments));
+ expectThrows(IllegalArgumentException.class, () ->
pruner.prune(request(invalidOperator), segments));
+
+ Expression unsupported = predicate("GREATER_THAN", "100");
+ assertEquals(pruner.prune(request(function("AND", predicate("IN", "0",
"1"), unsupported)), segments), segments);
+ assertEquals(pruner.prune(request(function("OR", predicate("EQUALS", "2"),
unsupported)), segments), segments);
+ assertEquals(pruner.prune(request(function("NOT", predicate("EQUALS",
"1"))), segments), segments);
+ assertEquals(pruner.prune(request(function("EQUALS",
RequestUtils.getIdentifierExpression("other"),
+ RequestUtils.getLiteralExpression("invalid-number"))), segments),
segments);
+ Expression transformedColumn = function("LOWER",
RequestUtils.getIdentifierExpression(COLUMN));
+ assertEquals(pruner.prune(request(function("EQUALS", transformedColumn,
+ RequestUtils.getLiteralExpression("invalid-number"))), segments),
segments);
+ BrokerRequest unfilteredRequest = new BrokerRequest();
+ unfilteredRequest.setPinotQuery(new PinotQuery());
+ assertSame(pruner.prune(unfilteredRequest, segments), segments);
+ }
+
+ @Test
+ public void testEmptyCandidatesDoNotEvaluateFilter() throws Exception {
+ SinglePartitionColumnSegmentPruner pruner = pruner(Map.of("one",
metadata("one", "Modulo", 8, Set.of(1), null)));
+ assertEquals(pruner.prune(request(function("INVALID_OPERATOR")),
Set.of()), Set.of());
+ assertEquals(pruner.prune(request(predicate("EQUALS", "invalid-number")),
Set.of()), Set.of());
+ }
+
+ @Test(dataProvider = "candidateCounts")
+ public void testUnknownMetadataIsConservativeAndDoesNotEvaluateFilter(int
numSegments,
+ @Nullable Integer queryMinSegments)
+ throws Exception {
+ SinglePartitionColumnSegmentPruner pruner = pruner(Map.of(),
queryMinSegments);
+ ZNRecord invalid = new ZNRecord("invalid");
+ invalid.setSimpleField(CommonConstants.Segment.PARTITION_METADATA,
"invalid-json");
+ pruner.init(null, null, List.of("missing", "empty", "invalid"),
+ Arrays.asList(null, new ZNRecord("empty"), invalid));
+ Set<String> segments = new HashSet<>(Set.of("missing", "empty", "invalid",
"not-initialized"));
+ for (int i = segments.size(); i < numSegments; i++) {
+ segments.add("not-initialized_" + i);
+ }
+ assertEquals(pruner.prune(request(predicate("EQUALS", "invalid-number")),
segments), segments);
+ assertEquals(pruner.prune(request(function("INVALID_OPERATOR")),
segments), segments);
+ }
+
+ @Test(dataProvider = "candidateCounts")
+ public void testRefreshUsesCurrentPartitionsAndConfiguration(int
numSegments, @Nullable Integer queryMinSegments)
+ throws Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ for (int i = 0; i < numSegments; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "PrunerCounting", 8, Set.of(3),
null));
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records,
queryMinSegments);
+ Set<String> segments = records.keySet();
+ BrokerRequest request = request(predicate("EQUALS", "3"));
+ assertEquals(pruner.prune(request, segments), segments);
+ for (String segment : segments) {
+ pruner.refreshSegment(segment, metadata(segment, "PrunerCounting", 8,
Set.of(4), null));
+ }
+ assertEquals(pruner.prune(request, segments), Set.of());
+ for (String segment : segments) {
+ pruner.refreshSegment(segment, metadata(segment, "PrunerCounting", 8,
Set.of(4), Map.of("offset", "1")));
+ }
+ assertEquals(pruner.prune(request, segments), segments);
+ for (String segment : segments) {
+ pruner.refreshSegment(segment, null);
+ }
+ assertEquals(pruner.prune(request(predicate("EQUALS", "invalid-number")),
segments), segments);
+ }
+
+ @Test
+ public void testConcurrentQueriesKeepSeparatePreparedValues() throws
Exception {
+ Map<String, ZNRecord> records = new LinkedHashMap<>();
+ for (int i = 0; i < 256; i++) {
+ String segment = "segment_" + i;
+ records.put(segment, metadata(segment, "Modulo", 8, Set.of(i % 8),
null));
+ }
+ SinglePartitionColumnSegmentPruner pruner = pruner(records);
+ ExecutorService executor = Executors.newFixedThreadPool(4);
+ try {
+ List<Callable<Void>> queries = new ArrayList<>();
+ for (int i = 0; i < 32; i++) {
+ int partition = i % 8;
+ Set<String> expected = new HashSet<>();
+ for (int segment = partition; segment < 256; segment += 8) {
+ expected.add("segment_" + segment);
+ }
+ queries.add(() -> {
+ assertEquals(pruner.prune(request(predicate("EQUALS",
Integer.toString(partition))), records.keySet()),
+ expected);
+ return null;
+ });
+ }
+ for (Future<Void> future : executor.invokeAll(queries)) {
+ future.get();
+ }
+ } finally {
+ executor.shutdownNow();
+ }
+ }
+
+ private static SinglePartitionColumnSegmentPruner pruner(Map<String,
ZNRecord> records) {
+ SinglePartitionColumnSegmentPruner pruner = new
SinglePartitionColumnSegmentPruner(TABLE, COLUMN);
+ pruner.init(null, null, new ArrayList<>(records.keySet()), new
ArrayList<>(records.values()));
+ return pruner;
+ }
+
+ private static SinglePartitionColumnSegmentPruner pruner(Map<String,
ZNRecord> records,
+ @Nullable Integer queryMinSegments) {
+ SinglePartitionColumnSegmentPruner pruner = new
SinglePartitionColumnSegmentPruner(TABLE, COLUMN) {
+ @Override
+ public Set<String> prune(BrokerRequest request, Set<String> segments) {
+ if (queryMinSegments != null) {
+ request.getPinotQuery().setQueryOptions(
+ Map.of(QueryOptionKey.PARTITION_PRUNING_PREPARATION_THRESHOLD,
queryMinSegments.toString()));
+ }
+ return super.prune(request, segments);
+ }
+ };
+ pruner.init(null, null, new ArrayList<>(records.keySet()), new
ArrayList<>(records.values()));
+ return pruner;
+ }
+
+ private static ZNRecord metadata(String segment, String function, int count,
Set<Integer> partitions,
+ @Nullable Map<String, String> config) throws Exception {
+ ZNRecord record = new ZNRecord(segment);
+ record.setSimpleField(CommonConstants.Segment.PARTITION_METADATA, new
SegmentPartitionMetadata(Map.of(COLUMN,
+ new ColumnPartitionMetadata(function, count, partitions,
config))).toJsonString());
+ return record;
+ }
+
+ private static Expression predicate(String operator, String... values) {
+ List<Expression> operands = new ArrayList<>();
+ operands.add(RequestUtils.getIdentifierExpression(COLUMN));
+ for (String value : values) {
+ operands.add(RequestUtils.getLiteralExpression(value));
+ }
+ return RequestUtils.getFunctionExpression(operator, operands);
+ }
+
+ private static Expression function(String operator, Expression... operands) {
+ return RequestUtils.getFunctionExpression(operator, List.of(operands));
+ }
+
+ private static BrokerRequest request(Expression filter) {
+ PinotQuery pinotQuery = new PinotQuery();
+ pinotQuery.setFilterExpression(filter);
+ BrokerRequest brokerRequest = new BrokerRequest();
+ brokerRequest.setPinotQuery(pinotQuery);
+ return brokerRequest;
+ }
+
+ /// Legacy plugin with identity equality and the default preparation
capability, registered by the factory scan.
+ public static class LegacyPartitionFunction implements PartitionFunction {
+ private static final long serialVersionUID = 1L;
+ private final int _numPartitions;
+
+ public LegacyPartitionFunction(int numPartitions, @Nullable Map<String,
String> config) {
+ _numPartitions = numPartitions;
+ }
+
+ @Override
+ public int getPartition(String value) {
+ CountingPartitionFunction.CALLS.incrementAndGet();
+ return Math.floorMod(Integer.parseInt(value), _numPartitions);
+ }
+
+ @Override
+ public String getName() {
+ return "PrunerLegacy";
+ }
+
+ @Override
+ public int getNumPartitions() {
+ return _numPartitions;
+ }
+
+ @Override
+ public PartitionIdNormalizer getPartitionIdNormalizer() {
+ return PartitionIdNormalizer.POSITIVE_MODULO;
+ }
+ }
+
+ /// Stateless partition function registered by the factory scan with
observable hash calls. Exact-count tests require
+ /// sequential test methods; the concurrent-query fixture deliberately uses
Modulo instead.
+ public static class CountingPartitionFunction implements PartitionFunction {
+ private static final long serialVersionUID = 1L;
+ private static final AtomicInteger CALLS = new AtomicInteger();
+ private final int _numPartitions;
+ private final int _offset;
+ private final Map<String, String> _functionConfig;
+
+ public CountingPartitionFunction(int numPartitions, @Nullable Map<String,
String> config) {
+ _numPartitions = numPartitions;
+ _offset = config == null ? 0 :
Integer.parseInt(config.getOrDefault("offset", "0"));
+ _functionConfig = config == null ? null : Map.copyOf(config);
+ }
+
+ @Override
+ public int getPartition(String value) {
+ CALLS.incrementAndGet();
+ return Math.floorMod(Integer.parseInt(value) + _offset, _numPartitions);
+ }
+
+ @Override
+ public String getName() {
+ return "PrunerCounting";
+ }
+
+ @Override
+ public int getNumPartitions() {
+ return _numPartitions;
+ }
+
+ @Override
+ public Map<String, String> getFunctionConfig() {
+ return _functionConfig;
+ }
+
+ @Override
+ public PartitionIdNormalizer getPartitionIdNormalizer() {
+ return PartitionIdNormalizer.POSITIVE_MODULO;
+ }
+
+ @Override
+ public boolean supportsPartitionIdPreparation() {
+ return true;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ CountingPartitionFunction that = (CountingPartitionFunction) other;
+ return _numPartitions == that._numPartitions && _offset == that._offset;
+ }
+
+ @Override
+ public int hashCode() {
+ return 31 * _numPartitions + _offset;
+ }
+ }
+}
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ByteArrayPartitionFunction.java
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ByteArrayPartitionFunction.java
index 90b1d030d86..138e0b083fd 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ByteArrayPartitionFunction.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ByteArrayPartitionFunction.java
@@ -64,6 +64,30 @@ public class ByteArrayPartitionFunction implements
PartitionFunction {
return _normalizer;
}
+ @Override
+ public boolean supportsPartitionIdPreparation() {
+ return getClass() == ByteArrayPartitionFunction.class;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ ByteArrayPartitionFunction that = (ByteArrayPartitionFunction) other;
+ return _numPartitions == that._numPartitions && _normalizer ==
that._normalizer;
+ }
+
+ @Override
+ public int hashCode() {
+ int result = _numPartitions;
+ result = 31 * result + _normalizer.hashCode();
+ return result;
+ }
+
// Keep it for backward-compatibility, use getName() instead
@Override
public String toString() {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/FnvPartitionFunction.java
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/FnvPartitionFunction.java
index 83fe5a8d5c0..ca288ec581f 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/FnvPartitionFunction.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/FnvPartitionFunction.java
@@ -102,6 +102,34 @@ public class FnvPartitionFunction implements
PartitionFunction {
return _normalizer;
}
+ @Override
+ public boolean supportsPartitionIdPreparation() {
+ return getClass() == FnvPartitionFunction.class;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ FnvPartitionFunction that = (FnvPartitionFunction) other;
+ return _numPartitions == that._numPartitions && _normalizer ==
that._normalizer
+ && _variant == that._variant
+ && _useRawBytes == that._useRawBytes;
+ }
+
+ @Override
+ public int hashCode() {
+ int result = _numPartitions;
+ result = 31 * result + _normalizer.hashCode();
+ result = 31 * result + _variant.hashCode();
+ result = 31 * result + Boolean.hashCode(_useRawBytes);
+ return result;
+ }
+
// Keep it for backward-compatibility, use getName() instead
@Override
public String toString() {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/HashCodePartitionFunction.java
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/HashCodePartitionFunction.java
index a6f5c5867d7..b36456bbb56 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/HashCodePartitionFunction.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/HashCodePartitionFunction.java
@@ -61,6 +61,30 @@ public class HashCodePartitionFunction implements
PartitionFunction {
return _normalizer;
}
+ @Override
+ public boolean supportsPartitionIdPreparation() {
+ return getClass() == HashCodePartitionFunction.class;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ HashCodePartitionFunction that = (HashCodePartitionFunction) other;
+ return _numPartitions == that._numPartitions && _normalizer ==
that._normalizer;
+ }
+
+ @Override
+ public int hashCode() {
+ int result = _numPartitions;
+ result = 31 * result + _normalizer.hashCode();
+ return result;
+ }
+
// Keep it for backward-compatibility, use getName() instead
@Override
public String toString() {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ModuloPartitionFunction.java
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ModuloPartitionFunction.java
index 8d34d855a75..48c2b35577f 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ModuloPartitionFunction.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/ModuloPartitionFunction.java
@@ -60,6 +60,30 @@ public class ModuloPartitionFunction implements
PartitionFunction {
return _normalizer;
}
+ @Override
+ public boolean supportsPartitionIdPreparation() {
+ return getClass() == ModuloPartitionFunction.class;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ ModuloPartitionFunction that = (ModuloPartitionFunction) other;
+ return _numPartitions == that._numPartitions && _normalizer ==
that._normalizer;
+ }
+
+ @Override
+ public int hashCode() {
+ int result = _numPartitions;
+ result = 31 * result + _normalizer.hashCode();
+ return result;
+ }
+
// Keep it for backward-compatibility, use getName() instead
@Override
public String toString() {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/Murmur3PartitionFunction.java
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/Murmur3PartitionFunction.java
index 19b35c1aae3..1cd9012e8f4 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/Murmur3PartitionFunction.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/Murmur3PartitionFunction.java
@@ -114,6 +114,35 @@ public class Murmur3PartitionFunction implements
PartitionFunction {
return _normalizer;
}
+ @Override
+ public boolean supportsPartitionIdPreparation() {
+ return getClass() == Murmur3PartitionFunction.class;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ Murmur3PartitionFunction that = (Murmur3PartitionFunction) other;
+ return _numPartitions == that._numPartitions && _normalizer ==
that._normalizer
+ && _seed == that._seed && _useX64 == that._useX64
+ && _useRawBytes == that._useRawBytes;
+ }
+
+ @Override
+ public int hashCode() {
+ int result = _numPartitions;
+ result = 31 * result + _normalizer.hashCode();
+ result = 31 * result + _seed;
+ result = 31 * result + Boolean.hashCode(_useX64);
+ result = 31 * result + Boolean.hashCode(_useRawBytes);
+ return result;
+ }
+
// Keep it for backward-compatibility, use getName() instead
@Override
public String toString() {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/MurmurPartitionFunction.java
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/MurmurPartitionFunction.java
index 5396fbe07cc..5a8383aa6a6 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/partition/function/MurmurPartitionFunction.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/partition/function/MurmurPartitionFunction.java
@@ -86,6 +86,32 @@ public class MurmurPartitionFunction implements
PartitionFunction {
return _normalizer;
}
+ @Override
+ public boolean supportsPartitionIdPreparation() {
+ return getClass() == MurmurPartitionFunction.class;
+ }
+
+ @Override
+ public boolean equals(Object other) {
+ if (this == other) {
+ return true;
+ }
+ if (other == null || getClass() != other.getClass()) {
+ return false;
+ }
+ MurmurPartitionFunction that = (MurmurPartitionFunction) other;
+ return _numPartitions == that._numPartitions && _normalizer ==
that._normalizer
+ && _useRawBytes == that._useRawBytes;
+ }
+
+ @Override
+ public int hashCode() {
+ int result = _numPartitions;
+ result = 31 * result + _normalizer.hashCode();
+ result = 31 * result + Boolean.hashCode(_useRawBytes);
+ return result;
+ }
+
// Keep it for backward-compatibility, use getName() instead
@Override
public String toString() {
diff --git
a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
index 3b0da6329eb..f5f8168ca5a 100644
---
a/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
+++
b/pinot-common/src/main/java/org/apache/pinot/common/utils/config/QueryOptionsUtils.java
@@ -353,6 +353,14 @@ public class QueryOptionsUtils {
queryOptions.get(QueryOptionKey.IN_PREDICATE_PRUNING_THRESHOLD));
}
+ /// Query-level override for the minimum candidate segments needed to
prepare the partition-pruning predicate.
+ /// A negative value disables preparation.
+ @Nullable
+ public static Integer getPartitionPruningPreparationThreshold(@Nullable
Map<String, String> queryOptions) {
+ return queryOptions != null ?
uncheckedParseInt(QueryOptionKey.PARTITION_PRUNING_PREPARATION_THRESHOLD,
+
queryOptions.get(QueryOptionKey.PARTITION_PRUNING_PREPARATION_THRESHOLD)) :
null;
+ }
+
/// Returns whether materialized-view rewrite is allowed for this query.
Defaults to `true`
/// (absence ⇒ rewrite allowed) for backward compatibility with the
pre-option behavior; the
/// MV minion executor sets it to `false` so a materialization query is
never rewritten back
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 6e55464c472..ffce3322e6b 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
@@ -28,10 +28,12 @@ import org.apache.pinot.spi.utils.BytesUtils;
import org.apache.pinot.spi.utils.JsonUtils;
import org.apache.pinot.spi.utils.hash.FnvHashFunctions;
import org.apache.pinot.spi.utils.hash.MurmurHashFunctions;
+import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
import static java.nio.charset.StandardCharsets.UTF_8;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
import static org.testng.Assert.assertNotEquals;
import static org.testng.Assert.assertTrue;
import static org.testng.Assert.expectThrows;
@@ -42,6 +44,117 @@ public class PartitionFunctionTest {
private static final int NUM_ROUNDS = 1000;
private static final int MAX_NUM_PARTITIONS = 100;
+ @Test
+ public void testPartitionIdReuse() {
+ PartitionFunction function = new MurmurPartitionFunction(8, null);
+ PartitionFunction emptyConfig = new MurmurPartitionFunction(8, Map.of());
+ assertTrue(function.equals(emptyConfig));
+ assertTrue(emptyConfig.equals(function));
+ PartitionFunction third = new MurmurPartitionFunction(8, Map.of());
+ assertTrue(emptyConfig.equals(third));
+ assertTrue(function.equals(third));
+ assertFalse(function.equals(new MurmurPartitionFunction(16, null)));
+ assertFalse(function.equals(new Murmur3PartitionFunction(8, null)));
+
+ PartitionFunction configured = new MurmurPartitionFunction(8,
Map.of("useRawBytes", "true"));
+ assertTrue(configured.equals(configured));
+ assertFalse(function.equals(configured));
+ assertFalse(configured.equals(function));
+ assertTrue(configured.equals(new MurmurPartitionFunction(8,
Map.of("useRawBytes", "true"))));
+
+ PartitionFunction modulo = new ModuloPartitionFunction(8, null);
+ PartitionFunction abs = new ModuloPartitionFunction(8,
Map.of("partitionIdNormalizer", "ABS"));
+ assertFalse(modulo.equals(abs));
+ assertFalse(abs.equals(modulo));
+ }
+
+ @DataProvider
+ public Object[][] reuseConfigurations() {
+ return new Object[][]{
+ {"Murmur", Map.of("useRawBytes", "true"), Map.of("useRawBytes",
"TRUE"),
+ Map.of("useRawBytes", "false")},
+ {"Murmur3", Map.of("seed", "7"), Map.of("seed", "07", "variant",
"x86_32"), Map.of("seed", "8")},
+ {"Murmur3", Map.of("variant", "x64_32"), Map.of("variant", "x64_32",
"seed", "0"),
+ Map.of("variant", "x86_32")},
+ {"Murmur3", Map.of("useRawBytes", "true"), Map.of("useRawBytes",
"TRUE"),
+ Map.of("useRawBytes", "false")},
+ {"FNV", Map.of("variant", "FNV1_64"), Map.of("variant", "FNV1_64",
"useRawBytes", "false"),
+ Map.of("variant", "FNV1A_64")},
+ {"FNV", Map.of("useRawBytes", "true"), Map.of("useRawBytes", "TRUE"),
Map.of("useRawBytes", "false")},
+ {"Modulo", Map.of("partitionIdNormalizer", "ABS"),
Map.of("partitionIdNormalizer", " abs "),
+ Map.of("partitionIdNormalizer", "POSITIVE_MODULO")},
+ {"HashCode", Map.of(), Map.of("partitionIdNormalizer",
"PRE_MODULO_ABS"),
+ Map.of("partitionIdNormalizer", "MASK")},
+ {"ByteArray", Map.of(), Map.of("partitionIdNormalizer",
"PRE_MODULO_ABS"),
+ Map.of("partitionIdNormalizer", "MASK")}
+ };
+ }
+
+ @Test
+ public void testComplexFunctionsAndSubclasses() {
+ Map<String, String> config = Map.of("columnValues", "a|b",
"columnValuesDelimiter", "|");
+ PartitionFunction bounded = new BoundedColumnValuePartitionFunction(3,
config);
+ assertFalse(bounded.supportsPartitionIdPreparation());
+ assertTrue(bounded.equals(bounded));
+ assertFalse(bounded.equals(new BoundedColumnValuePartitionFunction(3,
config)));
+ class CustomMurmur extends MurmurPartitionFunction {
+ CustomMurmur() {
+ super(8, null);
+ }
+ }
+ PartitionFunction subclass = new CustomMurmur();
+ assertFalse(subclass.supportsPartitionIdPreparation());
+ PartitionFunction another = new CustomMurmur();
+ assertTrue(subclass.equals(subclass));
+ assertTrue(subclass.equals(another));
+ assertFalse(subclass.equals(new MurmurPartitionFunction(8, null)));
+ assertFalse(new MurmurPartitionFunction(8, null).equals(subclass));
+ }
+
+ @Test
+ public void testHashCollisionDoesNotImplyEquality() {
+ // Increasing the partition count and decreasing the seed compensates in
the 31-based hash.
+ PartitionFunction first = new Murmur3PartitionFunction(8, Map.of("seed",
"961"));
+ PartitionFunction second = new Murmur3PartitionFunction(9, Map.of("seed",
"0"));
+ assertEquals(first.hashCode(), second.hashCode());
+ assertNotEquals(first, second);
+ assertFalse(first.equals(second));
+ }
+
+ @Test(dataProvider = "reuseConfigurations")
+ public void testConfiguredPartitionIdReuse(String name, Map<String, String>
config,
+ Map<String, String> equivalentConfig, Map<String, String>
differentConfig) {
+ Map<String, String> mutableConfig = new HashMap<>(config);
+ PartitionFunction first =
PartitionFunctionFactory.getPartitionFunction(name, 8, mutableConfig);
+ assertTrue(first.supportsPartitionIdPreparation());
+ PartitionFunction second =
PartitionFunctionFactory.getPartitionFunction(name, 8, equivalentConfig);
+ PartitionFunction third =
PartitionFunctionFactory.getPartitionFunction(name, 8, new
HashMap<>(equivalentConfig));
+ assertTrue(first.equals(first));
+ assertEquals(first, second);
+ assertEquals(first.hashCode(), second.hashCode());
+ assertFalse(first.equals(null));
+ assertTrue(first.equals(second));
+ assertTrue(second.equals(first));
+ assertTrue(second.equals(third));
+ assertTrue(first.equals(third));
+ for (String value : new String[]{"00", "1234", "abcd"}) {
+ if (!name.equals("Modulo") || !value.equals("abcd")) {
+ assertEquals(first.getPartition(value), second.getPartition(value));
+ }
+ }
+ PartitionFunction different =
PartitionFunctionFactory.getPartitionFunction(name, 8, differentConfig);
+ assertFalse(first.equals(different));
+ assertFalse(different.equals(first));
+
assertFalse(first.equals(PartitionFunctionFactory.getPartitionFunction(name,
16, config)));
+ Map<String, String> otherNormalizer = new HashMap<>(config);
+ otherNormalizer.put("partitionIdNormalizer",
+ first.getPartitionIdNormalizer().name().equals("MASK") ? "ABS" :
"MASK");
+
assertFalse(first.equals(PartitionFunctionFactory.getPartitionFunction(name, 8,
otherNormalizer)));
+ mutableConfig.clear();
+ mutableConfig.putAll(differentConfig);
+ assertTrue(first.equals(second), "Caller mutations must not change the
effective function");
+ }
+
/// Unit test for [ModuloPartitionFunction].
///
/// - Builds an instance of the [ModuloPartitionFunction].
@@ -143,6 +256,7 @@ public class PartitionFunctionTest {
// Create partition function with function config present but no seed
value present.
PartitionFunction partitionFunction2 =
PartitionFunctionFactory.getPartitionFunction(functionName,
numPartitions, functionConfig);
+ testBasicProperties(partitionFunction2, functionName, numPartitions,
functionConfig);
// Get partition number with random value.
int partitionNumWithNoSeedValue =
partitionFunction2.getPartition(valueTobeHashed);
@@ -156,12 +270,14 @@ public class PartitionFunctionTest {
// Create partition function with function config present but random
seed value present in function config.
PartitionFunction partitionFunction3 =
PartitionFunctionFactory.getPartitionFunction(functionName,
numPartitions, functionConfig);
+ testBasicProperties(partitionFunction3, functionName, numPartitions,
functionConfig);
// Create partition function with function config present with random
seed value
// and with variant provided as "x64_32" in function config.
functionConfig.put("variant", "x64_32");
PartitionFunction partitionFunction4 =
PartitionFunctionFactory.getPartitionFunction(functionName,
numPartitions, functionConfig);
+ testBasicProperties(partitionFunction4, functionName, numPartitions,
functionConfig);
// Put variant value as "x86_32" in function config.
functionConfig.put("variant", "x86_32");
@@ -172,6 +288,7 @@ public class PartitionFunctionTest {
// Create partition function with function config present with variant
provided as "x86_32" in function config.
PartitionFunction partitionFunction5 =
PartitionFunctionFactory.getPartitionFunction(functionName,
numPartitions, functionConfig);
+ testBasicProperties(partitionFunction5, functionName, numPartitions,
functionConfig);
// Partition number should be equal as partitionNumWithNullConfig and
partitionNumWithNoSeedValue as this is
// default behavior.
@@ -184,6 +301,7 @@ public class PartitionFunctionTest {
// value in functionConfig.
PartitionFunction partitionFunction6 =
PartitionFunctionFactory.getPartitionFunction(functionName,
numPartitions, functionConfig);
+ testBasicProperties(partitionFunction6, functionName, numPartitions,
functionConfig);
// Partition number should be equal as partitionNumWithNullConfig and
partitionNumWithNoSeedValue as this is
// default behavior.
@@ -201,11 +319,6 @@ public class PartitionFunctionTest {
assertEquals(partitionFunction7.getPartition(valueTobeHashed),
partitionNumWithNullConfig);
testBasicProperties(partitionFunction1, functionName, numPartitions);
- testBasicProperties(partitionFunction2, functionName, numPartitions,
functionConfig);
- testBasicProperties(partitionFunction3, functionName, numPartitions,
functionConfig);
- testBasicProperties(partitionFunction4, functionName, numPartitions,
functionConfig);
- testBasicProperties(partitionFunction5, functionName, numPartitions,
functionConfig);
- testBasicProperties(partitionFunction6, functionName, numPartitions,
functionConfig);
testBasicProperties(partitionFunction7, functionName, numPartitions,
functionConfig);
for (int j = 0; j < NUM_ROUNDS; j++) {
diff --git
a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
index 73d248c5c09..b432f47c948 100644
---
a/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
+++
b/pinot-common/src/test/java/org/apache/pinot/common/utils/config/QueryOptionsUtilsTest.java
@@ -120,6 +120,16 @@ public class QueryOptionsUtilsTest {
assertNull(QueryOptionsUtils.getInPredicatePruningThreshold(Map.of()));
}
+ @Test
+ public void shouldReadPartitionPruningPreparationThresholdOption() {
+ assertEquals(QueryOptionsUtils.getPartitionPruningPreparationThreshold(
+ Map.of(PARTITION_PRUNING_PREPARATION_THRESHOLD, "20")), 20);
+ assertEquals(QueryOptionsUtils.getPartitionPruningPreparationThreshold(
+ Map.of(PARTITION_PRUNING_PREPARATION_THRESHOLD, "-1")), -1);
+
assertNull(QueryOptionsUtils.getPartitionPruningPreparationThreshold(Map.of()));
+
assertNull(QueryOptionsUtils.getPartitionPruningPreparationThreshold(null));
+ }
+
@Test(expectedExceptions = IllegalArgumentException.class)
public void shouldRejectInvalidInPredicatePruningThreshold() {
QueryOptionsUtils.getInPredicatePruningThreshold(Map.of(IN_PREDICATE_PRUNING_THRESHOLD,
"invalid"));
diff --git
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunction.java
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunction.java
index 69e9c5cfffd..743105bc25a 100644
---
a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunction.java
+++
b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunction.java
@@ -30,6 +30,10 @@ import javax.annotation.Nullable;
/// Implementations of this interface are assumed not to be stateful. That is,
two invocations of
/// `PartitionFunction.getPartition(value)` with the same value are expected
to produce the same
/// result. Implementations must also be safe for concurrent invocation by
multiple threads.
+/// Functions with equal identity may reuse partition ids during broker
pruning. Implementations that override
+/// `equals` and `hashCode` must include every setting that affects partition
ids, including subclass state.
+/// The raw function-config map need not be compared when parsed fields
already capture all partition-ID-affecting
+/// settings; unrelated config entries should not prevent reuse.
public interface PartitionFunction extends Serializable {
/// Method to compute and return partition id for the given value.
@@ -58,6 +62,13 @@ public interface PartitionFunction extends Serializable {
/// @return Number of possible partitions.
int getNumPartitions();
+ /// Whether query-local predicate preparation is useful for this function.
Implementations must opt in only when
+ /// they provide cheap value equality across segments. The default preserves
per-segment pruning for legacy plugins.
+ default boolean supportsPartitionIdPreparation() {
+ return false;
+ }
+
+ /// Returns the configuration used to construct this function, or null if
none is exposed.
@Nullable
default Map<String, String> getFunctionConfig() {
return null;
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/RoutingConfig.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/RoutingConfig.java
index b3cde3ea82b..d88428b0a09 100644
---
a/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/RoutingConfig.java
+++
b/pinot-spi/src/main/java/org/apache/pinot/spi/config/table/RoutingConfig.java
@@ -42,16 +42,26 @@ public class RoutingConfig extends BaseJsonConfig {
private final List<String> _segmentPrunerTypes;
private final String _instanceSelectorType;
private final Boolean _useFixedReplica;
+ private final Integer _partitionPruningPreparationThreshold;
- @JsonCreator
public RoutingConfig(@JsonProperty("routingTableBuilderName") @Nullable
String routingTableBuilderName,
@JsonProperty("segmentPrunerTypes") @Nullable List<String>
segmentPrunerTypes,
@JsonProperty("instanceSelectorType") @Nullable String
instanceSelectorType,
@JsonProperty("useFixedReplica") @Nullable Boolean useFixedReplica) {
+ this(routingTableBuilderName, segmentPrunerTypes, instanceSelectorType,
useFixedReplica, null);
+ }
+
+ @JsonCreator
+ public RoutingConfig(@JsonProperty("routingTableBuilderName") @Nullable
String routingTableBuilderName,
+ @JsonProperty("segmentPrunerTypes") @Nullable List<String>
segmentPrunerTypes,
+ @JsonProperty("instanceSelectorType") @Nullable String
instanceSelectorType,
+ @JsonProperty("useFixedReplica") @Nullable Boolean useFixedReplica,
+ @JsonProperty("partitionPruningPreparationThreshold") @Nullable Integer
partitionPruningPreparationThreshold) {
_routingTableBuilderName = routingTableBuilderName;
_segmentPrunerTypes = segmentPrunerTypes;
_instanceSelectorType = instanceSelectorType;
_useFixedReplica = useFixedReplica;
+ _partitionPruningPreparationThreshold =
partitionPruningPreparationThreshold;
}
@Nullable
@@ -73,4 +83,11 @@ public class RoutingConfig extends BaseJsonConfig {
public Boolean getUseFixedReplica() {
return _useFixedReplica;
}
+
+ /// Returns the table-level override for the minimum candidate segments
needed to prepare a partition-pruning
+ /// predicate. A negative value disables preparation. Null means use the
cluster-level value.
+ @Nullable
+ public Integer getPartitionPruningPreparationThreshold() {
+ return _partitionPruningPreparationThreshold;
+ }
}
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
index 7ab92db3953..70b97dbc3f5 100644
--- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
+++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/CommonConstants.java
@@ -339,6 +339,10 @@ public class CommonConstants {
}
public static class Broker {
+ /// Minimum candidate segments for preparing a partition-pruning
predicate. A negative value disables preparation.
+ public static final String
CONFIG_OF_PARTITION_PRUNING_PREPARATION_THRESHOLD =
+ "pinot.broker.partition.pruning.preparation.threshold";
+ public static final int DEFAULT_PARTITION_PRUNING_PREPARATION_THRESHOLD =
32;
public static final String ROUTING_TABLE_CONFIG_PREFIX =
"pinot.broker.routing.table";
public static final String ACCESS_CONTROL_CONFIG_PREFIX =
"pinot.broker.access.control";
/// Namespace for service credentials used by the broker when invoking
Server admin APIs.
@@ -1029,6 +1033,10 @@ public class CommonConstants {
/// Query-level override for `inpredicate.threshold`. Negative means
always prune.
public static final String IN_PREDICATE_PRUNING_THRESHOLD =
"inPredicatePruningThreshold";
+ /// Query-level override for the minimum candidate segments needed to
prepare the partition-pruning predicate.
+ /// A negative value disables preparation, but does not disable
partition pruning itself.
+ public static final String PARTITION_PRUNING_PREPARATION_THRESHOLD =
"partitionPruningPreparationThreshold";
+
// When evaluating REGEXP_LIKE predicate on a dictionary encoded
column:
// - If dictionary size is smaller than this threshold, scan the
dictionary to get the matching dictionary ids
// first, where inverted index can be applied if exists
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
