This is an automated email from the ASF dual-hosted git repository.
Jackie-Jiang 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 afbf3049019 Cache segment to tier map at rebalance start (#19055)
afbf3049019 is described below
commit afbf3049019b7cc337287d6647a7f32902b5e250
Author: Jhow <[email protected]>
AuthorDate: Thu Jul 23 13:49:02 2026 -0700
Cache segment to tier map at rebalance start (#19055)
---
.../assignment/segment/BaseSegmentAssignment.java | 25 ++++++-
.../assignment/segment/SegmentAssignmentUtils.java | 82 +++++++++++++++-------
...NonReplicaGroupTieredSegmentAssignmentTest.java | 64 +++++++++++++++++
...NonReplicaGroupTieredSegmentAssignmentTest.java | 8 +++
.../StrictRealtimeSegmentAssignmentTest.java | 14 +++-
5 files changed, 164 insertions(+), 29 deletions(-)
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java
index 7a6ad7511d6..16ed00f9dc8 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/BaseSegmentAssignment.java
@@ -71,6 +71,20 @@ public abstract class BaseSegmentAssignment implements
SegmentAssignment {
protected TableConfig _tableConfig;
protected ControllerMetrics _controllerMetrics;
+ // Cache of segment name to the name of the tier it is eligible for (null
value means not eligible for any tier), to
+ // avoid reading each segment's ZK metadata on every rebalanceTiers()
invocation.
+ // NOTE:
+ // 1. This cache is used for table rebalance only. During rebalance,
rebalanceTable() (and thus rebalanceTiers()) can
+ // be invoked multiple times when the ideal state changes during the
rebalance process. The eligible tier of a
+ // segment is derived from its ZK metadata, which is not expected to
change during a rebalance, so it is resolved
+ // once (with a single bulk ZK read) and reused across invocations. A
fresh SegmentAssignment instance is created
+ // per rebalance (see SegmentAssignmentFactory), so the cache is scoped
to a single rebalance.
+ // 2. Segments that are added after the initial bulk read (e.g. realtime
segments that commit during a long
+ // rebalance) will be absent from the cache; they are resolved on demand
and merged in, so they are still routed
+ // to the correct tier.
+ @Nullable
+ private Map<String, String> _segmentToTierName;
+
@Override
public void init(HelixManager helixManager, TableConfig tableConfig,
@Nullable ControllerMetrics controllerMetrics) {
_helixManager = helixManager;
@@ -102,10 +116,17 @@ public abstract class BaseSegmentAssignment implements
SegmentAssignment {
_logger.info("Rebalancing tiers: {} for table: {} with bootstrap: {}",
tierInstancePartitionsMap.keySet(),
_tableNameWithType, bootstrap);
+ // Resolve the eligible tier of each segment once (single bulk ZK read)
and reuse it across the multiple
+ // rebalanceTiers() invocations of this rebalance, instead of reading each
segment's ZK metadata every time.
+ // Any new segments come later would not be in _segmentToTierName map, and
categorized in the null (default) tier.
+ if (_segmentToTierName == null) {
+ _segmentToTierName =
+ SegmentAssignmentUtils.getSegmentToTierNameMap(_helixManager,
_tableNameWithType, sortedTiers);
+ }
+
// Get tier to segment assignment map i.e. current assignments split by
tiers they are eligible for
SegmentAssignmentUtils.TierSegmentAssignment tierSegmentAssignment =
- new SegmentAssignmentUtils.TierSegmentAssignment(_helixManager,
_tableNameWithType, sortedTiers,
- currentAssignment);
+ new SegmentAssignmentUtils.TierSegmentAssignment(sortedTiers,
currentAssignment, _segmentToTierName);
Map<String, Map<String, Map<String, String>>>
tierNameToSegmentAssignmentMap =
tierSegmentAssignment.getTierNameToSegmentAssignmentMap();
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java
index dfc9eadb11b..500af00d0d4 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SegmentAssignmentUtils.java
@@ -452,6 +452,48 @@ public class SegmentAssignmentUtils {
}
}
+ /**
+ * Resolves, for each segment of the table, the name of the first tier (in
{@code sortedTiers} order) that the segment
+ * is eligible for. The value is the tier name, or {@code null} if the
segment (including COMMITTING segments) is not
+ * eligible for any tier. This performs a single bulk ZK read for all
segment ZK metadata, so it is meant to be
+ * computed once per rebalance and reused across the multiple {@code
rebalanceTable()} invocations of that rebalance
+ * rather than reading each segment's ZK metadata every time. Segments that
appear only after this map is computed
+ * (e.g. segments that commit during a long rebalance) are absent from the
map and treated as not eligible for any
+ * tier.
+ */
+ static Map<String, String> getSegmentToTierNameMap(HelixManager
helixManager, String tableNameWithType,
+ List<Tier> sortedTiers) {
+ ZkHelixPropertyStore<ZNRecord> propertyStore =
helixManager.getHelixPropertyStore();
+ List<SegmentZKMetadata> segmentsZKMetadata =
ZKMetadataProvider.getSegmentsZKMetadata(propertyStore,
+ tableNameWithType);
+ Map<String, String> segmentToTierName = new HashMap<>();
+ for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) {
+ segmentToTierName.put(segmentZKMetadata.getSegmentName(),
getTierName(tableNameWithType, sortedTiers,
+ segmentZKMetadata));
+ }
+ return segmentToTierName;
+ }
+
+ /**
+ * Returns the name of the first tier (in {@code sortedTiers} order) the
given segment is eligible for, or
+ * {@code null} if the segment is COMMITTING or not eligible for any tier.
+ */
+ @Nullable
+ private static String getTierName(String tableNameWithType, List<Tier>
sortedTiers,
+ SegmentZKMetadata segmentZKMetadata) {
+ // Skip COMMITTING segments
+ if (segmentZKMetadata.getStatus() == Status.COMMITTING) {
+ return null;
+ }
+ // Find an eligible tier for the segment, from the ordered list of tiers
+ for (Tier tier : sortedTiers) {
+ if (tier.getSegmentSelector().selectSegment(tableNameWithType,
segmentZKMetadata)) {
+ return tier.getName();
+ }
+ }
+ return null;
+ }
+
/**
* Takes a segment assignment and splits them up based on which tiers the
segments are eligible for. Only considers
* ONLINE segments.
@@ -463,44 +505,32 @@ public class SegmentAssignmentUtils {
private final Map<String, Map<String, String>> _nonTierSegmentAssignment =
new TreeMap<>();
/**
- * Creates a TierSegmentAssignment from the given segmentAssignment
- * @param tableNameWithType table to which the segment assignment belongs
+ * Creates a TierSegmentAssignment from the given segmentAssignment.
* @param sortedTiers list of tiers, pre-sorted as per desired order by
caller
* @param segmentAssignment segment assignment of the table
+ * @param segmentToTierName map from segment name to the name of the tier
it is eligible for (see
+ * {@link #getSegmentToTierNameMap}); segments
absent from this map are not eligible for
+ * any tier
*/
- TierSegmentAssignment(HelixManager helixManager, String tableNameWithType,
List<Tier> sortedTiers,
- Map<String, Map<String, String>> segmentAssignment) {
+ TierSegmentAssignment(List<Tier> sortedTiers, Map<String, Map<String,
String>> segmentAssignment,
+ Map<String, String> segmentToTierName) {
// initialize tier to segmentAssignment map
sortedTiers.forEach(t ->
_tierNameToSegmentAssignmentMap.put(t.getName(), new TreeMap<>()));
// iterate over all segments
- // TODO: Reduce ZK access
- ZkHelixPropertyStore<ZNRecord> propertyStore =
helixManager.getHelixPropertyStore();
for (Map.Entry<String, Map<String, String>> entry :
segmentAssignment.entrySet()) {
String segmentName = entry.getKey();
Map<String, String> instanceStateMap = entry.getValue();
- boolean selected = false;
-
- // only consider ONLINE segments for tiers
- if (instanceStateMap.containsValue(SegmentStateModel.ONLINE)) {
- // find an eligible tier for the segment, from the ordered list of
tiers
- SegmentZKMetadata segmentZKMetadata =
- ZKMetadataProvider.getSegmentZKMetadata(propertyStore,
tableNameWithType, segmentName);
- // Skip COMMITTING segments
- if (segmentZKMetadata != null && segmentZKMetadata.getStatus() !=
Status.COMMITTING) {
- for (Tier tier : sortedTiers) {
- if (tier.getSegmentSelector().selectSegment(tableNameWithType,
segmentZKMetadata)) {
-
_tierNameToSegmentAssignmentMap.get(tier.getName()).put(segmentName,
instanceStateMap);
- selected = true;
- break;
- }
- }
- }
- }
- // if segment not eligible for any tier, put in ordinary segments map
- if (!selected) {
+ // Only consider ONLINE segments for tiers. The eligible tier for each
segment is resolved once via
+ // segmentToTierName, avoiding a per-segment ZK read on every
rebalanceTable() invocation.
+ String tierName =
instanceStateMap.containsValue(SegmentStateModel.ONLINE)
+ ? segmentToTierName.get(segmentName) : null;
+ if (tierName != null) {
+ _tierNameToSegmentAssignmentMap.get(tierName).put(segmentName,
instanceStateMap);
+ } else {
+ // if segment not eligible for any tier, put in ordinary segments map
_nonTierSegmentAssignment.put(segmentName, instanceStateMap);
}
}
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineNonReplicaGroupTieredSegmentAssignmentTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineNonReplicaGroupTieredSegmentAssignmentTest.java
index 83ef265676f..18d8f1939b6 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineNonReplicaGroupTieredSegmentAssignmentTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/OfflineNonReplicaGroupTieredSegmentAssignmentTest.java
@@ -19,6 +19,7 @@
package org.apache.pinot.controller.helix.core.assignment.segment;
import com.google.common.collect.Lists;
+import java.util.ArrayList;
import java.util.Arrays;
import java.util.HashMap;
import java.util.List;
@@ -47,9 +48,12 @@ import org.testng.Assert;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
+import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
+import static org.mockito.Mockito.times;
+import static org.mockito.Mockito.verify;
import static org.mockito.Mockito.when;
import static org.testng.Assert.assertEquals;
import static org.testng.Assert.assertTrue;
@@ -116,6 +120,13 @@ public class
OfflineNonReplicaGroupTieredSegmentAssignmentTest {
String segmentName = path.substring(path.lastIndexOf('/') + 1);
return new ZNRecord(segmentName);
});
+ // Bulk read of all segment ZK metadata, used to resolve the eligible tier
of each segment once per rebalance
+ List<ZNRecord> segmentZNRecords = new ArrayList<>(NUM_SEGMENTS);
+ for (String segmentName : SEGMENTS) {
+ segmentZNRecords.add(new ZNRecord(segmentName));
+ }
+ when(propertyStore.getChildren(anyString(), eq(null),
eq(AccessOption.PERSISTENT), anyInt(), anyInt())).thenReturn(
+ segmentZNRecords);
//noinspection unchecked
when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore);
@@ -269,6 +280,59 @@ public class
OfflineNonReplicaGroupTieredSegmentAssignmentTest {
}
}
+ /**
+ * The segment to tier resolution is computed once (single bulk ZK read) at
the start of the rebalance and reused
+ * across the multiple rebalanceTable() invocations of that rebalance. A
segment that appears in the current
+ * assignment only in a later invocation (e.g. after the ideal state changes
during a long rebalance) is absent from
+ * the resolution map, so it is kept on the default instances rather than
triggering another ZK read.
+ */
+ @Test
+ public void testSegmentAddedAfterInitialBulkReadStaysOnDefaultTier() {
+ HelixManager helixManager = mock(HelixManager.class);
+ //noinspection rawtypes
+ ZkHelixPropertyStore propertyStore = mock(ZkHelixPropertyStore.class);
+ // The initial bulk read only returns the segments that existed when the
rebalance started (segment_0 .. 99)
+ List<ZNRecord> segmentZNRecords = new ArrayList<>(NUM_SEGMENTS);
+ for (String segmentName : SEGMENTS) {
+ segmentZNRecords.add(new ZNRecord(segmentName));
+ }
+ //noinspection unchecked
+ when(propertyStore.getChildren(anyString(), eq(null),
eq(AccessOption.PERSISTENT), anyInt(), anyInt())).thenReturn(
+ segmentZNRecords);
+ //noinspection unchecked
+ when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore);
+ SegmentAssignment segmentAssignment =
SegmentAssignmentFactory.getSegmentAssignment(helixManager,
+ new
TableConfigBuilder(TableType.OFFLINE).setTableName(RAW_TABLE_NAME).setNumReplicas(NUM_REPLICAS).build(),
+ null);
+
+ Map<String, Map<String, String>> currentAssignment = new TreeMap<>();
+ for (String segmentName : SEGMENTS) {
+ currentAssignment.put(segmentName,
+ SegmentAssignmentUtils.getInstanceStateMap(INSTANCES.subList(0,
NUM_REPLICAS), SegmentStateModel.ONLINE));
+ }
+
+ // First rebalance resolves the segment to tier map from the single bulk
read
+ segmentAssignment.rebalanceTable(currentAssignment,
_instancePartitionsMap, _sortedTiers,
+ _tierInstancePartitionsMap, new RebalanceConfig());
+ //noinspection unchecked
+ verify(propertyStore, times(1)).getChildren(anyString(), eq(null),
eq(AccessOption.PERSISTENT), anyInt(), anyInt());
+
+ // A new segment that would be eligible for tierC (segId >= 120) appears
in the assignment after the initial bulk
+ // read
+ String newSegment = SEGMENT_NAME_PREFIX + 200;
+ currentAssignment.put(newSegment,
+ SegmentAssignmentUtils.getInstanceStateMap(INSTANCES.subList(0,
NUM_REPLICAS), SegmentStateModel.ONLINE));
+ Map<String, Map<String, String>> newAssignment =
+ segmentAssignment.rebalanceTable(currentAssignment,
_instancePartitionsMap, _sortedTiers,
+ _tierInstancePartitionsMap, new RebalanceConfig());
+
+ // No second bulk read happens, and the new segment is kept on the default
instances (not relocated to tierC)
+ //noinspection unchecked
+ verify(propertyStore, times(1)).getChildren(anyString(), eq(null),
eq(AccessOption.PERSISTENT), anyInt(), anyInt());
+ Assert.assertEquals(newAssignment.get(newSegment).size(), NUM_REPLICAS);
+
Assert.assertTrue(INSTANCES.containsAll(newAssignment.get(newSegment).keySet()));
+ }
+
/**
* Selects segment_50 to segment_69 i.e. 20 segments
*/
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeNonReplicaGroupTieredSegmentAssignmentTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeNonReplicaGroupTieredSegmentAssignmentTest.java
index b2316ba1547..2f5cbf29a15 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeNonReplicaGroupTieredSegmentAssignmentTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeNonReplicaGroupTieredSegmentAssignmentTest.java
@@ -47,6 +47,7 @@ import org.apache.pinot.spi.utils.builder.TableConfigBuilder;
import org.testng.annotations.BeforeClass;
import org.testng.annotations.Test;
+import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
@@ -126,6 +127,13 @@ public class
RealtimeNonReplicaGroupTieredSegmentAssignmentTest {
String segmentName = path.substring(path.lastIndexOf('/') + 1);
return new ZNRecord(segmentName);
});
+ // Bulk read of all segment ZK metadata, used to resolve the eligible tier
of each segment once per rebalance
+ List<ZNRecord> segmentZNRecords = new ArrayList<>(NUM_SEGMENTS);
+ for (String segmentName : _segments) {
+ segmentZNRecords.add(new ZNRecord(segmentName));
+ }
+ when(propertyStore.getChildren(anyString(), eq(null),
eq(AccessOption.PERSISTENT), anyInt(), anyInt())).thenReturn(
+ segmentZNRecords);
//noinspection unchecked
when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore);
diff --git
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java
index e308b43bb8d..80968cf5a96 100644
---
a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java
+++
b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/StrictRealtimeSegmentAssignmentTest.java
@@ -54,6 +54,7 @@ import org.testng.annotations.BeforeClass;
import org.testng.annotations.DataProvider;
import org.testng.annotations.Test;
+import static org.mockito.ArgumentMatchers.anyInt;
import static org.mockito.ArgumentMatchers.anyString;
import static org.mockito.ArgumentMatchers.eq;
import static org.mockito.Mockito.mock;
@@ -78,7 +79,7 @@ public class StrictRealtimeSegmentAssignmentTest {
private static final String CONSUMING_INSTANCE_PARTITIONS_NAME =
InstancePartitionsType.CONSUMING.getInstancePartitionsName(RAW_TABLE_NAME);
- private List<String> _segments;
+ private static List<String> _segments;
private Map<InstancePartitionsType, InstancePartitions>
_instancePartitionsMap;
private InstancePartitions _newConsumingInstancePartitions;
@@ -493,6 +494,17 @@ public class StrictRealtimeSegmentAssignmentTest {
}
return record;
});
+ // Bulk read of all segment ZK metadata, used to resolve the eligible tier
of each segment once per rebalance
+ List<ZNRecord> segmentZNRecords = new ArrayList<>(_segments.size());
+ for (String segmentName : _segments) {
+ ZNRecord record = new ZNRecord(segmentName);
+ if (tieredSegments.contains(segmentName)) {
+ record.setSimpleField(Segment.TIER, "coldTier");
+ }
+ segmentZNRecords.add(record);
+ }
+ when(propertyStore.getChildren(anyString(), eq(null),
eq(AccessOption.PERSISTENT), anyInt(), anyInt())).thenReturn(
+ segmentZNRecords);
when(helixManager.getHelixPropertyStore()).thenReturn(propertyStore);
return helixManager;
}
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
