This is an automated email from the ASF dual-hosted git repository.
yashmayya 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 a080910ca48 Build TimeSegmentPruner's interval tree from flat arrays
and skip no-op rebuilds (#19643)
a080910ca48 is described below
commit a080910ca48f11cbf8b72a67f17b7fa10c070ed6
Author: Yash Mayya <[email protected]>
AuthorDate: Wed Sep 30 13:30:51 2026 -0700
Build TimeSegmentPruner's interval tree from flat arrays and skip no-op
rebuilds (#19643)
---
.../routing/segmentpruner/TimeSegmentPruner.java | 38 +++-
.../segmentpruner/interval/IntervalTree.java | 240 +++++++++------------
.../routing/intervalst/IntervalTreeTest.java | 60 ++++++
.../routing/segmentpruner/SegmentPrunerTest.java | 120 +++++++++++
.../pinot/perf/BenchmarkTimeSegmentPruner.java | 163 ++++++++++++++
5 files changed, 481 insertions(+), 140 deletions(-)
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java
index 19ef2ee0cc4..7c2833b12d6 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/TimeSegmentPruner.java
@@ -18,6 +18,7 @@
*/
package org.apache.pinot.broker.routing.segmentpruner;
+import com.google.common.annotations.VisibleForTesting;
import com.google.common.base.Preconditions;
import java.util.ArrayList;
import java.util.HashMap;
@@ -61,8 +62,16 @@ public class TimeSegmentPruner implements SegmentPruner {
private final String _timeColumn;
private final DateTimeFormatSpec _timeFormatSpec;
+ /// Rebuilt under the lock on this pruner and read by [#prune(BrokerRequest,
Set)] without locking. Each tree is
+ /// immutable, so a query always sees a complete tree.
private volatile IntervalTree<String> _intervalTree;
+ /// Guarded by the lock on this pruner, except in [#init], which runs before
the pruner is used.
private final Map<String, Interval> _intervalMap = new HashMap<>();
+ /// True while [#_intervalTree] may lag [#_intervalMap]: set before the map
is updated and cleared once the tree
+ /// matches it again. If an update throws part way (e.g. with an
OutOfMemoryError), the next update rebuilds the tree
+ /// even when it changes nothing, instead of skipping the rebuild and
leaving segments missing from the tree. Guarded
+ /// by the lock on this pruner.
+ private boolean _intervalTreeStale;
public TimeSegmentPruner(TableConfig tableConfig, DateTimeFieldSpec
timeFieldSpec) {
_tableNameWithType = tableConfig.getTableName();
@@ -107,20 +116,41 @@ public class TimeSegmentPruner implements SegmentPruner {
Set<String> onlineSegments, List<String> pulledSegments, List<ZNRecord>
znRecords) {
// NOTE: We don't update all the segment ZK metadata for every external
view change, but only the new added/removed
// ones. The refreshed segment ZK metadata change won't be picked up.
+ boolean treeWasStale = _intervalTreeStale;
+ _intervalTreeStale = true;
+ int numSegmentsBefore = _intervalMap.size();
for (int idx = 0; idx < pulledSegments.size(); idx++) {
String segment = pulledSegments.get(idx);
ZNRecord zNrecord = znRecords.get(idx);
_intervalMap.computeIfAbsent(segment, k ->
extractIntervalFromSegmentZKMetaZNRecord(k, zNrecord));
}
- _intervalMap.keySet().retainAll(onlineSegments);
- _intervalTree = new IntervalTree<>(_intervalMap);
+ // Only insertions can change the size because computeIfAbsent never
replaces an existing interval. An external
+ // view change that adds and removes no segment (e.g. a replica changing
state) leaves the tree correct as is.
+ boolean segmentsChanged = _intervalMap.size() != numSegmentsBefore;
+ segmentsChanged |= _intervalMap.keySet().retainAll(onlineSegments);
+ if (segmentsChanged || treeWasStale) {
+ _intervalTree = new IntervalTree<>(_intervalMap);
+ }
+ _intervalTreeStale = false;
}
@Override
public synchronized void refreshSegment(String segment, @Nullable ZNRecord
znRecord) {
Interval interval = extractIntervalFromSegmentZKMetaZNRecord(segment,
znRecord);
- _intervalMap.put(segment, interval);
- _intervalTree = new IntervalTree<>(_intervalMap);
+ boolean treeWasStale = _intervalTreeStale;
+ _intervalTreeStale = true;
+ Interval previousInterval = _intervalMap.put(segment, interval);
+ // A segment is commonly refreshed onto the time interval it already has
(e.g. an OFFLINE segment replaced with a
+ // new build of the same time range), which leaves the tree correct as is
+ if (treeWasStale || !interval.equals(previousInterval)) {
+ _intervalTree = new IntervalTree<>(_intervalMap);
+ }
+ _intervalTreeStale = false;
+ }
+
+ @VisibleForTesting
+ IntervalTree<String> getIntervalTree() {
+ return _intervalTree;
}
/// NOTE: Pruning is done by searching \_intervalTree based on request time
interval and check if the results
diff --git
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/interval/IntervalTree.java
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/interval/IntervalTree.java
index 59d6dcc95dd..f8d90684b47 100644
---
a/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/interval/IntervalTree.java
+++
b/pinot-broker/src/main/java/org/apache/pinot/broker/routing/segmentpruner/interval/IntervalTree.java
@@ -18,179 +18,147 @@
*/
package org.apache.pinot.broker.routing.segmentpruner.interval;
-import com.google.common.base.Preconditions;
+import it.unimi.dsi.fastutil.ints.IntArrays;
import java.util.ArrayList;
-import java.util.Collections;
-import java.util.HashMap;
-import java.util.LinkedList;
import java.util.List;
import java.util.Map;
-import org.apache.pinot.spi.utils.Pairs;
-
-
-/// The `IntervalTree` class represents read-only balanced binary interval
tree map (from intervals to values)
+import javax.annotation.Nullable;
+
+
+/// The `IntervalTree` class represents a read-only balanced binary interval
tree map (from intervals to values).
+///
+/// The distinct intervals are held sorted in [#_intervals], and the balanced
tree over them is implicit: the root of
+/// the index range `[start, end)` is `start + (end - start) / 2`, its left
child is the root of `[start, mid)` and its
+/// right child the root of `[mid + 1, end)`. A typical balanced tree:
+/// ```
+/// [10, 20]
+/// / \
+/// [8, 15] [12, 20]
+/// / /
+/// [5, 10] [10, 30]
+/// ```
+/// is held as the sorted array `{ [5, 10], [8, 15], [10, 20], [10, 30], [12,
20] }`.
+///
+/// The tree is held implicitly and its payload in flat arrays because a tree
is rebuilt whenever the segments of a
+/// table change, including on every REALTIME segment commit, and on a table
with hundreds of thousands of segments an
+/// object per interval dominates broker allocation.
+///
+/// Instances are immutable and safe to publish to readers through a volatile
field.
public class IntervalTree<VALUE> {
-
- // List representation of BST with root at index 0. For node with index x,
it's left child index is (2x+1), right
- // child index is (2x+2)
- private final List<Node> _nodes;
+ /// The distinct intervals, sorted ascending by `(min, max)`.
+ private final Interval[] _intervals;
+ /// Max interval end of the subtree rooted at each node, used to skip
subtrees that cannot match.
+ private final long[] _subtreeMaxs;
+ /// The values mapped to `_intervals[i]` are `_values[_valueOffsets[i]]`
(inclusive) through
+ /// `_values[_valueOffsets[i + 1]]` (exclusive). Holds one extra trailing
entry so the last node has an end offset.
+ private final int[] _valueOffsets;
+ /// Every key of the map the tree was built from, grouped by interval. Typed
as `Object[]` because a `VALUE[]`
+ /// cannot be created from a type parameter; only values put here by the
constructor are ever read back out.
+ private final Object[] _values;
public IntervalTree(Map<VALUE, Interval> valueToIntervalMap) {
- Map<Interval, List<VALUE>> intervalToValuesMap = new HashMap<>();
+ int numValues = valueToIntervalMap.size();
+ Object[] values = new Object[numValues];
+ Interval[] intervals = new Interval[numValues];
+ int index = 0;
for (Map.Entry<VALUE, Interval> entry : valueToIntervalMap.entrySet()) {
- intervalToValuesMap.putIfAbsent(entry.getValue(), new ArrayList<>());
- intervalToValuesMap.get(entry.getValue()).add(entry.getKey());
+ values[index] = entry.getKey();
+ intervals[index] = entry.getValue();
+ index++;
}
- List<Node<VALUE>> sortedNodes = new ArrayList<>();
- for (Map.Entry<Interval, List<VALUE>> entry :
intervalToValuesMap.entrySet()) {
- sortedNodes.add(new Node(entry.getKey(), entry.getValue()));
+ // Sort a permutation of the indexes rather than the entries themselves,
so that no object is allocated per entry.
+ // This sorts every value instead of grouping by interval first and
sorting only the distinct ones. Grouping first
+ // needs a hash structure over the values, which costs more than it saves
unless most values share an interval:
+ // measured over 250k values, grouping first is 35 ms / 12 MB against 20
ms / 3 MB when every interval is distinct,
+ // and 3 ms / 1 MB against 12 ms / 3 MB when 64 values share each interval.
+ int[] sortedIndexes = new int[numValues];
+ for (int i = 0; i < numValues; i++) {
+ sortedIndexes[i] = i;
}
- Collections.sort(sortedNodes);
- _nodes = buildIntervalTree(sortedNodes);
- buildAuxiliaryInfo();
- }
-
- /// Build interval bst by bfs, the root for each subtree will be the one
with median interval.
- /// A typical balanced tree:
- /// \[10, 20\]
- /// / \
- /// \[8, 15\] \[12, 20\]
- /// / /
- /// \[5, 10\] \[10, 30\]
- /// is represented as { \[10, 20\], \[8, 15\], \[12, 20\], \[5, 10\], null,
\[10, 30\] }
- private List<Node> buildIntervalTree(List<Node<VALUE>> sortedNodes) {
- List<Node> resNodes = new ArrayList<>();
- LinkedList<Pairs.IntPair> indexQueue = new LinkedList<>();
- indexQueue.add(new Pairs.IntPair(0, sortedNodes.size()));
- int count = 0;
- while (count < sortedNodes.size()) {
- Pairs.IntPair indexPair = indexQueue.pollFirst();
- int start = indexPair.getLeft();
- int end = indexPair.getRight();
-
- if (start < end) {
- int mid = start + (end - start) / 2;
- resNodes.add(sortedNodes.get(mid));
- count++;
- indexQueue.add(new Pairs.IntPair(start, mid));
- indexQueue.add(new Pairs.IntPair(mid + 1, end));
- } else {
- resNodes.add(null);
+ IntArrays.quickSort(sortedIndexes, (i, j) ->
intervals[i].compareTo(intervals[j]));
+
+ // Equal intervals are adjacent after the sort, so a single scan gives the
number of distinct intervals
+ int numIntervals = 0;
+ Interval previousInterval = null;
+ for (int i = 0; i < numValues; i++) {
+ Interval interval = intervals[sortedIndexes[i]];
+ if (!interval.equals(previousInterval)) {
+ numIntervals++;
+ previousInterval = interval;
}
}
- return resNodes;
- }
-
- private void buildAuxiliaryInfo() {
- // Build max info for the interval tree by dfs
- buildAuxiliaryInfo(0);
- }
- private void buildAuxiliaryInfo(int nodeIndex) {
- if (!hasNode(nodeIndex)) {
- return;
+ _intervals = new Interval[numIntervals];
+ _subtreeMaxs = new long[numIntervals];
+ _valueOffsets = new int[numIntervals + 1];
+ _values = new Object[numValues];
+ int numIntervalsAdded = 0;
+ previousInterval = null;
+ for (int i = 0; i < numValues; i++) {
+ int sortedIndex = sortedIndexes[i];
+ Interval interval = intervals[sortedIndex];
+ if (!interval.equals(previousInterval)) {
+ _intervals[numIntervalsAdded] = interval;
+ _valueOffsets[numIntervalsAdded] = i;
+ numIntervalsAdded++;
+ previousInterval = interval;
+ }
+ _values[i] = values[sortedIndex];
}
+ _valueOffsets[numIntervals] = numValues;
- int leftChildIndex = getLeftChildIndex(nodeIndex);
- int rightChildIndex = getRightChildIndex(nodeIndex);
-
- buildAuxiliaryInfo(leftChildIndex);
- buildAuxiliaryInfo(rightChildIndex);
-
- long max = _nodes.get(nodeIndex)._interval._max;
- max = Math.max(getMax(rightChildIndex), Math.max(max,
getMax(leftChildIndex)));
- _nodes.get(nodeIndex)._max = max;
+ buildSubtreeMaxs(0, numIntervals);
}
- private int getLeftChildIndex(int nodeIndex) {
- return nodeIndex * 2 + 1;
- }
-
- private int getRightChildIndex(int nodeIndex) {
- return nodeIndex * 2 + 2;
- }
-
- private long getMax(int index) {
- if (!hasNode(index)) {
+ /// Fills [#_subtreeMaxs] for the subtree covering `[start, end)` and
returns its max interval end.
+ private long buildSubtreeMaxs(int start, int end) {
+ if (start >= end) {
return Long.MIN_VALUE;
}
- return _nodes.get(index)._max;
+ int mid = start + (end - start) / 2;
+ long max = Math.max(_intervals[mid]._max, Math.max(buildSubtreeMaxs(start,
mid), buildSubtreeMaxs(mid + 1, end)));
+ _subtreeMaxs[mid] = max;
+ return max;
}
/// Find all values whose intervals intersect with the input interval.
///
/// @param searchInterval search interval
/// @return list of all qualified values.
- public List<VALUE> searchAll(Interval searchInterval) {
- List<VALUE> list = new ArrayList<>();
- if (searchInterval == null) {
- return list;
+ public List<VALUE> searchAll(@Nullable Interval searchInterval) {
+ List<VALUE> values = new ArrayList<>();
+ if (searchInterval != null) {
+ searchAll(0, _intervals.length, searchInterval, values);
}
- searchAll(0, searchInterval, list);
- return list;
+ return values;
}
- private void searchAll(int nodeIndex, Interval searchInterval, List<VALUE>
list) {
- if (!hasNode(nodeIndex)) {
+ private void searchAll(int start, int end, Interval searchInterval,
List<VALUE> values) {
+ if (start >= end) {
return;
}
+ int mid = start + (end - start) / 2;
- int leftChildIndex = getLeftChildIndex(nodeIndex);
- int rightChildIndex = getRightChildIndex(nodeIndex);
-
- if (hasNode(leftChildIndex) && getMax(leftChildIndex) >=
searchInterval._min) {
- searchAll(leftChildIndex, searchInterval, list);
+ // Search the left subtree unless every interval in it ends before the
search interval starts
+ if (start < mid && _subtreeMaxs[start + (mid - start) / 2] >=
searchInterval._min) {
+ searchAll(start, mid, searchInterval, values);
}
- Node<VALUE> node = _nodes.get(nodeIndex);
- Interval interval = node._interval;
+ Interval interval = _intervals[mid];
if (searchInterval.intersects(interval)) {
- list.addAll(node._values);
- }
-
- if (interval._min <= searchInterval._max) {
- searchAll(rightChildIndex, searchInterval, list);
- }
- }
-
- private boolean hasNode(int nodeIndex) {
- return nodeIndex < _nodes.size() && _nodes.get(nodeIndex) != null;
- }
-
- private class Node<VALUE> implements Comparable<Node> {
- private final Interval _interval;
- private final List<VALUE> _values;
- private long _max; // max interval right end of subtree rooted at this node
-
- Node(Interval interval, List<VALUE> values) {
- _interval = interval;
- _values = values;
- }
-
- @Override
- public int compareTo(Node o) {
- Preconditions.checkNotNull(o, "Compare to invalid node: null");
- return _interval.compareTo(o._interval);
- }
-
- @Override
- public boolean equals(Object o) {
- if (this == o) {
- return true;
+ int valueEndOffset = _valueOffsets[mid + 1];
+ for (int i = _valueOffsets[mid]; i < valueEndOffset; i++) {
+ @SuppressWarnings("unchecked")
+ VALUE value = (VALUE) _values[i];
+ values.add(value);
}
- if (o == null || getClass() != o.getClass()) {
- return false;
- }
-
- Node<?> node = (Node<?>) o;
-
- return _interval.equals(node._interval);
}
- @Override
- public int hashCode() {
- return _interval.hashCode();
+ // Intervals are sorted by start, so nothing in the right subtree can
match once this one starts after the search
+ // interval ends
+ if (interval._min <= searchInterval._max) {
+ searchAll(mid + 1, end, searchInterval, values);
}
}
}
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/intervalst/IntervalTreeTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/intervalst/IntervalTreeTest.java
index 30c082d8965..10e7cdbda6c 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/intervalst/IntervalTreeTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/intervalst/IntervalTreeTest.java
@@ -23,6 +23,8 @@ import java.util.HashMap;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
+import java.util.Random;
+import java.util.Set;
import org.apache.pinot.broker.routing.segmentpruner.interval.Interval;
import org.apache.pinot.broker.routing.segmentpruner.interval.IntervalTree;
import org.testng.Assert;
@@ -86,4 +88,62 @@ public class IntervalTreeTest {
Assert.assertEquals(new HashSet<>(intervalTree.searchAll(new Interval(20,
30))),
new HashSet<>(Arrays.asList(name8, name9, name10, name11, name12)));
}
+
+ @Test
+ public void testEmptyTree() {
+ IntervalTree<String> intervalTree = new IntervalTree<>(Map.of());
+ Assert.assertEquals(intervalTree.searchAll(new Interval(0,
Long.MAX_VALUE)), List.of());
+ Assert.assertEquals(intervalTree.searchAll(null), List.of());
+ }
+
+ @Test
+ public void testNullSearchInterval() {
+ IntervalTree<String> intervalTree = new IntervalTree<>(Map.of("value", new
Interval(5, 10)));
+ Assert.assertEquals(intervalTree.searchAll(null), List.of());
+ }
+
+ @Test
+ public void testAllValuesOnOneInterval() {
+ Map<String, Interval> valueToIntervalMap = new HashMap<>();
+ for (int i = 0; i < 100; i++) {
+ valueToIntervalMap.put("value" + i, new Interval(5, 10));
+ }
+ IntervalTree<String> intervalTree = new IntervalTree<>(valueToIntervalMap);
+ Assert.assertEquals(new HashSet<>(intervalTree.searchAll(new Interval(10,
20))), valueToIntervalMap.keySet());
+ Assert.assertEquals(intervalTree.searchAll(new Interval(11, 20)),
List.of());
+ }
+
+ /// Cross-checks the tree against a brute force scan over randomly generated
intervals. The tight value ranges make
+ /// many values share an interval, and the `[0, Long.MAX_VALUE]` entry
mirrors the interval
+ /// [org.apache.pinot.broker.routing.segmentpruner.TimeSegmentPruner] gives
a segment with no time range.
+ @Test
+ public void testAgainstBruteForceSearch() {
+ Random random = new Random(0);
+ for (int numValues : new int[]{1, 2, 3, 5, 13, 64, 500}) {
+ for (int valueRange : new int[]{1, 2, 5, 40, 1000}) {
+ Map<String, Interval> valueToIntervalMap = new HashMap<>();
+ for (int i = 0; i < numValues; i++) {
+ long min = random.nextInt(valueRange);
+ valueToIntervalMap.put("value" + i, new Interval(min, min +
random.nextInt(valueRange)));
+ }
+ valueToIntervalMap.put("noTimeRange", new Interval(0, Long.MAX_VALUE));
+
+ IntervalTree<String> intervalTree = new
IntervalTree<>(valueToIntervalMap);
+ for (int i = 0; i < 100; i++) {
+ long min = random.nextInt(valueRange + 2);
+ Interval searchInterval = new Interval(min, min +
random.nextInt(valueRange + 2));
+ Set<String> expected = new HashSet<>();
+ for (Map.Entry<String, Interval> entry :
valueToIntervalMap.entrySet()) {
+ if (entry.getValue().intersects(searchInterval)) {
+ expected.add(entry.getKey());
+ }
+ }
+ List<String> actual = intervalTree.searchAll(searchInterval);
+ Assert.assertEquals(new HashSet<>(actual), expected, "Mismatch for
search interval: " + searchInterval);
+ Assert.assertEquals(actual.size(), expected.size(), "Duplicate
values for search interval: "
+ + searchInterval);
+ }
+ }
+ }
+ }
}
diff --git
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerTest.java
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerTest.java
index 5e48a981ccc..283f148ba4e 100644
---
a/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerTest.java
+++
b/pinot-broker/src/test/java/org/apache/pinot/broker/routing/segmentpruner/SegmentPrunerTest.java
@@ -35,6 +35,7 @@ import org.apache.helix.zookeeper.datamodel.ZNRecord;
import org.apache.helix.zookeeper.datamodel.serializer.ZNRecordSerializer;
import org.apache.helix.zookeeper.impl.client.ZkClient;
import
org.apache.pinot.broker.routing.segmentmetadata.SegmentZkMetadataFetcher;
+import org.apache.pinot.broker.routing.segmentpruner.interval.IntervalTree;
import org.apache.pinot.common.metadata.ZKMetadataProvider;
import org.apache.pinot.common.metadata.segment.SegmentPartitionMetadata;
import org.apache.pinot.common.metadata.segment.SegmentZKMetadata;
@@ -61,9 +62,14 @@ import org.testng.annotations.AfterClass;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
+import static org.mockito.ArgumentMatchers.anyLong;
+import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.Mockito.mock;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertNotSame;
+import static org.testng.Assert.assertSame;
+import static org.testng.Assert.assertThrows;
import static org.testng.Assert.assertTrue;
@@ -501,6 +507,112 @@ public class SegmentPrunerTest extends ControllerTest {
assertEquals(segmentPruner.prune(brokerRequest9, input), Set.of()); //
Query with invalid range
}
+ /// The interval tree is rebuilt over every segment of the table, so it must
be rebuilt only when the set of segments
+ /// or a segment's time range actually changes. A refresh onto the same time
range (an OFFLINE segment replaced by a
+ /// new build of the same range) and an assignment change that adds and
removes no segment must keep the tree as is.
+ @Test
+ public void testTimeSegmentPrunerRebuildsOnlyOnChange() {
+ BrokerRequest selectAll =
CalciteSqlCompiler.compileToBrokerRequest(QUERY_1);
+ BrokerRequest between20And30 =
CalciteSqlCompiler.compileToBrokerRequest(TIME_QUERY_2);
+
+ TableConfig tableConfig =
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setTimeColumnName(TIME_COLUMN).build();
+ DateTimeFieldSpec timeFieldSpec = new DateTimeFieldSpec(TIME_COLUMN,
DataType.INT, "EPOCH|DAYS", "1:DAYS");
+ TimeSegmentPruner segmentPruner = new TimeSegmentPruner(tableConfig,
timeFieldSpec);
+
+ String segment0 = "segment0";
+ String segment1 = "segment1";
+ Set<String> input = Set.of(segment0, segment1);
+ segmentPruner.init(null, null, List.of(segment0, segment1),
+ List.of(createTimeRangeZNRecord(segment0, 10, 15),
createTimeRangeZNRecord(segment1, 25, 35)));
+ assertEquals(segmentPruner.prune(selectAll, input), input);
+ assertEquals(segmentPruner.prune(between20And30, input), Set.of(segment1));
+
+ // Refreshing onto the same time range keeps the tree
+ IntervalTree<String> intervalTree = segmentPruner.getIntervalTree();
+ segmentPruner.refreshSegment(segment0, createTimeRangeZNRecord(segment0,
10, 15));
+ segmentPruner.refreshSegment(segment1, createTimeRangeZNRecord(segment1,
25, 35));
+ assertSame(segmentPruner.getIntervalTree(), intervalTree);
+ assertEquals(segmentPruner.prune(between20And30, input), Set.of(segment1));
+
+ // An assignment change that adds and removes no segment keeps the tree
+ segmentPruner.onAssignmentChange(null, null, input, List.of(), List.of());
+ assertSame(segmentPruner.getIntervalTree(), intervalTree);
+
+ // Refreshing onto a new time range rebuilds the tree
+ segmentPruner.refreshSegment(segment0, createTimeRangeZNRecord(segment0,
20, 22));
+ segmentPruner.refreshSegment(segment1, createTimeRangeZNRecord(segment1,
40, 50));
+ assertNotSame(segmentPruner.getIntervalTree(), intervalTree);
+ assertEquals(segmentPruner.prune(between20And30, input), Set.of(segment0));
+
+ // Adding a segment rebuilds the tree
+ intervalTree = segmentPruner.getIntervalTree();
+ String segment2 = "segment2";
+ Set<String> inputWithSegment2 = Set.of(segment0, segment1, segment2);
+ segmentPruner.onAssignmentChange(null, null, inputWithSegment2,
List.of(segment2),
+ List.of(createTimeRangeZNRecord(segment2, 28, 29)));
+ assertNotSame(segmentPruner.getIntervalTree(), intervalTree);
+ assertEquals(segmentPruner.prune(between20And30, inputWithSegment2),
Set.of(segment0, segment2));
+
+ // Removing a segment rebuilds the tree, so the removed segment is no
longer selected
+ intervalTree = segmentPruner.getIntervalTree();
+ segmentPruner.onAssignmentChange(null, null, input, List.of(), List.of());
+ assertNotSame(segmentPruner.getIntervalTree(), intervalTree);
+ assertEquals(segmentPruner.prune(between20And30, inputWithSegment2),
Set.of(segment0));
+
+ // A segment whose ZK metadata went missing falls back to the full time
range and is not pruned
+ segmentPruner.refreshSegment(segment1, null);
+ assertEquals(segmentPruner.prune(between20And30, input), Set.of(segment0,
segment1));
+
+ // Replacing a segment adds one and removes one, so the segment count
stays the same, but the tree must be rebuilt
+ intervalTree = segmentPruner.getIntervalTree();
+ String segment3 = "segment3";
+ segmentPruner.onAssignmentChange(null, null, Set.of(segment0, segment3),
List.of(segment3),
+ List.of(createTimeRangeZNRecord(segment3, 24, 26)));
+ assertNotSame(segmentPruner.getIntervalTree(), intervalTree);
+ assertEquals(segmentPruner.prune(between20And30, Set.of(segment0,
segment1, segment3)), Set.of(segment0, segment3));
+ }
+
+ /// An update that throws part way (e.g. with an OutOfMemoryError while the
broker is under heap pressure) can leave
+ /// the interval map ahead of the tree. The next update must rebuild the
tree even when it changes nothing itself,
+ /// or the segments added before the failure stay missing from the tree and
are dropped from the routing.
+ @Test
+ public void testTimeSegmentPrunerRecoversFromFailedUpdate() {
+ BrokerRequest between20And30 =
CalciteSqlCompiler.compileToBrokerRequest(TIME_QUERY_2);
+
+ TableConfig tableConfig =
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setTimeColumnName(TIME_COLUMN).build();
+ DateTimeFieldSpec timeFieldSpec = new DateTimeFieldSpec(TIME_COLUMN,
DataType.INT, "EPOCH|DAYS", "1:DAYS");
+ TimeSegmentPruner segmentPruner = new TimeSegmentPruner(tableConfig,
timeFieldSpec);
+ String segment0 = "segment0";
+ segmentPruner.init(null, null, List.of(segment0),
List.of(createTimeRangeZNRecord(segment0, 20, 22)));
+ ZNRecord brokenZNRecord = mock(ZNRecord.class);
+ when(brokenZNRecord.getLongField(anyString(), anyLong())).thenThrow(new
IllegalStateException("broken"));
+ String brokenSegment = "brokenSegment";
+
+ // The update adds segment1, then fails on the broken segment before the
tree is rebuilt
+ String segment1 = "segment1";
+ assertThrows(IllegalStateException.class,
+ () -> segmentPruner.onAssignmentChange(null, null, Set.of(segment0,
segment1, brokenSegment),
+ List.of(segment1, brokenSegment),
List.of(createTimeRangeZNRecord(segment1, 25, 26), brokenZNRecord)));
+ Set<String> input = Set.of(segment0, segment1);
+ assertEquals(segmentPruner.prune(between20And30, input), Set.of(segment0));
+
+ // A refresh onto the same time range changes nothing, but must still
rebuild the stale tree
+ segmentPruner.refreshSegment(segment0, createTimeRangeZNRecord(segment0,
20, 22));
+ assertEquals(segmentPruner.prune(between20And30, input), Set.of(segment0,
segment1));
+
+ // The same failure again, this time repaired by an assignment change that
adds and removes no segment
+ String segment2 = "segment2";
+ assertThrows(IllegalStateException.class,
+ () -> segmentPruner.onAssignmentChange(null, null, Set.of(segment0,
segment1, segment2, brokenSegment),
+ List.of(segment2, brokenSegment),
List.of(createTimeRangeZNRecord(segment2, 27, 28), brokenZNRecord)));
+ Set<String> inputWithSegment2 = Set.of(segment0, segment1, segment2);
+ assertEquals(segmentPruner.prune(between20And30, inputWithSegment2),
Set.of(segment0, segment1));
+ segmentPruner.onAssignmentChange(null, null, inputWithSegment2, List.of(),
List.of());
+ assertEquals(segmentPruner.prune(between20And30, inputWithSegment2),
Set.of(segment0, segment1, segment2));
+ }
+
@Test
public void testTimeSegmentPrunerSimpleDateFormat() {
BrokerRequest brokerRequest1 =
CalciteSqlCompiler.compileToBrokerRequest(SDF_QUERY_1);
@@ -696,6 +808,14 @@ public class SegmentPrunerTest extends ControllerTest {
ZKMetadataProvider.setSegmentZKMetadata(_propertyStore, tableNameWithType,
segmentZKMetadata);
}
+ private static ZNRecord createTimeRangeZNRecord(String segment, long
startTime, long endTime) {
+ SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segment);
+ segmentZKMetadata.setStartTime(startTime);
+ segmentZKMetadata.setEndTime(endTime);
+ segmentZKMetadata.setTimeUnit(TimeUnit.DAYS);
+ return segmentZKMetadata.toZNRecord();
+ }
+
private void setSegmentZKTotalDocsMetadata(String tableNameWithType, String
segment, long totalDocs) {
SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segment);
segmentZKMetadata.setTotalDocs(totalDocs);
diff --git
a/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkTimeSegmentPruner.java
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkTimeSegmentPruner.java
new file mode 100644
index 00000000000..03097cd7bc4
--- /dev/null
+++
b/pinot-perf/src/main/java/org/apache/pinot/perf/BenchmarkTimeSegmentPruner.java
@@ -0,0 +1,163 @@
+/**
+ * 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.perf;
+
+import java.util.ArrayList;
+import java.util.HashSet;
+import java.util.List;
+import java.util.Set;
+import java.util.concurrent.TimeUnit;
+import org.apache.helix.zookeeper.datamodel.ZNRecord;
+import org.apache.pinot.broker.routing.segmentpruner.TimeSegmentPruner;
+import org.apache.pinot.common.request.BrokerRequest;
+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.FieldSpec.DataType;
+import org.apache.pinot.spi.utils.CommonConstants;
+import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
+import org.apache.pinot.sql.parsers.CalciteSqlCompiler;
+import org.openjdk.jmh.annotations.Benchmark;
+import org.openjdk.jmh.annotations.BenchmarkMode;
+import org.openjdk.jmh.annotations.Fork;
+import org.openjdk.jmh.annotations.Level;
+import org.openjdk.jmh.annotations.Measurement;
+import org.openjdk.jmh.annotations.Mode;
+import org.openjdk.jmh.annotations.OutputTimeUnit;
+import org.openjdk.jmh.annotations.Param;
+import org.openjdk.jmh.annotations.Scope;
+import org.openjdk.jmh.annotations.Setup;
+import org.openjdk.jmh.annotations.State;
+import org.openjdk.jmh.annotations.Warmup;
+import org.openjdk.jmh.infra.Blackhole;
+import org.openjdk.jmh.runner.Runner;
+import org.openjdk.jmh.runner.RunnerException;
+import org.openjdk.jmh.runner.options.OptionsBuilder;
+
+
+/// Benchmarks the broker-side segment metadata maintenance path of
[TimeSegmentPruner].
+///
+/// Every LLC REALTIME segment commit makes the controller broadcast a segment
refresh message to every broker, and each
+/// refresh that changes a segment's time range rebuilds the table's interval
tree. `refreshChangedInterval` is that
+/// cost: a committing REALTIME segment getting its time range.
`refreshUnchangedInterval` is an OFFLINE segment
+/// re-pushed with the same time range, which must skip the rebuild. `prune`
is the query path, which reads the tree
+/// but never builds it.
+///
+/// Run with `-prof gc` — allocation rate matters more than latency here,
because the garbage this path produces is
+/// promoted to old gen rather than dying in eden.
+@BenchmarkMode(Mode.AverageTime)
+@OutputTimeUnit(TimeUnit.MICROSECONDS)
+@Fork(1)
+@Warmup(iterations = 3, time = 3)
+@Measurement(iterations = 5, time = 3)
+@State(Scope.Benchmark)
+public class BenchmarkTimeSegmentPruner {
+ private static final String RAW_TABLE_NAME = "benchTable";
+ private static final String TIME_COLUMN = "tsMs";
+ private static final long BASE_TIME_MS = 1704067200000L; //
2024-01-01T00:00:00Z
+ private static final long SEGMENT_SPAN_MS = TimeUnit.HOURS.toMillis(1);
+
+ @Param({"10000", "100000", "250000"})
+ private int _numSegments;
+
+ /// Number of segments sharing one `[start, end]` interval. 1 models
REALTIME segments whose time range comes from
+ /// the data (effectively all distinct); 64 models one segment per stream
partition landing on the same coarse
+ /// time bucket.
+ @Param({"1", "64"})
+ private int _numSegmentsPerInterval;
+
+ private TimeSegmentPruner _segmentPruner;
+ private List<String> _segments;
+ private Set<String> _onlineSegments;
+ private BrokerRequest _brokerRequest;
+ private int _numIntervals;
+ private int _cursor;
+
+ @Setup(Level.Iteration)
+ public void setUp() {
+ TableConfig tableConfig = new
TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME)
+ .setTimeColumnName(TIME_COLUMN).build();
+ DateTimeFieldSpec timeFieldSpec =
+ new DateTimeFieldSpec(TIME_COLUMN, DataType.LONG,
"EPOCH|MILLISECONDS", "1:MILLISECONDS");
+ _segmentPruner = new TimeSegmentPruner(tableConfig, timeFieldSpec);
+
+ _segments = new ArrayList<>(_numSegments);
+ List<ZNRecord> znRecords = new ArrayList<>(_numSegments);
+ for (int i = 0; i < _numSegments; i++) {
+ String segment = RAW_TABLE_NAME + "__" + (i % 64) + "__" + (i / 64) +
"__20240101T0000Z";
+ long startTimeMs = BASE_TIME_MS + (long) (i / _numSegmentsPerInterval) *
SEGMENT_SPAN_MS;
+ _segments.add(segment);
+ znRecords.add(createZNRecord(segment, startTimeMs, startTimeMs +
SEGMENT_SPAN_MS - 1));
+ }
+ // NOTE: Ideal state and external view are not used by TimeSegmentPruner
+ _segmentPruner.init(null, null, _segments, znRecords);
+ _onlineSegments = new HashSet<>(_segments);
+ _numIntervals = (_numSegments + _numSegmentsPerInterval - 1) /
_numSegmentsPerInterval;
+
+ long queryStartMs = BASE_TIME_MS + (_numIntervals / 4L) * SEGMENT_SPAN_MS;
+ _brokerRequest = CalciteSqlCompiler.compileToBrokerRequest(
+ "SELECT * FROM " + RAW_TABLE_NAME + " WHERE " + TIME_COLUMN + "
BETWEEN " + queryStartMs + " AND " + (
+ queryStartMs + 24 * SEGMENT_SPAN_MS));
+ _cursor = 0;
+ }
+
+ private static ZNRecord createZNRecord(String segment, long startTimeMs,
long endTimeMs) {
+ ZNRecord znRecord = new ZNRecord(segment);
+ znRecord.setLongField(CommonConstants.Segment.START_TIME, startTimeMs);
+ znRecord.setLongField(CommonConstants.Segment.END_TIME, endTimeMs);
+ znRecord.setEnumField(CommonConstants.Segment.TIME_UNIT,
TimeUnit.MILLISECONDS);
+ return znRecord;
+ }
+
+ /// Refreshes a segment onto a time range it did not have before - the
REALTIME commit shape.
+ ///
+ /// The segment is moved one time bucket forward rather than past the end of
the table, so that the number of
+ /// distinct intervals and the selectivity of the benchmark query both stay
at the values `_numSegments` and
+ /// `_numSegmentsPerInterval` model, however long the benchmark runs.
+ @Benchmark
+ public void refreshChangedInterval() {
+ int index = _cursor++;
+ int segmentIndex = index % _numSegments;
+ String segment = _segments.get(segmentIndex);
+ int intervalIndex = (segmentIndex / _numSegmentsPerInterval + 1 + index /
_numSegments) % _numIntervals;
+ long startTimeMs = BASE_TIME_MS + (long) intervalIndex * SEGMENT_SPAN_MS;
+ _segmentPruner.refreshSegment(segment, createZNRecord(segment,
startTimeMs, startTimeMs + SEGMENT_SPAN_MS - 1));
+ }
+
+ /// Refreshes a segment onto the time range it already has - the OFFLINE
re-push shape.
+ @Benchmark
+ public void refreshUnchangedInterval() {
+ int index = _cursor++ % _numSegments;
+ String segment = _segments.get(index);
+ long startTimeMs = BASE_TIME_MS + (long) (index / _numSegmentsPerInterval)
* SEGMENT_SPAN_MS;
+ _segmentPruner.refreshSegment(segment, createZNRecord(segment,
startTimeMs, startTimeMs + SEGMENT_SPAN_MS - 1));
+ }
+
+ /// Query path only, to confirm pruning itself does not regress.
+ @Benchmark
+ public void prune(Blackhole blackhole) {
+ blackhole.consume(_segmentPruner.prune(_brokerRequest, _onlineSegments));
+ }
+
+ public static void main(String[] args)
+ throws RunnerException {
+ new Runner(new
OptionsBuilder().include(BenchmarkTimeSegmentPruner.class.getSimpleName()).addProfiler("gc").build())
+ .run();
+ }
+}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
