JingsongLi commented on code in PR #9182:
URL: https://github.com/apache/paimon/pull/9182#discussion_r3763201619
##########
paimon-core/src/main/java/org/apache/paimon/manifest/ManifestAvroReader.java:
##########
@@ -36,281 +40,915 @@
import java.io.IOException;
import java.io.InputStream;
import java.io.UncheckedIOException;
+import java.nio.ByteBuffer;
+import java.util.Arrays;
+import java.util.Iterator;
import java.util.NoSuchElementException;
-import static org.apache.paimon.utils.SerializationUtils.deserializeBinaryRow;
+/** Reader which exposes reusable raw blocks from a Manifest Avro file. */
+public final class ManifestAvroReader implements AutoCloseable {
-/**
- * Schema-aware Avro reader which projects fields and filters before decoding
data file metadata.
- */
-final class ManifestAvroReader implements CloseableIterator<InternalRow> {
+ private static final String[] TOP_LEVEL_FIELDS = {
+ ManifestSchemaUtils.FORMAT_IDENTIFIER_FIELD,
+ ManifestEntry.KIND,
+ ManifestEntry.PARTITION,
+ ManifestEntry.BUCKET,
+ ManifestEntry.TOTAL_BUCKETS,
+ ManifestEntry.FILE
+ };
- private final AvroBlockReader blockReader;
- private final AvroRecordDecoder decoder;
- private final ManifestRecordDecoder recordDecoder;
+ private final AvroBlockReader stream;
+ private final DecoderContext decoderContext;
+ private final boolean rawBlockCopySupported;
+ private final ReadStatistics readStatistics = new ReadStatistics();
- private long recordsRemaining;
- private @Nullable InternalRow next;
- private boolean nextReady;
- private boolean finished;
+ private long blockOrdinal = -1;
- ManifestAvroReader(
- InputStream input,
- RowType projectedType,
- @Nullable PartitionPredicate partitionFilter,
- @Nullable BucketFilter bucketFilter)
- throws IOException {
- AvroBlockReader blockReader = new AvroBlockReader(input);
+ ManifestAvroReader(InputStream input, AvroFileFormat avroFileFormat)
throws IOException {
+ AvroBlockReader stream = null;
try {
- AvroRecordDecoder decoder = blockReader.createRecordDecoder();
- this.recordDecoder =
- new ManifestRecordDecoder(
- decoder, projectedType, partitionFilter,
bucketFilter);
- this.decoder = decoder;
- this.blockReader = blockReader;
- } catch (RuntimeException | Error e) {
- IOUtils.closeQuietly(blockReader);
- throw e;
+ stream = new AvroBlockReader(input);
+ this.stream = stream;
+ this.decoderContext = new
DecoderContext(stream.createRecordDecoder());
+ this.rawBlockCopySupported =
+ stream.supportsRawBlockCopy(avroFileFormat,
ManifestEntry.MANIFEST_ROW_TYPE);
+ } catch (IOException | RuntimeException | Error failure) {
+ IOUtils.closeQuietly(stream == null ? input : stream);
+ throw failure;
+ }
+ }
+
+ /** Returns whether another raw Avro block is available. */
+ public boolean hasNext() throws IOException {
+ return stream.hasNextBlock();
+ }
+
+ /** Returns the next raw block without decompressing it. */
+ public RawBlock next() throws IOException {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
}
+ return new RawBlock(
+ decoderContext,
+ rawBlockCopySupported,
+ stream.nextBorrowedRawBlock(),
+ ++blockOrdinal);
+ }
+
+ /**
+ * Returns an iterator over projected rows from all remaining blocks.
+ *
+ * <p>The returned row is reused within each block and must be consumed
before the iterator
+ * advances. Closing the iterator closes this reader.
+ */
+ public CloseableIterator<InternalRow> read(
+ RowType projectedType,
+ @Nullable PartitionPredicate partitionFilter,
+ @Nullable BucketFilter bucketFilter) {
+ return new RowsIterator(projectedType, partitionFilter, bucketFilter);
+ }
+
+ long decodedDataFiles() {
+ return readStatistics.decodedDataFiles;
+ }
+
+ long skippedDataFiles() {
+ return readStatistics.skippedDataFiles;
}
@Override
- public boolean hasNext() {
- if (nextReady) {
- return true;
+ public void close() throws IOException {
+ stream.close();
+ }
+
+ /** Borrowed raw block which must be consumed before the enclosing reader
advances. */
+ public static final class RawBlock {
+
+ private final DecoderContext decoderContext;
+ private final boolean rawBlockCopySupported;
+ private final AvroRawBlock block;
+ private final long blockOrdinal;
+ private final long blockRecordCount;
+
+ private RawBlock(
+ DecoderContext decoderContext,
+ boolean rawBlockCopySupported,
+ AvroRawBlock block,
+ long blockOrdinal) {
+ this.decoderContext = decoderContext;
+ this.rawBlockCopySupported = rawBlockCopySupported;
+ this.block = block;
+ this.blockOrdinal = blockOrdinal;
+ this.blockRecordCount = block.recordCount();
}
- if (finished) {
- return false;
+
+ /** Lazily decompresses this block and returns an iterator over one
reusable row. */
+ public RowIterator toRows(RowType projectRowType) throws IOException {
+ return toRows(projectRowType, null, null, null);
}
- try {
- while (true) {
- if (recordsRemaining == 0) {
- if (decoder.isInitialized() && !decoder.isEnd()) {
- throw new IOException(
- "Manifest Avro block contains trailing
undecoded bytes.");
- }
- if (!blockReader.hasNextBlock()) {
- finished = true;
+ private RowIterator toRows(
+ RowType projectRowType,
+ @Nullable PartitionPredicate partitionFilter,
+ @Nullable BucketFilter bucketFilter,
+ @Nullable ReadStatistics readStatistics)
+ throws IOException {
+ ManifestEntryDecoder recordDecoder =
decoderContext.recordDecoder(projectRowType);
+
+ ByteBuffer decompressed = decoderContext.decompress(block);
+ decoderContext.decoder.reset(decompressed);
+ BlockRow row = recordDecoder.createRow(decompressed.array());
+ return new RowIterator(
+ blockRecordCount,
+ decoderContext.decoder,
+ recordDecoder,
+ row,
+ partitionFilter,
+ bucketFilter,
+ readStatistics);
+ }
+
+ public long blockOrdinal() {
+ return blockOrdinal;
+ }
+
+ public long recordCount() {
+ return blockRecordCount;
+ }
+
+ public boolean rawBlockCopySupported() {
+ return rawBlockCopySupported;
+ }
+
+ public AvroRawBlock encodedBlock() {
+ return block;
+ }
+ }
+
+ private final class RowsIterator implements CloseableIterator<InternalRow>
{
+
+ private final RowType projectedType;
+ private final @Nullable PartitionPredicate partitionFilter;
+ private final @Nullable BucketFilter bucketFilter;
+
+ private @Nullable RowIterator rows;
+ private boolean closed;
+
+ private RowsIterator(
+ RowType projectedType,
+ @Nullable PartitionPredicate partitionFilter,
+ @Nullable BucketFilter bucketFilter) {
+ this.projectedType = projectedType;
+ this.partitionFilter = partitionFilter;
+ this.bucketFilter = bucketFilter;
+ }
+
+ @Override
+ public boolean hasNext() {
+ if (closed) {
+ return false;
+ }
+ try {
+ while (rows == null || !rows.hasNext()) {
+ if (!ManifestAvroReader.this.hasNext()) {
return false;
}
- BorrowedBlock block = blockReader.nextBorrowedBlock();
- decoder.reset(block.bytes(), block.offset(),
block.length());
- recordsRemaining = block.recordCount();
+ rows =
+ ManifestAvroReader.this
+ .next()
+ .toRows(
+ projectedType,
+ partitionFilter,
+ bucketFilter,
+ readStatistics);
}
+ return true;
+ } catch (IOException e) {
+ throw new UncheckedIOException("Failed to decode Manifest Avro
block.", e);
+ }
+ }
- InternalRow candidate = recordDecoder.read(decoder);
- recordsRemaining--;
- if (candidate != null) {
- next = candidate;
- nextReady = true;
- return true;
- }
+ @Override
+ public InternalRow next() {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
}
- } catch (IOException e) {
- throw new UncheckedIOException("Failed to decode Manifest Avro
block.", e);
+ return rows.next();
}
- }
- @Override
- public InternalRow next() {
- if (!hasNext()) {
- throw new NoSuchElementException();
+ @Override
+ public void close() throws IOException {
+ closed = true;
+ ManifestAvroReader.this.close();
}
- InternalRow result = next;
- next = null;
- nextReady = false;
- return result;
}
- long decodedDataFiles() {
- return recordDecoder.decodedDataFiles;
+ /** Decoder state shared by the borrowed blocks produced by one reader. */
+ private static final class DecoderContext {
+
+ private final AvroRecordDecoder decoder;
+
+ private @Nullable ByteBuffer decompressionBuffer;
+ private RowType projectedRowType;
+ private ManifestEntryDecoder recordDecoder;
+
+ private DecoderContext(AvroRecordDecoder decoder) {
+ this.decoder = decoder;
+ }
+
+ private ByteBuffer decompress(AvroRawBlock block) throws IOException {
+ decompressionBuffer = block.decompress(decompressionBuffer);
+ return decompressionBuffer;
+ }
+
+ private ManifestEntryDecoder recordDecoder(RowType rowType) {
+ if (!rowType.equals(projectedRowType)) {
+ recordDecoder = new ManifestEntryDecoder(decoder, rowType);
+ projectedRowType = rowType;
+ }
+ return recordDecoder;
+ }
}
- long skippedDataFiles() {
- return recordDecoder.skippedDataFiles;
+ /** Iterator over the reusable row decoded from one borrowed block. */
+ public static final class RowIterator implements Iterator<BlockRow> {
+
+ private final AvroRecordDecoder decoder;
+ private final ManifestEntryDecoder recordDecoder;
+ private final BlockRow row;
+ private final @Nullable PartitionPredicate partitionFilter;
+ private final @Nullable BucketFilter bucketFilter;
+ private final @Nullable ReadStatistics readStatistics;
+ private final boolean filtered;
+
+ private long blockRemaining;
+ private long blockRecordIndex = -1;
+ private boolean nextReady;
+
+ private RowIterator(
+ long recordCount,
+ AvroRecordDecoder decoder,
+ ManifestEntryDecoder recordDecoder,
+ BlockRow row,
+ @Nullable PartitionPredicate partitionFilter,
+ @Nullable BucketFilter bucketFilter,
+ @Nullable ReadStatistics readStatistics) {
+ blockRemaining = recordCount;
+ this.decoder = decoder;
+ this.recordDecoder = recordDecoder;
+ this.row = row;
+ this.partitionFilter = partitionFilter;
+ this.bucketFilter = bucketFilter;
+ this.readStatistics = readStatistics;
+ this.filtered = partitionFilter != null || bucketFilter != null;
+ }
+
+ @Override
+ public boolean hasNext() {
+ try {
+ if (!filtered) {
+ ensureBlockFullyConsumed();
+ return blockRemaining > 0;
+ }
+ if (nextReady) {
+ return true;
+ }
+ while (blockRemaining > 0) {
+ blockRecordIndex++;
+ boolean selected =
+ recordDecoder.read(decoder, row, partitionFilter,
bucketFilter);
+ recordRead(selected);
+ blockRemaining--;
+ if (selected) {
+ nextReady = true;
+ return true;
+ }
+ }
+ ensureBlockFullyConsumed();
+ return false;
+ } catch (IOException e) {
+ throw new UncheckedIOException(
+ "Failed to decode projected Manifest Avro record.", e);
+ }
+ }
+
+ @Override
+ public BlockRow next() {
+ if (filtered) {
+ if (!hasNext()) {
+ throw new NoSuchElementException();
+ }
+ nextReady = false;
+ return row;
+ }
+ try {
+ if (blockRemaining == 0) {
+ throw new NoSuchElementException();
+ }
+ blockRecordIndex++;
+ recordDecoder.read(decoder, row, null, null);
+ recordRead(true);
+ blockRemaining--;
+ return row;
+ } catch (IOException e) {
+ throw new UncheckedIOException(
+ "Failed to decode projected Manifest Avro record.", e);
+ }
+ }
+
+ public long recordIndex() {
+ checkCurrentRecord();
+ return blockRecordIndex;
+ }
+
+ private void recordRead(boolean selected) {
+ if (readStatistics != null) {
+ if (selected && recordDecoder.fileDecoder != null) {
+ readStatistics.decodedDataFiles++;
+ } else {
+ readStatistics.skippedDataFiles++;
+ }
+ }
+ }
+
+ private void ensureBlockFullyConsumed() throws IOException {
+ if (blockRemaining == 0 && !decoder.isEnd()) {
+ throw new IOException("Manifest Avro block contains trailing
undecoded bytes.");
+ }
+ }
+
+ private void checkCurrentRecord() {
+ if (blockRecordIndex < 0) {
+ throw new IllegalStateException("No current Manifest Avro
record.");
+ }
+ }
}
- @Override
- public void close() throws IOException {
- next = null;
- nextReady = false;
- finished = true;
- blockReader.close();
+ private static final class ReadStatistics {
+
+ private long decodedDataFiles;
+ private long skippedDataFiles;
}
- private static class ManifestRecordDecoder {
+ /** Reusable {@link GenericRow} with borrowed views over fields in the
current Avro block. */
+ public static final class BlockRow extends GenericRow {
Review Comment:
This path is only for full projection.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
