diff --git a/paimon-common/src/main/java/org/apache/paimon/fs/ByteArraySeekableStream.java b/paimon-common/src/main/java/org/apache/paimon/fs/ByteArraySeekableStream.java index d6536927b100..2b83cd5b4636 100644 --- a/paimon-common/src/main/java/org/apache/paimon/fs/ByteArraySeekableStream.java +++ b/paimon-common/src/main/java/org/apache/paimon/fs/ByteArraySeekableStream.java @@ -92,7 +92,7 @@ public ByteArrayStream(byte[] buf) { } public void seek(int position) throws IOException { - if (position >= count) { + if (position > count) { throw new EOFException("Can't seek position: " + position + ", length is " + count); } pos = position; diff --git a/paimon-common/src/test/java/org/apache/paimon/fs/ByteArraySeekableStreamTest.java b/paimon-common/src/test/java/org/apache/paimon/fs/ByteArraySeekableStreamTest.java index 2df02e85db49..31d37a6a2f43 100644 --- a/paimon-common/src/test/java/org/apache/paimon/fs/ByteArraySeekableStreamTest.java +++ b/paimon-common/src/test/java/org/apache/paimon/fs/ByteArraySeekableStreamTest.java @@ -65,12 +65,25 @@ public void testBasic() throws IOException { } } + @Test + public void testSeekToEnd() throws IOException { + for (int length : new int[] {0, 10}) { + try (ByteArraySeekableStream stream = + new ByteArraySeekableStream(randomBytes(length))) { + stream.seek(length); + Assertions.assertThat(stream.getPos()).isEqualTo(length); + Assertions.assertThat(stream.available()).isZero(); + Assertions.assertThat(stream.read()).isEqualTo(-1); + } + } + } + @Test public void testThrow() { int bl = 10; byte[] b = randomBytes(bl); ByteArraySeekableStream byteArraySeekableStream = new ByteArraySeekableStream(b); - Assertions.assertThatCode(() -> byteArraySeekableStream.seek(10)) - .hasMessage("Can't seek position: 10, length is 10"); + Assertions.assertThatCode(() -> byteArraySeekableStream.seek(11)) + .hasMessage("Can't seek position: 11, length is 10"); } } diff --git a/paimon-core/src/main/java/org/apache/paimon/manifest/ManifestAvroReader.java b/paimon-core/src/main/java/org/apache/paimon/manifest/ManifestAvroReader.java index 995834c37e7e..4e84c690ee97 100644 --- a/paimon-core/src/main/java/org/apache/paimon/manifest/ManifestAvroReader.java +++ b/paimon-core/src/main/java/org/apache/paimon/manifest/ManifestAvroReader.java @@ -26,6 +26,7 @@ import org.apache.paimon.format.avro.AvroRecordDecoder; import org.apache.paimon.format.avro.AvroRecordDecoder.FieldDecoder; import org.apache.paimon.format.avro.AvroRecordDecoder.FieldType; +import org.apache.paimon.fs.SeekableInputStream; import org.apache.paimon.partition.PartitionPredicate; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.CloseableIterator; @@ -34,7 +35,6 @@ import javax.annotation.Nullable; import java.io.IOException; -import java.io.InputStream; import java.io.UncheckedIOException; import java.nio.ByteBuffer; import java.util.Arrays; @@ -56,7 +56,7 @@ public final class ManifestAvroReader implements AutoCloseable { private long blockOrdinal = -1; - ManifestAvroReader(InputStream input) throws IOException { + ManifestAvroReader(SeekableInputStream input) throws IOException { AvroBlockReader blockReader = null; try { blockReader = new AvroBlockReader(input); @@ -70,6 +70,21 @@ public final class ManifestAvroReader implements AutoCloseable { } } + /** Returns a copy of the complete OCF header, including schema, codec and sync marker. */ + public byte[] headerBytes() { + return blockReader.headerBytes(); + } + + /** Returns the physical block offset; read immediately after {@link #next()}. */ + public long blockOffset() { + return blockReader.blockOffset(); + } + + /** Returns the last-read block's encoded length, including its header and sync marker. */ + public long blockLength() { + return blockReader.blockLength(); + } + /** Returns whether another raw Avro block is available. */ public boolean hasNext() throws IOException { return blockReader.hasNextBlock(); diff --git a/paimon-core/src/test/java/org/apache/paimon/manifest/ManifestFileTest.java b/paimon-core/src/test/java/org/apache/paimon/manifest/ManifestFileTest.java index d8eab09df774..b5ff0fb00249 100644 --- a/paimon-core/src/test/java/org/apache/paimon/manifest/ManifestFileTest.java +++ b/paimon-core/src/test/java/org/apache/paimon/manifest/ManifestFileTest.java @@ -58,6 +58,7 @@ import java.io.IOException; import java.io.UncheckedIOException; import java.nio.ByteBuffer; +import java.nio.file.Files; import java.util.ArrayList; import java.util.Arrays; import java.util.Collections; @@ -1129,11 +1130,18 @@ void testBlockReaderReadsAcrossMultipleBlocks() throws Exception { ProjectedManifestEntry.Projection projection = projection(DataFileMeta.FILE_NAME); int blockCount = 0; int rowCount = 0; + byte[] bytes = Files.readAllBytes(tempDir.resolve("manifest").resolve(manifest.fileName())); try (ManifestAvroReader reader = openManifestReader(manifest)) { + byte[] header = reader.headerBytes(); + assertThat(header).isEqualTo(Arrays.copyOf(bytes, header.length)); + long nextOffset = header.length; while (reader.hasNext()) { - ManifestAvroReader.RowIterator rows = - reader.next().toRows(projection.projectedType()); + ManifestAvroReader.RawBlock block = reader.next(); + assertThat(reader.blockOffset()).isEqualTo(nextOffset); + assertThat(reader.blockLength()).isPositive(); + nextOffset += reader.blockLength(); + ManifestAvroReader.RowIterator rows = block.toRows(projection.projectedType()); assertThat(rows.hasNext()).isTrue(); while (rows.hasNext()) { rows.next(); @@ -1141,6 +1149,7 @@ void testBlockReaderReadsAcrossMultipleBlocks() throws Exception { } blockCount++; } + assertThat(nextOffset).isEqualTo(bytes.length); } assertThat(blockCount).isGreaterThan(1); diff --git a/paimon-format/src/main/java/org/apache/avro/file/RawBlockReader.java b/paimon-format/src/main/java/org/apache/avro/file/RawBlockReader.java index 43a68c54a44d..60ae9b195abb 100644 --- a/paimon-format/src/main/java/org/apache/avro/file/RawBlockReader.java +++ b/paimon-format/src/main/java/org/apache/avro/file/RawBlockReader.java @@ -18,26 +18,77 @@ package org.apache.avro.file; +import org.apache.paimon.fs.SeekableInputStream; +import org.apache.paimon.utils.IOUtils; + import org.apache.avro.Schema; import org.apache.avro.io.DatumReader; import org.apache.avro.io.Decoder; import java.io.IOException; -import java.io.InputStream; +import java.util.NoSuchElementException; /** Package bridge exposing Avro's compressed blocks without reflection. */ public final class RawBlockReader extends DataFileStream { - public RawBlockReader(InputStream input) throws IOException { + private final SeekableInputStream input; + private final byte[] headerBytes; + private long blockOffset; + private long blockLength; + private boolean pending; + + public RawBlockReader(SeekableInputStream input) throws IOException { + this(input, input.getPos()); + } + + private RawBlockReader(SeekableInputStream input, long headerOffset) throws IOException { super(input, new NoOpDatumReader()); + this.input = input; + this.headerBytes = new byte[Math.toIntExact(position() - headerOffset)]; + long resumePosition = input.getPos(); + input.seek(headerOffset); + IOUtils.readFully(input, headerBytes); + // Preserve the position past any bytes already buffered by the Avro decoder. + input.seek(resumePosition); + } + + /** Returns a copy of the complete OCF header, including schema, codec and sync marker. */ + public byte[] headerBytes() { + return headerBytes.clone(); + } + + /** + * Returns the physical block offset; read immediately after {@link #nextRawBlock(RawBlock)}. + */ + public long blockOffset() { + return blockOffset; + } + + /** Returns the last-read block's encoded length, including its header and sync marker. */ + public long blockLength() { + return blockLength; } - public boolean hasNextRawBlock() { - return super.hasNextBlock(); + private long position() throws IOException { + // This is the same read-ahead adjustment used by DataFileReader.blockFinished(). + return input.getPos() - vin.inputStream().available(); + } + + public boolean hasNextRawBlock() throws IOException { + if (!pending) { + blockOffset = position(); + pending = super.hasNextBlock(); + } + return pending; } public RawBlock nextRawBlock(RawBlock reuse) throws IOException { + if (!hasNextRawBlock()) { + throw new NoSuchElementException(); + } DataBlock raw = super.nextRawBlock(reuse == null ? null : reuse.dataBlock()); + blockLength = position() - blockOffset; + pending = false; return reuse == null ? new RawBlock(raw, resolveCodec(), getSchema()) : reuse.replace(raw, resolveCodec(), getSchema()); diff --git a/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBlockReader.java b/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBlockReader.java index c4359afcda43..61e387d580c8 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBlockReader.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroBlockReader.java @@ -18,6 +18,7 @@ package org.apache.paimon.format.avro; +import org.apache.paimon.fs.SeekableInputStream; import org.apache.paimon.types.RowType; import org.apache.paimon.utils.IOUtils; @@ -30,7 +31,6 @@ import java.io.Closeable; import java.io.IOException; -import java.io.InputStream; import java.util.Collections; /** @@ -45,7 +45,7 @@ public final class AvroBlockReader implements Closeable { private @Nullable AvroRawBlock borrowedRawBlock; - public AvroBlockReader(InputStream input) throws IOException { + public AvroBlockReader(SeekableInputStream input) throws IOException { try { this.reader = new RawBlockReader(input); } catch (IOException | RuntimeException | Error e) { @@ -54,6 +54,23 @@ public AvroBlockReader(InputStream input) throws IOException { } } + /** Returns a copy of the complete OCF header, including schema, codec and sync marker. */ + public byte[] headerBytes() { + return reader.headerBytes(); + } + + /** + * Returns the physical block offset; read immediately after {@link #nextBorrowedRawBlock()}. + */ + public long blockOffset() { + return reader.blockOffset(); + } + + /** Returns the last-read block's encoded length, including its header and sync marker. */ + public long blockLength() { + return reader.blockLength(); + } + /** Creates a record decoder from the writer schema stored in the Avro file header. */ public AvroRecordDecoder createRecordDecoder() { return new AvroRecordDecoder(reader.getSchema()); diff --git a/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroBlockReaderTest.java b/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroBlockReaderTest.java new file mode 100644 index 000000000000..c0faacd3eb8b --- /dev/null +++ b/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroBlockReaderTest.java @@ -0,0 +1,270 @@ +/* + * 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.paimon.format.avro; + +import org.apache.paimon.fs.ByteArraySeekableStream; +import org.apache.paimon.fs.Path; +import org.apache.paimon.fs.SeekableInputStream; +import org.apache.paimon.fs.SeekableInputStreamWrapper; +import org.apache.paimon.fs.local.LocalFileIO; + +import org.apache.avro.Schema; +import org.apache.avro.file.CodecFactory; +import org.apache.avro.file.DataFileStream; +import org.apache.avro.file.DataFileWriter; +import org.apache.avro.generic.GenericDatumReader; +import org.apache.avro.generic.GenericDatumWriter; +import org.junit.jupiter.api.Test; +import org.junit.jupiter.api.io.TempDir; +import org.junit.jupiter.params.ParameterizedTest; +import org.junit.jupiter.params.provider.ValueSource; + +import java.io.ByteArrayInputStream; +import java.io.ByteArrayOutputStream; +import java.io.IOException; +import java.nio.file.Files; +import java.util.ArrayList; +import java.util.Arrays; +import java.util.List; +import java.util.NoSuchElementException; +import java.util.concurrent.atomic.AtomicBoolean; + +import static org.assertj.core.api.Assertions.assertThat; +import static org.assertj.core.api.Assertions.assertThatThrownBy; + +/** Tests for physical Avro block metadata. */ +class AvroBlockReaderTest { + + private static final Schema SCHEMA = Schema.create(Schema.Type.LONG); + + @TempDir private java.nio.file.Path tempDir; + + @ParameterizedTest + @ValueSource(strings = {"null", "deflate", "snappy", "zstandard"}) + void blockMetadataMatchesWriterBoundaries(String codec) throws Exception { + long[][] values = {{0L, 1L, Long.MAX_VALUE}, {100L}, {1000L, 1001L}}; + long[] boundaries = new long[values.length + 1]; + ByteArrayOutputStream output = new ByteArrayOutputStream(); + try (DataFileWriter writer = new DataFileWriter<>(new GenericDatumWriter<>(SCHEMA))) { + writer.setCodec(CodecFactory.fromString(codec)); + // Exercise headers larger than the decoder's read-ahead buffer. + writer.setMeta("test.padding", new byte[20_000]); + writer.create(SCHEMA, output); + boundaries[0] = writer.sync(); + for (int i = 0; i < values.length; i++) { + for (long value : values[i]) { + writer.append(value); + } + boundaries[i + 1] = writer.sync(); + } + } + byte[] bytes = output.toByteArray(); + assertThat(boundaries[values.length]).isEqualTo(bytes.length); + + for (int maxRead : new int[] {1, 7, Integer.MAX_VALUE}) { + List seeks = new ArrayList<>(); + SeekableInputStream input = + new SeekableInputStreamWrapper(open(bytes)) { + @Override + public int read(byte[] data, int offset, int length) throws IOException { + return super.read(data, offset, Math.min(length, maxRead)); + } + + @Override + public void seek(long position) throws IOException { + seeks.add(position); + super.seek(position); + } + }; + try (AvroBlockReader reader = new AvroBlockReader(input)) { + long resumePosition = input.getPos(); + assertThat(seeks).containsExactly(0L, resumePosition); + byte[] header = reader.headerBytes(); + assertThat(header).isEqualTo(Arrays.copyOf(bytes, (int) boundaries[0])); + assertThat(input.getPos()).isEqualTo(resumePosition); + assertThat(seeks).containsExactly(0L, resumePosition); + byte[] anotherHeader = reader.headerBytes(); + anotherHeader[0] = 0; + assertThat(reader.headerBytes()).isEqualTo(header); + assertThat(seeks).hasSize(2); + AvroRawBlock previous = null; + for (int i = 0; i < values.length; i++) { + // Exercise next() both directly and after repeated look-ahead calls. + if (i > 0) { + assertThat(reader.hasNextBlock()).isTrue(); + assertThat(reader.hasNextBlock()).isTrue(); + } + AvroRawBlock block = reader.nextBorrowedRawBlock(); + if (previous != null) { + assertThat(block).isSameAs(previous); + } + previous = block; + assertThat(block.recordCount()).isEqualTo(values[i].length); + assertThat(reader.blockOffset()).isEqualTo(boundaries[i]); + assertThat(reader.blockLength()).isEqualTo(boundaries[i + 1] - boundaries[i]); + assertBlockReadable( + header, bytes, reader.blockOffset(), reader.blockLength(), values[i]); + } + assertThat(reader.hasNextBlock()).isFalse(); + assertThat(reader.hasNextBlock()).isFalse(); + assertThatThrownBy(reader::nextBorrowedRawBlock) + .isInstanceOf(NoSuchElementException.class); + } + } + } + + @Test + void emptyFileContainsOnlyTheHeader() throws IOException { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + try (DataFileWriter writer = new DataFileWriter<>(new GenericDatumWriter<>(SCHEMA))) { + writer.create(SCHEMA, output); + } + byte[] bytes = output.toByteArray(); + try (AvroBlockReader reader = new AvroBlockReader(open(bytes))) { + assertThat(reader.headerBytes()).isEqualTo(bytes); + assertThat(reader.hasNextBlock()).isFalse(); + assertThatThrownBy(reader::nextBorrowedRawBlock) + .isInstanceOf(NoSuchElementException.class); + } + } + + @Test + void failedHeaderReadClosesTheInput() throws IOException { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + try (DataFileWriter writer = new DataFileWriter<>(new GenericDatumWriter<>(SCHEMA))) { + writer.create(SCHEMA, output); + writer.append(11L); + } + byte[] bytes = output.toByteArray(); + AtomicBoolean closed = new AtomicBoolean(); + SeekableInputStream input = + new SeekableInputStreamWrapper(open(bytes)) { + private boolean readingHeader; + + @Override + public void seek(long position) throws IOException { + super.seek(position); + readingHeader = position == 0; + } + + @Override + public int read(byte[] data, int offset, int length) throws IOException { + if (readingHeader) { + throw new IOException("header read failed"); + } + return super.read(data, offset, length); + } + + @Override + public void close() throws IOException { + super.close(); + closed.set(true); + } + }; + assertThatThrownBy(() -> new AvroBlockReader(input)) + .isInstanceOf(IOException.class) + .hasMessage("header read failed"); + assertThat(closed.get()).isTrue(); + } + + @Test + void headerFromMemoryCanRestoreEof() throws IOException { + for (int records : new int[] {0, 1}) { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + long headerLength; + try (DataFileWriter writer = + new DataFileWriter<>(new GenericDatumWriter<>(SCHEMA))) { + writer.create(SCHEMA, output); + headerLength = writer.sync(); + if (records > 0) { + writer.append(17L); + } + } + byte[] bytes = output.toByteArray(); + ByteArraySeekableStream input = new ByteArraySeekableStream(bytes); + try (AvroBlockReader reader = new AvroBlockReader(input)) { + assertThat(input.getPos()).isEqualTo(bytes.length); + byte[] header = reader.headerBytes(); + assertThat(header).isEqualTo(Arrays.copyOf(bytes, (int) headerLength)); + assertThat(input.getPos()).isEqualTo(bytes.length); + if (records > 0) { + assertThat(reader.nextBorrowedRawBlock().recordCount()).isEqualTo(records); + assertBlockReadable( + header, + bytes, + reader.blockOffset(), + reader.blockLength(), + new long[] {17L}); + } + assertThat(reader.hasNextBlock()).isFalse(); + } + } + } + + @Test + void headerStartsAtInitialStreamPosition() throws IOException { + ByteArrayOutputStream output = new ByteArrayOutputStream(); + long headerLength; + try (DataFileWriter writer = new DataFileWriter<>(new GenericDatumWriter<>(SCHEMA))) { + writer.create(SCHEMA, output); + headerLength = writer.sync(); + writer.append(17L); + } + byte[] avro = output.toByteArray(); + int prefixLength = 13; + byte[] bytes = new byte[prefixLength + avro.length]; + System.arraycopy(avro, 0, bytes, prefixLength, avro.length); + SeekableInputStream input = open(bytes); + input.seek(prefixLength); + try (AvroBlockReader reader = new AvroBlockReader(input)) { + long resumePosition = input.getPos(); + byte[] header = reader.headerBytes(); + assertThat(header).isEqualTo(Arrays.copyOf(avro, (int) headerLength)); + assertThat(input.getPos()).isEqualTo(resumePosition); + assertThat(reader.nextBorrowedRawBlock().recordCount()).isEqualTo(1); + assertThat(reader.blockOffset()).isEqualTo(prefixLength + headerLength); + assertBlockReadable( + header, bytes, reader.blockOffset(), reader.blockLength(), new long[] {17L}); + assertThat(reader.hasNextBlock()).isFalse(); + } + } + + private SeekableInputStream open(byte[] bytes) throws IOException { + java.nio.file.Path file = Files.createTempFile(tempDir, "blocks-", ".avro"); + Files.write(file, bytes); + return LocalFileIO.create().newInputStream(new Path(file.toUri())); + } + + private static void assertBlockReadable( + byte[] header, byte[] file, long offset, long length, long[] expected) + throws IOException { + ByteArrayOutputStream selected = new ByteArrayOutputStream(); + selected.write(header); + selected.write(file, (int) offset, (int) length); + try (DataFileStream reader = + new DataFileStream<>( + new ByteArrayInputStream(selected.toByteArray()), + new GenericDatumReader<>())) { + for (long value : expected) { + assertThat(reader.next()).isEqualTo(value); + } + assertThat(reader.hasNext()).isFalse(); + } + } +}