This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new ef01353569 [python] Optimize BLOB write performance (#8625)
ef01353569 is described below
commit ef0135356952e4d3b5a614fbbce2aa5cd6d52ada
Author: umi <[email protected]>
AuthorDate: Thu Jul 16 22:12:12 2026 +0800
[python] Optimize BLOB write performance (#8625)
Improve Python BLOB write performance by reducing CRC and Arrow
row-selection overhead.
- Use `isal_zlib.crc32` when available, with `zlib` as the fallback.
- Reuse the original Arrow batch when all rows belong to one
partition/bucket.
- Use zero-copy Arrow `slice` for contiguous partition/bucket groups.
- Keep Arrow `take` for non-contiguous groups.
- Add coverage for CRC compatibility and all three Arrow grouping paths.
A 60 GB BLOB write test improved as follows:
| Data size | Before | After | Elapsed-time reduction |
| --- | ---: | ---: | ---: |
| 60 GB | 178.984 s | 132.871 s | 25.76% |
The elapsed time decreased by 46.113 seconds, from 178.984 seconds to
132.871 seconds.
---
.github/workflows/paimon-python-checks.yml | 2 +-
paimon-python/dev/requirements.txt | 1 +
paimon-python/pypaimon/tests/blob_test.py | 32 ++++++++++
.../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 ++++-
6 files changed, 120 insertions(+), 5 deletions(-)
diff --git a/.github/workflows/paimon-python-checks.yml
b/.github/workflows/paimon-python-checks.yml
index 5c6eb99cf4..8e92a9c1fb 100755
--- a/.github/workflows/paimon-python-checks.yml
+++ b/.github/workflows/paimon-python-checks.yml
@@ -128,7 +128,7 @@ jobs:
else
python -m pip install --upgrade pip
pip install torch --index-url https://download.pytorch.org/whl/cpu
- python -m pip install pyroaring readerwriterlock==1.0.9
fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0
fastavro==1.11.1 pyarrow==16.0.0 zstandard==0.24.0 polars==1.32.0 duckdb==1.3.2
numpy==1.24.3 pandas==2.0.3 pylance==0.39.0 cramjam flake8==4.0.1 pytest~=7.0
py4j==0.10.9.9 requests parameterized==0.9.0 'daft>=0.7.6' pypaimon-rust==0.2.0
'datafusion>=52'
+ python -m pip install pyroaring readerwriterlock==1.0.9
fsspec==2024.3.1 cachetools==5.3.3 ossfs==2023.12.0 ray==2.54.0
fastavro==1.11.1 'isal>=1.8,<2' pyarrow==16.0.0 zstandard==0.24.0
polars==1.32.0 duckdb==1.3.2 numpy==1.24.3 pandas==2.0.3 pylance==0.39.0
cramjam flake8==4.0.1 pytest~=7.0 py4j==0.10.9.9 requests parameterized==0.9.0
'daft>=0.7.6' pypaimon-rust==0.2.0 'datafusion>=52'
python -m pip install 'lumina-data>=${{ env.LUMINA_DATA_VERSION
}}' -i https://pypi.org/simple/
if python -c "import sys; sys.exit(0 if sys.version_info >= (3,
11) else 1)"; then
python -m pip install vortex-data==0.70.0
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 a743442b5b..9da03898e4 100644
--- a/paimon-python/pypaimon/tests/blob_test.py
+++ b/paimon-python/pypaimon/tests/blob_test.py
@@ -21,6 +21,7 @@ import shutil
import struct
import tempfile
import unittest
+import zlib
from pathlib import Path
from unittest.mock import patch
@@ -1893,6 +1894,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 ee2f75bec4..d078304d27 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 ad88295b64..bfcb4c4f3a 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):
