missed one util method that should derive partition id from uploaded …

…segment name (apache#13612)

* add one util method to derive partition id from uploaded segment names to avoid missing the segment name checks

* add tests

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

@@ -39,17 +39,10 @@ private SegmentUtils() {
   @Nullable
   public static Integer getRealtimeSegmentPartitionId(String segmentName, String realtimeTableName,
       HelixManager helixManager, @Nullable String partitionColumn) {
-    // A fast path if the segmentName is an LLC segment name: get the partition id from the name directly
-    LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName);
-    if (llcSegmentName != null) {
-      return llcSegmentName.getPartitionGroupId();
+    Integer partitionId = getPartitionIdFromRealtimeSegmentName(segmentName);
+    if (partitionId != null) {
+      return partitionId;
     }
-
-    UploadedRealtimeSegmentName uploadedRealtimeSegmentName = UploadedRealtimeSegmentName.of(segmentName);
-    if (uploadedRealtimeSegmentName != null) {
-      return uploadedRealtimeSegmentName.getPartitionId();
-    }
-
     // Otherwise, retrieve the partition id from the segment zk metadata.
     SegmentZKMetadata segmentZKMetadata =
         ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), realtimeTableName, segmentName);
@@ -61,13 +54,26 @@ public static Integer getRealtimeSegmentPartitionId(String segmentName, String r
   @Nullable
   public static Integer getRealtimeSegmentPartitionId(String segmentName, SegmentZKMetadata segmentZKMetadata,
       @Nullable String partitionColumn) {
-    // A fast path if the segmentName is an LLC segment name: get the partition id from the name directly
+    Integer partitionId = getPartitionIdFromRealtimeSegmentName(segmentName);
+    if (partitionId != null) {
+      return partitionId;
+    }
+    // Otherwise, retrieve the partition id from the segment zk metadata.
+    return getRealtimeSegmentPartitionId(segmentZKMetadata, partitionColumn);
+  }
+
+  @Nullable
+  private static Integer getPartitionIdFromRealtimeSegmentName(String segmentName) {
+    // A fast path to get partition id if the segmentName is in a known format like LLC.
     LLCSegmentName llcSegmentName = LLCSegmentName.of(segmentName);
     if (llcSegmentName != null) {
       return llcSegmentName.getPartitionGroupId();
     }
-    // Otherwise, retrieve the partition id from the segment zk metadata.
-    return getRealtimeSegmentPartitionId(segmentZKMetadata, partitionColumn);
+    UploadedRealtimeSegmentName uploadedRealtimeSegmentName = UploadedRealtimeSegmentName.of(segmentName);
+    if (uploadedRealtimeSegmentName != null) {
+      return uploadedRealtimeSegmentName.getPartitionId();
+    }
+    return null;
   }
 
   @Nullable

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

@@ -88,11 +88,20 @@ void testGetRealtimeSegmentPartitionIdForUploadedRealtimeSegment() {
     String segmentName = "uploaded__table_name__3__100__1716185755000";
 
     try {
+      // Check the util method that gets segmentZKMetadata via HelixManager for partition id.
       Integer partitionId =
           SegmentUtils.getRealtimeSegmentPartitionId(segmentName, "realtimeTableName", null, "partitionColumn");
       assertEquals(partitionId, 3);
     } catch (Exception e) {
       fail("Exception should not be thrown");
     }
+
+    try {
+      // Check the util method that has segmentZKMetadata passed in directly for partition id.
+      Integer partitionId = SegmentUtils.getRealtimeSegmentPartitionId(segmentName, null, "partitionColumn");
+      assertEquals(partitionId, 3);
+    } catch (Exception e) {
+      fail("Exception should not be thrown");
+    }
   }
 }