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 1b802b9b76 [python] Reduce memory usage for Parquet row-id updates
(#9840)
1b802b9b76 is described below
commit 1b802b9b76a059caed4fcb3b9842d43975aa36c8
Author: XiaoHongbo <[email protected]>
AuthorDate: Wed Sep 16 19:35:30 2026 +0800
[python] Reduce memory usage for Parquet row-id updates (#9840)
---
paimon-python/pypaimon/read/table_read.py | 6 +
.../tests/map_shared_shredding_write_test.py | 44 ++++
.../tests/table_update_by_row_id_chunked_test.py | 112 ++++++++-
.../pypaimon/tests/table_upsert_by_key_test.py | 250 +++++++++++++++++++++
paimon-python/pypaimon/write/table_update.py | 1 +
.../pypaimon/write/table_update_by_row_id.py | 227 ++++++++++++++++++-
paimon-python/pypaimon/write/writer/data_writer.py | 43 +++-
.../pypaimon/write/writer/single_file_writer.py | 106 +++++++++
8 files changed, 767 insertions(+), 22 deletions(-)
diff --git a/paimon-python/pypaimon/read/table_read.py
b/paimon-python/pypaimon/read/table_read.py
index 2b0d5a4ea3..5b8407874e 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -196,6 +196,12 @@ class TableRead:
self,
splits: List[Split],
blob_parallelism: Optional[int] = None):
+ """Return a closeable batch reader supporting context management.
+
+ Newer PyArrow versions use ``RecordBatchReader.from_stream``. Older
+ versions fall back to ``_ClosableArrowBatchReader``, which closes both
+ the batch iterator and its underlying reader.
+ """
reader, batch_iterator = self._new_arrow_batch_reader(
splits, blob_parallelism)
if (_RECORD_BATCH_READER_FROM_STREAM is not None
diff --git a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py
b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py
index eba20b4270..fe82a6d2ed 100644
--- a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py
+++ b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py
@@ -92,6 +92,50 @@ class MapSharedShreddingWriteTest(unittest.TestCase):
result.column("metrics_overflow").to_pylist(),
)
+ def test_row_id_update_preserves_shared_shredding(self):
+ table = self._create_table('parquet', 2, {
+ 'data-evolution.enabled': 'true',
+ 'row-tracking.enabled': 'true',
+ })
+ self._write(table, pa.Table.from_pydict({
+ 'id': [1, 2],
+ 'metrics': [[('hot', 1)], [('warm', 2)]],
+ }, schema=self.arrow_schema))
+
+ read_builder = table.new_read_builder().with_projection(['id',
'_ROW_ID'])
+ rows = read_builder.new_read().to_arrow(
+ read_builder.new_scan().plan().splits())
+ row_ids = dict(zip(rows.column('id').to_pylist(),
+ rows.column('_ROW_ID').to_pylist()))
+ row_id = row_ids[2]
+ update = pa.Table.from_pydict({
+ '_ROW_ID': [row_id],
+ 'metrics': [[('hot', 99), ('new', 7)]],
+ }, schema=pa.schema([
+ ('_ROW_ID', pa.int64()),
+ self.arrow_schema.field('metrics'),
+ ]))
+
+ builder = table.new_batch_write_builder()
+ messages = builder.new_update().with_update_type(
+ ['metrics']).update_by_arrow_with_row_id(update)
+ self.assertEqual(1, len(messages[0].new_files))
+ overlay = messages[0].new_files[0]
+ self.assertTrue(is_shared_shredding(
+ pq.read_schema(overlay.file_path).field('metrics')))
+ commit = builder.new_commit()
+ commit.commit(messages)
+ commit.close()
+
+ read_builder = table.new_read_builder().with_projection(['id',
'metrics'])
+ result = read_builder.new_read().to_arrow(
+ read_builder.new_scan().plan().splits())
+ self.assertEqual(
+ {1: [('hot', 1)], 2: [('hot', 99), ('new', 7)]},
+ dict(zip(result.column('id').to_pylist(),
+ result.column('metrics').to_pylist())),
+ )
+
def test_reject_orc(self):
with self.assertRaisesRegex(
ValueError,
diff --git
a/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py
b/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py
index c0902c5421..3ce93701ee 100644
--- a/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py
+++ b/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py
@@ -21,13 +21,76 @@ from unittest import mock
import pyarrow as pa
import pyarrow.compute as pc
+import pyarrow.parquet as pq
+from pypaimon.common.options.core_options import ChangelogProducer
from pypaimon.table.special_fields import SpecialFields
-from pypaimon.write.table_update_by_row_id import TableUpdateByRowId
+from pypaimon.write.table_update_by_row_id import (
+ TableUpdateByRowId,
+ _RowIdUpdateFileWriter,
+)
+from pypaimon.write.writer.single_file_writer import SingleFileWriter
class TableUpdateByRowIdChunkedTest(unittest.TestCase):
+ def test_streaming_overlay_rejects_sidecar_before_output(self):
+ table = mock.Mock()
+ table.is_primary_key_table = False
+ table.fields = []
+ options = table.options
+ options.file_format.return_value = 'parquet'
+ options.variant_shredding_enabled.return_value = False
+ options.data_evolution_row_sidecar_enabled.return_value = False
+ options.with_vector_format.return_value = False
+ options.changelog_producer.return_value = ChangelogProducer.NONE
+ self.assertTrue(_RowIdUpdateFileWriter.supports_table(table))
+
+ options.data_evolution_row_sidecar_enabled.return_value = True
+ with self.assertRaisesRegex(ValueError, 'requires plain Parquet'):
+ _RowIdUpdateFileWriter(table, (), ['id'])
+ table.file_io.new_output_stream.assert_not_called()
+
+ def test_streaming_overlay_rejects_shared_shredding_before_output(self):
+ table = mock.Mock()
+ table.is_primary_key_table = False
+ field = mock.Mock()
+ field.name = 'metrics'
+ table.fields = [field]
+ options = table.options
+ options.file_format.return_value = 'parquet'
+ options.variant_shredding_enabled.return_value = False
+ options.data_evolution_row_sidecar_enabled.return_value = False
+ options.with_vector_format.return_value = False
+ options.changelog_producer.return_value = ChangelogProducer.NONE
+ options.map_storage_layout.return_value = 'shared-shredding'
+
+ with self.assertRaisesRegex(ValueError, 'requires plain Parquet'):
+ _RowIdUpdateFileWriter(table, (), ['metrics'])
+ table.file_io.new_output_stream.assert_not_called()
+
+ def test_single_file_writer_rejects_other_formats_before_open(self):
+ file_io = mock.Mock()
+ with self.assertRaisesRegex(NotImplementedError, 'only supports
Parquet'):
+ SingleFileWriter(file_io, 'unused.orc', pa.schema([]), 'orc',
+ 'zstd', 1, [], mock.Mock())
+ file_io.new_output_stream.assert_not_called()
+
+ def test_single_file_writer_failed_open_only_deletes_owned_file(self):
+ file_io = mock.Mock()
+ file_io.new_output_stream.side_effect = FileExistsError('already
exists')
+ with self.assertRaises(FileExistsError):
+ SingleFileWriter(file_io, 'existing.parquet', pa.schema([]),
'parquet',
+ 'zstd', 1, [], mock.Mock())
+ file_io.delete_quietly.assert_not_called()
+
+ file_io.new_output_stream.side_effect = None
+ with mock.patch.object(pq, 'ParquetWriter',
side_effect=OSError('failed to open')):
+ with self.assertRaisesRegex(OSError, 'failed to open'):
+ SingleFileWriter(file_io, 'new.parquet', pa.schema([]),
'parquet',
+ 'zstd', 1, [], mock.Mock())
+ file_io.delete_quietly.assert_called_once_with('new.parquet')
+
@staticmethod
def _updater():
updater = TableUpdateByRowId.__new__(TableUpdateByRowId)
@@ -105,6 +168,53 @@ class TableUpdateByRowIdChunkedTest(unittest.TestCase):
self._updater()._merge_update_with_original(
original, updates, ["payload"], first_row_id=0)
+ def test_streaming_update_rejects_non_contiguous_original_row_ids(self):
+ updates = pa.table({
+ SpecialFields.ROW_ID.name: pa.array([0], type=pa.int64()),
+ "payload": pa.array([3]),
+ })
+
+ for row_ids in ([0, 2], [0, 2, 1]):
+ with self.subTest(row_ids=row_ids):
+ batch = pa.RecordBatch.from_pydict({
+ SpecialFields.ROW_ID.name:
+ pa.array(row_ids, type=pa.int64()),
+ "payload": pa.array(row_ids),
+ })
+ managed_reader = mock.MagicMock()
+ managed_reader.__enter__.return_value = iter([batch])
+ table_read = mock.Mock()
+ table_read._to_managed_arrow_batch_reader.return_value = (
+ managed_reader)
+ updater = self._updater()
+ updater._original_file_read = mock.Mock(
+ return_value=(table_read, mock.sentinel.split))
+
+ with self.assertRaisesRegex(
+ ValueError, "not contiguous at row ID 0"):
+ list(updater._merged_batches(0, updates, ["payload"]))
+
+ def test_streaming_update_rejects_row_ids_before_original_group(self):
+ updates = pa.table({
+ SpecialFields.ROW_ID.name: pa.array([9], type=pa.int64()),
+ "payload": pa.array([3]),
+ })
+ batch = pa.RecordBatch.from_pydict({
+ SpecialFields.ROW_ID.name: pa.array([10, 11], type=pa.int64()),
+ "payload": pa.array([1, 2]),
+ })
+ managed_reader = mock.MagicMock()
+ managed_reader.__enter__.return_value = iter([batch])
+ table_read = mock.Mock()
+ table_read._to_managed_arrow_batch_reader.return_value = managed_reader
+ updater = self._updater()
+ updater._original_file_read = mock.Mock(
+ return_value=(table_read, mock.sentinel.split))
+
+ with self.assertRaisesRegex(
+ ValueError, "precede the original file group"):
+ list(updater._merged_batches(10, updates, ["payload"]))
+
def test_total_offsets_over_int32_remain_in_separate_chunks(self):
child_length = 1_100_000_000
large_chunk = self._list_chunk([0, child_length])
diff --git a/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
index 7f06fcab42..564d61f3ed 100644
--- a/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
+++ b/paimon-python/pypaimon/tests/table_upsert_by_key_test.py
@@ -16,12 +16,16 @@
# under the License.
import os
+from contextlib import contextmanager
import unittest
from unittest import mock
import pyarrow as pa
+import pyarrow.parquet as pq
+from pypaimon.read.table_read import TableRead
from pypaimon.table.special_fields import SpecialFields
+from pypaimon.write.table_update_by_row_id import _RowIdUpdateFileWriter
from pypaimon.tests.data_evolution_test_helpers import (
BatchModeMixin,
DataEvolutionTestBase,
@@ -61,6 +65,197 @@ class _TableUpsertByKeyTestBase(DataEvolutionTestBase):
def _apply_upsert_rows(self, table_update, rows, upsert_keys, cid):
raise NotImplementedError
+ def test_upsert_row_groups_do_not_follow_read_batches(self):
+ schema = pa.schema([('id', pa.int32()), ('score', pa.int32())])
+ for read_size in (73, 1024):
+ with self.subTest(read_size=read_size):
+ table = self._create_table(pa_schema=schema, options={
+ **self.table_options, 'read.batch-size': str(read_size)})
+ original = pa.Table.from_pydict(
+ {'id': list(range(5000)), 'score': list(range(5000))},
schema=schema)
+ self._write_arrow(table, original)
+ updates = pa.Table.from_pydict({'id': [13], 'score': [-1]},
schema=schema)
+ messages = self._upsert(table, updates, ['id'], ['score'])
+ files = [f for message in messages for f in message.new_files]
+ self.assertEqual(len(files), 1)
+ metadata = pq.read_metadata(files[0].file_path)
+ self.assertEqual(metadata.num_row_groups, 1)
+ self.assertEqual(metadata.row_group(0).num_rows, 5000)
+ scores = self._read_all(table).to_pydict()['score']
+ expected = list(range(5000))
+ expected[13] = -1
+ self.assertEqual(scores, expected)
+
+ def test_row_group_byte_budget_and_oversized_row(self):
+ schema = pa.schema([('id', pa.int32())])
+ table = self._create_table(pa_schema=schema, options={
+ **self.table_options, 'file.block-size': '400 b'})
+ writer = _RowIdUpdateFileWriter(table, (), ['id'])
+ data = pa.Table.from_pydict({'id': list(range(250))}, schema=schema)
+ try:
+ for batch_size in (1, 73, 250):
+ groups = list(writer._row_groups(
+ data.to_batches(max_chunksize=batch_size)))
+ self.assertEqual([group.num_rows for group in groups], [100,
100, 50])
+ self.assertTrue(pa.concat_tables(groups).equals(data))
+ self.assertTrue(all(group.nbytes <= 400 for group in groups))
+ large = pa.Table.from_pydict({'text': ['a', 'x' * 500, 'b']})
+ groups = list(writer._row_groups(large.to_batches()))
+ self.assertEqual([group.num_rows for group in groups], [1, 1, 1])
+ self.assertTrue(pa.concat_tables(groups).equals(large))
+ # Empty batches, nulls and sliced variable-width columns must not
+ # lose rows or make the output depend on input batch boundaries.
+ mixed = pa.Table.from_pydict({
+ 'text': [
+ 'skip', None, '', 'a' * 300,
+ 'b' * 300, None, 'last', 'skip',
+ ],
+ }).slice(1, 6)
+ layouts = []
+ for batch_size in (1, 2, 6):
+ batches = mixed.to_batches(max_chunksize=batch_size)
+ batches.insert(0, batches[0].slice(0, 0))
+ groups = list(writer._row_groups(batches))
+ self.assertTrue(pa.concat_tables(groups).equals(mixed))
+ self.assertTrue(all(group.nbytes <= 400 for group in groups))
+ layouts.append([group.num_rows for group in groups])
+ self.assertTrue(all(layout == layouts[0] for layout in layouts))
+ self.assertEqual(list(writer._row_groups([])), [])
+ with mock.patch.object(_RowIdUpdateFileWriter,
'_ROW_GROUP_MAX_ROWS', 17):
+ groups = list(writer._row_groups(
+ data.to_batches(max_chunksize=73)))
+ self.assertEqual([group.num_rows for group in groups], [17] *
14 + [12])
+ self.assertTrue(pa.concat_tables(groups).equals(data))
+ finally:
+ writer.close()
+
+ def test_invalid_row_group_size_fails_before_opening_output(self):
+ schema = pa.schema([('id', pa.int32())])
+ table = self._create_table(pa_schema=schema, options={
+ **self.table_options, 'file.block-size': '0 b'})
+ with mock.patch.object(
+ table.file_io, 'new_output_stream') as output_stream:
+ with self.assertRaisesRegex(
+ ValueError, 'file.block-size must be positive'):
+ _RowIdUpdateFileWriter(table, (), ['id'])
+ output_stream.assert_not_called()
+
+ @mock.patch.object(_RowIdUpdateFileWriter, '_ROW_GROUP_MAX_ROWS', 2)
+ def test_partial_upsert_streams_original_file_group(self):
+ schema = pa.schema([
+ ('id', pa.int32()),
+ ('payload', pa.list_(pa.struct([('text', pa.string())]))),
+ ('score', pa.int32()),
+ ])
+ table = self._create_table(pa_schema=schema, options={
+ **self.table_options, 'metadata.stats-mode': 'full'})
+ expected = [{'id': i, 'payload': [{'text': str(i)}], 'score': i}
+ for i in range(12)]
+
+ def as_table(rows):
+ return pa.Table.from_pydict(
+ {name: [row[name] for row in rows] for name in schema.names},
+ schema=schema)
+
+ self._write_arrow(table, as_table(expected))
+ read_batches = TableRead._to_managed_arrow_batch_reader
+ write = pq.ParquetWriter.write_table
+ progress = {'read': 0, 'written': 0}
+ closed = []
+
+ @contextmanager
+ def bounded_reader(reader, splits, **kwargs):
+ source = read_batches(reader, splits, **kwargs)
+
+ def batches():
+ try:
+ for batch in source:
+ for start in range(0, batch.num_rows, 2):
+ self.assertEqual(progress['read'],
progress['written'])
+ piece = batch.slice(start, 2)
+ progress['read'] += piece.num_rows
+ yield piece
+ finally:
+ source.close()
+ closed.append(True)
+ iterator = batches()
+ try:
+ yield iterator
+ finally:
+ iterator.close()
+
+ def record_write(writer, batch, **kwargs):
+ result = write(writer, batch, **kwargs)
+ progress['written'] += batch.num_rows
+ return result
+
+ for replacements in ([9, 1, 5], [0, 11]):
+ updates = [{'id': i, 'payload': None if i == 5 else
+ [{'text': 'updated-' + str(i)}],
+ 'score': None if i == 5 else -i} for i in replacements]
+ progress.update(read=0, written=0)
+ with mock.patch.object(TableRead, 'to_arrow',
side_effect=AssertionError(
+ 'upsert must not materialize the original file group')):
+ with mock.patch.object(TableRead,
'_to_managed_arrow_batch_reader', bounded_reader):
+ with mock.patch.object(pq.ParquetWriter, 'write_table',
record_write):
+ messages = self._upsert(
+ table, as_table(updates),
+ ['id'], ['payload', 'score'])
+ self.assertEqual(progress, {'read': 12, 'written': 12})
+ for row in updates:
+ expected[row['id']] = row
+ self.assertEqual(self._read_all(table).to_pydict(),
as_table(expected).to_pydict())
+ files = [f for msg in messages for f in msg.new_files]
+ self.assertEqual(len(files), 1)
+ self.assertEqual((files[0].first_row_id, files[0].row_count), (0,
12))
+ scores = [r['score'] for r in expected if r['score'] is not None]
+ self.assertEqual(files[0].value_stats.min_values.values[1],
min(scores))
+ self.assertEqual(files[0].value_stats.max_values.values[1],
max(scores))
+ self.assertEqual(files[0].value_stats.null_counts, [1, 1])
+
+ # A failed streamed write must leave the committed table intact and
+ # remove the partially written overlay.
+ self.assertEqual(len(closed), 2)
+ progress.update(read=0, written=0)
+
+ def fail_second_write(writer, batch, **kwargs):
+ if progress['written']:
+ raise OSError('injected write failure')
+ return record_write(writer, batch, **kwargs)
+
+ with mock.patch.object(table.file_io, 'delete_quietly',
+ wraps=table.file_io.delete_quietly) as delete:
+ with mock.patch.object(pq.ParquetWriter, 'write_table',
+ fail_second_write):
+ with mock.patch.object(TableRead,
'_to_managed_arrow_batch_reader', bounded_reader):
+ with self.assertRaisesRegex(OSError, 'injected write
failure'):
+ self._upsert(table, as_table(updates),
+ ['id'], ['payload', 'score'])
+ self.assertEqual(len(closed), 3)
+ self.assertTrue(delete.called)
+ for call in delete.call_args_list:
+ self.assertFalse(table.file_io.exists(call[0][0]))
+ self.assertEqual(self._read_all(table).to_pydict(),
as_table(expected).to_pydict())
+
+ # Closing the format writer is part of the file transaction too.
+ close = pq.ParquetWriter.close
+
+ def fail_close(writer):
+ was_open = writer.is_open
+ close(writer)
+ if was_open:
+ raise OSError('injected close failure')
+
+ with mock.patch.object(table.file_io, 'delete_quietly',
+ wraps=table.file_io.delete_quietly) as delete:
+ with mock.patch.object(pq.ParquetWriter, 'close', fail_close):
+ with self.assertRaisesRegex(OSError, 'injected close failure'):
+ self._upsert(table, as_table(updates), ['id'], ['payload',
'score'])
+ self.assertTrue(delete.called)
+ for call in delete.call_args_list:
+ self.assertFalse(table.file_io.exists(call[0][0]))
+ self.assertEqual(self._read_all(table).to_pydict(),
as_table(expected).to_pydict())
+
# ------------------------------------------------------------------
# Helpers built on the primitives
# ------------------------------------------------------------------
@@ -953,6 +1148,61 @@ class _TableUpsertByKeyTestBase(DataEvolutionTestBase):
self.assertEqual((1, 'Alice', 99, 'NYC'), rows[0])
self.assertEqual((2, 'Bob', 88, 'LA'), rows[1])
+ def test_duplicate_update_cols_are_deduplicated(self):
+ table = self._create_table()
+ self._write_arrow(table, pa.Table.from_pydict({
+ 'id': [1],
+ 'name': ['Alice'],
+ 'age': [25],
+ 'city': ['NYC'],
+ }, schema=self.pa_schema))
+
+ messages = self._upsert(
+ table,
+ pa.Table.from_pydict({
+ 'id': [1],
+ 'name': ['ignored'],
+ 'age': [99],
+ 'city': ['ignored'],
+ }, schema=self.pa_schema),
+ upsert_keys=['id'],
+ # Matching the schema width must not mean "update all columns".
+ update_cols=['age'] * len(table.field_names),
+ )
+
+ self.assertEqual(
+ self._read_all(table).to_pydict(),
+ {'id': [1], 'name': ['Alice'], 'age': [99], 'city': ['NYC']},
+ )
+ files = [file for message in messages for file in message.new_files]
+ self.assertEqual([file.write_cols for file in files], [['age']])
+
+ def test_not_null_update_across_read_batches(self):
+ schema = pa.schema([
+ pa.field('id', pa.int32(), nullable=False),
+ pa.field('score', pa.int32(), nullable=False),
+ ])
+ table = self._create_table(pa_schema=schema, options={
+ **self.table_options, 'read.batch-size': '2'})
+ original = pa.Table.from_pydict({
+ 'id': list(range(4)),
+ 'score': list(range(4)),
+ }, schema=schema)
+ self._write_arrow(table, original)
+
+ updates = pa.Table.from_pydict({'id': [2], 'score': [99]},
schema=schema)
+ messages = self._upsert(table, updates, ['id'], ['score'])
+
+ expected = original.set_column(
+ 1,
+ schema.field('score'),
+ pa.array([0, 1, 99, 3], type=pa.int32()),
+ )
+ self.assertTrue(self._read_all(table).equals(expected))
+ files = [file for message in messages for file in message.new_files]
+ self.assertEqual(len(files), 1)
+
self.assertFalse(pq.read_schema(files[0].file_path).field('score').nullable)
+
# ==================================================================
# Duplicate-key dedup tests — parametrised
# ==================================================================
diff --git a/paimon-python/pypaimon/write/table_update.py
b/paimon-python/pypaimon/write/table_update.py
index c358644460..922b1bc8c3 100644
--- a/paimon-python/pypaimon/write/table_update.py
+++ b/paimon-python/pypaimon/write/table_update.py
@@ -126,6 +126,7 @@ class TableUpdate:
self.projection = None
def with_update_type(self, update_cols: List[str]):
+ update_cols = list(dict.fromkeys(update_cols))
for col in update_cols:
if col not in self.table.field_names:
raise ValueError(f"Column {col} is not in table schema.")
diff --git a/paimon-python/pypaimon/write/table_update_by_row_id.py
b/paimon-python/pypaimon/write/table_update_by_row_id.py
index 90c08d20d2..9226f8d823 100644
--- a/paimon-python/pypaimon/write/table_update_by_row_id.py
+++ b/paimon-python/pypaimon/write/table_update_by_row_id.py
@@ -16,6 +16,7 @@
# under the License.
import bisect
+import uuid
from dataclasses import dataclass, field
from typing import Any, Dict, List, Optional, Set, Tuple
@@ -23,8 +24,10 @@ import numpy as np
import pyarrow as pa
import pyarrow.compute as pc
+from pypaimon.common.options.core_options import ChangelogProducer, CoreOptions
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.manifest_entry import ManifestEntry
+from pypaimon.manifest.schema.simple_stats import SimpleStats
from pypaimon.read.scanner.data_evolution_split_generator import (
DataEvolutionSplitGenerator,
)
@@ -49,6 +52,13 @@ from pypaimon.write.row_utils import (
value_for_arrow,
)
from pypaimon.write.writer.blob_writer import BlobWriter
+from pypaimon.write.writer.append_only_data_writer import AppendOnlyDataWriter
+from pypaimon.write.writer.single_file_writer import SingleFileWriter
+from pypaimon.write.writer.write_buffer import WriteBuffer
+
+_ARROW_MAJOR = int(pa.__version__.split('.')[0])
+# Keep aligned with org.apache.parquet.hadoop.ParquetWriter.DEFAULT_BLOCK_SIZE.
+_DEFAULT_PARQUET_BLOCK_SIZE = 128 * 1024 * 1024
@dataclass(frozen=True)
@@ -66,6 +76,134 @@ class _FilesInfo:
valid_row_id_ranges: List[Range] = field(default_factory=list)
+class _RowIdUpdateFileWriter:
+ """Write one plain-Parquet update file for a row-id file group."""
+
+ _ROW_GROUP_MAX_ROWS = 1024 * 1024
+
+ @staticmethod
+ def supports_table(table):
+ options = table.options
+ return (not table.is_primary_key_table
+ and options.file_format(CoreOptions.FILE_FORMAT_PARQUET)
+ == CoreOptions.FILE_FORMAT_PARQUET
+ and not (options.variant_shredding_enabled()
+ and options.variant_shredding_schema())
+ and not options.data_evolution_row_sidecar_enabled(False)
+ and not options.with_vector_format()
+ and options.changelog_producer() == ChangelogProducer.NONE
+ and not any(options.map_storage_layout(f.name) ==
'shared-shredding'
+ for f in table.fields)
+ and not any(is_blob_file_field(f) for f in table.fields))
+
+ def __init__(self, table, partition, column_names):
+ if not self.supports_table(table):
+ raise ValueError('Row-id update file writer requires plain Parquet
without sidecars')
+ configured = table.options.file_block_size()
+ self._target_bytes = (configured.get_bytes() if configured is not None
+ else _DEFAULT_PARQUET_BLOCK_SIZE)
+ if self._target_bytes <= 0:
+ raise ValueError('file.block-size must be positive')
+ self._data_writer = AppendOnlyDataWriter(
+ table, partition, 0, 0, table.options, write_cols=column_names)
+
+ @staticmethod
+ def _row_group_slice(batch, offset, count):
+ piece = batch.slice(offset, count)
+ # Arrow 6 nbytes counts full backing buffers even for slices. Compact
+ # the slice there so both accounting and retained buffers stay bounded.
+ if _ARROW_MAJOR < 7:
+ piece = pa.RecordBatch.from_arrays(
+ [pa.concat_arrays([column]) for column in piece.columns],
schema=piece.schema)
+ return piece
+
+ def _row_groups(self, batches):
+ """Bound row-group buffering using Arrow-side bytes, not encoded
size."""
+ buffer = WriteBuffer(self._data_writer._merge_data)
+ try:
+ for batch in batches:
+ offset = 0
+ while offset < batch.num_rows:
+ count = min(batch.num_rows - offset,
+ self._ROW_GROUP_MAX_ROWS - buffer.num_rows)
+ piece = self._row_group_slice(batch, offset, count)
+ available = self._target_bytes - buffer.nbytes
+ if piece.nbytes > available:
+ low, high = 0, count
+ while low < high:
+ middle = (low + high + 1) // 2
+ if self._row_group_slice(piece, 0, middle).nbytes
<= available:
+ low = middle
+ else:
+ high = middle - 1
+ count = low
+ if count == 0 and buffer.num_rows:
+ yield buffer.take()
+ continue
+ count = max(1, count)
+ piece = self._row_group_slice(piece, 0, count)
+ buffer.append(pa.Table.from_batches([piece]))
+ offset += count
+ del piece
+ if buffer.nbytes >= self._target_bytes or buffer.num_rows
>= self._ROW_GROUP_MAX_ROWS:
+ yield buffer.take()
+ del batch
+ if buffer.num_rows:
+ yield buffer.take()
+ finally:
+ buffer.reset()
+
+ def write_batches(self, batches):
+ """Write one file while keeping input and output row groups bounded."""
+ writer = self._data_writer
+ file_name = '{}{}-0.parquet'.format(
+ writer.options.data_file_prefix(), uuid.uuid4())
+ file_path = writer._generate_file_path(file_name)
+ fields = []
+ groups = self._row_groups(batches)
+ file_writer = None
+ try:
+ for batch in groups:
+ if not batch.num_rows:
+ continue
+ if file_writer is None:
+ if writer.options.metadata_stats_enabled():
+ fields =
PyarrowFieldParser.to_paimon_schema(batch.schema)
+ file_writer = SingleFileWriter(
+ writer.file_io, file_path, batch.schema,
writer.file_format,
+ writer.compression, writer.zstd_level,
+ fields, writer._get_column_stats)
+ file_writer.write(batch, row_group_size=batch.num_rows)
+ del batch
+ if file_writer is None:
+ return []
+ file_writer.close()
+ meta = writer._create_data_file_meta(
+ file_name=file_name,
+ file_path=file_path,
+ row_count=file_writer.row_count,
+ min_key=GenericRow([], []), max_key=GenericRow([], []),
+ key_stats=SimpleStats.empty_stats(),
+ value_stats=writer._collect_value_stats(
+ None, fields, file_writer.column_stats),
+ min_sequence_number=0, max_sequence_number=0,
+ )
+ writer._finish_data_file(meta)
+ return [meta]
+ except Exception:
+ if file_writer is not None:
+ file_writer.abort()
+ raise
+ finally:
+ groups.close()
+
+ def abort(self):
+ self._data_writer.abort()
+
+ def close(self):
+ self._data_writer.close()
+
+
class TableUpdateByRowId:
"""
Table update for partial column updates (data evolution).
@@ -211,6 +349,7 @@ class TableUpdateByRowId:
if not column_names:
raise ValueError("column_names cannot be empty")
+ column_names = list(dict.fromkeys(column_names))
if SpecialFields.ROW_ID.name not in data.column_names:
raise ValueError(f"Input data must contain
{SpecialFields.ROW_ID.name} column")
@@ -257,6 +396,7 @@ class TableUpdateByRowId:
) -> List[CommitMessage]:
if not column_names:
raise ValueError("column_names cannot be empty")
+ column_names = list(dict.fromkeys(column_names))
if len(rows) != len(row_ids_by_row):
raise ValueError(
"rows and row_ids_by_row must have the same length: "
@@ -390,7 +530,7 @@ class TableUpdateByRowId:
group_blob_object_columns,
)
- def _read_original_file_data(self, first_row_id: int, column_names:
List[str]) -> Optional[pa.Table]:
+ def _read_original_file_data(self, first_row_id: int, column_names:
List[str]) -> pa.Table:
"""Read original file data for the given first_row_id.
Only reads columns that exist in the original file and need to be
updated.
@@ -402,16 +542,17 @@ class TableUpdateByRowId:
column_names: The column names to update
Returns:
- PyArrow Table containing the original data for columns that exist
in the file,
- or None if no columns need to be read from the original file.
+ PyArrow Table containing the original values for the requested
columns.
"""
+ table_read, origin_split = self._original_file_read(first_row_id,
column_names)
+ original = table_read.to_arrow([origin_split])
+ return original.select(column_names)
+
+ def _original_file_read(self, first_row_id, column_names):
wanted = set(column_names)
read_fields: List[DataField] = [
table_field for table_field in self.table.fields if
table_field.name in wanted
]
- if not read_fields:
- return None
-
entry = self._first_row_id_index.get(first_row_id)
if entry is None:
raise ValueError(f"No file found for first_row_id {first_row_id}")
@@ -432,8 +573,66 @@ class TableUpdateByRowId:
predicate=None,
read_type=read_fields + [SpecialFields.ROW_ID],
)
- original = table_read.to_arrow([origin_split])
- return original.select([field.name for field in read_fields])
+ return table_read, origin_split
+
+ def _merged_batches(self, first_row_id, data, column_names):
+ """Merge ordinary columns a batch at a time in physical row order."""
+ table_read, split = self._original_file_read(first_row_id,
column_names)
+ updates =
sorted(enumerate(data[SpecialFields.ROW_ID.name].to_pylist()),
+ key=lambda item: item[1])
+ update_index = 0
+ offset = first_row_id
+ with table_read._to_managed_arrow_batch_reader([split]) as reader:
+ for batch in reader:
+ if not batch.num_rows:
+ continue
+ row_ids = batch[SpecialFields.ROW_ID.name]
+ if row_ids.null_count:
+ raise ValueError(
+ 'Original file group contains null _ROW_ID values')
+ row_id_values = row_ids.to_numpy(zero_copy_only=True)
+ if (row_id_values[0] != offset
+ or (len(row_id_values) > 1
+ and not np.all(np.diff(row_id_values) == 1))):
+ raise ValueError(
+ f'Original file group is not contiguous at row ID
{offset}')
+ end = offset + batch.num_rows
+ selected = []
+ while update_index < len(updates) and updates[update_index][1]
< end:
+ if updates[update_index][1] < offset:
+ raise ValueError(
+ 'Update row IDs precede the original file group')
+ selected.append(updates[update_index][0])
+ update_index += 1
+ original = pa.Table.from_batches([batch]).select(column_names)
+ if selected:
+ merged, _ = self._merge_update_with_original(
+ original, data.take(selected), column_names, offset)
+ else:
+ merged = original
+ yield from merged.to_batches()
+ offset = end
+ del batch, original, merged
+ if update_index != len(updates):
+ raise ValueError('Update row IDs extend past the original file
group')
+
+ def _write_group_streaming(self, partition, first_row_id, data,
column_names):
+ writer = _RowIdUpdateFileWriter(
+ self.table, tuple(partition.values), column_names)
+ batches = self._merged_batches(first_row_id, data, column_names)
+ try:
+ files = writer.write_batches(batches)
+ self._assign_update_file_metadata(files, first_row_id,
column_names, {})
+ if files:
+ self.commit_messages.append(CommitMessage(
+ partition=tuple(partition.values), bucket=0,
new_files=files,
+ check_from_snapshot=self.snapshot_id))
+ except Exception:
+ writer.abort()
+ raise
+ finally:
+ batches.close()
+ writer.close()
def _merge_update_with_original(
self,
@@ -533,7 +732,13 @@ class TableUpdateByRowId:
merged_columns[col_name] = self._merge_chunked_column(
original_col, update_col, sorted_updates)
- merged_table = pa.table(merged_columns) if merged_columns else None
+ merged_table = None
+ if merged_columns:
+ merged_schema = pa.schema([
+ original_data.schema.field(name)
+ for name in merged_columns
+ ])
+ merged_table = pa.table(merged_columns, schema=merged_schema)
return merged_table, blob_columns
@@ -770,6 +975,10 @@ class TableUpdateByRowId:
Reads the original file data, merges in the update values, and
writes a single output file (rolling disabled) for the group.
"""
+ # Specialized writers still own their sidecars and physical encoding.
+ if _RowIdUpdateFileWriter.supports_table(self.table):
+ self._write_group_streaming(partition, first_row_id, data,
column_names)
+ return
original_data = self._read_original_file_data(first_row_id,
column_names)
_, target_files = self._first_row_id_index[first_row_id]
blob_columns_with_baseline = {
diff --git a/paimon-python/pypaimon/write/writer/data_writer.py
b/paimon-python/pypaimon/write/writer/data_writer.py
index 46db0578ce..71f5004f7e 100644
--- a/paimon-python/pypaimon/write/writer/data_writer.py
+++ b/paimon-python/pypaimon/write/writer/data_writer.py
@@ -343,9 +343,9 @@ class DataWriter(ABC):
min_seq = self.sequence_generator.start
max_seq = self.sequence_generator.current
creation_time = Timestamp.now()
- data_meta = DataFileMeta.create(
+ data_meta = self._create_data_file_meta(
file_name=file_name,
- file_size=self.file_io.get_file_size(file_path),
+ file_path=file_path,
row_count=data.num_rows,
min_key=GenericRow(min_key, self.trimmed_primary_keys_fields),
max_key=GenericRow(max_key, self.trimmed_primary_keys_fields),
@@ -353,17 +353,8 @@ class DataWriter(ABC):
value_stats=value_stats,
min_sequence_number=min_seq,
max_sequence_number=max_seq,
- schema_id=self.table.table_schema.id,
- level=0,
extra_files=extra_files,
creation_time=creation_time,
- delete_row_count=0,
- file_source=0,
- value_stats_cols=None if value_stats_enabled else [],
- external_path=external_path_str,
- first_row_id=None,
- write_cols=self.write_cols,
- file_path=file_path,
)
if self.changelog_producer == ChangelogProducer.INPUT:
@@ -378,8 +369,14 @@ class DataWriter(ABC):
self.file_io.delete_quietly(row_sidecar_path)
raise
+ self._finish_data_file(data_meta, changelog_meta,
shared_shredding_stats)
+
+ def _finish_data_file(self, data_meta, changelog_meta=None,
+ shared_shredding_stats=None):
+ """Record a fully written data file and its optional changelog."""
self.sequence_generator.start = self.sequence_generator.current
- self._map_shared_shredding.file_completed(shared_shredding_stats)
+ if shared_shredding_stats is not None:
+ self._map_shared_shredding.file_completed(shared_shredding_stats)
self.committed_files.append(data_meta)
if changelog_meta is not None:
self.committed_changelog_files.append(changelog_meta)
@@ -391,6 +388,28 @@ class DataWriter(ABC):
self.file_io.write_parquet(path, data, compression=self.compression,
zstd_level=self.zstd_level)
return {}
+ def _create_data_file_meta(self, file_name, file_path, row_count,
+ min_key, max_key, key_stats, value_stats,
+ min_sequence_number, max_sequence_number,
+ extra_files=None, creation_time=None):
+ """Common metadata finalization for buffered and incremental files."""
+ return DataFileMeta.create(
+ file_name=file_name,
+ file_size=self.file_io.get_file_size(file_path),
+ row_count=row_count,
+ min_key=min_key, max_key=max_key,
+ key_stats=key_stats, value_stats=value_stats,
+ min_sequence_number=min_sequence_number,
+ max_sequence_number=max_sequence_number,
+ schema_id=self.table.table_schema.id, level=0,
+ extra_files=extra_files if extra_files is not None else [],
+ creation_time=creation_time if creation_time is not None else
Timestamp.now(),
+ delete_row_count=0, file_source=0,
+ value_stats_cols=None if self.options.metadata_stats_enabled()
else [],
+ external_path=file_path if self.external_path_provider is not None
else None,
+ first_row_id=None, write_cols=self.write_cols, file_path=file_path,
+ )
+
def _apply_variant_shredding(self, data: pa.Table) -> pa.Table:
"""Transform VARIANT columns into shredded Parquet format.
diff --git a/paimon-python/pypaimon/write/writer/single_file_writer.py
b/paimon-python/pypaimon/write/writer/single_file_writer.py
new file mode 100644
index 0000000000..184d013822
--- /dev/null
+++ b/paimon-python/pypaimon/write/writer/single_file_writer.py
@@ -0,0 +1,106 @@
+# Licensed to the Apache Software Foundation (ASF) under one
+# or more contributor license agreements. See the NOTICE file
+# distributed with this work for additional information
+# regarding copyright ownership. The ASF licenses this file
+# to you under the Apache License, Version 2.0 (the
+# "License"); you may not use this file except in compliance
+# with the License. You may obtain a copy of the License at
+#
+# http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing,
+# software distributed under the License is distributed on an
+# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
+# KIND, either express or implied. See the License for the
+# specific language governing permissions and limitations
+# under the License.
+
+from contextlib import suppress
+
+import pyarrow.parquet as pq
+
+from pypaimon.common.options.core_options import CoreOptions
+
+
+class SingleFileWriter:
+ """Write one file incrementally; currently only Parquet is supported."""
+
+ def __init__(self, file_io, path, schema, file_format, compression,
zstd_level,
+ stats_fields, stats_collector):
+ if file_format != CoreOptions.FILE_FORMAT_PARQUET:
+ raise NotImplementedError(
+ 'SingleFileWriter only supports Parquet, got
{}'.format(file_format))
+ self._file_io = file_io
+ self._path = path
+ self._stats_fields = stats_fields
+ self._stats_collector = stats_collector
+ self._stream = None
+ self._writer = None
+ self._owns_file = False
+ self._closed = False
+ self.row_count = 0
+ self.column_stats = {}
+
+ kwargs = {'compression': compression}
+ if compression.lower() == 'zstd':
+ kwargs['compression_level'] = zstd_level
+ try:
+ self._stream = file_io.new_output_stream(path)
+ self._owns_file = True
+ self._writer = pq.ParquetWriter(self._stream, schema, **kwargs)
+ except Exception:
+ self.abort()
+ raise
+
+ def write(self, data, row_group_size=None):
+ if self._closed:
+ raise RuntimeError('Writer is already closed')
+ try:
+ kwargs = {}
+ if row_group_size is not None:
+ kwargs['row_group_size'] = row_group_size
+ self._writer.write_table(data, **kwargs)
+ self.row_count += data.num_rows
+ for field in self._stats_fields:
+ current = self._stats_collector(data, field.name)
+ previous = self.column_stats.get(field.name)
+ if previous is not None:
+ current['null_counts'] += previous['null_counts']
+ for key, choose in (
+ ('min_values', min), ('max_values', max)):
+ values = [
+ value for value in (previous[key], current[key])
+ if value is not None
+ ]
+ current[key] = choose(values) if values else None
+ self.column_stats[field.name] = current
+ except Exception:
+ self.abort()
+ raise
+
+ def close(self):
+ if self._closed:
+ return
+ try:
+ self._writer.close()
+ self._writer = None
+ self._stream.close()
+ self._stream = None
+ self._closed = True
+ except Exception:
+ self.abort()
+ raise
+
+ def abort(self):
+ if self._writer is not None:
+ with suppress(Exception):
+ self._writer.close()
+ self._writer = None
+ if self._stream is not None:
+ with suppress(Exception):
+ self._stream.close()
+ self._stream = None
+ self._closed = True
+ if self._owns_file:
+ self._file_io.delete_quietly(self._path)
+ self._owns_file = False
