This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit a27b548472658bb9311fe74fa4e088df44392029 Author: umi <[email protected]> AuthorDate: Fri Sep 18 12:55:49 2026 +0800 [format] Use buffered encoder for Avro writes (#9928) --- .../java/org/apache/paimon/format/FileFormat.java | 9 +- .../apache/paimon/format/FileFormatFactory.java | 25 ++++ .../apache/paimon/format/avro/AvroFileFormat.java | 21 ++- .../paimon/format/avro/AvroFileFormatTest.java | 153 +++++++++++++++++++++ 4 files changed, 205 insertions(+), 3 deletions(-) diff --git a/paimon-common/src/main/java/org/apache/paimon/format/FileFormat.java b/paimon-common/src/main/java/org/apache/paimon/format/FileFormat.java index f17c23b255..7a06e544ea 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/FileFormat.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/FileFormat.java @@ -74,6 +74,10 @@ public abstract class FileFormat { } public static FileFormat fromIdentifier(String identifier, Options options) { + return fromIdentifier(identifier, options, false); + } + + private static FileFormat fromIdentifier(String identifier, Options options, boolean manifest) { return fromIdentifier( normalizeFileFormat(identifier), new FormatContext( @@ -82,7 +86,8 @@ public abstract class FileFormat { options.get(CoreOptions.WRITE_BATCH_SIZE), options.get(CoreOptions.WRITE_BATCH_MEMORY), options.get(CoreOptions.FILE_COMPRESSION_ZSTD_LEVEL), - options.get(CoreOptions.FILE_BLOCK_SIZE))); + options.get(CoreOptions.FILE_BLOCK_SIZE), + manifest)); } /** Create a {@link FileFormat} from format identifier and format options. */ @@ -120,6 +125,6 @@ public abstract class FileFormat { // Manifest blocks retain the Avro default independently of data-file block sizing. Options manifestOptions = new Options(options.toMap()); manifestOptions.remove(CoreOptions.FILE_BLOCK_SIZE); - return FileFormat.fromIdentifier(CoreOptions.FILE_FORMAT_AVRO, manifestOptions); + return FileFormat.fromIdentifier(CoreOptions.FILE_FORMAT_AVRO, manifestOptions, true); } } diff --git a/paimon-common/src/main/java/org/apache/paimon/format/FileFormatFactory.java b/paimon-common/src/main/java/org/apache/paimon/format/FileFormatFactory.java index b5cba150fe..157f151ad2 100644 --- a/paimon-common/src/main/java/org/apache/paimon/format/FileFormatFactory.java +++ b/paimon-common/src/main/java/org/apache/paimon/format/FileFormatFactory.java @@ -41,6 +41,7 @@ public interface FileFormatFactory { private final MemorySize writeBatchMemory; private final int zstdLevel; @Nullable private final MemorySize blockSize; + private final boolean manifest; @VisibleForTesting public FormatContext(Options options, int readBatchSize, int writeBatchSize) { @@ -63,12 +64,31 @@ public interface FileFormatFactory { MemorySize writeBatchMemory, int zstdLevel, @Nullable MemorySize blockSize) { + this( + options, + readBatchSize, + writeBatchSize, + writeBatchMemory, + zstdLevel, + blockSize, + false); + } + + public FormatContext( + Options options, + int readBatchSize, + int writeBatchSize, + MemorySize writeBatchMemory, + int zstdLevel, + @Nullable MemorySize blockSize, + boolean manifest) { this.options = options; this.readBatchSize = readBatchSize; this.writeBatchSize = writeBatchSize; this.writeBatchMemory = writeBatchMemory; this.zstdLevel = zstdLevel; this.blockSize = blockSize; + this.manifest = manifest; } public Options options() { @@ -96,6 +116,11 @@ public interface FileFormatFactory { return blockSize; } + /** Whether the format is used for manifest metadata rather than data files. */ + public boolean isManifest() { + return manifest; + } + public boolean caseSensitive() { return options.getOptional(CatalogOptions.CASE_SENSITIVE).orElse(true); } diff --git a/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroFileFormat.java b/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroFileFormat.java index a23c66af21..845cce65d5 100644 --- a/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroFileFormat.java +++ b/paimon-format/src/main/java/org/apache/paimon/format/avro/AvroFileFormat.java @@ -39,6 +39,7 @@ import org.apache.paimon.types.RowType; import org.apache.avro.Schema; import org.apache.avro.file.CodecFactory; import org.apache.avro.file.DataFileWriter; +import org.apache.avro.io.EncoderFactory; import javax.annotation.Nullable; @@ -67,6 +68,7 @@ public class AvroFileFormat extends FileFormat { private final Options options; private final int zstdLevel; @Nullable private final MemorySize blockSize; + private final boolean useBufferedEncoder; public AvroFileFormat(FormatContext context) { super(IDENTIFIER); @@ -74,6 +76,7 @@ public class AvroFileFormat extends FileFormat { this.options = getIdentifierPrefixOptions(context.options()); this.zstdLevel = context.zstdLevel(); this.blockSize = context.blockSize(); + this.useBufferedEncoder = !context.isManifest(); } @Override @@ -91,10 +94,26 @@ public class AvroFileFormat extends FileFormat { public AvroBlockWriter createBlockWriter( PositionOutputStream out, RowType rowType, String compression) throws IOException { + // Retain Avro's direct encoder for pre-encoded manifest records. The buffered encoder + // copies array-backed ByteBuffers into a temporary byte array for each appendEncoded call. + return createBlockWriter(out, rowType, compression, false); + } + + private AvroBlockWriter createBlockWriter( + PositionOutputStream out, + RowType rowType, + String compression, + boolean useBufferedEncoder) + throws IOException { Schema schema = AvroSchemaConverter.convertToSchema(rowType, options.get(AVRO_ROW_NAME_MAPPING)); AvroRowDatumWriter datumWriter = new AvroRowDatumWriter(rowType); DataFileWriter<InternalRow> writer = new DataFileWriter<>(datumWriter); + if (useBufferedEncoder) { + // Batch data-file field encodings before writing them to the Avro block buffer. + writer.setEncoder( + outputStream -> EncoderFactory.get().binaryEncoder(outputStream, null)); + } writer.setCodec(createCodecFactory(compression)); if (blockSize != null) { writer.setSyncInterval(Math.toIntExact(blockSize.getBytes())); @@ -141,7 +160,7 @@ public class AvroFileFormat extends FileFormat { @Override public FormatWriter create(PositionOutputStream out, String compression) throws IOException { - return createBlockWriter(out, rowType, compression); + return createBlockWriter(out, rowType, compression, useBufferedEncoder); } } } diff --git a/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroFileFormatTest.java b/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroFileFormatTest.java index 70780f8d9b..251f3ca974 100644 --- a/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroFileFormatTest.java +++ b/paimon-format/src/test/java/org/apache/paimon/format/avro/AvroFileFormatTest.java @@ -44,7 +44,9 @@ import org.apache.avro.generic.GenericDatumWriter; import org.apache.avro.generic.GenericRecord; import org.apache.avro.io.BinaryDecoder; import org.apache.avro.io.BinaryEncoder; +import org.apache.avro.io.BufferedBinaryEncoder; import org.apache.avro.io.DecoderFactory; +import org.apache.avro.io.DirectBinaryEncoder; import org.apache.avro.io.EncoderFactory; import org.junit.jupiter.api.BeforeAll; import org.junit.jupiter.api.Test; @@ -176,6 +178,157 @@ public class AvroFileFormatTest { } } + @ParameterizedTest + @ValueSource(strings = {"null", "deflate", "snappy", "zstd"}) + void testCloseFlushesBufferedRecords(String compression) throws IOException { + RowType rowType = DataTypes.ROW(DataTypes.INT().notNull()).notNull(); + LocalFileIO fileIO = LocalFileIO.create(); + Path file = new Path(new Path(tempPath.toUri()), UUID.randomUUID().toString()); + + try (PositionOutputStream out = fileIO.newOutputStream(file, false); + FormatWriter writer = + fileFormat.createWriterFactory(rowType).create(out, compression)) { + assertThat(writer) + .extracting("writer") + .extracting("bufOut") + .isInstanceOf(BufferedBinaryEncoder.class); + // Leave a partial record-encoder buffer as well as a partial Avro block. + writer.addElement(GenericRow.of(42)); + } + + List<Integer> result = new ArrayList<>(); + try (RecordReader<InternalRow> reader = + fileFormat + .createReaderFactory(rowType, rowType, new ArrayList<>()) + .createReader( + new FormatReaderContext( + fileIO, file, fileIO.getFileSize(file), null, null))) { + reader.forEachRemaining(row -> result.add(row.getInt(0))); + } + assertThat(result).containsExactly(42); + } + + @ParameterizedTest + @CsvSource({ + "null, 16", + "null, 4096", + "deflate, 16", + "deflate, 4096", + "snappy, 16", + "snappy, 4096", + "zstd, 16", + "zstd, 4096" + }) + void testBufferedRecordSerializationFailureRollsBack(String compression, int payloadSize) + throws IOException { + RowType rowType = + DataTypes.ROW(DataTypes.BYTES().notNull(), DataTypes.INT().notNull()).notNull(); + LocalFileIO fileIO = LocalFileIO.create(); + Path file = new Path(new Path(tempPath.toUri()), UUID.randomUUID().toString()); + GenericRow before = GenericRow.of(new byte[] {1, 2, 3}, 42); + GenericRow after = GenericRow.of(new byte[] {4, 5, 6}, 43); + byte[] failedPayload = new byte[payloadSize]; + Arrays.fill(failedPayload, (byte) 7); + + try (PositionOutputStream out = fileIO.newOutputStream(file, false); + FormatWriter writer = + fileFormat.createWriterFactory(rowType).create(out, compression)) { + assertThat(writer) + .extracting("writer") + .extracting("bufOut") + .isInstanceOf(BufferedBinaryEncoder.class); + writer.addElement(before); + // The first field is written before the second field fails. Small payloads stay in + // the encoder buffer; large payloads flush it and reach the Avro block buffer. + assertThatThrownBy(() -> writer.addElement(GenericRow.of(failedPayload, "not an int"))) + .isInstanceOf(DataFileWriter.AppendWriteException.class) + .hasCauseInstanceOf(ClassCastException.class); + writer.addElement(after); + } + + List<GenericRow> result = new ArrayList<>(); + try (RecordReader<InternalRow> reader = + fileFormat + .createReaderFactory(rowType, rowType, new ArrayList<>()) + .createReader( + new FormatReaderContext( + fileIO, file, fileIO.getFileSize(file), null, null))) { + reader.forEachRemaining( + row -> result.add(GenericRow.of(row.getBinary(0), row.getInt(1)))); + } + assertThat(result).containsExactly(before, after); + } + + @ParameterizedTest + @ValueSource(strings = {"avro", "parquet", "orc"}) + void testManifestWriterUsesDirectEncoder(String dataFileFormat) throws IOException { + Options options = new Options(); + options.set(CoreOptions.FILE_FORMAT, dataFileFormat); + FileFormat manifestFormat = FileFormat.manifestFormat(new CoreOptions(options)); + RowType rowType = DataTypes.ROW(DataTypes.INT().notNull()).notNull(); + LocalFileIO fileIO = LocalFileIO.create(); + Path file = new Path(new Path(tempPath.toUri()), UUID.randomUUID().toString()); + + try (PositionOutputStream out = fileIO.newOutputStream(file, false); + FormatWriter writer = + manifestFormat.createWriterFactory(rowType).create(out, "null")) { + assertThat(writer) + .extracting("writer") + .extracting("bufOut") + .isInstanceOf(DirectBinaryEncoder.class); + writer.addElement(GenericRow.of(42)); + } + } + + @ParameterizedTest + @ValueSource(strings = {"null", "deflate", "snappy", "zstd"}) + void testRecordsAroundRawBlockCopy(String compression) throws IOException { + RowType rowType = DataTypes.ROW(DataTypes.INT().notNull()).notNull(); + LocalFileIO fileIO = LocalFileIO.create(); + Path source = new Path(new Path(tempPath.toUri()), UUID.randomUUID().toString()); + Path target = new Path(new Path(tempPath.toUri()), UUID.randomUUID().toString()); + + try (PositionOutputStream out = fileIO.newOutputStream(source, false); + FormatWriter writer = + fileFormat.createWriterFactory(rowType).create(out, "deflate")) { + writer.addElement(GenericRow.of(2)); + writer.addElement(GenericRow.of(3)); + } + + ByteArrayOutputStream encoded = new ByteArrayOutputStream(); + BinaryEncoder encoder = EncoderFactory.get().binaryEncoder(encoded, null); + encoder.writeInt(1); + encoder.flush(); + + try (AvroBlockReader reader = new AvroBlockReader(fileIO.newInputStream(source)); + PositionOutputStream out = fileIO.newOutputStream(target, false); + AvroBlockWriter writer = fileFormat.createBlockWriter(out, rowType, compression)) { + assertThat(writer) + .extracting("writer") + .extracting("bufOut") + .isInstanceOf(DirectBinaryEncoder.class); + writer.addElement(GenericRow.of(0)); + ByteBuffer record = ByteBuffer.wrap(encoded.toByteArray()); + writer.addEncoded(record); + assertThat(record.position()).isZero(); + assertThat(record.remaining()).isEqualTo(encoded.size()); + // Copying a raw block must first flush pending rows and encoded records. + writer.addEncodedBlock(reader.nextBorrowedRawBlock()); + writer.addElement(GenericRow.of(4)); + } + + List<Integer> result = new ArrayList<>(); + try (RecordReader<InternalRow> reader = + fileFormat + .createReaderFactory(rowType, rowType, new ArrayList<>()) + .createReader( + new FormatReaderContext( + fileIO, target, fileIO.getFileSize(target), null, null))) { + reader.forEachRemaining(row -> result.add(row.getInt(0))); + } + assertThat(result).containsExactly(0, 1, 2, 3, 4); + } + @Test public void testSupportedDataTypes() { ArrayList<DataField> dataFields = new ArrayList<>();
