Support segmentGranularity for auto-compaction

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

@@ -99,6 +99,7 @@ public void setup()
               null,
               null,
               null,
+              null,
               null
           )
       );

core/src/main/java/org/apache/druid/timeline/CompactionState.java

@@ -43,15 +43,20 @@ public class CompactionState
   // org.apache.druid.segment.IndexSpec cannot be used here because it's in the 'processing' module which
   // has a dependency on the 'core' module where this class is.
   private final Map<String, Object> indexSpec;
+  // org.apache.druid.segment.indexing.granularity.GranularitySpec cannot be used here because it's in the
+  // 'server' module which has a dependency on the 'core' module where this class is.
+  private final Map<String, Object> granularitySpec;
 
   @JsonCreator
   public CompactionState(
       @JsonProperty("partitionsSpec") PartitionsSpec partitionsSpec,
-      @JsonProperty("indexSpec") Map<String, Object> indexSpec
+      @JsonProperty("indexSpec") Map<String, Object> indexSpec,
+      @JsonProperty("granularitySpec") Map<String, Object> granularitySpec
   )
   {
     this.partitionsSpec = partitionsSpec;
     this.indexSpec = indexSpec;
+    this.granularitySpec = granularitySpec;
   }
 
   @JsonProperty
@@ -66,6 +71,12 @@ public Map<String, Object> getIndexSpec()
     return indexSpec;
   }
 
+  @JsonProperty
+  public Map<String, Object> getGranularitySpec()
+  {
+    return granularitySpec;
+  }
+
   @Override
   public boolean equals(Object o)
   {
@@ -77,13 +88,14 @@ public boolean equals(Object o)
     }
     CompactionState that = (CompactionState) o;
     return Objects.equals(partitionsSpec, that.partitionsSpec) &&
-           Objects.equals(indexSpec, that.indexSpec);
+           Objects.equals(indexSpec, that.indexSpec) &&
+           Objects.equals(granularitySpec, that.granularitySpec);
   }
 
   @Override
   public int hashCode()
   {
-    return Objects.hash(partitionsSpec, indexSpec);
+    return Objects.hash(partitionsSpec, indexSpec, granularitySpec);
   }
 
   @Override
@@ -92,6 +104,7 @@ public String toString()
     return "CompactionState{" +
            "partitionsSpec=" + partitionsSpec +
            ", indexSpec=" + indexSpec +
+           ", granularitySpec=" + granularitySpec +
            '}';
   }
 }

core/src/main/java/org/apache/druid/timeline/partition/PartitionChunk.java

@@ -58,7 +58,7 @@ public interface PartitionChunk<T> extends Comparable<PartitionChunk<T>>
    * Returns true if this chunk is the end of the partition.  Most commonly, that means it represents the range
    * [X, infinity] for some concrete X.
    *
-   * @return true if the chunk is the beginning of the partition
+   * @return true if the chunk is the end of the partition
    */
   boolean isEnd();
 

core/src/test/java/org/apache/druid/timeline/DataSegmentTest.java

@@ -120,6 +120,7 @@ public void testV1Serialization() throws Exception
         new NumberedShardSpec(3, 0),
         new CompactionState(
             new HashedPartitionsSpec(100000, null, ImmutableList.of("dim1")),
+            ImmutableMap.of(),
             ImmutableMap.of()
         ),
         TEST_VERSION,
@@ -231,7 +232,8 @@ public void testWithLastCompactionState()
   {
     final CompactionState compactionState = new CompactionState(
         new DynamicPartitionsSpec(null, null),
-        Collections.singletonMap("test", "map")
+        Collections.singletonMap("test", "map"),
+        Collections.singletonMap("test2", "map2")
     );
     final DataSegment segment1 = DataSegment.builder()
                                             .dataSource("foo")

indexing-service/src/main/java/org/apache/druid/indexing/common/task/AbstractBatchIndexTask.java

@@ -473,12 +473,14 @@ public static boolean isGuaranteedRollup(IndexIOConfig ioConfig, IndexTuningConf
   public static Function<Set<DataSegment>, Set<DataSegment>> compactionStateAnnotateFunction(
       boolean storeCompactionState,
       TaskToolbox toolbox,
-      IndexTuningConfig tuningConfig
+      IndexTuningConfig tuningConfig,
+      GranularitySpec granularitySpec
   )
   {
     if (storeCompactionState) {
       final Map<String, Object> indexSpecMap = tuningConfig.getIndexSpec().asMap(toolbox.getJsonMapper());
-      final CompactionState compactionState = new CompactionState(tuningConfig.getPartitionsSpec(), indexSpecMap);
+      final Map<String, Object> granularitySpecMap = granularitySpec.asMap(toolbox.getJsonMapper());
+      final CompactionState compactionState = new CompactionState(tuningConfig.getPartitionsSpec(), indexSpecMap, granularitySpecMap);
       return segments -> segments
           .stream()
           .map(s -> s.withLastCompactionState(compactionState))

indexing-service/src/main/java/org/apache/druid/indexing/common/task/CompactionTask.java

@@ -141,6 +141,8 @@ public class CompactionTask extends AbstractBatchIndexTask
   @Nullable
   private final Granularity segmentGranularity;
   @Nullable
+  private final GranularitySpec granularitySpec;
+  @Nullable
   private final ParallelIndexTuningConfig tuningConfig;
   @JsonIgnore
   private final SegmentProvider segmentProvider;
@@ -172,7 +174,8 @@ public CompactionTask(
       @JsonProperty("dimensions") @Nullable final DimensionsSpec dimensions,
       @JsonProperty("dimensionsSpec") @Nullable final DimensionsSpec dimensionsSpec,
       @JsonProperty("metricsSpec") @Nullable final AggregatorFactory[] metricsSpec,
-      @JsonProperty("segmentGranularity") @Nullable final Granularity segmentGranularity,
+      @JsonProperty("segmentGranularity") @Deprecated @Nullable final Granularity segmentGranularity,
+      @JsonProperty("granularitySpec") @Nullable final GranularitySpec granularitySpec,
       @JsonProperty("tuningConfig") @Nullable final TuningConfig tuningConfig,
       @JsonProperty("context") @Nullable final Map<String, Object> context,
       @JacksonInject SegmentLoaderFactory segmentLoaderFactory,
@@ -202,6 +205,16 @@ public CompactionTask(
     this.dimensionsSpec = dimensionsSpec == null ? dimensions : dimensionsSpec;
     this.metricsSpec = metricsSpec;
     this.segmentGranularity = segmentGranularity;
+    if (granularitySpec == null && segmentGranularity != null) {
+      this.granularitySpec = new UniformGranularitySpec(
+          segmentGranularity,
+          null,
+          null,
+          null
+      );
+    } else {
+      this.granularitySpec = granularitySpec;
+    }
     this.tuningConfig = tuningConfig != null ? getTuningConfig(tuningConfig) : null;
     this.segmentProvider = new SegmentProvider(dataSource, this.ioConfig.getInputSpec());
     this.partitionConfigurationManager = new PartitionConfigurationManager(this.tuningConfig);
@@ -288,7 +301,14 @@ public AggregatorFactory[] getMetricsSpec()
   @Override
   public Granularity getSegmentGranularity()
   {
-    return segmentGranularity;
+    return granularitySpec == null ? null : granularitySpec.getSegmentGranularity();
+  }
+
+  @JsonProperty
+  @Nullable
+  public GranularitySpec getGranularitySpec()
+  {
+    return granularitySpec;
   }
 
   @Nullable
@@ -348,7 +368,7 @@ public TaskStatus runTask(TaskToolbox toolbox) throws Exception
         partitionConfigurationManager,
         dimensionsSpec,
         metricsSpec,
-        segmentGranularity,
+        getSegmentGranularity(),
         toolbox.getCoordinatorClient(),
         segmentLoaderFactory,
         retryPolicyFactory
@@ -892,6 +912,8 @@ public static class Builder
     @Nullable
     private Granularity segmentGranularity;
     @Nullable
+    private GranularitySpec granularitySpec;
+    @Nullable
     private TuningConfig tuningConfig;
     @Nullable
     private Map<String, Object> context;
@@ -941,6 +963,12 @@ public Builder segmentGranularity(Granularity segmentGranularity)
       return this;
     }
 
+    public Builder granularitySpec(GranularitySpec granularitySpec)
+    {
+      this.granularitySpec = granularitySpec;
+      return this;
+    }
+
     public Builder tuningConfig(TuningConfig tuningConfig)
     {
       this.tuningConfig = tuningConfig;
@@ -966,6 +994,7 @@ public CompactionTask build()
           dimensionsSpec,
           metricsSpec,
           segmentGranularity,
+          granularitySpec,
           tuningConfig,
           context,
           segmentLoaderFactory,

indexing-service/src/main/java/org/apache/druid/indexing/common/task/IndexTask.java

@@ -904,7 +904,8 @@ private TaskStatus generateAndPublishSegments(
           compactionStateAnnotateFunction(
               storeCompactionState,
               toolbox,
-              ingestionSchema.getTuningConfig()
+              ingestionSchema.getTuningConfig(),
+              ingestionSchema.getDataSchema().getGranularitySpec()
           );
 
       // Probably we can publish atomicUpdateGroup along with segments.

indexing-service/src/main/java/org/apache/druid/indexing/common/task/batch/parallel/ParallelIndexSupervisorTask.java

@@ -929,7 +929,8 @@ private void publishSegments(
     final Function<Set<DataSegment>, Set<DataSegment>> annotateFunction = compactionStateAnnotateFunction(
         storeCompactionState,
         toolbox,
-        ingestionSchema.getTuningConfig()
+        ingestionSchema.getTuningConfig(),
+        ingestionSchema.getDataSchema().getGranularitySpec()
     );
     final TransactionalSegmentPublisher publisher = (segmentsToBeOverwritten, segmentsToPublish, commitMetadata) ->
         toolbox.getTaskActionClient().submit(

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

@@ -28,6 +28,7 @@
 import org.apache.druid.client.indexing.ClientCompactionIOConfig;
 import org.apache.druid.client.indexing.ClientCompactionIntervalSpec;
 import org.apache.druid.client.indexing.ClientCompactionTaskQuery;
+import org.apache.druid.client.indexing.ClientCompactionTaskQueryGranularitySpec;
 import org.apache.druid.client.indexing.ClientCompactionTaskQueryTuningConfig;
 import org.apache.druid.client.indexing.ClientTaskQuery;
 import org.apache.druid.client.indexing.IndexingServiceClient;
@@ -45,11 +46,13 @@
 import org.apache.druid.jackson.DefaultObjectMapper;
 import org.apache.druid.java.util.common.HumanReadableBytes;
 import org.apache.druid.java.util.common.Intervals;
+import org.apache.druid.java.util.common.granularity.Granularities;
 import org.apache.druid.segment.IndexSpec;
 import org.apache.druid.segment.data.BitmapSerde.DefaultBitmapSerdeFactory;
 import org.apache.druid.segment.data.CompressionFactory.LongEncodingStrategy;
 import org.apache.druid.segment.data.CompressionStrategy;
 import org.apache.druid.segment.incremental.RowIngestionMetersFactory;
+import org.apache.druid.segment.indexing.granularity.UniformGranularitySpec;
 import org.apache.druid.segment.realtime.appenderator.AppenderatorsManager;
 import org.apache.druid.segment.realtime.firehose.ChatHandlerProvider;
 import org.apache.druid.segment.realtime.firehose.NoopChatHandlerProvider;
@@ -113,6 +116,7 @@ public void testClientCompactionTaskQueryToCompactionTask() throws IOException
             1000,
             100
         ),
+        new ClientCompactionTaskQueryGranularitySpec(Granularities.DAY, Granularities.HOUR, true),
         ImmutableMap.of("key", "value")
     );
 
@@ -186,6 +190,18 @@ public void testClientCompactionTaskQueryToCompactionTask() throws IOException
         query.getTuningConfig().getTotalNumMergeTasks().intValue(),
         task.getTuningConfig().getTotalNumMergeTasks()
     );
+    Assert.assertEquals(
+        query.getGranularitySpec().getQueryGranularity(),
+        task.getGranularitySpec().getQueryGranularity()
+    );
+    Assert.assertEquals(
+        query.getGranularitySpec().getSegmentGranularity(),
+        task.getGranularitySpec().getSegmentGranularity()
+    );
+    Assert.assertEquals(
+        query.getGranularitySpec().isRollup(),
+        task.getGranularitySpec().isRollup()
+    );
     Assert.assertEquals(query.getContext(), task.getContext());
   }
 
@@ -243,6 +259,7 @@ public void testCompactionTaskToClientCompactionTaskQuery() throws IOException
                 null
             )
         )
+        .granularitySpec(new UniformGranularitySpec(Granularities.DAY, Granularities.HOUR, null))
         .build();
 
     final ClientCompactionTaskQuery expected = new ClientCompactionTaskQuery(
@@ -284,6 +301,7 @@ public void testCompactionTaskToClientCompactionTaskQuery() throws IOException
             1000,
             100
         ),
+        new ClientCompactionTaskQueryGranularitySpec(Granularities.DAY, Granularities.HOUR, true),
         new HashMap<>()
     );