From 5abcff00d2c102e1cdf310f80ba8090ef5ea29f2 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Thu, 2 Jul 2026 14:38:07 -0400 Subject: [PATCH 1/2] Introduce TopicPartitionId composite key for LLC segment partition identity (Part 1 of multi-topic segment naming) This is part 1 of the multi-topic segment name format work. It introduces TopicPartitionId as a composite key holding topicId and partitionId, and adds getTopicPartitionId() to LLCSegmentName returning it. The original getPartitionGroupId() returning int is preserved as deprecated for binary compatibility during rolling upgrades. All call sites are migrated to use getTopicPartitionId().getPartitionId(). --- .../pinot/common/utils/LLCSegmentName.java | 18 ++-- .../pinot/common/utils/SegmentUtils.java | 2 +- .../pinot/common/utils/TopicPartitionId.java | 81 ++++++++++++++ .../common/utils/LLCSegmentNameTest.java | 6 +- .../common/utils/TopicPartitionIdTest.java | 100 ++++++++++++++++++ .../PinotSegmentRestletResource.java | 4 +- .../controller/api/resources/TableViews.java | 3 +- .../helix/core/PinotHelixResourceManager.java | 2 +- .../BaseStrictRealtimeSegmentAssignment.java | 2 +- ...tiTierStrictRealtimeSegmentAssignment.java | 2 +- .../MissingConsumingSegmentFinder.java | 2 +- .../PinotLLCRealtimeSegmentManager.java | 54 ++++++---- .../helix/PinotResourceManagerTest.java | 2 +- ...inotHelixResourceManagerStatelessTest.java | 2 +- ...timeReplicaGroupSegmentAssignmentTest.java | 2 +- .../PinotLLCRealtimeSegmentManagerTest.java | 14 +-- .../core/rebalance/TableRebalancerTest.java | 2 +- .../manager/realtime/ConsumerCoordinator.java | 8 +- .../realtime/IngestionDelayTracker.java | 5 +- .../realtime/RealtimeSegmentDataManager.java | 2 +- .../realtime/RealtimeTableDataManager.java | 12 ++- .../RealtimeSegmentDataManagerTest.java | 2 +- .../tests/BaseClusterIntegrationTest.java | 2 +- .../LLCRealtimeClusterIntegrationTest.java | 4 +- ...CRealtimeKafka3ClusterIntegrationTest.java | 2 +- ...CRealtimeKafka4ClusterIntegrationTest.java | 2 +- ...alUpsertTableRebalanceIntegrationTest.java | 2 +- ...timeIngestionSegmentCommitFailureTest.java | 3 +- ...tionLLCRealtimeClusterIntegrationTest.java | 9 +- ...ertTableSegmentPreloadIntegrationTest.java | 2 +- ...sertTableSegmentUploadIntegrationTest.java | 2 +- ...tiTopicRealtimeClusterIntegrationTest.java | 4 +- .../utils/PauselessRealtimeTestUtils.java | 3 +- .../mergerollup/MergeRollupTaskGenerator.java | 15 +-- .../tasks/purge/PurgeTaskGenerator.java | 17 +-- ...ealtimeToOfflineSegmentsTaskGenerator.java | 2 +- .../rules/LeafStageWorkerAssignmentRule.java | 2 +- .../MultiColumnRealtimeLuceneTextIndex.java | 2 +- .../RealtimeLuceneTextIndex.java | 2 +- .../StatelessRealtimeSegmentWriter.java | 2 +- 40 files changed, 306 insertions(+), 98 deletions(-) create mode 100644 pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java create mode 100644 pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.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 fdff73c0318d..20d9ac61e03b 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,7 +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 _partitionGroupId; + private final TopicPartitionId _topicPartitionId; private final int _sequenceNumber; private final String _creationTime; private final String _segmentName; @@ -43,7 +43,7 @@ public LLCSegmentName(String segmentName) { String[] parts = StringUtils.splitByWholeSeparator(segmentName, SEPARATOR); Preconditions.checkArgument(parts.length == 4, "Invalid LLC segment name: %s", segmentName); _tableName = parts[0]; - _partitionGroupId = Integer.parseInt(parts[1]); + _topicPartitionId = new TopicPartitionId(Integer.parseInt(parts[1])); _sequenceNumber = Integer.parseInt(parts[2]); _creationTime = parts[3]; _segmentName = segmentName; @@ -52,7 +52,7 @@ public LLCSegmentName(String segmentName) { public LLCSegmentName(String tableName, int partitionGroupId, int sequenceNumber, long msSinceEpoch) { Preconditions.checkArgument(!tableName.contains(SEPARATOR), "Illegal table name: %s", tableName); _tableName = tableName; - _partitionGroupId = partitionGroupId; + _topicPartitionId = new TopicPartitionId(partitionGroupId); _sequenceNumber = sequenceNumber; // ISO8601 date: 20160120T1234Z _creationTime = DATE_FORMATTER.print(msSinceEpoch); @@ -101,8 +101,13 @@ public String getTableName() { return _tableName; } + public TopicPartitionId getTopicPartitionId() { + return _topicPartitionId; + } + + @Deprecated public int getPartitionGroupId() { - return _partitionGroupId; + return _topicPartitionId.getPartitionId(); } public int getSequenceNumber() { @@ -127,8 +132,9 @@ 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()); - if (_partitionGroupId != other._partitionGroupId) { - return Integer.compare(_partitionGroupId, other._partitionGroupId); + int cmp = _topicPartitionId.compareTo(other._topicPartitionId); + if (cmp != 0) { + return cmp; } return Integer.compare(_sequenceNumber, other._sequenceNumber); } diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java index 80e735fb2f80..fc9187338df6 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java @@ -72,7 +72,7 @@ public static Integer getSegmentPartitionId(SegmentZKMetadata segmentZKMetadata, public static Integer getPartitionIdFromSegmentName(String segmentName) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { - return llcSegmentName.getPartitionGroupId(); + return llcSegmentName.getTopicPartitionId().getPartitionId(); } UploadedRealtimeSegmentName uploadedRealtimeSegmentName = UploadedRealtimeSegmentName.of(segmentName); if (uploadedRealtimeSegmentName != null) { diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java new file mode 100644 index 000000000000..b2918881863f --- /dev/null +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java @@ -0,0 +1,81 @@ +/** + * 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.common.utils; + +import java.util.Objects; + + +/** + * Identifies a partition within a (possibly multi-topic) real-time table. + * + *

For single-topic tables, {@code topicId} is 0 and {@code partitionId} is the stream partition number. + * For multi-topic tables, {@code topicId} distinguishes the topic and {@code partitionId} is the stream partition + * number within that topic. + * + *

Thread-safe: instances are immutable. + */ +public final class TopicPartitionId implements Comparable { + private final int _topicId; + private final int _partitionId; + + public TopicPartitionId(int topicId, int partitionId) { + _topicId = topicId; + _partitionId = partitionId; + } + + public TopicPartitionId(int partitionId) { + this(0, partitionId); + } + + public int getTopicId() { + return _topicId; + } + + public int getPartitionId() { + return _partitionId; + } + + @Override + public int compareTo(TopicPartitionId other) { + int cmp = Integer.compare(_topicId, other._topicId); + return cmp != 0 ? cmp : Integer.compare(_partitionId, other._partitionId); + } + + @Override + public boolean equals(Object o) { + if (this == o) { + return true; + } + if (!(o instanceof TopicPartitionId)) { + return false; + } + TopicPartitionId that = (TopicPartitionId) o; + return _topicId == that._topicId && _partitionId == that._partitionId; + } + + @Override + public int hashCode() { + return Objects.hash(_topicId, _partitionId); + } + + @Override + public String toString() { + return _topicId + ":" + _partitionId; + } +} 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..30323ecc0897 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 @@ -39,7 +39,7 @@ public void testSegmentNameBuilder() { assertEquals(segmentName, "myTable__0__1__20160609T2142Z"); assertTrue(LLCSegmentName.isLLCSegment(segmentName)); assertEquals(llcSegmentName.getTableName(), "myTable"); - assertEquals(llcSegmentName.getPartitionGroupId(), 0); + assertEquals(llcSegmentName.getTopicPartitionId().getPartitionId(), 0); assertEquals(llcSegmentName.getSequenceNumber(), 1); // Invalid segment name @@ -58,7 +58,7 @@ public void testLLCSegmentName() { LLCSegmentName segName1 = new LLCSegmentName(tableName, partitionGroupId, sequenceNumber, msSinceEpoch); Assert.assertEquals(segName1.getSegmentName(), segmentName); - Assert.assertEquals(segName1.getPartitionGroupId(), partitionGroupId); + Assert.assertEquals(segName1.getTopicPartitionId().getPartitionId(), partitionGroupId); Assert.assertEquals(segName1.getCreationTime(), creationTime); Assert.assertEquals(segName1.getCreationTimeMs(), creationTimeInMs); Assert.assertEquals(segName1.getSequenceNumber(), sequenceNumber); @@ -66,7 +66,7 @@ public void testLLCSegmentName() { LLCSegmentName segName2 = new LLCSegmentName(segmentName); Assert.assertEquals(segName2.getSegmentName(), segmentName); - Assert.assertEquals(segName2.getPartitionGroupId(), partitionGroupId); + Assert.assertEquals(segName2.getTopicPartitionId().getPartitionId(), partitionGroupId); Assert.assertEquals(segName2.getCreationTime(), creationTime); Assert.assertEquals(segName2.getCreationTimeMs(), creationTimeInMs); Assert.assertEquals(segName2.getSequenceNumber(), sequenceNumber); diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java new file mode 100644 index 000000000000..fca10f36452e --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java @@ -0,0 +1,100 @@ +/** + * 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.common.utils; + +import java.util.HashMap; +import java.util.Map; +import org.testng.annotations.Test; + +import static org.testng.Assert.assertEquals; +import static org.testng.Assert.assertFalse; +import static org.testng.Assert.assertNotEquals; +import static org.testng.Assert.assertTrue; + + +public class TopicPartitionIdTest { + + @Test + public void testSingleArgConstructor() { + TopicPartitionId id = new TopicPartitionId(5); + assertEquals(id.getTopicId(), 0); + assertEquals(id.getPartitionId(), 5); + } + + @Test + public void testTwoArgConstructor() { + TopicPartitionId id = new TopicPartitionId(2, 7); + assertEquals(id.getTopicId(), 2); + assertEquals(id.getPartitionId(), 7); + } + + @Test + public void testEquality() { + TopicPartitionId id1 = new TopicPartitionId(1, 3); + TopicPartitionId id2 = new TopicPartitionId(1, 3); + TopicPartitionId id3 = new TopicPartitionId(1, 4); + TopicPartitionId id4 = new TopicPartitionId(2, 3); + + assertEquals(id1, id2); + assertEquals(id1.hashCode(), id2.hashCode()); + assertNotEquals(id1, id3); + assertNotEquals(id1, id4); + } + + @Test + public void testSingleArgEqualsZeroTopic() { + TopicPartitionId id1 = new TopicPartitionId(5); + TopicPartitionId id2 = new TopicPartitionId(0, 5); + assertEquals(id1, id2); + assertEquals(id1.hashCode(), id2.hashCode()); + } + + @Test + public void testCompareTo() { + TopicPartitionId a = new TopicPartitionId(0, 1); + TopicPartitionId b = new TopicPartitionId(0, 2); + TopicPartitionId c = new TopicPartitionId(1, 0); + TopicPartitionId d = new TopicPartitionId(1, 1); + + assertTrue(a.compareTo(b) < 0); + assertTrue(b.compareTo(a) > 0); + assertTrue(a.compareTo(c) < 0); + assertTrue(c.compareTo(d) < 0); + assertEquals(a.compareTo(new TopicPartitionId(0, 1)), 0); + } + + @Test + public void testAsMapKey() { + Map map = new HashMap<>(); + TopicPartitionId id1 = new TopicPartitionId(1, 3); + map.put(id1, "value1"); + + TopicPartitionId id2 = new TopicPartitionId(1, 3); + assertEquals(map.get(id2), "value1"); + + assertFalse(map.containsKey(new TopicPartitionId(1, 4))); + assertFalse(map.containsKey(new TopicPartitionId(2, 3))); + } + + @Test + public void testToString() { + TopicPartitionId id = new TopicPartitionId(2, 7); + assertEquals(id.toString(), "2:7"); + } +} diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java index cbda557d411a..8b00467dcd75 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java @@ -1124,7 +1124,7 @@ Map> getPartitionIdToSegmentsToDeleteMap( LOGGER.info("Skip segment: {} not in low-level consumer format", segmentName); continue; } - int partitionId = llcSegmentName.getPartitionGroupId(); + int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); LLCSegmentName oldestSegment = partitionToOldestSegment.get(partitionId); if (oldestSegment == null) { @@ -1163,7 +1163,7 @@ Map getPartitionIDToOldestSegment(List segments LOGGER.warn("Segment: {} is not present in the ideal state", segment); continue; } - int partitionId = llcSegmentName.getPartitionGroupId(); + int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); LLCSegmentName currentOldest = partitionToOldestSegment.get(partitionId); if (currentOldest == null || llcSegmentName.getSequenceNumber() < currentOldest.getSequenceNumber()) { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java index 02855d88fd92..35bad09dd210 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java @@ -217,7 +217,8 @@ public String getBadLLCSegmentsPerPartition( if (llcSegmentName == null) { continue; } - partitionIdToSegments.computeIfAbsent(llcSegmentName.getPartitionGroupId(), k -> new TreeSet<>()) + partitionIdToSegments.computeIfAbsent( + llcSegmentName.getTopicPartitionId().getPartitionId(), k -> new TreeSet<>()) .add(llcSegmentName); } 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..0da6670c2135 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 @@ -1164,7 +1164,7 @@ public Collection getLastLLCCompletedSegments(List getExistingAssignment(int partitionId, Map getExistingAssignment(int partitionId, Map latestLLCSegmentName.getSequenceNumber())) { latestLLCSegmentName = llcSegmentName; latestAssignment = entry.getValue().keySet(); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java index 9c17cc127dc4..d2d00130ae62 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java @@ -187,7 +187,7 @@ private void updateMaxDurationInfo(MissingSegmentInfo missingSegmentInfo, Intege private void updateMap(Map partitionGroupIdToLatestSegmentMap, LLCSegmentName llcSegmentName) { - int partitionGroupId = llcSegmentName.getPartitionGroupId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); partitionGroupIdToLatestSegmentMap.compute(partitionGroupId, (pid, existingSegment) -> { if (existingSegment == null) { return llcSegmentName; 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 9bc0d4e8a70e..118735f8249a 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 @@ -280,7 +280,7 @@ public List getPartitionGroupConsumptionStatusL if (llcSegmentName == null) { continue; } - int partitionGroupId = llcSegmentName.getPartitionGroupId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); partitionGroupIdToLatestSegment.compute(partitionGroupId, (k, latestSegment) -> { if (latestSegment == null) { return llcSegmentName; @@ -837,7 +837,7 @@ private String createNewSegmentMetadata(TableConfig tableConfig, String newConsumingSegmentName = null; LLCSegmentName committingLLCSegment = new LLCSegmentName(committingSegmentName); - int committingSegmentPartitionGroupId = committingLLCSegment.getPartitionGroupId(); + int committingSegmentPartitionGroupId = committingLLCSegment.getTopicPartitionId().getPartitionId(); List streamConfigs = IngestionConfigUtils.getStreamConfigs(tableConfig); PartitionIdsWithIdealState partitionIdsWithIdealState = getPartitionIdsWithIdealState(streamConfigs, @@ -1007,7 +1007,8 @@ private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig st // Handle offset auto reset String nextOffset = committingSegmentDescriptor.getNextOffset(); - String startOffset = computeStartOffset(nextOffset, streamConfig, newLLCSegmentName.getPartitionGroupId()); + String startOffset = computeStartOffset(nextOffset, streamConfig, + newLLCSegmentName.getTopicPartitionId().getPartitionId()); if (!startOffset.equals(nextOffset)) { _controllerMetrics.addMeteredTableValue(realtimeTableName, ControllerMeter.OFFSET_AUTO_RESET_SKIPPED_OFFSETS, Long.parseLong(startOffset) - Long.parseLong(nextOffset)); @@ -1016,7 +1017,7 @@ private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig st taskProperties.put(Constants.RESET_OFFSET_TO, startOffset); taskProperties.put(Constants.RESET_OFFSET_TOPIC_NAME, streamConfig.getTopicName()); taskProperties.put(Constants.RESET_OFFSET_TOPIC_PARTITION, - Integer.toString(newLLCSegmentName.getPartitionGroupId())); + Integer.toString(newLLCSegmentName.getTopicPartitionId().getPartitionId())); _helixResourceManager.invokeControllerPeriodicTask(streamConfig.getTableNameWithType(), Constants.REALTIME_OFFSET_AUTO_RESET_MANAGER, taskProperties); } @@ -1034,7 +1035,8 @@ private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig st // Add the partition metadata if available SegmentPartitionMetadata partitionMetadata = - getPartitionMetadataFromTableConfig(tableConfig, newLLCSegmentName.getPartitionGroupId(), numPartitions); + getPartitionMetadataFromTableConfig(tableConfig, + newLLCSegmentName.getTopicPartitionId().getPartitionId(), numPartitions); if (partitionMetadata != null) { newSegmentZKMetadata.setPartitionMetadata(partitionMetadata); } @@ -1429,17 +1431,18 @@ private Map getLatestSegmentZKMetadataMap(String rea Map latestLLCSegmentNameMap = new HashMap<>(); for (String segmentName : segments) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - latestLLCSegmentNameMap.compute(llcSegmentName.getPartitionGroupId(), (partitionId, latestLLCSegmentName) -> { - if (latestLLCSegmentName == null) { - return llcSegmentName; - } else { - if (llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { - return llcSegmentName; - } else { - return latestLLCSegmentName; - } - } - }); + latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId().getPartitionId(), + (partitionId, latestLLCSegmentName) -> { + if (latestLLCSegmentName == null) { + return llcSegmentName; + } else { + if (llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { + return llcSegmentName; + } else { + return latestLLCSegmentName; + } + } + }); } Map latestSegmentZKMetadataMap = new HashMap<>(); @@ -1581,7 +1584,8 @@ public static boolean isTopicPaused(IdealState idealState, String segmentName) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { return isTopicPaused(idealState, - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(llcSegmentName.getPartitionGroupId())); + IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId( + llcSegmentName.getTopicPartitionId().getPartitionId())); } return false; } @@ -1630,7 +1634,7 @@ private void updateInstanceStatesForNewConsumingSegment(Map> segmentEntry : instanceStatesMap.entrySet()) { if (segmentEntry.getValue().containsValue(SegmentStateModel.CONSUMING) - && new LLCSegmentName(segmentEntry.getKey()).getPartitionGroupId() == partitionId) { + && new LLCSegmentName(segmentEntry.getKey()).getTopicPartitionId().getPartitionId() == partitionId) { previousConsumingSegment = segmentEntry.getKey(); break; } @@ -2104,7 +2109,8 @@ private StreamPartitionMsgOffset selectStartOffset(OffsetCriteria offsetCriteria } private LLCSegmentName getNextLLCSegmentName(LLCSegmentName lastLLCSegmentName, long creationTimeMs) { - return new LLCSegmentName(lastLLCSegmentName.getTableName(), lastLLCSegmentName.getPartitionGroupId(), + return new LLCSegmentName(lastLLCSegmentName.getTableName(), + lastLLCSegmentName.getTopicPartitionId().getPartitionId(), lastLLCSegmentName.getSequenceNumber() + 1, creationTimeMs); } @@ -2705,7 +2711,8 @@ private Set filterSegmentsToCommit(Set allConsumingSegments, .map(Integer::parseInt) .collect(Collectors.toSet()); Set targetSegments = allConsumingSegments.stream() - .filter(segmentName -> partitionsToCommit.contains(new LLCSegmentName(segmentName).getPartitionGroupId())) + .filter(segmentName -> partitionsToCommit.contains( + new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId())) .collect(Collectors.toSet()); Preconditions.checkState(!targetSegments.isEmpty(), "Cannot find segments to commit for partitions: %s", partitionGroupIdsToCommitStr); @@ -2893,7 +2900,8 @@ private Set findConsumingSegmentsOfTopics(IdealState idealState, List { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null && !topicIndices.contains( - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(llcSegmentName.getPartitionGroupId()))) { + IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId( + llcSegmentName.getTopicPartitionId().getPartitionId()))) { return; } for (String state : instanceToStateMap.values()) { diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java index 1ed304740f5f..e300eb450629 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java @@ -215,7 +215,7 @@ public void testAddingRealtimeTableSegmentsWithPartitionIdInZkMetadata() { int partitionId; LLCSegmentName llcSegmentName = LLCSegmentName.of(segment); if (llcSegmentName != null) { - partitionId = llcSegmentName.getPartitionGroupId(); + partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); } else { partitionId = Integer.parseInt(segment.substring(1, 2)); } 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..ebbb96cc0fd5 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 @@ -1860,7 +1860,7 @@ public void testAddRealtimeTableWithConsumingMetadata() for (String segmentName : idealState.getPartitionSet()) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionGroupId = llcSegmentName.getPartitionGroupId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); if (partitionGroupId == 0) { assertEquals(llcSegmentName.getSequenceNumber(), 5); } else if (partitionGroupId == 1) { diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java index ffd37a2d54a1..a1802950f899 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java @@ -543,7 +543,7 @@ public void testSubsetPartitionAssignment() { // Extract partition ID from segment name LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionId = llcSegmentName.getPartitionGroupId(); + int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); // Verify the partition ID is one of our subset partitions boolean isSubsetPartition = false; 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..497dcc08224e 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 @@ -259,7 +259,7 @@ public void testSetUpNewTableWithExplicitSequenceNumbers() { // Verify segments are created with the explicit sequence numbers for (String segmentName : instanceStatesMap.keySet()) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionGroupId = llcSegmentName.getPartitionGroupId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); int sequence = llcSegmentName.getSequenceNumber(); if (partitionGroupId == 0) { assertEquals(sequence, 5); @@ -296,7 +296,8 @@ public void testSetUpNewTableDefaultSequenceNumberResolvesToZero() { for (String segmentName : instanceStatesMap.keySet()) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); assertEquals(llcSegmentName.getSequenceNumber(), 0, - "Default sequence number -1 should resolve to 0 for partition " + llcSegmentName.getPartitionGroupId()); + "Default sequence number -1 should resolve to 0 for partition " + + llcSegmentName.getTopicPartitionId().getPartitionId()); } } @@ -777,7 +778,8 @@ private void testSetUpNewPartitions(FakePinotLLCRealtimeSegmentManager segmentMa assertTrue(oldSegmentZKMetadataMap.containsKey(segmentName)); assertTrue(segmentZKMetadataMap.containsKey(segmentName)); assertEquals(segmentZKMetadataMap.get(segmentName), oldSegmentZKMetadataMap.get(segmentName)); - oldNumPartitions = Math.max(oldNumPartitions, new LLCSegmentName(segmentName).getPartitionGroupId() + 1); + oldNumPartitions = Math.max(oldNumPartitions, + new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId() + 1); } // Check that for new partition groups, each partition group should have exactly 1 new segment in CONSUMING @@ -786,7 +788,7 @@ private void testSetUpNewPartitions(FakePinotLLCRealtimeSegmentManager segmentMa Map> partitionGroupIdToSegmentsMap = new HashMap<>(); for (Map.Entry> entry : instanceStatesMap.entrySet()) { String segmentName = entry.getKey(); - int partitionGroupId = new LLCSegmentName(segmentName).getPartitionGroupId(); + int partitionGroupId = new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); partitionGroupIdToSegmentsMap.computeIfAbsent(partitionGroupId, k -> new ArrayList<>()).add(segmentName); } for (int partitionGroupId = oldNumPartitions; partitionGroupId < segmentManager._numPartitions; @@ -1154,7 +1156,7 @@ private void verifyRepairs(FakePinotLLCRealtimeSegmentManager segmentManager, Li if (instanceStateMap.containsValue(SegmentStateModel.ONLINE) || instanceStateMap.containsValue( SegmentStateModel.CONSUMING)) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionsId = llcSegmentName.getPartitionGroupId(); + int partitionsId = llcSegmentName.getTopicPartitionId().getPartitionId(); Map sequenceNumberToSegmentMap = partitionGroupIdToSegmentsMap.get(partitionsId); int sequenceNumber = llcSegmentName.getSequenceNumber(); assertFalse(sequenceNumberToSegmentMap.containsKey(sequenceNumber)); @@ -1996,7 +1998,7 @@ public void testBuildPartitionGroupConsumptionStatusFromZKMetadataMatchesOrigina Map latestSegmentZKMetadataMap = new HashMap<>(); for (Map.Entry entry : segmentManager._segmentZKMetadataMap.entrySet()) { LLCSegmentName llcSegmentName = new LLCSegmentName(entry.getKey()); - int partitionId = llcSegmentName.getPartitionGroupId(); + int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); latestSegmentZKMetadataMap.merge(partitionId, entry.getValue(), (existing, candidate) -> { int existingSeq = new LLCSegmentName(existing.getSegmentName()).getSequenceNumber(); int candidateSeq = new LLCSegmentName(candidate.getSegmentName()).getSequenceNumber(); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java index da98f32aba85..f51eaf890181 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java @@ -48,7 +48,7 @@ public class TableRebalancerTest { private static final TableRebalancer.PartitionIdFetcher DUMMY_PARTITION_FETCHER = segmentName -> 0; private static final TableRebalancer.PartitionIdFetcher SIMPLE_PARTITION_FETCHER = segmentName -> { LLCSegmentName name = LLCSegmentName.of(segmentName); - return name == null ? -1 : name.getPartitionGroupId(); + return name == null ? -1 : name.getTopicPartitionId().getPartitionId(); }; private static final TableRebalancer.DataLossRiskAssessor DEFAULT_DATA_LOSS_RISK_ASSESSOR = diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java index cc99a15e1679..f078e753262d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java @@ -102,11 +102,11 @@ public void acquire(LLCSegmentName llcSegmentName) LOGGER.warn("Failed to acquire consumer semaphore for segment: {} in: {}ms. Retrying.", segmentName, waitTimeMs); _serverMetrics.setValueOfPartitionGauge(_realtimeTableDataManager.getTableName(), - llcSegmentName.getPartitionGroupId(), ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, waitTimeMs); + llcSegmentName.getTopicPartitionId().getPartitionId(), ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, waitTimeMs); } } finally { _serverMetrics.setValueOfPartitionGauge(_realtimeTableDataManager.getTableName(), - llcSegmentName.getPartitionGroupId(), ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, 0); + llcSegmentName.getTopicPartitionId().getPartitionId(), ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, 0); } } @@ -197,7 +197,7 @@ int getPreviousSegmentSequenceNumberFromIdealState(LLCSegmentName currentSegment // so that it can always pass the check. int maxSequenceNumberBelowCurrentSegment = -1; String instanceId = _realtimeTableDataManager.getServerInstance(); - int partitionId = currentSegment.getPartitionGroupId(); + int partitionId = currentSegment.getTopicPartitionId().getPartitionId(); int currentSequenceNumber = currentSegment.getSequenceNumber(); for (Map.Entry> entry : getSegmentAssignment().entrySet()) { @@ -215,7 +215,7 @@ int getPreviousSegmentSequenceNumberFromIdealState(LLCSegmentName currentSegment continue; } - if (llcSegmentName.getPartitionGroupId() != partitionId) { + if (llcSegmentName.getTopicPartitionId().getPartitionId() != partitionId) { // ignore segments of different partitions. continue; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java index e26390b9e4fb..ed4b7236af9d 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java @@ -485,7 +485,7 @@ public Set stopTrackingAllPartitions() { */ public void stopTrackingPartition(String segmentName) { _segmentsToIgnore.put(segmentName, true); - removePartitionId(new LLCSegmentName(segmentName).getPartitionGroupId()); + removePartitionId(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); } /** @@ -535,7 +535,8 @@ public void markPartitionForVerification(String segmentName) { // Do not update the tracker state during server startup period or if the segment is marked to be ignored return; } - _partitionsMarkedForVerification.put(new LLCSegmentName(segmentName).getPartitionGroupId(), _clock.millis()); + _partitionsMarkedForVerification.put(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(), + _clock.millis()); } /** 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 0fbe75597d44..a922c84b2ff2 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 @@ -1784,7 +1784,7 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf String timeColumnName = tableConfig.getValidationConfig().getTimeColumnName(); // TODO Validate configs IndexingConfig indexingConfig = _tableConfig.getIndexingConfig(); - _partitionGroupId = llcSegmentName.getPartitionGroupId(); + _partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); List> streamConfigMaps = IngestionConfigUtils.getStreamConfigMaps(_tableConfig); int numStreams = streamConfigMaps.size(); if (numStreams == 1) { diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java index bb65dc48dbb5..3a720f4ce3c1 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java @@ -351,7 +351,8 @@ public void updateIngestionMetrics(String segmentName, int partitionId, long ing * {@code Long.MIN_VALUE} when it is not available. */ public long getPartitionIngestionTimeMs(String segmentName) { - return _ingestionDelayTracker.getPartitionIngestionTimeMs(new LLCSegmentName(segmentName).getPartitionGroupId()); + return _ingestionDelayTracker.getPartitionIngestionTimeMs( + new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); } /** @@ -372,7 +373,8 @@ public void removeIngestionMetrics(String segmentName) { @Override public void onConsumingToDropped(String segmentName) { // NOTE: No need to mark segment ignored here because it should have already been dropped. - _ingestionDelayTracker.stopTrackingPartition(new LLCSegmentName(segmentName).getPartitionGroupId()); + _ingestionDelayTracker.stopTrackingPartition( + new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); } /** @@ -413,7 +415,7 @@ public Set getHostedPartitionsGroupIds() { CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING); for (String segmentNameStr : segments) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - partitionsHostedByThisServer.add(segmentName.getPartitionGroupId()); + partitionsHostedByThisServer.add(segmentName.getTopicPartitionId().getPartitionId()); } return partitionsHostedByThisServer; } @@ -524,7 +526,7 @@ protected void doAddOnlineSegment(String segmentName) if (_enforceConsumptionInOrder) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { - getConsumerCoordinator(llcSegmentName.getPartitionGroupId()).register(llcSegmentName); + getConsumerCoordinator(llcSegmentName.getTopicPartitionId().getPartitionId()).register(llcSegmentName); } } } @@ -599,7 +601,7 @@ private void doAddConsumingSegment(String segmentName) // Generates only one semaphore for every partition LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionGroupId = llcSegmentName.getPartitionGroupId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); ConsumerCoordinator consumerCoordinator = getConsumerCoordinator(partitionGroupId); // Create the segment data manager and register it 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 41fb745acd40..4285c6bd9b1c 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 @@ -1200,7 +1200,7 @@ public FakeRealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, Table throws Exception { super(segmentZKMetadata, tableConfig, realtimeTableDataManager, resourceDataDir, new IndexLoadingConfig(makeInstanceDataManagerConfig(), tableConfig), schema, llcSegmentName, - consumerCoordinatorMap.get(llcSegmentName.getPartitionGroupId()), serverMetrics, null, null, + consumerCoordinatorMap.get(llcSegmentName.getTopicPartitionId().getPartitionId()), serverMetrics, null, null, () -> true); _state = RealtimeSegmentDataManager.class.getDeclaredField("_state"); _state.setAccessible(true); diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java index 28706764a530..3a11c38861ea 100644 --- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java +++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java @@ -1098,7 +1098,7 @@ private int getNumConsumingPartitions(String tableNameWithType) { Map stateMap = idealState.getInstanceStateMap(segmentName); if (stateMap != null && stateMap.containsValue(CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { - consumingPartitions.add(new LLCSegmentName(segmentName).getPartitionGroupId()); + consumingPartitions.add(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); } } return consumingPartitions.size(); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java index fb214936fad1..7af6cb915de9 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java @@ -152,7 +152,7 @@ private boolean isOffline(int partition, int seqNum) { for (String segmentNameStr : ev.getPartitionSet()) { if (LLCSegmentName.isLLCSegment(segmentNameStr)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getSequenceNumber() == seqNum && segmentName.getPartitionGroupId() == partition + if (segmentName.getSequenceNumber() == seqNum && segmentName.getTopicPartitionId().getPartitionId() == partition && ev.getStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.OFFLINE)) { isOffline = true; @@ -563,7 +563,7 @@ private int getSegmentSeqNum(int partition) { if (is.getInstanceStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getPartitionGroupId() == partition) { + if (segmentName.getTopicPartitionId().getPartitionId() == partition) { seqNum.set(segmentName.getSequenceNumber()); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java index 98cae2d48d69..cb6657826fa7 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java @@ -111,7 +111,7 @@ private int getSegmentSeqNum(int partition) { if (is.getInstanceStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getPartitionGroupId() == partition) { + if (segmentName.getTopicPartitionId().getPartitionId() == partition) { seqNum.set(segmentName.getSequenceNumber()); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java index 7b101d49d527..094d9b49bd50 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java @@ -111,7 +111,7 @@ private int getSegmentSeqNum(int partition) { if (is.getInstanceStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getPartitionGroupId() == partition) { + if (segmentName.getTopicPartitionId().getPartitionId() == partition) { seqNum.set(segmentName.getSequenceNumber()); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java index d636a61351f2..818161e2fa66 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java @@ -407,7 +407,7 @@ private static int getSegmentPartitionId(String segmentName) { case UPLOADED_SEGMENT_3: return 1; default: - return new LLCSegmentName(segmentName).getPartitionGroupId(); + return new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java index 5c578d20b5ec..392313a6f7ce 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java @@ -244,7 +244,8 @@ private Map getPartitionSegmentNumberToMetadataMap( Map segmentZKMetadataMap = new HashMap<>(); for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentZKMetadata.getSegmentName()); - String segmentKey = llcSegmentName.getPartitionGroupId() + "_" + llcSegmentName.getSequenceNumber(); + String segmentKey = + llcSegmentName.getTopicPartitionId().getPartitionId() + "_" + llcSegmentName.getSequenceNumber(); segmentZKMetadataMap.put(segmentKey, segmentZKMetadata); } return segmentZKMetadataMap; diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java index 201fd136e094..d3353b68d407 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java @@ -161,7 +161,8 @@ public void testPartitionMetadata() { assertNotNull(columnPartitionMetadata); assertTrue(columnPartitionMetadata.getFunctionName().equalsIgnoreCase("murmur")); assertEquals(columnPartitionMetadata.getNumPartitions(), 2); - int partitionGroupId = new LLCSegmentName(segmentZKMetadata.getSegmentName()).getPartitionGroupId(); + int partitionGroupId = + new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().getPartitionId(); assertEquals(columnPartitionMetadata.getPartitions(), Set.of(partitionGroupId)); numSegmentsForPartition[partitionGroupId]++; } @@ -268,7 +269,8 @@ public void testNonPartitionedStream() assertNotNull(columnPartitionMetadata); assertTrue(columnPartitionMetadata.getFunctionName().equalsIgnoreCase("murmur")); assertEquals(columnPartitionMetadata.getNumPartitions(), 2); - int partitionGroupId = new LLCSegmentName(segmentZKMetadata.getSegmentName()).getPartitionGroupId(); + int partitionGroupId = + new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().getPartitionId(); numSegmentsForPartition[partitionGroupId]++; if (segmentZKMetadata.getStatus() == Status.IN_PROGRESS) { @@ -345,7 +347,8 @@ public void testNonPartitionedStream() assertNotNull(columnPartitionMetadata); assertTrue(columnPartitionMetadata.getFunctionName().equalsIgnoreCase("murmur")); assertEquals(columnPartitionMetadata.getNumPartitions(), 2); - int partitionGroupId = new LLCSegmentName(segmentZKMetadata.getSegmentName()).getPartitionGroupId(); + int partitionGroupId = + new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().getPartitionId(); numSegmentsForPartition[partitionGroupId]++; if (segmentZKMetadata.getStatus() == Status.IN_PROGRESS) { diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java index 1dd1eb1f2422..8c78c861f8df 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java @@ -298,7 +298,7 @@ private static int getSegmentPartitionId(String segmentName) { case UPLOADED_SEGMENT_3: return 1; default: - return new LLCSegmentName(segmentName).getPartitionGroupId(); + return new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); } } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java index 2e9f472d9853..3f901841bb66 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java @@ -225,7 +225,7 @@ private static int getSegmentPartitionId(String segmentName) { case UPLOADED_SEGMENT_3: return 1; default: - return new LLCSegmentName(segmentName).getPartitionGroupId(); + return new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); } } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java index d2013f22399f..d04a559618c6 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java @@ -256,7 +256,7 @@ private int getNumConsumingPartitions(String tableNameWithType) { Map stateMap = idealState.getInstanceStateMap(segmentName); if (stateMap != null && stateMap.containsValue(CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { - consumingPartitions.add(new LLCSegmentName(segmentName).getPartitionGroupId()); + consumingPartitions.add(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); } } return consumingPartitions.size(); @@ -383,7 +383,7 @@ public void testSegmentsFromAllTopics() { Set topicIndicesSeen = new HashSet<>(); for (String segmentName : idealState.getPartitionSet()) { if (LLCSegmentName.isLLCSegment(segmentName)) { - int pgId = new LLCSegmentName(segmentName).getPartitionGroupId(); + int pgId = new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); topicIndicesSeen.add(IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(pgId)); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java index fc7f5264c016..d7b954b96c12 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java @@ -70,7 +70,8 @@ private static Map getPartitionSegmentNumberToMetadat Map segmentZKMetadataMap = new HashMap<>(); for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentZKMetadata.getSegmentName()); - String segmentKey = llcSegmentName.getPartitionGroupId() + "_" + llcSegmentName.getSequenceNumber(); + String segmentKey = + llcSegmentName.getTopicPartitionId().getPartitionId() + "_" + llcSegmentName.getSequenceNumber(); segmentZKMetadataMap.put(segmentKey, segmentZKMetadata); } return segmentZKMetadataMap; diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java index 839d58236664..b4d7218524a1 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java @@ -603,14 +603,15 @@ static List filterSegmentsforRealtimeTable(List { - if (latestSegment == null) { - return llcSegmentName; - } else { - return latestSegment.getSequenceNumber() > llcSegmentName.getSequenceNumber() + partitionIdToLatestCompletedSegment.compute(llcSegmentName.getTopicPartitionId().getPartitionId(), + (partId, latestSegment) -> { + if (latestSegment == null) { + return llcSegmentName; + } else { + return latestSegment.getSequenceNumber() > llcSegmentName.getSequenceNumber() ? latestSegment : llcSegmentName; - } - }); + } + }); } } Set filteredSegmentNames = new HashSet<>(); diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java index 16506881cefe..a2ab2de5f645 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java @@ -169,14 +169,15 @@ private Set getLastLLCSegmentPerPartition(List segmen // Skip UPLOADED segments that don't conform to the LLC segment name LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentZKMetadata.getSegmentName()); if (llcSegmentName != null) { - latestLLCSegmentNameMap.compute(llcSegmentName.getPartitionGroupId(), (k, latestLLCSegmentName) -> { - if (latestLLCSegmentName == null - || llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { - return llcSegmentName; - } else { - return latestLLCSegmentName; - } - }); + latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId().getPartitionId(), + (k, latestLLCSegmentName) -> { + if (latestLLCSegmentName == null + || llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { + return llcSegmentName; + } else { + return latestLLCSegmentName; + } + }); } } Set lastLLCSegmentPerPartition = new HashSet<>(); diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java index 2240974cec1a..166d4387a743 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java @@ -260,7 +260,7 @@ private void getCompletedSegmentsInfo(String realtimeTableName, List { if (latestLLCSegmentName == null diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java index d0ae144ae5e8..a0746bb48669 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java @@ -455,7 +455,7 @@ static Map> getInvalidSegmentsByInferredPartition(@Nullabl static int inferPartitionId(String segmentName, int numPartitions) { if (LLCSegmentName.isLLCSegment(segmentName)) { LLCSegmentName llc = LLCSegmentName.of(segmentName); - return llc != null ? (llc.getPartitionGroupId() % numPartitions) : -1; + return llc != null ? (llc.getTopicPartitionId().getPartitionId() % numPartitions) : -1; } else if (UploadedRealtimeSegmentName.isUploadedRealtimeSegmentName(segmentName)) { UploadedRealtimeSegmentName uploaded = UploadedRealtimeSegmentName.of(segmentName); return uploaded != null ? (uploaded.getPartitionId() % numPartitions) : -1; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java index 884227d31ec5..b8efa78d4beb 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java @@ -115,7 +115,7 @@ public MultiColumnRealtimeLuceneTextIndex( LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); _refreshListener = new RealtimeLuceneRefreshListener(llcSegmentName.getTableName(), segmentName, MultiColumnTextIndexConstants.INDEX_DIR_NAME, - llcSegmentName.getPartitionGroupId(), _indexCreator::getNumDocs); + llcSegmentName.getTopicPartitionId().getPartitionId(), _indexCreator::getNumDocs); _searcherManager.addListener(_refreshListener); _analyzer = _indexCreator.getIndexWriter().getConfig().getAnalyzer(); _queryParserClassConstructor = diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java index 47b7ef75f189..e0adc6f49d8a 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java @@ -92,7 +92,7 @@ public RealtimeLuceneTextIndex(String column, File segmentIndexDir, String segme LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); _refreshListener = new RealtimeLuceneRefreshListener(llcSegmentName.getTableName(), segmentName, column, - llcSegmentName.getPartitionGroupId(), _indexCreator::getNumDocs); + llcSegmentName.getTopicPartitionId().getPartitionId(), _indexCreator::getNumDocs); _searcherManager.addListener(_refreshListener); _analyzer = _indexCreator.getIndexWriter().getConfig().getAnalyzer(); _queryParserClassConstructor = diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java index 2d0db9781864..7577e4484e32 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java @@ -125,7 +125,7 @@ public StatelessRealtimeSegmentWriter(SegmentZKMetadata segmentZKMetadata, Index LLCSegmentName llcSegmentName = new LLCSegmentName(segmentZKMetadata.getSegmentName()); _segmentName = segmentZKMetadata.getSegmentName(); - _partitionGroupId = llcSegmentName.getPartitionGroupId(); + _partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); _segBuildSemaphore = segBuildSemaphore; _tableNameWithType = TableNameBuilder.forType(TableType.REALTIME).tableNameWithType(llcSegmentName.getTableName()); _segmentZKMetadata = segmentZKMetadata; From 528f11f5527fba94e5811977483d382826675129 Mon Sep 17 00:00:00 2001 From: Rekha Seethamraju Date: Thu, 9 Jul 2026 13:10:14 -0700 Subject: [PATCH 2/2] Use TopicPartitionId as primary partition identity across all internal code Add context-aware LLCSegmentName(String, boolean hasMultipleStreams) that decomposes old-format composite partition IDs (e.g. 10003 -> topic 1, partition 3) at parse time when the table has multiple streams. This ensures TopicPartitionId is always canonical regardless of segment format. Migrate all internal maps from Map to Map across controller, server, query planner, segment-local, and minion tasks. Change assignConsumingSegment, computeStartOffset, and other internal methods to accept TopicPartitionId directly. Thread hasMultipleStreams through RealtimeSegmentConfig, MutableIndexContext, and Lucene index constructors. Add TopicPartitionId.fromPartitionGroupMetadata() for SPI boundary wrapping and SegmentUtils.getTopicPartitionIdFromSegmentName() as the new API. Deprecate IngestionConfigUtils encode/decode helpers. --- .../pinot/common/utils/LLCSegmentName.java | 44 +++- .../pinot/common/utils/SegmentUtils.java | 20 +- .../pinot/common/utils/TopicPartitionId.java | 28 ++ .../common/utils/LLCSegmentNameTest.java | 18 ++ .../common/utils/TopicPartitionIdTest.java | 36 +++ .../PinotSegmentRestletResource.java | 33 +-- .../controller/api/resources/TableViews.java | 5 +- .../helix/core/PinotHelixResourceManager.java | 5 +- .../BaseStrictRealtimeSegmentAssignment.java | 30 ++- ...tiTierStrictRealtimeSegmentAssignment.java | 11 +- .../segment/RealtimeSegmentAssignment.java | 14 +- ...leTierStrictRealtimeSegmentAssignment.java | 3 +- .../MissingConsumingSegmentFinder.java | 21 +- .../PinotLLCRealtimeSegmentManager.java | 167 ++++++------ .../PinotSegmentRestletResourceTest.java | 29 ++- .../helix/PinotResourceManagerTest.java | 2 +- ...inotHelixResourceManagerStatelessTest.java | 2 +- ...timeReplicaGroupSegmentAssignmentTest.java | 2 +- .../MissingConsumingSegmentFinderTest.java | 37 +-- .../PinotLLCRealtimeSegmentManagerTest.java | 35 +-- .../core/rebalance/TableRebalancerTest.java | 2 +- .../manager/realtime/ConsumerCoordinator.java | 11 +- .../realtime/IngestionDelayTracker.java | 155 +++++++----- .../realtime/RealtimeSegmentDataManager.java | 39 ++- .../realtime/RealtimeTableDataManager.java | 58 +++-- .../realtime/ConsumerCoordinatorTest.java | 3 +- .../realtime/IngestionDelayTrackerTest.java | 239 +++++++++++------- .../RealtimeSegmentDataManagerTest.java | 3 +- .../tests/BaseClusterIntegrationTest.java | 2 +- .../LLCRealtimeClusterIntegrationTest.java | 5 +- ...CRealtimeKafka3ClusterIntegrationTest.java | 2 +- ...CRealtimeKafka4ClusterIntegrationTest.java | 2 +- ...alUpsertTableRebalanceIntegrationTest.java | 2 +- ...timeIngestionSegmentCommitFailureTest.java | 3 +- ...tionLLCRealtimeClusterIntegrationTest.java | 6 +- ...ertTableSegmentPreloadIntegrationTest.java | 2 +- ...sertTableSegmentUploadIntegrationTest.java | 2 +- ...tiTopicRealtimeClusterIntegrationTest.java | 4 +- .../utils/PauselessRealtimeTestUtils.java | 3 +- .../mergerollup/MergeRollupTaskGenerator.java | 14 +- .../tasks/purge/PurgeTaskGenerator.java | 14 +- ...ealtimeToOfflineSegmentsTaskGenerator.java | 25 +- .../MergeRollupTaskGeneratorTest.java | 6 +- .../rules/LeafStageWorkerAssignmentRule.java | 5 + .../mutable/MutableSegmentImpl.java | 3 +- .../realtime/impl/RealtimeSegmentConfig.java | 18 +- .../MultiColumnRealtimeLuceneTextIndex.java | 13 +- .../RealtimeLuceneTextIndex.java | 12 +- .../StatelessRealtimeSegmentWriter.java | 9 +- .../segment/index/text/TextIndexType.java | 2 +- .../mutable/provider/MutableIndexContext.java | 16 +- .../PartitionGroupConsumptionStatus.java | 3 + .../pinot/spi/utils/IngestionConfigUtils.java | 7 +- 53 files changed, 807 insertions(+), 425 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 20d9ac61e03b..3dd9a88b11e9 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 @@ -28,6 +28,10 @@ import org.joda.time.format.DateTimeFormatter; +/** + * Represents an LLC (Low-Level Consumer) segment name in the format: + * {@code {tableName}__{partitionGroupId}__{sequenceNumber}__{date}} + */ public class LLCSegmentName implements Comparable { private static final String SEPARATOR = "__"; private static final String DATE_FORMAT = "yyyyMMdd'T'HHmm'Z'"; @@ -39,39 +43,61 @@ public class LLCSegmentName implements Comparable { private final String _creationTime; private final String _segmentName; - public LLCSegmentName(String segmentName) { + /** + * Parses a segment name string. + * + *

When {@code hasMultipleStreams} is true, composite partition IDs (>= 10000) + * are decomposed into their topic and partition components. This ensures that + * segment {@code table__10003__5__date} produces {@code TopicPartitionId(1, 3)} + * instead of {@code TopicPartitionId(0, 10003)}. + */ + public LLCSegmentName(String segmentName, boolean hasMultipleStreams) { String[] parts = StringUtils.splitByWholeSeparator(segmentName, SEPARATOR); Preconditions.checkArgument(parts.length == 4, "Invalid LLC segment name: %s", segmentName); _tableName = parts[0]; - _topicPartitionId = new TopicPartitionId(Integer.parseInt(parts[1])); + int rawId = Integer.parseInt(parts[1]); + if (hasMultipleStreams && rawId >= TopicPartitionId.PARTITION_PADDING_OFFSET) { + _topicPartitionId = TopicPartitionId.fromMultiTopicPinotPartitionId(rawId); + } else { + _topicPartitionId = new TopicPartitionId(rawId); + } _sequenceNumber = Integer.parseInt(parts[2]); _creationTime = parts[3]; _segmentName = segmentName; } + /** @deprecated Use {@link #LLCSegmentName(String, boolean)} to provide multi-stream context. */ + @Deprecated + public LLCSegmentName(String segmentName) { + this(segmentName, false); + } + + @Deprecated public LLCSegmentName(String tableName, int partitionGroupId, int sequenceNumber, long msSinceEpoch) { Preconditions.checkArgument(!tableName.contains(SEPARATOR), "Illegal table name: %s", tableName); _tableName = tableName; _topicPartitionId = new TopicPartitionId(partitionGroupId); _sequenceNumber = sequenceNumber; - // ISO8601 date: 20160120T1234Z _creationTime = DATE_FORMATTER.print(msSinceEpoch); _segmentName = tableName + 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. - */ @Nullable - public static LLCSegmentName of(String segmentName) { + public static LLCSegmentName of(String segmentName, boolean hasMultipleStreams) { try { - return new LLCSegmentName(segmentName); + return new LLCSegmentName(segmentName, hasMultipleStreams); } catch (Exception e) { return null; } } + /** @deprecated Use {@link #of(String, boolean)} to provide multi-stream context. */ + @Deprecated + @Nullable + public static LLCSegmentName of(String segmentName) { + return of(segmentName, false); + } + /** * Returns whether the given segment name represents an LLC segment. */ diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java index fc9187338df6..b443d1fae61c 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/SegmentUtils.java @@ -68,11 +68,13 @@ public static Integer getSegmentPartitionId(SegmentZKMetadata segmentZKMetadata, /// Returns the partition id of a segment based on segment name. /// Can return `null` if the partition id cannot be determined. + /// @deprecated Use {@link #getTopicPartitionIdFromSegmentName(String, boolean)} instead. + @Deprecated @Nullable public static Integer getPartitionIdFromSegmentName(String segmentName) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { - return llcSegmentName.getTopicPartitionId().getPartitionId(); + return llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); } UploadedRealtimeSegmentName uploadedRealtimeSegmentName = UploadedRealtimeSegmentName.of(segmentName); if (uploadedRealtimeSegmentName != null) { @@ -81,6 +83,22 @@ public static Integer getPartitionIdFromSegmentName(String segmentName) { return null; } + /// Returns the topic-partition ID of a segment based on segment name. + /// Can return {@code null} if the partition ID cannot be determined. + @Nullable + public static TopicPartitionId getTopicPartitionIdFromSegmentName( + String segmentName, boolean hasMultipleStreams) { + LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName, hasMultipleStreams); + if (llcSegmentName != null) { + return llcSegmentName.getTopicPartitionId(); + } + UploadedRealtimeSegmentName uploadedRealtimeSegmentName = UploadedRealtimeSegmentName.of(segmentName); + if (uploadedRealtimeSegmentName != null) { + return new TopicPartitionId(uploadedRealtimeSegmentName.getPartitionId()); + } + return null; + } + /// Returns the partition id of a segment based on segment ZK metadata. /// Can return `null` if the partition id cannot be determined. @Nullable diff --git a/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java b/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java index b2918881863f..24c9133eea8f 100644 --- a/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java @@ -18,6 +18,7 @@ */ package org.apache.pinot.common.utils; +import com.google.common.base.Preconditions; import java.util.Objects; @@ -31,6 +32,9 @@ *

Thread-safe: instances are immutable. */ public final class TopicPartitionId implements Comparable { + // Must match IngestionConfigUtils.PARTITION_PADDING_OFFSET in pinot-spi + public static final int PARTITION_PADDING_OFFSET = 10000; + private final int _topicId; private final int _partitionId; @@ -74,6 +78,30 @@ public int hashCode() { return Objects.hash(_topicId, _partitionId); } + /** Recomposes the composite Pinot partition ID: {@code topicId * 10000 + partitionId}. */ + public int toMultiTopicPinotPartitionId() { + return _topicId * PARTITION_PADDING_OFFSET + _partitionId; + } + + /** Decomposes a composite Pinot partition ID into its topic and partition components. */ + public static TopicPartitionId fromMultiTopicPinotPartitionId(int pinotPartitionId) { + Preconditions.checkArgument(pinotPartitionId >= 0, "Negative Pinot partition ID: %s", pinotPartitionId); + return new TopicPartitionId(pinotPartitionId / PARTITION_PADDING_OFFSET, + pinotPartitionId % PARTITION_PADDING_OFFSET); + } + + /** + * Wraps a partition group ID from stream metadata into a {@link TopicPartitionId}. + * When {@code hasMultipleStreams} is true and the ID is composite-encoded (>= 10000), + * decomposes it into topic and partition components. + */ + public static TopicPartitionId fromPartitionGroupMetadata(int partitionGroupId, boolean hasMultipleStreams) { + if (hasMultipleStreams && partitionGroupId >= PARTITION_PADDING_OFFSET) { + return fromMultiTopicPinotPartitionId(partitionGroupId); + } + return new TopicPartitionId(partitionGroupId); + } + @Override public String toString() { return _topicId + ":" + _partitionId; 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 30323ecc0897..465af504d2ce 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 @@ -95,4 +95,22 @@ public void testLLCSegmentName() { Arrays.sort(testSorted); Assert.assertEquals(testSorted, new LLCSegmentName[]{segName5, segName1, segName6, segName3, segName4}); } + + @Test + public void testContextAwareParsing() { + // Old format with hasMultipleStreams=true decomposes composite + LLCSegmentName withContext = new LLCSegmentName("myTable__10003__5__20250101T0000Z", true); + assertEquals(withContext.getTopicPartitionId().getTopicId(), 1); + assertEquals(withContext.getTopicPartitionId().getPartitionId(), 3); + + // Old format with hasMultipleStreams=false keeps raw partition + LLCSegmentName withoutContext = new LLCSegmentName("myTable__10003__5__20250101T0000Z", false); + assertEquals(withoutContext.getTopicPartitionId().getTopicId(), 0); + assertEquals(withoutContext.getTopicPartitionId().getPartitionId(), 10003); + + // Small partition ID stays unchanged even with hasMultipleStreams=true + LLCSegmentName small = new LLCSegmentName("myTable__5__3__20250101T0000Z", true); + assertEquals(small.getTopicPartitionId().getTopicId(), 0); + assertEquals(small.getTopicPartitionId().getPartitionId(), 5); + } } diff --git a/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java b/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java index fca10f36452e..46c41e3a870b 100644 --- a/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java @@ -97,4 +97,40 @@ public void testToString() { TopicPartitionId id = new TopicPartitionId(2, 7); assertEquals(id.toString(), "2:7"); } + + @Test + public void testToMultiTopicPinotPartitionId() { + assertEquals(new TopicPartitionId(0, 5).toMultiTopicPinotPartitionId(), 5); + assertEquals(new TopicPartitionId(1, 3).toMultiTopicPinotPartitionId(), 10003); + assertEquals(new TopicPartitionId(2, 0).toMultiTopicPinotPartitionId(), 20000); + assertEquals(new TopicPartitionId(5).toMultiTopicPinotPartitionId(), 5); + } + + @Test + public void testFromMultiTopicPinotPartitionId() { + TopicPartitionId id1 = TopicPartitionId.fromMultiTopicPinotPartitionId(10003); + assertEquals(id1.getTopicId(), 1); + assertEquals(id1.getPartitionId(), 3); + + TopicPartitionId id2 = TopicPartitionId.fromMultiTopicPinotPartitionId(3); + assertEquals(id2.getTopicId(), 0); + assertEquals(id2.getPartitionId(), 3); + + TopicPartitionId id3 = TopicPartitionId.fromMultiTopicPinotPartitionId(20000); + assertEquals(id3.getTopicId(), 2); + assertEquals(id3.getPartitionId(), 0); + } + + @Test + public void testRoundTrip() { + TopicPartitionId original = new TopicPartitionId(1, 3); + TopicPartitionId roundTripped = + TopicPartitionId.fromMultiTopicPinotPartitionId(original.toMultiTopicPinotPartitionId()); + assertEquals(roundTripped, original); + + TopicPartitionId singleTopic = new TopicPartitionId(0, 42); + TopicPartitionId roundTripped2 = + TopicPartitionId.fromMultiTopicPinotPartitionId(singleTopic.toMultiTopicPinotPartitionId()); + assertEquals(roundTripped2, singleTopic); + } } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java index 8b00467dcd75..e5df900a1cea 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResource.java @@ -73,6 +73,7 @@ import org.apache.pinot.common.utils.DatabaseUtils; import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.common.utils.PauselessConsumptionUtils; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.common.utils.URIUtils; import org.apache.pinot.common.utils.UploadedRealtimeSegmentName; import org.apache.pinot.controller.ControllerConf; @@ -1046,17 +1047,17 @@ public String deleteSegmentsFromSequenceNum( Preconditions.checkState(idealState != null, "Ideal State does not exist for table " + tableNameWithType); Set idealStateSegmentsSet = idealState.getRecord().getMapFields().keySet(); - Map partitionToOldestSegment = + Map partitionToOldestSegment = getPartitionIDToOldestSegment(segments, idealStateSegmentsSet); - Map partitionIdToLatestSegment = new HashMap<>(); - Map> partitionIdToSegmentsToDeleteMap = - getPartitionIdToSegmentsToDeleteMap(partitionToOldestSegment, idealStateSegmentsSet, - partitionIdToLatestSegment); + Map partitionIdToLatestSegment = new HashMap<>(); + Map> partitionIdToSegmentsToDeleteMap = + getPartitionIdToSegmentsToDeleteMap(partitionToOldestSegment, + idealStateSegmentsSet, partitionIdToLatestSegment); Map response = new HashMap<>(); - Map partitionDetails = new HashMap<>(); + Map partitionDetails = new HashMap<>(); - for (Integer partitionID : partitionToOldestSegment.keySet()) { + for (TopicPartitionId partitionID : partitionToOldestSegment.keySet()) { Set segmentsToDeleteForPartition = partitionIdToSegmentsToDeleteMap.get(partitionID); LLCSegmentName oldestSegment = partitionToOldestSegment.get(partitionID); LLCSegmentName latestSegment = partitionIdToLatestSegment.get(partitionID); @@ -1111,12 +1112,13 @@ public String deleteSegmentsFromSequenceNum( * for that particular partition. */ @VisibleForTesting - Map> getPartitionIdToSegmentsToDeleteMap( - Map partitionToOldestSegment, - Set idealStateSegmentsSet, Map partitionIdToLatestSegment) { + Map> getPartitionIdToSegmentsToDeleteMap( + Map partitionToOldestSegment, + Set idealStateSegmentsSet, + Map partitionIdToLatestSegment) { // Find segments to delete (those with higher sequence numbers) - Map> partitionToSegmentsToDelete = new HashMap<>(); + Map> partitionToSegmentsToDelete = new HashMap<>(); for (String segmentName : idealStateSegmentsSet) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); @@ -1124,7 +1126,7 @@ Map> getPartitionIdToSegmentsToDeleteMap( LOGGER.info("Skip segment: {} not in low-level consumer format", segmentName); continue; } - int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionId = llcSegmentName.getTopicPartitionId(); LLCSegmentName oldestSegment = partitionToOldestSegment.get(partitionId); if (oldestSegment == null) { @@ -1148,8 +1150,9 @@ Map> getPartitionIdToSegmentsToDeleteMap( } @VisibleForTesting - Map getPartitionIDToOldestSegment(List segments, Set idealStateSegmentsSet) { - Map partitionToOldestSegment = new HashMap<>(); + Map getPartitionIDToOldestSegment( + List segments, Set idealStateSegmentsSet) { + Map partitionToOldestSegment = new HashMap<>(); for (String segment : segments) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segment); @@ -1163,7 +1166,7 @@ Map getPartitionIDToOldestSegment(List segments LOGGER.warn("Segment: {} is not present in the ideal state", segment); continue; } - int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionId = llcSegmentName.getTopicPartitionId(); LLCSegmentName currentOldest = partitionToOldestSegment.get(partitionId); if (currentOldest == null || llcSegmentName.getSequenceNumber() < currentOldest.getSequenceNumber()) { diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java index 35bad09dd210..cedf0e31cd74 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/TableViews.java @@ -57,6 +57,7 @@ import org.apache.pinot.common.lineage.SegmentLineageUtils; import org.apache.pinot.common.utils.DatabaseUtils; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.controller.api.exception.ControllerApplicationException; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.core.auth.Actions; @@ -210,7 +211,7 @@ public String getBadLLCSegmentsPerPartition( return JsonUtils.objectToPrettyString(new HashMap<>()); } - Map> partitionIdToSegments = new HashMap<>(); + Map> partitionIdToSegments = new HashMap<>(); for (SegmentStatusInfo segmentStatusInfo : segmentStatusInfoList) { String segmentName = segmentStatusInfo.getSegmentName(); LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); @@ -218,7 +219,7 @@ public String getBadLLCSegmentsPerPartition( continue; } partitionIdToSegments.computeIfAbsent( - llcSegmentName.getTopicPartitionId().getPartitionId(), k -> new TreeSet<>()) + llcSegmentName.getTopicPartitionId(), k -> new TreeSet<>()) .add(llcSegmentName); } 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 0da6670c2135..be62c388d553 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 @@ -131,6 +131,7 @@ import org.apache.pinot.common.utils.HashUtil; import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.common.utils.LogicalTableConfigUtils; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.common.utils.ZkStarter; import org.apache.pinot.common.utils.config.AccessControlUserConfigUtils; import org.apache.pinot.common.utils.config.InstanceUtils; @@ -1156,7 +1157,7 @@ public Collection getLastLLCCompletedSegments(String tableNameWithType) /// segments of a table and want to derive the last-completed LLC segment per partition without /// re-reading the property store). public Collection getLastLLCCompletedSegments(List segmentZKMetadataList) { - Map partitionIdToLastLLCCompletedSegmentMap = new HashMap<>(); + Map partitionIdToLastLLCCompletedSegmentMap = new HashMap<>(); for (SegmentZKMetadata zkMetadata : segmentZKMetadataList) { if (zkMetadata.getStatus() == CommonConstants.Segment.Realtime.Status.DONE) { LLCSegmentName llcName = LLCSegmentName.of(zkMetadata.getSegmentName()); @@ -1164,7 +1165,7 @@ public Collection getLastLLCCompletedSegments(List _segmentPartitionIdMap = new Object2IntOpenHashMap<>(); + private final Map _segmentPartitionIdMap = new HashMap<>(); @Override public List assignSegment(String segmentName, Map> currentAssignment, @@ -78,8 +80,9 @@ public List assignSegment(String segmentName, Map instancesAssigned = assignConsumingSegment(partitionId, instancePartitions); + TopicPartitionId partitionId = getPartitionId(segmentName); + List instancesAssigned = assignConsumingSegment( + partitionId, instancePartitions); Set existingAssignment = getExistingAssignment(partitionId, currentAssignment); // Check if the candidate assignment is consistent with existing assignment. Use existing assignment if not. if (existingAssignment == null) { @@ -103,7 +106,8 @@ public List assignSegment(String segmentName, Map getExistingAssignment(int partitionId, Map> currentAssignment) { + protected Set getExistingAssignment(TopicPartitionId partitionId, + Map> currentAssignment) { List uploadedSegments = new ArrayList<>(); for (Map.Entry> entry : currentAssignment.entrySet()) { // Skip OFFLINE segments as they are not rebalanced, so their assignment in idealState can be stale. @@ -115,13 +119,13 @@ protected Set getExistingAssignment(int partitionId, Map instanceStateMap) { /** * Returns the partition id of the given segment. */ - private int getPartitionId(String segmentName) { + private TopicPartitionId getPartitionId(String segmentName) { Integer partitionId = SegmentUtils.getSegmentPartitionId(segmentName, _tableNameWithType, _helixManager, _partitionColumn); Preconditions.checkState(partitionId != null, "Failed to find partition id for segment: %s of table: %s", segmentName, _tableNameWithType); - return partitionId; + boolean hasMultipleStreams = + IngestionConfigUtils.hasMultipleStreams(_tableConfig); + return TopicPartitionId.fromPartitionGroupMetadata( + partitionId, hasMultipleStreams); } /** @@ -157,7 +164,8 @@ private boolean isSameAssignment(Set idealAssignment, List insta /** * Returns the partition id of the given segment, using cached partition id if exists. */ - protected int getPartitionIdUsingCache(String segmentName) { - return _segmentPartitionIdMap.computeIntIfAbsent(segmentName, this::getPartitionId); + protected TopicPartitionId getPartitionIdUsingCache(String segmentName) { + return _segmentPartitionIdMap.computeIfAbsent(segmentName, + this::getPartitionId); } } diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/MultiTierStrictRealtimeSegmentAssignment.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/MultiTierStrictRealtimeSegmentAssignment.java index 99295611f34c..1209303a4b98 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/MultiTierStrictRealtimeSegmentAssignment.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/MultiTierStrictRealtimeSegmentAssignment.java @@ -32,6 +32,7 @@ import org.apache.pinot.common.restlet.resources.RebalanceConfig; import org.apache.pinot.common.tier.Tier; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.spi.config.table.assignment.InstancePartitionsType; import org.apache.pinot.spi.utils.CommonConstants; @@ -72,7 +73,9 @@ public Map> rebalanceTable(Map instancesAssigned = - assignConsumingSegment(getPartitionIdUsingCache(segmentName), instancePartitions); + assignConsumingSegment( + getPartitionIdUsingCache(segmentName), + instancePartitions); String state = instanceStateMap.containsValue(CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING) ? CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING : CommonConstants.Helix.StateModel.SegmentStateModel.ONLINE; @@ -96,7 +99,8 @@ public Map> rebalanceTable(Map getExistingAssignment(int partitionId, Map> currentAssignment) { + protected Set getExistingAssignment(TopicPartitionId partitionId, + Map> currentAssignment) { LLCSegmentName latestLLCSegmentName = null; Set latestAssignment = null; for (Map.Entry> entry : currentAssignment.entrySet()) { @@ -107,7 +111,8 @@ protected Set getExistingAssignment(int partitionId, Map latestLLCSegmentName.getSequenceNumber())) { latestLLCSegmentName = llcSegmentName; latestAssignment = entry.getValue().keySet(); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java index d4b991e54b7a..84570bc79bed 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeSegmentAssignment.java @@ -33,6 +33,7 @@ import org.apache.pinot.common.tier.Tier; import org.apache.pinot.common.utils.PauselessConsumptionUtils; import org.apache.pinot.common.utils.SegmentUtils; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategy; import org.apache.pinot.controller.helix.core.assignment.segment.strategy.SegmentAssignmentStrategyFactory; import org.apache.pinot.controller.helix.core.realtime.PinotLLCRealtimeSegmentManager; @@ -112,14 +113,15 @@ public List assignSegment(String segmentName, Map assignConsumingSegment(String segmentName, InstancePartitions instancePartitions) { int segmentPartitionId = SegmentUtils.getSegmentPartitionIdOrDefault(segmentName, _tableNameWithType, _helixManager, _partitionColumn); - return assignConsumingSegment(segmentPartitionId, instancePartitions); + boolean hasMultipleStreams = IngestionConfigUtils.hasMultipleStreams(_tableConfig); + TopicPartitionId topicPartitionId = + TopicPartitionId.fromPartitionGroupMetadata(segmentPartitionId, hasMultipleStreams); + return assignConsumingSegment(topicPartitionId, instancePartitions); } - protected List assignConsumingSegment(int segmentPartitionId, InstancePartitions instancePartitions) { - // For multi-stream tables, Pinot partition IDs are encoded as (streamIndex * 10000 + streamPartitionId). - // Extract the stream-level partition id before computing the instance index to avoid incorrect slot mapping. - segmentPartitionId = - IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_tableConfig, segmentPartitionId); + protected List assignConsumingSegment(TopicPartitionId topicPartitionId, + InstancePartitions instancePartitions) { + int segmentPartitionId = topicPartitionId.getPartitionId(); int numReplicaGroups = instancePartitions.getNumReplicaGroups(); int numPartitions = instancePartitions.getNumPartitions(); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SingleTierStrictRealtimeSegmentAssignment.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SingleTierStrictRealtimeSegmentAssignment.java index c879e951b90c..10e318b16fda 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SingleTierStrictRealtimeSegmentAssignment.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/assignment/segment/SingleTierStrictRealtimeSegmentAssignment.java @@ -64,7 +64,8 @@ public Map> rebalanceTable(Map instancesAssigned = - assignConsumingSegment(getPartitionIdUsingCache(segmentName), instancePartitions); + assignConsumingSegment(getPartitionIdUsingCache(segmentName), + instancePartitions); String state = instanceStateMap.containsValue(SegmentStateModel.CONSUMING) ? SegmentStateModel.CONSUMING : SegmentStateModel.ONLINE; newAssignment.put(segmentName, SegmentAssignmentUtils.getInstanceStateMap(instancesAssigned, state)); diff --git a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java index d2d00130ae62..781bfea438a5 100644 --- a/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java +++ b/pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinder.java @@ -37,6 +37,7 @@ import org.apache.pinot.common.metrics.ControllerMeter; import org.apache.pinot.common.metrics.ControllerMetrics; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.controller.helix.core.PinotTableIdealStateBuilder; import org.apache.pinot.spi.config.table.PauseState; import org.apache.pinot.spi.stream.OffsetCriteria; @@ -63,7 +64,7 @@ public class MissingConsumingSegmentFinder { private final String _realtimeTableName; private final SegmentMetadataFetcher _segmentMetadataFetcher; - private final Map _partitionGroupIdToLargestStreamOffsetMap; + private final Map _partitionGroupIdToLargestStreamOffsetMap; private final StreamPartitionMsgOffsetFactory _streamPartitionMsgOffsetFactory; private ControllerMetrics _controllerMetrics; @@ -87,8 +88,12 @@ public MissingConsumingSegmentFinder(String realtimeTableName, ZkHelixPropertySt PinotTableIdealStateBuilder.getStreamMetadataList(streamConfigs, List.of(), pauseState == null ? new ArrayList<>() : pauseState.getIndexOfInactiveTopics(), false) .forEach(streamMetadata -> { + boolean hasMultipleStreams = streamConfigs.size() > 1; for (PartitionGroupMetadata metadata : streamMetadata.getPartitionGroupMetadataList()) { - _partitionGroupIdToLargestStreamOffsetMap.put(metadata.getPartitionGroupId(), metadata.getStartOffset()); + _partitionGroupIdToLargestStreamOffsetMap.put( + TopicPartitionId.fromPartitionGroupMetadata( + metadata.getPartitionGroupId(), hasMultipleStreams), + metadata.getStartOffset()); } }); } catch (Exception e) { @@ -101,7 +106,7 @@ public MissingConsumingSegmentFinder(String realtimeTableName, ZkHelixPropertySt @VisibleForTesting MissingConsumingSegmentFinder(String realtimeTableName, SegmentMetadataFetcher segmentMetadataFetcher, - Map partitionGroupIdToLargestStreamOffsetMap, + Map partitionGroupIdToLargestStreamOffsetMap, StreamPartitionMsgOffsetFactory streamPartitionMsgOffsetFactory) { _realtimeTableName = realtimeTableName; _segmentMetadataFetcher = segmentMetadataFetcher; @@ -124,8 +129,8 @@ public void findAndEmitMetrics(IdealState idealState) { @VisibleForTesting MissingSegmentInfo findMissingSegments(Map> idealStateMap, Instant now) { // create the maps - Map partitionGroupIdToLatestConsumingSegmentMap = new HashMap<>(); - Map partitionGroupIdToLatestCompletedSegmentMap = new HashMap<>(); + Map partitionGroupIdToLatestConsumingSegmentMap = new HashMap<>(); + Map partitionGroupIdToLatestCompletedSegmentMap = new HashMap<>(); idealStateMap.forEach((segmentName, instanceToStatusMap) -> { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { // Skip the uploaded realtime segments that don't conform to llc naming @@ -176,7 +181,7 @@ MissingSegmentInfo findMissingSegments(Map> idealSta return missingSegmentInfo; } - private void updateMaxDurationInfo(MissingSegmentInfo missingSegmentInfo, Integer partitionGroupId, + private void updateMaxDurationInfo(MissingSegmentInfo missingSegmentInfo, TopicPartitionId partitionGroupId, long segmentCompletionTimeMillis, Instant now) { long duration = Duration.between(Instant.ofEpochMilli(segmentCompletionTimeMillis), now).toMinutes(); if (duration > missingSegmentInfo._maxDurationInMinutes) { @@ -185,9 +190,9 @@ private void updateMaxDurationInfo(MissingSegmentInfo missingSegmentInfo, Intege LOGGER.warn("PartitionGroupId {} hasn't had a consuming segment for {} minutes!", partitionGroupId, duration); } - private void updateMap(Map partitionGroupIdToLatestSegmentMap, + private void updateMap(Map partitionGroupIdToLatestSegmentMap, LLCSegmentName llcSegmentName) { - int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionGroupId = llcSegmentName.getTopicPartitionId(); partitionGroupIdToLatestSegmentMap.compute(partitionGroupId, (pid, existingSegment) -> { if (existingSegment == null) { return llcSegmentName; 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 118735f8249a..608acd8c3a4b 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 @@ -85,6 +85,7 @@ import org.apache.pinot.common.utils.FileUploadDownloadClient; import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.common.utils.PauselessConsumptionUtils; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.common.utils.URIUtils; import org.apache.pinot.common.utils.helix.HelixHelper; import org.apache.pinot.common.utils.helix.IdealStateSingleCommit; @@ -272,7 +273,7 @@ public List getPartitionGroupConsumptionStatusL List partitionGroupConsumptionStatusList = new ArrayList<>(); // From all segment names in the ideal state, find unique partition group ids and their latest segment - Map partitionGroupIdToLatestSegment = new HashMap<>(); + Map partitionGroupIdToLatestSegment = new HashMap<>(); for (String segment : idealState.getRecord().getMapFields().keySet()) { // With Pinot upsert table allowing uploads of segments, the segment name of an upsert table segment may not // conform to LLCSegment format. We can skip such segments because they are NOT the consuming segments. @@ -280,7 +281,7 @@ public List getPartitionGroupConsumptionStatusL if (llcSegmentName == null) { continue; } - int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionGroupId = llcSegmentName.getTopicPartitionId(); partitionGroupIdToLatestSegment.compute(partitionGroupId, (k, latestSegment) -> { if (latestSegment == null) { return llcSegmentName; @@ -301,12 +302,13 @@ public List getPartitionGroupConsumptionStatusL StreamConfig streamConfig = streamConfigs.get(0); StreamPartitionMsgOffsetFactory offsetFactory = StreamConsumerFactoryProvider.create(streamConfig).createStreamMsgOffsetFactory(); - for (Map.Entry entry : partitionGroupIdToLatestSegment.entrySet()) { - int partitionGroupId = entry.getKey(); + for (Map.Entry entry : partitionGroupIdToLatestSegment.entrySet()) { + TopicPartitionId partitionGroupId = entry.getKey(); LLCSegmentName llcSegmentName = entry.getValue(); SegmentZKMetadata segmentZKMetadata = getSegmentZKMetadata(tableNameWithType, llcSegmentName.getSegmentName()); PartitionGroupConsumptionStatus partitionGroupConsumptionStatus = - new PartitionGroupConsumptionStatus(partitionGroupId, llcSegmentName.getSequenceNumber(), + new PartitionGroupConsumptionStatus(partitionGroupId.toMultiTopicPinotPartitionId(), + llcSegmentName.getSequenceNumber(), offsetFactory.create(segmentZKMetadata.getStartOffset()), segmentZKMetadata.getEndOffset() != null ? offsetFactory.create(segmentZKMetadata.getEndOffset()) : null, segmentZKMetadata.getStatus().toString()); @@ -315,10 +317,10 @@ public List getPartitionGroupConsumptionStatusL } else { // Multiple streams 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); + for (Map.Entry entry : partitionGroupIdToLatestSegment.entrySet()) { + TopicPartitionId partitionGroupId = entry.getKey(); + int index = partitionGroupId.getTopicId(); + int streamPartitionId = partitionGroupId.getPartitionId(); LLCSegmentName llcSegmentName = entry.getValue(); SegmentZKMetadata segmentZKMetadata = getSegmentZKMetadata(tableNameWithType, llcSegmentName.getSegmentName()); StreamPartitionMsgOffsetFactory offsetFactory = offsetFactories[index]; @@ -327,7 +329,8 @@ public List getPartitionGroupConsumptionStatusL offsetFactories[index] = offsetFactory; } PartitionGroupConsumptionStatus partitionGroupConsumptionStatus = - new PartitionGroupConsumptionStatus(partitionGroupId, streamPartitionId, llcSegmentName.getSequenceNumber(), + new PartitionGroupConsumptionStatus(partitionGroupId.toMultiTopicPinotPartitionId(), streamPartitionId, + llcSegmentName.getSequenceNumber(), offsetFactory.create(segmentZKMetadata.getStartOffset()), segmentZKMetadata.getEndOffset() != null ? offsetFactory.create(segmentZKMetadata.getEndOffset()) : null, segmentZKMetadata.getStatus().toString()); @@ -837,12 +840,12 @@ private String createNewSegmentMetadata(TableConfig tableConfig, String newConsumingSegmentName = null; LLCSegmentName committingLLCSegment = new LLCSegmentName(committingSegmentName); - int committingSegmentPartitionGroupId = committingLLCSegment.getTopicPartitionId().getPartitionId(); + TopicPartitionId committingSegmentPartitionGroupId = committingLLCSegment.getTopicPartitionId(); List streamConfigs = IngestionConfigUtils.getStreamConfigs(tableConfig); PartitionIdsWithIdealState partitionIdsWithIdealState = getPartitionIdsWithIdealState(streamConfigs, () -> getIdealState(realtimeTableName)); - Set partitionIds = partitionIdsWithIdealState._partitionIds; + Set partitionIds = partitionIdsWithIdealState._partitionIds; if (partitionIds.contains(committingSegmentPartitionGroupId)) { IdealState idealState = partitionIdsWithIdealState._idealState; @@ -857,11 +860,11 @@ private String createNewSegmentMetadata(TableConfig tableConfig, } String rawTableName = TableNameBuilder.extractRawTableName(realtimeTableName); long newSegmentCreationTimeMs = getCurrentTimeMs(); - LLCSegmentName newLLCSegment = new LLCSegmentName(rawTableName, committingSegmentPartitionGroupId, + LLCSegmentName newLLCSegment = new LLCSegmentName(rawTableName, + committingSegmentPartitionGroupId.toMultiTopicPinotPartitionId(), committingLLCSegment.getSequenceNumber() + 1, newSegmentCreationTimeMs); - StreamConfig streamConfig = - IngestionConfigUtils.getStreamConfigFromPinotPartitionId(streamConfigs, committingSegmentPartitionGroupId); + StreamConfig streamConfig = streamConfigs.get(committingSegmentPartitionGroupId.getTopicId()); createNewSegmentZKMetadata(tableConfig, streamConfig, newLLCSegment, newSegmentCreationTimeMs, committingSegmentDescriptor, committingSegmentZKMetadata, instancePartitions, partitionIds.size(), numReplicas); @@ -1008,7 +1011,7 @@ private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig st // Handle offset auto reset String nextOffset = committingSegmentDescriptor.getNextOffset(); String startOffset = computeStartOffset(nextOffset, streamConfig, - newLLCSegmentName.getTopicPartitionId().getPartitionId()); + newLLCSegmentName.getTopicPartitionId()); if (!startOffset.equals(nextOffset)) { _controllerMetrics.addMeteredTableValue(realtimeTableName, ControllerMeter.OFFSET_AUTO_RESET_SKIPPED_OFFSETS, Long.parseLong(startOffset) - Long.parseLong(nextOffset)); @@ -1036,7 +1039,7 @@ private void createNewSegmentZKMetadata(TableConfig tableConfig, StreamConfig st // Add the partition metadata if available SegmentPartitionMetadata partitionMetadata = getPartitionMetadataFromTableConfig(tableConfig, - newLLCSegmentName.getTopicPartitionId().getPartitionId(), numPartitions); + newLLCSegmentName.getTopicPartitionId(), numPartitions); if (partitionMetadata != null) { newSegmentZKMetadata.setPartitionMetadata(partitionMetadata); } @@ -1074,7 +1077,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, TopicPartitionId topicPartitionId) { if (!streamConfig.isEnableOffsetAutoReset() || streamConfig.isBackfillTopic()) { return nextOffset; } @@ -1087,7 +1090,7 @@ private String computeStartOffset(String nextOffset, StreamConfig streamConfig, timeThreshold, offsetThreshold); return nextOffset; } - int streamPartitionId = IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(partitionId); + int streamPartitionId = topicPartitionId.getPartitionId(); String clientId = getTableTopicUniqueClientId(streamConfig); StreamConsumerFactory consumerFactory = StreamConsumerFactoryProvider.create(streamConfig); StreamPartitionMsgOffsetFactory offsetFactory = consumerFactory.createStreamMsgOffsetFactory(); @@ -1137,8 +1140,8 @@ private String computeStartOffset(String nextOffset, StreamConfig streamConfig, @Nullable @VisibleForTesting - SegmentPartitionMetadata getPartitionMetadataFromTableConfig(TableConfig tableConfig, int partitionId, - int numPartitionGroups) { + SegmentPartitionMetadata getPartitionMetadataFromTableConfig(TableConfig tableConfig, + TopicPartitionId topicPartitionId, int numPartitionGroups) { SegmentPartitionConfig partitionConfig = tableConfig.getIndexingConfig().getSegmentPartitionConfig(); if (partitionConfig == null) { return null; @@ -1147,10 +1150,7 @@ SegmentPartitionMetadata getPartitionMetadataFromTableConfig(TableConfig tableCo if (columnPartitionMap.size() == 1) { Map.Entry entry = columnPartitionMap.entrySet().iterator().next(); ColumnPartitionConfig columnPartitionConfig = entry.getValue(); - // For multi-stream tables, convert Pinot partition ID (which includes padding offset) to stream partition ID. - // This ensures the partition metadata stored in ZK matches what the broker's partition function computes - // during query pruning. For example, stream 1 partition 5 has Pinot partition ID 10005, but should store 5. - int streamPartitionId = IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(tableConfig, partitionId); + int streamPartitionId = topicPartitionId.getPartitionId(); // If there are multiple streams, we assume the partition groups are evenly distributed across the streams, and // compute the per-stream partition group count accordingly. // This is needed for the partition function to correctly compute the stream partition id for pruning. @@ -1235,16 +1235,18 @@ public Map getPartitionCountMap(List streamConfi } @VisibleForTesting - Set getPartitionIds(List streamConfigs, IdealState idealState) { + Set getPartitionIds(List streamConfigs, + IdealState idealState) { return getPartitionIdsWithIdealState(streamConfigs, () -> idealState)._partitionIds; } private static class PartitionIdsWithIdealState { - private final Set _partitionIds; + private final Set _partitionIds; @Nullable private final IdealState _idealState; - private PartitionIdsWithIdealState(Set partitionIds, @Nullable IdealState idealState) { + private PartitionIdsWithIdealState(Set partitionIds, + @Nullable IdealState idealState) { _partitionIds = partitionIds; _idealState = idealState; } @@ -1252,7 +1254,7 @@ private PartitionIdsWithIdealState(Set partitionIds, @Nullable IdealSta private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List streamConfigs, Supplier idealStateSupplier) { - Set partitionIds = new HashSet<>(); + Set partitionIds = new HashSet<>(); boolean allPartitionIdsFetched = true; int numStreams = streamConfigs.size(); if (numStreams == 1) { @@ -1261,7 +1263,9 @@ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(partitionId, index)) + .map(partitionId -> new TopicPartitionId(index, partitionId)) .collect(Collectors.toSet())); } catch (UnsupportedOperationException ignored) { allPartitionIdsFetched = false; @@ -1303,7 +1307,8 @@ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List 1)); } } return new PartitionIdsWithIdealState(partitionIds, idealState); @@ -1425,13 +1430,14 @@ public void reduceSegmentSizeAndReset(LLCSegmentName llcSegmentName, int prevNum * @param realtimeTableName Realtime table name * @return Map from partition group id to the latest LLC realtime segment ZK metadata */ - private Map getLatestSegmentZKMetadataMap(String realtimeTableName) { + private Map getLatestSegmentZKMetadataMap( + String realtimeTableName) { List segments = getLLCSegments(realtimeTableName); - Map latestLLCSegmentNameMap = new HashMap<>(); + Map latestLLCSegmentNameMap = new HashMap<>(); for (String segmentName : segments) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId().getPartitionId(), + latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId(), (partitionId, latestLLCSegmentName) -> { if (latestLLCSegmentName == null) { return llcSegmentName; @@ -1445,8 +1451,8 @@ private Map getLatestSegmentZKMetadataMap(String rea }); } - Map latestSegmentZKMetadataMap = new HashMap<>(); - for (Map.Entry entry : latestLLCSegmentNameMap.entrySet()) { + Map latestSegmentZKMetadataMap = new HashMap<>(); + for (Map.Entry entry : latestLLCSegmentNameMap.entrySet()) { SegmentZKMetadata latestSegmentZKMetadata = getSegmentZKMetadata(realtimeTableName, entry.getValue().getSegmentName()); latestSegmentZKMetadataMap.put(entry.getKey(), latestSegmentZKMetadata); @@ -1584,8 +1590,7 @@ public static boolean isTopicPaused(IdealState idealState, String segmentName) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { return isTopicPaused(idealState, - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId( - llcSegmentName.getTopicPartitionId().getPartitionId())); + llcSegmentName.getTopicPartitionId().getTopicId()); } return false; } @@ -1634,7 +1639,7 @@ private void updateInstanceStatesForNewConsumingSegment(Map latestSegmentZKMetadataMap = getLatestSegmentZKMetadataMap(realtimeTableName); + Map latestSegmentZKMetadataMap = + getLatestSegmentZKMetadataMap(realtimeTableName); // Create a map from partition id to start offset // TODO: Directly return map from StreamMetadataProvider - Map partitionIdToStartOffset = Maps.newHashMapWithExpectedSize(numPartitions); + boolean hasMultipleStreams = streamConfigs.size() > 1; + Map partitionIdToStartOffset = + Maps.newHashMapWithExpectedSize(numPartitions); for (StreamMetadata streamMetadata : streamMetadataList) { for (PartitionGroupMetadata metadata : streamMetadata.getPartitionGroupMetadataList()) { - partitionIdToStartOffset.put(metadata.getPartitionGroupId(), metadata.getStartOffset()); + partitionIdToStartOffset.put(TopicPartitionId.fromPartitionGroupMetadata( + metadata.getPartitionGroupId(), hasMultipleStreams), metadata.getStartOffset()); } } // Create a map from partition id to the smallest stream offset - Map partitionIdToSmallestOffset = null; + Map partitionIdToSmallestOffset = null; if (offsetCriteria != null && offsetCriteria.equals(OffsetCriteria.SMALLEST_OFFSET_CRITERIA)) { partitionIdToSmallestOffset = partitionIdToStartOffset; } @@ -1828,12 +1837,12 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List entry : latestSegmentZKMetadataMap.entrySet()) { - int partitionId = entry.getKey(); + for (Map.Entry entry : latestSegmentZKMetadataMap.entrySet()) { + TopicPartitionId partitionId = entry.getKey(); SegmentZKMetadata latestSegmentZKMetadata = entry.getValue(); String latestSegmentName = latestSegmentZKMetadata.getSegmentName(); LLCSegmentName latestLLCSegmentName = new LLCSegmentName(latestSegmentName); - int streamConfigIdx = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(partitionId); + int streamConfigIdx = partitionId.getTopicId(); Map instanceStateMap = instanceStatesMap.get(latestSegmentName); if (instanceStateMap != null) { @@ -1852,7 +1861,8 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List> segmentEntry : instanceStatesMap.entrySet()) { if (segmentEntry.getValue().containsValue(SegmentStateModel.CONSUMING) - && new LLCSegmentName(segmentEntry.getKey()).getTopicPartitionId().getPartitionId() == partitionId) { + && new LLCSegmentName(segmentEntry.getKey()).getTopicPartitionId() + .equals(partitionId)) { previousConsumingSegment = segmentEntry.getKey(); break; } @@ -1983,7 +1994,8 @@ && new LLCSegmentName(segmentEntry.getKey()).getTopicPartitionId().getPartitionI // Set up new partitions if not exist for (StreamMetadata streamMetadata : streamMetadataList) { for (PartitionGroupMetadata partitionGroupMetadata : streamMetadata.getPartitionGroupMetadataList()) { - int partitionId = partitionGroupMetadata.getPartitionGroupId(); + TopicPartitionId partitionId = TopicPartitionId.fromPartitionGroupMetadata( + partitionGroupMetadata.getPartitionGroupId(), hasMultipleStreams); if (!latestSegmentZKMetadataMap.containsKey(partitionId)) { String newSegmentName = setupNewPartitionGroup(tableConfig, streamMetadata.getStreamConfig(), partitionGroupMetadata, @@ -2004,7 +2016,8 @@ private void createNewConsumingSegment(TableConfig tableConfig, StreamConfig str Map instancePartitionsMap, StreamPartitionMsgOffset startOffset) { int numReplicas = getNumReplicas(tableConfig, instancePartitions); LLCSegmentName latestLLCSegmentName = new LLCSegmentName(latestSegmentZKMetadata.getSegmentName()); - LLCSegmentName newLLCSegmentName = getNextLLCSegmentName(latestLLCSegmentName, currentTimeMs); + LLCSegmentName newLLCSegmentName = + getNextLLCSegmentName(latestLLCSegmentName, currentTimeMs); CommittingSegmentDescriptor committingSegmentDescriptor = new CommittingSegmentDescriptor(latestSegmentZKMetadata.getSegmentName(), startOffset.toString(), 0); createNewSegmentZKMetadata(tableConfig, streamConfig, newLLCSegmentName, currentTimeMs, committingSegmentDescriptor, @@ -2014,13 +2027,14 @@ private void createNewConsumingSegment(TableConfig tableConfig, StreamConfig str instancePartitionsMap); } - private Map fetchPartitionGroupIdToSmallestOffset(List streamConfigs, - IdealState idealState, Map latestSegmentZKMetadataMap) { + private Map fetchPartitionGroupIdToSmallestOffset( + List streamConfigs, IdealState idealState, + Map latestSegmentZKMetadataMap) { // Build consumption status from pre-computed ZK metadata map instead of rescanning IdealState (O(1) vs O(N)) List currentPartitionGroupConsumptionStatusList = buildPartitionGroupConsumptionStatusFromZKMetadata(latestSegmentZKMetadataMap, streamConfigs); - Map partitionGroupIdToSmallestOffset = new HashMap<>(); + Map partitionGroupIdToSmallestOffset = new HashMap<>(); for (StreamConfig streamConfig : streamConfigs) { OffsetCriteria originalOffsetCriteria = streamConfig.getOffsetCriteria(); streamConfig.setOffsetCriteria(OffsetCriteria.SMALLEST_OFFSET_CRITERIA); @@ -2034,7 +2048,8 @@ private Map fetchPartitionGroupIdToSmallestOf streamConfig.setOffsetCriteria(originalOffsetCriteria); for (StreamMetadata streamMetadata : streamMetadataList) { for (PartitionGroupMetadata metadata : streamMetadata.getPartitionGroupMetadataList()) { - partitionGroupIdToSmallestOffset.put(metadata.getPartitionGroupId(), metadata.getStartOffset()); + partitionGroupIdToSmallestOffset.put(TopicPartitionId.fromPartitionGroupMetadata( + metadata.getPartitionGroupId(), streamConfigs.size() > 1), metadata.getStartOffset()); } } } @@ -2047,27 +2062,29 @@ private Map fetchPartitionGroupIdToSmallestOf */ @VisibleForTesting List buildPartitionGroupConsumptionStatusFromZKMetadata( - Map latestSegmentZKMetadataMap, List streamConfigs) { + Map latestSegmentZKMetadataMap, + List streamConfigs) { List result = new ArrayList<>(latestSegmentZKMetadataMap.size()); int numStreams = streamConfigs.size(); if (numStreams == 1) { StreamPartitionMsgOffsetFactory offsetFactory = StreamConsumerFactoryProvider.create(streamConfigs.get(0)).createStreamMsgOffsetFactory(); - for (Map.Entry entry : latestSegmentZKMetadataMap.entrySet()) { - int partitionGroupId = entry.getKey(); + for (Map.Entry entry : latestSegmentZKMetadataMap.entrySet()) { + TopicPartitionId 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.toMultiTopicPinotPartitionId(), + llcSegmentName.getSequenceNumber(), offsetFactory.create(zkMetadata.getStartOffset()), zkMetadata.getEndOffset() != null ? offsetFactory.create(zkMetadata.getEndOffset()) : null, zkMetadata.getStatus().toString())); } } else { 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); + for (Map.Entry entry : latestSegmentZKMetadataMap.entrySet()) { + TopicPartitionId partitionGroupId = entry.getKey(); + int index = partitionGroupId.getTopicId(); + int streamPartitionId = partitionGroupId.getPartitionId(); SegmentZKMetadata zkMetadata = entry.getValue(); LLCSegmentName llcSegmentName = new LLCSegmentName(zkMetadata.getSegmentName()); StreamPartitionMsgOffsetFactory offsetFactory = offsetFactories[index]; @@ -2076,7 +2093,8 @@ List buildPartitionGroupConsumptionStatusFromZK StreamConsumerFactoryProvider.create(streamConfigs.get(index)).createStreamMsgOffsetFactory(); offsetFactories[index] = offsetFactory; } - result.add(new PartitionGroupConsumptionStatus(partitionGroupId, streamPartitionId, + result.add(new PartitionGroupConsumptionStatus( + partitionGroupId.toMultiTopicPinotPartitionId(), streamPartitionId, llcSegmentName.getSequenceNumber(), offsetFactory.create(zkMetadata.getStartOffset()), zkMetadata.getEndOffset() != null ? offsetFactory.create(zkMetadata.getEndOffset()) : null, @@ -2086,9 +2104,11 @@ List buildPartitionGroupConsumptionStatusFromZK return result; } - private StreamPartitionMsgOffset selectStartOffset(OffsetCriteria offsetCriteria, int partitionGroupId, - Map partitionGroupIdToStartOffset, - Map partitionGroupIdToSmallestStreamOffset, String tableName, + private StreamPartitionMsgOffset selectStartOffset(OffsetCriteria offsetCriteria, + TopicPartitionId partitionGroupId, + Map partitionGroupIdToStartOffset, + Map partitionGroupIdToSmallestStreamOffset, + String tableName, StreamPartitionMsgOffsetFactory offsetFactory, String startOffsetInSegmentZkMetadataStr) { if (offsetCriteria != null) { // use the fetched offset according to offset criteria @@ -2706,13 +2726,15 @@ private Set filterSegmentsToCommit(Set allConsumingSegments, } // partitionGroupIdsToCommitStr != null - Set partitionsToCommit = Arrays.stream(partitionGroupIdsToCommitStr.split(",")) + Set partitionsToCommit = Arrays.stream( + partitionGroupIdsToCommitStr.split(",")) .map(String::trim) - .map(Integer::parseInt) + .map(s -> TopicPartitionId.fromMultiTopicPinotPartitionId( + Integer.parseInt(s))) .collect(Collectors.toSet()); Set targetSegments = allConsumingSegments.stream() .filter(segmentName -> partitionsToCommit.contains( - new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId())) + new LLCSegmentName(segmentName).getTopicPartitionId())) .collect(Collectors.toSet()); Preconditions.checkState(!targetSegments.isEmpty(), "Cannot find segments to commit for partitions: %s", partitionGroupIdsToCommitStr); @@ -2900,8 +2922,7 @@ private Set findConsumingSegmentsOfTopics(IdealState idealState, List { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null && !topicIndices.contains( - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId( - llcSegmentName.getTopicPartitionId().getPartitionId()))) { + llcSegmentName.getTopicPartitionId().getTopicId())) { return; } for (String state : instanceToStateMap.values()) { diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java index 8cc2a397c27f..27115443d665 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/api/resources/PinotSegmentRestletResourceTest.java @@ -30,6 +30,7 @@ import org.apache.pinot.common.exception.TableNotFoundException; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; import org.apache.pinot.spi.utils.CommonConstants; import org.mockito.InjectMocks; @@ -87,31 +88,31 @@ public void testGetPartitionIdToSegmentsToDeleteMap() { when(znRecord.getMapFields()).thenReturn(segmentsToInstanceState); // Create the partition to oldest segment map - Map partitionToOldestSegment = Map.of( - 0, new LLCSegmentName(tableName, 0, 3, currentTime), - 1, new LLCSegmentName(tableName, 1, 5, currentTime) + Map partitionToOldestSegment = Map.of( + new TopicPartitionId(0), new LLCSegmentName(tableName, 0, 3, currentTime), + new TopicPartitionId(1), new LLCSegmentName(tableName, 1, 5, currentTime) ); // This map will be populated by the method - Map partitionIdToLatestSegment = new HashMap<>(); + Map partitionIdToLatestSegment = new HashMap<>(); // Create the expected response map - Map> expectedResponse = new HashMap<>(); - expectedResponse.put(0, + Map> expectedResponse = new HashMap<>(); + expectedResponse.put(new TopicPartitionId(0), getSegmentForPartition(tableName, 0, 3, 7, currentTime).stream().collect(Collectors.toSet())); - expectedResponse.put(1, + expectedResponse.put(new TopicPartitionId(1), getSegmentForPartition(tableName, 1, 5, 5, currentTime).stream().collect(Collectors.toSet())); // Call the method and check the result - Map> result = _pinotSegmentRestletResource.getPartitionIdToSegmentsToDeleteMap( + Map> result = _pinotSegmentRestletResource.getPartitionIdToSegmentsToDeleteMap( partitionToOldestSegment, segmentsToInstanceState.keySet(), partitionIdToLatestSegment); assertEquals(expectedResponse, result); // Verify that partitionIdToLatestSegment has been populated with the latest segment for each partition assertEquals(2, partitionIdToLatestSegment.size()); - assertEquals(9, partitionIdToLatestSegment.get(0).getSequenceNumber()); - assertEquals(9, partitionIdToLatestSegment.get(1).getSequenceNumber()); + assertEquals(9, partitionIdToLatestSegment.get(new TopicPartitionId(0)).getSequenceNumber()); + assertEquals(9, partitionIdToLatestSegment.get(new TopicPartitionId(1)).getSequenceNumber()); } @Test @@ -132,12 +133,12 @@ public void testGetPartitionIDToOldestSegment() { getSegmentForPartition(tableName + "fake", 0, 1, 3, currentTime)); // Segments with seq 1,2,3 for partition 0 // Create expected result map - Map expectedResult = new HashMap<>(); - expectedResult.put(0, new LLCSegmentName(tableName, 0, 3, currentTime)); - expectedResult.put(1, new LLCSegmentName(tableName, 1, 4, currentTime)); + Map expectedResult = new HashMap<>(); + expectedResult.put(new TopicPartitionId(0), new LLCSegmentName(tableName, 0, 3, currentTime)); + expectedResult.put(new TopicPartitionId(1), new LLCSegmentName(tableName, 1, 4, currentTime)); // Call the method and check the result - Map result = + Map result = _pinotSegmentRestletResource.getPartitionIDToOldestSegment(segments, idealStateSegmentSet); assertEquals(expectedResult, result); diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java index e300eb450629..b060730ceae3 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/PinotResourceManagerTest.java @@ -215,7 +215,7 @@ public void testAddingRealtimeTableSegmentsWithPartitionIdInZkMetadata() { int partitionId; LLCSegmentName llcSegmentName = LLCSegmentName.of(segment); if (llcSegmentName != null) { - partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); + partitionId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); } else { partitionId = Integer.parseInt(segment.substring(1, 2)); } 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 ebbb96cc0fd5..09a95f68b101 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 @@ -1860,7 +1860,7 @@ public void testAddRealtimeTableWithConsumingMetadata() for (String segmentName : idealState.getPartitionSet()) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); if (partitionGroupId == 0) { assertEquals(llcSegmentName.getSequenceNumber(), 5); } else if (partitionGroupId == 1) { diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java index a1802950f899..eab31f61f99f 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/assignment/segment/RealtimeReplicaGroupSegmentAssignmentTest.java @@ -543,7 +543,7 @@ public void testSubsetPartitionAssignment() { // Extract partition ID from segment name LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); + int partitionId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); // Verify the partition ID is one of our subset partitions boolean isSubsetPartition = false; diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinderTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinderTest.java index 223760f952a1..df1a0040b07f 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinderTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/realtime/MissingConsumingSegmentFinderTest.java @@ -22,6 +22,7 @@ import java.util.HashMap; import java.util.Map; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.LongMsgOffsetFactory; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; @@ -90,11 +91,11 @@ public void noMissingConsumingSegmentsScenario2() { idealStateMap.put("tableA__3__1__20220601T1200Z", Map.of("ServerX", "ONLINE", "ServerY", "ONLINE")); idealStateMap.put("tableA__3__2__20220601T1500Z", Map.of("ServerX", "CONSUMING", "ServerY", "CONSUMING")); - Map partitionGroupIdToLargestStreamOffsetMap = Map.of( - 0, new LongMsgOffset(1000), - 1, new LongMsgOffset(1001), - 2, new LongMsgOffset(1002), - 3, new LongMsgOffset(1003) + Map partitionGroupIdToLargestStreamOffsetMap = Map.of( + new TopicPartitionId(0), new LongMsgOffset(1000), + new TopicPartitionId(1), new LongMsgOffset(1001), + new TopicPartitionId(2), new LongMsgOffset(1002), + new TopicPartitionId(3), new LongMsgOffset(1003) ); Instant now = Instant.parse("2022-06-01T18:00:00.00Z"); @@ -127,11 +128,11 @@ public void noMissingConsumingSegmentsScenario3() { idealStateMap.put("tableA__3__0__20220601T0900Z", Map.of("ServerX", "ONLINE", "ServerY", "ONLINE")); idealStateMap.put("tableA__3__1__20220601T1200Z", Map.of("ServerX", "ONLINE", "ServerY", "ONLINE")); - Map partitionGroupIdToLargestStreamOffsetMap = Map.of( - 0, new LongMsgOffset(1000), - 1, new LongMsgOffset(701), - 2, new LongMsgOffset(1002), - 3, new LongMsgOffset(703) + Map partitionGroupIdToLargestStreamOffsetMap = Map.of( + new TopicPartitionId(0), new LongMsgOffset(1000), + new TopicPartitionId(1), new LongMsgOffset(701), + new TopicPartitionId(2), new LongMsgOffset(1002), + new TopicPartitionId(3), new LongMsgOffset(703) ); // setup segment metadata fetcher @@ -228,14 +229,14 @@ public void missingConsumingSegments() { idealStateMap.put("tableA__5__2__20220601T1500Z", Map.of("ServerX", "CONSUMING", "ServerY", "CONSUMING")); // partition 6 is a new partition and there's no consuming segment in ideal states for it - Map partitionGroupIdToLargestStreamOffsetMap = new HashMap<>(); - partitionGroupIdToLargestStreamOffsetMap.put(0, new LongMsgOffset(1000)); - partitionGroupIdToLargestStreamOffsetMap.put(1, new LongMsgOffset(1001)); - partitionGroupIdToLargestStreamOffsetMap.put(2, new LongMsgOffset(1002)); - partitionGroupIdToLargestStreamOffsetMap.put(3, new LongMsgOffset(1003)); - partitionGroupIdToLargestStreamOffsetMap.put(4, new LongMsgOffset(1004)); - partitionGroupIdToLargestStreamOffsetMap.put(5, new LongMsgOffset(1005)); - partitionGroupIdToLargestStreamOffsetMap.put(6, new LongMsgOffset(16)); + Map partitionGroupIdToLargestStreamOffsetMap = new HashMap<>(); + partitionGroupIdToLargestStreamOffsetMap.put(new TopicPartitionId(0), new LongMsgOffset(1000)); + partitionGroupIdToLargestStreamOffsetMap.put(new TopicPartitionId(1), new LongMsgOffset(1001)); + partitionGroupIdToLargestStreamOffsetMap.put(new TopicPartitionId(2), new LongMsgOffset(1002)); + partitionGroupIdToLargestStreamOffsetMap.put(new TopicPartitionId(3), new LongMsgOffset(1003)); + partitionGroupIdToLargestStreamOffsetMap.put(new TopicPartitionId(4), new LongMsgOffset(1004)); + partitionGroupIdToLargestStreamOffsetMap.put(new TopicPartitionId(5), new LongMsgOffset(1005)); + partitionGroupIdToLargestStreamOffsetMap.put(new TopicPartitionId(6), new LongMsgOffset(16)); // setup segment metadata fetcher SegmentZKMetadata m1 = mock(SegmentZKMetadata.class); 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 497dcc08224e..4c3f9010d071 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 @@ -65,6 +65,7 @@ import org.apache.pinot.common.restlet.resources.TableLLCSegmentUploadResponse; import org.apache.pinot.common.utils.FileUploadDownloadClient; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.common.utils.URIUtils; import org.apache.pinot.controller.ControllerConf; import org.apache.pinot.controller.helix.core.PinotHelixResourceManager; @@ -259,7 +260,7 @@ public void testSetUpNewTableWithExplicitSequenceNumbers() { // Verify segments are created with the explicit sequence numbers for (String segmentName : instanceStatesMap.keySet()) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); int sequence = llcSegmentName.getSequenceNumber(); if (partitionGroupId == 0) { assertEquals(sequence, 5); @@ -297,7 +298,7 @@ public void testSetUpNewTableDefaultSequenceNumberResolvesToZero() { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); assertEquals(llcSegmentName.getSequenceNumber(), 0, "Default sequence number -1 should resolve to 0 for partition " - + llcSegmentName.getTopicPartitionId().getPartitionId()); + + llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId()); } } @@ -779,7 +780,7 @@ private void testSetUpNewPartitions(FakePinotLLCRealtimeSegmentManager segmentMa assertTrue(segmentZKMetadataMap.containsKey(segmentName)); assertEquals(segmentZKMetadataMap.get(segmentName), oldSegmentZKMetadataMap.get(segmentName)); oldNumPartitions = Math.max(oldNumPartitions, - new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId() + 1); + new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId() + 1); } // Check that for new partition groups, each partition group should have exactly 1 new segment in CONSUMING @@ -788,7 +789,7 @@ private void testSetUpNewPartitions(FakePinotLLCRealtimeSegmentManager segmentMa Map> partitionGroupIdToSegmentsMap = new HashMap<>(); for (Map.Entry> entry : instanceStatesMap.entrySet()) { String segmentName = entry.getKey(); - int partitionGroupId = new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); + int partitionGroupId = new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId(); partitionGroupIdToSegmentsMap.computeIfAbsent(partitionGroupId, k -> new ArrayList<>()).add(segmentName); } for (int partitionGroupId = oldNumPartitions; partitionGroupId < segmentManager._numPartitions; @@ -1156,7 +1157,7 @@ private void verifyRepairs(FakePinotLLCRealtimeSegmentManager segmentManager, Li if (instanceStateMap.containsValue(SegmentStateModel.ONLINE) || instanceStateMap.containsValue( SegmentStateModel.CONSUMING)) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionsId = llcSegmentName.getTopicPartitionId().getPartitionId(); + int partitionsId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); Map sequenceNumberToSegmentMap = partitionGroupIdToSegmentsMap.get(partitionsId); int sequenceNumber = llcSegmentName.getSequenceNumber(); assertFalse(sequenceNumberToSegmentMap.containsKey(sequenceNumber)); @@ -1954,7 +1955,7 @@ public void testGetPartitionIds() segmentManager._numPartitions = 2; // Test empty ideal state - Set partitionIds = segmentManager.getPartitionIds(streamConfigs, idealState); + Set partitionIds = segmentManager.getPartitionIds(streamConfigs, idealState); Assert.assertEquals(partitionIds.size(), 2); partitionIds.clear(); @@ -1995,10 +1996,10 @@ public void testBuildPartitionGroupConsumptionStatusFromZKMetadataMatchesOrigina } // Build latestSegmentZKMetadataMap from the fake ZK metadata (same logic as getLatestSegmentZKMetadataMap) - Map latestSegmentZKMetadataMap = new HashMap<>(); + Map latestSegmentZKMetadataMap = new HashMap<>(); for (Map.Entry entry : segmentManager._segmentZKMetadataMap.entrySet()) { LLCSegmentName llcSegmentName = new LLCSegmentName(entry.getKey()); - int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionId = llcSegmentName.getTopicPartitionId(); latestSegmentZKMetadataMap.merge(partitionId, entry.getValue(), (existing, candidate) -> { int existingSeq = new LLCSegmentName(existing.getSegmentName()).getSequenceNumber(); int candidateSeq = new LLCSegmentName(candidate.getSegmentName()).getSequenceNumber(); @@ -2374,11 +2375,12 @@ public void testGetPartitionMetadataFromTableConfig() { // No SegmentPartitionConfig → null TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME) .setStreamConfigs(singleStreamConfigMap).build(); - assertNull(segmentManager.getPartitionMetadataFromTableConfig(tableConfig, 2, 4)); + assertNull(segmentManager.getPartitionMetadataFromTableConfig(tableConfig, new TopicPartitionId(2), 4)); // Single-stream: perStreamNumPartitions = numPartitionGroups tableConfig.getIndexingConfig().setSegmentPartitionConfig(partitionConfig); - SegmentPartitionMetadata metadata = segmentManager.getPartitionMetadataFromTableConfig(tableConfig, 2, 4); + SegmentPartitionMetadata metadata = + segmentManager.getPartitionMetadataFromTableConfig(tableConfig, new TopicPartitionId(2), 4); assertNotNull(metadata); ColumnPartitionMetadata colMetadata = metadata.getColumnPartitionMap().get("col"); assertEquals(colMetadata.getNumPartitions(), 4); @@ -2394,16 +2396,17 @@ public void testGetPartitionMetadataFromTableConfig() { .build(); multiStreamTableConfig.getIndexingConfig().setSegmentPartitionConfig(partitionConfig); - // Stream 0, partition 2: Pinot partition ID = 0 * 10000 + 2 = 2 - metadata = segmentManager.getPartitionMetadataFromTableConfig(multiStreamTableConfig, 2, 8); + // Stream 0, partition 2: TopicPartitionId(topicId=0, partitionId=2) + metadata = segmentManager.getPartitionMetadataFromTableConfig(multiStreamTableConfig, + new TopicPartitionId(0, 2), 8); assertNotNull(metadata); colMetadata = metadata.getColumnPartitionMap().get("col"); assertEquals(colMetadata.getNumPartitions(), 4, "Multi-stream partition count must be per-stream (numPartitionGroups / numStreams), not total"); assertEquals(colMetadata.getPartitions(), Set.of(2)); - // Stream 1, partition 3: Pinot partition ID = 1 * 10000 + 3 = 10003 - int stream1Partition3 = IngestionConfigUtils.PARTITION_PADDING_OFFSET + 3; + // Stream 1, partition 3: TopicPartitionId(topicId=1, partitionId=3) + TopicPartitionId stream1Partition3 = new TopicPartitionId(1, 3); metadata = segmentManager.getPartitionMetadataFromTableConfig(multiStreamTableConfig, stream1Partition3, 8); assertNotNull(metadata); colMetadata = metadata.getColumnPartitionMap().get("col"); @@ -2411,7 +2414,9 @@ public void testGetPartitionMetadataFromTableConfig() { assertEquals(colMetadata.getPartitions(), Set.of(3)); // Multi-stream, uneven distribution: numPartitionGroups not divisible by numStreams → null - assertNull(segmentManager.getPartitionMetadataFromTableConfig(multiStreamTableConfig, 0, 7), + assertNull( + segmentManager.getPartitionMetadataFromTableConfig( + multiStreamTableConfig, new TopicPartitionId(0, 0), 7), "Uneven partition distribution across streams must return null"); } diff --git a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java index f51eaf890181..fb2498c0db8a 100644 --- a/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java +++ b/pinot-controller/src/test/java/org/apache/pinot/controller/helix/core/rebalance/TableRebalancerTest.java @@ -48,7 +48,7 @@ public class TableRebalancerTest { private static final TableRebalancer.PartitionIdFetcher DUMMY_PARTITION_FETCHER = segmentName -> 0; private static final TableRebalancer.PartitionIdFetcher SIMPLE_PARTITION_FETCHER = segmentName -> { LLCSegmentName name = LLCSegmentName.of(segmentName); - return name == null ? -1 : name.getTopicPartitionId().getPartitionId(); + return name == null ? -1 : name.getTopicPartitionId().toMultiTopicPinotPartitionId(); }; private static final TableRebalancer.DataLossRiskAssessor DEFAULT_DATA_LOSS_RISK_ASSESSOR = diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java index f078e753262d..9e63818449c5 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinator.java @@ -33,6 +33,7 @@ import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.common.metrics.ServerTimer; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.common.utils.helix.HelixHelper; import org.apache.pinot.spi.config.table.ingestion.StreamIngestionConfig; import org.apache.pinot.spi.utils.CommonConstants.Helix.StateModel.SegmentStateModel; @@ -102,11 +103,13 @@ public void acquire(LLCSegmentName llcSegmentName) LOGGER.warn("Failed to acquire consumer semaphore for segment: {} in: {}ms. Retrying.", segmentName, waitTimeMs); _serverMetrics.setValueOfPartitionGauge(_realtimeTableDataManager.getTableName(), - llcSegmentName.getTopicPartitionId().getPartitionId(), ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, waitTimeMs); + llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(), + ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, waitTimeMs); } } finally { _serverMetrics.setValueOfPartitionGauge(_realtimeTableDataManager.getTableName(), - llcSegmentName.getTopicPartitionId().getPartitionId(), ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, 0); + llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(), + ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, 0); } } @@ -197,7 +200,7 @@ int getPreviousSegmentSequenceNumberFromIdealState(LLCSegmentName currentSegment // so that it can always pass the check. int maxSequenceNumberBelowCurrentSegment = -1; String instanceId = _realtimeTableDataManager.getServerInstance(); - int partitionId = currentSegment.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionId = currentSegment.getTopicPartitionId(); int currentSequenceNumber = currentSegment.getSequenceNumber(); for (Map.Entry> entry : getSegmentAssignment().entrySet()) { @@ -215,7 +218,7 @@ int getPreviousSegmentSequenceNumberFromIdealState(LLCSegmentName currentSegment continue; } - if (llcSegmentName.getTopicPartitionId().getPartitionId() != partitionId) { + if (!llcSegmentName.getTopicPartitionId().equals(partitionId)) { // ignore segments of different partitions. continue; } diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java index ed4b7236af9d..8355b0ad975e 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTracker.java @@ -24,6 +24,7 @@ import java.io.IOException; import java.time.Clock; import java.util.ArrayList; +import java.util.HashMap; import java.util.HashSet; import java.util.List; import java.util.Map; @@ -40,6 +41,7 @@ import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.common.metrics.ServerTimer; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.StreamConfig; import org.apache.pinot.spi.stream.StreamConsumerFactory; @@ -140,21 +142,22 @@ void update(long ingestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffs // This map is accessed by: // 1. _ingestionDelayTrackingScheduler thread. // 2. All threads which can removes metrics - Consumer thread, Helix Thread, Server API Thread, etc. - private final Map _partitionsTracked = new ConcurrentHashMap<>(); + private final Map _partitionsTracked = new ConcurrentHashMap<>(); // Map to hold the ingestion info reported by the consumer. // This map is accessed by: // 1. _ingestionDelayTrackingScheduler thread. // 2. All threads which can removes metrics - Consumer thread, Helix Thread, Server API Thread, etc. // 3. Consumer thread when it updates the ingestion info of the partition. - private final Map _ingestionInfoMap = new ConcurrentHashMap<>(); - private final Map _partitionsMarkedForVerification = new ConcurrentHashMap<>(); + private final Map _ingestionInfoMap = new ConcurrentHashMap<>(); + private final Map _partitionsMarkedForVerification = new ConcurrentHashMap<>(); private final ScheduledExecutorService _metricsRemovalScheduler; private final ScheduledExecutorService _ingestionDelayTrackingScheduler; private Clock _clock = Clock.systemUTC(); - protected volatile Map _partitionIdToLatestOffset; - protected volatile ConcurrentHashMap _partitionsHostedByThisServer = new ConcurrentHashMap<>(); + protected volatile Map _partitionIdToLatestOffset; + protected volatile ConcurrentHashMap _partitionsHostedByThisServer = + new ConcurrentHashMap<>(); // Map of StreamMetadataProvider to fetch upstream latest stream offset (Table can have multiple upstream topics) // This map is accessed by: // 1. _ingestionDelayTrackingScheduler thread. @@ -232,7 +235,8 @@ private void trackIngestionDelay() { // or once the table data manager has been shutdown. return; } - Set partitionsHosted = new HashSet<>(_partitionsHostedByThisServer.keySet()); + Set partitionsHosted = + new HashSet<>(_partitionsHostedByThisServer.keySet()); if (_ingestionInfoMap.size() > partitionsHosted.size()) { // In-case new partition got assigned to the server before _partitionsHostedByThisServer was updated. @@ -245,7 +249,7 @@ private void trackIngestionDelay() { updateLatestStreamOffset(partitionsHosted); - for (Integer partitionId : partitionsHosted) { + for (TopicPartitionId partitionId : partitionsHosted) { _partitionsTracked.computeIfAbsent(partitionId, k -> { // Check below condition in-case partition was stopped being tracked. Eg: due to manual operations like - // calling remove ingestion metrics API @@ -267,9 +271,12 @@ private void trackIngestionDelay() { } @VisibleForTesting - void updateLatestStreamOffset(Set partitionsHosted) { - Map> streamIndexToStreamPartitionIds = - IngestionConfigUtils.getStreamConfigIndexToStreamPartitions(partitionsHosted); + void updateLatestStreamOffset(Set partitionsHosted) { + Map> streamIndexToStreamPartitionIds = new HashMap<>(); + for (TopicPartitionId tpId : partitionsHosted) { + streamIndexToStreamPartitionIds.computeIfAbsent(tpId.getTopicId(), k -> new HashSet<>()) + .add(tpId.getPartitionId()); + } if (streamIndexToStreamPartitionIds.size() > _streamConfigIndexToStreamMetadataProvider.size()) { // There might be a new stream config added or need to retry creation of streamMetadataProvider for a stream @@ -293,15 +300,19 @@ void updateLatestStreamOffset(Set partitionsHosted) { streamMetadataProvider.fetchLatestStreamOffset(streamPartitionIds, LATEST_STREAM_OFFSET_FETCH_TIMEOUT_MS); if (streamIndex > 0) { - // Need to convert stream partition Ids back to pinot partition Ids. + // Need to convert stream partition Ids back to TopicPartitionId keys. for (Map.Entry latestOffsetEntry : streamPartitionIdToLatestOffset.entrySet()) { _partitionIdToLatestOffset.put( - IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(latestOffsetEntry.getKey(), - streamIndex), latestOffsetEntry.getValue()); + new TopicPartitionId(streamIndex, latestOffsetEntry.getKey()), + latestOffsetEntry.getValue()); } } else { - _partitionIdToLatestOffset.putAll(streamPartitionIdToLatestOffset); + for (Map.Entry e + : streamPartitionIdToLatestOffset.entrySet()) { + _partitionIdToLatestOffset.put( + new TopicPartitionId(e.getKey()), e.getValue()); + } } } catch (Exception e) { _serverMetrics.addMeteredTableValue(_realTimeTableDataManager.getTableName(), @@ -331,7 +342,7 @@ public Thread newThread(Runnable r) { * * @param partitionId partition ID which we should stop tracking. */ - private void removePartitionId(int partitionId) { + private void removePartitionId(TopicPartitionId partitionId) { _partitionsHostedByThisServer.remove(partitionId); _ingestionInfoMap.remove(partitionId); _partitionsTracked.computeIfPresent(partitionId, (k, v) -> { @@ -347,9 +358,10 @@ private void removePartitionId(int partitionId) { * Helper functions that creates a list of all the partitions that are marked for verification and whose * timeouts are expired. This helps us optimize checks of the ideal state. */ - private List getPartitionsToBeVerified() { - List partitionsToVerify = new ArrayList<>(); - for (Map.Entry entry : _partitionsMarkedForVerification.entrySet()) { + private List getPartitionsToBeVerified() { + List partitionsToVerify = new ArrayList<>(); + for (Map.Entry entry + : _partitionsMarkedForVerification.entrySet()) { long timeMarked = _clock.millis() - entry.getValue(); if (timeMarked > PARTITION_TIMEOUT_MS) { // Partition must be verified @@ -370,45 +382,57 @@ void setClock(Clock clock) { } @VisibleForTesting - void createMetrics(int partitionId) { - int streamConfigIndex = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(partitionId); - StreamMetadataProvider streamMetadataProvider = _streamConfigIndexToStreamMetadataProvider.get(streamConfigIndex); + void createMetrics(TopicPartitionId partitionId) { + int streamConfigIndex = partitionId.getTopicId(); + StreamMetadataProvider streamMetadataProvider = + _streamConfigIndexToStreamMetadataProvider.get(streamConfigIndex); + int pinotPartitionId = partitionId.toMultiTopicPinotPartitionId(); if (streamMetadataProvider != null && streamMetadataProvider.supportsOffsetLag()) { - _serverMetrics.setOrUpdatePartitionGauge(_metricName, partitionId, ServerGauge.REALTIME_INGESTION_OFFSET_LAG, + _serverMetrics.setOrUpdatePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_OFFSET_LAG, () -> getPartitionIngestionOffsetLag(partitionId)); - _serverMetrics.setOrUpdatePartitionGauge(_metricName, partitionId, - ServerGauge.REALTIME_INGESTION_CONSUMING_OFFSET, () -> getPartitionIngestionConsumingOffset(partitionId)); + _serverMetrics.setOrUpdatePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_CONSUMING_OFFSET, + () -> getPartitionIngestionConsumingOffset(partitionId)); - _serverMetrics.setOrUpdatePartitionGauge(_metricName, partitionId, - ServerGauge.REALTIME_INGESTION_UPSTREAM_OFFSET, () -> getLatestPartitionOffset(partitionId)); + _serverMetrics.setOrUpdatePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_UPSTREAM_OFFSET, + () -> getLatestPartitionOffset(partitionId)); } - _serverMetrics.setOrUpdatePartitionGauge(_metricName, partitionId, ServerGauge.REALTIME_INGESTION_DELAY_MS, + _serverMetrics.setOrUpdatePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_DELAY_MS, () -> getPartitionIngestionDelayMs(partitionId)); - _serverMetrics.setOrUpdatePartitionGauge(_metricName, partitionId, - ServerGauge.END_TO_END_REALTIME_INGESTION_DELAY_MS, () -> getPartitionEndToEndIngestionDelayMs(partitionId)); - _serverMetrics.setOrUpdatePartitionGauge(_metricName, partitionId, - ServerGauge.REALTIME_INGESTION_DELAY_REPORTING_STATUS, () -> getPartitionIngestionReportingStatus(partitionId)); + _serverMetrics.setOrUpdatePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.END_TO_END_REALTIME_INGESTION_DELAY_MS, + () -> getPartitionEndToEndIngestionDelayMs(partitionId)); + _serverMetrics.setOrUpdatePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_DELAY_REPORTING_STATUS, + () -> getPartitionIngestionReportingStatus(partitionId)); LOGGER.info("Successfully created ingestion metrics for partition id: {}", partitionId); } - private void removeMetrics(int partitionId) { - int streamConfigIndex = IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(partitionId); + private void removeMetrics(TopicPartitionId partitionId) { + int streamConfigIndex = partitionId.getTopicId(); + int pinotPartitionId = partitionId.toMultiTopicPinotPartitionId(); StreamMetadataProvider streamMetadataProvider = _streamConfigIndexToStreamMetadataProvider.get(streamConfigIndex); // Remove all metrics associated with this partition if (streamMetadataProvider != null && streamMetadataProvider.supportsOffsetLag()) { - _serverMetrics.removePartitionGauge(_metricName, partitionId, ServerGauge.REALTIME_INGESTION_OFFSET_LAG); - _serverMetrics.removePartitionGauge(_metricName, partitionId, ServerGauge.REALTIME_INGESTION_UPSTREAM_OFFSET); - _serverMetrics.removePartitionGauge(_metricName, partitionId, + _serverMetrics.removePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_OFFSET_LAG); + _serverMetrics.removePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_UPSTREAM_OFFSET); + _serverMetrics.removePartitionGauge(_metricName, pinotPartitionId, ServerGauge.REALTIME_INGESTION_CONSUMING_OFFSET); } - _serverMetrics.removePartitionGauge(_metricName, partitionId, ServerGauge.REALTIME_INGESTION_DELAY_MS); - _serverMetrics.removePartitionGauge(_metricName, partitionId, + _serverMetrics.removePartitionGauge(_metricName, pinotPartitionId, + ServerGauge.REALTIME_INGESTION_DELAY_MS); + _serverMetrics.removePartitionGauge(_metricName, pinotPartitionId, ServerGauge.END_TO_END_REALTIME_INGESTION_DELAY_MS); - _serverMetrics.removePartitionGauge(_metricName, partitionId, + _serverMetrics.removePartitionGauge(_metricName, pinotPartitionId, ServerGauge.REALTIME_INGESTION_DELAY_REPORTING_STATUS); LOGGER.info("Successfully removed ingestion metrics for partition id: {}", partitionId); @@ -425,7 +449,8 @@ private void removeMetrics(int partitionId) { * {@link StreamMessageMetadata}) * @param currentOffset offset of the last consumed message (from {@link StreamMessageMetadata}) */ - public void updateMetrics(String segmentName, int partitionId, long ingestionTimeMs, + public void updateMetrics(String segmentName, TopicPartitionId partitionId, + long ingestionTimeMs, long firstStreamIngestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffset) { if (!_isServerReadyToServeQueries.getAsBoolean() || _realTimeTableDataManager.isShutDown()) { // Do not update the ingestion delay metrics during server startup period @@ -460,7 +485,7 @@ public void updateMetrics(String segmentName, int partitionId, long ingestionTim * * @param partitionId partition id that we should stop tracking. */ - public void stopTrackingPartition(int partitionId) { + public void stopTrackingPartition(TopicPartitionId partitionId) { removePartitionId(partitionId); } @@ -470,9 +495,10 @@ public void stopTrackingPartition(int partitionId) { * * @return Set of partitionIds for which ingestion metrics were removed. */ - public Set stopTrackingAllPartitions() { - Set removedPartitionIds = new HashSet<>(_ingestionInfoMap.keySet()); - for (Integer partitionId : _ingestionInfoMap.keySet()) { + public Set stopTrackingAllPartitions() { + Set removedPartitionIds = + new HashSet<>(_ingestionInfoMap.keySet()); + for (TopicPartitionId partitionId : _ingestionInfoMap.keySet()) { removePartitionId(partitionId); } return removedPartitionIds; @@ -485,7 +511,8 @@ public Set stopTrackingAllPartitions() { */ public void stopTrackingPartition(String segmentName) { _segmentsToIgnore.put(segmentName, true); - removePartitionId(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); + removePartitionId( + new LLCSegmentName(segmentName).getTopicPartitionId()); } /** @@ -501,27 +528,31 @@ public void timeoutInactivePartitions() { } // Check if we have any partition to verify, else don't make the call to check ideal state as that // involves network traffic and may be inefficient. - List partitionsToVerify = getPartitionsToBeVerified(); + List partitionsToVerify = getPartitionsToBeVerified(); if (partitionsToVerify.isEmpty()) { // Don't make the call to getHostedPartitionsGroupIds() as it involves checking ideal state. return; } - Set partitionsHostedByThisServer; + Set partitionsHostedByThisServer; try { - partitionsHostedByThisServer = _realTimeTableDataManager.getHostedPartitionsGroupIds(); + partitionsHostedByThisServer = + _realTimeTableDataManager.getHostedPartitionsGroupIds(); } catch (Exception e) { - LOGGER.error("Failed to get partitions hosted by this server, table={}, exception={}:{}", _tableNameWithType, - e.getClass(), e.getMessage()); + LOGGER.error( + "Failed to get partitions hosted by this server, table={}, " + + "exception={}:{}", + _tableNameWithType, e.getClass(), e.getMessage()); return; } - for (int partitionId : partitionsToVerify) { + for (TopicPartitionId partitionId : partitionsToVerify) { if (!partitionsHostedByThisServer.contains(partitionId)) { // Partition is not hosted in this server anymore, stop tracking it removePartitionId(partitionId); } } - ConcurrentHashMap newMap = new ConcurrentHashMap<>(); + ConcurrentHashMap newMap = + new ConcurrentHashMap<>(); partitionsHostedByThisServer.forEach(p -> newMap.put(p, true)); _partitionsHostedByThisServer = newMap; } @@ -535,7 +566,8 @@ public void markPartitionForVerification(String segmentName) { // Do not update the tracker state during server startup period or if the segment is marked to be ignored return; } - _partitionsMarkedForVerification.put(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(), + _partitionsMarkedForVerification.put( + new LLCSegmentName(segmentName).getTopicPartitionId(), _clock.millis()); } @@ -546,7 +578,7 @@ public void markPartitionForVerification(String segmentName) { * * @return ingestion delay timestamp in milliseconds for the given partition ID. */ - public long getPartitionIngestionTimeMs(int partitionId) { + public long getPartitionIngestionTimeMs(TopicPartitionId partitionId) { IngestionInfo ingestionInfo = _ingestionInfoMap.get(partitionId); return ingestionInfo != null ? ingestionInfo._ingestionTimeMs : Long.MIN_VALUE; } @@ -560,7 +592,7 @@ public long getPartitionIngestionTimeMs(int partitionId) { * or null if first stream ingestion time is not available for the partition. */ @Nullable - public Long getPartitionEndToEndIngestionDelayMs(int partitionId) { + public Long getPartitionEndToEndIngestionDelayMs(TopicPartitionId partitionId) { IngestionInfo ingestionInfo = _ingestionInfoMap.get(partitionId); if (ingestionInfo == null || ingestionInfo._firstStreamIngestionTimeMs < 0) { return null; @@ -580,7 +612,7 @@ public Long getPartitionEndToEndIngestionDelayMs(int partitionId) { * or null if ingestion time is not available for the partition. */ @Nullable - public Long getPartitionIngestionDelayMs(int partitionId) { + public Long getPartitionIngestionDelayMs(TopicPartitionId partitionId) { IngestionInfo ingestionInfo = _ingestionInfoMap.get(partitionId); if (ingestionInfo == null || ingestionInfo._ingestionTimeMs < 0) { return null; @@ -591,6 +623,7 @@ public Long getPartitionIngestionDelayMs(int partitionId) { return Math.max(agedIngestionDelayMs, 0); } + /** /** * Method to get if ingestion delay data is available for the given partition (i.e. * ingestion info has been reported) @@ -599,7 +632,7 @@ public Long getPartitionIngestionDelayMs(int partitionId) { * * @return 1 if ingestion delay data is available, 0 otherwise */ - public long getPartitionIngestionReportingStatus(int partitionId) { + public long getPartitionIngestionReportingStatus(TopicPartitionId partitionId) { IngestionInfo ingestionInfo = _ingestionInfoMap.get(partitionId); if (ingestionInfo == null) { return 0; @@ -616,7 +649,7 @@ public long getPartitionIngestionReportingStatus(int partitionId) { * @param partitionId partition for which the ingestion lag is computed * @return offset lag for the given partition */ - public long getPartitionIngestionOffsetLag(int partitionId) { + public long getPartitionIngestionOffsetLag(TopicPartitionId partitionId) { StreamPartitionMsgOffset latestOffset = _partitionIdToLatestOffset.get(partitionId); if (latestOffset == null) { return 0; @@ -645,7 +678,7 @@ public long getPartitionIngestionOffsetLag(int partitionId) { * @param partitionId partition for which the consuming offset was retrieved * @return consuming offset value for the given partition, or {@code 0} if no ingestion info is available */ - public long getPartitionIngestionConsumingOffset(int partitionId) { + public long getPartitionIngestionConsumingOffset(TopicPartitionId partitionId) { IngestionInfo ingestionInfo = _ingestionInfoMap.get(partitionId); if (ingestionInfo == null) { return 0; @@ -668,7 +701,7 @@ public long getPartitionIngestionConsumingOffset(int partitionId) { * @param partitionId partition for which the latest upstream offset is retrieved * @return latest offset value for the given partition, or {@code 0} if not available */ - public long getLatestPartitionOffset(int partitionId) { + public long getLatestPartitionOffset(TopicPartitionId partitionId) { StreamPartitionMsgOffset latestOffset = _partitionIdToLatestOffset.get(partitionId); if (latestOffset == null) { return 0; @@ -699,7 +732,7 @@ public void shutdown() { return; } // Remove partitions so their related metrics get uninstalled. - for (Integer partitionId : _ingestionInfoMap.keySet()) { + for (TopicPartitionId partitionId : _ingestionInfoMap.keySet()) { removePartitionId(partitionId); } } 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 a922c84b2ff2..242958a7892b 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 @@ -52,6 +52,7 @@ import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.common.utils.PauselessConsumptionUtils; import org.apache.pinot.common.utils.TarCompressionUtils; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.core.data.manager.SegmentOperationsTaskContext; import org.apache.pinot.core.data.manager.SegmentOperationsTaskType; import org.apache.pinot.core.data.manager.realtime.RealtimeConsumptionRateManager.ConsumptionRateLimiter; @@ -306,13 +307,13 @@ public void deleteSegmentFile() { private static final int MAX_TIME_FOR_CONSUMING_TO_ONLINE_IN_SECONDS = 31; private Thread _consumerThread; - // _partitionGroupId represents the Pinot's internal partition number which will eventually be used as part of - // segment name. + // _partitionGroupId is a TopicPartitionId that encodes both the topic index and stream partition. It is used + // to derive the Pinot partition number (for segment names) via toMultiTopicPinotPartitionId(). // _streamPartitionId represents the partition number in the stream topic, which could be derived from the // _partitionGroupId and identify which partition of the stream topic this consumer is consuming from. // Note that in traditional single topic ingestion mode, those two concepts were identical which got separated // in multi-topic ingestion mode. - private final int _partitionGroupId; + private final TopicPartitionId _partitionGroupId; private final int _streamPartitionId; private final PartitionGroupConsumptionStatus _partitionGroupConsumptionStatus; final String _clientId; @@ -1129,7 +1130,7 @@ protected void buildSegmentForCommit(long buildTimeLeaseMs) * Returns the current offset for the partition group. */ public Map getPartitionToCurrentOffset() { - return Map.of(String.valueOf(_partitionGroupId), _currentOffset.toString()); + return Map.of(_partitionGroupId.toString(), _currentOffset.toString()); } /** @@ -1151,7 +1152,7 @@ public long getLastConsumedTimestamp() { */ public Map getConsumerPartitionState( @Nullable StreamPartitionMsgOffset latestMsgOffset) { - String partitionGroupId = String.valueOf(_partitionGroupId); + String partitionGroupId = _partitionGroupId.toString(); return Map.of(partitionGroupId, new ConsumerPartitionState(partitionGroupId, getCurrentOffset(), getLastConsumedTimestamp(), latestMsgOffset, _lastRowMetadata)); @@ -1784,21 +1785,17 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf String timeColumnName = tableConfig.getValidationConfig().getTimeColumnName(); // TODO Validate configs IndexingConfig indexingConfig = _tableConfig.getIndexingConfig(); - _partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); + _partitionGroupId = llcSegmentName.getTopicPartitionId(); List> streamConfigMaps = IngestionConfigUtils.getStreamConfigMaps(_tableConfig); int numStreams = streamConfigMaps.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. - _streamPartitionId = _partitionGroupId; + _streamPartitionId = _partitionGroupId.getPartitionId(); _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); + _streamPartitionId = _partitionGroupId.getPartitionId(); + int index = _partitionGroupId.getTopicId(); + Preconditions.checkState(numStreams > index, + "Cannot find stream config of index: %s for table: %s", index, _tableNameWithType); _streamConfig = new StreamConfig(_tableNameWithType, streamConfigMaps.get(index)); } _streamConsumerFactory = StreamConsumerFactoryProvider.create(_streamConfig); @@ -1806,7 +1803,8 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf String streamTopic = _streamConfig.getTopicName(); _segmentNameStr = _segmentZKMetadata.getSegmentName(); _partitionGroupConsumptionStatus = - new PartitionGroupConsumptionStatus(_partitionGroupId, _streamPartitionId, llcSegmentName.getSequenceNumber(), + new PartitionGroupConsumptionStatus(_partitionGroupId.toMultiTopicPinotPartitionId(), + _streamPartitionId, llcSegmentName.getSequenceNumber(), _streamPartitionMsgOffsetFactory.create(_segmentZKMetadata.getStartOffset()), _segmentZKMetadata.getEndOffset() == null ? null : _streamPartitionMsgOffsetFactory.create(_segmentZKMetadata.getEndOffset()), @@ -1875,7 +1873,8 @@ public RealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, TableConf .setTextIndexConfig(consumingIndexLoadingConfig.getMultiColTextIndexConfig()) .setDropRecordOnPartitionMismatch(ingestionConfig != null && ingestionConfig.getStreamIngestionConfig() != null - && ingestionConfig.getStreamIngestionConfig().isDropRecordOnPartitionMismatch()); + && ingestionConfig.getStreamIngestionConfig().isDropRecordOnPartitionMismatch()) + .setHasMultipleStreams(IngestionConfigUtils.hasMultipleStreams(_tableConfig)); // Create message decoder Set fieldsToRead = IngestionUtils.getFieldsForRecordExtractor(_tableConfig, _schema); @@ -2082,7 +2081,7 @@ private void setPartitionParameters(RealtimeSegmentConfig.Builder realtimeSegmen realtimeSegmentConfigBuilder.setPartitionFunction( PartitionFunctionFactory.getPartitionFunction(partitionFunctionName, numPartitions, columnPartitionConfig.getFunctionConfig())); - realtimeSegmentConfigBuilder.setPartitionId(_partitionGroupId); + realtimeSegmentConfigBuilder.setPartitionId(_partitionGroupId.toMultiTopicPinotPartitionId()); } else { _segmentLogger.warn("Cannot partition on multiple columns: {}", columnPartitionMap.keySet()); } @@ -2157,8 +2156,8 @@ private void updateIngestionMetrics(StreamMessageMetadata metadata) { */ private void setIngestionDelayToZero() { long currentTimeMs = System.currentTimeMillis(); - _realtimeTableDataManager.updateIngestionMetrics(_segmentNameStr, _partitionGroupId, currentTimeMs, currentTimeMs, - null); + _realtimeTableDataManager.updateIngestionMetrics(_segmentNameStr, _partitionGroupId, + currentTimeMs, currentTimeMs, null); } // This should be done during commit? We may not always commit when we build a segment.... diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java index 3a720f4ce3c1..7fd35cdf3c98 100644 --- a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java +++ b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/realtime/RealtimeTableDataManager.java @@ -53,6 +53,7 @@ import org.apache.pinot.common.restlet.resources.SegmentErrorInfo; import org.apache.pinot.common.utils.LLCSegmentName; import org.apache.pinot.common.utils.SegmentUtils; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.core.data.manager.BaseTableDataManager; import org.apache.pinot.core.util.PeerServerSegmentFinder; import org.apache.pinot.segment.local.data.manager.SegmentDataManager; @@ -108,7 +109,8 @@ public class RealtimeTableDataManager extends BaseTableDataManager { // allows releasing the lock from a different thread. // The consumer coordinators will stay in the map even if the consuming partitions moved to a different server. We // expect a small number of consumer coordinators, so it should be fine to not remove them. - private final Map _partitionIdToConsumerCoordinatorMap = new ConcurrentHashMap<>(); + private final Map _partitionIdToConsumerCoordinatorMap = + new ConcurrentHashMap<>(); // The old name of the stats file used to be stats.ser which we changed when we moved all packages // from com.linkedin to org.apache because of not being able to deserialize the old files using the newer classes private static final String STATS_FILE_NAME = "segment-stats.ser"; @@ -332,15 +334,16 @@ protected void doShutdown() { * Updates the ingestion metrics for the given partition. * * @param segmentName name of the consuming segment - * @param partitionId partition id of the consuming segment (directly passed in to avoid parsing - * the segment name) + * @param partitionId partition id of the consuming segment (directly passed + * in to avoid parsing the segment name) * @param ingestionTimeMs ingestion time of the last consumed message (from * {@link StreamMessageMetadata}) * @param firstStreamIngestionTimeMs ingestion time of the last consumed message in the first stream (from * {@link StreamMessageMetadata}) * @param currentOffset offset of the last consumed message (from {@link StreamMessageMetadata}) */ - public void updateIngestionMetrics(String segmentName, int partitionId, long ingestionTimeMs, + public void updateIngestionMetrics(String segmentName, TopicPartitionId partitionId, + long ingestionTimeMs, long firstStreamIngestionTimeMs, @Nullable StreamPartitionMsgOffset currentOffset) { _ingestionDelayTracker.updateMetrics(segmentName, partitionId, ingestionTimeMs, firstStreamIngestionTimeMs, currentOffset); @@ -352,7 +355,7 @@ public void updateIngestionMetrics(String segmentName, int partitionId, long ing */ public long getPartitionIngestionTimeMs(String segmentName) { return _ingestionDelayTracker.getPartitionIngestionTimeMs( - new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); + new LLCSegmentName(segmentName).getTopicPartitionId()); } /** @@ -374,7 +377,7 @@ public void removeIngestionMetrics(String segmentName) { public void onConsumingToDropped(String segmentName) { // NOTE: No need to mark segment ignored here because it should have already been dropped. _ingestionDelayTracker.stopTrackingPartition( - new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); + new LLCSegmentName(segmentName).getTopicPartitionId()); } /** @@ -409,13 +412,13 @@ public StreamMetadataProvider getStreamMetadataProvider(RealtimeSegmentDataManag * Returns all partitionGroupIds for the partitions hosted by this server for current table. * @apiNote this involves Zookeeper read and should not be used frequently due to efficiency concerns. */ - public Set getHostedPartitionsGroupIds() { - Set partitionsHostedByThisServer = new HashSet<>(); + public Set getHostedPartitionsGroupIds() { + Set partitionsHostedByThisServer = new HashSet<>(); List segments = TableStateUtils.getSegmentsInGivenStateForThisInstance(_helixManager, _tableNameWithType, CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING); for (String segmentNameStr : segments) { - LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - partitionsHostedByThisServer.add(segmentName.getTopicPartitionId().getPartitionId()); + partitionsHostedByThisServer.add( + new LLCSegmentName(segmentNameStr).getTopicPartitionId()); } return partitionsHostedByThisServer; } @@ -526,7 +529,8 @@ protected void doAddOnlineSegment(String segmentName) if (_enforceConsumptionInOrder) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { - getConsumerCoordinator(llcSegmentName.getTopicPartitionId().getPartitionId()).register(llcSegmentName); + getConsumerCoordinator(llcSegmentName.getTopicPartitionId()) + .register(llcSegmentName); } } } @@ -551,12 +555,24 @@ public void addConsumingSegment(String segmentName) } } - public Set stopTrackingPartitionIngestionDelay(@Nullable Set partitionIds) { + /** + * Stops tracking ingestion delay for the given partition IDs (as raw integers). + * Kept for backward compatibility with the REST API which receives partition IDs + * as integers from query parameters. + */ + public Set stopTrackingPartitionIngestionDelay( + @Nullable Set partitionIds) { if (CollectionUtils.isEmpty(partitionIds)) { - return _ingestionDelayTracker.stopTrackingAllPartitions(); + Set removed = _ingestionDelayTracker.stopTrackingAllPartitions(); + Set result = new HashSet<>(); + for (TopicPartitionId tpId : removed) { + result.add(tpId.toMultiTopicPinotPartitionId()); + } + return result; } - for (Integer partitionId: partitionIds) { - _ingestionDelayTracker.stopTrackingPartition(partitionId); + for (Integer partitionId : partitionIds) { + _ingestionDelayTracker.stopTrackingPartition( + TopicPartitionId.fromMultiTopicPinotPartitionId(partitionId)); } return partitionIds; } @@ -601,15 +617,19 @@ private void doAddConsumingSegment(String segmentName) // Generates only one semaphore for every partition LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionGroupId = llcSegmentName.getTopicPartitionId(); ConsumerCoordinator consumerCoordinator = getConsumerCoordinator(partitionGroupId); // Create the segment data manager and register it PartitionUpsertMetadataManager partitionUpsertMetadataManager = - _tableUpsertMetadataManager != null ? _tableUpsertMetadataManager.getOrCreatePartitionManager(partitionGroupId) + _tableUpsertMetadataManager != null + ? _tableUpsertMetadataManager.getOrCreatePartitionManager( + partitionGroupId.getPartitionId()) : null; PartitionDedupMetadataManager partitionDedupMetadataManager = - _tableDedupMetadataManager != null ? _tableDedupMetadataManager.getOrCreatePartitionManager(partitionGroupId) + _tableDedupMetadataManager != null + ? _tableDedupMetadataManager.getOrCreatePartitionManager( + partitionGroupId.getPartitionId()) : null; RealtimeSegmentDataManager realtimeSegmentDataManager = createRealtimeSegmentDataManager(zkMetadata, tableConfig, indexLoadingConfig, schema, llcSegmentName, @@ -853,7 +873,7 @@ public StreamIngestionConfig getStreamIngestionConfig() { } @VisibleForTesting - ConsumerCoordinator getConsumerCoordinator(int partitionId) { + ConsumerCoordinator getConsumerCoordinator(TopicPartitionId partitionId) { return _partitionIdToConsumerCoordinatorMap.computeIfAbsent(partitionId, k -> new ConsumerCoordinator(_enforceConsumptionInOrder, this)); } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinatorTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinatorTest.java index 0650f059c321..bf2a11188460 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinatorTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/ConsumerCoordinatorTest.java @@ -27,6 +27,7 @@ import java.util.concurrent.locks.ReentrantLock; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.spi.config.table.ingestion.StreamIngestionConfig; import org.apache.pinot.spi.utils.CommonConstants.Segment.Realtime.Status; import org.apache.pinot.util.TestUtils; @@ -55,7 +56,7 @@ public FakeRealtimeTableDataManager(Semaphore segmentBuildSemaphore, } @Override - ConsumerCoordinator getConsumerCoordinator(int partitionId) { + ConsumerCoordinator getConsumerCoordinator(TopicPartitionId partitionId) { return _consumerCoordinator; } diff --git a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTrackerTest.java b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTrackerTest.java index 80d5444e3ed9..9472b2bacd88 100644 --- a/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTrackerTest.java +++ b/pinot-core/src/test/java/org/apache/pinot/core/data/manager/realtime/IngestionDelayTrackerTest.java @@ -37,13 +37,13 @@ import org.apache.pinot.common.metrics.ServerGauge; import org.apache.pinot.common.metrics.ServerMetrics; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.config.table.ingestion.IngestionConfig; import org.apache.pinot.spi.config.table.ingestion.StreamIngestionConfig; import org.apache.pinot.spi.stream.LongMsgOffset; import org.apache.pinot.spi.stream.StreamPartitionMsgOffset; -import org.apache.pinot.spi.utils.IngestionConfigUtils; import org.apache.pinot.spi.utils.builder.TableConfigBuilder; import org.apache.pinot.spi.utils.builder.TableNameBuilder; import org.apache.pinot.util.TestUtils; @@ -84,7 +84,7 @@ private static Map getStreamConfigs() { } private static class MockIngestionDelayTracker extends IngestionDelayTracker { - private volatile Map>> _partitionToMetricToValues; + private volatile Map>> _partitionToMetricToValues; private ScheduledExecutorService _scheduledExecutorService; public MockIngestionDelayTracker(ServerMetrics serverMetrics, String tableNameWithType, @@ -101,7 +101,7 @@ public MockIngestionDelayTracker(ServerMetrics serverMetrics, String tableNameWi } @Override - public void createMetrics(int partitionId) { + public void createMetrics(TopicPartitionId partitionId) { synchronized (this) { if (_partitionToMetricToValues == null) { _partitionToMetricToValues = new ConcurrentHashMap<>(); @@ -137,7 +137,8 @@ public void createMetrics(int partitionId) { }), 0, 10, TimeUnit.MILLISECONDS); } - void updatePartitionIdToLatestOffset(Map partitionIdToLatestOffset) { + void updatePartitionIdToLatestOffset( + Map partitionIdToLatestOffset) { _partitionIdToLatestOffset = partitionIdToLatestOffset; } @@ -163,21 +164,22 @@ public void testTrackerConstructors() { Clock clock = Clock.systemUTC(); ingestionDelayTracker.setClock(clock); + TopicPartitionId tp0 = new TopicPartitionId(0); // Untracked partitions should return null (metric not emitted) - Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(0)); - Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(0)); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(0), Long.MIN_VALUE); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionReportingStatus(0), 0); + Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(tp0)); + Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tp0)); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(tp0), Long.MIN_VALUE); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionReportingStatus(tp0), 0); ingestionDelayTracker.shutdown(); // Test constructor with timer arguments ingestionDelayTracker = new MockIngestionDelayTracker(_serverMetrics, REALTIME_TABLE_NAME, REALTIME_TABLE_DATA_MANAGER, METRICS_CLEANUP_TICK_INTERVAL_MS, METRICS_TRACKING_TICK_INTERVAL_MS, () -> true); // Untracked partitions should return null (metric not emitted) - Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(0)); - Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(0)); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(0), Long.MIN_VALUE); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionReportingStatus(0), 0); + Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(tp0)); + Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tp0)); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(tp0), Long.MIN_VALUE); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionReportingStatus(tp0), 0); // Test bad timer args to the constructor try { new IngestionDelayTracker(_serverMetrics, REALTIME_TABLE_NAME, REALTIME_TABLE_DATA_MANAGER, 0, 0, () -> true); @@ -192,10 +194,12 @@ public void testTrackerConstructors() { @Test public void testRecordIngestionDelayWithNoAging() { final long maxTestDelay = 100; - final int partition0 = 0; - final String segment0 = new LLCSegmentName(RAW_TABLE_NAME, partition0, 0, 123).getSegmentName(); - final int partition1 = 1; - final String segment1 = new LLCSegmentName(RAW_TABLE_NAME, partition1, 0, 234).getSegmentName(); + final TopicPartitionId partition0 = new TopicPartitionId(0); + final String segment0 = + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, 123).getSegmentName(); + final TopicPartitionId partition1 = new TopicPartitionId(1); + final String segment1 = + new LLCSegmentName(RAW_TABLE_NAME, 1, 0, 234).getSegmentName(); IngestionDelayTracker ingestionDelayTracker = createTracker(); // Use fixed clock so samples dont age @@ -261,14 +265,16 @@ public void testRecordIngestionDelayWithNoAging() { @Test public void testRecordIngestionDelayWithAging() { - final int partition0 = 0; - final String segment0 = new LLCSegmentName(RAW_TABLE_NAME, partition0, 0, 123).getSegmentName(); + final TopicPartitionId partition0 = new TopicPartitionId(0); + final String segment0 = + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, 123).getSegmentName(); final long partition0Delay0 = 1000; final long partition0Delay1 = 10; // record lower delay to make sure max gets reduced final long partition0Offset0Ms = 300; final long partition0Offset1Ms = 1000; - final int partition1 = 1; - final String segment1 = new LLCSegmentName(RAW_TABLE_NAME, partition1, 0, 234).getSegmentName(); + final TopicPartitionId partition1 = new TopicPartitionId(1); + final String segment1 = + new LLCSegmentName(RAW_TABLE_NAME, 1, 0, 234).getSegmentName(); final long partition1Delay0 = 11; final long partition1Offset0Ms = 150; @@ -335,22 +341,34 @@ public void testStopTrackingIngestionDelay() { // Record a number of partitions with delay equal to partition id for (int partitionId = 0; partitionId <= maxTestDelay; partitionId++) { - String segmentName = new LLCSegmentName(RAW_TABLE_NAME, partitionId, 0, 123).getSegmentName(); + TopicPartitionId tpId = new TopicPartitionId(partitionId); + String segmentName = + new LLCSegmentName(RAW_TABLE_NAME, partitionId, 0, 123) + .getSegmentName(); long ingestionTimeMs = clock.millis() - partitionId; - ingestionDelayTracker.updateMetrics(segmentName, partitionId, ingestionTimeMs, ingestionTimeMs, null); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionDelayMs(partitionId), partitionId); - Assert.assertEquals(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(partitionId), partitionId); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(partitionId), ingestionTimeMs); + ingestionDelayTracker.updateMetrics( + segmentName, tpId, ingestionTimeMs, ingestionTimeMs, null); + Assert.assertEquals( + ingestionDelayTracker.getPartitionIngestionDelayMs(tpId), + partitionId); + Assert.assertEquals( + ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tpId), + partitionId); + Assert.assertEquals( + ingestionDelayTracker.getPartitionIngestionTimeMs(tpId), + ingestionTimeMs); } for (int partitionId = maxPartition; partitionId >= 0; partitionId--) { - ingestionDelayTracker.stopTrackingPartition(partitionId); + ingestionDelayTracker.stopTrackingPartition( + new TopicPartitionId(partitionId)); } for (int partitionId = 0; partitionId <= maxTestDelay; partitionId++) { + TopicPartitionId tpId = new TopicPartitionId(partitionId); // Untracked partitions must return null (metric not emitted) - Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(partitionId)); - Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(partitionId)); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(partitionId), Long.MIN_VALUE); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionReportingStatus(partitionId), 0); + Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(tpId)); + Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tpId)); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(tpId), Long.MIN_VALUE); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionReportingStatus(tpId), 0); } } @@ -363,23 +381,30 @@ public void testStopTrackingIngestionDelayWithSegment() { Clock clock = Clock.fixed(now, zoneId); ingestionDelayTracker.setClock(clock); - String segmentName = new LLCSegmentName(RAW_TABLE_NAME, 0, 0, 123).getSegmentName(); + TopicPartitionId tp0 = new TopicPartitionId(0); + String segmentName = + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, 123).getSegmentName(); long ingestionTimeMs = clock.millis() - 10; - ingestionDelayTracker.updateMetrics(segmentName, 0, ingestionTimeMs, ingestionTimeMs, null); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionDelayMs(0), 10); - Assert.assertEquals(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(0), 10); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(0), ingestionTimeMs); + ingestionDelayTracker.updateMetrics( + segmentName, tp0, ingestionTimeMs, ingestionTimeMs, null); + Assert.assertEquals( + ingestionDelayTracker.getPartitionIngestionDelayMs(tp0), 10); + Assert.assertEquals( + ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tp0), 10); + Assert.assertEquals( + ingestionDelayTracker.getPartitionIngestionTimeMs(tp0), + ingestionTimeMs); ingestionDelayTracker.stopTrackingPartition(segmentName); - Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(0)); - Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(0)); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(0), Long.MIN_VALUE); + Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(tp0)); + Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tp0)); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(tp0), Long.MIN_VALUE); // Should not update metrics for removed segment - ingestionDelayTracker.updateMetrics(segmentName, 0, ingestionTimeMs, ingestionTimeMs, null); - Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(0)); - Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(0)); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(0), Long.MIN_VALUE); + ingestionDelayTracker.updateMetrics(segmentName, tp0, ingestionTimeMs, ingestionTimeMs, null); + Assert.assertNull(ingestionDelayTracker.getPartitionIngestionDelayMs(tp0)); + Assert.assertNull(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tp0)); + Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionTimeMs(tp0), Long.MIN_VALUE); } @Test @@ -395,11 +420,19 @@ public void testShutdown() { // Test Shutdown with partitions active for (int partitionId = 0; partitionId <= maxTestDelay; partitionId++) { - String segmentName = new LLCSegmentName(RAW_TABLE_NAME, partitionId, 0, 123).getSegmentName(); + TopicPartitionId tpId = new TopicPartitionId(partitionId); + String segmentName = + new LLCSegmentName(RAW_TABLE_NAME, partitionId, 0, 123) + .getSegmentName(); long ingestionTimeMs = clock.millis() - partitionId; - ingestionDelayTracker.updateMetrics(segmentName, partitionId, ingestionTimeMs, ingestionTimeMs, null); - Assert.assertEquals(ingestionDelayTracker.getPartitionIngestionDelayMs(partitionId), partitionId); - Assert.assertEquals(ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(partitionId), partitionId); + ingestionDelayTracker.updateMetrics( + segmentName, tpId, ingestionTimeMs, ingestionTimeMs, null); + Assert.assertEquals( + ingestionDelayTracker.getPartitionIngestionDelayMs(tpId), + partitionId); + Assert.assertEquals( + ingestionDelayTracker.getPartitionEndToEndIngestionDelayMs(tpId), + partitionId); } ingestionDelayTracker.shutdown(); @@ -410,13 +443,16 @@ public void testShutdown() { @Test public void testRecordIngestionDelayOffset() { - final int partition0 = 0; - final String segment0 = new LLCSegmentName(RAW_TABLE_NAME, partition0, 0, 123).getSegmentName(); - final int partition1 = 1; - final String segment1 = new LLCSegmentName(RAW_TABLE_NAME, partition1, 0, 234).getSegmentName(); + final TopicPartitionId partition0 = new TopicPartitionId(0); + final String segment0 = + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, 123).getSegmentName(); + final TopicPartitionId partition1 = new TopicPartitionId(1); + final String segment1 = + new LLCSegmentName(RAW_TABLE_NAME, 1, 0, 234).getSegmentName(); IngestionDelayTracker ingestionDelayTracker = createTracker(); - Map partitionMsgOffsetMap = new HashMap<>(); + Map partitionMsgOffsetMap = + new HashMap<>(); ((MockIngestionDelayTracker) ingestionDelayTracker).updatePartitionIdToLatestOffset(partitionMsgOffsetMap); // Test tracking offset lag for a single partition @@ -455,40 +491,56 @@ public void testRecordIngestionDelayOffset() { @Test public void testUpdateLatestStreamOffset() { IngestionDelayTracker ingestionDelayTracker = createTracker(); - Set partitionsHosted = new HashSet<>(); - partitionsHosted.add(0); - partitionsHosted.add(1); + Set partitionsHosted = new HashSet<>(); + partitionsHosted.add(new TopicPartitionId(0)); + partitionsHosted.add(new TopicPartitionId(1)); ingestionDelayTracker.updateLatestStreamOffset(partitionsHosted); - Assert.assertEquals(ingestionDelayTracker._partitionIdToLatestOffset.size(), 2); - Assert.assertEquals(((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset.get(0))).getOffset(), + Assert.assertEquals( + ingestionDelayTracker._partitionIdToLatestOffset.size(), 2); + Assert.assertEquals( + ((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset + .get(new TopicPartitionId(0)))).getOffset(), Integer.MAX_VALUE); - Assert.assertEquals(((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset.get(1))).getOffset(), + Assert.assertEquals( + ((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset + .get(new TopicPartitionId(1)))).getOffset(), Integer.MAX_VALUE); IngestionConfig ingestionConfig = new IngestionConfig(); List> streamConfigMaps = new ArrayList<>(); streamConfigMaps.add(getStreamConfigs()); streamConfigMaps.add(getStreamConfigs()); - StreamIngestionConfig streamIngestionConfig = new StreamIngestionConfig(streamConfigMaps); + StreamIngestionConfig streamIngestionConfig = + new StreamIngestionConfig(streamConfigMaps); ingestionConfig.setStreamIngestionConfig(streamIngestionConfig); - TableConfig tableConfig = new TableConfigBuilder(TableType.REALTIME).setTableName(RAW_TABLE_NAME) - .setTimeColumnName("ts") - .setNullHandlingEnabled(true) - .setIngestionConfig(ingestionConfig) - .setStreamConfigs(getStreamConfigs()) - .build(); - when(REALTIME_TABLE_DATA_MANAGER.getCachedTableConfigAndSchema()).thenReturn(Pair.of(tableConfig, null)); + TableConfig tableConfig = + new TableConfigBuilder(TableType.REALTIME) + .setTableName(RAW_TABLE_NAME) + .setTimeColumnName("ts") + .setNullHandlingEnabled(true) + .setIngestionConfig(ingestionConfig) + .setStreamConfigs(getStreamConfigs()) + .build(); + when(REALTIME_TABLE_DATA_MANAGER.getCachedTableConfigAndSchema()) + .thenReturn(Pair.of(tableConfig, null)); ingestionDelayTracker = createTracker(); - partitionsHosted.add(IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(0, 1)); + partitionsHosted.add(new TopicPartitionId(1, 0)); ingestionDelayTracker.updateLatestStreamOffset(partitionsHosted); - Assert.assertEquals(ingestionDelayTracker._partitionIdToLatestOffset.size(), 3); - Assert.assertEquals(((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset.get(0))).getOffset(), + Assert.assertEquals( + ingestionDelayTracker._partitionIdToLatestOffset.size(), 3); + Assert.assertEquals( + ((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset + .get(new TopicPartitionId(0)))).getOffset(), Integer.MAX_VALUE); - Assert.assertEquals(((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset.get(1))).getOffset(), + Assert.assertEquals( + ((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset + .get(new TopicPartitionId(1)))).getOffset(), + Integer.MAX_VALUE); + Assert.assertEquals( + ((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset + .get(new TopicPartitionId(1, 0)))).getOffset(), Integer.MAX_VALUE); - Assert.assertEquals(((LongMsgOffset) (ingestionDelayTracker._partitionIdToLatestOffset.get( - IngestionConfigUtils.getPinotPartitionIdFromStreamPartitionId(0, 1)))).getOffset(), Integer.MAX_VALUE); } @Test @@ -497,26 +549,37 @@ public void testIngestionDelay() { new MockIngestionDelayTracker(_serverMetrics, REALTIME_TABLE_NAME, REALTIME_TABLE_DATA_MANAGER, METRICS_CLEANUP_TICK_INTERVAL_MS, METRICS_TRACKING_TICK_INTERVAL_MS, () -> true); - final int partition0 = 0; - final String segment0 = new LLCSegmentName(RAW_TABLE_NAME, partition0, 0, 123).getSegmentName(); - final int partition1 = 1; - - ingestionDelayTracker._partitionsHostedByThisServer.put(partition0, true); - ingestionDelayTracker._partitionsHostedByThisServer.put(partition1, true); - ingestionDelayTracker.updateMetrics(segment0, partition0, System.currentTimeMillis(), System.currentTimeMillis(), + final TopicPartitionId partition0 = new TopicPartitionId(0); + final String segment0 = + new LLCSegmentName(RAW_TABLE_NAME, 0, 0, 123).getSegmentName(); + final TopicPartitionId partition1 = new TopicPartitionId(1); + + ingestionDelayTracker._partitionsHostedByThisServer.put( + partition0, true); + ingestionDelayTracker._partitionsHostedByThisServer.put( + partition1, true); + ingestionDelayTracker.updateMetrics(segment0, partition0, + System.currentTimeMillis(), System.currentTimeMillis(), new LongMsgOffset(50)); - ScheduledExecutorService scheduledExecutorService = Executors.newSingleThreadScheduledExecutor(); - Map partitionIdVsLatestOffset = new HashMap<>(); + ScheduledExecutorService scheduledExecutorService = + Executors.newSingleThreadScheduledExecutor(); + Map + partitionIdVsLatestOffset = new HashMap<>(); partitionIdVsLatestOffset.put(partition0, new LongMsgOffset(200)); partitionIdVsLatestOffset.put(partition1, new LongMsgOffset(200)); - ingestionDelayTracker.updatePartitionIdToLatestOffset(partitionIdVsLatestOffset); + ingestionDelayTracker.updatePartitionIdToLatestOffset( + partitionIdVsLatestOffset); scheduledExecutorService.scheduleWithFixedDelay(() -> { partitionIdVsLatestOffset.put(partition0, - new LongMsgOffset(((LongMsgOffset) (partitionIdVsLatestOffset.get(partition0))).getOffset() + 50)); + new LongMsgOffset( + ((LongMsgOffset) (partitionIdVsLatestOffset.get(partition0))) + .getOffset() + 50)); partitionIdVsLatestOffset.put(partition1, - new LongMsgOffset(((LongMsgOffset) (partitionIdVsLatestOffset.get(partition1))).getOffset() + 50)); + new LongMsgOffset( + ((LongMsgOffset) (partitionIdVsLatestOffset.get(partition1))) + .getOffset() + 50)); }, 0, 10, TimeUnit.MILLISECONDS); TestUtils.waitForCondition((aVoid) -> { @@ -531,10 +594,14 @@ public void testIngestionDelay() { ingestionDelayTracker.shutdown(); } - private void verifyMetrics(Map>> partitionToMetricToValues) { + private void verifyMetrics( + Map>> + partitionToMetricToValues) { Assert.assertEquals(partitionToMetricToValues.size(), 2); - verifyPartition0(partitionToMetricToValues.get(0)); - verifyPartition1(partitionToMetricToValues.get(1)); + verifyPartition0( + partitionToMetricToValues.get(new TopicPartitionId(0))); + verifyPartition1( + partitionToMetricToValues.get(new TopicPartitionId(1))); } private void verifyPartition0(Map> metrics) { 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 4285c6bd9b1c..ef26e1cb1322 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 @@ -1200,7 +1200,8 @@ public FakeRealtimeSegmentDataManager(SegmentZKMetadata segmentZKMetadata, Table throws Exception { super(segmentZKMetadata, tableConfig, realtimeTableDataManager, resourceDataDir, new IndexLoadingConfig(makeInstanceDataManagerConfig(), tableConfig), schema, llcSegmentName, - consumerCoordinatorMap.get(llcSegmentName.getTopicPartitionId().getPartitionId()), serverMetrics, null, null, + consumerCoordinatorMap.get(llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId()), + serverMetrics, null, null, () -> true); _state = RealtimeSegmentDataManager.class.getDeclaredField("_state"); _state.setAccessible(true); diff --git a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java index 3a11c38861ea..07ef6c592ace 100644 --- a/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java +++ b/pinot-integration-test-base/src/test/java/org/apache/pinot/integration/tests/BaseClusterIntegrationTest.java @@ -1098,7 +1098,7 @@ private int getNumConsumingPartitions(String tableNameWithType) { Map stateMap = idealState.getInstanceStateMap(segmentName); if (stateMap != null && stateMap.containsValue(CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { - consumingPartitions.add(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); + consumingPartitions.add(new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId()); } } return consumingPartitions.size(); diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java index 7af6cb915de9..107fc22b9fb3 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeClusterIntegrationTest.java @@ -152,7 +152,8 @@ private boolean isOffline(int partition, int seqNum) { for (String segmentNameStr : ev.getPartitionSet()) { if (LLCSegmentName.isLLCSegment(segmentNameStr)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getSequenceNumber() == seqNum && segmentName.getTopicPartitionId().getPartitionId() == partition + if (segmentName.getSequenceNumber() == seqNum + && segmentName.getTopicPartitionId().toMultiTopicPinotPartitionId() == partition && ev.getStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.OFFLINE)) { isOffline = true; @@ -563,7 +564,7 @@ private int getSegmentSeqNum(int partition) { if (is.getInstanceStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getTopicPartitionId().getPartitionId() == partition) { + if (segmentName.getTopicPartitionId().toMultiTopicPinotPartitionId() == partition) { seqNum.set(segmentName.getSequenceNumber()); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java index cb6657826fa7..3afab5661a01 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka3ClusterIntegrationTest.java @@ -111,7 +111,7 @@ private int getSegmentSeqNum(int partition) { if (is.getInstanceStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getTopicPartitionId().getPartitionId() == partition) { + if (segmentName.getTopicPartitionId().toMultiTopicPinotPartitionId() == partition) { seqNum.set(segmentName.getSequenceNumber()); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java index 094d9b49bd50..62fc8bbe3be5 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/LLCRealtimeKafka4ClusterIntegrationTest.java @@ -111,7 +111,7 @@ private int getSegmentSeqNum(int partition) { if (is.getInstanceStateMap(segmentNameStr).values().contains( CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { LLCSegmentName segmentName = new LLCSegmentName(segmentNameStr); - if (segmentName.getTopicPartitionId().getPartitionId() == partition) { + if (segmentName.getTopicPartitionId().toMultiTopicPinotPartitionId() == partition) { seqNum.set(segmentName.getSequenceNumber()); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java index 818161e2fa66..07d199fb1fdc 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PartialUpsertTableRebalanceIntegrationTest.java @@ -407,7 +407,7 @@ private static int getSegmentPartitionId(String segmentName) { case UPLOADED_SEGMENT_3: return 1; default: - return new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); + return new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId(); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java index 392313a6f7ce..ecdee8d076f6 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/PauselessRealtimeIngestionSegmentCommitFailureTest.java @@ -245,7 +245,8 @@ private Map getPartitionSegmentNumberToMetadataMap( for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentZKMetadata.getSegmentName()); String segmentKey = - llcSegmentName.getTopicPartitionId().getPartitionId() + "_" + llcSegmentName.getSequenceNumber(); + llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId() + + "_" + llcSegmentName.getSequenceNumber(); segmentZKMetadataMap.put(segmentKey, segmentZKMetadata); } return segmentZKMetadataMap; diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java index d3353b68d407..bdd6c39ff97f 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/SegmentPartitionLLCRealtimeClusterIntegrationTest.java @@ -162,7 +162,7 @@ public void testPartitionMetadata() { assertTrue(columnPartitionMetadata.getFunctionName().equalsIgnoreCase("murmur")); assertEquals(columnPartitionMetadata.getNumPartitions(), 2); int partitionGroupId = - new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().getPartitionId(); + new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().toMultiTopicPinotPartitionId(); assertEquals(columnPartitionMetadata.getPartitions(), Set.of(partitionGroupId)); numSegmentsForPartition[partitionGroupId]++; } @@ -270,7 +270,7 @@ public void testNonPartitionedStream() assertTrue(columnPartitionMetadata.getFunctionName().equalsIgnoreCase("murmur")); assertEquals(columnPartitionMetadata.getNumPartitions(), 2); int partitionGroupId = - new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().getPartitionId(); + new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().toMultiTopicPinotPartitionId(); numSegmentsForPartition[partitionGroupId]++; if (segmentZKMetadata.getStatus() == Status.IN_PROGRESS) { @@ -348,7 +348,7 @@ public void testNonPartitionedStream() assertTrue(columnPartitionMetadata.getFunctionName().equalsIgnoreCase("murmur")); assertEquals(columnPartitionMetadata.getNumPartitions(), 2); int partitionGroupId = - new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().getPartitionId(); + new LLCSegmentName(segmentZKMetadata.getSegmentName()).getTopicPartitionId().toMultiTopicPinotPartitionId(); numSegmentsForPartition[partitionGroupId]++; if (segmentZKMetadata.getStatus() == Status.IN_PROGRESS) { diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java index 8c78c861f8df..93ce8c107db4 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentPreloadIntegrationTest.java @@ -298,7 +298,7 @@ private static int getSegmentPartitionId(String segmentName) { case UPLOADED_SEGMENT_3: return 1; default: - return new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); + return new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId(); } } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java index 3f901841bb66..be1fa8f8bfd7 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/UpsertTableSegmentUploadIntegrationTest.java @@ -225,7 +225,7 @@ private static int getSegmentPartitionId(String segmentName) { case UPLOADED_SEGMENT_3: return 1; default: - return new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); + return new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId(); } } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java index d04a559618c6..590690c1a903 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/custom/MultiTopicRealtimeClusterIntegrationTest.java @@ -256,7 +256,7 @@ private int getNumConsumingPartitions(String tableNameWithType) { Map stateMap = idealState.getInstanceStateMap(segmentName); if (stateMap != null && stateMap.containsValue(CommonConstants.Helix.StateModel.SegmentStateModel.CONSUMING)) { - consumingPartitions.add(new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId()); + consumingPartitions.add(new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId()); } } return consumingPartitions.size(); @@ -383,7 +383,7 @@ public void testSegmentsFromAllTopics() { Set topicIndicesSeen = new HashSet<>(); for (String segmentName : idealState.getPartitionSet()) { if (LLCSegmentName.isLLCSegment(segmentName)) { - int pgId = new LLCSegmentName(segmentName).getTopicPartitionId().getPartitionId(); + int pgId = new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId(); topicIndicesSeen.add(IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(pgId)); } } diff --git a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java index d7b954b96c12..db9642fadc47 100644 --- a/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java +++ b/pinot-integration-tests/src/test/java/org/apache/pinot/integration/tests/realtime/utils/PauselessRealtimeTestUtils.java @@ -71,7 +71,8 @@ private static Map getPartitionSegmentNumberToMetadat for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { LLCSegmentName llcSegmentName = new LLCSegmentName(segmentZKMetadata.getSegmentName()); String segmentKey = - llcSegmentName.getTopicPartitionId().getPartitionId() + "_" + llcSegmentName.getSequenceNumber(); + llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId() + + "_" + llcSegmentName.getSequenceNumber(); segmentZKMetadataMap.put(segmentKey, segmentZKMetadata); } return segmentZKMetadataMap; diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java index b4d7218524a1..eec030f5d3c5 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGenerator.java @@ -43,6 +43,7 @@ import org.apache.pinot.common.metrics.ControllerMetrics; import org.apache.pinot.common.minion.MergeRollupTaskMetadata; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.controller.helix.core.minion.generator.BaseTaskGenerator; import org.apache.pinot.controller.helix.core.minion.generator.PinotTaskGenerator; import org.apache.pinot.controller.helix.core.minion.generator.TaskGeneratorUtils; @@ -168,7 +169,8 @@ public List generateTasks(List tableConfigs) { tableConfig.getTableType() == TableType.OFFLINE ? getSegmentsZKMetadataForTable(tableNameWithType) : filterSegmentsforRealtimeTable( - getNonConsumingSegmentsZKMetadataForRealtimeTable(tableNameWithType)); + getNonConsumingSegmentsZKMetadataForRealtimeTable(tableNameWithType), + tableConfig); // Select current segment snapshot based on lineage, filter out empty segments SegmentLineage segmentLineage = _clusterInfoAccessor.getSegmentLineage(tableNameWithType); @@ -577,7 +579,8 @@ public void validateTaskConfigs(TableConfig tableConfig, Schema schema, Map filterSegmentsforRealtimeTable(List allSegments) { + static List filterSegmentsforRealtimeTable(List allSegments, + TableConfig tableConfig) { // For realtime table, don't process // 1. in-progress segments (Segment.Realtime.Status.IN_PROGRESS), this has been taken care of in // getNonConsumingSegmentsZKMetadataForRealtimeTable() @@ -598,12 +601,13 @@ static List filterSegmentsforRealtimeTable(List partitionIdToLatestCompletedSegment = new HashMap<>(); + boolean hasMultipleStreams = IngestionConfigUtils.hasMultipleStreams(tableConfig); + Map partitionIdToLatestCompletedSegment = new HashMap<>(); for (SegmentZKMetadata segmentZKMetadata : allSegments) { String segmentName = segmentZKMetadata.getSegmentName(); if (LLCSegmentName.isLLCSegment(segmentName)) { - LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - partitionIdToLatestCompletedSegment.compute(llcSegmentName.getTopicPartitionId().getPartitionId(), + LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName, hasMultipleStreams); + partitionIdToLatestCompletedSegment.compute(llcSegmentName.getTopicPartitionId(), (partId, latestSegment) -> { if (latestSegment == null) { return llcSegmentName; diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java index a2ab2de5f645..3592c8bb3580 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/purge/PurgeTaskGenerator.java @@ -30,6 +30,7 @@ import org.apache.pinot.common.data.Segment; import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.controller.helix.core.minion.generator.BaseTaskGenerator; import org.apache.pinot.controller.helix.core.minion.generator.TaskGeneratorUtils; import org.apache.pinot.core.common.MinionConstants; @@ -39,6 +40,7 @@ import org.apache.pinot.spi.config.table.TableConfig; import org.apache.pinot.spi.config.table.TableTaskConfig; import org.apache.pinot.spi.config.table.TableType; +import org.apache.pinot.spi.utils.IngestionConfigUtils; import org.apache.pinot.spi.utils.TimeUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -108,7 +110,7 @@ public List generateTasks(List tableConfigs) { // For realtime tables, build a map of partition to latest segment to avoid deleting last segments Set lastLLCSegmentPerPartition = new HashSet<>(); if (tableConfig.getTableType() == TableType.REALTIME) { - lastLLCSegmentPerPartition = getLastLLCSegmentPerPartition(segmentsZKMetadata); + lastLLCSegmentPerPartition = getLastLLCSegmentPerPartition(segmentsZKMetadata, tableConfig); } for (SegmentZKMetadata segmentZKMetadata : notpurgedSegmentsZKMetadata) { String segmentName = segmentZKMetadata.getSegmentName(); @@ -163,13 +165,15 @@ public List generateTasks(List tableConfigs) { return pinotTaskConfigs; } - private Set getLastLLCSegmentPerPartition(List segmentsZKMetadata) { - Map latestLLCSegmentNameMap = new HashMap<>(); + private Set getLastLLCSegmentPerPartition(List segmentsZKMetadata, + TableConfig tableConfig) { + boolean hasMultipleStreams = IngestionConfigUtils.hasMultipleStreams(tableConfig); + Map latestLLCSegmentNameMap = new HashMap<>(); for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { // Skip UPLOADED segments that don't conform to the LLC segment name - LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentZKMetadata.getSegmentName()); + LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentZKMetadata.getSegmentName(), hasMultipleStreams); if (llcSegmentName != null) { - latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId().getPartitionId(), + latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId(), (k, latestLLCSegmentName) -> { if (latestLLCSegmentName == null || llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java index 166d4387a743..bf010f432a31 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/main/java/org/apache/pinot/plugin/minion/tasks/realtimetoofflinesegments/RealtimeToOfflineSegmentsTaskGenerator.java @@ -32,6 +32,7 @@ import org.apache.pinot.common.metadata.segment.SegmentZKMetadata; import org.apache.pinot.common.minion.RealtimeToOfflineSegmentsTaskMetadata; import org.apache.pinot.common.utils.LLCSegmentName; +import org.apache.pinot.common.utils.TopicPartitionId; import org.apache.pinot.controller.helix.core.minion.generator.BaseTaskGenerator; import org.apache.pinot.controller.helix.core.minion.generator.PinotTaskGenerator; import org.apache.pinot.controller.helix.core.minion.generator.TaskGeneratorUtils; @@ -48,6 +49,7 @@ import org.apache.pinot.spi.config.table.TableType; import org.apache.pinot.spi.config.table.UpsertConfig; import org.apache.pinot.spi.data.Schema; +import org.apache.pinot.spi.utils.IngestionConfigUtils; import org.apache.pinot.spi.utils.TimeUtils; import org.slf4j.Logger; import org.slf4j.LoggerFactory; @@ -117,10 +119,10 @@ public List generateTasks(List tableConfigs) { // Get all segment metadata for completed segments (DONE/UPLOADED status). List completedSegmentsZKMetadata = new ArrayList<>(); - Map partitionToLatestLLCSegmentName = new HashMap<>(); - Set allPartitions = new HashSet<>(); - getCompletedSegmentsInfo(realtimeTableName, completedSegmentsZKMetadata, partitionToLatestLLCSegmentName, - allPartitions); + Map partitionToLatestLLCSegmentName = new HashMap<>(); + Set allPartitions = new HashSet<>(); + getCompletedSegmentsInfo(realtimeTableName, tableConfig, completedSegmentsZKMetadata, + partitionToLatestLLCSegmentName, allPartitions); if (completedSegmentsZKMetadata.isEmpty()) { LOGGER.info("No realtime-completed segments found for table: {}, skipping task generation: {}", realtimeTableName, taskType); @@ -245,22 +247,25 @@ public List generateTasks(List tableConfigs) { * Fetch completed (DONE/UPLOADED) segment and partition information * * @param realtimeTableName the realtime table name + * @param tableConfig the table config for multi-stream context * @param completedSegmentsZKMetadata list for collecting the completed (DONE/UPLOADED) segments ZK metadata * @param partitionToLatestLLCSegmentName map for collecting the partitionId to the latest LLC segment name * @param allPartitions set for collecting all partition ids */ - private void getCompletedSegmentsInfo(String realtimeTableName, List completedSegmentsZKMetadata, - Map partitionToLatestLLCSegmentName, Set allPartitions) { + private void getCompletedSegmentsInfo(String realtimeTableName, TableConfig tableConfig, + List completedSegmentsZKMetadata, + Map partitionToLatestLLCSegmentName, Set allPartitions) { + boolean hasMultipleStreams = IngestionConfigUtils.hasMultipleStreams(tableConfig); List segmentsZKMetadata = getNonConsumingSegmentsZKMetadataForRealtimeTable(realtimeTableName); - Map latestLLCSegmentNameMap = new HashMap<>(); + Map latestLLCSegmentNameMap = new HashMap<>(); for (SegmentZKMetadata segmentZKMetadata : segmentsZKMetadata) { completedSegmentsZKMetadata.add(segmentZKMetadata); // Skip UPLOADED segments that don't conform to the LLC segment name - LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentZKMetadata.getSegmentName()); + LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentZKMetadata.getSegmentName(), hasMultipleStreams); if (llcSegmentName != null) { - int partitionId = llcSegmentName.getTopicPartitionId().getPartitionId(); + TopicPartitionId partitionId = llcSegmentName.getTopicPartitionId(); allPartitions.add(partitionId); latestLLCSegmentNameMap.compute(partitionId, (k, latestLLCSegmentName) -> { if (latestLLCSegmentName == null @@ -273,7 +278,7 @@ private void getCompletedSegmentsInfo(String realtimeTableName, List entry : latestLLCSegmentNameMap.entrySet()) { + for (Map.Entry entry : latestLLCSegmentNameMap.entrySet()) { partitionToLatestLLCSegmentName.put(entry.getKey(), entry.getValue().getSegmentName()); } } diff --git a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGeneratorTest.java b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGeneratorTest.java index ac388bebe08f..6bd4453e245b 100644 --- a/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGeneratorTest.java +++ b/pinot-plugins/pinot-minion-tasks/pinot-minion-builtin-tasks/src/test/java/org/apache/pinot/plugin/minion/tasks/mergerollup/MergeRollupTaskGeneratorTest.java @@ -404,7 +404,8 @@ public void testGenerateTasksCheckConfigs() { // Skip task generation, if the table is a realtime table and all segments are skipped // We don't test realtime REFRESH table because this combination does not make sense assertTrue(MergeRollupTaskGenerator.filterSegmentsforRealtimeTable( - Lists.newArrayList(realtimeTableSegmentMetadata1, realtimeTableSegmentMetadata2) + Lists.newArrayList(realtimeTableSegmentMetadata1, realtimeTableSegmentMetadata2), + getTableConfig(TableType.REALTIME, new HashMap<>()) ).isEmpty()); TableConfig realtimeTableConfig = getTableConfig(TableType.REALTIME, new HashMap<>()); List pinotTaskConfigs = generator.generateTasks(Lists.newArrayList(realtimeTableConfig)); @@ -454,7 +455,8 @@ public void testFilterSegmentsforRealtimeTable() { List filterResult = MergeRollupTaskGenerator.filterSegmentsforRealtimeTable( Lists.newArrayList(realtimeTableSegmentMetadata1, realtimeTableSegmentMetadata2, - realtimeTableSegmentMetadata3)); + realtimeTableSegmentMetadata3), + getTableConfig(TableType.REALTIME, new HashMap<>())); assertEquals(filterResult.size(), 1); assertEquals(filterResult.get(0).getSegmentName(), "testTable__0__0__20250224T0900Z"); } diff --git a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java index a0746bb48669..120089d2991f 100644 --- a/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java +++ b/pinot-query-planner/src/main/java/org/apache/pinot/query/planner/physical/v2/opt/rules/LeafStageWorkerAssignmentRule.java @@ -453,6 +453,11 @@ static Map> getInvalidSegmentsByInferredPartition(@Nullabl @VisibleForTesting static int inferPartitionId(String segmentName, int numPartitions) { + // TODO: Thread hasMultipleStreams through the call chain from _tableCache so that + // multi-stream composite partition IDs are properly decomposed into raw stream partition IDs. + // Currently uses the deprecated no-arg parse, which treats composite IDs as single-stream + // partition IDs. This is correct for single-stream tables; for multi-stream tables the + // modulo result may differ but worker assignment remains deterministic. if (LLCSegmentName.isLLCSegment(segmentName)) { LLCSegmentName llc = LLCSegmentName.of(segmentName); return llc != null ? (llc.getTopicPartitionId().getPartitionId() % numPartitions) : -1; diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java index 0de51c032cb7..be96b1e4f5ad 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/indexsegment/mutable/MutableSegmentImpl.java @@ -335,6 +335,7 @@ public MultiColumnTextMetadata getMultiColumnTextMetadata() { .withEstimatedColSize(_statsHistory.getEstimatedAvgColSize(column)) .withAvgNumMultiValues(_statsHistory.getEstimatedAvgColSize(column)) .withConsumerDir(_consumerDir) + .withHasMultipleStreams(config.hasMultipleStreams()) .withFixedLengthBytes(fixedByteSize).build(); // Partition info @@ -481,7 +482,7 @@ public MultiColumnTextMetadata getMultiColumnTextMetadata() { } _multiColumnTextIndex = new MultiColumnRealtimeLuceneTextIndex(textColumns, columnsSV, _consumerDir, config.getSegmentName(), - textConfig); + textConfig, config.hasMultipleStreams()); _multiColumnPos = _multiColumnTextIndex.getMapping(); _multiColumnValues = new ArrayList<>(_multiColumnPos.size()); for (int i = 0; i < _multiColumnPos.size(); i++) { diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/RealtimeSegmentConfig.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/RealtimeSegmentConfig.java index 580d34495d7b..fcd936515565 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/RealtimeSegmentConfig.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/RealtimeSegmentConfig.java @@ -72,6 +72,7 @@ public class RealtimeSegmentConfig { @Nullable private final MultiColumnTextIndexConfig _multiColIndexConfig; private final boolean _dropRecordOnPartitionMismatch; + private final boolean _hasMultipleStreams; // TODO: Clean up this constructor. Most of these things can be extracted from tableConfig. @@ -101,7 +102,8 @@ private RealtimeSegmentConfig( @Nullable PartitionDedupMetadataManager partitionDedupMetadataManager, String consumerDir, @Nullable MultiColumnTextIndexConfig textIndexConfig, - boolean dropRecordOnPartitionMismatch) { + boolean dropRecordOnPartitionMismatch, + boolean hasMultipleStreams) { _tableNameWithType = tableNameWithType; _segmentName = segmentName; _streamName = streamName; @@ -125,6 +127,7 @@ private RealtimeSegmentConfig( _consumerDir = consumerDir; _multiColIndexConfig = textIndexConfig; _dropRecordOnPartitionMismatch = dropRecordOnPartitionMismatch; + _hasMultipleStreams = hasMultipleStreams; } public String getTableNameWithType() { @@ -224,6 +227,10 @@ public boolean isDropRecordOnPartitionMismatch() { return _dropRecordOnPartitionMismatch; } + public boolean hasMultipleStreams() { + return _hasMultipleStreams; + } + public static class Builder { private String _tableNameWithType; private String _segmentName; @@ -252,6 +259,7 @@ public static class Builder { private String _consumerDir; private MultiColumnTextIndexConfig _textIndexConfig; private boolean _dropRecordOnPartitionMismatch; + private boolean _hasMultipleStreams; public Builder() { _indexConfigByCol = new HashMap<>(); @@ -429,6 +437,11 @@ public Builder setDropRecordOnPartitionMismatch(boolean dropRecordOnPartitionMis return this; } + public Builder setHasMultipleStreams(boolean hasMultipleStreams) { + _hasMultipleStreams = hasMultipleStreams; + return this; + } + public RealtimeSegmentConfig build() { Map indexConfigByCol = Maps.newHashMapWithExpectedSize(_indexConfigByCol.size()); for (Map.Entry entry : _indexConfigByCol.entrySet()) { @@ -439,7 +452,8 @@ public RealtimeSegmentConfig build() { _capacity, _avgNumMultiValues, Collections.unmodifiableMap(indexConfigByCol), _segmentZKMetadata, _offHeap, _memoryManager, _statsHistory, _partitionColumn, _partitionFunction, _partitionId, _aggregateMetrics, _ingestionAggregationConfigs, _defaultNullHandlingEnabled, _partitionUpsertMetadataManager, - _partitionDedupMetadataManager, _consumerDir, _textIndexConfig, _dropRecordOnPartitionMismatch); + _partitionDedupMetadataManager, _consumerDir, _textIndexConfig, _dropRecordOnPartitionMismatch, + _hasMultipleStreams); } } } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java index b8efa78d4beb..48afd57ce04b 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/MultiColumnRealtimeLuceneTextIndex.java @@ -83,13 +83,15 @@ public class MultiColumnRealtimeLuceneTextIndex implements MultiColumnTextIndexR * @param segmentIndexDir realtime segment consumer dir * @param segmentName realtime segment name * @param mcTextConfig the table index config + * @param hasMultipleStreams whether the table has multiple stream configs */ public MultiColumnRealtimeLuceneTextIndex( List columns, BooleanList columnsSV, File segmentIndexDir, String segmentName, - MultiColumnTextIndexConfig mcTextConfig) { + MultiColumnTextIndexConfig mcTextConfig, + boolean hasMultipleStreams) { _columns = columns; _segmentName = segmentName; try { @@ -112,7 +114,7 @@ public MultiColumnRealtimeLuceneTextIndex( IndexWriter indexWriter = _indexCreator.getIndexWriter(); _searcherManager = new SearcherManager(indexWriter, false, false, null); - LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); + LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName, hasMultipleStreams); _refreshListener = new RealtimeLuceneRefreshListener(llcSegmentName.getTableName(), segmentName, MultiColumnTextIndexConstants.INDEX_DIR_NAME, llcSegmentName.getTopicPartitionId().getPartitionId(), _indexCreator::getNumDocs); @@ -135,6 +137,13 @@ public MultiColumnRealtimeLuceneTextIndex( } } + /** @deprecated Use the constructor with {@code hasMultipleStreams} parameter. */ + @Deprecated + public MultiColumnRealtimeLuceneTextIndex(List columns, BooleanList columnsSV, File segmentIndexDir, + String segmentName, MultiColumnTextIndexConfig mcTextConfig) { + this(columns, columnsSV, segmentIndexDir, segmentName, mcTextConfig, false); + } + public void add(List values) { _indexCreator.add(values); } diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java index e0adc6f49d8a..8531f4e72898 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/impl/invertedindex/RealtimeLuceneTextIndex.java @@ -71,8 +71,10 @@ public class RealtimeLuceneTextIndex implements MutableTextIndex { * @param segmentIndexDir realtime segment consumer dir * @param segmentName realtime segment name * @param config the table index config + * @param hasMultipleStreams whether the table has multiple stream configs */ - public RealtimeLuceneTextIndex(String column, File segmentIndexDir, String segmentName, TextIndexConfig config) { + public RealtimeLuceneTextIndex(String column, File segmentIndexDir, String segmentName, TextIndexConfig config, + boolean hasMultipleStreams) { _column = column; _segmentName = segmentName; try { @@ -90,7 +92,7 @@ public RealtimeLuceneTextIndex(String column, File segmentIndexDir, String segme IndexWriter indexWriter = _indexCreator.getIndexWriter(); _searcherManager = new SearcherManager(indexWriter, false, false, null); - LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); + LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName, hasMultipleStreams); _refreshListener = new RealtimeLuceneRefreshListener(llcSegmentName.getTableName(), segmentName, column, llcSegmentName.getTopicPartitionId().getPartitionId(), _indexCreator::getNumDocs); _searcherManager.addListener(_refreshListener); @@ -111,6 +113,12 @@ public RealtimeLuceneTextIndex(String column, File segmentIndexDir, String segme } } + /** @deprecated Use {@link #RealtimeLuceneTextIndex(String, File, String, TextIndexConfig, boolean)}. */ + @Deprecated + public RealtimeLuceneTextIndex(String column, File segmentIndexDir, String segmentName, TextIndexConfig config) { + this(column, segmentIndexDir, segmentName, config, false); + } + /** * Adds a new document. */ diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java index 7577e4484e32..9fae4181efe6 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/realtime/writer/StatelessRealtimeSegmentWriter.java @@ -122,10 +122,12 @@ public StatelessRealtimeSegmentWriter(SegmentZKMetadata segmentZKMetadata, Index throws Exception { Preconditions.checkNotNull(indexLoadingConfig.getTableConfig(), "Table config must be set in index loading config"); Preconditions.checkNotNull(indexLoadingConfig.getSchema(), "Schema must be set in index loading config"); - LLCSegmentName llcSegmentName = new LLCSegmentName(segmentZKMetadata.getSegmentName()); + LLCSegmentName llcSegmentName = + new LLCSegmentName(segmentZKMetadata.getSegmentName(), + IngestionConfigUtils.hasMultipleStreams(indexLoadingConfig.getTableConfig())); _segmentName = segmentZKMetadata.getSegmentName(); - _partitionGroupId = llcSegmentName.getTopicPartitionId().getPartitionId(); + _partitionGroupId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); _segBuildSemaphore = segBuildSemaphore; _tableNameWithType = TableNameBuilder.forType(TableType.REALTIME).tableNameWithType(llcSegmentName.getTableName()); _segmentZKMetadata = segmentZKMetadata; @@ -196,7 +198,8 @@ public StatelessRealtimeSegmentWriter(SegmentZKMetadata segmentZKMetadata, Index .setConsumerDir(_resourceDataDir.getAbsolutePath()) .setDropRecordOnPartitionMismatch(ingestionConfig != null && ingestionConfig.getStreamIngestionConfig() != null - && ingestionConfig.getStreamIngestionConfig().isDropRecordOnPartitionMismatch()); + && ingestionConfig.getStreamIngestionConfig().isDropRecordOnPartitionMismatch()) + .setHasMultipleStreams(IngestionConfigUtils.hasMultipleStreams(_tableConfig)); setPartitionParameters(realtimeSegmentConfigBuilder, _tableConfig.getIndexingConfig().getSegmentPartitionConfig()); diff --git a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/TextIndexType.java b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/TextIndexType.java index 0324d8e950e3..0c6aaa433b53 100644 --- a/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/TextIndexType.java +++ b/pinot-segment-local/src/main/java/org/apache/pinot/segment/local/segment/index/text/TextIndexType.java @@ -174,6 +174,6 @@ public MutableIndex createMutableIndex(MutableIndexContext context, TextIndexCon return null; } return new RealtimeLuceneTextIndex(context.getFieldSpec().getName(), context.getConsumerDir(), - context.getSegmentName(), config); + context.getSegmentName(), config, context.hasMultipleStreams()); } } diff --git a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/mutable/provider/MutableIndexContext.java b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/mutable/provider/MutableIndexContext.java index 57d71aa27882..693de3baa3fe 100644 --- a/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/mutable/provider/MutableIndexContext.java +++ b/pinot-segment-spi/src/main/java/org/apache/pinot/segment/spi/index/mutable/provider/MutableIndexContext.java @@ -36,10 +36,11 @@ public class MutableIndexContext { private final String _segmentName; private final PinotDataBufferMemoryManager _memoryManager; private final File _consumerDir; + private final boolean _hasMultipleStreams; public MutableIndexContext(FieldSpec fieldSpec, int fixedLengthBytes, boolean hasDictionary, String segmentName, PinotDataBufferMemoryManager memoryManager, int capacity, boolean offHeap, int estimatedColSize, - int estimatedCardinality, int avgNumMultiValues, File consumerDir) { + int estimatedCardinality, int avgNumMultiValues, File consumerDir, boolean hasMultipleStreams) { _fieldSpec = fieldSpec; _fixedLengthBytes = fixedLengthBytes; _hasDictionary = hasDictionary; @@ -51,6 +52,7 @@ public MutableIndexContext(FieldSpec fieldSpec, int fixedLengthBytes, boolean ha _estimatedCardinality = estimatedCardinality; _avgNumMultiValues = avgNumMultiValues; _consumerDir = consumerDir; + _hasMultipleStreams = hasMultipleStreams; } public PinotDataBufferMemoryManager getMemoryManager() { @@ -97,6 +99,10 @@ public File getConsumerDir() { return _consumerDir; } + public boolean hasMultipleStreams() { + return _hasMultipleStreams; + } + public static Builder builder() { return new Builder(); } @@ -113,6 +119,7 @@ public static class Builder { private int _estimatedCardinality; private int _avgNumMultiValues; private File _consumerDir; + private boolean _hasMultipleStreams; public Builder withMemoryManager(PinotDataBufferMemoryManager memoryManager) { _memoryManager = memoryManager; @@ -164,6 +171,11 @@ public Builder withConsumerDir(File consumerDir) { return this; } + public Builder withHasMultipleStreams(boolean hasMultipleStreams) { + _hasMultipleStreams = hasMultipleStreams; + return this; + } + public Builder withFixedLengthBytes(int fixedLengthBytes) { _fixedLengthBytes = fixedLengthBytes; return this; @@ -172,7 +184,7 @@ public Builder withFixedLengthBytes(int fixedLengthBytes) { public MutableIndexContext build() { return new MutableIndexContext(Objects.requireNonNull(_fieldSpec), _fixedLengthBytes, _hasDictionary, Objects.requireNonNull(_segmentName), Objects.requireNonNull(_memoryManager), _capacity, _offHeap, - _estimatedColSize, _estimatedCardinality, _avgNumMultiValues, _consumerDir); + _estimatedColSize, _estimatedCardinality, _avgNumMultiValues, _consumerDir, _hasMultipleStreams); } } } 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..2f5ee241c456 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 @@ -33,6 +33,9 @@ * 6. status - the consumption status IN_PROGRESS/DONE * * This information is needed by the stream, when grouping the partitions/shards into new partition groups. + * + * TODO: Add an overloaded constructor that accepts TopicPartitionId instead of int partitionGroupId, + * to avoid callers needing to decompose/recompose composite partition IDs. */ public class PartitionGroupConsumptionStatus { private final int _partitionGroupId; diff --git a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java index b0731ad38315..6aa57670d1d5 100644 --- a/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java +++ b/pinot-spi/src/main/java/org/apache/pinot/spi/utils/IngestionConfigUtils.java @@ -124,7 +124,9 @@ public static int getPinotPartitionIdFromStreamPartitionId(int partitionId, int } /// Returns the stream partition id from the Pinot segment partition id. - /// Returns {@code partitionId} unchanged for tables without multiple streams. + /// Safe to call for any table type: returns `partitionId` unchanged for OFFLINE tables or REALTIME tables + /// that lack stream configs (treat as single-stream). + @Deprecated public static int getStreamPartitionIdFromPinotPartitionId(TableConfig tableConfig, int partitionId) { return hasMultipleStreams(tableConfig) ? getStreamPartitionIdFromPinotPartitionId(partitionId) : partitionId; } @@ -132,6 +134,7 @@ public static int getStreamPartitionIdFromPinotPartitionId(TableConfig tableConf /// Returns the stream partition id from the Pinot segment partition id. /// NOTE: First verify if there are multiple stream configs before invoking this method. User might plug in a stream /// that generates large partition id. + @Deprecated public static int getStreamPartitionIdFromPinotPartitionId(int partitionId) { return partitionId % PARTITION_PADDING_OFFSET; } @@ -145,6 +148,7 @@ public static StreamConfig getStreamConfigFromPinotPartitionId(List> getStreamConfigIndexToStreamPartitions(Set pinotPartitionIds) { Map> streamIndexToPartitions = new HashMap<>(); for (Integer partition : pinotPartitionIds) {