This is an automated email from the ASF dual-hosted git repository. discivigour pushed a commit to branch feat/improveWritePerf in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 72e64d381f7865a720f46479f273201b057a1740 Author: umi <[email protected]> AuthorDate: Tue Jul 14 15:31:06 2026 +0800 [python] Optimize BLOB write performance --- paimon-python/dev/requirements.txt | 1 + paimon-python/pypaimon/tests/blob_test.py | 33 +++++++++++ .../pypaimon/tests/write/table_write_test.py | 68 ++++++++++++++++++++++ paimon-python/pypaimon/write/blob_format_writer.py | 8 ++- paimon-python/pypaimon/write/table_write.py | 14 ++++- 5 files changed, 120 insertions(+), 4 deletions(-) diff --git a/paimon-python/dev/requirements.txt b/paimon-python/dev/requirements.txt index 0bd9f01b76..3415427248 100644 --- a/paimon-python/dev/requirements.txt +++ b/paimon-python/dev/requirements.txt @@ -19,6 +19,7 @@ cachetools>=4.2,<6; python_version=="3.6" cachetools>=5,<6; python_version>"3.6" dataclasses>=0.8; python_version < "3.7" fastavro>=1.4,<2 +isal>=1.8,<2; python_version >= "3.9" and (platform_machine == "x86_64" or platform_machine == "AMD64" or platform_machine == "aarch64" or platform_machine == "arm64") fsspec>=2021.10,<2026; python_version<"3.8" fsspec>=2023,<2026; python_version>="3.8" packaging>=21,<26 diff --git a/paimon-python/pypaimon/tests/blob_test.py b/paimon-python/pypaimon/tests/blob_test.py index 86c7b5daa0..33fc5cbdd9 100644 --- a/paimon-python/pypaimon/tests/blob_test.py +++ b/paimon-python/pypaimon/tests/blob_test.py @@ -21,7 +21,9 @@ import shutil import struct import tempfile import unittest +import zlib from pathlib import Path +from unittest.mock import patch import pyarrow as pa @@ -1599,6 +1601,37 @@ class BlobEndToEndTest(unittest.TestCase): self.assertEqual(writer.lengths, [-1]) self.assertEqual(writer.position, 0) + @staticmethod + def _write_blob_record_with_crc_backend(backend): + from pypaimon.write import blob_format_writer + + output = io.BytesIO() + payload = b'blob-crc-payload' * 1024 + with patch.object(blob_format_writer, 'crc_backend', backend): + writer = blob_format_writer.BlobFormatWriter(output) + writer.add_blob('blob_field', BlobData(payload)) + return output.getvalue(), payload + + def test_blob_crc_fallback_matches_zlib(self): + from pypaimon.write.blob_format_writer import BlobFormatWriter + + record, payload = self._write_blob_record_with_crc_backend(zlib) + expected_crc = zlib.crc32( + struct.pack('<I', BlobFormatWriter.MAGIC_NUMBER)) + expected_crc = zlib.crc32(payload, expected_crc) & 0xffffffff + actual_crc = struct.unpack('<I', record[-4:])[0] + self.assertEqual(expected_crc, actual_crc) + + def test_blob_crc_isal_matches_zlib(self): + try: + from isal import isal_zlib + except ImportError: + self.skipTest('isal is not available on this platform') + + zlib_record, _ = self._write_blob_record_with_crc_backend(zlib) + isal_record, _ = self._write_blob_record_with_crc_backend(isal_zlib) + self.assertEqual(zlib_record, isal_record) + def test_null_blob_read(self): from pypaimon.write.blob_format_writer import BlobFormatWriter diff --git a/paimon-python/pypaimon/tests/write/table_write_test.py b/paimon-python/pypaimon/tests/write/table_write_test.py index f6eb97c5ee..e1fccdec08 100644 --- a/paimon-python/pypaimon/tests/write/table_write_test.py +++ b/paimon-python/pypaimon/tests/write/table_write_test.py @@ -22,6 +22,7 @@ import shutil import tempfile import unittest +from unittest.mock import Mock, patch from pypaimon import CatalogFactory, Schema import pyarrow as pa @@ -30,6 +31,7 @@ from parameterized import parameterized from pypaimon.common.json_util import JSON from pypaimon.common.options.core_options import CoreOptions from pypaimon.manifest.manifest_list_manager import ManifestListManager +from pypaimon.write.table_write import TableWrite from pypaimon.write.writer.append_only_data_writer import AppendOnlyDataWriter @@ -89,6 +91,72 @@ class TableWriteTest(unittest.TestCase): return read_builder.new_read().to_arrow( read_builder.new_scan().plan().splits()).sort_by(sort_keys) + @staticmethod + def _mock_table_write(partitions, buckets): + table_write = object.__new__(TableWrite) + table_write._validate_pyarrow_schema = Mock() + table_write.row_key_extractor = Mock() + table_write.file_store_write = Mock() + table_write.row_key_extractor.extract_partition_bucket_batch.return_value = ( + partitions, buckets) + return table_write + + def test_write_arrow_batch_reuses_full_batch(self): + data = pa.RecordBatch.from_pydict({ + 'id': [0, 1], + 'payload': [b'a', b'b'], + }) + table_write = self._mock_table_write( + [('p1',), ('p1',)], [0, 0]) + + with patch.object(pa.compute, 'take', wraps=pa.compute.take) as take: + table_write.write_arrow_batch(data) + + take.assert_not_called() + written = table_write.file_store_write.write.call_args[0][2] + self.assertIs(data, written) + + def test_write_arrow_batch_uses_zero_copy_for_contiguous_groups(self): + data = pa.RecordBatch.from_pydict({ + 'id': [0, 1, 2, 3], + 'payload': [b'a', b'b', b'c', b'd'], + }) + table_write = self._mock_table_write( + [('p1',), ('p1',), ('p2',), ('p2',)], + [0, 0, 1, 1]) + with patch.object(pa.compute, 'take', wraps=pa.compute.take) as take: + table_write.write_arrow_batch(data) + + take.assert_not_called() + calls = table_write.file_store_write.write.call_args_list + self.assertEqual(2, len(calls)) + self.assertEqual({'id': [0, 1], 'payload': [b'a', b'b']}, + calls[0][0][2].to_pydict()) + self.assertEqual({'id': [2, 3], 'payload': [b'c', b'd']}, + calls[1][0][2].to_pydict()) + self.assertEqual( + data.column(1).buffers()[2].address, + calls[0][0][2].column(1).buffers()[2].address) + + def test_write_arrow_batch_uses_take_for_non_contiguous_groups(self): + data = pa.RecordBatch.from_pydict({ + 'id': [0, 1, 2, 3], + 'payload': [b'a', b'b', b'c', b'd'], + }) + table_write = self._mock_table_write( + [('p1',), ('p2',), ('p1',), ('p2',)], + [0, 1, 0, 1]) + + with patch.object(pa.compute, 'take', wraps=pa.compute.take) as take: + table_write.write_arrow_batch(data) + + self.assertEqual(2, take.call_count) + calls = table_write.file_store_write.write.call_args_list + self.assertEqual({'id': [0, 2], 'payload': [b'a', b'c']}, + calls[0][0][2].to_pydict()) + self.assertEqual({'id': [1, 3], 'payload': [b'b', b'd']}, + calls[1][0][2].to_pydict()) + def test_write_snapshot(self): schema = Schema.from_pyarrow_schema(self.pa_schema, partition_keys=['dt']) self.catalog.create_table('default.test_write_snapshot', schema, False) diff --git a/paimon-python/pypaimon/write/blob_format_writer.py b/paimon-python/pypaimon/write/blob_format_writer.py index e85eebccf6..36c1044443 100644 --- a/paimon-python/pypaimon/write/blob_format_writer.py +++ b/paimon-python/pypaimon/write/blob_format_writer.py @@ -16,9 +16,13 @@ # under the License. import struct -import zlib from typing import BinaryIO, List, Optional +try: + from isal import isal_zlib as crc_backend +except ImportError: + import zlib as crc_backend + from pypaimon.schema.data_types import is_array_blob_type, is_blob_file_type from pypaimon.table.row.blob import Blob, BlobData, BlobDescriptor, BlobConsumer from pypaimon.common.delta_varint_compressor import DeltaVarintCompressor @@ -182,7 +186,7 @@ class BlobFormatWriter: return blob_pos, self.position - blob_pos, crc32 def _write_with_crc(self, data: bytes, crc32: int) -> int: - crc32 = zlib.crc32(data, crc32) + crc32 = crc_backend.crc32(data, crc32) self.output_stream.write(data) self.position += len(data) return crc32 diff --git a/paimon-python/pypaimon/write/table_write.py b/paimon-python/pypaimon/write/table_write.py index 629383409f..59846bd59e 100644 --- a/paimon-python/pypaimon/write/table_write.py +++ b/paimon-python/pypaimon/write/table_write.py @@ -56,8 +56,18 @@ class TableWrite: partition_bucket_groups[(tuple(partitions[i]), buckets[i])].append(i) for (partition, bucket), row_indices in partition_bucket_groups.items(): - indices_array = pa.array(row_indices, type=pa.int64()) - sub_table = pa.compute.take(data, indices_array) + if len(row_indices) == data.num_rows: + # Every input row belongs to the same partition/bucket. Passing the + # original batch through avoids copying large BLOB values through + # Arrow take before the dedicated BLOB writer consumes them. + sub_table = data + elif row_indices[-1] - row_indices[0] + 1 == len(row_indices): + # Contiguous groups can share the original Arrow buffers instead of + # gathering their rows into newly allocated buffers with take. + sub_table = data.slice(row_indices[0], len(row_indices)) + else: + indices_array = pa.array(row_indices, type=pa.int64()) + sub_table = pa.compute.take(data, indices_array) self.file_store_write.write(partition, bucket, sub_table) def write_row(self, row):
