[python] Optimize BLOB write performance

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

paimon-python/pypaimon/tests/blob_test.py

@@ -21,7 +21,9 @@
 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 @@ def test_null_blob_write(self):
         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
 

paimon-python/pypaimon/tests/write/table_write_test.py

@@ -22,6 +22,7 @@
 
 import tempfile
 import unittest
+from unittest.mock import Mock, patch
 
 from pypaimon import CatalogFactory, Schema
 import pyarrow as pa
@@ -30,6 +31,7 @@
 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 @@ def _read_sorted(table, sort_keys):
         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)

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 @@ def _write_blob_data(self, blob_value: Blob, crc32: int):
         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

paimon-python/pypaimon/write/table_write.py

@@ -56,8 +56,18 @@ def write_arrow_batch(self, data: pa.RecordBatch):
             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):