Fix status check for downloads to also include UPLOADED status

apache · Jan 10, 2025 · 81ea9ee · 81ea9ee
1 parent f4b1420
commit 81ea9ee
Showing 1 changed file with 7 additions and 3 deletions.
diff --git a/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java b/pinot-core/src/main/java/org/apache/pinot/core/data/manager/BaseTableDataManager.java
@@ -653,7 +653,7 @@ public void reloadSegment(String segmentName, IndexLoadingConfig indexLoadingCon
       - Continue loading the segment from the index directory.
       */
       boolean shouldDownload =
-          forceDownload || (zkMetadata.getStatus() == CommonConstants.Segment.Realtime.Status.DONE && !hasSameCRC(
+          forceDownload || (isSegmentStatusCompleted(zkMetadata) && !hasSameCRC(
               zkMetadata, localMetadata));
       if (shouldDownload) {
         // Create backup directory to handle failure of segment reloading.
@@ -717,6 +717,11 @@ public void reloadSegment(String segmentName, IndexLoadingConfig indexLoadingCon
     _logger.info("Reloaded segment: {}", segmentName);
   }
 
+  private boolean isSegmentStatusCompleted(SegmentZKMetadata zkMetadata) {
+    return zkMetadata.getStatus() == CommonConstants.Segment.Realtime.Status.DONE
+        || zkMetadata.getStatus() == CommonConstants.Segment.Realtime.Status.UPLOADED;
+  }
+
   private boolean canReuseExistingDirectoryForReload(SegmentZKMetadata segmentZKMetadata, String currentSegmentTier,
       SegmentDirectory segmentDirectory, IndexLoadingConfig indexLoadingConfig, Schema schema)
       throws Exception {
@@ -1011,8 +1016,7 @@ public boolean tryLoadExistingSegment(SegmentZKMetadata zkMetadata, IndexLoading
     Then:
     We need to fall back to downloading the segment from deep storage to load it.
     */
-    if (segmentMetadata == null || (zkMetadata.getStatus() == CommonConstants.Segment.Realtime.Status.DONE
-        && !hasSameCRC(zkMetadata, segmentMetadata))) {
+    if (segmentMetadata == null || (isSegmentStatusCompleted(zkMetadata) && !hasSameCRC(zkMetadata, segmentMetadata))) {
       if (segmentMetadata == null) {
         _logger.info("Segment: {} does not exist", segmentName);
       } else if (!hasSameCRC(zkMetadata, segmentMetadata)) {