@Nullable
  public static Integer getRealtimeSegmentPartitionId(String segmentName, String realtimeTableName,
      HelixManager helixManager, @Nullable String partitionColumn) {
    Integer partitionId = getPartitionIdFromRealtimeSegmentName(segmentName);
    if (partitionId != null) {
      return partitionId;
    }
    // Otherwise, retrieve the partition id from the segment zk metadata.
    SegmentZKMetadata segmentZKMetadata =
        ZKMetadataProvider.getSegmentZKMetadata(helixManager.getHelixPropertyStore(), realtimeTableName, segmentName);
    Preconditions.checkState(segmentZKMetadata != null,
        "Failed to find segment ZK metadata for segment: %s of table: %s", segmentName, realtimeTableName);
    return getRealtimeSegmentPartitionId(segmentZKMetadata, partitionColumn);
  }
@Test
  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");
    }
  }