Revert "Avoid expensive findEntry call in segment metadata query (#10892)"

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

@@ -19,7 +19,7 @@
 
 package org.apache.druid.timeline;
 
-import org.apache.druid.timeline.partition.PartitionChunk;
+import org.apache.druid.timeline.partition.PartitionHolder;
 import org.joda.time.Interval;
 
 import javax.annotation.Nullable;
@@ -51,9 +51,5 @@ public interface TimelineLookup<VersionType, ObjectType extends Overshadowable<O
    */
   List<TimelineObjectHolder<VersionType, ObjectType>> lookupWithIncompletePartitions(Interval interval);
 
-  /**
-   * Finds the {@link PartitionChunk} for the given time interval, version and chunk number.
-   */
-  @Nullable
-  PartitionChunk<ObjectType> findChunk(Interval interval, VersionType version, int partitionNum);
+  @Nullable PartitionHolder<ObjectType> findEntry(Interval interval, VersionType version);
 }

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

@@ -20,6 +20,7 @@
 package org.apache.druid.timeline;
 
 import com.google.common.annotations.VisibleForTesting;
+import com.google.common.base.Function;
 import com.google.common.base.Preconditions;
 import com.google.common.collect.FluentIterable;
 import com.google.common.collect.Iterators;
@@ -116,14 +117,10 @@ public static void addSegments(
   )
   {
     timeline.addAll(
-        Iterators.transform(
-            segments,
-            segment -> new PartitionChunkEntry<>(
-                segment.getInterval(),
-                segment.getVersion(),
-                segment.getShardSpec().createChunk(segment)
-            )
-        ));
+        Iterators.transform(segments, segment -> segment.getShardSpec().createChunk(segment)),
+        DataSegment::getInterval,
+        DataSegment::getVersion
+    );
   }
 
   public Map<Interval, TreeMap<VersionType, TimelineEntry>> getAllTimelineEntries()
@@ -186,11 +183,13 @@ public Set<ObjectType> findNonOvershadowedObjectsInInterval(Interval interval, P
 
   public void add(final Interval interval, VersionType version, PartitionChunk<ObjectType> object)
   {
-    addAll(Iterators.singletonIterator(new PartitionChunkEntry<>(interval, version, object)));
+    addAll(Iterators.singletonIterator(object), o -> interval, o -> version);
   }
 
-  public void addAll(
-      final Iterator<PartitionChunkEntry<VersionType, ObjectType>> objects
+  private void addAll(
+      final Iterator<PartitionChunk<ObjectType>> objects,
+      final Function<ObjectType, Interval> intervalFunction,
+      final Function<ObjectType, VersionType> versionFunction
   )
   {
     lock.writeLock().lock();
@@ -199,10 +198,9 @@ public void addAll(
       final IdentityHashMap<TimelineEntry, Interval> allEntries = new IdentityHashMap<>();
 
       while (objects.hasNext()) {
-        PartitionChunkEntry<VersionType, ObjectType> chunkEntry = objects.next();
-        PartitionChunk<ObjectType> object = chunkEntry.getChunk();
-        Interval interval = chunkEntry.getInterval();
-        VersionType version = chunkEntry.getVersion();
+        PartitionChunk<ObjectType> object = objects.next();
+        Interval interval = intervalFunction.apply(object.getObject());
+        VersionType version = versionFunction.apply(object.getObject());
         Map<VersionType, TimelineEntry> exists = allTimelineEntries.get(interval);
         TimelineEntry entry;
 
@@ -286,15 +284,15 @@ public PartitionChunk<ObjectType> remove(Interval interval, VersionType version,
 
   @Override
   @Nullable
-  public PartitionChunk<ObjectType> findChunk(Interval interval, VersionType version, int partitionNum)
+  public PartitionHolder<ObjectType> findEntry(Interval interval, VersionType version)
   {
     lock.readLock().lock();
     try {
       for (Entry<Interval, TreeMap<VersionType, TimelineEntry>> entry : allTimelineEntries.entrySet()) {
         if (entry.getKey().equals(interval) || entry.getKey().contains(interval)) {
           TimelineEntry foundEntry = entry.getValue().get(version);
           if (foundEntry != null) {
-            return foundEntry.getPartitionHolder().getChunk(partitionNum);
+            return foundEntry.getPartitionHolder().asImmutable();
           }
         }
       }
@@ -851,41 +849,4 @@ public int hashCode()
       return Objects.hash(trueInterval, version, partitionHolder);
     }
   }
-
-  /**
-   * Stores a {@link PartitionChunk} for a given interval and version. The
-   * interval corresponds to the {@link LogicalSegment#getInterval()}
-   */
-  public static class PartitionChunkEntry<VersionType, ObjectType>
-  {
-    private final Interval interval;
-    private final VersionType version;
-    private final PartitionChunk<ObjectType> chunk;
-
-    public PartitionChunkEntry(
-        Interval interval,
-        VersionType version,
-        PartitionChunk<ObjectType> chunk
-    )
-    {
-      this.interval = interval;
-      this.version = version;
-      this.chunk = chunk;
-    }
-
-    public Interval getInterval()
-    {
-      return interval;
-    }
-
-    public VersionType getVersion()
-    {
-      return version;
-    }
-
-    public PartitionChunk<ObjectType> getChunk()
-    {
-      return chunk;
-    }
-  }
 }

core/src/main/java/org/apache/druid/timeline/partition/ImmutablePartitionHolder.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.timeline.partition;
+
+import org.apache.druid.timeline.Overshadowable;
+
+/**
+ */
+public class ImmutablePartitionHolder<T extends Overshadowable<T>> extends PartitionHolder<T>
+{
+  protected ImmutablePartitionHolder(OvershadowableManager<T> overshadowableManager)
+  {
+    super(overshadowableManager);
+  }
+
+  @Override
+  public PartitionChunk<T> remove(PartitionChunk<T> tPartitionChunk)
+  {
+    throw new UnsupportedOperationException();
+  }
+
+  @Override
+  public boolean add(PartitionChunk<T> tPartitionChunk)
+  {
+    throw new UnsupportedOperationException();
+  }
+}

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

@@ -65,6 +65,11 @@ protected PartitionHolder(OvershadowableManager<T> overshadowableManager)
     this.overshadowableManager = overshadowableManager;
   }
 
+  public ImmutablePartitionHolder<T> asImmutable()
+  {
+    return new ImmutablePartitionHolder<>(OvershadowableManager.copyVisible(overshadowableManager));
+  }
+
   public boolean add(PartitionChunk<T> chunk)
   {
     return overshadowableManager.addChunk(chunk);

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

@@ -21,7 +21,6 @@
 
 /**
  */
-@Deprecated
 public class SingleElementPartitionChunk<T> implements PartitionChunk<T>
 {
   private final T element;

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

@@ -25,6 +25,7 @@
 import org.apache.druid.java.util.common.Intervals;
 import org.apache.druid.timeline.partition.IntegerPartitionChunk;
 import org.apache.druid.timeline.partition.OvershadowableInteger;
+import org.apache.druid.timeline.partition.PartitionHolder;
 import org.joda.time.DateTime;
 import org.joda.time.Days;
 import org.joda.time.Hours;
@@ -220,64 +221,36 @@ public void testRemove()
   }
 
   @Test
-  public void testFindChunk()
+  public void testFindEntry()
   {
-    assertSingleElementChunks(
-        makeSingle("1", 1),
-        timeline.findChunk(Intervals.of("2011-10-01/2011-10-02"), "1", 0)
-    );
-
-    assertSingleElementChunks(
-        makeSingle("1", 1),
-        timeline.findChunk(Intervals.of("2011-10-01/2011-10-01T10"), "1", 0)
-    );
-
-    assertSingleElementChunks(
-        makeSingle("1", 1),
-        timeline.findChunk(Intervals.of("2011-10-01T02/2011-10-02"), "1", 0)
-    );
-
-    assertSingleElementChunks(
-        makeSingle("1", 1),
-        timeline.findChunk(Intervals.of("2011-10-01T04/2011-10-01T17"), "1", 0)
+    Assert.assertEquals(
+        new PartitionHolder<>(makeSingle("1", 1)).asImmutable(),
+        timeline.findEntry(Intervals.of("2011-10-01/2011-10-02"), "1")
     );
 
-    IntegerPartitionChunk<OvershadowableInteger> expected = IntegerPartitionChunk.make(
-        10,
-        null,
-        1,
-        new OvershadowableInteger(
-            "3",
-            1,
-            21
-        )
-    );
-    IntegerPartitionChunk<OvershadowableInteger> actual = (IntegerPartitionChunk<OvershadowableInteger>) timeline.findChunk(
-        Intervals.of("2011-10-02/2011-10-03"),
-        "3",
-        1
+    Assert.assertEquals(
+        new PartitionHolder<>(makeSingle("1", 1)).asImmutable(),
+        timeline.findEntry(Intervals.of("2011-10-01/2011-10-01T10"), "1")
     );
-    Assert.assertEquals(expected, actual);
-    Assert.assertEquals(expected.getObject(), actual.getObject());
 
     Assert.assertEquals(
-        null,
-        timeline.findChunk(Intervals.of("2011-10-01T04/2011-10-01T17"), "1", 1)
+        new PartitionHolder<>(makeSingle("1", 1)).asImmutable(),
+        timeline.findEntry(Intervals.of("2011-10-01T02/2011-10-02"), "1")
     );
 
     Assert.assertEquals(
-        null,
-        timeline.findChunk(Intervals.of("2011-10-01T04/2011-10-01T17"), "2", 0)
+        new PartitionHolder<>(makeSingle("1", 1)).asImmutable(),
+        timeline.findEntry(Intervals.of("2011-10-01T04/2011-10-01T17"), "1")
     );
 
     Assert.assertEquals(
         null,
-        timeline.findChunk(Intervals.of("2011-10-01T04/2011-10-02T17"), "1", 0)
+        timeline.findEntry(Intervals.of("2011-10-01T04/2011-10-01T17"), "2")
     );
 
     Assert.assertEquals(
         null,
-        timeline.findChunk(Intervals.of("2011-10-01T04/2011-10-02T17"), "1", 0)
+        timeline.findEntry(Intervals.of("2011-10-01T04/2011-10-02T17"), "1")
     );
   }
 

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

@@ -51,14 +51,14 @@ public void setUp()
   }
 
   @Test
-  public void testFindChunkWithOverlap()
+  public void testFindEntryWithOverlap()
   {
     add("2011-01-01/2011-01-10", "1", 1);
     add("2011-01-02/2011-01-05", "2", 1);
 
-    assertSingleElementChunks(
-        makeSingle("1", 1),
-        timeline.findChunk(Intervals.of("2011-01-02T02/2011-01-04"), "1", 0)
+    Assert.assertEquals(
+        new PartitionHolder<>(makeSingle("1", 1)).asImmutable(),
+        timeline.findEntry(Intervals.of("2011-01-02T02/2011-01-04"), "1")
     );
   }
 

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

@@ -98,16 +98,6 @@ static void assertValues(
     Assert.assertEquals(expected, actualSet);
   }
 
-  static void assertSingleElementChunks(
-      PartitionChunk<OvershadowableInteger> expected,
-      PartitionChunk<OvershadowableInteger> actual
-  )
-  {
-    SingleElementPartitionChunk<OvershadowableInteger> expectedSingle = (SingleElementPartitionChunk<OvershadowableInteger>) expected;
-    SingleElementPartitionChunk<OvershadowableInteger> actualSingle = (SingleElementPartitionChunk<OvershadowableInteger>) actual;
-    Assert.assertEquals(expectedSingle.getObject(), actualSingle.getObject());
-  }
-
   static VersionedIntervalTimeline<String, OvershadowableInteger> makeStringIntegerTimeline()
   {
     return new VersionedIntervalTimeline<>(Ordering.natural());