[python] Support target file row num for blob and vector writes

paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py

@@ -24,6 +24,8 @@
 import pyarrow as pa
 
 from pypaimon import CatalogFactory, Schema
+from pypaimon.common.uri_reader import FileUriReader
+from pypaimon.table.row.blob import Blob
 
 
 class DataEvolutionRowRollingTest(unittest.TestCase):
@@ -34,6 +36,19 @@ class DataEvolutionRowRollingTest(unittest.TestCase):
         ('id', pa.int32()),
         ('name', pa.string()),
     ])
+    blob_schema = pa.schema([
+        ('id', pa.int32()),
+        ('payload', pa.large_binary()),
+    ])
+    vector_schema = pa.schema([
+        ('id', pa.int32()),
+        ('embedding', pa.list_(pa.float32(), 3)),
+    ])
+    blob_vector_schema = pa.schema([
+        ('id', pa.int32()),
+        ('payload', pa.large_binary()),
+        ('embedding', pa.list_(pa.float32(), 3)),
+    ])
     de_options = {
         'row-tracking.enabled': 'true',
         'data-evolution.enabled': 'true',
@@ -57,11 +72,49 @@ def _create(self, options):
             False)
         return self.catalog.get_table(name)
 
+    def _create_with_schema(self, pa_schema, options):
+        name = f'default.roll_{uuid.uuid4().hex[:8]}'
+        self.catalog.create_table(
+            name, Schema.from_pyarrow_schema(pa_schema, options=options),
+            False)
+        return self.catalog.get_table(name)
+
     def _rows(self, n):
         return pa.Table.from_pydict(
             {'id': list(range(n)), 'name': [f'n{i}' for i in range(n)]},
             schema=self.pa_schema)
 
+    def _blob_rows(self, n):
+        return pa.Table.from_pydict(
+            {
+                'id': list(range(n)),
+                'payload': [f'blob-{i}'.encode() for i in range(n)],
+            },
+            schema=self.blob_schema)
+
+    def _vector_rows(self, n):
+        return pa.Table.from_pydict(
+            {
+                'id': list(range(n)),
+                'embedding': [
+                    [float(i), float(i + 1), float(i + 2)]
+                    for i in range(n)
+                ],
+            },
+            schema=self.vector_schema)
+
+    def _blob_vector_rows(self, n):
+        return pa.Table.from_pydict(
+            {
+                'id': list(range(n)),
+                'payload': [f'blob-{i}'.encode() for i in range(n)],
+                'embedding': [
+                    [float(i), float(i + 1), float(i + 2)]
+                    for i in range(n)
+                ],
+            },
+            schema=self.blob_vector_schema)
+
     def _write_files(self, table, data):
         """Write one Arrow table and return the committed DataFileMeta list."""
         wb = table.new_batch_write_builder()
@@ -74,7 +127,7 @@ def _write_files(self, table, data):
         return files
 
     def _read_ids(self, table):
-        rb = table.new_read_builder()
+        rb = table.new_read_builder().with_projection(['id'])
         return sorted(
             rb.new_read().to_arrow(rb.new_scan().plan().splits())
             ['id'].to_pylist())
@@ -126,6 +179,92 @@ def test_non_de_table_still_fails_fast(self):
                 NotImplementedError, 'row-count based file rolling'):
             tw.write_arrow(self._rows(4))
 
+    def test_blob_writer_supports_target_file_row_num(self):
+        table = self._create_with_schema(
+            self.blob_schema,
+            {**self.de_options, 'target-file-row-num': '3'})
+
+        files = self._write_files(table, self._blob_rows(7))
+
+        data_rows = sorted(
+            f.row_count for f in files
+            if not f.file_name.endswith('.blob'))
+        blob_rows = sorted(
+            f.row_count for f in files
+            if f.file_name.endswith('.blob'))
+        self.assertEqual([1, 3, 3], data_rows)
+        self.assertEqual([1, 3, 3], blob_rows)
+        self.assertEqual(list(range(7)), self._read_ids(table))
+
+    def test_blob_consumer_descriptors_survive_abort_after_rolling(self):
+        table = self._create_with_schema(
+            self.blob_schema,
+            {**self.de_options, 'target-file-row-num': '3'})
+        descriptors = []
+
+        def consume(_, descriptor):
+            if descriptor is not None:
+                descriptors.append(descriptor)
+            return True
+
+        writer = table.new_batch_write_builder().new_write()
+        writer.with_blob_consumer(consume)
+        writer.write_arrow(self._blob_rows(7))
+        writer.abort()
+
+        self.assertEqual(7, len(descriptors))
+        uri_reader = FileUriReader(table.file_io)
+        for index, descriptor in enumerate(descriptors):
+            self.assertEqual(
+                f'blob-{index}'.encode(),
+                Blob.from_descriptor(uri_reader, descriptor).to_data())
+
+    def test_vector_writer_supports_target_file_row_num(self):
+        table = self._create_with_schema(
+            self.vector_schema,
+            {
+                **self.de_options,
+                'target-file-row-num': '3',
+                'vector.file.format': 'parquet',
+            })
+
+        files = self._write_files(table, self._vector_rows(7))
+
+        data_rows = sorted(
+            f.row_count for f in files
+            if '.vector.' not in f.file_name)
+        vector_rows = sorted(
+            f.row_count for f in files
+            if '.vector.' in f.file_name)
+        self.assertEqual([1, 3, 3], data_rows)
+        self.assertEqual([1, 3, 3], vector_rows)
+        self.assertEqual(list(range(7)), self._read_ids(table))
+
+    def test_dedicated_writer_rolls_blob_and_vector_together(self):
+        table = self._create_with_schema(
+            self.blob_vector_schema,
+            {
+                **self.de_options,
+                'target-file-row-num': '3',
+                'vector.file.format': 'parquet',
+            })
+
+        files = self._write_files(table, self._blob_vector_rows(7))
+
+        data_rows = sorted(
+            f.row_count for f in files
+            if not f.file_name.endswith('.blob') and '.vector.' not in f.file_name)
+        blob_rows = sorted(
+            f.row_count for f in files
+            if f.file_name.endswith('.blob'))
+        vector_rows = sorted(
+            f.row_count for f in files
+            if '.vector.' in f.file_name)
+        self.assertEqual([1, 3, 3], data_rows)
+        self.assertEqual([1, 3, 3], blob_rows)
+        self.assertEqual([1, 3, 3], vector_rows)
+        self.assertEqual(list(range(7)), self._read_ids(table))
+
 
 if __name__ == '__main__':
     unittest.main()

paimon-python/pypaimon/write/file_store_write.py

@@ -145,20 +145,14 @@ def _create_data_writer(self, partition: Tuple, bucket: int, options: CoreOption
             raise ValueError(
                 f"target-file-row-num should be at most {max_value}")
         if row_limit != max_value:
-            # Row-count rolling is implemented in the base append writer only.
-            # DE (data-evolution) append tables are the target; primary-key,
-            # blob and vector writers override rolling and are not supported yet.
             row_rolling_supported = (
                 self.table.options.data_evolution_enabled()
-                and not self.table.is_primary_key_table
-                and not self._has_blob_columns()
-                and not (self._has_vector_columns()
-                         and options.with_vector_format()))
+                and not self.table.is_primary_key_table)
             if not row_rolling_supported:
                 raise NotImplementedError(
                     "target-file-row-num is set on this table but pypaimon supports row-count "
-                    "based file rolling only for data-evolution append tables (no primary key, "
-                    "blob or vector columns); unset it or write with Java/Flink/Spark.")
+                    "based file rolling only for data-evolution append tables (no primary key); "
+                    "unset it or write with Java/Flink/Spark.")
 
         def max_seq_number():
             return self._seq_number_stats(partition).get(bucket, 1)

paimon-python/pypaimon/write/writer/blob_writer.py

@@ -113,7 +113,10 @@ def rolling_file(self) -> bool:
         if self.current_writer is None:
             return False
 
-        return self.current_writer.reach_target_size(self.blob_target_file_size)
+        return (
+            self.current_writer.row_count >= self.target_file_row_num
+            or self.current_writer.reach_target_size(self.blob_target_file_size)
+        )
 
     def close_current_writer(self):
         """Close current writer and create metadata."""
@@ -237,6 +240,9 @@ def close(self):
         # Call parent to handle pending_data fallback.
         super().close()
 
+    def delete_file_upon_abort(self) -> bool:
+        return self._blob_consumer is None
+
     def abort(self):
         if self.current_writer is not None:
             try:
@@ -245,7 +251,7 @@ def abort(self):
                 logger.warning(f"Error aborting blob writer: {e}", exc_info=e)
             self.current_writer = None
             self.current_file_path = None
-        if self._blob_consumer is not None:
+        if not self.delete_file_upon_abort():
             self.pending_data = None
             self.committed_files.clear()
         else:

paimon-python/pypaimon/write/writer/data_vector_writer.py

@@ -110,27 +110,43 @@ def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) -> pa.Table:
 
     def write(self, data: pa.RecordBatch):
         try:
-            normal_data, vector_data = self._split_data(data)
-
-            processed_normal = pa.Table.from_batches([normal_data]) if normal_data is not None else None
-            if self.pending_normal_data is None:
-                self.pending_normal_data = processed_normal
-            elif processed_normal is not None:
-                self.pending_normal_data = pa.concat_tables([self.pending_normal_data, processed_normal])
-
-            if self.vector_writer is not None and vector_data is not None and vector_data.num_rows > 0:
-                self.vector_writer.write(vector_data)
-
-            self.record_count += data.num_rows
-
-            if self._should_roll_normal():
-                self._close_current_writers()
+            offset = 0
+            # _write_batch keeps normal and vector pending rows in lockstep
+            # and closes both writers when the shared row limit is reached.
+            while offset < data.num_rows:
+                capacity = self.target_file_row_num - self._current_row_count()
+                if capacity <= 0:
+                    self._close_current_writers()
+                    capacity = self.target_file_row_num
+                length = min(capacity, data.num_rows - offset)
+                self._write_batch(data.slice(offset, length))
+                offset += length
 
         except Exception as e:
             logger.error("Exception occurs when writing data. Cleaning up.", exc_info=e)
             self.abort()
             raise e
 
+    def _write_batch(self, data: pa.RecordBatch):
+        if data.num_rows == 0:
+            return
+
+        normal_data, vector_data = self._split_data(data)
+
+        processed_normal = pa.Table.from_batches([normal_data]) if normal_data is not None else None
+        if self.pending_normal_data is None:
+            self.pending_normal_data = processed_normal
+        elif processed_normal is not None:
+            self.pending_normal_data = pa.concat_tables([self.pending_normal_data, processed_normal])
+
+        if self.vector_writer is not None and vector_data is not None and vector_data.num_rows > 0:
+            self.vector_writer.write(vector_data)
+
+        self.record_count += data.num_rows
+
+        if self._should_roll_normal():
+            self._close_current_writers()
+
     def prepare_commit(self) -> List[DataFileMeta]:
         self._close_current_writers()
         return self.committed_files.copy()
@@ -174,10 +190,19 @@ def _split_data(self, data: pa.RecordBatch) -> Tuple[pa.RecordBatch, pa.RecordBa
     def _should_roll_normal(self) -> bool:
         if self.pending_normal_data is None:
             return False
+        if self.pending_normal_data.num_rows >= self.target_file_row_num:
+            return True
         if self.record_count % self.CHECK_ROLLING_RECORD_CNT != 0:
             return False
         return self.pending_normal_data.nbytes > self.target_file_size
 
+    def _current_row_count(self) -> int:
+        if self.pending_normal_data is not None:
+            return self.pending_normal_data.num_rows
+        if self.vector_writer is not None and self.vector_writer.pending_data is not None:
+            return self.vector_writer.pending_data.num_rows
+        return 0
+
     def _close_current_writers(self):
         has_normal = self.pending_normal_data is not None and self.pending_normal_data.num_rows > 0
 
@@ -195,6 +220,7 @@ def _close_current_writers(self):
             self.vector_writer.committed_files.clear()
 
         self.pending_normal_data = None
+        self.record_count = 0
 
     def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]:
         if data.num_rows == 0: