Use TopicPartitionId as primary partition identity across all internal code

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<LLCSegmentName> {
   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.
+   *
+   * <p>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);
   }

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

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.
+ *
+ * <p>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.
+ *
+ * <p>Thread-safe: instances are immutable.
+ */
+public final class TopicPartitionId implements Comparable<TopicPartitionId> {
+  // 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;
+  }
+}

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,15 +58,15 @@ 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);
     Assert.assertEquals(segName1.getTableName(), tableName);
 
     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);
+  }
 }