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..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,50 +28,76 @@ 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'"; 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; - 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]; - _partitionGroupId = 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; - _partitionGroupId = partitionGroupId; + _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. */ @@ -101,8 +127,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 +158,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..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.getPartitionGroupId(); + 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 new file mode 100644 index 000000000000..24c9133eea8f --- /dev/null +++ b/pinot-common/src/main/java/org/apache/pinot/common/utils/TopicPartitionId.java @@ -0,0 +1,109 @@ +/** + * 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 com.google.common.base.Preconditions; +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 { + // 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; + + 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); + } + + /** 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 04b2a9e2fbf2..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 @@ -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); @@ -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 new file mode 100644 index 000000000000..46c41e3a870b --- /dev/null +++ b/pinot-common/src/test/java/org/apache/pinot/common/utils/TopicPartitionIdTest.java @@ -0,0 +1,136 @@ +/** + * 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"); + } + + @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 cbda557d411a..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.getPartitionGroupId(); + 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.getPartitionGroupId(); + 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 02855d88fd92..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,14 +211,15 @@ 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); if (llcSegmentName == null) { continue; } - partitionIdToSegments.computeIfAbsent(llcSegmentName.getPartitionGroupId(), k -> new TreeSet<>()) + partitionIdToSegments.computeIfAbsent( + 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 605ade8a2d3b..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 631c6609b114..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 9c17cc127dc4..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.getPartitionGroupId(); + 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 9bc0d4e8a70e..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.getPartitionGroupId(); + 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.getPartitionGroupId(); + 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); @@ -1007,7 +1010,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()); if (!startOffset.equals(nextOffset)) { _controllerMetrics.addMeteredTableValue(realtimeTableName, ControllerMeter.OFFSET_AUTO_RESET_SKIPPED_OFFSETS, Long.parseLong(startOffset) - Long.parseLong(nextOffset)); @@ -1016,7 +1020,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 +1038,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(), numPartitions); if (partitionMetadata != null) { newSegmentZKMetadata.setPartitionMetadata(partitionMetadata); } @@ -1072,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; } @@ -1085,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(); @@ -1135,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; @@ -1145,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. @@ -1233,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; } @@ -1250,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) { @@ -1259,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; @@ -1301,7 +1307,8 @@ private PartitionIdsWithIdealState getPartitionIdsWithIdealState(List 1)); } } return new PartitionIdsWithIdealState(partitionIds, idealState); @@ -1423,27 +1430,29 @@ 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.getPartitionGroupId(), (partitionId, latestLLCSegmentName) -> { - if (latestLLCSegmentName == null) { - return llcSegmentName; - } else { - if (llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { - return llcSegmentName; - } else { - return latestLLCSegmentName; - } - } - }); + latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId(), + (partitionId, latestLLCSegmentName) -> { + if (latestLLCSegmentName == null) { + return llcSegmentName; + } else { + if (llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { + return llcSegmentName; + } else { + return latestLLCSegmentName; + } + } + }); } - 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); @@ -1581,7 +1590,7 @@ public static boolean isTopicPaused(IdealState idealState, String segmentName) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { return isTopicPaused(idealState, - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(llcSegmentName.getPartitionGroupId())); + llcSegmentName.getTopicPartitionId().getTopicId()); } return false; } @@ -1630,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; } @@ -1823,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) { @@ -1847,7 +1861,8 @@ IdealState ensureAllPartitionsConsuming(TableConfig tableConfig, List> segmentEntry : instanceStatesMap.entrySet()) { if (segmentEntry.getValue().containsValue(SegmentStateModel.CONSUMING) - && new LLCSegmentName(segmentEntry.getKey()).getPartitionGroupId() == partitionId) { + && new LLCSegmentName(segmentEntry.getKey()).getTopicPartitionId() + .equals(partitionId)) { previousConsumingSegment = segmentEntry.getKey(); break; } @@ -1978,7 +1994,8 @@ && new LLCSegmentName(segmentEntry.getKey()).getPartitionGroupId() == partitionI // 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, @@ -1999,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, @@ -2009,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); @@ -2029,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()); } } } @@ -2042,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]; @@ -2071,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, @@ -2081,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 @@ -2104,7 +2129,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); } @@ -2700,12 +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).getPartitionGroupId())) + .filter(segmentName -> partitionsToCommit.contains( + new LLCSegmentName(segmentName).getTopicPartitionId())) .collect(Collectors.toSet()); Preconditions.checkState(!targetSegments.isEmpty(), "Cannot find segments to commit for partitions: %s", partitionGroupIdsToCommitStr); @@ -2893,7 +2922,7 @@ private Set findConsumingSegmentsOfTopics(IdealState idealState, List { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null && !topicIndices.contains( - IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(llcSegmentName.getPartitionGroupId()))) { + 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 1ed304740f5f..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.getPartitionGroupId(); + 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 9e90ef8f2a2e..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.getPartitionGroupId(); + 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 ffd37a2d54a1..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.getPartitionGroupId(); + 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 45e85b57e640..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.getPartitionGroupId(); + int partitionGroupId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); int sequence = llcSegmentName.getSequenceNumber(); if (partitionGroupId == 0) { assertEquals(sequence, 5); @@ -296,7 +297,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().toMultiTopicPinotPartitionId()); } } @@ -777,7 +779,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().toMultiTopicPinotPartitionId() + 1); } // Check that for new partition groups, each partition group should have exactly 1 new segment in CONSUMING @@ -786,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).getPartitionGroupId(); + int partitionGroupId = new LLCSegmentName(segmentName).getTopicPartitionId().toMultiTopicPinotPartitionId(); partitionGroupIdToSegmentsMap.computeIfAbsent(partitionGroupId, k -> new ArrayList<>()).add(segmentName); } for (int partitionGroupId = oldNumPartitions; partitionGroupId < segmentManager._numPartitions; @@ -1154,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.getPartitionGroupId(); + int partitionsId = llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(); Map sequenceNumberToSegmentMap = partitionGroupIdToSegmentsMap.get(partitionsId); int sequenceNumber = llcSegmentName.getSequenceNumber(); assertFalse(sequenceNumberToSegmentMap.containsKey(sequenceNumber)); @@ -1952,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(); @@ -1993,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.getPartitionGroupId(); + 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(); @@ -2372,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); @@ -2392,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"); @@ -2409,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 da98f32aba85..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.getPartitionGroupId(); + 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 cc99a15e1679..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.getPartitionGroupId(), ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, waitTimeMs); + llcSegmentName.getTopicPartitionId().toMultiTopicPinotPartitionId(), + ServerGauge.CONSUMER_LOCK_WAIT_TIME_MS, waitTimeMs); } } finally { _serverMetrics.setValueOfPartitionGauge(_realtimeTableDataManager.getTableName(), - llcSegmentName.getPartitionGroupId(), 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.getPartitionGroupId(); + 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.getPartitionGroupId() != 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 e26390b9e4fb..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).getPartitionGroupId()); + 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,9 @@ 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(), + _clock.millis()); } /** @@ -545,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; } @@ -559,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; @@ -579,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; @@ -590,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) @@ -598,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; @@ -615,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; @@ -644,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; @@ -667,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; @@ -698,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 0fbe75597d44..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.getPartitionGroupId(); + _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 bb65dc48dbb5..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); @@ -351,7 +354,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()); } /** @@ -372,7 +376,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()); } /** @@ -407,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.getPartitionGroupId()); + partitionsHostedByThisServer.add( + new LLCSegmentName(segmentNameStr).getTopicPartitionId()); } return partitionsHostedByThisServer; } @@ -524,7 +529,8 @@ protected void doAddOnlineSegment(String segmentName) if (_enforceConsumptionInOrder) { LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName); if (llcSegmentName != null) { - getConsumerCoordinator(llcSegmentName.getPartitionGroupId()).register(llcSegmentName); + getConsumerCoordinator(llcSegmentName.getTopicPartitionId()) + .register(llcSegmentName); } } } @@ -549,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; } @@ -599,15 +617,19 @@ private void doAddConsumingSegment(String segmentName) // Generates only one semaphore for every partition LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName); - int partitionGroupId = llcSegmentName.getPartitionGroupId(); + 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, @@ -851,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 41fb745acd40..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.getPartitionGroupId()), 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 28706764a530..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).getPartitionGroupId()); + 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 fb214936fad1..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.getPartitionGroupId() == 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.getPartitionGroupId() == 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 98cae2d48d69..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.getPartitionGroupId() == 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 7b101d49d527..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.getPartitionGroupId() == 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 d636a61351f2..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).getPartitionGroupId(); + 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 5c578d20b5ec..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 @@ -244,7 +244,9 @@ 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().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 201fd136e094..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 @@ -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().toMultiTopicPinotPartitionId(); 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().toMultiTopicPinotPartitionId(); 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().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 1dd1eb1f2422..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).getPartitionGroupId(); + 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 2e9f472d9853..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).getPartitionGroupId(); + 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 d2013f22399f..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).getPartitionGroupId()); + 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).getPartitionGroupId(); + 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 fc7f5264c016..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 @@ -70,7 +70,9 @@ 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().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 839d58236664..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,19 +601,21 @@ 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.getPartitionGroupId(), (partId, latestSegment) -> { - if (latestSegment == null) { - return llcSegmentName; - } else { - return latestSegment.getSequenceNumber() > llcSegmentName.getSequenceNumber() + LLCSegmentName llcSegmentName = new LLCSegmentName(segmentName, hasMultipleStreams); + partitionIdToLatestCompletedSegment.compute(llcSegmentName.getTopicPartitionId(), + (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..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,20 +165,23 @@ 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.getPartitionGroupId(), (k, latestLLCSegmentName) -> { - if (latestLLCSegmentName == null - || llcSegmentName.getSequenceNumber() > latestLLCSegmentName.getSequenceNumber()) { - return llcSegmentName; - } else { - return latestLLCSegmentName; - } - }); + latestLLCSegmentNameMap.compute(llcSegmentName.getTopicPartitionId(), + (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..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.getPartitionGroupId(); + 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 d0ae144ae5e8..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,9 +453,14 @@ 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.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/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 884227d31ec5..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,10 +114,10 @@ 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.getPartitionGroupId(), _indexCreator::getNumDocs); + llcSegmentName.getTopicPartitionId().getPartitionId(), _indexCreator::getNumDocs); _searcherManager.addListener(_refreshListener); _analyzer = _indexCreator.getIndexWriter().getConfig().getAnalyzer(); _queryParserClassConstructor = @@ -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 47b7ef75f189..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,9 +92,9 @@ 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.getPartitionGroupId(), _indexCreator::getNumDocs); + llcSegmentName.getTopicPartitionId().getPartitionId(), _indexCreator::getNumDocs); _searcherManager.addListener(_refreshListener); _analyzer = _indexCreator.getIndexWriter().getConfig().getAnalyzer(); _queryParserClassConstructor = @@ -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 2d0db9781864..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.getPartitionGroupId(); + _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) {