gortiz commented on code in PR #19539:
URL: https://github.com/apache/pinot/pull/19539#discussion_r4063350625
##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunction.java:
##########
@@ -63,6 +63,26 @@ default Map<String, String> getFunctionConfig() {
return null;
}
+ /// Returns whether the exposed function configuration is null or empty.
This does not imply that other settings,
+ /// such as the partition id normalizer, have their default values.
+ @JsonIgnore
+ default boolean hasEmptyConfig() {
+ Map<String, String> config = getFunctionConfig();
+ return config == null || config.isEmpty();
+ }
+
+ /// Returns whether partition ids computed by this function can be reused
for the non-null `other` function.
+ /// A true result must guarantee identical results for every input value.
False is conservative, not proof that the
+ /// functions differ. Implementations with additional output-affecting state
must account for it in this method.
+ ///
+ /// The default only permits matching functions with empty exposed
configurations, without comparing config contents.
+ /// Configured implementations may override this method to compare their
effective settings.
+ default boolean canReusePartitionIds(PartitionFunction other) {
+ return hasEmptyConfig() && other.hasEmptyConfig() && getClass() ==
other.getClass()
Review Comment:
The default is opt-OUT for existing third-party PartitionFunction plugins:
any impl whose behaviour is not fully captured by (class, name, numPartitions,
normalizer) and that returns null from getFunctionConfig() now silently shares
partition ids.
All seven OSS impls (Murmur, Murmur3, Modulo, Fnv, HashCode, ByteArray,
BoundedColumnValue) derive every output-affecting field from the function
config, so hasEmptyConfig() plus class/count/normalizer really is sufficient
for them - I checked each one. But a plugin that loads a lookup table from
disk, or takes state from a table-config field it does not echo back in
getFunctionConfig(), will be recompiled against this default and start reusing
ids across segments that disagree. The javadoc warns implementers, but they
have to know to go read it. Safer shape is opt-IN (default false, built-ins
override) or config-content equality: getClass() == other.getClass() &&
Objects.equals(getFunctionConfig(), other.getFunctionConfig()) &&
getNumPartitions() == ... && getPartitionIdNormalizer() == ... . The second
form is strictly more capable too - it would let the BoundedColumnValue 100 KB
case share ids, which today is the one workload that gets no benefit (-3.1%,
i.e. noise), and a confi
g == other.getFunctionConfig() identity fast path makes the common interned
case free.
##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunction.java:
##########
@@ -63,6 +63,26 @@ default Map<String, String> getFunctionConfig() {
return null;
}
+ /// Returns whether the exposed function configuration is null or empty.
This does not imply that other settings,
+ /// such as the partition id normalizer, have their default values.
+ @JsonIgnore
+ default boolean hasEmptyConfig() {
+ Map<String, String> config = getFunctionConfig();
+ return config == null || config.isEmpty();
+ }
+
+ /// Returns whether partition ids computed by this function can be reused
for the non-null `other` function.
+ /// A true result must guarantee identical results for every input value.
False is conservative, not proof that the
+ /// functions differ. Implementations with additional output-affecting state
must account for it in this method.
+ ///
+ /// The default only permits matching functions with empty exposed
configurations, without comparing config contents.
+ /// Configured implementations may override this method to compare their
effective settings.
+ default boolean canReusePartitionIds(PartitionFunction other) {
Review Comment:
This is the load-bearing contract of the whole PR: a wrong 'true' silently
drops segments and loses rows with no error anywhere.
The pruner picks the first valid segment's function as a pivot and then
fills one shared partition-id cache from whichever segment answers true. For
that to be safe, the relation must be symmetric AND transitive over the pivot's
compatibility class. The javadoc's phrasing 'ids computed by THIS function can
be reused FOR other' reads directional, while the guarantee sentence
('identical results for every input value') is bidirectional. An implementer
reading only the first sentence could write a legitimately asymmetric override
(e.g. 'other is coarser / a superset'), and the pruner would then cache ids
produced by a non-pivot function and hand them to a third segment. Worth
stating the required algebra explicitly: canReusePartitionIds must be
reflexive, symmetric and transitive, and true implies getPartition(v) is equal
for both functions for every v.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java:
##########
@@ -94,6 +98,13 @@ public Set<String> prune(BrokerRequest brokerRequest,
Set<String> segments) {
if (filterExpression == null) {
return segments;
}
+ int numSegments = segments.size();
+ if (numSegments == 0) {
+ return segments;
+ }
+ if (numSegments >= MIN_SEGMENTS_FOR_PREPARATION) {
Review Comment:
Hardcoded 256-segment cutoff with no escape hatch, and it creates a second
evaluator that must stay semantically identical to isPartitionMatch() forever.
Two concerns. (1) No broker config or query option can force the prepared
path on or off. The PR's own numbers show the interleaved case is bimodal in
both versions and landed 14.7% slower on the pooled mean; if that shows up in a
real cluster there is no knob to turn. A `private static final int` read from a
CommonConstants broker config, defaulting to 256, would cost nothing. (2) The
old loop and isPartitionMatch() stay, so there are now two implementations of
the same predicate semantics - including the subtle bits (identifier must equal
_partitionColumn, unsupported kinds return true, OR short-circuits before
FilterKind.valueOf can throw). The tests do cover both sides via the
candidateCounts data provider, but only for some of the cases; a single
data-provider-driven equivalence test that runs every fixture at 255 and 256
candidates and asserts the two paths return the same set would pin this down
permanently.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java:
##########
@@ -105,6 +116,30 @@ public Set<String> prune(BrokerRequest brokerRequest,
Set<String> segments) {
return selectedSegments;
}
+ private Set<String> pruneWithPreparedPredicate(Expression filterExpression,
Set<String> segments) {
+ Set<String> selectedSegments = new HashSet<>();
+ PreparedPredicate predicate = null;
+ PartitionFunction cachedFunction = 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 (predicate == null) {
+ predicate = new PreparedPredicate(filterExpression);
+ cachedFunction = function;
+ }
+ // Reuse the filter structure even when the functions cannot share
partition ids.
+ boolean reusePartitionIds =
cachedFunction.canReusePartitionIds(function);
+ if (predicate.matches(partitionInfo.getPartitions(), function,
reusePartitionIds)) {
Review Comment:
Suggestion: once the prepared tree is built and no leaf touched the
partition column, the answer is 'all segments' for every remaining segment -
you can return `segments` and skip the loop entirely.
For the unrelated-predicate workload the current code still pays a
getPartitionFunction() read plus a virtual canReusePartitionIds() call per
segment, and the description measures only -7.7% there. If every leaf ended up
with _partitionIds == null the predicate is a constant true, so the whole prune
degenerates to O(1). Careful with two things: segments with null/INVALID
partition info must still be kept (they are, since the answer is 'keep
everything'), and the exception behaviour must not change - FilterKind.valueOf
on a bad operator has to still throw on the first evaluated segment, which it
does because the tree is built before you can know the answer.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java:
##########
@@ -152,4 +187,82 @@ private boolean isPartitionMatch(Expression
filterExpression, SegmentPartitionIn
return true;
}
}
+
+ /// Lazily prepares only visited expressions and values. Instances belong to
one prune call, never shared by queries.
+ private final class PreparedPredicate {
+ private final Expression _expression;
+ private FilterKind _kind;
+ private List<Expression> _operands;
+ private PreparedPredicate[] _children;
+ private List<Integer> _partitionIds;
+
+ private PreparedPredicate(Expression expression) {
+ _expression = expression;
+ }
+
+ private boolean matches(Set<Integer> partitions, PartitionFunction
partitionFunction, boolean reusePartitionIds) {
+ if (_kind == null) {
+ Function function = _expression.getFunctionCall();
+ _kind = FilterKind.valueOf(function.getOperator());
+ _operands = function.getOperands();
+ if (_kind == FilterKind.AND || _kind == FilterKind.OR) {
+ _children = new PreparedPredicate[_operands.size()];
+ for (int i = 0; i < _children.length; i++) {
+ _children[i] = new PreparedPredicate(_operands.get(i));
+ }
+ } else if (_kind == FilterKind.EQUALS || _kind == FilterKind.IN) {
+ Identifier identifier = _operands.get(0).getIdentifier();
+ if (identifier != null &&
identifier.getName().equals(_partitionColumn)) {
+ _partitionIds = new ArrayList<>(1);
Review Comment:
new ArrayList<Integer>(1) for an IN list of up to N values: repeated array
growth plus a boxed Integer per value, on the path this PR exists to speed up.
_operands.size() - 1 is already known here, so size it (or use an Integer[]
with a fill count - the index is always dense and appended at position size(),
so an array is a drop-in). Keeping boxed Integer objects is right, since
`partitions` is a Set<Integer> and re-boxing per segment would undo part of the
win, but growing 1 -> 2 -> 3 -> ... -> 128 does ~8 copies for the 128-value IN
benchmark.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java:
##########
@@ -152,4 +187,82 @@ private boolean isPartitionMatch(Expression
filterExpression, SegmentPartitionIn
return true;
}
}
+
+ /// Lazily prepares only visited expressions and values. Instances belong to
one prune call, never shared by queries.
+ private final class PreparedPredicate {
+ private final Expression _expression;
+ private FilterKind _kind;
+ private List<Expression> _operands;
+ private PreparedPredicate[] _children;
+ private List<Integer> _partitionIds;
+
+ private PreparedPredicate(Expression expression) {
+ _expression = expression;
+ }
+
+ private boolean matches(Set<Integer> partitions, PartitionFunction
partitionFunction, boolean reusePartitionIds) {
+ if (_kind == null) {
+ Function function = _expression.getFunctionCall();
+ _kind = FilterKind.valueOf(function.getOperator());
+ _operands = function.getOperands();
+ if (_kind == FilterKind.AND || _kind == FilterKind.OR) {
+ _children = new PreparedPredicate[_operands.size()];
+ for (int i = 0; i < _children.length; i++) {
+ _children[i] = new PreparedPredicate(_operands.get(i));
+ }
+ } else if (_kind == FilterKind.EQUALS || _kind == FilterKind.IN) {
+ Identifier identifier = _operands.get(0).getIdentifier();
+ if (identifier != null &&
identifier.getName().equals(_partitionColumn)) {
+ _partitionIds = new ArrayList<>(1);
+ }
+ }
+ }
+ switch (_kind) {
+ case AND:
+ for (PreparedPredicate child : _children) {
+ if (!child.matches(partitions, partitionFunction,
reusePartitionIds)) {
+ return false;
+ }
+ }
+ return true;
+ case OR:
+ for (PreparedPredicate child : _children) {
+ if (child.matches(partitions, partitionFunction,
reusePartitionIds)) {
+ return true;
+ }
+ }
+ return false;
+ case EQUALS:
+ case IN:
+ if (_partitionIds != null) {
+ int numValues = _kind == FilterKind.EQUALS ? 1 : _operands.size()
- 1;
+ if (!reusePartitionIds) {
Review Comment:
Not a finding - recording that I verified this, because it is the part of
the diff most likely to be misread as a bug.
numValues maps i in [0, numValues) to _operands.get(i + 1), which matches
the original loop over [1, numOperands). Early `return true` on a hit leaves
_partitionIds partially filled, but the next segment resumes at exactly i ==
_partitionIds.size() and appends there, so index i always holds the id for
operand i + 1. The `if (!reusePartitionIds)` branch never reads or writes the
cache, so an incompatible segment cannot poison or consume the prefix - the
testDuplicateInPartitionsAndIncrementalEvaluation fixture is exactly the right
test for this. Noting it here because it is the part of the diff most likely to
be misread as a bug.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java:
##########
@@ -105,6 +116,30 @@ public Set<String> prune(BrokerRequest brokerRequest,
Set<String> segments) {
return selectedSegments;
}
+ private Set<String> pruneWithPreparedPredicate(Expression filterExpression,
Set<String> segments) {
Review Comment:
Scope: MultiPartitionColumnsSegmentPruner has the same per-segment-per-value
hashing loop and gets none of this.
It hashes at lines 142 and 157 in exactly the same shape, so a
multi-column-partitioned table with 4096 segments and a 128-value IN still pays
the full 100k hashes this PR is about. Fine to leave for a follow-up, but worth
saying so in the description so it does not look like the problem is solved
table-wide. If a follow-up does extend it, the PreparedPredicate tree would
need to be per partition column, which is a good reason to check now that the
current shape can be generalised rather than refactored again.
##########
pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPrunerTest.java:
##########
@@ -0,0 +1,393 @@
+/**
+ * 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.utils.CommonConstants;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.expectThrows;
+
+
+/// Exercises query-local partition ID caching with real metadata
initialization and refresh, without ZooKeeper.
+public class SinglePartitionColumnSegmentPrunerTest {
+ private static final String COLUMN = "memberId";
+ private static final String TABLE = "testTable_OFFLINE";
+
+ @Test
+ public void testHashesOnceAcrossDistinctMetadataInstancesPerQuery() 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, "Computed hashes
must not survive a prune call");
+ }
+
+ @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(), 4,
+ "Only compatible default functions reuse IDs; configured functions
never compare configuration contents");
+ 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(), 3,
+ "Different partition counts must not reuse partition IDs");
+ }
+
+ @Test
+ public void testDuplicateInPartitionsAndIncrementalEvaluation() 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);
+ // The configured segment must not consume or extend the prefix cached by
the first segment.
+ assertEquals(pruner.prune(request(predicate("IN", "1", "9", "17", "2")),
records.keySet()),
+ expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 8);
+ // 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
+ public void testLargeConfigurationsAndUnrelatedPredicates() 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);
+ SinglePartitionColumnSegmentPruner pruner = pruner(Map.of("first", first,
"second", second));
+ assertEquals(pruner.prune(request(predicate("EQUALS", "first")),
Set.of("first", "second")), Set.of("first"));
+ assertEquals(pruner.prune(request(function("EQUALS",
RequestUtils.getIdentifierExpression("other"),
+ RequestUtils.getLiteralExpression("value"))), Set.of("first",
"second")), Set.of("first", "second"));
+ pruner.refreshSegment("first", metadata("first", "BoundedColumnValue", 3,
Set.of(2), config));
+ assertEquals(pruner.prune(request(predicate("EQUALS", "first")),
Set.of("first", "second")), Set.of());
+ }
+
+ @Test
+ public void testConfigurationHashCollisionsDoNotReusePartitionIds() throws
Exception {
Review Comment:
This test is vestigial - it guards against a hash collision in a
config-hashing scheme the current implementation no longer has.
With the interface-based canReusePartitionIds there is no
PartitionFunctionKey and configs are never hashed or compared, so 'Fixture must
exercise a configuration hash collision' asserts a property of an earlier
design. It also only has 2 candidate segments, which is below
MIN_SEGMENTS_FOR_PREPARATION, so it exercises the legacy loop rather than the
new code. Harmless, but the name and the assertion message will mislead the
next reader. Either drop it or re-point it at whatever the current design
actually risks (two configured BoundedColumnValue functions with equal
numPartitions and normalizer).
##########
pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPrunerTest.java:
##########
@@ -0,0 +1,393 @@
+/**
+ * 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.utils.CommonConstants;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.expectThrows;
+
+
+/// Exercises query-local partition ID caching with real metadata
initialization and refresh, without ZooKeeper.
+public class SinglePartitionColumnSegmentPrunerTest {
+ private static final String COLUMN = "memberId";
+ private static final String TABLE = "testTable_OFFLINE";
+
+ @Test
+ public void testHashesOnceAcrossDistinctMetadataInstancesPerQuery() 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, "Computed hashes
must not survive a prune call");
+ }
+
+ @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(), 4,
+ "Only compatible default functions reuse IDs; configured functions
never compare configuration contents");
+ 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(), 3,
+ "Different partition counts must not reuse partition IDs");
+ }
+
+ @Test
+ public void testDuplicateInPartitionsAndIncrementalEvaluation() 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);
+ // The configured segment must not consume or extend the prefix cached by
the first segment.
+ assertEquals(pruner.prune(request(predicate("IN", "1", "9", "17", "2")),
records.keySet()),
+ expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 8);
+ // 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
+ public void testLargeConfigurationsAndUnrelatedPredicates() 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);
+ SinglePartitionColumnSegmentPruner pruner = pruner(Map.of("first", first,
"second", second));
+ assertEquals(pruner.prune(request(predicate("EQUALS", "first")),
Set.of("first", "second")), Set.of("first"));
+ assertEquals(pruner.prune(request(function("EQUALS",
RequestUtils.getIdentifierExpression("other"),
+ RequestUtils.getLiteralExpression("value"))), Set.of("first",
"second")), Set.of("first", "second"));
+ pruner.refreshSegment("first", metadata("first", "BoundedColumnValue", 3,
Set.of(2), config));
+ assertEquals(pruner.prune(request(predicate("EQUALS", "first")),
Set.of("first", "second")), Set.of());
+ }
+
+ @Test
+ public void testConfigurationHashCollisionsDoNotReusePartitionIds() throws
Exception {
+ Map<String, String> firstConfig = Map.of("columnValues", "Aa|BB",
"columnValuesDelimiter", "|");
+ Map<String, String> secondConfig = Map.of("columnValues", "BB|Aa",
"columnValuesDelimiter", "|");
+ assertEquals(firstConfig.hashCode(), secondConfig.hashCode(),
+ "Fixture must exercise a configuration hash collision");
+ ZNRecord first = metadata("first", "BoundedColumnValue", 3, Set.of(1),
firstConfig);
+ ZNRecord second = metadata("second", "BoundedColumnValue", 3, Set.of(2),
secondConfig);
+ Map<String, ZNRecord> records = Map.of("first", first, "second", second);
+ assertEquals(pruner(records).prune(request(predicate("EQUALS", "Aa")),
records.keySet()), records.keySet());
+ }
+
+ @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());
+ }
+
+ @DataProvider
+ public Object[][] candidateCounts() {
+ return new Object[][]{{1}, {2}, {256}};
+ }
+
+ @Test(dataProvider = "candidateCounts")
+ public void testAndOrUnsupportedPredicatesAndLazyInValues(int numSegments)
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);
+ 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
+ public void testUnknownMetadataIsConservativeAndDoesNotEvaluateFilter()
throws Exception {
+ SinglePartitionColumnSegmentPruner pruner = new
SinglePartitionColumnSegmentPruner(TABLE, COLUMN);
+ 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 = Set.of("missing", "empty", "invalid",
"not-initialized");
+ 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) 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);
+ 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 testConcurrentQueriesKeepSeparateCachedValues() 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 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;
+ }
+
+ /// Stateless partition function with observable hash calls.
+ public static class CountingPartitionFunction implements PartitionFunction {
+ private static final long serialVersionUID = 1L;
+ private static final AtomicInteger CALLS = new AtomicInteger();
Review Comment:
The exact-call-count assertions on this static counter are safe today, but
only because of the repo's surefire settings - worth a comment in the test so
it is not lost.
Three tests reset CALLS and assert exact values. The root pom sets
<parallel>false</parallel> and reuseForks=false, so methods run sequentially
and this class gets its own JVM. If the module ever enabled TestNG method-level
parallelism these would interleave and flake. Also note
CountingPartitionFunction is public and in org.apache.pinot, so
PartitionFunctionFactory's classpath scan registers 'PrunerCounting' globally
for every test in this fork - intentional and needed here, just a global side
effect to be aware of.
##########
pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/partition/PartitionFunction.java:
##########
@@ -63,6 +63,26 @@ default Map<String, String> getFunctionConfig() {
return null;
}
+ /// Returns whether the exposed function configuration is null or empty.
This does not imply that other settings,
+ /// such as the partition id normalizer, have their default values.
+ @JsonIgnore
+ default boolean hasEmptyConfig() {
Review Comment:
hasEmptyConfig() becomes permanent public SPI surface, but it exists only to
support the default canReusePartitionIds. A private interface method would do
the same job without the commitment.
Java 9+ private interface methods are callable on any receiver from inside
the interface body, so `private boolean hasEmptyConfig()` still lets the
default impl call other.hasEmptyConfig(). SPI method names are forever in
Pinot; I would rather not add one whose only caller is two lines below unless
overriders are expected to reuse it. If it is meant as a public helper for
overriders, say so in the javadoc. The @JsonIgnore is correct and necessary, by
the way - Jackson would otherwise expose it as an 'emptyConfig' property on
serialised partition metadata.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java:
##########
@@ -105,6 +116,30 @@ public Set<String> prune(BrokerRequest brokerRequest,
Set<String> segments) {
return selectedSegments;
}
+ private Set<String> pruneWithPreparedPredicate(Expression filterExpression,
Set<String> segments) {
+ Set<String> selectedSegments = new HashSet<>();
+ PreparedPredicate predicate = null;
+ PartitionFunction cachedFunction = 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 (predicate == null) {
+ predicate = new PreparedPredicate(filterExpression);
+ cachedFunction = function;
+ }
+ // Reuse the filter structure even when the functions cannot share
partition ids.
+ boolean reusePartitionIds =
cachedFunction.canReusePartitionIds(function);
Review Comment:
The pivot is whichever segment HashSet iteration happens to visit first, and
for any configured partition function the first call is a self-comparison that
returns false - so the optimization silently disables itself for the whole
query.
cachedFunction is pinned to the first valid segment, then compared against
itself on that same iteration. With the default canReusePartitionIds,
f.canReusePartitionIds(f) is FALSE whenever f has a non-empty functionConfig,
because hasEmptyConfig() fails on both sides. So the default impl is not
reflexive, which is surprising for a relation named 'canReuse...', and it means
a configured table gets zero id reuse - matching the BOUNDED_100KB rows in the
description. Two follow-ons: (a) worth a `function == cachedFunction` identity
short-circuit before the virtual call, both for speed and so self-comparison is
trivially true; (b) one outlier segment arriving first (a stale numPartitions
after a table-config change, say) disables reuse for the other N-1 segments. If
that matters, re-pivot when the pivot yields no reuse for the first few
segments, or key a small map by function instead of one pivot.
##########
pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPruner.java:
##########
@@ -94,6 +98,13 @@ public Set<String> prune(BrokerRequest brokerRequest,
Set<String> segments) {
if (filterExpression == null) {
return segments;
}
+ int numSegments = segments.size();
+ if (numSegments == 0) {
Review Comment:
This branch is unreachable in production and it changes what the method
returns: the caller's set instead of a fresh one.
BaseBrokerRoutingManager.selectThenPrune() guards the whole pruner chain
with `if (!selectedSegments.isEmpty())`, so prune() is never called with an
empty set on the query path. The old code already never evaluated the filter
for an empty input (the loop body just did not run), so the branch buys nothing
behaviourally. What it does change is aliasing: getPrunedSegments() documents
that 'the pruners return a new set rather than editing the one they are
handed'. The filterExpression == null branch above already breaks that, so this
is not new, but I would drop the branch rather than add a second exception to a
convention a neighbouring comment relies on.
##########
pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SinglePartitionColumnSegmentPrunerTest.java:
##########
@@ -0,0 +1,393 @@
+/**
+ * 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.utils.CommonConstants;
+import org.testng.annotations.DataProvider;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotEquals;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.expectThrows;
+
+
+/// Exercises query-local partition ID caching with real metadata
initialization and refresh, without ZooKeeper.
+public class SinglePartitionColumnSegmentPrunerTest {
+ private static final String COLUMN = "memberId";
+ private static final String TABLE = "testTable_OFFLINE";
+
+ @Test
+ public void testHashesOnceAcrossDistinctMetadataInstancesPerQuery() 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, "Computed hashes
must not survive a prune call");
+ }
+
+ @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(), 4,
+ "Only compatible default functions reuse IDs; configured functions
never compare configuration contents");
+ 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(), 3,
+ "Different partition counts must not reuse partition IDs");
+ }
+
+ @Test
+ public void testDuplicateInPartitionsAndIncrementalEvaluation() 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);
+ // The configured segment must not consume or extend the prefix cached by
the first segment.
+ assertEquals(pruner.prune(request(predicate("IN", "1", "9", "17", "2")),
records.keySet()),
+ expected);
+ assertEquals(CountingPartitionFunction.CALLS.get(), 8);
+ // 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
+ public void testLargeConfigurationsAndUnrelatedPredicates() 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);
+ SinglePartitionColumnSegmentPruner pruner = pruner(Map.of("first", first,
"second", second));
+ assertEquals(pruner.prune(request(predicate("EQUALS", "first")),
Set.of("first", "second")), Set.of("first"));
+ assertEquals(pruner.prune(request(function("EQUALS",
RequestUtils.getIdentifierExpression("other"),
+ RequestUtils.getLiteralExpression("value"))), Set.of("first",
"second")), Set.of("first", "second"));
+ pruner.refreshSegment("first", metadata("first", "BoundedColumnValue", 3,
Set.of(2), config));
+ assertEquals(pruner.prune(request(predicate("EQUALS", "first")),
Set.of("first", "second")), Set.of());
+ }
+
+ @Test
+ public void testConfigurationHashCollisionsDoNotReusePartitionIds() throws
Exception {
+ Map<String, String> firstConfig = Map.of("columnValues", "Aa|BB",
"columnValuesDelimiter", "|");
+ Map<String, String> secondConfig = Map.of("columnValues", "BB|Aa",
"columnValuesDelimiter", "|");
+ assertEquals(firstConfig.hashCode(), secondConfig.hashCode(),
+ "Fixture must exercise a configuration hash collision");
+ ZNRecord first = metadata("first", "BoundedColumnValue", 3, Set.of(1),
firstConfig);
+ ZNRecord second = metadata("second", "BoundedColumnValue", 3, Set.of(2),
secondConfig);
+ Map<String, ZNRecord> records = Map.of("first", first, "second", second);
+ assertEquals(pruner(records).prune(request(predicate("EQUALS", "Aa")),
records.keySet()), records.keySet());
+ }
+
+ @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());
+ }
+
+ @DataProvider
+ public Object[][] candidateCounts() {
+ return new Object[][]{{1}, {2}, {256}};
+ }
+
+ @Test(dataProvider = "candidateCounts")
+ public void testAndOrUnsupportedPredicatesAndLazyInValues(int numSegments)
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);
+ 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
+ public void testUnknownMetadataIsConservativeAndDoesNotEvaluateFilter()
throws Exception {
Review Comment:
Test gap: the all-invalid-metadata case only runs on the legacy path (4
candidates), yet it is the one case where the prepared path never builds the
predicate at all.
In pruneWithPreparedPredicate the PreparedPredicate is constructed only when
the first VALID partition info is found, so when every segment is null/INVALID
the filter is never decoded and FilterKind.valueOf never throws. That is the
same conservative behaviour as the original short-circuit, but it is a distinct
code path and this test does not reach it. Padding the fixture past 256
candidates (or adding it to the candidateCounts data provider) would cover it.
Same thought for testLargeConfigurationsAndUnrelatedPredicates, which has 2
segments and therefore never touches the prepared evaluator it is named after.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
