[python] Normalize inline BLOB rows before postpone bucket planning

Author: zhuanshenbsj1

paimon-python/pypaimon/execution/python/write/file_store_write.py

@@ -25,7 +25,7 @@
 from pypaimon.schema.data_types import is_blob_file_field, is_blob_type
 from pypaimon.write.commit_message import CommitMessage
 from pypaimon.write.row_utils import (
-    inline_blob_value,
+    normalize_inline_blob_row,
     require_columns,
     row_values_to_arrow_table,
 )
@@ -135,14 +135,14 @@ def write_row(
             descriptor_fields = CoreOptions.blob_descriptor_fields(self.options)
             view_fields = CoreOptions.blob_view_fields(self.options)
             inline_columns = [name for name in column_names if name in descriptor_fields or name in view_fields]
-            if inline_columns:
-                values_by_name = dict(values_by_name)
-                for name in inline_columns:
-                    values_by_name[name] = inline_blob_value(
-                        values_by_name[name], name in descriptor_fields, name in view_fields)
-                from pypaimon.table.row.generic_row import GenericRow
-                row = GenericRow([values_by_name[name] for name in column_names],
-                                 [self.table.field_dict[name] for name in column_names], row.get_row_kind())
+            row, values_by_name = normalize_inline_blob_row(
+                row,
+                values_by_name,
+                self.table.field_dict,
+                column_names,
+                descriptor_fields,
+                view_fields,
+            )
             if not self.table.is_primary_key_table:
                 if writer is None:
                     ordinary_columns = [

paimon-python/pypaimon/execution/python/write/postpone_batch_table_write.py

@@ -18,12 +18,14 @@
 
 import pyarrow as pa
 
+from pypaimon.common.options.core_options import CoreOptions
 from pypaimon.write.commit_message import CommitMessage
 from pypaimon.execution.python.write.file_store_write import PostponeFixedBucketFileStoreWrite
 from pypaimon.execution.python.write.postpone_bucket import PythonPostponeBucketPlanner as PostponeBucketPlanner
 from pypaimon.execution.python.write.table_write import PythonTableWrite
 from pypaimon.execution.python.write.row_key_extractor import PostponeFixedBucketRowKeyExtractor
 from pypaimon.write.row_utils import (
+    normalize_inline_blob_row,
     require_columns,
     row_to_named_values,
     row_values_to_arrow_table,
@@ -129,6 +131,14 @@ def write_row(self, row):
         if self._bucket_plan.contains(partition):
             return super().write_row(row)
 
+        row, values_by_name = normalize_inline_blob_row(
+            row,
+            values_by_name,
+            self.table.field_dict,
+            column_names,
+            CoreOptions.blob_descriptor_fields(self.table.options),
+            CoreOptions.blob_view_fields(self.table.options),
+        )
         arrow_row = row_values_to_arrow_table(
             values_by_name, self.table.table_schema.fields, column_names)
         size = self._planner.input_partition_stats(arrow_row)[partition][1]

paimon-python/pypaimon/tests/blob_table_test.py

@@ -6027,6 +6027,149 @@ def test_primary_key_managed_blob_requires_native_writer(self):
             }, schema=pa_schema))
         writer.close()
 
+    def test_postpone_fixed_bucket_pk_inline_blob_write_row_new_partition(self):
+        from pypaimon import Schema
+        from pypaimon.table.row.blob import Blob, BlobDescriptor, BlobViewStruct
+        from pypaimon.table.row.generic_row import GenericRow
+
+        pk_pa_schema = pa.schema([
+            pa.field('pt', pa.string(), nullable=False),
+            pa.field('id', pa.int32(), nullable=False),
+            pa.field('picture', pa.large_binary()),
+        ])
+        with self.subTest(field='blob-descriptor-field'):
+            blob_file = os.path.join(self.temp_dir, 'postpone_pk_descriptor.bin')
+            with open(blob_file, 'wb') as f:
+                f.write(b'postpone-desc-0postpone-desc-1')
+            desc_0 = BlobDescriptor(blob_file, 0, 15)
+            desc_1 = BlobDescriptor(blob_file, 15, 15)
+
+            desc_schema = Schema.from_pyarrow_schema(
+                pk_pa_schema,
+                primary_keys=['pt', 'id'],
+                partition_keys=['pt'],
+                options={
+                    'bucket': '-2',
+                    'postpone.batch-write-fixed-bucket': 'true',
+                    'blob-descriptor-field': 'picture',
+                    'write.native.enabled': 'false',
+                },
+            )
+            self.catalog.create_table('test_db.pk_postpone_blob_descriptor', desc_schema, False)
+            desc_table = self.catalog.get_table('test_db.pk_postpone_blob_descriptor')
+
+            desc_wb = desc_table.new_postpone_fixed_bucket_write_builder()
+            desc_writer = desc_wb.new_write()
+            try:
+                desc_writer.write_row(GenericRow(
+                    [
+                        'p0',
+                        1,
+                        Blob.from_descriptor_bytes(
+                            desc_0.serialize(), file_io=desc_table.file_io
+                        ),
+                    ],
+                    desc_table.fields,
+                ))
+                desc_writer.write_row(GenericRow(
+                    [
+                        'p1',
+                        2,
+                        Blob.from_descriptor_bytes(
+                            desc_1.serialize(), file_io=desc_table.file_io
+                        ),
+                    ],
+                    desc_table.fields,
+                ))
+                messages = desc_writer.prepare_commit()
+                self.assertEqual({message.partition for message in messages}, {('p0',), ('p1',)})
+                for message in messages:
+                    self.assertEqual(message.total_buckets, 1)
+                desc_wb.new_commit().commit(messages)
+            finally:
+                desc_writer.close()
+
+            desc_rb = desc_table.new_read_builder()
+            desc_result = desc_rb.new_read().to_arrow(
+                desc_rb.new_scan().plan().splits()
+            ).sort_by('id')
+            self.assertEqual(desc_result.column('pt').to_pylist(), ['p0', 'p1'])
+            self.assertEqual(desc_result.column('id').to_pylist(), [1, 2])
+            self.assertEqual(
+                desc_result.column('picture').to_pylist(),
+                [b'postpone-desc-0', b'postpone-desc-1'],
+            )
+
+        with self.subTest(field='blob-view-field'):
+            source_schema = pa.schema([
+                ('id', pa.int32()),
+                ('picture', pa.large_binary()),
+            ])
+            source = Schema.from_pyarrow_schema(
+                source_schema,
+                options={
+                    'row-tracking.enabled': 'true',
+                    'data-evolution.enabled': 'true',
+                },
+            )
+            self.catalog.create_table('test_db.pk_postpone_view_source', source, False)
+            source_table = self.catalog.get_table('test_db.pk_postpone_view_source')
+            src_wb = source_table.new_batch_write_builder()
+            src_writer = src_wb.new_write()
+            src_writer.write_arrow(pa.Table.from_pydict({
+                'id': [10, 20],
+                'picture': [b'postpone-view-0', b'postpone-view-1'],
+            }, schema=source_schema))
+            src_wb.new_commit().commit(src_writer.prepare_commit())
+            src_writer.close()
+
+            view_0 = BlobViewStruct.of('test_db.pk_postpone_view_source', source_table, 'picture', 0)
+            view_1 = BlobViewStruct.of('test_db.pk_postpone_view_source', source_table, 'picture', 1)
+
+            view_schema = Schema.from_pyarrow_schema(
+                pk_pa_schema,
+                primary_keys=['pt', 'id'],
+                partition_keys=['pt'],
+                options={
+                    'bucket': '-2',
+                    'postpone.batch-write-fixed-bucket': 'true',
+                    'blob-view-field': 'picture',
+                    'write.native.enabled': 'false',
+                },
+            )
+            self.catalog.create_table('test_db.pk_postpone_blob_view', view_schema, False)
+            view_table = self.catalog.get_table('test_db.pk_postpone_blob_view')
+
+            view_wb = view_table.new_postpone_fixed_bucket_write_builder()
+            view_writer = view_wb.new_write()
+            try:
+                view_writer.write_row(GenericRow(
+                    ['p0', 1, Blob.from_view(view_0)],
+                    view_table.fields,
+                ))
+                view_writer.write_row(GenericRow(
+                    ['p1', 2, Blob.from_view(view_1)],
+                    view_table.fields,
+                ))
+                messages = view_writer.prepare_commit()
+                self.assertEqual({message.partition for message in messages}, {('p0',), ('p1',)})
+                for message in messages:
+                    self.assertEqual(message.total_buckets, 1)
+                view_wb.new_commit().commit(messages)
+            finally:
+                view_writer.close()
+
+            view_rb = view_table.new_read_builder()
+            view_result = view_rb.new_read().to_arrow(
+                view_rb.new_scan().plan().splits()
+            ).sort_by('id')
+            self.assertEqual(view_result.column('pt').to_pylist(), ['p0', 'p1'])
+            self.assertEqual(view_result.column('id').to_pylist(), [1, 2])
+            self.assertEqual(
+                view_result.column('picture').to_pylist(),
+                [b'postpone-view-0', b'postpone-view-1'],
+            )
+
 
 class GetBlobTest(unittest.TestCase):
 

paimon-python/pypaimon/write/row_utils.py

@@ -90,6 +90,38 @@ def inline_blob_value(value: Any, descriptor: bool, view: bool) -> Any:
     return value
 
 
+def normalize_inline_blob_row(
+        row: InternalRow,
+        values_by_name: Dict[str, Any],
+        field_dict: Dict[str, DataField],
+        column_names: List[str],
+        descriptor_fields: Iterable[str],
+        view_fields: Iterable[str]):
+    """Normalize inline BLOB reference fields in a row and value map to bytes."""
+    descriptor_set = set(descriptor_fields)
+    view_set = set(view_fields)
+    inline_columns = [
+        name for name in column_names
+        if name in descriptor_set or name in view_set
+    ]
+    if not inline_columns:
+        return row, values_by_name
+    values_by_name = dict(values_by_name)
+    for name in inline_columns:
+        values_by_name[name] = inline_blob_value(
+            values_by_name[name],
+            name in descriptor_set,
+            name in view_set,
+        )
+    from pypaimon.table.row.generic_row import GenericRow
+    row = GenericRow(
+        [values_by_name[name] for name in column_names],
+        [field_dict[name] for name in column_names],
+        row.get_row_kind(),
+    )
+    return row, values_by_name
+
+
 def value_for_arrow(value: Any, field: DataField) -> Any:
     if isinstance(value, Vector):
         return value.to_list()