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 e6613e9ac8 [python] Fix fragmented Parquet row groups in MAP
shared-shredding writes (#10386)
e6613e9ac8 is described below
commit e6613e9ac8a59878072a46b75c26a11e2c32987d
Author: XiaoHongbo <[email protected]>
AuthorDate: Mon Oct 5 20:40:37 2026 +0800
[python] Fix fragmented Parquet row groups in MAP shared-shredding writes
(#10386)
---
.../tests/map_shared_shredding_write_test.py | 86 +++++++++++++++++--
.../pypaimon/write/map_shared_shredding_writer.py | 98 +++++++++++++++++++---
2 files changed, 165 insertions(+), 19 deletions(-)
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 2ca511d09e..2cbc44cace 100644
--- a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py
+++ b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py
@@ -277,18 +277,15 @@ class MapSharedShreddingWriteTest(unittest.TestCase):
}, schema=self.arrow_schema)
convert = _MapFieldConverter.convert
write_table = pq.ParquetWriter.write_table
- pending = []
sizes = []
+ written = []
def convert_batch(converter, column):
- self.assertFalse(pending, 'physical batches were retained before
writing')
- pending.append(len(column))
sizes.append(len(column))
return convert(converter, column)
def write_batch(writer, physical, *args, **kwargs):
- self.assertEqual([physical.num_rows], pending)
- pending.clear()
+ written.append(physical.num_rows)
return write_table(writer, physical, *args, **kwargs)
with patch.object(_MapFieldConverter, 'convert', convert_batch), \
@@ -296,6 +293,7 @@ class MapSharedShreddingWriteTest(unittest.TestCase):
messages = self._write(table, data)
self.assertEqual(5000, sum(sizes))
self.assertLessEqual(max(sizes), 1024)
+ self.assertEqual([5000], written)
self.assertEqual(1, len(messages[0].new_files))
reader = table.new_read_builder().with_projection(["metrics['a']",
"metrics['late']"])
result = reader.new_read().to_arrow(reader.new_scan().plan().splits())
@@ -305,7 +303,7 @@ class MapSharedShreddingWriteTest(unittest.TestCase):
def test_failed_stream_removes_partial_file(self):
from pypaimon.write.map_shared_shredding_writer import
MapSharedShreddingWriter
- table = self._create_table('parquet', 2)
+ table = self._create_table('parquet', 2, {'file.block-size': '1 kb'})
converter = MapSharedShreddingWriter(table.fields, table.options,
'parquet', None)
data = pa.Table.from_pydict({
'id': list(range(2048)), 'metrics': [[('a', 1)]] * 2048,
@@ -351,13 +349,87 @@ class MapSharedShreddingWriteTest(unittest.TestCase):
metadata = pq.read_metadata(files[0].file_path)
layouts.append([metadata.row_group(i).num_rows
for i in range(metadata.num_row_groups)])
- self.assertEqual((count + 1023) // 1024,
metadata.num_row_groups)
+ self.assertEqual(1, metadata.num_row_groups)
reader = table.new_read_builder().with_projection(['id',
"metrics['a']"])
result =
reader.new_read().to_arrow(reader.new_scan().plan().splits())
self.assertEqual(list(range(count)),
result.column('id').to_pylist())
self.assertEqual(list(range(count)),
result.column('metrics_a').to_pylist())
self.assertEqual(layouts[0], layouts[1])
+ def test_physical_row_group_byte_and_row_limits(self):
+ from pypaimon.write.map_shared_shredding_writer import
_physical_row_groups
+
+ data = pa.table({'id': pa.array(range(25), type=pa.int64())})
+ for chunks in ([data], [data.slice(i, 3) for i in range(0, 25, 3)]):
+ for target, max_rows, expected in (
+ (80, 100, [10, 10, 5]), (800, 6, [6, 6, 6, 6, 1])):
+ groups = list(_physical_row_groups(iter(chunks), target,
max_rows))
+ self.assertEqual(expected, [g.num_rows for g in groups])
+ self.assertTrue(all(g.nbytes <= target for g in groups))
+ self.assertTrue(pa.concat_tables(groups).equals(data))
+
+ def test_physical_row_groups_empty_nulls_and_oversized_row(self):
+ from pypaimon.write.map_shared_shredding_writer import
_physical_row_groups
+
+ self.assertEqual([], list(_physical_row_groups(iter([]), 100)))
+ data = pa.table({'value': [None, 'x' * 300, '', 'tail']})
+ groups = list(_physical_row_groups(iter([data.slice(0, 0), data]), 64))
+ self.assertTrue(pa.concat_tables(groups).equals(data))
+ self.assertTrue(all(g.nbytes <= 64 or g.num_rows == 1 for g in groups))
+ self.assertEqual([1, 1, 2], [g.num_rows for g in groups])
+
+ def test_physical_row_groups_release_input_on_close(self):
+ from pypaimon.write.map_shared_shredding_writer import
_physical_row_groups
+
+ closed = []
+
+ def batches():
+ try:
+ while True:
+ yield pa.table({'id': [1, 2]})
+ finally:
+ closed.append(True)
+
+ groups = _physical_row_groups(batches(), 16)
+ next(groups)
+ groups.close()
+ self.assertEqual([True], closed)
+
+ def test_physical_row_groups_preserve_nested_shredding_metadata(self):
+ from pypaimon.write.map_shared_shredding_writer import
_physical_row_groups
+
+ table = self._create_table('parquet', 2, {
+ 'write.native.enabled': 'false', 'read.native.enabled': 'false',
+ 'file.block-size': '4 kb',
+ })
+ maps = [[('a', i), ('b', None), ('overflow', i + 1)] for i in
range(2500)]
+ data = pa.Table.from_pydict(
+ {'id': list(range(2500)), 'metrics': maps},
schema=self.arrow_schema)
+ captured = []
+
+ def collect(batches, target_bytes):
+ for group in _physical_row_groups(batches, target_bytes):
+ captured.append((group.num_rows, group.nbytes))
+
self.assertTrue(is_shared_shredding(group.schema.field('metrics')))
+ yield group
+
+ with
patch('pypaimon.write.map_shared_shredding_writer._physical_row_groups',
collect):
+ messages = self._write(table, data)
+ self.assertTrue(all(size <= 4096 or rows == 1 for rows, size in
captured))
+ metadata = pq.read_metadata(messages[0].new_files[0].file_path)
+ self.assertEqual([rows for rows, _ in captured],
+ [metadata.row_group(i).num_rows for i in
range(metadata.num_row_groups)])
+ reader = table.new_read_builder()
+ actual = reader.new_read().to_arrow(reader.new_scan().plan().splits())
+ self.assertEqual(maps, actual['metrics'].to_pylist())
+
+ def test_reject_nonpositive_row_group_target(self):
+ from pypaimon.write.map_shared_shredding_writer import
MapSharedShreddingWriter
+
+ table = self._create_table('parquet', 2, {'file.block-size': '0 b'})
+ with self.assertRaisesRegex(ValueError, 'file.block-size must be
positive'):
+ MapSharedShreddingWriter(table.fields, table.options, 'parquet',
None)
+
def test_reject_postpone_with_fixed_output_bucket(self):
from pypaimon.write.writer.append_only_data_writer import
AppendOnlyDataWriter
diff --git a/paimon-python/pypaimon/write/map_shared_shredding_writer.py
b/paimon-python/pypaimon/write/map_shared_shredding_writer.py
index 539e9a0425..0302459e15 100644
--- a/paimon-python/pypaimon/write/map_shared_shredding_writer.py
+++ b/paimon-python/pypaimon/write/map_shared_shredding_writer.py
@@ -45,6 +45,9 @@ _PHYSICAL_COLUMN_PREFIX = "__col_"
_FIELD_ID_BASE = 2147483647 // 4
_FIELD_ID_DEPTH_LIMIT = 1 << 10
_CONVERSION_BYTES = 8 * 1024 * 1024
+_ROW_GROUP_MAX_ROWS = 1024 * 1024
+# Match the default Parquet block target, measured here in physical Arrow
bytes.
+_DEFAULT_ROW_GROUP_BYTES = 128 * 1024 * 1024
class MapSharedShreddingWriter:
@@ -90,6 +93,11 @@ class MapSharedShreddingWriter:
self._validate_format("file.format", file_format)
self._validate_format("changelog.file.format", changelog_format)
self._parquet_writer_options = create_parquet_writer_options(options)
+ block_size = options.file_block_size()
+ self._row_group_bytes = (block_size.get_bytes() if block_size is not
None
+ else _DEFAULT_ROW_GROUP_BYTES)
+ if self._row_group_bytes <= 0:
+ raise ValueError('file.block-size must be positive')
self._validate_compression("file.compression",
options.file_compression())
if options.bucket() == BucketMode.POSTPONE_BUCKET.value:
raise ValueError(
@@ -107,7 +115,7 @@ class MapSharedShreddingWriter:
return bool(self._fields)
def write_parquet(self, file_io, path, data, compression, zstd_level):
- """Plan key metadata, then write bounded physical batches without
retaining them."""
+ """Convert small batches, buffering physical rows independently of
conversion."""
import pyarrow.parquet as pq
fields = list(data.schema)
@@ -163,23 +171,28 @@ class MapSharedShreddingWriter:
kwargs = dict(self._parquet_writer_options, compression=compression)
if compression.lower() == 'zstd':
kwargs['compression_level'] = zstd_level
+
+ def physical_batches():
+ for offset in range(0, data.num_rows, batch_rows):
+ for bounded in _bounded_batches(data.slice(offset,
batch_rows)):
+ batch = bounded.combine_chunks().to_batches()[0]
+ columns = list(batch.columns)
+ for index, converter in converters.items():
+ columns[index] = converter.convert(columns[index])
+ yield pa.Table.from_arrays(columns, schema=schema)
+
+ groups = _physical_row_groups(physical_batches(),
self._row_group_bytes)
try:
with file_io.new_output_stream(path) as stream:
with pq.ParquetWriter(stream, schema, **kwargs) as writer:
- for offset in range(0, data.num_rows, batch_rows):
- for bounded in _bounded_batches(data.slice(offset,
batch_rows)):
- # Coalesce only this bounded logical window, not
the
- # file. Input calls must not define row groups.
- batch = bounded.combine_chunks().to_batches()[0]
- columns = list(batch.columns)
- for index, converter in converters.items():
- columns[index] =
converter.convert(columns[index])
- physical = pa.Table.from_arrays(columns,
schema=schema)
- writer.write_table(physical)
- del physical, columns
+ for physical in groups:
+ writer.write_table(physical,
row_group_size=physical.num_rows)
+ del physical
except Exception:
file_io.delete_quietly(path)
raise
+ finally:
+ groups.close()
return completed
def file_completed(self, completed):
@@ -396,6 +409,67 @@ class _MapFieldConverter:
return selected
+def _physical_row_groups(batches, target_bytes, max_rows=_ROW_GROUP_MAX_ROWS):
+ """Group converted batches without copying them into one contiguous Arrow
table.
+
+ Bytes are an Arrow-side target, not compressed Parquet size or a process
RSS
+ limit. One indivisible row may exceed it. In addition to the pending group,
+ the current conversion batch and Parquet encoder retain their own buffers.
+ """
+ pending = []
+ rows = size = 0
+ try:
+ for batch in batches:
+ offset = 0
+ while offset < batch.num_rows:
+ count = min(batch.num_rows - offset, max_rows - rows)
+ piece = _physical_slice(batch, offset, count)
+ available = target_bytes - size
+ if piece.nbytes > available:
+ low, high = 0, count
+ while low < high:
+ middle = (low + high + 1) // 2
+ if _physical_slice(batch, offset, middle).nbytes <=
available:
+ low = middle
+ else:
+ high = middle - 1
+ if low == 0 and rows:
+ yield pa.concat_tables(pending)
+ pending.clear()
+ rows = size = 0
+ continue
+ count = max(1, low)
+ piece = _physical_slice(batch, offset, count)
+ pending.append(piece)
+ rows += count
+ size += piece.nbytes
+ offset += count
+ del piece
+ if rows >= max_rows or size >= target_bytes:
+ yield pa.concat_tables(pending)
+ pending.clear()
+ rows = size = 0
+ del batch
+ if rows:
+ yield pa.concat_tables(pending)
+ finally:
+ pending.clear()
+ close = getattr(batches, 'close', None)
+ if close is not None:
+ close()
+
+
+def _physical_slice(batch, offset, count):
+ piece = batch.slice(offset, count)
+ # Arrow 6 accounts full backing buffers for slices. Compact partial slices
+ # there, as the row-ID update writer does, for bounded retained-byte
accounting.
+ if int(pa.__version__.split('.')[0]) < 7:
+ piece = pa.Table.from_arrays([
+ pa.concat_arrays(column.chunks) for column in piece.columns
+ ], schema=piece.schema)
+ return piece
+
+
def _bounded_batches(batch):
# One oversized row is indivisible; bound the other logical batches before
# expanding values into Python objects. The caller's input buffer is
unchanged.
