Allow append to existing datasources when dynamic partitioning is used

core/src/main/java/org/apache/druid/segment/SegmentUtils.java

@@ -29,16 +29,20 @@
 import org.apache.druid.java.util.common.IOE;
 import org.apache.druid.java.util.common.StringUtils;
 import org.apache.druid.timeline.DataSegment;
+import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
 import java.io.File;
 import java.io.FileInputStream;
 import java.io.IOException;
 import java.io.InputStream;
 import java.nio.charset.StandardCharsets;
+import java.util.ArrayList;
 import java.util.Collection;
 import java.util.Collections;
+import java.util.HashMap;
 import java.util.List;
+import java.util.Map;
 
 /**
  * Utility methods useful for implementing deep storage extensions.
@@ -93,6 +97,15 @@ public static Object commaSeparatedIdentifiers(@Nullable final Collection<DataSe
     return Collections2.transform(segments, DataSegment::getId);
   }
 
+  public static Map<Interval, List<DataSegment>> groupSegmentsByInterval(Collection<DataSegment> segments)
+  {
+    final Map<Interval, List<DataSegment>> intervalToSegments = new HashMap<>();
+    segments.forEach(
+        segment -> intervalToSegments.computeIfAbsent(segment.getInterval(), k -> new ArrayList<>()).add(segment)
+    );
+    return intervalToSegments;
+  }
+
   private SegmentUtils()
   {
     // no instantiation

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

@@ -168,7 +168,7 @@ public DataSegment(
       List<String> dimensions,
       List<String> metrics,
       ShardSpec shardSpec,
-      CompactionState lastCompactionState,
+      @Nullable CompactionState lastCompactionState,
       Integer binaryVersion,
       long size
   )

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

@@ -97,10 +97,4 @@ default boolean possibleInDomain(Map<String, RangeSet<String>> domain)
   {
     throw new UnsupportedOperationException();
   }
-
-  @Override
-  default boolean isCompatible(Class<? extends ShardSpec> other)
-  {
-    throw new UnsupportedOperationException();
-  }
 }

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

@@ -94,10 +94,4 @@ default boolean possibleInDomain(Map<String, RangeSet<String>> domain)
   {
     throw new UnsupportedOperationException();
   }
-
-  @Override
-  default boolean isCompatible(Class<? extends ShardSpec> other)
-  {
-    throw new UnsupportedOperationException();
-  }
 }

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

@@ -68,30 +68,18 @@ public int getNumBuckets()
   }
 
   @Override
-  public ShardSpec complete(ObjectMapper objectMapper, @Nullable ShardSpec specOfPreviousMaxPartitionId)
+  public ShardSpec complete(ObjectMapper objectMapper, int partitionId, int numCorePartitions)
   {
-    // The shardSpec is created by the Overlord.
-    // For batch tasks, this code can be executed only with segment locking (forceTimeChunkLock = false).
-    // In this mode, you can have 2 or more tasks concurrently ingesting into the same time chunk of
-    // the same datasource. Since there is no restriction for those tasks in segment allocation, the
-    // allocated IDs for each task can interleave. As a result, the core partition set cannot be
-    // represented as a range. We always set 0 for the core partition set size if this is an initial segment.
     return new HashBasedNumberedShardSpec(
-        specOfPreviousMaxPartitionId == null ? 0 : specOfPreviousMaxPartitionId.getPartitionNum() + 1,
-        specOfPreviousMaxPartitionId == null ? 0 : specOfPreviousMaxPartitionId.getNumCorePartitions(),
+        partitionId,
+        numCorePartitions,
         bucketId,
         numBuckets,
         partitionDimensions,
         objectMapper
     );
   }
 
-  @Override
-  public ShardSpec complete(ObjectMapper objectMapper, int partitionId)
-  {
-    return new HashBasedNumberedShardSpec(partitionId, 0, bucketId, numBuckets, partitionDimensions, objectMapper);
-  }
-
   @Override
   public Class<? extends ShardSpec> getShardSpecClass()
   {

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

@@ -54,7 +54,7 @@ public class HashBasedNumberedShardSpec extends NumberedShardSpec
 
   @JsonCreator
   public HashBasedNumberedShardSpec(
-      @JsonProperty("partitionNum") int partitionNum,    // partitionId, hash bucketId
+      @JsonProperty("partitionNum") int partitionNum,    // partitionId
       @JsonProperty("partitions") int partitions,        // core partition set size
       @JsonProperty("bucketId") @Nullable Integer bucketId, // nullable for backward compatibility
       @JsonProperty("numBuckets") @Nullable Integer numBuckets, // nullable for backward compatibility
@@ -90,12 +90,6 @@ public List<String> getPartitionDimensions()
     return partitionDimensions;
   }
 
-  @Override
-  public boolean isCompatible(Class<? extends ShardSpec> other)
-  {
-    return other == HashBasedNumberedShardSpec.class;
-  }
-
   @Override
   public boolean isInChunk(long timestamp, InputRow inputRow)
   {

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

@@ -21,8 +21,6 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 
-import javax.annotation.Nullable;
-
 public class LinearPartialShardSpec implements PartialShardSpec
 {
   private static final LinearPartialShardSpec INSTANCE = new LinearPartialShardSpec();
@@ -37,16 +35,9 @@ private LinearPartialShardSpec()
   }
 
   @Override
-  public ShardSpec complete(ObjectMapper objectMapper, @Nullable ShardSpec specOfPreviousMaxPartitionId)
-  {
-    return new LinearShardSpec(
-        specOfPreviousMaxPartitionId == null ? 0 : specOfPreviousMaxPartitionId.getPartitionNum() + 1
-    );
-  }
-
-  @Override
-  public ShardSpec complete(ObjectMapper objectMapper, int partitionId)
+  public ShardSpec complete(ObjectMapper objectMapper, int partitionId, int numCorePartitions)
   {
+    // numCorePartitions is ignored
     return new LinearShardSpec(partitionId);
   }
 

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

@@ -73,12 +73,6 @@ public boolean possibleInDomain(Map<String, RangeSet<String>> domain)
     return true;
   }
 
-  @Override
-  public boolean isCompatible(Class<? extends ShardSpec> other)
-  {
-    return other == LinearShardSpec.class;
-  }
-
   @Override
   public <T> PartitionChunk<T> createChunk(T obj)
   {

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

@@ -89,12 +89,6 @@ public boolean possibleInDomain(Map<String, RangeSet<String>> domain)
     return true;
   }
 
-  @Override
-  public boolean isCompatible(Class<? extends ShardSpec> other)
-  {
-    return other == NoneShardSpec.class;
-  }
-
   @Override
   public boolean equals(Object obj)
   {

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

@@ -22,8 +22,7 @@
 import com.fasterxml.jackson.annotation.JsonCreator;
 import com.fasterxml.jackson.annotation.JsonProperty;
 import com.fasterxml.jackson.databind.ObjectMapper;
-
-import javax.annotation.Nullable;
+import com.google.common.annotations.VisibleForTesting;
 
 public class NumberedOverwritePartialShardSpec implements PartialShardSpec
 {
@@ -43,6 +42,12 @@ public NumberedOverwritePartialShardSpec(
     this.minorVersion = minorVersion;
   }
 
+  @VisibleForTesting
+  public NumberedOverwritePartialShardSpec(int startRootPartitionId, int endRootPartitionId, int minorVersion)
+  {
+    this(startRootPartitionId, endRootPartitionId, (short) minorVersion);
+  }
+
   @JsonProperty
   public int getStartRootPartitionId()
   {
@@ -62,29 +67,25 @@ public short getMinorVersion()
   }
 
   @Override
-  public ShardSpec complete(ObjectMapper objectMapper, @Nullable ShardSpec specOfPreviousMaxPartitionId)
+  public ShardSpec complete(ObjectMapper objectMapper, int partitionId, int numCorePartitions)
   {
-    // specOfPreviousMaxPartitionId is the max partitionId of the same shardSpec
-    // and could be null if all existing segments are first-generation segments.
     return new NumberedOverwriteShardSpec(
-        specOfPreviousMaxPartitionId == null
-        ? PartitionIds.NON_ROOT_GEN_START_PARTITION_ID
-        : specOfPreviousMaxPartitionId.getPartitionNum() + 1,
+        partitionId,
         startRootPartitionId,
         endRootPartitionId,
         minorVersion
     );
   }
 
   @Override
-  public ShardSpec complete(ObjectMapper objectMapper, int partitionId)
+  public Class<? extends ShardSpec> getShardSpecClass()
   {
-    return new NumberedOverwriteShardSpec(partitionId, startRootPartitionId, endRootPartitionId, minorVersion);
+    return NumberedOverwriteShardSpec.class;
   }
 
   @Override
-  public Class<? extends ShardSpec> getShardSpecClass()
+  public boolean useNonRootGenerationPartitionSpace()
   {
-    return NumberedOverwriteShardSpec.class;
+    return true;
   }
 }

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

@@ -205,12 +205,6 @@ public boolean possibleInDomain(Map<String, RangeSet<String>> domain)
     return true;
   }
 
-  @Override
-  public boolean isCompatible(Class<? extends ShardSpec> other)
-  {
-    return other == NumberedOverwriteShardSpec.class || other == NumberedShardSpec.class;
-  }
-
   @Override
   public boolean equals(Object o)
   {

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

@@ -21,8 +21,6 @@
 
 import com.fasterxml.jackson.databind.ObjectMapper;
 
-import javax.annotation.Nullable;
-
 public class NumberedPartialShardSpec implements PartialShardSpec
 {
   private static final NumberedPartialShardSpec INSTANCE = new NumberedPartialShardSpec();
@@ -37,27 +35,9 @@ private NumberedPartialShardSpec()
   }
 
   @Override
-  public ShardSpec complete(ObjectMapper objectMapper, @Nullable ShardSpec specOfPreviousMaxPartitionId)
-  {
-    if (specOfPreviousMaxPartitionId == null) {
-      // The shardSpec is created by the Overlord.
-      // - For streaming ingestion tasks, the core partition set is always 0.
-      // - For batch tasks, this code is executed only with segment locking (forceTimeChunkLock = false).
-      //   In this mode, you can have 2 or more tasks concurrently ingesting into the same time chunk of
-      //   the same datasource. Since there is no restriction for those tasks in segment allocation, the
-      //   allocated IDs for each task can interleave. As a result, the core partition set cannot be
-      //   represented as a range. We always set 0 for the core partition set size.
-      return new NumberedShardSpec(0, 0);
-    } else {
-      final NumberedShardSpec prevSpec = (NumberedShardSpec) specOfPreviousMaxPartitionId;
-      return new NumberedShardSpec(prevSpec.getPartitionNum() + 1, prevSpec.getNumCorePartitions());
-    }
-  }
-
-  @Override
-  public ShardSpec complete(ObjectMapper objectMapper, int partitionId)
+  public ShardSpec complete(ObjectMapper objectMapper, int partitionId, int numCorePartitions)
   {
-    return new NumberedShardSpec(partitionId, 0);
+    return new NumberedShardSpec(partitionId, numCorePartitions);
   }
 
   @Override

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

@@ -93,12 +93,6 @@ public boolean possibleInDomain(Map<String, RangeSet<String>> domain)
     return true;
   }
 
-  @Override
-  public boolean isCompatible(Class<? extends ShardSpec> other)
-  {
-    return other == NumberedShardSpec.class || other == NumberedOverwriteShardSpec.class;
-  }
-
   @Override
   @JsonProperty("partitions")
   public int getNumCorePartitions()

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

@@ -45,4 +45,16 @@ default OverwriteShardSpec withAtomicUpdateGroupSize(int atomicUpdateGroupSize)
   }
 
   OverwriteShardSpec withAtomicUpdateGroupSize(short atomicUpdateGroupSize);
+
+  /**
+   * Returns true if this shardSpec and the given {@link PartialShardSpec} share the same partition space.
+   * This shardSpec uses non-root-generation partition space and thus does not share the space with other shardSpecs.
+   *
+   * @see PartitionIds
+   */
+  @Override
+  default boolean sharePartitionSpace(PartialShardSpec partialShardSpec)
+  {
+    return partialShardSpec.useNonRootGenerationPartitionSpace();
+  }
 }

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

@@ -25,8 +25,6 @@
 import com.fasterxml.jackson.annotation.JsonTypeInfo;
 import com.fasterxml.jackson.databind.ObjectMapper;
 
-import javax.annotation.Nullable;
-
 /**
  * This interface is used in the segment allocation protocol when it is coordinated by the Overlord; when appending
  * segments to an existing datasource (either streaming ingestion or batch append) or any case when segment
@@ -45,21 +43,31 @@
 public interface PartialShardSpec
 {
   /**
-   * Creates a new ShardSpec based on {@code specOfPreviousMaxPartitionId}. If it's null, it assumes that this is the
-   * first call for the time chunk where the new segment is created.
-   * Note that {@code specOfPreviousMaxPartitionId} can also be null for {@link OverwriteShardSpec} if all segments
-   * in the timeChunk are first-generation segments.
-   */
-  ShardSpec complete(ObjectMapper objectMapper, @Nullable ShardSpec specOfPreviousMaxPartitionId);
-
-  /**
-   * Creates a new shardSpec having the given partitionId.
+   * Creates a new ShardSpec with given partitionId and numCorePartitions.
+   *
+   * @param objectMapper      jsonMapper used only for {@link HashBasedNumberedShardSpec}
+   * @param partitionId       partitionId of the shardSpec. must be carefully chosen to be unique in a time chunk
+   * @param numCorePartitions the core partition set size. Should be set properly to determine if this segment belongs
+   *                          to the core partitions.
    */
-  ShardSpec complete(ObjectMapper objectMapper, int partitionId);
+  ShardSpec complete(ObjectMapper objectMapper, int partitionId, int numCorePartitions);
 
   /**
    * Returns the class of the shardSpec created by this factory.
    */
   @JsonIgnore
   Class<? extends ShardSpec> getShardSpecClass();
+
+  /**
+   * Returns true if this partialShardSpec needs a partitionId of a non-root generation.
+   * Any partialShardSpec to overwrite a subset of segments in a time chunk such as
+   * {@link NumberedOverwritePartialShardSpec} should return true.
+   *
+   *
+   * @see PartitionIds
+   */
+  default boolean useNonRootGenerationPartitionSpace()
+  {
+    return false;
+  }
 }

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

@@ -125,8 +125,14 @@ default short getAtomicUpdateGroupSize()
   boolean possibleInDomain(Map<String, RangeSet<String>> domain);
 
   /**
-   * Returns true if two segments of this and other shardSpecs can exist in the same time chunk.
+   * Returns true if this shardSpec and the given {@link PartialShardSpec} share the same partition space.
+   * All shardSpecs except {@link OverwriteShardSpec} use the root-generation partition space and thus share the same
+   * space.
+   *
+   * @see PartitionIds
    */
-  @JsonIgnore
-  boolean isCompatible(Class<? extends ShardSpec> other);
+  default boolean sharePartitionSpace(PartialShardSpec partialShardSpec)
+  {
+    return !partialShardSpec.useNonRootGenerationPartitionSpace();
+  }
 }