diff --git a/core/src/main/java/org/apache/iceberg/puffin/PuffinReader.java b/core/src/main/java/org/apache/iceberg/puffin/PuffinReader.java index e30b6e1ee6ef..e0f694f79fe3 100644 --- a/core/src/main/java/org/apache/iceberg/puffin/PuffinReader.java +++ b/core/src/main/java/org/apache/iceberg/puffin/PuffinReader.java @@ -35,6 +35,7 @@ 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; @@ -42,6 +43,9 @@ 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> readAll(List 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 ordered = Lists.newArrayList(blobs); + ordered.sort(Comparator.comparingLong(BlobMetadata::offset)); + List> groups = coalesce(ordered, MAX_COALESCED_READ_SIZE); + + return () -> groups.stream().flatMap(group -> readGroup(group).stream()).iterator(); + } + + private List> readGroup(List 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> 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> coalesce(List orderedByOffset, long maxReadSize) { + List> groups = Lists.newArrayList(); + List 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) { diff --git a/core/src/test/java/org/apache/iceberg/puffin/TestPuffinReader.java b/core/src/test/java/org/apache/iceberg/puffin/TestPuffinReader.java index 094fee618dfa..d5165d7e2fcc 100644 --- a/core/src/test/java/org/apache/iceberg/puffin/TestPuffinReader.java +++ b/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,78 @@ public void testValidateFooterSizeValue() throws Exception { .isEqualTo(ImmutableMap.of("created-by", "Test 1234")); } } + + @Test + void coalescesContiguousBlobsAndSplitsOnGaps() { + BlobMetadata blobA = blob("a", 0, 10); + BlobMetadata blobB = blob("b", 10, 8); + BlobMetadata blobC = blob("c", 15, 10); // overlaps blobB: same read + BlobMetadata blobD = blob("d", 100, 4); // gap: separate read + + List> groups = + PuffinReader.coalesce(ImmutableList.of(blobA, blobB, blobC, blobD), Long.MAX_VALUE); + + assertThat(groups).hasSize(2); + assertThat(groups.get(0)).containsExactly(blobA, blobB, blobC); + assertThat(groups.get(1)).containsExactly(blobD); + } + + @Test + void splitsContiguousRunThatExceedsMaxReadSize() { + BlobMetadata blobA = blob("a", 0, 10); + BlobMetadata blobB = blob("b", 10, 10); + BlobMetadata blobC = blob("c", 20, 10); // would exceed the 25-byte cap: separate read + + List> groups = + PuffinReader.coalesce(ImmutableList.of(blobA, blobB, blobC), 25); + + assertThat(groups).hasSize(2); + assertThat(groups.get(0)).containsExactly(blobA, blobB); + assertThat(groups.get(1)).containsExactly(blobC); + } + + @Test + void keepsBlobLargerThanMaxReadSizeInItsOwnGroup() { + BlobMetadata big = blob("big", 0, 100); // exceeds the cap on its own + BlobMetadata next = blob("next", 100, 10); + + List> 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 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()); + } }