[bugfix] Fixing instance partition assignments for multi-stream realtime tables

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.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 @@ private List<String> assignConsumingSegment(String segmentName, InstancePartitio
   }
 
   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();
 

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/strategy/ReplicaGroupSegmentAssignmentStrategy.java

@@ -33,6 +33,7 @@
 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,13 +42,15 @@ 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;
 
   @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 List<String> assignSegment(String segmentName, Map<String, Map<String, St
     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 Map<String, Map<String, String>> reassignSegments(Map<String, Map<String,
     } 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 Map<String, Map<String, String>> reassignSegments(Map<String, Map<String,
     }
   }
 
+  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

pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java

@@ -117,8 +117,19 @@ public static int getPinotPartitionIdFromStreamPartitionId(int partitionId, int
   }
 
   /// Returns the stream partition id from the Pinot segment partition id.
+  /// Safe to call for any table type: returns `partitionId` unchanged for OFFLINE tables or REALTIME tables
+  /// that lack stream configs (treat as single-stream).
   public static int getStreamPartitionIdFromPinotPartitionId(TableConfig tableConfig, int partitionId) {
-    return hasMultipleStreams(tableConfig) ? getStreamPartitionIdFromPinotPartitionId(partitionId) : partitionId;
+    if (tableConfig.getTableType() != TableType.REALTIME) {
+      return partitionId;
+    }
+    IngestionConfig ingestionConfig = tableConfig.getIngestionConfig();
+    if (ingestionConfig == null || ingestionConfig.getStreamIngestionConfig() == null) {
+      return partitionId;
+    }
+    List<Map<String, String>> streamConfigMaps = ingestionConfig.getStreamIngestionConfig().getStreamConfigMaps();
+    return streamConfigMaps != null && streamConfigMaps.size() > 1
+        ? getStreamPartitionIdFromPinotPartitionId(partitionId) : partitionId;
   }
 
   /// Returns the stream partition id from the Pinot segment partition id.