From 554a5df037be262e7faaa7fdd7aca6464e0fa2b3 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Mon, 13 Jul 2026 15:00:04 -0700 Subject: [PATCH 01/12] Added topic name in LLCSegmentName --- .../pinot/common/utils/LLCSegmentName.java | 56 +++++++++++++-- .../common/utils/LLCSegmentNameTest.java | 68 +++++++++++++++++++ 2 files changed, 119 insertions(+), 5 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java index fdff73c0318d..daeb0ca1cdf0 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java @@ -34,6 +34,7 @@ public class LLCSegmentName implements Comparable { private static final DateTimeFormatter DATE_FORMATTER = DateTimeFormat.forPattern(DATE_FORMAT).withZoneUTC(); private final String _tableName; + private final int _topicId; private final int _partitionGroupId; private final int _sequenceNumber; private final String _creationTime; @@ -41,17 +42,26 @@ public class LLCSegmentName implements Comparable { public LLCSegmentName(String segmentName) { String[] parts = StringUtils.splitByWholeSeparator(segmentName, SEPARATOR); - Preconditions.checkArgument(parts.length == 4, "Invalid LLC segment name: %s", segmentName); + Preconditions.checkArgument(parts.length == 4 || parts.length == 5, "Invalid LLC segment name: %s", segmentName); _tableName = parts[0]; - _partitionGroupId = Integer.parseInt(parts[1]); - _sequenceNumber = Integer.parseInt(parts[2]); - _creationTime = parts[3]; + if (parts.length == 4) { + _topicId = 0; + _partitionGroupId = Integer.parseInt(parts[1]); + _sequenceNumber = Integer.parseInt(parts[2]); + _creationTime = parts[3]; + } else { + _topicId = Integer.parseInt(parts[1]); + _partitionGroupId = Integer.parseInt(parts[2]); + _sequenceNumber = Integer.parseInt(parts[3]); + _creationTime = parts[4]; + } _segmentName = segmentName; } public LLCSegmentName(String tableName, int partitionGroupId, int sequenceNumber, long msSinceEpoch) { Preconditions.checkArgument(!tableName.contains(SEPARATOR), "Illegal table name: %s", tableName); _tableName = tableName; + _topicId = 0; _partitionGroupId = partitionGroupId; _sequenceNumber = sequenceNumber; // ISO8601 date: 20160120T1234Z @@ -59,6 +69,17 @@ public LLCSegmentName(String tableName, int partitionGroupId, int sequenceNumber _segmentName = tableName + SEPARATOR + partitionGroupId + SEPARATOR + sequenceNumber + SEPARATOR + _creationTime; } + public LLCSegmentName(String tableName, int configId, int partitionGroupId, int sequenceNumber, long msSinceEpoch) { + Preconditions.checkArgument(!tableName.contains(SEPARATOR), "Illegal table name: %s", tableName); + _tableName = tableName; + _topicId = configId; + _partitionGroupId = partitionGroupId; + _sequenceNumber = sequenceNumber; + _creationTime = DATE_FORMATTER.print(msSinceEpoch); + _segmentName = tableName + SEPARATOR + configId + SEPARATOR + partitionGroupId + SEPARATOR + sequenceNumber + + SEPARATOR + _creationTime; + } + /** * Returns the {@link LLCSegmentName} for the given segment name, or {@code null} if the given segment name does not * represent an LLC segment. @@ -82,7 +103,26 @@ public static boolean isLLCSegment(String segmentName) { numSeparators++; index += 2; // SEPARATOR.length() } - return numSeparators == 3; + return numSeparators == 3 || numSeparators == 4; + } + + /** + * Returns whether the given segment name is a multi-topic LLC segment (5-part format where + * parts[1] and parts[3] are integers, distinguishing it from UploadedRealtimeSegmentName). + */ + public static boolean isMultiTopicLLCSegment(String segmentName) { + String[] parts = StringUtils.splitByWholeSeparator(segmentName, SEPARATOR); + if (parts.length != 5) { + return false; + } + try { + Integer.parseInt(parts[1]); // topicId + Integer.parseInt(parts[2]); // partitionGroupId + Integer.parseInt(parts[3]); // sequenceNumber + return true; + } catch (NumberFormatException e) { + return false; + } } @Deprecated @@ -101,6 +141,10 @@ public String getTableName() { return _tableName; } + public int getTopicId() { + return _topicId; + } + public int getPartitionGroupId() { return _partitionGroupId; } @@ -127,6 +171,8 @@ public String getSegmentName() { public int compareTo(LLCSegmentName other) { Preconditions.checkArgument(_tableName.equals(other._tableName), "Cannot compare segment names from different table: %s, %s", _segmentName, other.getSegmentName()); + Preconditions.checkArgument(_topicId == other._topicId, + "Cannot compare segment names from different topic ids: %s, %s", _segmentName, other.getSegmentName()); if (_partitionGroupId != other._partitionGroupId) { return Integer.compare(_partitionGroupId, other._partitionGroupId); } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/LLCSegmentNameTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/LLCSegmentNameTest.java index 04b2a9e2fbf2..3209523f34e1 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/LLCSegmentNameTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/LLCSegmentNameTest.java @@ -31,6 +31,8 @@ * Tests for the realtime segment name builder. */ public class LLCSegmentNameTest { + private static final String TABLE_NAME = "myTable"; + private static final long TIMESTAMP_MS = 1465508537069L; @Test public void testSegmentNameBuilder() { @@ -95,4 +97,70 @@ public void testLLCSegmentName() { Arrays.sort(testSorted); Assert.assertEquals(testSorted, new LLCSegmentName[]{segName5, segName1, segName6, segName3, segName4}); } + + @Test + public void testSegmentNameWithoutTopicIdDefaultsTopicIdToZero() { + LLCSegmentName segmentName = new LLCSegmentName("myTable__4__27__20160617T2150Z"); + assertEquals(segmentName.getTopicId(), 0); + } + + @Test + public void testSegmentNameWithTopicIdParsesTopicIdAndOtherFields() { + LLCSegmentName segmentName = new LLCSegmentName("myTable__2__4__27__20160617T2150Z"); + assertEquals(segmentName.getTopicId(), 2); + assertEquals(segmentName.getTableName(), "myTable"); + assertEquals(segmentName.getPartitionGroupId(), 4); + assertEquals(segmentName.getSequenceNumber(), 27); + assertEquals(segmentName.getCreationTime(), "20160617T2150Z"); + assertTrue(LLCSegmentName.isLLCSegment(segmentName.getSegmentName())); + } + + @Test + public void testConstructorWithTopicIdBuildsFivePartSegmentName() { + LLCSegmentName segmentName = new LLCSegmentName(TABLE_NAME, 2, 0, 1, TIMESTAMP_MS); + assertEquals(segmentName.getSegmentName(), "myTable__2__0__1__20160609T2142Z"); + assertEquals(segmentName.getTopicId(), 2); + } + + @Test + public void testConstructorWithoutTopicIdDefaultsTopicIdToZero() { + LLCSegmentName segmentName = new LLCSegmentName(TABLE_NAME, 0, 1, TIMESTAMP_MS); + assertEquals(segmentName.getTopicId(), 0); + } + + @Test + public void testIsMultiTopicLLCSegmentRecognizesFivePartNumericSegmentName() { + assertTrue(LLCSegmentName.isMultiTopicLLCSegment("myTable__2__4__27__20160617T2150Z")); + } + + @Test + public void testIsMultiTopicLLCSegmentRejectsFourPartSegmentName() { + assertFalse(LLCSegmentName.isMultiTopicLLCSegment("myTable__4__27__20160617T2150Z")); + } + + @Test + public void testIsMultiTopicLLCSegmentRejectsNonNumericParts() { + assertFalse(LLCSegmentName.isMultiTopicLLCSegment("myTable__uploaded__4__27__20160617T2150Z")); + } + + @Test + public void testCompareToThrowsWhenTopicIdsDiffer() { + LLCSegmentName segmentInTopic1 = new LLCSegmentName(TABLE_NAME, 1, 0, 1, TIMESTAMP_MS); + LLCSegmentName segmentInTopic2 = new LLCSegmentName(TABLE_NAME, 2, 0, 1, TIMESTAMP_MS); + try { + segmentInTopic1.compareTo(segmentInTopic2); + Assert.fail("Expected compareTo to throw for segments from different topic ids"); + } catch (IllegalArgumentException e) { + // expected + } + } + + @Test + public void testCompareToOrdersSegmentsWithSameNonZeroTopicId() { + LLCSegmentName segment1 = new LLCSegmentName(TABLE_NAME, 3, 0, 1, TIMESTAMP_MS); + LLCSegmentName segment2 = new LLCSegmentName(TABLE_NAME, 3, 0, 2, TIMESTAMP_MS); + assertTrue(segment1.compareTo(segment2) < 0); + assertTrue(segment2.compareTo(segment1) > 0); + assertEquals(segment1.compareTo(new LLCSegmentName(TABLE_NAME, 3, 0, 1, TIMESTAMP_MS)), 0); + } } From 0bc1c092750b4ea96674719d3ad5636211425aa4 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Mon, 13 Jul 2026 17:20:37 -0700 Subject: [PATCH 02/12] creating new segment name --- .../pinot/common/utils/LLCSegmentName.java | 9 ++ .../PinotLLCRealtimeSegmentManager.java | 114 +++++++++++------- 2 files changed, 82 insertions(+), 41 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java index daeb0ca1cdf0..b38f820abe11 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java @@ -39,6 +39,7 @@ public class LLCSegmentName implements Comparable { private final int _sequenceNumber; private final String _creationTime; private final String _segmentName; + private final boolean _isMultiTopicFormat; public LLCSegmentName(String segmentName) { String[] parts = StringUtils.splitByWholeSeparator(segmentName, SEPARATOR); @@ -49,11 +50,13 @@ public LLCSegmentName(String segmentName) { _partitionGroupId = Integer.parseInt(parts[1]); _sequenceNumber = Integer.parseInt(parts[2]); _creationTime = parts[3]; + _isMultiTopicFormat = false; } else { _topicId = Integer.parseInt(parts[1]); _partitionGroupId = Integer.parseInt(parts[2]); _sequenceNumber = Integer.parseInt(parts[3]); _creationTime = parts[4]; + _isMultiTopicFormat = true; } _segmentName = segmentName; } @@ -66,6 +69,7 @@ public LLCSegmentName(String tableName, int partitionGroupId, int sequenceNumber _sequenceNumber = sequenceNumber; // ISO8601 date: 20160120T1234Z _creationTime = DATE_FORMATTER.print(msSinceEpoch); + _isMultiTopicFormat = false; _segmentName = tableName + SEPARATOR + partitionGroupId + SEPARATOR + sequenceNumber + SEPARATOR + _creationTime; } @@ -76,6 +80,7 @@ public LLCSegmentName(String tableName, int configId, int partitionGroupId, int _partitionGroupId = partitionGroupId; _sequenceNumber = sequenceNumber; _creationTime = DATE_FORMATTER.print(msSinceEpoch); + _isMultiTopicFormat = false; _segmentName = tableName + SEPARATOR + configId + SEPARATOR + partitionGroupId + SEPARATOR + sequenceNumber + SEPARATOR + _creationTime; } @@ -162,6 +167,10 @@ public long getCreationTimeMs() { return dateTime.getMillis(); } + public boolean isMultiTopicFormat() { + return _isMultiTopicFormat; + } + @JsonValue public String getSegmentName() { return _segmentName; diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index d1d1f3bb3062..be7244401ac6 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -408,14 +408,16 @@ public void setUpNewTable(TableConfig tableConfig, IdealState idealState, List> instanceStatesMap = idealState.getRecord().getMapFields(); - for (StreamMetadata streamMetadata : streamMetadataList) { + boolean isMultiTopic = streamMetadataList.size() > 1; + for (int topicId = 0; topicId < streamMetadataList.size(); topicId++) { + StreamMetadata streamMetadata = streamMetadataList.get(topicId); StreamConfig streamConfig = streamMetadata.getStreamConfig(); for (PartitionGroupMetadata metadata : streamMetadata.getPartitionGroupMetadataList()) { int sequenceNumber = metadata.getSequenceNumber() >= 0 ? metadata.getSequenceNumber() : STARTING_SEQUENCE_NUMBER; String segmentName = setupNewPartitionGroup(tableConfig, streamConfig, metadata, sequenceNumber, currentTimeMs, - instancePartitions, numPartitionGroups, numReplicas); + instancePartitions, numPartitionGroups, numReplicas, topicId, isMultiTopic); updateInstanceStatesForNewConsumingSegment(instanceStatesMap, null, segmentName, segmentAssignment, instancePartitionsMap); } @@ -833,6 +835,7 @@ private String createNewSegmentMetadata(TableConfig tableConfig, int committingSegmentPartitionGroupId = committingLLCSegment.getPartitionGroupId(); List streamConfigs = IngestionConfigUtils.getStreamConfigs(tableConfig); + boolean isMultiTopic = streamConfigs.size() > 1; PartitionIdsWithIdealState partitionIdsWithIdealState = getPartitionIdsWithIdealState(streamConfigs, () -> getIdealState(realtimeTableName)); Set partitionIds = partitionIdsWithIdealState._partitionIds; @@ -850,8 +853,25 @@ private String createNewSegmentMetadata(TableConfig tableConfig, } String rawTableName = TableNameBuilder.extractRawTableName(realtimeTableName); long newSegmentCreationTimeMs = getCurrentTimeMs(); - LLCSegmentName newLLCSegment = new LLCSegmentName(rawTableName, committingSegmentPartitionGroupId, - committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); + LLCSegmentName newLLCSegment; + if (isMultiTopic) { + if (committingLLCSegment.isMultiTopicFormat()) { + // old segment is in the new format, only seq number needs to be updated for new segment + newLLCSegment = new LLCSegmentName(rawTableName, committingLLCSegment.getTopicId(), + committingSegmentPartitionGroupId, + committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); + } else { + // old segment is in the old format (index * 10000 + partition number), + // New segment should be created with new format (_topicId_partition number) + int topicId = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(committingSegmentPartitionGroupId); + newLLCSegment = new LLCSegmentName(rawTableName, topicId, + IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(committingSegmentPartitionGroupId), + committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); + } + } else { + newLLCSegment = new LLCSegmentName(rawTableName, committingSegmentPartitionGroupId, + committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); + } StreamConfig streamConfig = IngestionConfigUtils.getStreamConfigFromPinotPartitionId(streamConfigs, committingSegmentPartitionGroupId); @@ -1231,11 +1251,11 @@ Set getPartitionIds(List streamConfigs, IdealState idealS } private static class PartitionIdsWithIdealState { - private final Set _partitionIds; + private final Map> _partitionIds; @Nullable private final IdealState _idealState; - private PartitionIdsWithIdealState(Set partitionIds, @Nullable IdealState idealState) { + private PartitionIdsWithIdealState(Map> partitionIds, @Nullable IdealState idealState) { _partitionIds = partitionIds; _idealState = idealState; } @@ -1243,32 +1263,15 @@ private PartitionIdsWithIdealState(Set partitionIds, @Nullable IdealSta private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamConfigs, Supplier idealStateSupplier) { - Set partitionIds = new HashSet<>(); + Map> partitionIds = new HashMap<>(); boolean allPartitionIdsFetched = true; int numStreams = streamConfigs.size(); - if (numStreams == 1) { - // Single stream - // NOTE: We skip partition id translation logic to handle cases where custom stream might return partition id - // larger than 10000. - StreamConfig streamConfig = streamConfigs.get(0); - try { - partitionIds = getPartitionIds(streamConfig); - } catch (UnsupportedOperationException ignored) { - allPartitionIdsFetched = false; - // Stream does not support fetching partition ids. There is a log in the fallback code which is sufficient - } catch (Exception e) { - allPartitionIdsFetched = false; - LOGGER.warn("Failed to fetch partition ids for stream: {}", streamConfig.getTopicName(), e); - } - } else { // Multiple streams for (int i = 0; i < numStreams; i++) { StreamConfig streamConfig = streamConfigs.get(i); int index = i; try { - partitionIds.addAll(getPartitionIds(streamConfig).stream() - .map(partitionId -> IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(partitionId, index)) - .collect(Collectors.toSet())); + partitionIds.put(i, new HashSet<>(getPartitionIds(streamConfig))); } catch (UnsupportedOperationException ignored) { allPartitionIdsFetched = false; // Stream does not support fetching partition ids. There is a log in the fallback code which is sufficient @@ -1277,7 +1280,6 @@ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List instanceStateMap, Stri IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List streamConfigs, IdealState idealState, List streamMetadataList, OffsetCriteria offsetCriteria) { String realtimeTableName = tableConfig.getTableName(); + boolean isMultiTopic = streamConfigs.size() > 1; InstancePartitions instancePartitions = getConsumingInstancePartitions(tableConfig); int numReplicas = getNumReplicas(tableConfig, instancePartitions); @@ -1840,7 +1843,8 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List> instanceStatesMap, SegmentAssignment segmentAssignment, - Map instancePartitionsMap, StreamPartitionMsgOffset startOffset) { + Map instancePartitionsMap, StreamPartitionMsgOffset startOffset, + boolean isMultiTopic) { int numReplicas = getNumReplicas(tableConfig, instancePartitions); LLCSegmentName latestLLCSegmentName = new LLCSegmentName(latestSegmentZKMetadata.getSegmentName()); - LLCSegmentName newLLCSegmentName = getNextLLCSegmentName(latestLLCSegmentName, currentTimeMs); + LLCSegmentName newLLCSegmentName = getNextLLCSegmentName(latestLLCSegmentName, currentTimeMs, isMultiTopic); CommittingSegmentDescriptor committingSegmentDescriptor = new CommittingSegmentDescriptor(latestSegmentZKMetadata.getSegmentName(), startOffset.toString(), 0); createNewSegmentZKMetadata(tableConfig, streamConfig, newLLCSegmentName, currentTimeMs, committingSegmentDescriptor, @@ -2096,25 +2102,47 @@ private StreamPartitionMsgOffset selectStartOffset(OffsetCriteria offsetCriteria } } - private LLCSegmentName getNextLLCSegmentName(LLCSegmentName lastLLCSegmentName, long creationTimeMs) { - return new LLCSegmentName(lastLLCSegmentName.getTableName(), lastLLCSegmentName.getPartitionGroupId(), - lastLLCSegmentName.getSequenceNumber() + 1, creationTimeMs); + private LLCSegmentName getNextLLCSegmentName(LLCSegmentName lastLLCSegmentName, long creationTimeMs, + boolean isMultiTopic) { + LLCSegmentName newLLCSegment; + int lastSegmentPartitionGroupId = lastLLCSegmentName.getPartitionGroupId(); + String rawTableName = lastLLCSegmentName.getTableName(); + if (isMultiTopic) { + if (lastLLCSegmentName.isMultiTopicFormat()) { + // old segment is in the new format, only seq number needs to be updated for new segment + newLLCSegment = new LLCSegmentName(rawTableName, lastLLCSegmentName.getTopicId(), + lastSegmentPartitionGroupId, + lastLLCSegmentName.getSequenceNumber() + 1, creationTimeMs); + } else { + // old segment is in the old format (index * 10000 + partition number), + // New segment should be created with new format (_topicId_partition number) + int topicId = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(lastSegmentPartitionGroupId); + newLLCSegment = new LLCSegmentName(rawTableName, topicId, + IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(lastSegmentPartitionGroupId), + lastLLCSegmentName.getSequenceNumber() + 1, creationTimeMs); + } + } else { + newLLCSegment = new LLCSegmentName(rawTableName, lastSegmentPartitionGroupId, + lastLLCSegmentName.getSequenceNumber() + 1, creationTimeMs); + } + return newLLCSegment; } + /** * Sets up a new partition group. *

Persists the ZK metadata for the first CONSUMING segment, and returns the segment name. */ private String setupNewPartitionGroup(TableConfig tableConfig, StreamConfig streamConfig, PartitionGroupMetadata partitionGroupMetadata, long creationTimeMs, InstancePartitions instancePartitions, - int numPartitions, int numReplicas) { + int numPartitions, int numReplicas, int topicId, boolean isMultiTopic) { return setupNewPartitionGroup(tableConfig, streamConfig, partitionGroupMetadata, STARTING_SEQUENCE_NUMBER, - creationTimeMs, instancePartitions, numPartitions, numReplicas); + creationTimeMs, instancePartitions, numPartitions, numReplicas, topicId, isMultiTopic); } private String setupNewPartitionGroup(TableConfig tableConfig, StreamConfig streamConfig, PartitionGroupMetadata partitionGroupMetadata, int sequence, long creationTimeMs, - InstancePartitions instancePartitions, int numPartitions, int numReplicas) { + InstancePartitions instancePartitions, int numPartitions, int numReplicas, int topicId, boolean isMultiTopic) { Preconditions.checkArgument(sequence >= 0, "Sequence number must be >= 0, got: %s", sequence); String realtimeTableName = tableConfig.getTableName(); int partitionGroupId = partitionGroupMetadata.getPartitionGroupId(); @@ -2123,8 +2151,12 @@ private String setupNewPartitionGroup(TableConfig tableConfig, StreamConfig stre partitionGroupId, realtimeTableName, sequence, startOffset); String rawTableName = TableNameBuilder.extractRawTableName(realtimeTableName); - LLCSegmentName newLLCSegmentName = - new LLCSegmentName(rawTableName, partitionGroupId, sequence, creationTimeMs); + LLCSegmentName newLLCSegmentName; + if (isMultiTopic) { + newLLCSegmentName = new LLCSegmentName(rawTableName, topicId, partitionGroupId, sequence, creationTimeMs); + } else { + newLLCSegmentName = new LLCSegmentName(rawTableName, partitionGroupId, sequence, creationTimeMs); + } String newSegmentName = newLLCSegmentName.getSegmentName(); CommittingSegmentDescriptor committingSegmentDescriptor = new CommittingSegmentDescriptor(null, From a64782738d5f556e16f1264cf9a7120aaf3dfbca Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Mon, 13 Jul 2026 19:23:13 -0700 Subject: [PATCH 03/12] creating new segment name --- .../PinotLLCRealtimeSegmentManager.java | 87 +++++++++++-------- 1 file changed, 51 insertions(+), 36 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index be7244401ac6..d3caeddf1d09 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -838,9 +838,20 @@ private String createNewSegmentMetadata(TableConfig tableConfig, boolean isMultiTopic = streamConfigs.size() > 1; PartitionIdsWithIdealState partitionIdsWithIdealState = getPartitionIdsWithIdealState(streamConfigs, () -> getIdealState(realtimeTableName)); - Set partitionIds = partitionIdsWithIdealState._partitionIds; - - if (partitionIds.contains(committingSegmentPartitionGroupId)) { + Map> partitionIdsByTopic = partitionIdsWithIdealState._partitionIdsByTopic; + + // Normalize the committing segment's partition into (topicId, streamPartitionId), regardless of whether the + // committing segment name uses the old (index * 10000 + partitionId) or new (explicit topicId) format. + int committingTopicId = isMultiTopic + ? (committingLLCSegment.isMultiTopicFormat() ? committingLLCSegment.getTopicId() + : IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(committingSegmentPartitionGroupId)) + : 0; + int committingStreamPartitionId = isMultiTopic && !committingLLCSegment.isMultiTopicFormat() + ? IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(committingSegmentPartitionGroupId) + : committingSegmentPartitionGroupId; + + Set streamPartitionIds = partitionIdsByTopic.get(committingTopicId); + if (streamPartitionIds != null && streamPartitionIds.contains(committingStreamPartitionId)) { IdealState idealState = partitionIdsWithIdealState._idealState; if (idealState == null) { idealState = getIdealState(realtimeTableName); @@ -855,28 +866,18 @@ private String createNewSegmentMetadata(TableConfig tableConfig, long newSegmentCreationTimeMs = getCurrentTimeMs(); LLCSegmentName newLLCSegment; if (isMultiTopic) { - if (committingLLCSegment.isMultiTopicFormat()) { - // old segment is in the new format, only seq number needs to be updated for new segment - newLLCSegment = new LLCSegmentName(rawTableName, committingLLCSegment.getTopicId(), - committingSegmentPartitionGroupId, - committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); - } else { - // old segment is in the old format (index * 10000 + partition number), - // New segment should be created with new format (_topicId_partition number) - int topicId = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(committingSegmentPartitionGroupId); - newLLCSegment = new LLCSegmentName(rawTableName, topicId, - IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(committingSegmentPartitionGroupId), - committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); - } + // Always create the new segment in the new (explicit topicId) format. + newLLCSegment = new LLCSegmentName(rawTableName, committingTopicId, committingStreamPartitionId, + committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); } else { newLLCSegment = new LLCSegmentName(rawTableName, committingSegmentPartitionGroupId, committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); } - StreamConfig streamConfig = - IngestionConfigUtils.getStreamConfigFromPinotPartitionId(streamConfigs, committingSegmentPartitionGroupId); + StreamConfig streamConfig = streamConfigs.get(committingTopicId); + int totalPartitionCount = partitionIdsByTopic.values().stream().mapToInt(Set::size).sum(); createNewSegmentZKMetadata(tableConfig, streamConfig, newLLCSegment, newSegmentCreationTimeMs, - committingSegmentDescriptor, committingSegmentZKMetadata, instancePartitions, partitionIds.size(), + committingSegmentDescriptor, committingSegmentZKMetadata, instancePartitions, totalPartitionCount, numReplicas); newConsumingSegmentName = newLLCSegment.getSegmentName(); LOGGER.info("Created new segment metadata for segment: {} with status: {}.", newConsumingSegmentName, @@ -1247,39 +1248,44 @@ public Map getPartitionCountMap(List streamConfi @VisibleForTesting Set getPartitionIds(List streamConfigs, IdealState idealState) { - return getPartitionIdsWithIdealState(streamConfigs, () -> idealState)._partitionIds; + Set partitionIds = new HashSet<>(); + getPartitionIdsWithIdealState(streamConfigs, () -> idealState)._partitionIdsByTopic.values() + .forEach(partitionIds::addAll); + return partitionIds; } private static class PartitionIdsWithIdealState { - private final Map> _partitionIds; + private final Map> _partitionIdsByTopic; @Nullable private final IdealState _idealState; private PartitionIdsWithIdealState(Map> partitionIds, @Nullable IdealState idealState) { - _partitionIds = partitionIds; + _partitionIdsByTopic = partitionIds; _idealState = idealState; } } + /** + * Fetches, per stream config index (topic id), the set of raw (unpadded) stream partition ids currently known + * to that stream. Values are always raw stream partition ids, never padded/encoded pinot partition ids. + */ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamConfigs, Supplier idealStateSupplier) { Map> partitionIds = new HashMap<>(); boolean allPartitionIdsFetched = true; int numStreams = streamConfigs.size(); - // Multiple streams - for (int i = 0; i < numStreams; i++) { - StreamConfig streamConfig = streamConfigs.get(i); - int index = i; - try { - partitionIds.put(i, new HashSet<>(getPartitionIds(streamConfig))); - } catch (UnsupportedOperationException ignored) { - allPartitionIdsFetched = false; - // Stream does not support fetching partition ids. There is a log in the fallback code which is sufficient - } catch (Exception e) { - allPartitionIdsFetched = false; - LOGGER.warn("Failed to fetch partition ids for stream: {}", streamConfig.getTopicName(), e); - } + for (int i = 0; i < numStreams; i++) { + StreamConfig streamConfig = streamConfigs.get(i); + try { + partitionIds.put(i, new HashSet<>(getPartitionIds(streamConfig))); + } catch (UnsupportedOperationException ignored) { + allPartitionIdsFetched = false; + // Stream does not support fetching partition ids. There is a log in the fallback code which is sufficient + } catch (Exception e) { + allPartitionIdsFetched = false; + LOGGER.warn("Failed to fetch partition ids for stream: {}", streamConfig.getTopicName(), e); } + } // If it is failing to fetch partition ids from stream (usually transient due to stream metadata service outage), // we need to use the existing partition information from ideal state to keep same ingestion behavior. @@ -1294,9 +1300,18 @@ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamMetadataList = getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, idealState); + // NOTE: streamMetadataList may be shorter than streamConfigs (a permanently-failed topic is skipped rather + // than represented by a placeholder), so its list position is NOT the topic id. Each PartitionGroupMetadata's + // partitionGroupId already encodes the true topic id via padding (see PartitionGroupMetadataFetcher), so + // decode from that instead of trusting list position. for (StreamMetadata streamMetadata : streamMetadataList) { for (PartitionGroupMetadata partitionGroupMetadata : streamMetadata.getPartitionGroupMetadataList()) { - partitionIds.add(partitionGroupMetadata.getPartitionGroupId()); + int pinotPartitionId = partitionGroupMetadata.getPartitionGroupId(); + int topicId = numStreams > 1 + ? IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(pinotPartitionId) : 0; + int streamPartitionId = numStreams > 1 + ? IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(pinotPartitionId) : pinotPartitionId; + partitionIds.computeIfAbsent(topicId, k -> new HashSet<>()).add(streamPartitionId); } } return new PartitionIdsWithIdealState(partitionIds, idealState); From f0ce917cfa2ed7c39c68db7fdc138f2a4584753c Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Mon, 13 Jul 2026 23:48:33 -0700 Subject: [PATCH 04/12] moved some stuff to LLCSegmentName --- .../pinot/common/utils/LLCSegmentName.java | 19 +++++++++++ .../PinotLLCRealtimeSegmentManager.java | 32 ++++++------------- 2 files changed, 28 insertions(+), 23 deletions(-) diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java index b38f820abe11..985de7ecb3fa 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java @@ -23,6 +23,7 @@ import java.util.Objects; import javax.annotation.Nullable; import org.apache.commons.lang3.StringUtils; +import org.apache.pinot.spi.utils.IngestionConfigUtils; import org.joda.time.DateTime; import org.joda.time.format.DateTimeFormat; import org.joda.time.format.DateTimeFormatter; @@ -150,6 +151,15 @@ public int getTopicId() { return _topicId; } + public int getTopicId(boolean hasMultiTopic) { + if (hasMultiTopic) { + return _isMultiTopicFormat ? _topicId : + IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(_partitionGroupId); + } else { + return _topicId; + } + } + public int getPartitionGroupId() { return _partitionGroupId; } @@ -176,6 +186,15 @@ public String getSegmentName() { return _segmentName; } + public int getStreamPartitionGroupId(boolean hasMultiTopics) { + if (hasMultiTopics) { + return _isMultiTopicFormat ? _partitionGroupId : + IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_partitionGroupId); + } else { + return _partitionGroupId; + } + } + @Override public int compareTo(LLCSegmentName other) { Preconditions.checkArgument(_tableName.equals(other._tableName), diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index d3caeddf1d09..ed1dd337a439 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -317,9 +317,9 @@ public List getPartitionGroupConsumptionStatusL StreamPartitionMsgOffsetFactory[] offsetFactories = new StreamPartitionMsgOffsetFactory[numStreams]; for (Map.Entry entry : partitionGroupIdToLatestSegment.entrySet()) { int partitionGroupId = entry.getKey(); - int index = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(partitionGroupId); - int streamPartitionId = IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(partitionGroupId); LLCSegmentName llcSegmentName = entry.getValue(); + int index = llcSegmentName.getTopicId(true); + int streamPartitionId = llcSegmentName.getStreamPartitionGroupId(true); SegmentZKMetadata segmentZKMetadata = getSegmentZKMetadata(tableNameWithType, llcSegmentName.getSegmentName()); StreamPartitionMsgOffsetFactory offsetFactory = offsetFactories[index]; if (offsetFactory == null) { @@ -842,13 +842,8 @@ private String createNewSegmentMetadata(TableConfig tableConfig, // Normalize the committing segment's partition into (topicId, streamPartitionId), regardless of whether the // committing segment name uses the old (index * 10000 + partitionId) or new (explicit topicId) format. - int committingTopicId = isMultiTopic - ? (committingLLCSegment.isMultiTopicFormat() ? committingLLCSegment.getTopicId() - : IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(committingSegmentPartitionGroupId)) - : 0; - int committingStreamPartitionId = isMultiTopic && !committingLLCSegment.isMultiTopicFormat() - ? IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(committingSegmentPartitionGroupId) - : committingSegmentPartitionGroupId; + int committingTopicId = committingLLCSegment.getTopicId(isMultiTopic); + int committingStreamPartitionId = committingLLCSegment.getStreamPartitionGroupId(isMultiTopic); Set streamPartitionIds = partitionIdsByTopic.get(committingTopicId); if (streamPartitionIds != null && streamPartitionIds.contains(committingStreamPartitionId)) { @@ -1018,10 +1013,11 @@ private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig st int numReplicas) { String realtimeTableName = tableConfig.getTableName(); String segmentName = newLLCSegmentName.getSegmentName(); - + boolean isMultiTopic = IngestionConfigUtils.hasMultipleStreams(tableConfig); // Handle offset auto reset String nextOffset = committingSegmentDescriptor.getNextOffset(); - String startOffset = computeStartOffset(nextOffset, streamConfig, newLLCSegmentName.getPartitionGroupId()); + String startOffset = computeStartOffset(nextOffset, streamConfig, + newLLCSegmentName.getStreamPartitionGroupId(isMultiTopic)); if (!startOffset.equals(nextOffset)) { _controllerMetrics.addMeteredTableValue(realtimeTableName, ControllerMeter.OFFSET_AUTO_RESET_SKIPPED_OFFSETS, Long.parseLong(startOffset) - Long.parseLong(nextOffset)); @@ -1086,7 +1082,7 @@ private void updateFlushThresholdGauge(StreamConfig streamConfig, int flushThres ControllerGauge.NUM_ROWS_THRESHOLD_WITH_TOPIC, flushThreshold); } - private String computeStartOffset(String nextOffset, StreamConfig streamConfig, int partitionId) { + private String computeStartOffset(String nextOffset, StreamConfig streamConfig, int streamPartitionId) { if (!streamConfig.isEnableOffsetAutoReset() || streamConfig.isBackfillTopic()) { return nextOffset; } @@ -1099,7 +1095,6 @@ private String computeStartOffset(String nextOffset, StreamConfig streamConfig, timeThreshold, offsetThreshold); return nextOffset; } - int streamPartitionId = IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(partitionId); String clientId = getTableTopicUniqueClientId(streamConfig); StreamConsumerFactory consumerFactory = StreamConsumerFactoryProvider.create(streamConfig); StreamPartitionMsgOffsetFactory offsetFactory = consumerFactory.createStreamMsgOffsetFactory(); @@ -1300,18 +1295,9 @@ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamMetadataList = getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, idealState); - // NOTE: streamMetadataList may be shorter than streamConfigs (a permanently-failed topic is skipped rather - // than represented by a placeholder), so its list position is NOT the topic id. Each PartitionGroupMetadata's - // partitionGroupId already encodes the true topic id via padding (see PartitionGroupMetadataFetcher), so - // decode from that instead of trusting list position. for (StreamMetadata streamMetadata : streamMetadataList) { for (PartitionGroupMetadata partitionGroupMetadata : streamMetadata.getPartitionGroupMetadataList()) { - int pinotPartitionId = partitionGroupMetadata.getPartitionGroupId(); - int topicId = numStreams > 1 - ? IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(pinotPartitionId) : 0; - int streamPartitionId = numStreams > 1 - ? IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(pinotPartitionId) : pinotPartitionId; - partitionIds.computeIfAbsent(topicId, k -> new HashSet<>()).add(streamPartitionId); + partitionIds.add(partitionGroupMetadata.getPartitionGroupId()); } } return new PartitionIdsWithIdealState(partitionIds, idealState); From 9aa99094fc9772f9078fcc3ec9de7fdd98e05bc2 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 01:37:57 -0700 Subject: [PATCH 05/12] Fix multi-topic partition id decoding in getPartitionIdsWithIdealState fallback The fallback branch (used when fetching partition ids directly from the stream fails) called Map>.add(...), which does not compile. Restore the fallback to decode each PartitionGroupMetadata's padded pinot partition id into (topicId, streamPartitionId) using the existing IngestionConfigUtils helpers, matching the convention already used in createNewSegmentMetadata/getNextLLCSegmentName, with a pass-through shortcut for single-stream tables. --- .../PinotLLCRealtimeSegmentManager.java | 21 +++++-- .../PinotLLCRealtimeSegmentManagerTest.java | 62 +++++++++++++++++++ 2 files changed, 79 insertions(+), 4 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index ed1dd337a439..047cb275566e 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -1244,12 +1244,13 @@ public Map getPartitionCountMap(List streamConfi @VisibleForTesting Set getPartitionIds(List streamConfigs, IdealState idealState) { Set partitionIds = new HashSet<>(); - getPartitionIdsWithIdealState(streamConfigs, () -> idealState)._partitionIdsByTopic.values() + getPartitionIdsWithIdealState(streamConfigs, () -> idealState).getPartitionIdsByTopic().values() .forEach(partitionIds::addAll); return partitionIds; } - private static class PartitionIdsWithIdealState { + @VisibleForTesting + static class PartitionIdsWithIdealState { private final Map> _partitionIdsByTopic; @Nullable private final IdealState _idealState; @@ -1258,13 +1259,19 @@ private PartitionIdsWithIdealState(Map> partitionIds, @Nul _partitionIdsByTopic = partitionIds; _idealState = idealState; } + + @VisibleForTesting + Map> getPartitionIdsByTopic() { + return _partitionIdsByTopic; + } } /** * Fetches, per stream config index (topic id), the set of raw (unpadded) stream partition ids currently known * to that stream. Values are always raw stream partition ids, never padded/encoded pinot partition ids. */ - private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamConfigs, + @VisibleForTesting + PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamConfigs, Supplier idealStateSupplier) { Map> partitionIds = new HashMap<>(); boolean allPartitionIdsFetched = true; @@ -1295,9 +1302,15 @@ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamMetadataList = getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, idealState); + boolean isMultiTopic = numStreams > 1; for (StreamMetadata streamMetadata : streamMetadataList) { for (PartitionGroupMetadata partitionGroupMetadata : streamMetadata.getPartitionGroupMetadataList()) { - partitionIds.add(partitionGroupMetadata.getPartitionGroupId()); + int pinotPartitionId = partitionGroupMetadata.getPartitionGroupId(); + int topicId = isMultiTopic + ? IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(pinotPartitionId) : 0; + int streamPartitionId = isMultiTopic + ? IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(pinotPartitionId) : pinotPartitionId; + partitionIds.computeIfAbsent(topicId, k -> new HashSet<>()).add(streamPartitionId); } } return new PartitionIdsWithIdealState(partitionIds, idealState); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java index 45e85b57e640..ae5c437cf136 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java @@ -1973,6 +1973,68 @@ public void testGetPartitionIds() Assert.assertEquals(partitionIds.size(), 2); } + @Test + public void testGetPartitionIdsWithIdealStateSingleStreamFallbackLargePartitionId() + throws Exception { + List streamConfigs = List.of(FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs()); + IdealState idealState = new IdealState(REALTIME_TABLE_NAME); + + // Simulate the case where getPartitionIds(StreamConfig) throws, forcing the fallback path + PinotLLCRealtimeSegmentManager segmentManagerSpy = spy(new FakePinotLLCRealtimeSegmentManager()); + doThrow(new RuntimeException()).when(segmentManagerSpy).getPartitionIds(any(StreamConfig.class)); + doReturn(List.of()).when(segmentManagerSpy).getPartitionGroupConsumptionStatusList(idealState, streamConfigs); + + // A single-stream table is not multi-topic, so a partition id above PARTITION_PADDING_OFFSET must be passed + // through unchanged rather than decoded into (topicId, streamPartitionId). + int largePartitionId = 20003; + List streamMetadataList = List.of(new StreamMetadata(streamConfigs.get(0), 2, + List.of(new PartitionGroupMetadata(largePartitionId, new LongMsgOffset(1))))); + doReturn(streamMetadataList).when(segmentManagerSpy) + .getNewStreamMetadataList(eq(streamConfigs), any(), eq(idealState)); + + PinotLLCRealtimeSegmentManager.PartitionIdsWithIdealState result = + segmentManagerSpy.getPartitionIdsWithIdealState(streamConfigs, () -> idealState); + Assert.assertEquals(result.getPartitionIdsByTopic().get(0), Set.of(largePartitionId)); + } + + @Test + public void testGetPartitionIdsWithIdealStateMultiTopicFallback() + throws Exception { + List streamConfigs = + List.of(FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs(), FakeStreamConfigUtils + .getDefaultLowLevelStreamConfigs()); + IdealState idealState = new IdealState(REALTIME_TABLE_NAME); + + // Simulate the case where getPartitionIds(StreamConfig) throws for both streams, forcing the fallback path + PinotLLCRealtimeSegmentManager segmentManagerSpy = spy(new FakePinotLLCRealtimeSegmentManager()); + doThrow(new RuntimeException()).when(segmentManagerSpy).getPartitionIds(any(StreamConfig.class)); + doReturn(List.of()).when(segmentManagerSpy).getPartitionGroupConsumptionStatusList(idealState, streamConfigs); + + // Topic 0, raw partitions 0 and 3 -> padded pinot partition ids 0 and 3 (two groups for the same topic, to + // verify they accumulate into one Set instead of the second overwriting the first) + // Topic 1, raw partition 5 -> padded pinot partition id 10005 + // topicId and streamPartitionId are chosen to differ per entry so a transposed topicId/streamPartitionId + // formula would produce a visibly wrong result rather than accidentally matching. + List streamMetadataList = List.of( + new StreamMetadata(streamConfigs.get(0), 2, + List.of(new PartitionGroupMetadata( + IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(0, 0), new LongMsgOffset(1)), + new PartitionGroupMetadata( + IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(3, 0), new LongMsgOffset(1)))), + new StreamMetadata(streamConfigs.get(1), 2, + List.of(new PartitionGroupMetadata( + IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(5, 1), new LongMsgOffset(1))))); + doReturn(streamMetadataList).when(segmentManagerSpy) + .getNewStreamMetadataList(eq(streamConfigs), any(), eq(idealState)); + + PinotLLCRealtimeSegmentManager.PartitionIdsWithIdealState result = + segmentManagerSpy.getPartitionIdsWithIdealState(streamConfigs, () -> idealState); + // Asserting on the per-topic map (not the flattened union) so a transposed topicId/streamPartitionId formula + // would be caught even though it could coincidentally produce the same flattened set. + Assert.assertEquals(result.getPartitionIdsByTopic().get(0), Set.of(0, 3)); + Assert.assertEquals(result.getPartitionIdsByTopic().get(1), Set.of(5)); + } + /** * Verifies that {@code buildPartitionGroupConsumptionStatusFromZKMetadata} produces the same results as * {@code getPartitionGroupConsumptionStatusList} for the common case where IdealState and ZK metadata are in sync. From 0796e310c95be0c6df3603806b400fc9a2db91d7 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 06:58:03 -0700 Subject: [PATCH 06/12] WIP: add streamPartitionGroupId and topicId fields to PartitionGroupMetadata Fields not yet wired through constructors/getters. --- .../org/apache/pinot/spi/stream/PartitionGroupMetadata.java | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java index 3b7103c82848..43718f29dd88 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java @@ -34,6 +34,8 @@ public class PartitionGroupMetadata { private final int _partitionGroupId; private final StreamPartitionMsgOffset _startOffset; private final int _sequenceNumber; + private final int _streamPartitionGroupId; + private final int _topicId; public PartitionGroupMetadata(int partitionGroupId, StreamPartitionMsgOffset startOffset) { this(partitionGroupId, startOffset, DEFAULT_SEQUENCE_NUMBER); From 80c04131355065fc2052d641c1028fd6fbfca3da Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 09:03:51 -0700 Subject: [PATCH 07/12] Add topicId field to PartitionGroupMetadata and PartitionGroupConsumptionStatus Additive, non-breaking: new constructor overloads accept an explicit topicId identifying which stream config (topic) a partition group belongs to, defaulting to 0 when omitted. Wires topicId through PartitionGroupMetadataFetcher.fetchMultipleStreams using the per-stream index. No existing callers are changed yet. --- .../spi/stream/PartitionGroupConsumptionStatus.java | 12 ++++++++++++ .../pinot/spi/stream/PartitionGroupMetadata.java | 12 +++++++++++- .../spi/stream/PartitionGroupMetadataFetcher.java | 2 +- .../stream/PartitionGroupMetadataFetcherTest.java | 9 +++++++++ 4 files changed, 33 insertions(+), 2 deletions(-) diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupConsumptionStatus.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupConsumptionStatus.java index d0405906cdae..ac21b38884e1 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupConsumptionStatus.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupConsumptionStatus.java @@ -31,12 +31,14 @@ * and recorded the end * offset) * 6. status - the consumption status IN_PROGRESS/DONE + * 7. topicId - The id of the topic (stream config) this partition group belongs to, for tables with multiple topics * * This information is needed by the stream, when grouping the partitions/shards into new partition groups. */ public class PartitionGroupConsumptionStatus { private final int _partitionGroupId; private final int _streamPartitionId; + private final int _topicId; private int _sequenceNumber; private StreamPartitionMsgOffset _startOffset; private StreamPartitionMsgOffset _endOffset; @@ -44,8 +46,14 @@ public class PartitionGroupConsumptionStatus { public PartitionGroupConsumptionStatus(int partitionGroupId, int streamPartitionId, int sequenceNumber, StreamPartitionMsgOffset startOffset, StreamPartitionMsgOffset endOffset, String status) { + this(partitionGroupId, streamPartitionId, 0, sequenceNumber, startOffset, endOffset, status); + } + + public PartitionGroupConsumptionStatus(int partitionGroupId, int streamPartitionId, int topicId, + int sequenceNumber, StreamPartitionMsgOffset startOffset, StreamPartitionMsgOffset endOffset, String status) { _partitionGroupId = partitionGroupId; _streamPartitionId = streamPartitionId; + _topicId = topicId; _sequenceNumber = sequenceNumber; _startOffset = startOffset; _endOffset = endOffset; @@ -65,6 +73,10 @@ public int getStreamPartitionGroupId() { return _streamPartitionId; } + public int getTopicId() { + return _topicId; + } + public int getSequenceNumber() { return _sequenceNumber; } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java index 43718f29dd88..36e6ae878466 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadata.java @@ -26,6 +26,7 @@ * 1. A unique partition group id for this partition group * 2. The start offset to begin consumption for this partition group * 3. The sequence number for the consuming segment (used when creating segments with designated offsets/sequences) + * 4. The id of the topic (stream config) this partition group belongs to, for tables with multiple topics */ public class PartitionGroupMetadata { @@ -34,7 +35,6 @@ public class PartitionGroupMetadata { private final int _partitionGroupId; private final StreamPartitionMsgOffset _startOffset; private final int _sequenceNumber; - private final int _streamPartitionGroupId; private final int _topicId; public PartitionGroupMetadata(int partitionGroupId, StreamPartitionMsgOffset startOffset) { @@ -42,7 +42,13 @@ public PartitionGroupMetadata(int partitionGroupId, StreamPartitionMsgOffset sta } public PartitionGroupMetadata(int partitionGroupId, StreamPartitionMsgOffset startOffset, int sequenceNumber) { + this(partitionGroupId, 0, startOffset, sequenceNumber); + } + + public PartitionGroupMetadata(int partitionGroupId, int topicId, StreamPartitionMsgOffset startOffset, + int sequenceNumber) { _partitionGroupId = partitionGroupId; + _topicId = topicId; _startOffset = startOffset; _sequenceNumber = sequenceNumber; } @@ -51,6 +57,10 @@ public int getPartitionGroupId() { return _partitionGroupId; } + public int getTopicId() { + return _topicId; + } + public StreamPartitionMsgOffset getStartOffset() { return _startOffset; } diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java index edf6d6ccdd46..6ef5cd000c43 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java @@ -143,7 +143,7 @@ private Boolean fetchMultipleStreams() .stream() .map(metadata -> new PartitionGroupMetadata( IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(metadata.getPartitionGroupId(), - index), metadata.getStartOffset(), metadata.getSequenceNumber())) + index), index, metadata.getStartOffset(), metadata.getSequenceNumber())) .collect(Collectors.toList()); _streamMetadataList.add( new StreamMetadata(streamConfig, partitionGroupMetadataList.size(), partitionGroupMetadataList)); diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java index 4983ea4dc265..860c2134b6cb 100644 --- a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java @@ -160,6 +160,15 @@ public void testFetchMultipleStreams() .collect(Collectors.toList()); Assert.assertEquals(partitionIds, Arrays.asList(0, 1, 10000, 10001)); + + // Verify the topic ids are correctly assigned: topic1 (index 0) -> 0, 0; topic2 (index 1) -> 1, 1 + List topicIds = streamMetadataList.stream() + .flatMap(sm -> sm.getPartitionGroupMetadataList().stream()) + .map(PartitionGroupMetadata::getTopicId) + .sorted() + .collect(Collectors.toList()); + + Assert.assertEquals(topicIds, Arrays.asList(0, 0, 1, 1)); } } From b9d47681beac2bbf335525cafd24d2bc9b1cb74a Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 12:07:48 -0700 Subject: [PATCH 08/12] Use LLCSegmentName topicId/streamPartitionId accessors in PinotLLCRealtimeSegmentManager buildPartitionGroupConsumptionStatusFromZKMetadata and getPartitionGroupConsumptionStatusList relied on IngestionConfigUtils padded-id arithmetic to decode topicId/streamPartitionId, which is wrong for multi-topic-format segment names where these values are embedded directly and unpadded. Use LLCSegmentName.getTopicId(true)/ getStreamPartitionGroupId(true) instead, mirroring the pattern already used elsewhere. --- .../PinotLLCRealtimeSegmentManager.java | 30 +++-- .../PinotLLCRealtimeSegmentManagerTest.java | 123 ++++++++++++++++++ 2 files changed, 140 insertions(+), 13 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index 047cb275566e..8e0999db5e9a 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -306,7 +306,8 @@ public List getPartitionGroupConsumptionStatusL LLCSegmentName llcSegmentName = entry.getValue(); SegmentZKMetadata segmentZKMetadata = getSegmentZKMetadata(tableNameWithType, llcSegmentName.getSegmentName()); PartitionGroupConsumptionStatus partitionGroupConsumptionStatus = - new PartitionGroupConsumptionStatus(partitionGroupId, llcSegmentName.getSequenceNumber(), + new PartitionGroupConsumptionStatus(partitionGroupId, partitionGroupId, 0, + llcSegmentName.getSequenceNumber(), offsetFactory.create(segmentZKMetadata.getStartOffset()), segmentZKMetadata.getEndOffset() != null ? offsetFactory.create(segmentZKMetadata.getEndOffset()) : null, segmentZKMetadata.getStatus().toString()); @@ -318,16 +319,18 @@ public List getPartitionGroupConsumptionStatusL for (Map.Entry entry : partitionGroupIdToLatestSegment.entrySet()) { int partitionGroupId = entry.getKey(); LLCSegmentName llcSegmentName = entry.getValue(); - int index = llcSegmentName.getTopicId(true); + int topicId = llcSegmentName.getTopicId(true); int streamPartitionId = llcSegmentName.getStreamPartitionGroupId(true); SegmentZKMetadata segmentZKMetadata = getSegmentZKMetadata(tableNameWithType, llcSegmentName.getSegmentName()); - StreamPartitionMsgOffsetFactory offsetFactory = offsetFactories[index]; + StreamPartitionMsgOffsetFactory offsetFactory = offsetFactories[topicId]; if (offsetFactory == null) { - offsetFactory = StreamConsumerFactoryProvider.create(streamConfigs.get(index)).createStreamMsgOffsetFactory(); - offsetFactories[index] = offsetFactory; + offsetFactory = + StreamConsumerFactoryProvider.create(streamConfigs.get(topicId)).createStreamMsgOffsetFactory(); + offsetFactories[topicId] = offsetFactory; } PartitionGroupConsumptionStatus partitionGroupConsumptionStatus = - new PartitionGroupConsumptionStatus(partitionGroupId, streamPartitionId, llcSegmentName.getSequenceNumber(), + new PartitionGroupConsumptionStatus(partitionGroupId, streamPartitionId, topicId, + llcSegmentName.getSequenceNumber(), offsetFactory.create(segmentZKMetadata.getStartOffset()), segmentZKMetadata.getEndOffset() != null ? offsetFactory.create(segmentZKMetadata.getEndOffset()) : null, segmentZKMetadata.getStatus().toString()); @@ -2065,7 +2068,8 @@ List buildPartitionGroupConsumptionStatusFromZK int partitionGroupId = entry.getKey(); SegmentZKMetadata zkMetadata = entry.getValue(); LLCSegmentName llcSegmentName = new LLCSegmentName(zkMetadata.getSegmentName()); - result.add(new PartitionGroupConsumptionStatus(partitionGroupId, llcSegmentName.getSequenceNumber(), + result.add(new PartitionGroupConsumptionStatus(partitionGroupId, partitionGroupId, 0, + llcSegmentName.getSequenceNumber(), offsetFactory.create(zkMetadata.getStartOffset()), zkMetadata.getEndOffset() != null ? offsetFactory.create(zkMetadata.getEndOffset()) : null, zkMetadata.getStatus().toString())); @@ -2074,17 +2078,17 @@ List buildPartitionGroupConsumptionStatusFromZK StreamPartitionMsgOffsetFactory[] offsetFactories = new StreamPartitionMsgOffsetFactory[numStreams]; for (Map.Entry entry : latestSegmentZKMetadataMap.entrySet()) { int partitionGroupId = entry.getKey(); - int index = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(partitionGroupId); - int streamPartitionId = IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(partitionGroupId); SegmentZKMetadata zkMetadata = entry.getValue(); LLCSegmentName llcSegmentName = new LLCSegmentName(zkMetadata.getSegmentName()); - StreamPartitionMsgOffsetFactory offsetFactory = offsetFactories[index]; + int topicId = llcSegmentName.getTopicId(true); + int streamPartitionId = llcSegmentName.getStreamPartitionGroupId(true); + StreamPartitionMsgOffsetFactory offsetFactory = offsetFactories[topicId]; if (offsetFactory == null) { offsetFactory = - StreamConsumerFactoryProvider.create(streamConfigs.get(index)).createStreamMsgOffsetFactory(); - offsetFactories[index] = offsetFactory; + StreamConsumerFactoryProvider.create(streamConfigs.get(topicId)).createStreamMsgOffsetFactory(); + offsetFactories[topicId] = offsetFactory; } - result.add(new PartitionGroupConsumptionStatus(partitionGroupId, streamPartitionId, + result.add(new PartitionGroupConsumptionStatus(partitionGroupId, streamPartitionId, topicId, llcSegmentName.getSequenceNumber(), offsetFactory.create(zkMetadata.getStartOffset()), zkMetadata.getEndOffset() != null ? offsetFactory.create(zkMetadata.getEndOffset()) : null, diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java index ae5c437cf136..a49b5f00f00e 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java @@ -2098,6 +2098,129 @@ public void testBuildPartitionGroupConsumptionStatusFromZKMetadataMatchesOrigina assertEquals(zkEnd, isEnd, "End offset mismatch for partition " + isStatus.getPartitionGroupId()); assertEquals(zkStatus.getStatus(), isStatus.getStatus(), "Status mismatch for partition " + isStatus.getPartitionGroupId()); + assertEquals(isStatus.getTopicId(), 0, "Expected topicId 0 for single-stream table, partition " + + isStatus.getPartitionGroupId()); + assertEquals(zkStatus.getTopicId(), 0, "Expected topicId 0 for single-stream table, partition " + + isStatus.getPartitionGroupId()); + } + } + + /** + * Verifies that {@code getPartitionGroupConsumptionStatusList} and + * {@code buildPartitionGroupConsumptionStatusFromZKMetadata} correctly populate {@code topicId} (rather than + * silently defaulting it to 0) for a multi-stream table, by decoding it from the padded partition group id. + */ + @Test + public void testPartitionGroupConsumptionStatusTopicIdMultiStream() { + FakePinotLLCRealtimeSegmentManager segmentManager = new FakePinotLLCRealtimeSegmentManager(); + Map streamConfigMap = FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs().getStreamConfigsMap(); + IngestionConfig ingestionConfig = new IngestionConfig(); + ingestionConfig.setStreamIngestionConfig( + new StreamIngestionConfig(Arrays.asList(streamConfigMap, streamConfigMap))); + TableConfig multiStreamTableConfig = + new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setIngestionConfig(ingestionConfig) + .build(); + List streamConfigs = IngestionConfigUtils.getStreamConfigs(multiStreamTableConfig); + + // Topic 0, raw partition 1 -> padded pinot partition id 1 + // Topic 1, raw partition 2 -> padded pinot partition id 10002 + int topic0PartitionGroupId = IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(1, 0); + int topic1PartitionGroupId = IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(2, 1); + + IdealState idealState = new IdealState(REALTIME_TABLE_NAME); + Map latestSegmentZKMetadataMap = new HashMap<>(); + for (int partitionGroupId : new int[]{topic0PartitionGroupId, topic1PartitionGroupId}) { + String segmentName = + new LLCSegmentName(RAW_TABLE_NAME, partitionGroupId, 0, CURRENT_TIME_MS).getSegmentName(); + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentName); + segmentZKMetadata.setStatus(Status.IN_PROGRESS); + segmentZKMetadata.setStartOffset(new LongMsgOffset(0).toString()); + segmentManager._segmentZKMetadataMap.put(segmentName, segmentZKMetadata); + latestSegmentZKMetadataMap.put(partitionGroupId, segmentZKMetadata); + idealState.getRecord().setMapField(segmentName, Map.of("server0", SegmentStateModel.CONSUMING)); + } + + List fromIdealState = + segmentManager.getPartitionGroupConsumptionStatusList(idealState, streamConfigs); + List fromZKMetadata = + segmentManager.buildPartitionGroupConsumptionStatusFromZKMetadata(latestSegmentZKMetadataMap, streamConfigs); + + for (List statusList : List.of(fromIdealState, fromZKMetadata)) { + Map byPartitionGroupId = new HashMap<>(); + for (PartitionGroupConsumptionStatus status : statusList) { + byPartitionGroupId.put(status.getPartitionGroupId(), status); + } + assertEquals(byPartitionGroupId.get(topic0PartitionGroupId).getTopicId(), 0, + "Expected topicId 0 for partition group " + topic0PartitionGroupId); + assertEquals(byPartitionGroupId.get(topic0PartitionGroupId).getStreamPartitionGroupId(), 1, + "Expected streamPartitionId 1 for partition group " + topic0PartitionGroupId); + assertEquals(byPartitionGroupId.get(topic1PartitionGroupId).getTopicId(), 1, + "Expected topicId 1 for partition group " + topic1PartitionGroupId); + assertEquals(byPartitionGroupId.get(topic1PartitionGroupId).getStreamPartitionGroupId(), 2, + "Expected streamPartitionId 2 for partition group " + topic1PartitionGroupId); + } + } + + /** + * Same as {@link #testPartitionGroupConsumptionStatusTopicIdMultiStream}, but using the newer multi-topic-format + * segment name (topicId embedded directly in the segment name, e.g. {@code table__1__5__100__...}) instead of the + * legacy padded-partition-id encoding, to verify {@code LLCSegmentName.getTopicId(true)}/ + * {@code getStreamPartitionGroupId(true)} are used (rather than always re-deriving from the padded id, which + * would be wrong for multi-topic-format segment names). + */ + @Test + public void testPartitionGroupConsumptionStatusTopicIdMultiStreamMultiTopicFormat() { + FakePinotLLCRealtimeSegmentManager segmentManager = new FakePinotLLCRealtimeSegmentManager(); + Map streamConfigMap = FakeStreamConfigUtils.getDefaultLowLevelStreamConfigs().getStreamConfigsMap(); + IngestionConfig ingestionConfig = new IngestionConfig(); + ingestionConfig.setStreamIngestionConfig( + new StreamIngestionConfig(Arrays.asList(streamConfigMap, streamConfigMap))); + TableConfig multiStreamTableConfig = + new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME).setIngestionConfig(ingestionConfig) + .build(); + List streamConfigs = IngestionConfigUtils.getStreamConfigs(multiStreamTableConfig); + + // Multi-topic-format segment names encode topicId and (unpadded) streamPartitionId directly, so use + // deliberately "unaligned" values (e.g. streamPartitionId 7 for topic 0) that a padded-id-arithmetic decode + // would get wrong, to make sure the test would actually catch the regression. topicId values must stay within + // [0, numStreams) since they're used to index into the configured stream configs. + int topicId0 = 0; + int streamPartitionId0 = 7; + int topicId1 = 1; + int streamPartitionId1 = 2; + + IdealState idealState = new IdealState(REALTIME_TABLE_NAME); + Map latestSegmentZKMetadataMap = new HashMap<>(); + for (int[] topicAndPartition : new int[][]{{topicId0, streamPartitionId0}, {topicId1, streamPartitionId1}}) { + // Round-trip through the segment name string, since only the string-parsing constructor sets + // isMultiTopicFormat=true. + String segmentNameStr = + new LLCSegmentName(RAW_TABLE_NAME, topicAndPartition[0], topicAndPartition[1], 0, CURRENT_TIME_MS) + .getSegmentName(); + LLCSegmentName llcSegmentName = new LLCSegmentName(segmentNameStr); + Assert.assertTrue(llcSegmentName.isMultiTopicFormat()); + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentNameStr); + segmentZKMetadata.setStatus(Status.IN_PROGRESS); + segmentZKMetadata.setStartOffset(new LongMsgOffset(0).toString()); + segmentManager._segmentZKMetadataMap.put(segmentNameStr, segmentZKMetadata); + latestSegmentZKMetadataMap.put(llcSegmentName.getPartitionGroupId(), segmentZKMetadata); + idealState.getRecord().setMapField(segmentNameStr, Map.of("server0", SegmentStateModel.CONSUMING)); + } + + List fromIdealState = + segmentManager.getPartitionGroupConsumptionStatusList(idealState, streamConfigs); + List fromZKMetadata = + segmentManager.buildPartitionGroupConsumptionStatusFromZKMetadata(latestSegmentZKMetadataMap, streamConfigs); + + for (List statusList : List.of(fromIdealState, fromZKMetadata)) { + Map byStreamPartitionId = new HashMap<>(); + for (PartitionGroupConsumptionStatus status : statusList) { + byStreamPartitionId.put(status.getStreamPartitionGroupId(), status); + } + assertEquals(byStreamPartitionId.get(streamPartitionId0).getTopicId(), topicId0, + "Expected topicId " + topicId0 + " for stream partition " + streamPartitionId0); + assertEquals(byStreamPartitionId.get(streamPartitionId1).getTopicId(), topicId1, + "Expected topicId " + topicId1 + " for stream partition " + streamPartitionId1); } } From e40782340ccc241077000a07849a7249df32dd71 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 12:08:09 -0700 Subject: [PATCH 09/12] Use LLCSegmentName topicId/streamPartitionId accessors in RealtimeSegmentDataManager The multi-stream constructor branch used IngestionConfigUtils padded-id arithmetic to decode topicId/streamPartitionId, which is wrong for multi-topic-format segment names. Use LLCSegmentName.getTopicId(true)/ getStreamPartitionGroupId(true) instead, consistent with PinotLLCRealtimeSegmentManager. Also fix the Preconditions.checkState message to reference topicId instead of the old "index" terminology. --- .../realtime/RealtimeSegmentDataManager.java | 15 ++- .../RealtimeSegmentDataManagerTest.java | 105 ++++++++++++++++++ 2 files changed, 114 insertions(+), 6 deletions(-) diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java index 47695c74ec28..d550d4d9a0a8 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManager.java @@ -1861,26 +1861,29 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf _partitionGroupId = llcSegmentName.getPartitionGroupId(); List> streamConfigMaps = IngestionConfigUtils.getStreamConfigMaps(_tableConfig); int numStreams = streamConfigMaps.size(); + int topicId; if (numStreams == 1) { // Single stream // NOTE: We skip partition id translation logic to handle cases where custom stream might return partition id // larger than 10000. _streamPartitionId = _partitionGroupId; + topicId = 0; _streamConfig = new StreamConfig(_tableNameWithType, streamConfigMaps.get(0)); } else { // Multiple streams - _streamPartitionId = IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_partitionGroupId); - int index = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(_partitionGroupId); - Preconditions.checkState(numStreams > index, "Cannot find stream config of index: %s for table: %s", index, - _tableNameWithType); - _streamConfig = new StreamConfig(_tableNameWithType, streamConfigMaps.get(index)); + _streamPartitionId = llcSegmentName.getStreamPartitionGroupId(true); + topicId = llcSegmentName.getTopicId(true); + Preconditions.checkState(numStreams > topicId, "Cannot find stream config of topicId: %s for table: %s", + topicId, _tableNameWithType); + _streamConfig = new StreamConfig(_tableNameWithType, streamConfigMaps.get(topicId)); } _streamConsumerFactory = StreamConsumerFactoryProvider.create(_streamConfig); _streamPartitionMsgOffsetFactory = _streamConsumerFactory.createStreamMsgOffsetFactory(); String streamTopic = _streamConfig.getTopicName(); _segmentNameStr = _segmentZKMetadata.getSegmentName(); _partitionGroupConsumptionStatus = - new PartitionGroupConsumptionStatus(_partitionGroupId, _streamPartitionId, llcSegmentName.getSequenceNumber(), + new PartitionGroupConsumptionStatus(_partitionGroupId, _streamPartitionId, topicId, + llcSegmentName.getSequenceNumber(), _streamPartitionMsgOffsetFactory.create(_segmentZKMetadata.getStartOffset()), _segmentZKMetadata.getEndOffset() == null ? null : _streamPartitionMsgOffsetFactory.create(_segmentZKMetadata.getEndOffset()), diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java index 95a64108511e..f0391f4e23e5 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/RealtimeSegmentDataManagerTest.java @@ -67,10 +67,12 @@ import org.apache.pinot.spi.metrics.PinotMetricUtils; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.LongMsgOffsetFactory; +import org.apache.pinot.spi.stream.PartitionGroupConsumptionStatus; import org.apache.pinot.spi.stream.PermanentConsumerException; import org.apache.pinot.spi.stream.StreamConfigProperties; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; import org.apache.pinot.spi.utils.CommonConstants; +import org.apache.pinot.spi.utils.IngestionConfigUtils; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.testng.Assert; import org.testng.annotations.AfterClass; @@ -140,6 +142,109 @@ private FakeRealtimeSegmentDataManager createFakeSegmentManager() return createFakeSegmentManager(false, new TimeSupplier(), null, null, null); } + // Verifies that for a multi-stream table, the constructor decodes and threads the topicId (rather than silently + // defaulting it to 0) into the resulting PartitionGroupConsumptionStatus. + @Test + public void testMultiStreamConstructorSetsTopicId() + throws Exception { + TableConfig tableConfig = createTableConfig(); + Map streamConfigMap = tableConfig.getIndexingConfig().getStreamConfigs(); + IngestionConfig ingestionConfig = new IngestionConfig(); + ingestionConfig.setStreamIngestionConfig( + new StreamIngestionConfig(List.of(streamConfigMap, streamConfigMap))); + tableConfig.setIngestionConfig(ingestionConfig); + + // Topic 1, raw stream partition 1 -> padded pinot partition id 10001 + int rawStreamPartitionId = 1; + int topicId = 1; + int paddedPartitionGroupId = + IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(rawStreamPartitionId, topicId); + LLCSegmentName llcSegmentName = + new LLCSegmentName(RAW_TABLE_NAME, paddedPartitionGroupId, SEQUENCE_ID, SEG_TIME_MS); + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(llcSegmentName.getSegmentName()); + segmentZKMetadata.setStartOffset(START_OFFSET.toString()); + segmentZKMetadata.setCreationTime(System.currentTimeMillis()); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS); + + RealtimeTableDataManager tableDataManager = createTableDataManager(tableConfig); + _partitionGroupIdToConsumerCoordinatorMap.putIfAbsent(paddedPartitionGroupId, + new ConsumerCoordinator(false, tableDataManager)); + Schema schema = Fixtures.createSchema(); + ServerMetrics serverMetrics = new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()); + try (FakeRealtimeSegmentDataManager segmentDataManager = new FakeRealtimeSegmentDataManager(segmentZKMetadata, + tableConfig, tableDataManager, new File(TEMP_DIR, REALTIME_TABLE_NAME).getAbsolutePath(), schema, + llcSegmentName, _partitionGroupIdToConsumerCoordinatorMap, serverMetrics, new TimeSupplier())) { + Field partitionGroupConsumptionStatusField = + RealtimeSegmentDataManager.class.getDeclaredField("_partitionGroupConsumptionStatus"); + partitionGroupConsumptionStatusField.setAccessible(true); + PartitionGroupConsumptionStatus partitionGroupConsumptionStatus = + (PartitionGroupConsumptionStatus) partitionGroupConsumptionStatusField.get(segmentDataManager); + Assert.assertEquals(partitionGroupConsumptionStatus.getTopicId(), topicId); + Assert.assertEquals(partitionGroupConsumptionStatus.getStreamPartitionGroupId(), rawStreamPartitionId); + } + } + + // Verifies that for a single-stream table, the constructor sets topicId to 0. + @Test + public void testSingleStreamConstructorSetsTopicIdToZero() + throws Exception { + try (FakeRealtimeSegmentDataManager segmentDataManager = createFakeSegmentManager()) { + Field partitionGroupConsumptionStatusField = + RealtimeSegmentDataManager.class.getDeclaredField("_partitionGroupConsumptionStatus"); + partitionGroupConsumptionStatusField.setAccessible(true); + PartitionGroupConsumptionStatus partitionGroupConsumptionStatus = + (PartitionGroupConsumptionStatus) partitionGroupConsumptionStatusField.get(segmentDataManager); + Assert.assertEquals(partitionGroupConsumptionStatus.getTopicId(), 0); + } + } + + // Same as testMultiStreamConstructorSetsTopicId, but using the newer multi-topic-format segment name (topicId + // embedded directly in the segment name) instead of the legacy padded-partition-id encoding, to verify + // LLCSegmentName.getTopicId(true)/getStreamPartitionGroupId(true) are used (rather than always re-deriving from + // the padded id, which would be wrong for multi-topic-format segment names). + @Test + public void testMultiStreamConstructorSetsTopicIdMultiTopicFormat() + throws Exception { + TableConfig tableConfig = createTableConfig(); + Map streamConfigMap = tableConfig.getIndexingConfig().getStreamConfigs(); + IngestionConfig ingestionConfig = new IngestionConfig(); + ingestionConfig.setStreamIngestionConfig( + new StreamIngestionConfig(List.of(streamConfigMap, streamConfigMap))); + tableConfig.setIngestionConfig(ingestionConfig); + + // Deliberately "unaligned" values (streamPartitionId 1 for topic 3) that a padded-id-arithmetic decode would + // get wrong, to make sure the test would actually catch the regression. + int topicId = 1; + int rawStreamPartitionId = 3; + // Round-trip through the segment name string, since only the string-parsing constructor sets + // isMultiTopicFormat=true. + String segmentNameStr = + new LLCSegmentName(RAW_TABLE_NAME, topicId, rawStreamPartitionId, SEQUENCE_ID, SEG_TIME_MS).getSegmentName(); + LLCSegmentName llcSegmentName = new LLCSegmentName(segmentNameStr); + Assert.assertTrue(llcSegmentName.isMultiTopicFormat()); + SegmentZKMetadata segmentZKMetadata = new SegmentZKMetadata(segmentNameStr); + segmentZKMetadata.setStartOffset(START_OFFSET.toString()); + segmentZKMetadata.setCreationTime(System.currentTimeMillis()); + segmentZKMetadata.setStatus(CommonConstants.Segment.Realtime.Status.IN_PROGRESS); + + RealtimeTableDataManager tableDataManager = createTableDataManager(tableConfig); + _partitionGroupIdToConsumerCoordinatorMap.putIfAbsent(llcSegmentName.getPartitionGroupId(), + new ConsumerCoordinator(false, tableDataManager)); + Schema schema = Fixtures.createSchema(); + ServerMetrics serverMetrics = new ServerMetrics(PinotMetricUtils.getPinotMetricsRegistry()); + try (FakeRealtimeSegmentDataManager segmentDataManager = new FakeRealtimeSegmentDataManager(segmentZKMetadata, + tableConfig, tableDataManager, new File(TEMP_DIR, REALTIME_TABLE_NAME).getAbsolutePath(), schema, + llcSegmentName, _partitionGroupIdToConsumerCoordinatorMap, serverMetrics, new TimeSupplier())) { + Field partitionGroupConsumptionStatusField = + RealtimeSegmentDataManager.class.getDeclaredField("_partitionGroupConsumptionStatus"); + partitionGroupConsumptionStatusField.setAccessible(true); + PartitionGroupConsumptionStatus partitionGroupConsumptionStatus = + (PartitionGroupConsumptionStatus) partitionGroupConsumptionStatusField.get(segmentDataManager); + Assert.assertEquals(partitionGroupConsumptionStatus.getTopicId(), topicId); + Assert.assertEquals(partitionGroupConsumptionStatus.getStreamPartitionGroupId(), rawStreamPartitionId); + } + } + @Test public void testInitializationErrorStopMsgSkippedWhenDifferentSegmentManagerRegistered() throws Exception { From 91108d5c046822a954635c74b2c08c9afc61f89f Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 13:37:15 -0700 Subject: [PATCH 10/12] Unpad partition group ids in multi-topic metadata fetch fetchMultipleStreams() now sets PartitionGroupMetadata's topicId and partitionGroupId directly from the loop index and raw stream partition id, instead of encoding them into a padded composite id via IngestionConfigUtils. This avoids relying on the padded-id scheme, which the LLCSegmentName multi-topic format has already moved away from. --- .../stream/PartitionGroupMetadataFetcher.java | 9 +++----- .../PartitionGroupMetadataFetcherTest.java | 23 +++++++++++-------- 2 files changed, 17 insertions(+), 15 deletions(-) diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java index 6ef5cd000c43..e854f4f50b2a 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcher.java @@ -23,7 +23,6 @@ import java.util.List; import java.util.concurrent.Callable; import java.util.stream.Collectors; -import org.apache.pinot.spi.utils.IngestionConfigUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -130,8 +129,7 @@ private Boolean fetchMultipleStreams() int index = i; List topicPartitionGroupConsumptionStatusList = _partitionGroupConsumptionStatusList.stream() - .filter(partitionGroupConsumptionStatus -> IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId( - partitionGroupConsumptionStatus.getPartitionGroupId()) == index) + .filter(partitionGroupConsumptionStatus -> partitionGroupConsumptionStatus.getTopicId() == index) .collect(Collectors.toList()); try (StreamMetadataProvider streamMetadataProvider = streamConsumerFactory.createStreamMetadataProvider( StreamConsumerFactory.getUniqueClientId(clientId))) { @@ -141,9 +139,8 @@ private Boolean fetchMultipleStreams() /*maxWaitTimeMs=*/METADATA_FETCH_TIMEOUT_MS, _forceGetOffsetFromStream) .stream() - .map(metadata -> new PartitionGroupMetadata( - IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(metadata.getPartitionGroupId(), - index), index, metadata.getStartOffset(), metadata.getSequenceNumber())) + .map(metadata -> new PartitionGroupMetadata(metadata.getPartitionGroupId(), index, + metadata.getStartOffset(), metadata.getSequenceNumber())) .collect(Collectors.toList()); _streamMetadataList.add( new StreamMetadata(streamConfig, partitionGroupMetadataList.size(), partitionGroupMetadataList)); diff --git a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java index 860c2134b6cb..d4ecd6098bbe 100644 --- a/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java +++ b/pinot-spi/src/test/java/org/apache/pinot/spi/stream/PartitionGroupMetadataFetcherTest.java @@ -116,8 +116,10 @@ public void testFetchMultipleStreams() StreamConfig streamConfig2 = createMockStreamConfig("topic2", "test-table", false); List streamConfigs = Arrays.asList(streamConfig1, streamConfig2); - PartitionGroupConsumptionStatus status1 = new PartitionGroupConsumptionStatus(0, 0, null, null, "IN_PROGRESS"); - PartitionGroupConsumptionStatus status2 = new PartitionGroupConsumptionStatus(1, 1, null, null, "IN_PROGRESS"); + PartitionGroupConsumptionStatus status1 = + new PartitionGroupConsumptionStatus(0, 0, 0, 0, null, null, "IN_PROGRESS"); + PartitionGroupConsumptionStatus status2 = + new PartitionGroupConsumptionStatus(1, 1, 1, 0, null, null, "IN_PROGRESS"); List statusList = Arrays.asList(status1, status2); PartitionGroupMetadata mockedMetadata1 = new PartitionGroupMetadata(0, mock(StreamPartitionMsgOffset.class)); @@ -152,14 +154,14 @@ public void testFetchMultipleStreams() Assert.assertEquals(streamMetadataList.get(1).getNumPartitions(), 2); Assert.assertEquals(streamMetadataList.get(1).getPartitionGroupMetadataList().size(), 2); - // Verify the correct partition group IDs: 0, 1, 10000, 10001 + // Verify the partition group IDs are the raw, unpadded per-topic stream partition ids: 0, 0, 1, 1 List partitionIds = streamMetadataList.stream() .flatMap(sm -> sm.getPartitionGroupMetadataList().stream()) .map(PartitionGroupMetadata::getPartitionGroupId) .sorted() .collect(Collectors.toList()); - Assert.assertEquals(partitionIds, Arrays.asList(0, 1, 10000, 10001)); + Assert.assertEquals(partitionIds, Arrays.asList(0, 0, 1, 1)); // Verify the topic ids are correctly assigned: topic1 (index 0) -> 0, 0; topic2 (index 1) -> 1, 1 List topicIds = streamMetadataList.stream() @@ -213,14 +215,14 @@ public void testFetchMultipleStreamsWithPause() Assert.assertEquals(streamMetadataList.size(), 2); Assert.assertNull(fetcher.getException()); - // Verify the correct partition group IDs + // Verify the partition group IDs are the raw, unpadded per-topic stream partition ids List partitionIds = streamMetadataList.stream() .flatMap(sm -> sm.getPartitionGroupMetadataList().stream()) .map(PartitionGroupMetadata::getPartitionGroupId) .sorted() .collect(Collectors.toList()); - Assert.assertEquals(partitionIds, Arrays.asList(0, 1, 20000, 20001)); + Assert.assertEquals(partitionIds, Arrays.asList(0, 0, 1, 1)); } } @@ -332,11 +334,14 @@ public void testSequenceNumberPreservedInMultiStreamRemap() List streamMetadataList = fetcher.getStreamMetadataList(); Assert.assertEquals(streamMetadataList.size(), 2); - // Second stream's partitions should have remapped IDs but preserved sequence numbers + // Second stream's partitions should preserve raw (unpadded) partition ids and sequence numbers, + // with topicId set to the stream's index List stream1Partitions = streamMetadataList.get(1).getPartitionGroupMetadataList(); - Assert.assertEquals(stream1Partitions.get(0).getPartitionGroupId(), 10000); + Assert.assertEquals(stream1Partitions.get(0).getPartitionGroupId(), 0); + Assert.assertEquals(stream1Partitions.get(0).getTopicId(), 1); Assert.assertEquals(stream1Partitions.get(0).getSequenceNumber(), 7); - Assert.assertEquals(stream1Partitions.get(1).getPartitionGroupId(), 10001); + Assert.assertEquals(stream1Partitions.get(1).getPartitionGroupId(), 1); + Assert.assertEquals(stream1Partitions.get(1).getTopicId(), 1); Assert.assertEquals(stream1Partitions.get(1).getSequenceNumber(), 3); } } From 14b6535fbadc19e8c66707477014c86459bd77dc Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 13:37:38 -0700 Subject: [PATCH 11/12] Replace padded partition id decoding with LLCSegmentName accessors Several call sites in PinotLLCRealtimeSegmentManager decoded topic and stream partition ids out of a padded composite pinotPartitionId via IngestionConfigUtils, even though a richer LLCSegmentName or PartitionGroupMetadata object was already available with direct, unpadded accessors. Switch those sites to use LLCSegmentName#getTopicId/getStreamPartitionGroupId and PartitionGroupMetadata#getTopicId/getPartitionGroupId instead. Leaves untouched the two call sites that legitimately need IngestionConfigUtils: decoding a genuinely old-format (padded) LLCSegmentName in getNextLLCSegmentName, and getPartitionMetadataFromTableConfig, which only has a raw partition id and TableConfig, not a richer object. --- .../realtime/PinotLLCRealtimeSegmentManager.java | 16 +++++----------- .../PinotLLCRealtimeSegmentManagerTest.java | 15 ++++++--------- 2 files changed, 11 insertions(+), 20 deletions(-) diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java index 8e0999db5e9a..c11b07f236ab 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManager.java @@ -1305,14 +1305,10 @@ PartitionIdsWithIdealState getPartitionIdsWithIdealState(List stre getPartitionGroupConsumptionStatusList(idealState, streamConfigs); List streamMetadataList = getNewStreamMetadataList(streamConfigs, currentPartitionGroupConsumptionStatusList, idealState); - boolean isMultiTopic = numStreams > 1; for (StreamMetadata streamMetadata : streamMetadataList) { for (PartitionGroupMetadata partitionGroupMetadata : streamMetadata.getPartitionGroupMetadataList()) { - int pinotPartitionId = partitionGroupMetadata.getPartitionGroupId(); - int topicId = isMultiTopic - ? IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(pinotPartitionId) : 0; - int streamPartitionId = isMultiTopic - ? IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(pinotPartitionId) : pinotPartitionId; + int topicId = partitionGroupMetadata.getTopicId(); + int streamPartitionId = partitionGroupMetadata.getPartitionGroupId(); partitionIds.computeIfAbsent(topicId, k -> new HashSet<>()).add(streamPartitionId); } } @@ -1592,8 +1588,7 @@ public static boolean isTopicPaused(IdealState idealState, int topicIndex) { public static boolean isTopicPaused(IdealState idealState, String segmentName) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { - return isTopicPaused(idealState, - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(llcSegmentName.getPartitionGroupId())); + return isTopicPaused(idealState, llcSegmentName.getTopicId(true)); } return false; } @@ -1841,7 +1836,7 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List instanceStateMap = instanceStatesMap.get(latestSegmentName); if (instanceStateMap != null) { @@ -2935,8 +2930,7 @@ private Set findConsumingSegmentsOfTopics(IdealState idealState, List consumingSegments = new TreeSet<>(); idealState.getRecord().getMapFields().forEach((segmentName, instanceToStateMap) -> { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); - if (llcSegmentName != null && !topicIndices.contains( - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(llcSegmentName.getPartitionGroupId()))) { + if (llcSegmentName != null && !topicIndices.contains(llcSegmentName.getTopicId(true))) { return; } for (String state : instanceToStateMap.values()) { diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java index a49b5f00f00e..1196126dd8cc 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/PinotLLCRealtimeSegmentManagerTest.java @@ -2010,20 +2010,17 @@ public void testGetPartitionIdsWithIdealStateMultiTopicFallback() doThrow(new RuntimeException()).when(segmentManagerSpy).getPartitionIds(any(StreamConfig.class)); doReturn(List.of()).when(segmentManagerSpy).getPartitionGroupConsumptionStatusList(idealState, streamConfigs); - // Topic 0, raw partitions 0 and 3 -> padded pinot partition ids 0 and 3 (two groups for the same topic, to - // verify they accumulate into one Set instead of the second overwriting the first) - // Topic 1, raw partition 5 -> padded pinot partition id 10005 + // Topic 0, raw partitions 0 and 3 (two groups for the same topic, to verify they accumulate into one Set + // instead of the second overwriting the first) + // Topic 1, raw partition 5 // topicId and streamPartitionId are chosen to differ per entry so a transposed topicId/streamPartitionId // formula would produce a visibly wrong result rather than accidentally matching. List streamMetadataList = List.of( new StreamMetadata(streamConfigs.get(0), 2, - List.of(new PartitionGroupMetadata( - IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(0, 0), new LongMsgOffset(1)), - new PartitionGroupMetadata( - IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(3, 0), new LongMsgOffset(1)))), + List.of(new PartitionGroupMetadata(0, 0, new LongMsgOffset(1), 0), + new PartitionGroupMetadata(3, 0, new LongMsgOffset(1), 0))), new StreamMetadata(streamConfigs.get(1), 2, - List.of(new PartitionGroupMetadata( - IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(5, 1), new LongMsgOffset(1))))); + List.of(new PartitionGroupMetadata(5, 1, new LongMsgOffset(1), 0)))); doReturn(streamMetadataList).when(segmentManagerSpy) .getNewStreamMetadataList(eq(streamConfigs), any(), eq(idealState)); From 96ba3703ccc5a5b843dd1d329316ae9b6a815909 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Tue, 14 Jul 2026 14:52:39 -0700 Subject: [PATCH 12/12] Add topicId to WatermarkInductionResult.Watermark for copy-table Watermark previously had no topicId field, forcing getStreamMetadataList to decode the stream config index from a padded composite partitionGroupId that is no longer guaranteed to be padded. Add an explicit topicId field and thread it through getConsumerWatermarks/getStreamMetadataList, using the decoded stream partition id (not the raw, possibly still-padded one) when constructing watermarks from PartitionGroupConsumptionStatus. --- .../pinot/common/utils/LLCSegmentName.java | 8 +-- .../resources/PinotTableRestletResource.java | 9 ++- .../helix/core/PinotHelixResourceManager.java | 3 +- .../helix/core/WatermarkInductionResult.java | 37 ++++++++++-- .../PinotTableRestletResourceTest.java | 8 +-- ...inotHelixResourceManagerStatelessTest.java | 22 +++++++ .../core/WatermarkInductionResultTest.java | 60 +++++++++++++++++++ 7 files changed, 129 insertions(+), 18 deletions(-) create mode 100644 pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/WatermarkInductionResultTest.java diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java index 985de7ecb3fa..180fa871a324 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java @@ -153,8 +153,8 @@ public int getTopicId() { public int getTopicId(boolean hasMultiTopic) { if (hasMultiTopic) { - return _isMultiTopicFormat ? _topicId : - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(_partitionGroupId); + return _isMultiTopicFormat + ? _topicId : IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(_partitionGroupId); } else { return _topicId; } @@ -188,8 +188,8 @@ public String getSegmentName() { public int getStreamPartitionGroupId(boolean hasMultiTopics) { if (hasMultiTopics) { - return _isMultiTopicFormat ? _partitionGroupId : - IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_partitionGroupId); + return _isMultiTopicFormat + ? _partitionGroupId : IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_partitionGroupId); } else { return _partitionGroupId; } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java index d723fae45f5d..4105e57d837a 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java @@ -393,14 +393,13 @@ List getStreamMetadataList(List streamConfigs, _pinotHelixResourceManager.getRealtimeSegmentManager().getPartitionCountMap(streamConfigs); Map> partitionGroupMetadataByStreamConfigIndex = new HashMap<>(); for (WatermarkInductionResult.Watermark watermark : watermarkInductionResult.getWatermarks()) { - int streamConfigIndex = - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(watermark.getPartitionGroupId()); + int streamConfigIndex = watermark.getTopicId(); Preconditions.checkArgument(streamConfigIndex >= 0 && streamConfigIndex < streamConfigs.size(), - "Invalid stream config index %s from watermark partition ID %s. Expected index in range [0, %s)", + "Invalid stream config index %s from watermark (partitionGroupId=%s). Expected index in range [0, %s)", streamConfigIndex, watermark.getPartitionGroupId(), streamConfigs.size()); partitionGroupMetadataByStreamConfigIndex.computeIfAbsent(streamConfigIndex, ignored -> new ArrayList<>()).add( - new PartitionGroupMetadata(watermark.getPartitionGroupId(), new LongMsgOffset(watermark.getOffset()), - watermark.getSequenceNumber())); + new PartitionGroupMetadata(watermark.getPartitionGroupId(), streamConfigIndex, + new LongMsgOffset(watermark.getOffset()), watermark.getSequenceNumber())); } // Iterate in order by streamConfigIndex to ensure deterministic ordering diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java index 605ade8a2d3b..1a3d68f52a46 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java @@ -5460,7 +5460,8 @@ public WatermarkInductionResult getConsumerWatermarks(String tableName) } catch (NumericException e) { throw new RuntimeException(e); } - return new WatermarkInductionResult.Watermark(status.getPartitionGroupId(), seq, startOffset); + return new WatermarkInductionResult.Watermark(status.getStreamPartitionGroupId(), status.getTopicId(), seq, + startOffset); }).collect(Collectors.toList()); return new WatermarkInductionResult(watermarks); } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/WatermarkInductionResult.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/WatermarkInductionResult.java index 4d6a4aa8da0c..93b58c8dee0e 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/WatermarkInductionResult.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/WatermarkInductionResult.java @@ -57,6 +57,7 @@ public List getWatermarks() { */ public static class Watermark { private int _partitionGroupId; + private int _topicId; private int _sequenceNumber; private long _offset; @@ -65,28 +66,56 @@ public static class Watermark { * a WaterMark instance from a JSON object. The @JsonProperty annotations * map the keys in the JSON object to the constructor parameters. * - * @param partitionGroupId The ID of the partition group. + * @param partitionGroupId The stream partition id of the partition group, decoded from any legacy + * padded/composite id. Defaults topicId to 0 (single-topic tables). + * @param sequenceNumber The segment sequence number of the consuming segment. + * @param offset The first Kafka offset whose corresponding record has not yet sealed in Pinot + */ + public Watermark(int partitionGroupId, int sequenceNumber, long offset) { + this(partitionGroupId, 0, sequenceNumber, offset); + } + + /** + * The @JsonCreator annotation tells Jackson to use this constructor to create + * a WaterMark instance from a JSON object. The @JsonProperty annotations + * map the keys in the JSON object to the constructor parameters. + * + * @param partitionGroupId The stream partition id of the partition group (i.e. the partition id within its + * topic/stream config), not the raw/composite id used by legacy padded encodings. + * @param topicId The id of the topic (stream config) this partition group belongs to, for tables with + * multiple topics. * @param sequenceNumber The segment sequence number of the consuming segment. * @param offset The first Kafka offset whose corresponding record has not yet sealed in Pinot */ @JsonCreator - public Watermark(@JsonProperty("partitionGroupId") int partitionGroupId, + public Watermark(@JsonProperty("partitionGroupId") int partitionGroupId, @JsonProperty("topicId") int topicId, @JsonProperty("sequenceNumber") int sequenceNumber, @JsonProperty("offset") long offset) { _partitionGroupId = partitionGroupId; + _topicId = topicId; _sequenceNumber = sequenceNumber; _offset = offset; } /** - * Gets the partition group ID. + * Gets the stream partition id of the partition group. * - * @return The partition group ID. + * @return The stream partition id. */ @JsonGetter("partitionGroupId") public int getPartitionGroupId() { return _partitionGroupId; } + /** + * Gets the topic ID (stream config index) this partition group belongs to. + * + * @return The topic ID. + */ + @JsonGetter("topicId") + public int getTopicId() { + return _topicId; + } + /** * Gets the segment sequence number of the most recent consuming segment. * diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTableRestletResourceTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTableRestletResourceTest.java index dde748ef6221..afd9ac552cab 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTableRestletResourceTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotTableRestletResourceTest.java @@ -71,9 +71,9 @@ public void testGetStreamMetadataList() List streamMetadataList = resource.getStreamMetadataList(List.of(streamConfig0, streamConfig1), new WatermarkInductionResult(List.of( - new WatermarkInductionResult.Watermark(1, 3, 101L), - new WatermarkInductionResult.Watermark(0, 2, 100L), - new WatermarkInductionResult.Watermark(10000, 5, 200L)))); + new WatermarkInductionResult.Watermark(1, 0, 3, 101L), + new WatermarkInductionResult.Watermark(0, 0, 2, 100L), + new WatermarkInductionResult.Watermark(0, 1, 5, 200L)))); assertEquals(streamMetadataList.size(), 2); @@ -95,7 +95,7 @@ public void testGetStreamMetadataList() assertEquals(streamMetadata1.getStreamConfig(), streamConfig1); assertEquals(streamMetadata1.getNumPartitions(), 8); assertEquals(streamMetadata1.getPartitionGroupMetadataList().size(), 1); - assertEquals(streamMetadata1.getPartitionGroupMetadataList().get(0).getPartitionGroupId(), 10000); + assertEquals(streamMetadata1.getPartitionGroupMetadataList().get(0).getPartitionGroupId(), 0); assertEquals(((LongMsgOffset) streamMetadata1.getPartitionGroupMetadataList().get(0).getStartOffset()).getOffset(), 200L); assertEquals(streamMetadata1.getPartitionGroupMetadataList().get(0).getSequenceNumber(), 5); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java index 9e90ef8f2a2e..da74276c81a0 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManagerStatelessTest.java @@ -1769,6 +1769,28 @@ public void testGetConsumerWatermarks() assertEquals(inProgressWatermark.getSequenceNumber(), 200L); assertEquals(inProgressWatermark.getOffset(), 789L); + // Test legacy-format multi-topic segments, where the raw partitionGroupId is still a padded composite id + // (partitionGroupId != streamPartitionId): the watermark must use the decoded streamPartitionId, not the + // raw padded partitionGroupId. Covers both the IN_PROGRESS and DONE branches. + PartitionGroupConsumptionStatus legacyFormatInProgressStatus = new PartitionGroupConsumptionStatus(10001, 1, 1, + 300, new LongMsgOffset(111), null, "IN_PROGRESS"); + PartitionGroupConsumptionStatus legacyFormatDoneStatus = new PartitionGroupConsumptionStatus(20002, 2, 2, 400, + new LongMsgOffset(222), new LongMsgOffset(333), "done"); + when(mockSegmentManager.getPartitionGroupConsumptionStatusList(any(), any())) + .thenReturn(List.of(legacyFormatInProgressStatus, legacyFormatDoneStatus)); + WatermarkInductionResult legacyFormatResult = _helixResourceManager.getConsumerWatermarks(rawTableName); + assertEquals(legacyFormatResult.getWatermarks().size(), 2); + WatermarkInductionResult.Watermark legacyFormatInProgressWatermark = legacyFormatResult.getWatermarks().get(0); + assertEquals(legacyFormatInProgressWatermark.getPartitionGroupId(), 1); + assertEquals(legacyFormatInProgressWatermark.getTopicId(), 1); + assertEquals(legacyFormatInProgressWatermark.getSequenceNumber(), 300L); + assertEquals(legacyFormatInProgressWatermark.getOffset(), 111L); + WatermarkInductionResult.Watermark legacyFormatDoneWatermark = legacyFormatResult.getWatermarks().get(1); + assertEquals(legacyFormatDoneWatermark.getPartitionGroupId(), 2); + assertEquals(legacyFormatDoneWatermark.getTopicId(), 2); + assertEquals(legacyFormatDoneWatermark.getSequenceNumber(), 401L); + assertEquals(legacyFormatDoneWatermark.getOffset(), 333L); + // recover the original values helixAdminField.set(_helixResourceManager, originalHelixAdmin); llcManagerField.set(_helixResourceManager, originalLlcManager); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/WatermarkInductionResultTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/WatermarkInductionResultTest.java new file mode 100644 index 000000000000..c61cb8deb9a5 --- /dev/null +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/WatermarkInductionResultTest.java @@ -0,0 +1,60 @@ +/** + * 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.controller.helix.core; + +import java.util.List; +import org.apache.pinot.spi.utils.JsonUtils; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; + + +public class WatermarkInductionResultTest { + + @Test + public void testTopicIdRoundTrip() + throws Exception { + WatermarkInductionResult result = + new WatermarkInductionResult(List.of(new WatermarkInductionResult.Watermark(1, 2, 3, 100L))); + String json = JsonUtils.objectToString(result); + WatermarkInductionResult deserialized = JsonUtils.stringToObject(json, WatermarkInductionResult.class); + + assertEquals(deserialized.getWatermarks().size(), 1); + WatermarkInductionResult.Watermark watermark = deserialized.getWatermarks().get(0); + assertEquals(watermark.getPartitionGroupId(), 1); + assertEquals(watermark.getTopicId(), 2); + assertEquals(watermark.getSequenceNumber(), 3); + assertEquals(watermark.getOffset(), 100L); + } + + @Test + public void testTopicIdDefaultsToZeroWhenAbsentFromJson() + throws Exception { + // Simulates a pre-migration payload (e.g. from an older controller) that has no topicId field. + String legacyJson = "{\"watermarks\":[{\"partitionGroupId\":5,\"sequenceNumber\":6,\"offset\":200}]}"; + WatermarkInductionResult deserialized = JsonUtils.stringToObject(legacyJson, WatermarkInductionResult.class); + + assertEquals(deserialized.getWatermarks().size(), 1); + WatermarkInductionResult.Watermark watermark = deserialized.getWatermarks().get(0); + assertEquals(watermark.getPartitionGroupId(), 5); + assertEquals(watermark.getTopicId(), 0); + assertEquals(watermark.getSequenceNumber(), 6); + assertEquals(watermark.getOffset(), 200L); + } +}