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 0ace824e6e4 [bugfix] Fixing instance partition assignments for
multi-stream realtime tables (#18433)
0ace824e6e4 is described below
commit 0ace824e6e40953e0fb02b73052fb5c261fbbb18
Author: Shaurya Chaturvedi <[email protected]>
AuthorDate: Mon Jun 15 11:45:17 2026 -0700
[bugfix] Fixing instance partition assignments for multi-stream realtime
tables (#18433)
---
.../segment/RealtimeSegmentAssignment.java | 6 +++
.../ReplicaGroupSegmentAssignmentStrategy.java | 21 +++++++---
.../pinot/spi/utils/IngestionConfigUtils.java | 12 +++++-
.../pinot/spi/utils/IngestionConfigUtilsTest.java | 48 ++++++++++++++++++++++
4 files changed, 79 insertions(+), 8 deletions(-)
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java
index 62821c252fd..d4b991e54b7 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java
@@ -38,6 +38,7 @@ import
org.apache.pinot.controller.helix.core.assignment.segment.strategy.Segmen
import
org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager;
import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType;
import
org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
/**
@@ -115,6 +116,11 @@ public class RealtimeSegmentAssignment extends
BaseSegmentAssignment {
}
protected List<String> assignConsumingSegment(int segmentPartitionId,
InstancePartitions instancePartitions) {
+ // For multi-stream tables, Pinot partition IDs are encoded as
(streamIndex * 10000 + streamPartitionId).
+ // Extract the stream-level partition id before computing the instance
index to avoid incorrect slot mapping.
+ segmentPartitionId =
+
IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_tableConfig,
segmentPartitionId);
+
int numReplicaGroups = instancePartitions.getNumReplicaGroups();
int numPartitions = instancePartitions.getNumPartitions();
diff --git
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java
index 15047747a73..70d897b187f 100644
---
a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java
+++
b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java
@@ -33,6 +33,7 @@ import
org.apache.pinot.controller.helix.core.assignment.segment.SegmentAssignme
import org.apache.pinot.segment.local.utils.TableConfigUtils;
import org.apache.pinot.spi.config.table.SegmentsValidationAndRetentionConfig;
import org.apache.pinot.spi.config.table.TableConfig;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
import org.slf4j.Logger;
import org.slf4j.LoggerFactory;
@@ -41,6 +42,7 @@ public class ReplicaGroupSegmentAssignmentStrategy implements
SegmentAssignmentS
private static final Logger LOGGER =
LoggerFactory.getLogger(ReplicaGroupSegmentAssignmentStrategy.class);
protected HelixManager _helixManager;
+ protected TableConfig _tableConfig;
protected String _tableName;
protected String _partitionColumn;
protected int _replication;
@@ -48,6 +50,7 @@ public class ReplicaGroupSegmentAssignmentStrategy implements
SegmentAssignmentS
@Override
public void init(HelixManager helixManager, TableConfig tableConfig) {
_helixManager = helixManager;
+ _tableConfig = tableConfig;
_tableName = tableConfig.getTableName();
SegmentsValidationAndRetentionConfig validationAndRetentionConfig =
tableConfig.getValidationConfig();
Preconditions.checkState(validationAndRetentionConfig != null,
"segmentsConfig is null");
@@ -77,9 +80,7 @@ public class ReplicaGroupSegmentAssignmentStrategy implements
SegmentAssignmentS
if (numPartitions == 1) {
partitionId = 0;
} else {
- partitionId =
- SegmentUtils.getSegmentPartitionIdOrDefault(segmentName, _tableName,
_helixManager, _partitionColumn)
- % numPartitions;
+ partitionId = getPartitionIdFromSegmentName(segmentName, numPartitions);
}
return
SegmentAssignmentUtils.assignSegmentWithReplicaGroup(currentAssignment,
instancePartitions, partitionId);
}
@@ -106,9 +107,7 @@ public class ReplicaGroupSegmentAssignmentStrategy
implements SegmentAssignmentS
} else {
Map<Integer, List<String>> instancePartitionIdToSegmentsMap =
Maps.newHashMapWithExpectedSize(numPartitions);
for (String segmentName : currentAssignment.keySet()) {
- int instancePartitionId =
- SegmentUtils.getSegmentPartitionIdOrDefault(segmentName,
_tableName, _helixManager, _partitionColumn)
- % numPartitions;
+ int instancePartitionId = getPartitionIdFromSegmentName(segmentName,
numPartitions);
instancePartitionIdToSegmentsMap.computeIfAbsent(instancePartitionId,
k -> new ArrayList<>()).add(segmentName);
}
@@ -125,6 +124,16 @@ public class ReplicaGroupSegmentAssignmentStrategy
implements SegmentAssignmentS
}
}
+ private int getPartitionIdFromSegmentName(String segmentName, int
numPartitions) {
+ int rawPartitionId =
+ SegmentUtils.getSegmentPartitionIdOrDefault(segmentName, _tableName,
_helixManager, _partitionColumn);
+ // For multi-stream realtime tables, translate the Pinot partition ID
(which encodes stream index and stream
+ // partition as streamIndex * 10000 + streamPartitionId) to the stream
partition id before computing the slot.
+ // getStreamPartitionIdFromPinotPartitionId is a no-op for offline tables
and single-stream realtime tables.
+ rawPartitionId =
IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_tableConfig,
rawPartitionId);
+ return rawPartitionId % numPartitions;
+ }
+
/**
* Helper method to check whether the number of replica-groups matches the
table replication for replica-group based
* instance partitions. Log a warning if they do not match and use the one
inside the instance partitions. The
diff --git
a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java
index 40f647c429c..b0731ad3831 100644
---
a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java
+++
b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java
@@ -102,9 +102,16 @@ public final class IngestionConfigUtils {
return new StreamConfig(tableConfig.getTableName(),
getFirstStreamConfigMap(tableConfig));
}
- /// Returns `true` if the table contains multiple streams.
+ /// Returns {@code true} if the table has multiple streams configured.
+ /// Safe to call for any table type: returns {@code false} for OFFLINE
tables and REALTIME tables that lack stream
+ /// configs.
public static boolean hasMultipleStreams(TableConfig tableConfig) {
- return getStreamConfigMaps(tableConfig).size() > 1;
+ IngestionConfig ingestionConfig = tableConfig.getIngestionConfig();
+ if (ingestionConfig == null || ingestionConfig.getStreamIngestionConfig()
== null) {
+ return false;
+ }
+ List<Map<String, String>> streamConfigMaps =
ingestionConfig.getStreamIngestionConfig().getStreamConfigMaps();
+ return streamConfigMaps != null && streamConfigMaps.size() > 1;
}
/**
@@ -117,6 +124,7 @@ public final class IngestionConfigUtils {
}
/// Returns the stream partition id from the Pinot segment partition id.
+ /// Returns {@code partitionId} unchanged for tables without multiple
streams.
public static int getStreamPartitionIdFromPinotPartitionId(TableConfig
tableConfig, int partitionId) {
return hasMultipleStreams(tableConfig) ?
getStreamPartitionIdFromPinotPartitionId(partitionId) : partitionId;
}
diff --git
a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/IngestionConfigUtilsTest.java
b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/IngestionConfigUtilsTest.java
index 3390126b2ff..b55fb22e71c 100644
---
a/pinot-spi/src/test/java/org/apache/pinot/spi/utils/IngestionConfigUtilsTest.java
+++
b/pinot-spi/src/test/java/org/apache/pinot/spi/utils/IngestionConfigUtilsTest.java
@@ -160,6 +160,54 @@ public class IngestionConfigUtilsTest {
Assert.assertEquals(2,
IngestionConfigUtils.getConfigMapWithPrefix(testMap, "k1.").size());
}
+ @Test
+ public void testHasMultipleStreams() {
+ // OFFLINE table — must not throw, must return false
+ TableConfig offlineTable = new
TableConfigBuilder(TableType.OFFLINE).setTableName("myTable").build();
+ Assert.assertFalse(IngestionConfigUtils.hasMultipleStreams(offlineTable));
+
+ // REALTIME, no ingestionConfig
+ TableConfig realtimeTable =
+ new
TableConfigBuilder(TableType.REALTIME).setTableName("myTable").setTimeColumnName("timeColumn").build();
+ Assert.assertFalse(IngestionConfigUtils.hasMultipleStreams(realtimeTable));
+
+ // REALTIME, single stream
+ Map<String, String> streamConfigMap1 =
Collections.singletonMap("streamType", "kafka");
+ IngestionConfig ingestionConfig = new IngestionConfig();
+ ingestionConfig.setStreamIngestionConfig(new
StreamIngestionConfig(Collections.singletonList(streamConfigMap1)));
+ realtimeTable.setIngestionConfig(ingestionConfig);
+ Assert.assertFalse(IngestionConfigUtils.hasMultipleStreams(realtimeTable));
+
+ // REALTIME, two streams
+ Map<String, String> streamConfigMap2 =
Collections.singletonMap("streamType", "kinesis");
+ ingestionConfig.setStreamIngestionConfig(new
StreamIngestionConfig(List.of(streamConfigMap1, streamConfigMap2)));
+ Assert.assertTrue(IngestionConfigUtils.hasMultipleStreams(realtimeTable));
+ }
+
+ @Test
+ public void testGetStreamPartitionIdFromPinotPartitionId() {
+ // OFFLINE table — must return partitionId unchanged
+ TableConfig offlineTable = new
TableConfigBuilder(TableType.OFFLINE).setTableName("myTable").build();
+
Assert.assertEquals(IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(offlineTable,
42), 42);
+
+ // REALTIME, single stream — must return partitionId unchanged
+ Map<String, String> streamConfigMap =
Collections.singletonMap("streamType", "kafka");
+ IngestionConfig ingestionConfig = new IngestionConfig();
+ ingestionConfig.setStreamIngestionConfig(new
StreamIngestionConfig(Collections.singletonList(streamConfigMap)));
+ TableConfig realtimeTable =
+ new
TableConfigBuilder(TableType.REALTIME).setTableName("myTable").setTimeColumnName("timeColumn").build();
+ realtimeTable.setIngestionConfig(ingestionConfig);
+
Assert.assertEquals(IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(realtimeTable,
42), 42);
+
+ // REALTIME, multi-stream — encoded partition id 10003 (stream 1,
partition 3) must decode to 3
+ Map<String, String> streamConfigMap2 =
Collections.singletonMap("streamType", "kinesis");
+ ingestionConfig.setStreamIngestionConfig(
+ new StreamIngestionConfig(List.of(streamConfigMap, streamConfigMap2)));
+ int pinotPartitionId =
IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(3, 1);
+
Assert.assertEquals(IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(realtimeTable,
pinotPartitionId),
+ 3);
+ }
+
@Test
public void testGetStreamConfigIndexToStreamPartitions() {
Set<Integer> pinotPartitionIds = new HashSet<>(2);
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
