Add DatasourceCompactibleSegmentIterator

benchmarks/src/test/java/org/apache/druid/server/coordinator/NewestSegmentFirstPolicyBenchmark.java

@@ -141,7 +141,7 @@ public void setup()
   @Benchmark
   public void measureNewestSegmentFirstPolicy(Blackhole blackhole)
   {
-    final CompactionSegmentIterator iterator = policy.reset(compactionConfigs, dataSources, Collections.emptyMap());
+    final CompactionSegmentIterator iterator = policy.createIterator(compactionConfigs, dataSources, Collections.emptyMap());
     for (int i = 0; i < numCompactionTaskSlots && iterator.hasNext(); i++) {
       blackhole.consume(iterator.next());
     }

extensions-contrib/kubernetes-overlord-extensions/src/main/java/org/apache/druid/k8s/overlord/KubernetesTaskRunner.java

@@ -443,11 +443,17 @@ public Collection<TaskRunnerWorkItem> getPendingTasks()
   @Override
   public TaskLocation getTaskLocation(String taskId)
   {
-    final KubernetesWorkItem workItem = tasks.get(taskId);
-    if (workItem == null) {
+    try {
+      final KubernetesWorkItem workItem = tasks.get(taskId);
+      if (workItem == null) {
+        return TaskLocation.unknown();
+      } else {
+        return workItem.getLocation();
+      }
+    }
+    catch (Exception e) {
+      log.warn("Unable to find location for task [%s]", taskId);
       return TaskLocation.unknown();
-    } else {
-      return workItem.getLocation();
     }
   }
 

extensions-contrib/kubernetes-overlord-extensions/src/test/java/org/apache/druid/k8s/overlord/KubernetesTaskRunnerTest.java

@@ -654,6 +654,24 @@ public TaskLocation getLocation()
     Assert.assertEquals(TaskLocation.create("host", 0, 1, false), taskLocation);
   }
 
+  @Test
+  public void test_getTaskLocation_throws()
+  {
+    KubernetesWorkItem workItem = new KubernetesWorkItem(task, null)
+    {
+      @Override
+      public TaskLocation getLocation()
+      {
+        throw new RuntimeException();
+      }
+    };
+
+    runner.tasks.put(task.getId(), workItem);
+
+    TaskLocation taskLocation = runner.getTaskLocation(task.getId());
+    Assert.assertEquals(TaskLocation.unknown(), taskLocation);
+  }
+
   @Test
   public void test_getTaskLocation_noTaskFound()
   {

extensions-core/multi-stage-query/src/test/java/org/apache/druid/msq/indexing/MSQCompactionRunnerTest.java

@@ -39,6 +39,7 @@
 import org.apache.druid.indexer.partitions.PartitionsSpec;
 import org.apache.druid.indexing.common.task.CompactionIntervalSpec;
 import org.apache.druid.indexing.common.task.CompactionTask;
+import org.apache.druid.indexing.common.task.TuningConfigBuilder;
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.java.util.common.granularity.Granularities;
@@ -397,39 +398,17 @@ private static CompactionTask.CompactionTuningConfig createTuningConfig(
       PartitionsSpec partitionsSpec
   )
   {
-    return new CompactionTask.CompactionTuningConfig(
-        null,
-        null, // null to compute maxRowsPerSegment automatically
-        null,
-        500000,
-        1000000L,
-        null,
-        null,
-        null,
-        null,
-        partitionsSpec,
-        indexSpec,
-        null,
-        null,
-        !(partitionsSpec instanceof DynamicPartitionsSpec),
-        false,
-        5000L,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null,
-        null
-    );
+    return TuningConfigBuilder
+        .forCompactionTask()
+        .withMaxRowsInMemory(500000)
+        .withMaxBytesInMemory(1000000L)
+        .withMaxTotalRows(Long.MAX_VALUE)
+        .withPartitionsSpec(partitionsSpec)
+        .withIndexSpec(indexSpec)
+        .withForceGuaranteedRollup(!(partitionsSpec instanceof DynamicPartitionsSpec))
+        .withReportParseExceptions(false)
+        .withPushTimeout(5000L)
+        .build();
   }
 
   private static IndexSpec createIndexSpec()

extensions-core/stats/src/test/java/org/apache/druid/query/aggregation/variance/sql/VarianceSqlAggregatorTest.java

@@ -712,8 +712,7 @@ public void testOverWindow()
             + "group by dim4, dim5, mod(m1, 3)")
         .queryContext(ImmutableMap.of(
             PlannerContext.CTX_ENABLE_WINDOW_FNS, true,
-            QueryContexts.ENABLE_DEBUG, true,
-            QueryContexts.WINDOWING_STRICT_VALIDATION, false
+            QueryContexts.ENABLE_DEBUG, true
         ))
         .expectedResults(ImmutableList.of(
             new Object[]{"a", "aa", 1.0D, 0.0D},

indexing-service/src/main/java/org/apache/druid/indexing/appenderator/ActionBasedPublishedSegmentRetriever.java

@@ -79,7 +79,7 @@ public Set<DataSegment> findPublishedSegments(Set<SegmentId> segmentIds) throws
     catch (Exception e) {
       log.warn(
           e,
-          "Could not retrieve published segment IDs[%s] using task action[segmentListById]."
+          "Could not retrieve published segment IDs[%s] using task action[retrieveSegmentsById]."
           + " Overlord maybe on an older version, retrying with action[segmentListUsed]."
           + " This task may fail to publish segments if there is a concurrent replace happening.",
           serializedSegmentIds

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

@@ -0,0 +1,90 @@
+/*
+ * 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.druid.indexing.common.actions;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.type.TypeReference;
+import org.apache.druid.indexing.common.task.Task;
+
+import java.util.Set;
+
+/**
+ * Task action to retrieve the segment IDs from which a given set of segments were upgraded.
+ */
+public class RetrieveUpgradedFromSegmentIdsAction implements TaskAction<UpgradedFromSegmentsResponse>
+{
+  private final String dataSource;
+  private final Set<String> segmentIds;
+
+  @JsonCreator
+  public RetrieveUpgradedFromSegmentIdsAction(
+      @JsonProperty("dataSource") String dataSource,
+      @JsonProperty("segmentIds") Set<String> segmentIds
+  )
+  {
+    this.dataSource = dataSource;
+    this.segmentIds = segmentIds;
+  }
+
+  @JsonProperty
+  public String getDataSource()
+  {
+    return dataSource;
+  }
+
+  @JsonProperty
+  public Set<String> getSegmentIds()
+  {
+    return segmentIds;
+  }
+
+  @Override
+  public TypeReference<UpgradedFromSegmentsResponse> getReturnTypeReference()
+  {
+    return new TypeReference<UpgradedFromSegmentsResponse>()
+    {
+    };
+  }
+
+  @Override
+  public UpgradedFromSegmentsResponse perform(Task task, TaskActionToolbox toolbox)
+  {
+    return new UpgradedFromSegmentsResponse(
+        toolbox.getIndexerMetadataStorageCoordinator()
+               .retrieveUpgradedFromSegmentIds(dataSource, segmentIds)
+    );
+  }
+
+  @Override
+  public boolean isAudited()
+  {
+    return false;
+  }
+
+  @Override
+  public String toString()
+  {
+    return getClass().getSimpleName() + "{" +
+           "dataSource='" + dataSource + '\'' +
+           ", segmentIds=" + segmentIds +
+           '}';
+  }
+}

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

@@ -0,0 +1,95 @@
+/*
+ * 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.druid.indexing.common.actions;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+import com.fasterxml.jackson.core.type.TypeReference;
+import org.apache.druid.indexing.common.task.Task;
+
+import java.util.Set;
+
+/**
+ * Task action to determine the set of all segments containing the same load spec given the parent id. <br/>
+ * Returns a map from a segment ID to a set containing:
+ * <ol>
+ * <li> all segment IDs that were upgraded from it AND are still present in the metadata store </li>
+ * <li> the segment ID itself if and only if it is still present in the metadata store </li>
+ * </ol>
+ */
+public class RetrieveUpgradedToSegmentIdsAction implements TaskAction<UpgradedToSegmentsResponse>
+{
+  private final String dataSource;
+  private final Set<String> segmentIds;
+
+  @JsonCreator
+  public RetrieveUpgradedToSegmentIdsAction(
+      @JsonProperty("dataSource") String dataSource,
+      @JsonProperty("segmentIds") Set<String> segmentIds
+  )
+  {
+    this.dataSource = dataSource;
+    this.segmentIds = segmentIds;
+  }
+
+  @JsonProperty
+  public String getDataSource()
+  {
+    return dataSource;
+  }
+
+  @JsonProperty
+  public Set<String> getSegmentIds()
+  {
+    return segmentIds;
+  }
+
+  @Override
+  public TypeReference<UpgradedToSegmentsResponse> getReturnTypeReference()
+  {
+    return new TypeReference<UpgradedToSegmentsResponse>()
+    {
+    };
+  }
+
+  @Override
+  public UpgradedToSegmentsResponse perform(Task task, TaskActionToolbox toolbox)
+  {
+    return new UpgradedToSegmentsResponse(
+        toolbox.getIndexerMetadataStorageCoordinator()
+               .retrieveUpgradedToSegmentIds(dataSource, segmentIds)
+    );
+  }
+
+  @Override
+  public boolean isAudited()
+  {
+    return false;
+  }
+
+  @Override
+  public String toString()
+  {
+    return getClass().getSimpleName() + "{" +
+           "dataSource='" + dataSource + '\'' +
+           ", segmentIds=" + segmentIds +
+           '}';
+  }
+}

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

@@ -39,6 +39,8 @@
     @JsonSubTypes.Type(name = "segmentTransactionalAppend", value = SegmentTransactionalAppendAction.class),
     @JsonSubTypes.Type(name = "segmentTransactionalReplace", value = SegmentTransactionalReplaceAction.class),
     @JsonSubTypes.Type(name = "retrieveSegmentsById", value = RetrieveSegmentsByIdAction.class),
+    @JsonSubTypes.Type(name = "retrieveUpgradedFromSegmentIds", value = RetrieveUpgradedFromSegmentIdsAction.class),
+    @JsonSubTypes.Type(name = "retrieveUpgradedToSegmentIds", value = RetrieveUpgradedToSegmentIdsAction.class),
     @JsonSubTypes.Type(name = "segmentListUsed", value = RetrieveUsedSegmentsAction.class),
     @JsonSubTypes.Type(name = "segmentListUnused", value = RetrieveUnusedSegmentsAction.class),
     @JsonSubTypes.Type(name = "markSegmentsAsUnused", value = MarkSegmentsAsUnusedAction.class),

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

@@ -0,0 +1,44 @@
+/*
+ * 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.druid.indexing.common.actions;
+
+import com.fasterxml.jackson.annotation.JsonCreator;
+import com.fasterxml.jackson.annotation.JsonProperty;
+
+import java.util.Map;
+
+public class UpgradedFromSegmentsResponse
+{
+  private final Map<String, String> upgradedFromSegmentIds;
+
+  @JsonCreator
+  public UpgradedFromSegmentsResponse(
+      @JsonProperty("upgradedFromSegmentIds") Map<String, String> upgradedFromSegmentIds
+  )
+  {
+    this.upgradedFromSegmentIds = upgradedFromSegmentIds;
+  }
+
+  @JsonProperty
+  public Map<String, String> getUpgradedFromSegmentIds()
+  {
+    return upgradedFromSegmentIds;
+  }
+}