Core: Coalesce nearby blob reads in PuffinReader.readAll

Resolve the TODO from #4537 by reading contiguous blobs in a single
  request, handed back as no-copy ByteBuffer views. A contiguous run is
  split into bounded reads (capped at MANIFEST_TARGET_SIZE, 8 MiB) so a
  large blob section never allocates one oversized buffer.

Author: vishnuprakaz

core/src/main/java/org/apache/iceberg/puffin/PuffinReader.java

@@ -35,13 +35,17 @@
 import org.apache.iceberg.puffin.PuffinFormat.Flag;
 import org.apache.iceberg.relocated.com.google.common.base.Preconditions;
 import org.apache.iceberg.relocated.com.google.common.collect.ImmutableList;
+import org.apache.iceberg.relocated.com.google.common.collect.Lists;
 import org.apache.iceberg.relocated.com.google.common.io.ByteStreams;
 import org.apache.iceberg.util.Pair;
 
 public class PuffinReader implements Closeable {
   // Must not be modified
   private static final byte[] MAGIC = PuffinFormat.getMagic();
 
+  // Bound on a single coalesced read, so a long contiguous run is split into bounded buffers.
+  private static final long MAX_COALESCED_READ_SIZE = 8 * 1024 * 1024;
+
   private final long fileSize;
   private final SeekableInputStream input;
   private Integer knownFooterSize;
@@ -125,27 +129,61 @@ public Iterable<Pair<BlobMetadata, ByteBuffer>> readAll(List<BlobMetadata> blobs
       return ImmutableList.of();
     }
 
-    // TODO inspect blob offsets and coalesce read regions close to each other
-
-    return () ->
-        blobs.stream()
-            .sorted(Comparator.comparingLong(BlobMetadata::offset))
-            .map(
-                (BlobMetadata blobMetadata) -> {
-                  try {
-                    input.seek(blobMetadata.offset());
-                    byte[] bytes = new byte[Math.toIntExact(blobMetadata.length())];
-                    ByteStreams.readFully(input, bytes);
-                    ByteBuffer rawData = ByteBuffer.wrap(bytes);
-                    PuffinCompressionCodec codec =
-                        PuffinCompressionCodec.forName(blobMetadata.compressionCodec());
-                    ByteBuffer data = PuffinFormat.decompress(codec, rawData);
-                    return Pair.of(blobMetadata, data);
-                  } catch (IOException e) {
-                    throw new UncheckedIOException(e);
-                  }
-                })
-            .iterator();
+    List<BlobMetadata> ordered = Lists.newArrayList(blobs);
+    ordered.sort(Comparator.comparingLong(BlobMetadata::offset));
+    List<List<BlobMetadata>> groups = coalesce(ordered, MAX_COALESCED_READ_SIZE);
+
+    return () -> groups.stream().flatMap(group -> readGroup(group).stream()).iterator();
+  }
+
+  private List<Pair<BlobMetadata, ByteBuffer>> readGroup(List<BlobMetadata> group) {
+    long start = group.get(0).offset();
+    long end = start;
+    for (BlobMetadata blob : group) {
+      end = Math.max(end, blob.offset() + blob.length());
+    }
+
+    byte[] region;
+    try {
+      region = readInput(start, Math.toIntExact(end - start));
+    } catch (IOException e) {
+      throw new UncheckedIOException(e);
+    }
+
+    List<Pair<BlobMetadata, ByteBuffer>> data = Lists.newArrayListWithCapacity(group.size());
+    for (BlobMetadata blob : group) {
+      int position = Math.toIntExact(blob.offset() - start);
+      // slice so the blob's bytes align to the buffer's array offset (needed by decompress)
+      ByteBuffer rawData =
+          ByteBuffer.wrap(region, position, Math.toIntExact(blob.length())).slice();
+      PuffinCompressionCodec codec = PuffinCompressionCodec.forName(blob.compressionCodec());
+      data.add(Pair.of(blob, PuffinFormat.decompress(codec, rawData)));
+    }
+
+    return data;
+  }
+
+  static List<List<BlobMetadata>> coalesce(List<BlobMetadata> orderedByOffset, long maxReadSize) {
+    List<List<BlobMetadata>> groups = Lists.newArrayList();
+    List<BlobMetadata> current = null;
+    long currentStart = 0;
+    long currentEnd = 0;
+    for (BlobMetadata blob : orderedByOffset) {
+      long blobEnd = blob.offset() + blob.length();
+      boolean gap = current != null && blob.offset() > currentEnd;
+      boolean tooLarge = current != null && blobEnd - currentStart > maxReadSize;
+      if (current == null || gap || tooLarge) {
+        current = Lists.newArrayList();
+        groups.add(current);
+        currentStart = blob.offset();
+        currentEnd = blob.offset();
+      }
+
+      current.add(blob);
+      currentEnd = Math.max(currentEnd, blobEnd);
+    }
+
+    return groups;
   }
 
   private static void checkMagic(byte[] data, int offset) {

core/src/test/java/org/apache/iceberg/puffin/TestPuffinReader.java

@@ -28,6 +28,7 @@
 import static org.assertj.core.api.Assertions.assertThat;
 import static org.assertj.core.api.Assertions.assertThatThrownBy;
 
+import java.util.List;
 import java.util.Map;
 import javax.annotation.Nullable;
 import org.apache.iceberg.inmemory.InMemoryInputFile;
@@ -154,4 +155,77 @@ public void testValidateFooterSizeValue() throws Exception {
           .isEqualTo(ImmutableMap.of("created-by", "Test 1234"));
     }
   }
+
+  @Test
+  void coalescesContiguousBlobsAndSplitsOnGaps() {
+    BlobMetadata a = blob("a", 0, 10);
+    BlobMetadata b = blob("b", 10, 8);
+    BlobMetadata c = blob("c", 15, 10); // overlaps b: same read
+    BlobMetadata d = blob("d", 100, 4); // gap: separate read
+
+    List<List<BlobMetadata>> groups =
+        PuffinReader.coalesce(ImmutableList.of(a, b, c, d), Long.MAX_VALUE);
+
+    assertThat(groups).hasSize(2);
+    assertThat(groups.get(0)).containsExactly(a, b, c);
+    assertThat(groups.get(1)).containsExactly(d);
+  }
+
+  @Test
+  void splitsContiguousRunThatExceedsMaxReadSize() {
+    BlobMetadata a = blob("a", 0, 10);
+    BlobMetadata b = blob("b", 10, 10);
+    BlobMetadata c = blob("c", 20, 10); // would exceed the 25-byte cap: separate read
+
+    List<List<BlobMetadata>> groups = PuffinReader.coalesce(ImmutableList.of(a, b, c), 25);
+
+    assertThat(groups).hasSize(2);
+    assertThat(groups.get(0)).containsExactly(a, b);
+    assertThat(groups.get(1)).containsExactly(c);
+  }
+
+  @Test
+  void keepsBlobLargerThanMaxReadSizeInItsOwnGroup() {
+    BlobMetadata big = blob("big", 0, 100); // exceeds the cap on its own
+    BlobMetadata next = blob("next", 100, 10);
+
+    List<List<BlobMetadata>> groups = PuffinReader.coalesce(ImmutableList.of(big, next), 25);
+
+    assertThat(groups).hasSize(2);
+    assertThat(groups.get(0)).containsExactly(big);
+    assertThat(groups.get(1)).containsExactly(next);
+  }
+
+  @Test
+  void readAllReadsFarApartBlobsInSeparateRegions() throws Exception {
+    byte[] first = "first blob".getBytes(UTF_8);
+    byte[] second = "second blob far away".getBytes(UTF_8);
+    int secondOffset = 2_000_000; // far from the first blob: separate read
+
+    // build the raw layout directly; the writer always packs blobs contiguously
+    byte[] bytes = new byte[secondOffset + second.length];
+    System.arraycopy(first, 0, bytes, 0, first.length);
+    System.arraycopy(second, 0, bytes, secondOffset, second.length);
+
+    BlobMetadata firstBlob = blob("first", 0, first.length);
+    BlobMetadata secondBlob = blob("second", secondOffset, second.length);
+
+    InMemoryInputFile inputFile = new InMemoryInputFile(bytes);
+    try (PuffinReader reader = Puffin.read(inputFile).withFileSize(bytes.length).build()) {
+      Map<BlobMetadata, byte[]> read =
+          Streams.stream(reader.readAll(ImmutableList.of(firstBlob, secondBlob)))
+              .collect(toImmutableMap(Pair::first, pair -> ByteBuffers.toByteArray(pair.second())));
+
+      assertThat(read)
+          .as("read")
+          .containsOnlyKeys(firstBlob, secondBlob)
+          .containsEntry(firstBlob, first)
+          .containsEntry(secondBlob, second);
+    }
+  }
+
+  private static BlobMetadata blob(String type, long offset, long length) {
+    return new BlobMetadata(
+        type, ImmutableList.of(1), 1L, 1L, offset, length, null, ImmutableMap.of());
+  }
 }