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 a00fe7a9dc [python] Support target file row num for blob and vector
writes (#9317)
a00fe7a9dc is described below
commit a00fe7a9dcd370a779105c62f377ef4b544b0a24
Author: Yann Byron <[email protected]>
AuthorDate: Fri Aug 21 17:29:21 2026 +0800
[python] Support target file row num for blob and vector writes (#9317)
---
.../tests/data_evolution_row_rolling_test.py | 141 ++++++++++++++++++++-
paimon-python/pypaimon/write/file_store_write.py | 12 +-
paimon-python/pypaimon/write/writer/blob_writer.py | 10 +-
.../pypaimon/write/writer/data_vector_writer.py | 56 +++++---
.../write/writer/dedicated_format_writer.py | 100 ++++++++++-----
5 files changed, 259 insertions(+), 60 deletions(-)
diff --git a/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
b/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
index 2dffcf493b..63bebc7173 100644
--- a/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
+++ b/paimon-python/pypaimon/tests/data_evolution_row_rolling_test.py
@@ -24,6 +24,8 @@ import uuid
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 @@ class DataEvolutionRowRollingTest(unittest.TestCase):
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 @@ class DataEvolutionRowRollingTest(unittest.TestCase):
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 @@ class DataEvolutionRowRollingTest(unittest.TestCase):
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()
diff --git a/paimon-python/pypaimon/write/file_store_write.py
b/paimon-python/pypaimon/write/file_store_write.py
index 9a1dd10721..df3022b37d 100644
--- a/paimon-python/pypaimon/write/file_store_write.py
+++ b/paimon-python/pypaimon/write/file_store_write.py
@@ -145,20 +145,14 @@ class FileStoreWrite:
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)
diff --git a/paimon-python/pypaimon/write/writer/blob_writer.py
b/paimon-python/pypaimon/write/writer/blob_writer.py
index c5bf143728..40d301c2f1 100644
--- a/paimon-python/pypaimon/write/writer/blob_writer.py
+++ b/paimon-python/pypaimon/write/writer/blob_writer.py
@@ -113,7 +113,10 @@ class BlobWriter(AppendOnlyDataWriter):
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 @@ class BlobWriter(AppendOnlyDataWriter):
# 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 @@ class BlobWriter(AppendOnlyDataWriter):
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:
diff --git a/paimon-python/pypaimon/write/writer/data_vector_writer.py
b/paimon-python/pypaimon/write/writer/data_vector_writer.py
index b474957870..72291cb738 100644
--- a/paimon-python/pypaimon/write/writer/data_vector_writer.py
+++ b/paimon-python/pypaimon/write/writer/data_vector_writer.py
@@ -110,27 +110,43 @@ class DataVectorWriter(DataWriter):
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 @@ class DataVectorWriter(DataWriter):
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 @@ class DataVectorWriter(DataWriter):
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:
diff --git a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
index b7118e0507..ae73880fd0 100644
--- a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
+++ b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py
@@ -136,6 +136,7 @@ class DedicatedFormatWriter(DataWriter):
# Track pending data for normal data only
self.pending_normal_data: Optional[pa.Table] = None
+ self._committed_files_to_delete_on_abort: List[DataFileMeta] = []
# Initialize blob writers for each blob-file column.
from pypaimon.write.writer.blob_writer import BlobWriter
@@ -200,39 +201,55 @@ class DedicatedFormatWriter(DataWriter):
def write(self, data: pa.RecordBatch):
try:
- # Split data into normal, blob, and vector parts
- normal_data, blob_data_map, vector_data = self._split_data(data)
- self._validate_inline_stored_fields_input(data)
-
- # Process and accumulate normal data (may be None for partial
writes)
- processed_normal = self._process_normal_data(normal_data)
- if processed_normal is not None:
- if self.pending_normal_data is None:
- self.pending_normal_data = processed_normal
- else:
- self.pending_normal_data =
self._merge_normal_data(self.pending_normal_data, processed_normal)
-
- # Write blob-file columns to dedicated blob writers.
- for blob_column, blob_data in blob_data_map.items():
- if blob_data is not None and blob_data.num_rows > 0:
- self.blob_writers[blob_column].write(blob_data)
-
- # Write vector columns to dedicated vector writer.
- 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
-
- # Check if normal data rolling is needed
- if self._should_roll_normal():
- # When normal data rolls, close both writers and fetch blob
metadata
- self._close_current_writers()
+ offset = 0
+ # _write_batch keeps normal/blob/vector pending rows in lockstep
+ # and closes all 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
+
+ # Split data into normal, blob, and vector parts
+ normal_data, blob_data_map, vector_data = self._split_data(data)
+ self._validate_inline_stored_fields_input(data)
+
+ # Process and accumulate normal data (may be None for partial writes)
+ processed_normal = self._process_normal_data(normal_data)
+ if processed_normal is not None:
+ if self.pending_normal_data is None:
+ self.pending_normal_data = processed_normal
+ else:
+ self.pending_normal_data =
self._merge_normal_data(self.pending_normal_data, processed_normal)
+
+ # Write blob-file columns to dedicated blob writers.
+ for blob_column, blob_data in blob_data_map.items():
+ if blob_data is not None and blob_data.num_rows > 0:
+ self.blob_writers[blob_column].write(blob_data)
+
+ # Write vector columns to dedicated vector writer.
+ 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
+
+ # Check if normal data rolling is needed
+ if self._should_roll_normal():
+ # When normal data rolls, close both writers and fetch blob
metadata
+ self._close_current_writers()
+
def write_row(self, row):
try:
values_by_name = row_to_named_values(
@@ -335,14 +352,11 @@ class DedicatedFormatWriter(DataWriter):
blob_writer.abort()
if self.vector_writer is not None:
self.vector_writer.abort()
- committed_non_blob_files = [
- file_meta for file_meta in self.committed_files
- if not DataFileMeta.is_blob_file(file_meta.file_name)
- ]
- self._delete_committed_files(committed_non_blob_files)
+ self._delete_committed_files(self._committed_files_to_delete_on_abort)
self.pending_normal_data = None
self.pending_data = None
self.committed_files.clear()
+ self._committed_files_to_delete_on_abort.clear()
def _split_data(self, data: pa.RecordBatch) -> Tuple[
Optional[pa.RecordBatch], Dict[str, pa.RecordBatch],
Optional[pa.RecordBatch]]:
@@ -454,6 +468,9 @@ class DedicatedFormatWriter(DataWriter):
if self.pending_normal_data is None:
return False
+ if self.pending_normal_data.num_rows >= self.target_file_row_num:
+ return True
+
# Check rolling condition periodically (every CHECK_ROLLING_RECORD_CNT
records)
if self.record_count % self.CHECK_ROLLING_RECORD_CNT != 0:
return False
@@ -462,19 +479,34 @@ class DedicatedFormatWriter(DataWriter):
current_size = self.pending_normal_data.nbytes
return current_size > 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
+ for blob_writer in self.blob_writers.values():
+ if blob_writer.current_writer is not None:
+ return blob_writer.current_writer.row_count
+ 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):
"""Close normal, blob, and vector writers; add metadata in order:
normal, blob, vector."""
normal_meta = None
if self.pending_normal_data is not None and
self.pending_normal_data.num_rows > 0:
normal_meta =
self._write_normal_data_to_file(self.pending_normal_data)
self.committed_files.append(normal_meta)
+ self._committed_files_to_delete_on_abort.append(normal_meta)
blob_metas = []
for blob_column in self.blob_file_column_names:
- writer_metas = self.blob_writers[blob_column].prepare_commit()
+ blob_writer = self.blob_writers[blob_column]
+ writer_metas = blob_writer.prepare_commit()
if normal_meta is not None:
self._validate_consistency(normal_meta, writer_metas,
blob_column)
blob_metas.extend(writer_metas)
+ if blob_writer.delete_file_upon_abort():
+ self._committed_files_to_delete_on_abort.extend(writer_metas)
+ blob_writer.committed_files.clear()
self.committed_files.extend(blob_metas)
vector_metas = []
@@ -483,9 +515,11 @@ class DedicatedFormatWriter(DataWriter):
if vector_metas and normal_meta is not None:
self._validate_consistency(normal_meta, vector_metas, 'vector')
self.committed_files.extend(vector_metas)
+ self._committed_files_to_delete_on_abort.extend(vector_metas)
self.vector_writer.committed_files.clear()
self.pending_normal_data = None
+ self.record_count = 0
if normal_meta is not None or blob_metas or vector_metas:
normal_name = normal_meta.file_name if normal_meta is not None
else '<none>'
