Add time-range based segment reload to controller API (#17627)

* Add time-range segment reload API

- Add controller endpoint to reload segments in a time range
- Select segments via metadata and dispatch per instance
- Add unit tests for time range reload behavior

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>

* Unify reload job ID handling: use single shared job ID for instanceToSegmentsMap path

The existing `instanceToSegmentsMap` reload path was creating per-instance ZK entries
with per-instance Helix message IDs, while the new time-range reload path used a single
shared UUID across all servers. Unify them by generating one UUID per logical reload
operation in reloadSegmentsForTable and storing a single ZK entry, consistent with the
time-range path.

Remove the now-unused 3-arg reloadSegments overload in PinotHelixResourceManager.

Co-Authored-By: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

---------

Co-authored-by: Claude Opus 4.6 <noreply@anthropic.com>

Author: xiangfu0

pinot-common/src/main/java/org/apache/pinot/common/messages/SegmentReloadMessage.java

@@ -56,6 +56,19 @@ public SegmentReloadMessage(String tableNameWithType, boolean forceDownload) {
    * @param forceDownload     whether to download segments from deep store when reloading.
    */
   public SegmentReloadMessage(String tableNameWithType, @Nullable List<String> segmentNames, boolean forceDownload) {
+    this(tableNameWithType, segmentNames, forceDownload, null);
+  }
+
+  /**
+   * This msg asks server to reload a list of specified segments with an explicit reload job id.
+   *
+   * @param tableNameWithType the table where the segments are from.
+   * @param segmentNames      a list of specified segments to reload, or null for all segments.
+   * @param forceDownload     whether to download segments from deep store when reloading.
+   * @param reloadJobId       reload job id to associate with the reload operation
+   */
+  public SegmentReloadMessage(String tableNameWithType, @Nullable List<String> segmentNames, boolean forceDownload,
+      @Nullable String reloadJobId) {
     super(MessageType.USER_DEFINE_MSG, UUID.randomUUID().toString());
     setResourceName(tableNameWithType);
     setMsgSubType(RELOAD_SEGMENT_MSG_SUB_TYPE);
@@ -66,6 +79,9 @@ public SegmentReloadMessage(String tableNameWithType, @Nullable List<String> seg
 
     // Persisting the msg ID as the reload job ID.
     znRecord.setSimpleField(RELOAD_JOB_ID_KEY, getMsgId());
+    if (reloadJobId != null) {
+      znRecord.setSimpleField(RELOAD_JOB_ID_KEY, reloadJobId);
+    }
 
     znRecord.setBooleanField(FORCE_DOWNLOAD_KEY, forceDownload);
     if (CollectionUtils.isNotEmpty(segmentNames)) {

pinot-common/src/main/java/org/apache/pinot/common/restlet/resources/PinotControllerJobMetadataDto.java

@@ -41,6 +41,7 @@ public class PinotControllerJobMetadataDto {
   private int _messageCount;
   private String _segmentName;
   private String _instanceName;
+  private String _instanceToSegmentsMap;
 
   public String getJobId() {
     return _jobId;
@@ -106,4 +107,14 @@ public PinotControllerJobMetadataDto setInstanceName(@Nullable String instanceNa
     _instanceName = instanceName;
     return this;
   }
+
+  @Nullable
+  public String getInstanceToSegmentsMap() {
+    return _instanceToSegmentsMap;
+  }
+
+  public PinotControllerJobMetadataDto setInstanceToSegmentsMap(@Nullable String instanceToSegmentsMap) {
+    _instanceToSegmentsMap = instanceToSegmentsMap;
+    return this;
+  }
 }

pinot-common/src/test/java/org/apache/pinot/common/messages/SegmentReloadMessageTest.java

@@ -0,0 +1,57 @@
+/**
+ * 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.messages;
+
+import java.util.List;
+import org.apache.helix.model.Message;
+import org.testng.annotations.Test;
+
+import static org.testng.Assert.assertEquals;
+import static org.testng.Assert.assertFalse;
+import static org.testng.Assert.assertThrows;
+import static org.testng.Assert.assertTrue;
+
+public class SegmentReloadMessageTest {
+
+  @Test
+  public void testReloadJobIdDefaultToMessageId() {
+    SegmentReloadMessage message = new SegmentReloadMessage("testTable_OFFLINE", false);
+
+    assertFalse(message.shouldForceDownload());
+    assertEquals(message.getReloadJobId(), message.getMsgId());
+  }
+
+  @Test
+  public void testReloadJobIdOverrideAndSegments() {
+    List<String> segments = List.of("seg_1", "seg_2");
+    SegmentReloadMessage message = new SegmentReloadMessage("testTable_OFFLINE", segments, true, "job-123");
+
+    assertTrue(message.shouldForceDownload());
+    assertEquals(message.getSegmentList(), segments);
+    assertEquals(message.getReloadJobId(), "job-123");
+  }
+
+  @Test
+  public void testInvalidSubtypeThrows() {
+    Message rawMessage = new Message(Message.MessageType.USER_DEFINE_MSG, "message-id");
+    rawMessage.setMsgSubType("INVALID_SUBTYPE");
+
+    assertThrows(IllegalArgumentException.class, () -> new SegmentReloadMessage(rawMessage));
+  }
+}

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

@@ -62,7 +62,7 @@
  *     POST requests:
  *     <ul>
  *       <li>"/segments/{tableName}/{segmentName}/reload": reload a specific segment</li>
- *       <li>"/segments/{tableName}/reload": reload all segments in a table</li>
+ *       <li>"/segments/{tableName}/reload": reload all segments in a table (optionally within a time range)</li>
  *     </ul>
  *   </li>
  *   <li>
@@ -128,7 +128,11 @@ public SuccessResponse reloadSegment(
   @Authenticate(AccessType.UPDATE)
   @Produces(MediaType.APPLICATION_JSON)
   @ApiOperation(value = "Reload all segments in a table",
-      notes = "Reloads all segments for the specified table. Supports filtering by type, instance, or custom mapping.")
+      notes = "Reloads all segments for the specified table. Supports filtering by type, instance, "
+          + "custom mapping, or time range. Time range params are in milliseconds and the range is "
+          + "[startTimestamp, endTimestamp). Either timestamp can be omitted to make that side of the "
+          + "range unbounded. When using time range parameters, do not specify targetInstance or "
+          + "instanceToSegmentsMap; these options are mutually exclusive with time range filters.")
   @ApiResponses(value = {
       @ApiResponse(code = 200, message = "Reload jobs submitted successfully"),
       @ApiResponse(code = 404, message = "No segments found")
@@ -140,13 +144,19 @@ public SuccessResponse reloadAllSegments(
       String tableTypeStr,
       @ApiParam(value = "Force server to re-download segments from deep store", defaultValue = "false")
       @QueryParam("forceDownload") @DefaultValue("false") boolean forceDownload,
+      @ApiParam(value = "Start timestamp (inclusive) in milliseconds. Omit for an unbounded lower bound.")
+      @QueryParam("startTimestamp") String startTimestampStr,
+      @ApiParam(value = "End timestamp (exclusive) in milliseconds. Omit for an unbounded upper bound.")
+      @QueryParam("endTimestamp") String endTimestampStr,
+      @ApiParam(value = "Whether to exclude segments overlapping the time range", defaultValue = "false")
+      @QueryParam("excludeOverlapping") @DefaultValue("false") boolean excludeOverlapping,
       @ApiParam(value = "Target specific server instance") @QueryParam("targetInstance") @Nullable
       String targetInstance,
       @ApiParam(value = "JSON map of instance to segment lists (overrides targetInstance)")
       @QueryParam("instanceToSegmentsMap") @Nullable String instanceToSegmentsMapInJson, @Context HttpHeaders headers)
       throws IOException {
     return _service.reloadAllSegments(tableName, tableTypeStr, forceDownload, targetInstance,
-        instanceToSegmentsMapInJson, headers);
+        instanceToSegmentsMapInJson, startTimestampStr, endTimestampStr, excludeOverlapping, headers);
   }
 
   @GET

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

@@ -2608,6 +2608,23 @@ public Map<String, Map<String, String>> getAllJobs(Set<ControllerJobType> jobTyp
    */
   public boolean addNewReloadSegmentJob(String tableNameWithType, String segmentNames, @Nullable String instanceName,
       String jobId, long jobSubmissionTimeMs, int numMessagesSent) {
+    return addNewReloadSegmentJob(tableNameWithType, segmentNames, instanceName, jobId, jobSubmissionTimeMs,
+        numMessagesSent, null);
+  }
+
+  /**
+   * Adds a new reload segment job metadata into ZK
+   * @param tableNameWithType Table for which job is to be added
+   * @param segmentNames Name of the segments being reloaded, separated by comma
+   * @param instanceName Name of the instance doing the segment reloading, optional.
+   * @param jobId job's UUID
+   * @param jobSubmissionTimeMs time at which the job was submitted
+   * @param numMessagesSent number of messages that were sent to servers. Saved as metadata
+   * @param instanceToSegmentsMapJson exact instance-to-segments mapping targeted by the job, optional.
+   * @return boolean representing success / failure of the ZK write step
+   */
+  public boolean addNewReloadSegmentJob(String tableNameWithType, String segmentNames, @Nullable String instanceName,
+      String jobId, long jobSubmissionTimeMs, int numMessagesSent, @Nullable String instanceToSegmentsMapJson) {
     Map<String, String> jobMetadata = new HashMap<>();
     jobMetadata.put(CommonConstants.ControllerJob.JOB_ID, jobId);
     jobMetadata.put(CommonConstants.ControllerJob.TABLE_NAME_WITH_TYPE, tableNameWithType);
@@ -2618,6 +2635,10 @@ public boolean addNewReloadSegmentJob(String tableNameWithType, String segmentNa
     if (instanceName != null) {
       jobMetadata.put(CommonConstants.ControllerJob.SEGMENT_RELOAD_JOB_INSTANCE_NAME, instanceName);
     }
+    if (instanceToSegmentsMapJson != null) {
+      jobMetadata.put(CommonConstants.ControllerJob.SEGMENT_RELOAD_JOB_INSTANCE_TO_SEGMENTS_MAP,
+          instanceToSegmentsMapJson);
+    }
     return addControllerJobToZK(jobId, jobMetadata, ControllerJobTypes.RELOAD_SEGMENT);
   }
 
@@ -2976,16 +2997,17 @@ public void refreshSegment(String tableNameWithType, SegmentMetadata segmentMeta
     sendSegmentRefreshMessage(tableNameWithType, segmentName, true, true);
   }
 
-  public Map<String, Pair<Integer, String>> reloadSegments(String tableNameWithType, boolean forceDownload,
-      Map<String, List<String>> instanceToSegmentsMap) {
+  public Map<String, Integer> reloadSegments(String tableNameWithType, boolean forceDownload,
+      Map<String, List<String>> instanceToSegmentsMap, String reloadJobId) {
     LOGGER.info("Sending reload messages for table: {} with forceDownload: {}, and instanceToSegmentsMap: {}",
         tableNameWithType, forceDownload, instanceToSegmentsMap);
 
     ClusterMessagingService messagingService = _helixZkManager.getMessagingService();
-    Map<String, Pair<Integer, String>> instanceMsgInfoMap = new HashMap<>();
+    Map<String, Integer> instanceMsgInfoMap = new HashMap<>();
     for (Map.Entry<String, List<String>> entry : instanceToSegmentsMap.entrySet()) {
       String targetInstance = entry.getKey();
-      SegmentReloadMessage message = new SegmentReloadMessage(tableNameWithType, entry.getValue(), forceDownload);
+      SegmentReloadMessage message =
+          new SegmentReloadMessage(tableNameWithType, entry.getValue(), forceDownload, reloadJobId);
       int numMessagesSent =
           MessagingServiceUtils.send(messagingService, message, tableNameWithType, null, targetInstance);
       if (numMessagesSent > 0) {
@@ -2994,7 +3016,7 @@ public Map<String, Pair<Integer, String>> reloadSegments(String tableNameWithTyp
       } else {
         LOGGER.warn("No reload message sent to instance: {} for table: {}", targetInstance, tableNameWithType);
       }
-      instanceMsgInfoMap.put(targetInstance, Pair.of(numMessagesSent, message.getMsgId()));
+      instanceMsgInfoMap.put(targetInstance, numMessagesSent);
     }
     return instanceMsgInfoMap;
   }
@@ -3014,7 +3036,7 @@ public Pair<Integer, String> reloadAllSegments(String tableNameWithType, boolean
       LOGGER.warn("No reload message sent for table: {}", tableNameWithType);
     }
 
-    return Pair.of(numMessagesSent, message.getMsgId());
+    return Pair.of(numMessagesSent, message.getReloadJobId());
   }
 
   public Pair<Integer, String> reloadSegment(String tableNameWithType, String segmentName, boolean forceDownload,
@@ -3033,7 +3055,7 @@ public Pair<Integer, String> reloadSegment(String tableNameWithType, String segm
       LOGGER.warn("No reload message sent for segment: {} in table: {}", segmentName, tableNameWithType);
     }
 
-    return Pair.of(numMessagesSent, message.getMsgId());
+    return Pair.of(numMessagesSent, message.getReloadJobId());
   }
 
   /**