perf: Speed up VersionedIntervalTimeline operations

Author: jtuglu1

benchmarks/src/test/java/org/apache/druid/timeline/VersionedIntervalTimelineBenchmark.java

@@ -26,6 +26,7 @@
 import org.apache.druid.java.util.common.granularity.GranularityType;
 import org.apache.druid.timeline.partition.NumberedOverwriteShardSpec;
 import org.apache.druid.timeline.partition.NumberedShardSpec;
+import org.apache.druid.timeline.partition.PartitionChunk;
 import org.apache.druid.timeline.partition.PartitionIds;
 import org.apache.druid.timeline.partition.ShardSpec;
 import org.joda.time.DateTime;
@@ -53,7 +54,7 @@
 import java.util.concurrent.ThreadLocalRandom;
 
 @State(Scope.Benchmark)
-@Fork(value = 1, jvmArgsAppend = {"-XX:+UseG1GC"})
+@Fork(value = 1, jvmArgsAppend = {"-XX:+UseG1GC", "-Xmx4g"})
 @Warmup(iterations = 10)
 @Measurement(iterations = 10)
 @BenchmarkMode({Mode.Throughput})
@@ -73,7 +74,7 @@ public class VersionedIntervalTimelineBenchmark
   @Param({"false", "true"})
   private boolean useSegmentLock;
 
-  @Param({"MONTH", "DAY"})
+  @Param({"MONTH", "DAY", "HOUR", "MINUTE"})
   private GranularityType segmentGranularity;
 
   private List<Interval> intervals;
@@ -216,16 +217,22 @@ public void benchAdd(Blackhole blackhole)
     }
   }
 
+  /**
+   * Measures the cost of draining 10% of the timeline in one shot by rebuilding the timeline fresh
+   * each invocation. This is the "bulk removal" case — use {@link Mode#AverageTime} so the result
+   * is interpretable as "seconds per full 10%-drain" rather than a near-zero throughput figure.
+   */
   @Benchmark
+  @BenchmarkMode(Mode.AverageTime)
   public void benchRemove(Blackhole blackhole)
   {
     final List<DataSegment> segmentsCopy = new ArrayList<>(segments);
-    final SegmentTimeline timeline = SegmentTimeline.forSegments(segmentsCopy);
+    final SegmentTimeline localTimeline = SegmentTimeline.forSegments(segmentsCopy);
     final int numTests = (int) (segmentsCopy.size() * 0.1);
     for (int i = 0; i < numTests; i++) {
       final DataSegment segment = segmentsCopy.remove(ThreadLocalRandom.current().nextInt(segmentsCopy.size()));
       blackhole.consume(
-          timeline.remove(
+          localTimeline.remove(
               segment.getInterval(),
               segment.getVersion(),
               segment.getShardSpec().createChunk(segment)
@@ -234,6 +241,27 @@ public void benchRemove(Blackhole blackhole)
     }
   }
 
+  /**
+   * Measures per-remove throughput against the pre-built timeline. One invocation = one remove + one
+   * re-add to keep the timeline stable.
+   */
+  @Benchmark
+  @BenchmarkMode(Mode.Throughput)
+  public void benchRemoveWithReplacement(Blackhole blackhole)
+  {
+    final DataSegment segment = segments.get(ThreadLocalRandom.current().nextInt(segments.size()));
+    final PartitionChunk<DataSegment> chunk = segment.getShardSpec().createChunk(segment);
+    final PartitionChunk<DataSegment> removed = timeline.remove(
+        segment.getInterval(),
+        segment.getVersion(),
+        chunk
+    );
+    if (removed != null) {
+      timeline.add(segment.getInterval(), segment.getVersion(), chunk);
+    }
+    blackhole.consume(removed);
+  }
+
   @Benchmark
   public void benchLookup(Blackhole blackhole)
   {
@@ -245,6 +273,26 @@ public void benchLookup(Blackhole blackhole)
     blackhole.consume(timeline.lookup(queryInterval));
   }
 
+  /**
+   * Looks up a single-interval window — the case that benefits most from the O(log N + K) range scan since only a
+   * tiny fraction of the timeline is touched regardless of overall size.
+   */
+  @Benchmark
+  public void benchLookupSingleInterval(Blackhole blackhole)
+  {
+    final int intervalIndex = ThreadLocalRandom.current().nextInt(intervals.size());
+    blackhole.consume(timeline.lookup(intervals.get(intervalIndex)));
+  }
+
+  /**
+   * Full-range lookup — scans all segments regardless of optimization; provides an upper-bound baseline.
+   */
+  @Benchmark
+  public void benchLookupFullRange(Blackhole blackhole)
+  {
+    blackhole.consume(timeline.lookup(TOTAL_INTERVAL));
+  }
+
   @Benchmark
   public void benchIsOvershadowed(Blackhole blackhole)
   {

processing/src/main/java/org/apache/druid/timeline/VersionedIntervalTimeline.java

@@ -88,7 +88,10 @@ public class VersionedIntervalTimeline<VersionType, ObjectType extends Overshado
       Comparators.intervalsByStartThenEnd()
   );
   // true interval -> version -> timelineEntry
-  private final Map<Interval, TreeMap<VersionType, TimelineEntry>> allTimelineEntries = new HashMap<>();
+  // NavigableMap so the private remove() can range-scan for overlapping intervals instead of a full O(N) pass.
+  private final NavigableMap<Interval, TreeMap<VersionType, TimelineEntry>> allTimelineEntries = new TreeMap<>(
+      Comparators.intervalsByStartThenEnd()
+  );
   private final AtomicInteger numObjects = new AtomicInteger();
 
   private final Comparator<? super VersionType> versionComparator;
@@ -118,7 +121,7 @@ public static <VersionType, ObjectType extends Overshadowable<ObjectType>> Itera
                .iterator();
   }
 
-  public Map<Interval, TreeMap<VersionType, TimelineEntry>> getAllTimelineEntries()
+  public NavigableMap<Interval, TreeMap<VersionType, TimelineEntry>> getAllTimelineEntries()
   {
     return allTimelineEntries;
   }
@@ -715,7 +718,15 @@ private void remove(
   {
     timeline.remove(interval);
 
-    for (Entry<Interval, TreeMap<VersionType, TimelineEntry>> versionEntry : allTimelineEntries.entrySet()) {
+    // allTimelineEntries stores TRUE intervals which CAN overlap each other (e.g. a wide interval
+    // [Jan01/Dec31] and a narrow one [Jun01/Jun30] may both start before the removed interval).
+    // We therefore cannot use floorKey — entries starting well before interval.start may still extend
+    // into the interval and must be re-added. Use headMap to skip entries whose start >= interval.end
+    // (those can never overlap), then filter for actual overlap. Reduces the scan from O(M) to O(M')
+    // where M' is the count of entries with start < interval.end.
+    final Interval headBound = new Interval(interval.getEnd(), interval.getEnd());
+    for (Entry<Interval, TreeMap<VersionType, TimelineEntry>> versionEntry :
+        allTimelineEntries.headMap(headBound, false).entrySet()) {
       if (versionEntry.getKey().overlap(interval) != null) {
         if (incompleteOk) {
           add(timeline, versionEntry.getKey(), versionEntry.getValue().lastEntry().getValue());
@@ -747,8 +758,27 @@ private List<TimelineObjectHolder<VersionType, ObjectType>> lookup(Interval inte
       timeline = completePartitionsTimeline;
     }
 
-    for (Entry<Interval, TimelineEntry> entry : timeline.entrySet()) {
+    // Both completePartitionsTimeline and incompletePartitionsTimeline contain non-overlapping adjusted intervals
+    // sorted by (start, end). To find entries overlapping [interval.start, interval.end) we only need to scan
+    // the range [floorKey(interval.start), first key with start >= interval.end). The floor handles the one
+    // entry that may have started before interval.start but still extends into the query window.
+    final Interval searchStart = new Interval(interval.getStart(), DateTimes.MAX);
+    Interval startKey = timeline.floorKey(searchStart);
+    if (startKey == null) {
+      // No entry starts at or before interval.start; begin from the first entry in the map.
+      startKey = timeline.isEmpty() ? null : timeline.firstKey();
+    }
+    if (startKey == null) {
+      return retVal;
+    }
+
+    for (Entry<Interval, TimelineEntry> entry : timeline.tailMap(startKey, true).entrySet()) {
       Interval timelineInterval = entry.getKey();
+      // All entries from here forward start at or after startKey.start. Once we reach an entry whose
+      // start is at or past the query end, there can be no more overlaps.
+      if (timelineInterval.getStartMillis() >= interval.getEndMillis()) {
+        break;
+      }
       TimelineEntry val = entry.getValue();
 
       // exclude empty partition holders (i.e. tombstones) since they do not add value