Add 5 part multi topic segment name

pinot-common/src/main/java/org/apache/pinot/common/utils/LLCSegmentName.java

@@ -23,6 +23,7 @@
 import java.util.Objects;
 import javax.annotation.Nullable;
 import org.apache.commons.lang3.StringUtils;
+import org.apache.pinot.spi.utils.IngestionConfigUtils;
 import org.joda.time.DateTime;
 import org.joda.time.format.DateTimeFormat;
 import org.joda.time.format.DateTimeFormatter;
@@ -34,31 +35,57 @@ public class LLCSegmentName implements Comparable<LLCSegmentName> {
   private static final DateTimeFormatter DATE_FORMATTER = DateTimeFormat.forPattern(DATE_FORMAT).withZoneUTC();
 
   private final String _tableName;
+  private final int _topicId;
   private final int _partitionGroupId;
   private final int _sequenceNumber;
   private final String _creationTime;
   private final String _segmentName;
+  private final boolean _isMultiTopicFormat;
 
   public LLCSegmentName(String segmentName) {
     String[] parts = StringUtils.splitByWholeSeparator(segmentName, SEPARATOR);
-    Preconditions.checkArgument(parts.length == 4, "Invalid LLC segment name: %s", segmentName);
+    Preconditions.checkArgument(parts.length == 4 || parts.length == 5, "Invalid LLC segment name: %s", segmentName);
     _tableName = parts[0];
-    _partitionGroupId = Integer.parseInt(parts[1]);
-    _sequenceNumber = Integer.parseInt(parts[2]);
-    _creationTime = parts[3];
+    if (parts.length == 4) {
+      _topicId = 0;
+      _partitionGroupId = Integer.parseInt(parts[1]);
+      _sequenceNumber = Integer.parseInt(parts[2]);
+      _creationTime = parts[3];
+      _isMultiTopicFormat = false;
+    } else {
+      _topicId = Integer.parseInt(parts[1]);
+      _partitionGroupId = Integer.parseInt(parts[2]);
+      _sequenceNumber = Integer.parseInt(parts[3]);
+      _creationTime = parts[4];
+      _isMultiTopicFormat = true;
+    }
     _segmentName = segmentName;
   }
 
   public LLCSegmentName(String tableName, int partitionGroupId, int sequenceNumber, long msSinceEpoch) {
     Preconditions.checkArgument(!tableName.contains(SEPARATOR), "Illegal table name: %s", tableName);
     _tableName = tableName;
+    _topicId = 0;
     _partitionGroupId = partitionGroupId;
     _sequenceNumber = sequenceNumber;
     // ISO8601 date: 20160120T1234Z
     _creationTime = DATE_FORMATTER.print(msSinceEpoch);
+    _isMultiTopicFormat = false;
     _segmentName = tableName + SEPARATOR + partitionGroupId + SEPARATOR + sequenceNumber + SEPARATOR + _creationTime;
   }
 
+  public LLCSegmentName(String tableName, int configId, int partitionGroupId, int sequenceNumber, long msSinceEpoch) {
+    Preconditions.checkArgument(!tableName.contains(SEPARATOR), "Illegal table name: %s", tableName);
+    _tableName = tableName;
+    _topicId = configId;
+    _partitionGroupId = partitionGroupId;
+    _sequenceNumber = sequenceNumber;
+    _creationTime = DATE_FORMATTER.print(msSinceEpoch);
+    _isMultiTopicFormat = false;
+    _segmentName = tableName + SEPARATOR + configId + SEPARATOR + partitionGroupId + SEPARATOR + sequenceNumber
+        + SEPARATOR + _creationTime;
+  }
+
   /**
    * Returns the {@link LLCSegmentName} for the given segment name, or {@code null} if the given segment name does not
    * represent an LLC segment.
@@ -82,7 +109,26 @@ public static boolean isLLCSegment(String segmentName) {
       numSeparators++;
       index += 2; // SEPARATOR.length()
     }
-    return numSeparators == 3;
+    return numSeparators == 3 || numSeparators == 4;
+  }
+
+  /**
+   * Returns whether the given segment name is a multi-topic LLC segment (5-part format where
+   * parts[1] and parts[3] are integers, distinguishing it from UploadedRealtimeSegmentName).
+   */
+  public static boolean isMultiTopicLLCSegment(String segmentName) {
+    String[] parts = StringUtils.splitByWholeSeparator(segmentName, SEPARATOR);
+    if (parts.length != 5) {
+      return false;
+    }
+    try {
+      Integer.parseInt(parts[1]); // topicId
+      Integer.parseInt(parts[2]); // partitionGroupId
+      Integer.parseInt(parts[3]); // sequenceNumber
+      return true;
+    } catch (NumberFormatException e) {
+      return false;
+    }
   }
 
   @Deprecated
@@ -101,6 +147,19 @@ public String getTableName() {
     return _tableName;
   }
 
+  public int getTopicId() {
+    return _topicId;
+  }
+
+  public int getTopicId(boolean hasMultiTopic) {
+    if (hasMultiTopic) {
+      return _isMultiTopicFormat
+          ? _topicId : IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(_partitionGroupId);
+    } else {
+      return _topicId;
+    }
+  }
+
   public int getPartitionGroupId() {
     return _partitionGroupId;
   }
@@ -118,15 +177,30 @@ public long getCreationTimeMs() {
     return dateTime.getMillis();
   }
 
+  public boolean isMultiTopicFormat() {
+    return _isMultiTopicFormat;
+  }
+
   @JsonValue
   public String getSegmentName() {
     return _segmentName;
   }
 
+  public int getStreamPartitionGroupId(boolean hasMultiTopics) {
+    if (hasMultiTopics) {
+      return _isMultiTopicFormat
+          ? _partitionGroupId : IngestionConfigUtils.getStreamPartitionIdFromPinotPartitionId(_partitionGroupId);
+    } else {
+      return _partitionGroupId;
+    }
+  }
+
   @Override
   public int compareTo(LLCSegmentName other) {
     Preconditions.checkArgument(_tableName.equals(other._tableName),
         "Cannot compare segment names from different table: %s, %s", _segmentName, other.getSegmentName());
+    Preconditions.checkArgument(_topicId == other._topicId,
+        "Cannot compare segment names from different topic ids: %s, %s", _segmentName, other.getSegmentName());
     if (_partitionGroupId != other._partitionGroupId) {
       return Integer.compare(_partitionGroupId, other._partitionGroupId);
     }

pinot-common/src/test/java/org/apache/pinot/common/utils/LLCSegmentNameTest.java

@@ -31,6 +31,8 @@
  * Tests for the realtime segment name builder.
  */
 public class LLCSegmentNameTest {
+  private static final String TABLE_NAME = "myTable";
+  private static final long TIMESTAMP_MS = 1465508537069L;
 
   @Test
   public void testSegmentNameBuilder() {
@@ -95,4 +97,70 @@ public void testLLCSegmentName() {
     Arrays.sort(testSorted);
     Assert.assertEquals(testSorted, new LLCSegmentName[]{segName5, segName1, segName6, segName3, segName4});
   }
+
+  @Test
+  public void testSegmentNameWithoutTopicIdDefaultsTopicIdToZero() {
+    LLCSegmentName segmentName = new LLCSegmentName("myTable__4__27__20160617T2150Z");
+    assertEquals(segmentName.getTopicId(), 0);
+  }
+
+  @Test
+  public void testSegmentNameWithTopicIdParsesTopicIdAndOtherFields() {
+    LLCSegmentName segmentName = new LLCSegmentName("myTable__2__4__27__20160617T2150Z");
+    assertEquals(segmentName.getTopicId(), 2);
+    assertEquals(segmentName.getTableName(), "myTable");
+    assertEquals(segmentName.getPartitionGroupId(), 4);
+    assertEquals(segmentName.getSequenceNumber(), 27);
+    assertEquals(segmentName.getCreationTime(), "20160617T2150Z");
+    assertTrue(LLCSegmentName.isLLCSegment(segmentName.getSegmentName()));
+  }
+
+  @Test
+  public void testConstructorWithTopicIdBuildsFivePartSegmentName() {
+    LLCSegmentName segmentName = new LLCSegmentName(TABLE_NAME, 2, 0, 1, TIMESTAMP_MS);
+    assertEquals(segmentName.getSegmentName(), "myTable__2__0__1__20160609T2142Z");
+    assertEquals(segmentName.getTopicId(), 2);
+  }
+
+  @Test
+  public void testConstructorWithoutTopicIdDefaultsTopicIdToZero() {
+    LLCSegmentName segmentName = new LLCSegmentName(TABLE_NAME, 0, 1, TIMESTAMP_MS);
+    assertEquals(segmentName.getTopicId(), 0);
+  }
+
+  @Test
+  public void testIsMultiTopicLLCSegmentRecognizesFivePartNumericSegmentName() {
+    assertTrue(LLCSegmentName.isMultiTopicLLCSegment("myTable__2__4__27__20160617T2150Z"));
+  }
+
+  @Test
+  public void testIsMultiTopicLLCSegmentRejectsFourPartSegmentName() {
+    assertFalse(LLCSegmentName.isMultiTopicLLCSegment("myTable__4__27__20160617T2150Z"));
+  }
+
+  @Test
+  public void testIsMultiTopicLLCSegmentRejectsNonNumericParts() {
+    assertFalse(LLCSegmentName.isMultiTopicLLCSegment("myTable__uploaded__4__27__20160617T2150Z"));
+  }
+
+  @Test
+  public void testCompareToThrowsWhenTopicIdsDiffer() {
+    LLCSegmentName segmentInTopic1 = new LLCSegmentName(TABLE_NAME, 1, 0, 1, TIMESTAMP_MS);
+    LLCSegmentName segmentInTopic2 = new LLCSegmentName(TABLE_NAME, 2, 0, 1, TIMESTAMP_MS);
+    try {
+      segmentInTopic1.compareTo(segmentInTopic2);
+      Assert.fail("Expected compareTo to throw for segments from different topic ids");
+    } catch (IllegalArgumentException e) {
+      // expected
+    }
+  }
+
+  @Test
+  public void testCompareToOrdersSegmentsWithSameNonZeroTopicId() {
+    LLCSegmentName segment1 = new LLCSegmentName(TABLE_NAME, 3, 0, 1, TIMESTAMP_MS);
+    LLCSegmentName segment2 = new LLCSegmentName(TABLE_NAME, 3, 0, 2, TIMESTAMP_MS);
+    assertTrue(segment1.compareTo(segment2) < 0);
+    assertTrue(segment2.compareTo(segment1) > 0);
+    assertEquals(segment1.compareTo(new LLCSegmentName(TABLE_NAME, 3, 0, 1, TIMESTAMP_MS)), 0);
+  }
 }

pinot-controller/src/main/java/org/apache/pinot/controller/api/resources/PinotTableRestletResource.java

@@ -393,14 +393,13 @@ List<StreamMetadata> getStreamMetadataList(List<StreamConfig> streamConfigs,
         _pinotHelixResourceManager.getRealtimeSegmentManager().getPartitionCountMap(streamConfigs);
     Map<Integer, List<PartitionGroupMetadata>> partitionGroupMetadataByStreamConfigIndex = new HashMap<>();
     for (WatermarkInductionResult.Watermark watermark : watermarkInductionResult.getWatermarks()) {
-      int streamConfigIndex =
-          IngestionConfigUtils.getStreamConfigIndexFromPinotPartitionId(watermark.getPartitionGroupId());
+      int streamConfigIndex = watermark.getTopicId();
       Preconditions.checkArgument(streamConfigIndex >= 0 && streamConfigIndex < streamConfigs.size(),
-          "Invalid stream config index %s from watermark partition ID %s. Expected index in range [0, %s)",
+          "Invalid stream config index %s from watermark (partitionGroupId=%s). Expected index in range [0, %s)",
           streamConfigIndex, watermark.getPartitionGroupId(), streamConfigs.size());
       partitionGroupMetadataByStreamConfigIndex.computeIfAbsent(streamConfigIndex, ignored -> new ArrayList<>()).add(
-          new PartitionGroupMetadata(watermark.getPartitionGroupId(), new LongMsgOffset(watermark.getOffset()),
-              watermark.getSequenceNumber()));
+          new PartitionGroupMetadata(watermark.getPartitionGroupId(), streamConfigIndex,
+              new LongMsgOffset(watermark.getOffset()), watermark.getSequenceNumber()));
     }
 
     // Iterate in order by streamConfigIndex to ensure deterministic ordering

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/PinotHelixResourceManager.java

@@ -5460,7 +5460,8 @@ public WatermarkInductionResult getConsumerWatermarks(String tableName)
       } catch (NumericException e) {
         throw new RuntimeException(e);
       }
-      return new WatermarkInductionResult.Watermark(status.getPartitionGroupId(), seq, startOffset);
+      return new WatermarkInductionResult.Watermark(status.getStreamPartitionGroupId(), status.getTopicId(), seq,
+          startOffset);
     }).collect(Collectors.toList());
     return new WatermarkInductionResult(watermarks);
   }

pinot-controller/src/main/java/org/apache/pinot/controller/helix/core/WatermarkInductionResult.java

@@ -57,6 +57,7 @@ public List<Watermark> getWatermarks() {
    */
   public static class Watermark {
     private int _partitionGroupId;
+    private int _topicId;
     private int _sequenceNumber;
     private long _offset;
 
@@ -65,28 +66,56 @@ public static class Watermark {
      * a WaterMark instance from a JSON object. The @JsonProperty annotations
      * map the keys in the JSON object to the constructor parameters.
      *
-     * @param partitionGroupId The ID of the partition group.
+     * @param partitionGroupId The stream partition id of the partition group, decoded from any legacy
+     *                         padded/composite id. Defaults topicId to 0 (single-topic tables).
+     * @param sequenceNumber The segment sequence number of the consuming segment.
+      * @param offset The first Kafka offset whose corresponding record has not yet sealed in Pinot
+     */
+    public Watermark(int partitionGroupId, int sequenceNumber, long offset) {
+      this(partitionGroupId, 0, sequenceNumber, offset);
+    }
+
+    /**
+     * The @JsonCreator annotation tells Jackson to use this constructor to create
+     * a WaterMark instance from a JSON object. The @JsonProperty annotations
+     * map the keys in the JSON object to the constructor parameters.
+     *
+     * @param partitionGroupId The stream partition id of the partition group (i.e. the partition id within its
+     *                         topic/stream config), not the raw/composite id used by legacy padded encodings.
+     * @param topicId The id of the topic (stream config) this partition group belongs to, for tables with
+     *                multiple topics.
      * @param sequenceNumber The segment sequence number of the consuming segment.
       * @param offset The first Kafka offset whose corresponding record has not yet sealed in Pinot
      */
     @JsonCreator
-    public Watermark(@JsonProperty("partitionGroupId") int partitionGroupId,
+    public Watermark(@JsonProperty("partitionGroupId") int partitionGroupId, @JsonProperty("topicId") int topicId,
         @JsonProperty("sequenceNumber") int sequenceNumber, @JsonProperty("offset") long offset) {
       _partitionGroupId = partitionGroupId;
+      _topicId = topicId;
       _sequenceNumber = sequenceNumber;
       _offset = offset;
     }
 
     /**
-     * Gets the partition group ID.
+     * Gets the stream partition id of the partition group.
      *
-     * @return The partition group ID.
+     * @return The stream partition id.
      */
     @JsonGetter("partitionGroupId")
     public int getPartitionGroupId() {
       return _partitionGroupId;
     }
 
+    /**
+     * Gets the topic ID (stream config index) this partition group belongs to.
+     *
+     * @return The topic ID.
+     */
+    @JsonGetter("topicId")
+    public int getTopicId() {
+      return _topicId;
+    }
+
     /**
      * Gets the segment sequence number of the most recent consuming segment.
      *