Add versions to markUsed and markUnused APIs

docs/api-reference/data-management-api.md

@@ -206,20 +206,22 @@ Marks the state of a group of segments as unused, using an array of segment IDs
 Pass the array of segment IDs or interval as a JSON object in the request body.
 
 For the interval, specify the start and end times as ISO 8601 strings to identify segments inclusive of the start time and exclusive of the end time.
-Druid only updates the segments completely contained within the specified interval; partially overlapping segments are not affected.
+Optionally, specify an array of segment versions with interval. Druid updates only the segments completely contained
+within the specified interval that match the optional list of versions; partially overlapping segments are not affected.
 
 #### URL
 
 <code class="postAPI">POST</code> <code>/druid/coordinator/v1/datasources/:datasource/markUnused</code>
 
 #### Request body
 
-The group of segments is sent as a JSON request payload that accepts one of the following properties:
+The group of segments is sent as a JSON request payload that accepts the following properties:
 
-|Property|Description|Example|
-|----------|-------------|---------|
-|`interval`|ISO 8601 segments interval.|`"2015-09-12T03:00:00.000Z/2015-09-12T05:00:00.000Z"`|
-|`segmentIds`|Array of segment IDs.|`["segmentId1", "segmentId2"]`|
+|Property|Description|Required|Example|
+|----------|-------------|---------|---------|
+|`interval`|ISO 8601 segments interval.|Yes, if `segmentIds` is not specified.|`"2015-09-12T03:00:00.000Z/2015-09-12T05:00:00.000Z"`|
+|`segmentIds`|List of segment IDs.|Yes, if `interval` is not specified.|`["segmentId1", "segmentId2"]`|
+|`versions`|List of segment versions. Must be provided with `interval`.|No.|`["2024-03-14T16:00:04.086Z", ""2024-03-12T16:00:04.086Z"]`|
 
 #### Responses
 
@@ -306,20 +308,22 @@ Marks the state of a group of segments as used, using an array of segment IDs or
 Pass the array of segment IDs or interval as a JSON object in the request body.
 
 For the interval, specify the start and end times as ISO 8601 strings to identify segments inclusive of the start time and exclusive of the end time.
-Druid only updates the segments completely contained within the specified interval; partially overlapping segments are not affected.
+Optionally, specify an array of segment versions with interval. Druid updates only the segments completely contained
+within the specified interval that match the optional list of versions; partially overlapping segments are not affected.
 
 #### URL
 
 <code class="postAPI">POST</code> <code>/druid/coordinator/v1/datasources/:datasource/markUsed</code>
 
 #### Request body
 
-The group of segments is sent as a JSON request payload that accepts one of the following properties:
+The group of segments is sent as a JSON request payload that accepts the following properties:
 
-|Property|Description|Example|
-|----------|-------------|---------|
-|`interval`| ISO 8601 segments interval.|`"2015-09-12T03:00:00.000Z/2015-09-12T05:00:00.000Z"`|
-|`segmentIds`|Array of segment IDs.|`["segmentId1", "segmentId2"]`|
+|Property|Description|Required|Example|
+|----------|-------------|---------|---------|
+|`interval`|ISO 8601 segments interval.|Yes, if `segmentIds` is not specified.|`"2015-09-12T03:00:00.000Z/2015-09-12T05:00:00.000Z"`|
+|`segmentIds`|List of segment IDs.|Yes, if `interval` is not specified.|`["segmentId1", "segmentId2"]`|
+|`versions`|List of segment versions. Must be provided with `interval`.|No.|`["2024-03-14T16:00:04.086Z", ""2024-03-12T16:00:04.086Z"]`|
 
 #### Responses
 

indexing-service/src/main/java/org/apache/druid/indexing/common/actions/MarkSegmentsAsUnusedAction.java

@@ -63,9 +63,8 @@ public TypeReference<Integer> getReturnTypeReference()
   @Override
   public Integer perform(Task task, TaskActionToolbox toolbox)
   {
-    int numMarked = toolbox.getIndexerMetadataStorageCoordinator()
-                           .markSegmentsAsUnusedWithinInterval(dataSource, interval);
-    return numMarked;
+    return toolbox.getIndexerMetadataStorageCoordinator()
+                  .markSegmentsAsUnusedWithinInterval(dataSource, interval);
   }
 
   @Override

indexing-service/src/test/java/org/apache/druid/indexing/common/task/KillUnusedSegmentsTaskTest.java

@@ -387,7 +387,8 @@ public void testKillBatchSizeOneAndLimit4() throws Exception
         segments.size(),
         getSegmentsMetadataManager().markAsUnusedSegmentsInInterval(
             DATA_SOURCE,
-            Intervals.of("2018-01-01/2020-01-01")
+            Intervals.of("2018-01-01/2020-01-01"),
+            null
         )
     );
 
@@ -434,23 +435,26 @@ public void testKillMultipleUnusedSegmentsWithNullMaxUsedStatusLastUpdatedTime()
         1,
         getSegmentsMetadataManager().markAsUnusedSegmentsInInterval(
             DATA_SOURCE,
-            segment1.getInterval()
+            segment1.getInterval(),
+            null
         )
     );
 
     Assert.assertEquals(
         1,
         getSegmentsMetadataManager().markAsUnusedSegmentsInInterval(
             DATA_SOURCE,
-            segment4.getInterval()
+            segment4.getInterval(),
+            null
         )
     );
 
     Assert.assertEquals(
         1,
         getSegmentsMetadataManager().markAsUnusedSegmentsInInterval(
             DATA_SOURCE,
-            segment3.getInterval()
+            segment3.getInterval(),
+            null
         )
     );
 
@@ -508,15 +512,17 @@ public void testKillMultipleUnusedSegmentsWithDifferentMaxUsedStatusLastUpdatedT
         1,
         getSegmentsMetadataManager().markAsUnusedSegmentsInInterval(
             DATA_SOURCE,
-            segment1.getInterval()
+            segment1.getInterval(),
+            null
         )
     );
 
     Assert.assertEquals(
         1,
         getSegmentsMetadataManager().markAsUnusedSegmentsInInterval(
             DATA_SOURCE,
-            segment4.getInterval()
+            segment4.getInterval(),
+            null
         )
     );
 
@@ -529,7 +535,8 @@ public void testKillMultipleUnusedSegmentsWithDifferentMaxUsedStatusLastUpdatedT
         1,
         getSegmentsMetadataManager().markAsUnusedSegmentsInInterval(
             DATA_SOURCE,
-            segment3.getInterval()
+            segment3.getInterval(),
+            null
         )
     );
 

server/src/main/java/org/apache/druid/metadata/SegmentsMetadataManager.java

@@ -53,7 +53,13 @@ public interface SegmentsMetadataManager
    */
   int markAsUsedAllNonOvershadowedSegmentsInDataSource(String dataSource);
 
-  int markAsUsedNonOvershadowedSegmentsInInterval(String dataSource, Interval interval);
+  /**
+   * Marks non-overshadowed unused segments for the given interval and optional list of versions
+   * as used. If versions are not specified, all versions of non-overshadowed unused segments in the interval
+   * will be marked as used.
+   * @return Number of segments updated
+   */
+  int markAsUsedNonOvershadowedSegmentsInInterval(String dataSource, Interval interval, @Nullable List<String> versions);
 
   /**
    * Marks the given segment IDs as "used" only if there are not already overshadowed
@@ -81,7 +87,13 @@ public interface SegmentsMetadataManager
    */
   int markAsUnusedAllSegmentsInDataSource(String dataSource);
 
-  int markAsUnusedSegmentsInInterval(String dataSource, Interval interval);
+  /**
+   * Marks segments as unused that are <b>fully contained</b> in the given interval for an optional list of versions.
+   * If versions are not specified, all versions of segments in the interval will be marked as unused.
+   * Segments that are already marked as unused are not updated.
+   * @return The number of segments updated
+   */
+  int markAsUnusedSegmentsInInterval(String dataSource, Interval interval, @Nullable List<String> versions);
 
   int markSegmentsAsUnused(Set<SegmentId> segmentIds);
 

server/src/main/java/org/apache/druid/metadata/SqlSegmentsMetadataManager.java

@@ -654,21 +654,25 @@ public boolean markSegmentAsUsed(final String segmentId)
   @Override
   public int markAsUsedAllNonOvershadowedSegmentsInDataSource(final String dataSource)
   {
-    return doMarkAsUsedNonOvershadowedSegments(dataSource, null);
+    return doMarkAsUsedNonOvershadowedSegments(dataSource, null, null);
   }
 
   @Override
-  public int markAsUsedNonOvershadowedSegmentsInInterval(final String dataSource, final Interval interval)
+  public int markAsUsedNonOvershadowedSegmentsInInterval(
+      final String dataSource,
+      final Interval interval,
+      @Nullable final List<String> versions
+  )
   {
     Preconditions.checkNotNull(interval);
-    return doMarkAsUsedNonOvershadowedSegments(dataSource, interval);
+    return doMarkAsUsedNonOvershadowedSegments(dataSource, interval, versions);
   }
 
-  /**
-   * Implementation for both {@link #markAsUsedAllNonOvershadowedSegmentsInDataSource} (if the given interval is null)
-   * and {@link #markAsUsedNonOvershadowedSegmentsInInterval}.
-   */
-  private int doMarkAsUsedNonOvershadowedSegments(String dataSourceName, @Nullable Interval interval)
+  private int doMarkAsUsedNonOvershadowedSegments(
+      final String dataSourceName,
+      final @Nullable Interval interval,
+      final @Nullable List<String> versions
+  )
   {
     final List<DataSegment> unusedSegments = new ArrayList<>();
     final SegmentTimeline timeline = new SegmentTimeline();
@@ -682,12 +686,12 @@ private int doMarkAsUsedNonOvershadowedSegments(String dataSourceName, @Nullable
               interval == null ? Intervals.ONLY_ETERNITY : Collections.singletonList(interval);
 
           try (final CloseableIterator<DataSegment> iterator =
-                   queryTool.retrieveUsedSegments(dataSourceName, intervals)) {
+                   queryTool.retrieveUsedSegments(dataSourceName, intervals, versions)) {
             timeline.addSegments(iterator);
           }
 
           try (final CloseableIterator<DataSegment> iterator =
-                   queryTool.retrieveUnusedSegments(dataSourceName, intervals, null, null, null, null, null)) {
+                   queryTool.retrieveUnusedSegments(dataSourceName, intervals, versions, null, null, null, null)) {
             while (iterator.hasNext()) {
               final DataSegment dataSegment = iterator.next();
               timeline.addSegments(Iterators.singletonIterator(dataSegment));
@@ -796,7 +800,7 @@ private CloseableIterator<DataSegment> retrieveUsedSegmentsOverlappingIntervals(
   private int markSegmentsAsUsed(final List<SegmentId> segmentIds)
   {
     if (segmentIds.isEmpty()) {
-      log.info("No segments found to update!");
+      log.info("No segments found to mark as used.");
       return 0;
     }
 
@@ -856,13 +860,18 @@ public int markSegmentsAsUnused(Set<SegmentId> segmentIds)
   }
 
   @Override
-  public int markAsUnusedSegmentsInInterval(String dataSourceName, Interval interval)
+  public int markAsUnusedSegmentsInInterval(
+      final String dataSource,
+      final Interval interval,
+      @Nullable final List<String> versions
+  )
   {
+    Preconditions.checkNotNull(interval);
     try {
       return connector.getDBI().withHandle(
           handle ->
               SqlSegmentsMetadataQuery.forHandle(handle, connector, dbTables.get(), jsonMapper)
-                                      .markSegmentsUnused(dataSourceName, interval)
+                                      .markSegmentsUnused(dataSource, interval, versions)
       );
     }
     catch (Exception e) {