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 2dffcf493b41..63bebc717310 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 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() diff --git a/paimon-python/pypaimon/write/file_store_write.py b/paimon-python/pypaimon/write/file_store_write.py index 9a1dd10721a8..df3022b37d2d 100644 --- a/paimon-python/pypaimon/write/file_store_write.py +++ b/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) diff --git a/paimon-python/pypaimon/write/writer/blob_writer.py b/paimon-python/pypaimon/write/writer/blob_writer.py index c5bf143728a7..40d301c2f151 100644 --- a/paimon-python/pypaimon/write/writer/blob_writer.py +++ b/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: diff --git a/paimon-python/pypaimon/write/writer/data_vector_writer.py b/paimon-python/pypaimon/write/writer/data_vector_writer.py index b47495787032..72291cb738f3 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 @@ 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: diff --git a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py index b7118e0507a7..ae73880fd091 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 @@ def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, op # 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 @@ def _merge_data(self, existing_data: pa.Table, new_data: pa.Table) -> pa.Table: 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 @@ def abort(self): 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 @@ 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 + # 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 @@ def _should_roll_normal(self) -> bool: 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 @@ def _close_current_writers(self): 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 ''