From 92927a2cd5d967062773fc73f4be468d6e7f28ac Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 15 Sep 2026 00:41:01 -0700 Subject: [PATCH 1/5] [python] Support writing shared-shredding MAPs in Parquet --- .../java/org/apache/paimon/JavaPyE2ETest.java | 45 ++ paimon-python/dev/run_mixed_tests.sh | 15 + .../pypaimon/common/options/core_options.py | 44 ++ .../pypaimon/data/map_shared_shredding.py | 51 +- .../tests/e2e/java_py_read_write_test.py | 35 ++ .../tests/map_shared_shredding_write_test.py | 164 +++++++ .../write/map_shared_shredding_writer.py | 463 ++++++++++++++++++ .../pypaimon/write/writer/data_writer.py | 26 +- 8 files changed, 837 insertions(+), 6 deletions(-) create mode 100644 paimon-python/pypaimon/tests/map_shared_shredding_write_test.py create mode 100644 paimon-python/pypaimon/write/map_shared_shredding_writer.py diff --git a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java index 524344ed54e1..3e375da1b760 100644 --- a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -1620,6 +1620,51 @@ public void testJavaWriteSharedShreddingMapTable() throws Exception { } } + /** Java reads shared-shredding MAP columns written by Python. */ + @Test + @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") + public void testJavaReadSharedShreddingMapTable() throws Exception { + FileStoreTable table = + (FileStoreTable) + catalog.getTable(identifier("shared_shredding_map_python_test_parquet")); + Map> rows = new HashMap<>(); + List splits = new ArrayList<>(table.newSnapshotReader().read().dataSplits()); + try (org.apache.paimon.reader.RecordReader reader = + table.newRead().createReader(splits)) { + reader.forEachRemaining( + row -> { + int id = row.getInt(0); + if (row.isNullAt(1)) { + rows.put(id, null); + return; + } + InternalMap map = row.getMap(1); + InternalArray keys = map.keyArray(); + InternalArray values = map.valueArray(); + Map converted = new LinkedHashMap<>(); + for (int i = 0; i < map.size(); i++) { + converted.put( + keys.getString(i).toString(), + values.isNullAt(i) ? null : values.getLong(i)); + } + rows.put(id, converted); + }); + } + + assertThat(rows).containsOnlyKeys(1, 2, 3, 4); + assertThat(rows.get(1)) + .containsOnlyKeys("hot", "warm", "overflow") + .containsEntry("hot", 10L) + .containsEntry("warm", 20L) + .containsEntry("overflow", 30L); + assertThat(rows.get(2)) + .containsOnlyKeys("hot", "new") + .containsEntry("hot", null) + .containsEntry("new", 40L); + assertThat(rows.get(3)).isEmpty(); + assertThat(rows.get(4)).isNull(); + } + private Map> readMapBlobRows(FileStoreTable table) throws Exception { Map> rows = new HashMap<>(); diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh index 620c4d989a03..54fef2a4e1c0 100755 --- a/paimon-python/dev/run_mixed_tests.sh +++ b/paimon-python/dev/run_mixed_tests.sh @@ -1033,6 +1033,21 @@ run_shared_shredding_map_test() { return 1 fi echo -e "${GREEN}✓ Python shared-shredding MAP read test completed successfully${NC}" + + echo "Running Python shared-shredding MAP write test..." + if ! python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest::test_write_shared_shredding_map_for_java -v; then + echo -e "${RED}✗ Python shared-shredding MAP write test failed${NC}" + return 1 + fi + echo -e "${GREEN}✓ Python shared-shredding MAP write test completed successfully${NC}" + + cd "$PROJECT_ROOT" + echo "Running Maven test for JavaPyE2ETest.testJavaReadSharedShreddingMapTable..." + if ! mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testJavaReadSharedShreddingMapTable -pl paimon-core -q -Drun.e2e.tests=true; then + echo -e "${RED}✗ Java shared-shredding MAP read test failed${NC}" + return 1 + fi + echo -e "${GREEN}✓ Java shared-shredding MAP read test completed successfully${NC}" } # Function to run VARIANT test (Java write, Python read) diff --git a/paimon-python/pypaimon/common/options/core_options.py b/paimon-python/pypaimon/common/options/core_options.py index ce77d21f5a84..a22957388eb2 100644 --- a/paimon-python/pypaimon/common/options/core_options.py +++ b/paimon-python/pypaimon/common/options/core_options.py @@ -158,6 +158,10 @@ class CoreOptions: NESTED_SEQUENCE_FIELD = "nested-sequence-field" COUNT_LIMIT = "count-limit" MERGE_MAP_TS_FIELD = "ts-field" + MAP_STORAGE_LAYOUT = "map.storage-layout" + MAP_SHARED_SHREDDING_MAX_COLUMNS = "map.shared-shredding.max-columns" + MAP_SHARED_SHREDDING_COLUMN_PLACEMENT_POLICY = \ + "map.shared-shredding.column-placement-policy" # Basic options AUTO_CREATE: ConfigOption[bool] = ( @@ -1873,6 +1877,46 @@ def field_merge_map_ts_field(self, field_name: str) -> str: .no_default_value() ) + def map_storage_layout(self, field_name: str) -> str: + return self.options.get( + ConfigOptions.key( + f'{CoreOptions.FIELDS_PREFIX}.{field_name}.{CoreOptions.MAP_STORAGE_LAYOUT}' + ) + .string_type() + .default_value('default') + ).lower() + + def map_shared_shredding_max_columns(self, field_name: str) -> int: + value = self.options.get( + ConfigOptions.key( + f'{CoreOptions.FIELDS_PREFIX}.{field_name}.' + f'{CoreOptions.MAP_SHARED_SHREDDING_MAX_COLUMNS}' + ) + .int_type() + .default_value(256) + ) + if value <= 0: + raise ValueError( + '{} must be greater than 0'.format( + CoreOptions.MAP_SHARED_SHREDDING_MAX_COLUMNS)) + return value + + def map_shared_shredding_column_placement_policy( + self, field_name: str) -> str: + value = self.options.get( + ConfigOptions.key( + f'{CoreOptions.FIELDS_PREFIX}.{field_name}.' + f'{CoreOptions.MAP_SHARED_SHREDDING_COLUMN_PLACEMENT_POLICY}' + ) + .string_type() + .default_value('lru') + ).lower() + if value not in ('plain', 'sequential', 'lru'): + raise ValueError( + "Unsupported shared-shredding column placement policy: {}".format( + value)) + return value + @property def query_auth_enabled(self) -> bool: return self.options.get(CoreOptions.QUERY_AUTH_ENABLED) diff --git a/paimon-python/pypaimon/data/map_shared_shredding.py b/paimon-python/pypaimon/data/map_shared_shredding.py index a3d6c58bac55..13c15df75fdc 100644 --- a/paimon-python/pypaimon/data/map_shared_shredding.py +++ b/paimon-python/pypaimon/data/map_shared_shredding.py @@ -14,7 +14,7 @@ # See the License for the specific language governing permissions and # limitations under the License. -"""Read support for Paimon's shared-shredding MAP storage layout.""" +"""Metadata and read support for the shared-shredding MAP layout.""" import json import struct @@ -33,6 +33,7 @@ _NUM_COLUMNS = b"paimon.map.shared-shredding.num-columns" _FIELD_COLUMNS = b"paimon.map.shared-shredding.field-columns" _OVERFLOW_SET = b"paimon.map.shared-shredding.overflow-set" +_MAX_ROW_WIDTH = b"paimon.map.shared-shredding.max-row-width" _FIELD_MAPPING = "__field_mapping" _OVERFLOW = "__overflow" _PHYSICAL_COLUMN_PREFIX = "__col_" @@ -104,6 +105,40 @@ def parse_shared_shredding_selection_metadata(field: pa.Field): return name_by_id, field_to_columns, set(overflow_json), num_columns +def shared_shredding_metadata( + name_to_id, field_to_columns, overflow_set, num_columns, + max_row_width, compression): + """Build Java-compatible Arrow field metadata for one data file.""" + compression = compression.lower() + if compression not in ("none", "lz4", "zstd"): + raise ValueError( + "MAP shared-shredding only supports none/lz4/zstd compression, " + "but is {}.".format(compression)) + field_dict = json.dumps( + dict(sorted(name_to_id.items())), + separators=(",", ":"), + sort_keys=True, + ).encode("utf-8") + encoded_dict = _compress(field_dict, compression) + columns = { + str(field_id): sorted(column_ids) + for field_id, column_ids in sorted(field_to_columns.items()) + } + return { + _STORAGE_LAYOUT: b"shared-shredding", + _VERSION: b"1", + _FIELD_DICT: encoded_dict.decode("latin-1").encode("utf-8"), + _FIELD_DICT_COMPRESSION: compression.encode("utf-8"), + _FIELD_DICT_ORIGINAL_SIZE: str(len(field_dict)).encode("utf-8"), + _FIELD_COLUMNS: json.dumps( + columns, separators=(",", ":"), sort_keys=True).encode("utf-8"), + _OVERFLOW_SET: json.dumps( + sorted(overflow_set), separators=(",", ":")).encode("utf-8"), + _NUM_COLUMNS: str(num_columns).encode("utf-8"), + _MAX_ROW_WIDTH: str(max_row_width).encode("utf-8"), + } + + def map_selected_keys(description: str) -> List[str]: if not description or not description.startswith(_SELECTED_KEYS_PREFIX): raise ValueError("Invalid selected-key MAP metadata: {}".format( @@ -589,6 +624,20 @@ def _decompress(data: bytes, original_size: int, compression: str) -> bytes: return result +def _compress(data: bytes, compression: str) -> bytes: + if compression == "none": + return data + if compression == "zstd": + import zstandard as zstd + return zstd.ZstdCompressor(level=1).compress(data) + if compression == "lz4": + payload = bytes(pa.Codec("lz4_raw").compress(data)) + return struct.pack("= 0: + self.field_to_columns.setdefault(field_id, set()).add(column_id) + if overflow: + row[_OVERFLOW] = [(field_id, values[field_id]) + for field_id in overflow] + self.overflow_set.update(overflow) + else: + row[_OVERFLOW] = None + self.max_row_width = max(self.max_row_width, len(field_ids)) + return row + + def _allocate(self, field_ids): + if self.policy == "plain": + ordered = field_ids + return self._leading(ordered) + if self.policy == "sequential": + return self._leading(sorted(field_ids)) + return self._lru(field_ids) + + def _leading(self, field_ids): + mapping = [-1] * self.num_columns + for index, field_id in enumerate(field_ids[:self.num_columns]): + mapping[index] = field_id + return mapping, list(field_ids[self.num_columns:]) + + def _lru(self, field_ids): + mapping = [-1] * self.num_columns + next_resident = list(self._resident) + used = [False] * self.num_columns + unassigned = [] + for field_id in sorted(field_ids): + try: + column_id = self._resident.index(field_id) + except ValueError: + unassigned.append(field_id) + continue + used[column_id] = True + mapping[column_id] = field_id + + overflow = [] + for field_id in unassigned: + column_id = self._select_lru_column(used, next_resident) + if column_id < 0: + overflow.append(field_id) + continue + used[column_id] = True + mapping[column_id] = field_id + next_resident[column_id] = field_id + + touched = False + for column_id, field_id in enumerate(mapping): + if field_id >= 0: + self._last_used[column_id] = self._clock + touched = True + if touched: + self._clock += 1 + self._resident = next_resident + return mapping, overflow + + def _select_lru_column(self, used, resident): + selected = -1 + selected_last_used = None + for column_id in range(self.num_columns): + if used[column_id]: + continue + if resident[column_id] < 0: + return column_id + last_used = self._last_used[column_id] + if selected < 0 or last_used < selected_last_used: + selected = column_id + selected_last_used = last_used + return selected + + +def _physical_struct_type(num_columns, item_type, logical_item_type): + mapping_type = pa.list_(pa.field( + "item", + pa.int32(), + metadata=_field_id_metadata(_array_element_id(0, 1)), + )) + fields = [pa.field( + _FIELD_MAPPING, + mapping_type, + metadata=_field_id_metadata(0), + )] + for column_id in range(num_columns): + fields.append(_field_with_ids( + _PHYSICAL_COLUMN_PREFIX + str(column_id), + item_type, + logical_item_type.nullable, + logical_item_type, + column_id + 1, + )) + overflow_id = num_columns + 1 + overflow_type = _map_type_with_ids( + pa.int32(), item_type, logical_item_type, overflow_id, 0) + fields.append(pa.field( + _OVERFLOW, + overflow_type, + metadata=_field_id_metadata(overflow_id), + )) + return pa.struct(fields) + + +def _field_with_ids(name, arrow_type, nullable, logical_type, field_id): + return pa.field( + name, + _type_with_ids(arrow_type, logical_type, field_id, 0), + nullable=nullable, + metadata=_field_id_metadata(field_id), + ) + + +def _type_with_ids(arrow_type, logical_type, field_id, depth): + if isinstance(logical_type, RowType) and pa.types.is_struct(arrow_type): + arrow_by_name = {field.name: field for field in arrow_type} + return pa.struct([ + _field_with_ids( + field.name, + arrow_by_name[field.name].type, + arrow_by_name[field.name].nullable, + field.type, + field.id, + ) + for field in logical_type.fields + ]) + if isinstance(logical_type, (ArrayType, VectorType)): + child_id = _array_element_id(field_id, depth + 1) + value_field = pa.field( + "item", + _type_with_ids( + arrow_type.value_type, + logical_type.element, + field_id, + depth + 1, + ), + nullable=logical_type.element.nullable, + metadata=_field_id_metadata(child_id), + ) + if isinstance(logical_type, VectorType): + return pa.list_(value_field, logical_type.length) + return pa.list_(value_field) + if isinstance(logical_type, MapType): + return _map_type_with_ids( + arrow_type.key_type, + arrow_type.item_type, + logical_type.value, + field_id, + depth, + ) + return arrow_type + + +def _map_type_with_ids( + key_type, item_type, logical_item_type, field_id, depth): + key = pa.field( + "key", + key_type, + nullable=False, + metadata=_field_id_metadata(_map_key_id(field_id, depth + 1)), + ) + item = pa.field( + "value", + _type_with_ids( + item_type, logical_item_type, field_id, depth + 1), + nullable=logical_item_type.nullable, + metadata=_field_id_metadata(_map_value_id(field_id, depth + 1)), + ) + return pa.map_(key, item) + + +def _field_id_metadata(field_id): + value = str(field_id).encode("utf-8") + return {b"PARQUET:field_id": value, b"paimon.id": value} + + +def _array_element_id(field_id, depth): + return _FIELD_ID_BASE + field_id * _FIELD_ID_DEPTH_LIMIT + depth + + +def _map_key_id(field_id, depth): + return _FIELD_ID_BASE - field_id * _FIELD_ID_DEPTH_LIMIT - depth + + +def _map_value_id(field_id, depth): + return _FIELD_ID_BASE + field_id * _FIELD_ID_DEPTH_LIMIT + depth + + +def _contains_type(data_type, predicate): + if predicate(data_type): + return True + if isinstance(data_type, RowType): + return any(_contains_type(field.type, predicate) + for field in data_type.fields) + if isinstance(data_type, (ArrayType, VectorType, MultisetType)): + return _contains_type(data_type.element, predicate) + if isinstance(data_type, MapType): + return (_contains_type(data_type.key, predicate) + or _contains_type(data_type.value, predicate)) + return False + + +def _is_variant(data_type): + return (isinstance(data_type, AtomicType) + and data_type.type.upper() == "VARIANT") + + +def _is_blob(data_type): + return (isinstance(data_type, AtomicType) + and data_type.type.upper() == "BLOB") diff --git a/paimon-python/pypaimon/write/writer/data_writer.py b/paimon-python/pypaimon/write/writer/data_writer.py index 24a67186c3d8..aa9bbc6d8906 100644 --- a/paimon-python/pypaimon/write/writer/data_writer.py +++ b/paimon-python/pypaimon/write/writer/data_writer.py @@ -29,6 +29,7 @@ from pypaimon.schema.data_types import PyarrowFieldParser from pypaimon.table.bucket_mode import BucketMode from pypaimon.table.row.generic_row import GenericRow +from pypaimon.write.map_shared_shredding_writer import MapSharedShreddingWriter from pypaimon.write.writer.mosaic_writer_options import create_mosaic_writer_options from pypaimon.write.writer.write_buffer import WriteBuffer @@ -104,6 +105,14 @@ def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, op # Paimon field id map, used by _apply_variant_shredding; built once since # the table schema is fixed for the lifetime of this writer. self._paimon_field_id: Dict[str, int] = {pf.name: pf.id for pf in self.table.fields} + self._map_shared_shredding = MapSharedShreddingWriter( + self.table.fields, + self.options, + self.file_format, + self.changelog_file_format + if self.changelog_producer == ChangelogProducer.INPUT else None, + self.bucket, + ) # Set by the composite writers when a flush landed its normal data file but a # later phase of the same flush failed; see their ``_close_current_writers``. @@ -259,8 +268,12 @@ def _write_data_to_file(self, data: pa.Table): extra_files = [] row_sidecar_path = None changelog_meta = None + shared_shredding_stats = {} if self._variant_shredding: data = self._apply_variant_shredding(data) + if self._map_shared_shredding.is_active(): + data, shared_shredding_stats = \ + self._map_shared_shredding.convert(data) # One data file means up to three files on disk -- the data file, its row # sidecar and its changelog -- and none of them is committed until all of @@ -300,7 +313,7 @@ def _write_data_to_file(self, data: pa.Table): # min key & max key - selected_table = data.select(self.trimmed_primary_keys) + selected_table = logical_data.select(self.trimmed_primary_keys) key_columns_batch = selected_table.to_batches()[0] min_key_row_batch = key_columns_batch.slice(0, 1) max_key_row_batch = key_columns_batch.slice(key_columns_batch.num_rows - 1, 1) @@ -311,21 +324,23 @@ def _write_data_to_file(self, data: pa.Table): value_stats_enabled = self.options.metadata_stats_enabled() if value_stats_enabled: stats_fields = self.table.fields if self.table.is_primary_key_table \ - else PyarrowFieldParser.to_paimon_schema(data.schema) + else PyarrowFieldParser.to_paimon_schema(logical_data.schema) else: stats_fields = self.table.trimmed_primary_keys_fields column_stats = { - field.name: self._get_column_stats(data, field.name) + field.name: self._get_column_stats(logical_data, field.name) for field in stats_fields } key_fields = self.trimmed_primary_keys_fields - key_stats = self._collect_value_stats(data, key_fields, column_stats) + key_stats = self._collect_value_stats( + logical_data, key_fields, column_stats) if not self.options.primary_key_nullable() and not all( count == 0 for count in key_stats.null_counts): raise RuntimeError("Primary key should not be null") value_fields = stats_fields if value_stats_enabled else [] - value_stats = self._collect_value_stats(data, value_fields, column_stats) + value_stats = self._collect_value_stats( + logical_data, value_fields, column_stats) # Read the range without advancing it: the advance belongs with the # append below, so a retried flush derives the same range. @@ -368,6 +383,7 @@ def _write_data_to_file(self, data: pa.Table): raise self.sequence_generator.start = self.sequence_generator.current + 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) From ee2d901f50a6347c6277ad85bf02ba674f6ff371 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 15 Sep 2026 02:46:59 -0700 Subject: [PATCH 2/5] [python] Fix shared-shredding writer edge cases --- .../java/org/apache/paimon/JavaPyE2ETest.java | 16 ++++ .../tests/e2e/java_py_read_write_test.py | 11 +++ .../tests/map_shared_shredding_write_test.py | 73 +++++++++++++++++++ .../write/map_shared_shredding_writer.py | 26 +++++-- .../write/writer/data_vector_writer.py | 9 ++- .../pypaimon/write/writer/data_writer.py | 1 - .../write/writer/dedicated_format_writer.py | 7 +- 7 files changed, 133 insertions(+), 10 deletions(-) diff --git a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java index 3e375da1b760..8a59c4332612 100644 --- a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -1634,6 +1634,22 @@ public void testJavaReadSharedShreddingMapTable() throws Exception { reader.forEachRemaining( row -> { int id = row.getInt(0); + for (int column = 2; column <= 3; column++) { + if (id == 3) { + assertThat(row.isNullAt(column)).isTrue(); + } else { + InternalMap required = row.getMap(column); + assertThat(required.size()).isEqualTo(id == 2 ? 0 : 1); + if (id != 2) { + InternalArray values = required.valueArray(); + long value = + column == 2 + ? values.getLong(0) + : values.getRow(0, 1).getLong(0); + assertThat(value).isEqualTo(id == 1 ? 1L : 2L); + } + } + } if (row.isNullAt(1)) { rows.put(id, null); return; diff --git a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py index ab3c88134f06..d85020627785 100644 --- a/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py +++ b/paimon-python/pypaimon/tests/e2e/java_py_read_write_test.py @@ -1640,6 +1640,11 @@ def test_write_shared_shredding_map_for_java(self): schema = pa.schema([ pa.field('id', pa.int32()), pa.field('metrics', pa.map_(pa.string(), pa.int64())), + pa.field('required_values', pa.map_(pa.string(), pa.field( + 'value', pa.int64(), nullable=False))), + pa.field('required_rows', pa.map_(pa.string(), pa.field( + 'value', pa.struct([pa.field('score', pa.int64(), nullable=False)]), + nullable=False))), ]) rows = pa.Table.from_pydict({ 'id': [1, 2, 3, 4], @@ -1649,6 +1654,8 @@ def test_write_shared_shredding_map_for_java(self): [], None, ], + 'required_values': [[('a', 1)], [], None, [('b', 2)]], + 'required_rows': [[('a', {'score': 1})], [], None, [('b', {'score': 2})]], }, schema=schema) table_name = 'default.shared_shredding_map_python_test_parquet' @@ -1661,6 +1668,10 @@ def test_write_shared_shredding_map_for_java(self): 'write-only': 'true', 'fields.metrics.map.storage-layout': 'shared-shredding', 'fields.metrics.map.shared-shredding.max-columns': '2', + 'fields.required_values.map.storage-layout': 'shared-shredding', + 'fields.required_values.map.shared-shredding.max-columns': '2', + 'fields.required_rows.map.storage-layout': 'shared-shredding', + 'fields.required_rows.map.shared-shredding.max-columns': '2', }), False, ) 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 57e72419c739..bbc5be506d78 100644 --- a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py +++ b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py @@ -100,6 +100,79 @@ def test_reject_orc(self): "metrics": [[("key", 1)]], }, schema=self.arrow_schema)) + def test_non_nullable_values(self): + for value_type, value in [ + (pa.int64(), 1), + (pa.struct([pa.field('score', pa.int64(), nullable=False)]), + {'score': 1})]: + with self.subTest(value_type=value_type): + self.arrow_schema = pa.schema([ + pa.field('id', pa.int32()), + pa.field('metrics', pa.map_( + pa.string(), pa.field('value', value_type, nullable=False))), + ]) + expected = [[('a', value)], [], None] + data = pa.Table.from_pydict({ + 'id': [1, 2, 3], 'metrics': expected, + }, schema=self.arrow_schema) + table = self._create_table('parquet', 2) + messages = self._write(table, data) + field = pq.read_schema(messages[0].new_files[0].file_path).field('metrics') + self.assertTrue(field.type['__col_0'].nullable) + self.assertTrue(field.type['__col_1'].nullable) + reader = table.new_read_builder() + result = reader.new_read().to_arrow(reader.new_scan().plan().splits()) + # Arrow 6 cannot convert non-nullable MAP values to scalars. + actual = result.column('metrics').cast(pa.map_(pa.string(), value_type)) + self.assertEqual(expected, actual.to_pylist()) + + def test_dedicated_columns_preserve_shredding(self): + for blob, vector in [(True, False), (False, True), (True, True)]: + with self.subTest(blob=blob, vector=vector): + fields = [pa.field('id', pa.int32()), + pa.field('metrics', pa.map_(pa.string(), pa.int64()))] + values = {'id': [1, 2], 'metrics': [[('a', 1)], []]} + options = {'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true'} + if blob: + fields.append(pa.field('payload', pa.large_binary())) + values['payload'] = [b'one', b'two'] + if vector: + fields.append(pa.field('embedding', pa.list_(pa.float32(), 2))) + values['embedding'] = [[1., 2.], [3., 4.]] + options['vector.file.format'] = 'parquet' + self.arrow_schema = pa.schema(fields) + table = self._create_table('parquet', 2, options) + messages = self._write(table, pa.Table.from_pydict( + values, schema=self.arrow_schema)) + normal_files = [f for m in messages for f in m.new_files + if f.file_name.endswith('.parquet') + and '.vector.' not in f.file_name] + self.assertTrue(normal_files) + for file in normal_files: + self.assertTrue(is_shared_shredding( + pq.read_schema(file.file_path).field('metrics'))) + reader = table.new_read_builder().with_projection(["metrics['a']"]) + result = reader.new_read().to_arrow(reader.new_scan().plan().splits()) + self.assertEqual([1, None], result.column('metrics_a').to_pylist()) + + def test_reject_default_layout_on_non_map(self): + table = self._create_table('parquet', 2, { + 'fields.metrics.map.storage-layout': 'default', + 'fields.id.map.storage-layout': 'default', + }) + with self.assertRaisesRegex(ValueError, 'its type is not MAP'): + self._write(table, pa.Table.from_pydict({ + 'id': [1], 'metrics': [[('a', 1)]], + }, schema=self.arrow_schema)) + + def test_reject_postpone_with_fixed_output_bucket(self): + from pypaimon.write.writer.append_only_data_writer import AppendOnlyDataWriter + + table = self._create_table('parquet', 2, {'bucket': '-2'}) + with self.assertRaisesRegex(ValueError, 'postpone bucket'): + AppendOnlyDataWriter(table, (), 0, 0, table.options) + def test_adapts_physical_column_count_between_files(self): table = self._create_table( "parquet", max_columns=4, diff --git a/paimon-python/pypaimon/write/map_shared_shredding_writer.py b/paimon-python/pypaimon/write/map_shared_shredding_writer.py index 25fcd885b15d..432fffaf6edc 100644 --- a/paimon-python/pypaimon/write/map_shared_shredding_writer.py +++ b/paimon-python/pypaimon/write/map_shared_shredding_writer.py @@ -22,7 +22,10 @@ import pyarrow as pa -from pypaimon.data.map_shared_shredding import shared_shredding_metadata +from pypaimon.data.map_shared_shredding import ( + _normalized_offsets, + shared_shredding_metadata, +) from pypaimon.schema.data_types import ( ArrayType, AtomicType, @@ -44,7 +47,7 @@ class MapSharedShreddingWriter: """Converts configured logical MAP columns before each file write.""" - def __init__(self, fields, options, file_format, changelog_format, bucket): + def __init__(self, fields, options, file_format, changelog_format): self._options = options field_by_name = {field.name: field for field in fields} configured = self._configured_fields() @@ -54,6 +57,10 @@ def __init__(self, fields, options, file_format, changelog_format, bucket): self._recent_widths = {} for name, field in field_by_name.items(): + if name in configured and not isinstance(field.type, MapType): + raise ValueError( + "Column '{}' is configured with map.storage-layout but " + "its type is not MAP.".format(name)) layout = options.map_storage_layout(name) if layout not in ("default", "shared-shredding"): raise ValueError( @@ -80,7 +87,7 @@ def __init__(self, fields, options, file_format, changelog_format, bucket): self._validate_format("file.format", file_format) self._validate_format("changelog.file.format", changelog_format) self._validate_compression("file.compression", options.file_compression()) - if bucket == BucketMode.POSTPONE_BUCKET.value: + if options.bucket() == BucketMode.POSTPONE_BUCKET.value: raise ValueError( "MAP shared-shredding does not support postpone bucket mode.") if any(_contains_type(field.type, _is_variant) @@ -228,9 +235,16 @@ def __init__(self, num_columns, policy, item_type, logical_item_type): num_columns, item_type, logical_item_type) def convert(self, column): + # Avoid MapScalar conversion: older Arrow versions cannot represent + # MAP scalars with non-nullable values. + offsets, start, end = _normalized_offsets(column) + keys = column.keys.slice(start, end - start).to_pylist() + values = column.items.slice(start, end - start).to_pylist() rows = [] - for value in column.to_pylist(): - rows.append(None if value is None else self._convert_map(value)) + for index, is_null in enumerate(column.is_null().to_pylist()): + start, end = offsets[index:index + 2] + rows.append(None if is_null else self._convert_map( + zip(keys[start:end], values[start:end]))) return pa.array(rows, type=self.physical_type) def _convert_map(self, value): @@ -340,7 +354,7 @@ def _physical_struct_type(num_columns, item_type, logical_item_type): fields.append(_field_with_ids( _PHYSICAL_COLUMN_PREFIX + str(column_id), item_type, - logical_item_type.nullable, + True, logical_item_type, column_id + 1, )) diff --git a/paimon-python/pypaimon/write/writer/data_vector_writer.py b/paimon-python/pypaimon/write/writer/data_vector_writer.py index 9294f67d92cb..0942a7bfd5a5 100644 --- a/paimon-python/pypaimon/write/writer/data_vector_writer.py +++ b/paimon-python/pypaimon/write/writer/data_vector_writer.py @@ -261,6 +261,9 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: if data.num_rows == 0: return None + logical_data = data + data, shredding_stats = self._map_shared_shredding.convert(data) + file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}" file_path = self._generate_file_path(file_name) @@ -286,11 +289,11 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: metadata_stats_enabled = self.options.metadata_stats_enabled() stats_columns = self.normal_columns if metadata_stats_enabled else [] - value_stats = self._collect_value_stats(data, stats_columns) + value_stats = self._collect_value_stats(logical_data, stats_columns) min_seq, max_seq = self._append_file_sequence_range(data.num_rows) - return DataFileMeta.create( + meta = DataFileMeta.create( file_name=file_name, file_size=self.file_io.get_file_size(file_path), row_count=data.num_rows, @@ -311,6 +314,8 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: file_path=file_path, write_cols=self.write_cols, ) + self._map_shared_shredding.file_completed(shredding_stats) + return meta def _validate_consistency( self, normal_meta: DataFileMeta, vector_metas: List[DataFileMeta]): diff --git a/paimon-python/pypaimon/write/writer/data_writer.py b/paimon-python/pypaimon/write/writer/data_writer.py index aa9bbc6d8906..4dc6693a9f6d 100644 --- a/paimon-python/pypaimon/write/writer/data_writer.py +++ b/paimon-python/pypaimon/write/writer/data_writer.py @@ -111,7 +111,6 @@ def __init__(self, table, partition: Tuple, bucket: int, max_seq_number: int, op self.file_format, self.changelog_file_format if self.changelog_producer == ChangelogProducer.INPUT else None, - self.bucket, ) # Set by the composite writers when a flush landed its normal data file but a diff --git a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py index c894182fbe0c..04d8fda8a21f 100644 --- a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py +++ b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py @@ -703,6 +703,9 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: if data.num_rows == 0: return None + logical_data = data + data, shredding_stats = self._map_shared_shredding.convert(data) + file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}" file_path = self._generate_file_path(file_name) @@ -728,7 +731,9 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: is_external_path = self.external_path_provider is not None external_path_str = file_path if is_external_path else None - return self._create_data_file_meta(file_name, file_path, data, external_path_str) + meta = self._create_data_file_meta(file_name, file_path, logical_data, external_path_str) + self._map_shared_shredding.file_completed(shredding_stats) + return meta def _create_data_file_meta(self, file_name: str, file_path: str, data: pa.Table, external_path: Optional[str] = None) -> DataFileMeta: From 7076ba7c14da472492e3a0f49169009c91eab066 Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 15 Sep 2026 03:10:01 -0700 Subject: [PATCH 3/5] [python] Avoid nested MAP scalar conversion in shared-shredding writes --- .../tests/map_shared_shredding_write_test.py | 49 +++++++++++++++++++ .../write/map_shared_shredding_writer.py | 42 ++++++++++++---- 2 files changed, 81 insertions(+), 10 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 bbc5be506d78..94bf83082075 100644 --- a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py +++ b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py @@ -16,6 +16,8 @@ # under the License. import shutil +import subprocess +import sys import tempfile import unittest @@ -126,6 +128,53 @@ def test_non_nullable_values(self): actual = result.column('metrics').cast(pa.map_(pa.string(), value_type)) self.assertEqual(expected, actual.to_pylist()) + def test_nested_non_nullable_map_values_in_subprocess(self): + # A native Arrow assertion aborts the process, not a Python exception. + result = subprocess.run([ + sys.executable, '-c', + 'from pypaimon.tests.map_shared_shredding_write_test import ' + 'MapSharedShreddingWriteTest\n' + 'test = MapSharedShreddingWriteTest()\n' + 'test.setUp()\n' + 'try:\n' + ' test.check_nested_non_nullable_map_values()\n' + 'finally:\n' + ' test.tearDown()\n', + ], stdout=subprocess.PIPE, stderr=subprocess.PIPE, timeout=60) + self.assertEqual(0, result.returncode, result.stderr.decode('utf-8', 'replace')) + + def check_nested_non_nullable_map_values(self): + from pypaimon.write.map_shared_shredding_writer import _to_python_values + + inner = pa.map_(pa.string(), pa.field('value', pa.int64(), nullable=False)) + for value_type, value in [ + (inner, [('x', 1)]), + (pa.struct([pa.field('nested', inner)]), {'nested': [('x', 1)]}), + (pa.list_(inner), [[('x', 1)], [], None])]: + self.arrow_schema = pa.schema([ + pa.field('id', pa.int32()), + pa.field('metrics', pa.map_(pa.string(), value_type)), + ]) + expected = [[('a', value)], [], None, [('b', value)]] + array = pa.array([[('unused', value)]] + expected + [[]], + type=self.arrow_schema.field('metrics').type) + # Keep nonzero offsets and multiple chunks in the conversion check. + chunks = [array.slice(1, 2), array.slice(3, 2)] + self.assertEqual(expected, [v for chunk in chunks + for v in _to_python_values(chunk)]) + data = pa.Table.from_arrays([ + pa.array([1, 2, 3, 4], type=pa.int32()), + pa.chunked_array(chunks), + ], schema=self.arrow_schema) + table = self._create_table('parquet', 2) + messages = self._write(table, data) + self.assertTrue(is_shared_shredding(pq.read_schema( + messages[0].new_files[0].file_path).field('metrics'))) + reader = table.new_read_builder() + result = reader.new_read().to_arrow(reader.new_scan().plan().splits()) + self.assertEqual(expected, [v for chunk in result.column('metrics').chunks + for v in _to_python_values(chunk)]) + def test_dedicated_columns_preserve_shredding(self): for blob, vector in [(True, False), (False, True), (True, True)]: with self.subTest(blob=blob, vector=vector): diff --git a/paimon-python/pypaimon/write/map_shared_shredding_writer.py b/paimon-python/pypaimon/write/map_shared_shredding_writer.py index 432fffaf6edc..79192229287f 100644 --- a/paimon-python/pypaimon/write/map_shared_shredding_writer.py +++ b/paimon-python/pypaimon/write/map_shared_shredding_writer.py @@ -235,16 +235,8 @@ def __init__(self, num_columns, policy, item_type, logical_item_type): num_columns, item_type, logical_item_type) def convert(self, column): - # Avoid MapScalar conversion: older Arrow versions cannot represent - # MAP scalars with non-nullable values. - offsets, start, end = _normalized_offsets(column) - keys = column.keys.slice(start, end - start).to_pylist() - values = column.items.slice(start, end - start).to_pylist() - rows = [] - for index, is_null in enumerate(column.is_null().to_pylist()): - start, end = offsets[index:index + 2] - rows.append(None if is_null else self._convert_map( - zip(keys[start:end], values[start:end]))) + rows = [None if value is None else self._convert_map(value) + for value in _to_python_values(column)] return pa.array(rows, type=self.physical_type) def _convert_map(self, value): @@ -339,6 +331,36 @@ def _select_lru_column(self, used, resident): return selected +def _to_python_values(column): + """Avoid Arrow 6 MAP scalars, including MAPs nested in ROW/ARRAY values.""" + data_type = column.type + if pa.types.is_struct(data_type): + children = [_to_python_values(column.field(i)) + for i in range(len(data_type))] + names = [field.name for field in data_type] + values = [dict(zip(names, (child[i] for child in children))) + for i in range(len(column))] + elif (pa.types.is_map(data_type) or pa.types.is_list(data_type) + or pa.types.is_large_list(data_type)): + offsets, start, end = _normalized_offsets(column) + if pa.types.is_map(data_type): + keys = _to_python_values(column.keys.slice(start, end - start)) + items = _to_python_values(column.items.slice(start, end - start)) + children = list(zip(keys, items)) + else: + children = _to_python_values(column.values.slice(start, end - start)) + values = [children[offsets[i]:offsets[i + 1]] for i in range(len(column))] + elif pa.types.is_fixed_size_list(data_type): + size = data_type.list_size + children = _to_python_values(column.values.slice( + column.offset * size, len(column) * size)) + values = [children[i * size:(i + 1) * size] for i in range(len(column))] + else: + return column.to_pylist() + return [None if is_null else value + for is_null, value in zip(column.is_null().to_pylist(), values)] + + def _physical_struct_type(num_columns, item_type, logical_item_type): mapping_type = pa.list_(pa.field( "item", From 95ef49cf7ff3b6e7116835f419ad8e9ac7daa86c Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 15 Sep 2026 04:06:45 -0700 Subject: [PATCH 4/5] [python] Bound shared-shredding conversion batches and reject CHAR keys --- .../tests/map_shared_shredding_write_test.py | 67 +++++++++ .../write/map_shared_shredding_writer.py | 133 ++++++++++++++---- .../write/writer/data_vector_writer.py | 7 +- .../pypaimon/write/writer/data_writer.py | 15 +- .../write/writer/dedicated_format_writer.py | 7 +- 5 files changed, 184 insertions(+), 45 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 94bf83082075..b9a91c6ce1fd 100644 --- a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py +++ b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py @@ -20,6 +20,7 @@ import sys import tempfile import unittest +from unittest.mock import patch import pyarrow as pa import pyarrow.parquet as pq @@ -215,6 +216,72 @@ def test_reject_default_layout_on_non_map(self): 'id': [1], 'metrics': [[('a', 1)]], }, schema=self.arrow_schema)) + def test_reject_char_keys(self): + from pypaimon.schema.data_types import AtomicType, DataField, MapType + from pypaimon.write.map_shared_shredding_writer import MapSharedShreddingWriter + + with self.assertRaisesRegex(ValueError, 'STRING keys'): + MapSharedShreddingWriter._validate_field(DataField( + 0, 'metrics', MapType(True, AtomicType('CHAR(4)', False), AtomicType('BIGINT')))) + + def test_streams_physical_batches(self): + from pypaimon.write.map_shared_shredding_writer import _MapFieldConverter + + table = self._create_table('parquet', 256) + data = pa.Table.from_pydict({ + 'id': list(range(5000)), 'metrics': [[('a', 1)]] * 4999 + [[('late', 2)]], + }, schema=self.arrow_schema) + convert = _MapFieldConverter.convert + write_table = pq.ParquetWriter.write_table + pending = [] + sizes = [] + + 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() + return write_table(writer, physical, *args, **kwargs) + + with patch.object(_MapFieldConverter, 'convert', convert_batch), \ + patch.object(pq.ParquetWriter, 'write_table', write_batch): + messages = self._write(table, data) + self.assertEqual(5000, sum(sizes)) + self.assertLessEqual(max(sizes), 1024) + 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()) + self.assertEqual([1] * 4999 + [None], result.column('metrics_a').to_pylist()) + self.assertEqual([None] * 4999 + [2], result.column('metrics_late').to_pylist()) + + def test_failed_stream_removes_partial_file(self): + from pypaimon.write.map_shared_shredding_writer import MapSharedShreddingWriter + + table = self._create_table('parquet', 2) + converter = MapSharedShreddingWriter(table.fields, table.options, 'parquet', None) + data = pa.Table.from_pydict({ + 'id': list(range(2048)), 'metrics': [[('a', 1)]] * 2048, + }, schema=self.arrow_schema) + path = self.temp_dir + '/partial.parquet' + write_table = pq.ParquetWriter.write_table + calls = [] + + def fail_second_batch(writer, physical, *args, **kwargs): + calls.append(physical.num_rows) + if len(calls) == 2: + raise OSError('injected write failure') + return write_table(writer, physical, *args, **kwargs) + + with patch.object(pq.ParquetWriter, 'write_table', fail_second_batch): + with self.assertRaisesRegex(OSError, 'injected write failure'): + converter.write_parquet(table.file_io, path, data, 'zstd', 1) + self.assertEqual(2, len(calls)) + self.assertFalse(table.file_io.exists(path)) + 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 79192229287f..3ad4e7789c3b 100644 --- a/paimon-python/pypaimon/write/map_shared_shredding_writer.py +++ b/paimon-python/pypaimon/write/map_shared_shredding_writer.py @@ -18,6 +18,7 @@ """Java-compatible shared-shredding MAP conversion for data files.""" import math +from array import array from collections import deque import pyarrow as pa @@ -42,6 +43,7 @@ _PHYSICAL_COLUMN_PREFIX = "__col_" _FIELD_ID_BASE = 2147483647 // 4 _FIELD_ID_DEPTH_LIMIT = 1 << 10 +_CONVERSION_BYTES = 8 * 1024 * 1024 class MapSharedShreddingWriter: @@ -102,12 +104,13 @@ def __init__(self, fields, options, file_format, changelog_format): def is_active(self): return bool(self._fields) - def convert(self, data): - """Return physical data and completed-file statistics.""" - columns = list(data.columns) + def write_parquet(self, file_io, path, data, compression, zstd_level): + """Plan key metadata, then write bounded physical batches without retaining them.""" + import pyarrow.parquet as pq + fields = list(data.schema) completed = {} - changed = False + converters = {} field_index = {field.name: index for index, field in enumerate(fields)} for name, logical_field in self._fields.items(): @@ -125,7 +128,11 @@ def convert(self, data): input_field.type.item_type, logical_field.type.value, ) - chunks = [converter.convert(chunk) for chunk in columns[index].chunks] + # Parquet's Arrow schema metadata must be known before opening the + # writer. Scan only keys; reset placement for the actual write pass. + for chunk in data.column(index).chunks: + for start in range(0, len(chunk), 1024): + converter.analyze(chunk.slice(start, 1024)) metadata = dict(input_field.metadata or {}) metadata.update(shared_shredding_metadata( converter.name_to_id, @@ -142,14 +149,33 @@ def convert(self, data): nullable=input_field.nullable, metadata=metadata, ) - columns[index] = pa.chunked_array( - chunks, type=converter.physical_type) completed[name] = converter.max_row_width - changed = True - - if not changed: - return data, completed - return pa.Table.from_arrays(columns, schema=pa.schema(fields)), completed + converters[index] = _MapFieldConverter( + num_columns, self._policies[name], input_field.type.item_type, + logical_field.type.value) + + schema = pa.schema(fields, metadata=data.schema.metadata) + # Bound scratch space for slot mappings/indices as well as logical values. + slots = sum(c.num_columns for c in converters.values()) + batch_rows = max(1, min(1024, _CONVERSION_BYTES // max(1, slots * 16))) + kwargs = {'compression': compression} + if compression.lower() == 'zstd': + kwargs['compression_level'] = zstd_level + try: + with file_io.new_output_stream(path) as stream: + with pq.ParquetWriter(stream, schema, **kwargs) as writer: + for batch in data.to_batches(max_chunksize=batch_rows): + for bounded in _bounded_batches(batch): + columns = list(bounded.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 + except Exception: + file_io.delete_quietly(path) + raise + return completed def file_completed(self, completed): for name, max_row_width in completed.items(): @@ -189,7 +215,7 @@ def _validate_field(field): key_type = field.type.key if (not isinstance(key_type, AtomicType) or not key_type.type.upper().startswith( - ("STRING", "VARCHAR", "CHAR"))): + ("STRING", "VARCHAR"))): raise ValueError( "Shared-shredding column '{}' must use STRING keys.".format( field.name)) @@ -235,38 +261,72 @@ def __init__(self, num_columns, policy, item_type, logical_item_type): num_columns, item_type, logical_item_type) def convert(self, column): - rows = [None if value is None else self._convert_map(value) - for value in _to_python_values(column)] - return pa.array(rows, type=self.physical_type) + offsets, start, end = _normalized_offsets(column) + keys = column.keys.slice(start, end - start).to_pylist() + values = _to_python_values(column.items.slice(start, end - start)) + nulls = column.is_null().to_pylist() + mappings = array('i') + mapping_offsets = [0] + slots = {} + overflows = [] + for row, is_null in enumerate(nulls): + start, end = offsets[row:row + 2] + if is_null: + mapping_offsets.append(len(mappings)) + overflows.append(None) + continue + field_ids, mapping, overflow = self._place(keys[start:end]) + items = dict(zip(field_ids, values[start:end])) + mappings.extend(mapping) + mapping_offsets.append(len(mappings)) + for column_id, field_id in enumerate(mapping): + if field_id >= 0: + if column_id not in slots: + slots[column_id] = [None] * len(column) + slots[column_id][row] = items[field_id] + overflows.append([(field_id, items[field_id]) for field_id in overflow] + if overflow else None) + children = [pa.ListArray.from_arrays( + pa.array(mapping_offsets, type=pa.int32()), + pa.array(mappings, type=pa.int32()))] + empty = None + for column_id in range(self.num_columns): + value_type = self.physical_type[column_id + 1].type + if column_id in slots: + children.append(pa.array(slots.pop(column_id), type=value_type)) + else: + if empty is None: + empty = pa.array([None] * len(column), type=value_type) + children.append(empty) + children.append(pa.array(overflows, type=self.physical_type[-1].type)) + return pa.StructArray.from_arrays( + children, fields=list(self.physical_type), + mask=column.is_null() if column.null_count else None) + + def analyze(self, column): + offsets, start, end = _normalized_offsets(column) + keys = column.keys.slice(start, end - start).to_pylist() + for row, is_null in enumerate(column.is_null().to_pylist()): + if not is_null: + self._place(keys[offsets[row]:offsets[row + 1]]) - def _convert_map(self, value): - entries = list(value.items()) if isinstance(value, dict) else list(value) + def _place(self, keys): field_ids = [] - values = {} - for key, item in entries: + for key in keys: if key is None: raise ValueError("Shared-shredding MAP keys cannot be null") if not isinstance(key, str): raise TypeError("Shared-shredding MAP keys must be strings") field_id = self.name_to_id.setdefault(key, len(self.name_to_id)) field_ids.append(field_id) - values[field_id] = item mapping, overflow = self._allocate(field_ids) - row = {_FIELD_MAPPING: mapping} for column_id, field_id in enumerate(mapping): - row[_PHYSICAL_COLUMN_PREFIX + str(column_id)] = ( - None if field_id < 0 else values[field_id]) if field_id >= 0: self.field_to_columns.setdefault(field_id, set()).add(column_id) - if overflow: - row[_OVERFLOW] = [(field_id, values[field_id]) - for field_id in overflow] - self.overflow_set.update(overflow) - else: - row[_OVERFLOW] = None + self.overflow_set.update(overflow) self.max_row_width = max(self.max_row_width, len(field_ids)) - return row + return field_ids, mapping, overflow def _allocate(self, field_ids): if self.policy == "plain": @@ -331,6 +391,17 @@ def _select_lru_column(self, used, resident): return selected +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. + if batch.nbytes > _CONVERSION_BYTES and batch.num_rows > 1: + middle = batch.num_rows // 2 + yield from _bounded_batches(batch.slice(0, middle)) + yield from _bounded_batches(batch.slice(middle)) + else: + yield batch + + def _to_python_values(column): """Avoid Arrow 6 MAP scalars, including MAPs nested in ROW/ARRAY values.""" data_type = column.type diff --git a/paimon-python/pypaimon/write/writer/data_vector_writer.py b/paimon-python/pypaimon/write/writer/data_vector_writer.py index 0942a7bfd5a5..4825b22af138 100644 --- a/paimon-python/pypaimon/write/writer/data_vector_writer.py +++ b/paimon-python/pypaimon/write/writer/data_vector_writer.py @@ -261,14 +261,13 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: if data.num_rows == 0: return None - logical_data = data - data, shredding_stats = self._map_shared_shredding.convert(data) + shredding_stats = {} file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}" file_path = self._generate_file_path(file_name) if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: - self.file_io.write_parquet(file_path, data, compression=self.compression, zstd_level=self.zstd_level) + shredding_stats = self._write_parquet_data(file_path, data) elif self.file_format == CoreOptions.FILE_FORMAT_ORC: self.file_io.write_orc(file_path, data, compression=self.compression, zstd_level=self.zstd_level) elif self.file_format == CoreOptions.FILE_FORMAT_AVRO: @@ -289,7 +288,7 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: metadata_stats_enabled = self.options.metadata_stats_enabled() stats_columns = self.normal_columns if metadata_stats_enabled else [] - value_stats = self._collect_value_stats(logical_data, stats_columns) + value_stats = self._collect_value_stats(data, stats_columns) min_seq, max_seq = self._append_file_sequence_range(data.num_rows) diff --git a/paimon-python/pypaimon/write/writer/data_writer.py b/paimon-python/pypaimon/write/writer/data_writer.py index 4dc6693a9f6d..46db0578ceb8 100644 --- a/paimon-python/pypaimon/write/writer/data_writer.py +++ b/paimon-python/pypaimon/write/writer/data_writer.py @@ -270,9 +270,6 @@ def _write_data_to_file(self, data: pa.Table): shared_shredding_stats = {} if self._variant_shredding: data = self._apply_variant_shredding(data) - if self._map_shared_shredding.is_active(): - data, shared_shredding_stats = \ - self._map_shared_shredding.convert(data) # One data file means up to three files on disk -- the data file, its row # sidecar and its changelog -- and none of them is committed until all of @@ -282,7 +279,7 @@ def _write_data_to_file(self, data: pa.Table): # already covers. try: if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: - self.file_io.write_parquet(file_path, data, compression=self.compression, zstd_level=self.zstd_level) + shared_shredding_stats = self._write_parquet_data(file_path, data) elif self.file_format == CoreOptions.FILE_FORMAT_ORC: self.file_io.write_orc(file_path, data, compression=self.compression, zstd_level=self.zstd_level) elif self.file_format == CoreOptions.FILE_FORMAT_AVRO: @@ -387,6 +384,13 @@ def _write_data_to_file(self, data: pa.Table): if changelog_meta is not None: self.committed_changelog_files.append(changelog_meta) + def _write_parquet_data(self, path, data): + if self._map_shared_shredding.is_active(): + return self._map_shared_shredding.write_parquet( + self.file_io, path, data, self.compression, self.zstd_level) + self.file_io.write_parquet(path, data, compression=self.compression, zstd_level=self.zstd_level) + return {} + def _apply_variant_shredding(self, data: pa.Table) -> pa.Table: """Transform VARIANT columns into shredded Parquet format. @@ -429,8 +433,7 @@ def _write_changelog_file(self, data, min_key, max_key, key_stats, value_stats, try: if cl_fmt == CoreOptions.FILE_FORMAT_PARQUET: - self.file_io.write_parquet(changelog_file_path, data, compression=self.compression, - zstd_level=self.zstd_level) + self._write_parquet_data(changelog_file_path, data) elif cl_fmt == CoreOptions.FILE_FORMAT_ORC: self.file_io.write_orc(changelog_file_path, data, compression=self.compression, zstd_level=self.zstd_level) diff --git a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py index 04d8fda8a21f..f79369d1cdff 100644 --- a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py +++ b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py @@ -703,15 +703,14 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: if data.num_rows == 0: return None - logical_data = data - data, shredding_stats = self._map_shared_shredding.convert(data) + shredding_stats = {} file_name = f"{CoreOptions.data_file_prefix(self.options)}{uuid.uuid4()}-0.{self.file_format}" file_path = self._generate_file_path(file_name) # Write file based on format if self.file_format == CoreOptions.FILE_FORMAT_PARQUET: - self.file_io.write_parquet(file_path, data, compression=self.compression, zstd_level=self.zstd_level) + shredding_stats = self._write_parquet_data(file_path, data) elif self.file_format == CoreOptions.FILE_FORMAT_ORC: self.file_io.write_orc(file_path, data, compression=self.compression, zstd_level=self.zstd_level) elif self.file_format == CoreOptions.FILE_FORMAT_AVRO: @@ -731,7 +730,7 @@ def _write_normal_data_to_file(self, data: pa.Table) -> Optional[DataFileMeta]: is_external_path = self.external_path_provider is not None external_path_str = file_path if is_external_path else None - meta = self._create_data_file_meta(file_name, file_path, logical_data, external_path_str) + meta = self._create_data_file_meta(file_name, file_path, data, external_path_str) self._map_shared_shredding.file_completed(shredding_stats) return meta From fec9335eb28cb27f6b1d984c210ebe7f3074f0bd Mon Sep 17 00:00:00 2001 From: xiaohongbo Date: Tue, 15 Sep 2026 23:52:23 -0700 Subject: [PATCH 5/5] [python] Coalesce bounded MAP input chunks before writing row groups --- .../tests/map_shared_shredding_write_test.py | 32 +++++++++++++++++++ .../write/map_shared_shredding_writer.py | 9 ++++-- 2 files changed, 38 insertions(+), 3 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 b9a91c6ce1fd..eba20b427037 100644 --- a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py +++ b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py @@ -282,6 +282,38 @@ def fail_second_batch(writer, physical, *args, **kwargs): self.assertEqual(2, len(calls)) self.assertFalse(table.file_io.exists(path)) + def test_row_groups_do_not_follow_input_calls(self): + from pypaimon.table.row.generic_row import GenericRow + + for count in (1000, 2500): + layouts = [] + for by_row in (False, True): + table = self._create_table('parquet', 256) + builder = table.new_batch_write_builder() + writer = builder.new_write() + if by_row: + for i in range(count): + writer.write_row(GenericRow([i, [('a', i)]], table.fields)) + else: + writer.write_arrow(pa.Table.from_pydict({ + 'id': list(range(count)), + 'metrics': [[('a', i)] for i in range(count)], + }, schema=self.arrow_schema)) + messages = writer.prepare_commit() + builder.new_commit().commit(messages) + writer.close() + files = [f for m in messages for f in m.new_files] + self.assertEqual(1, len(files)) + 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) + 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_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 3ad4e7789c3b..07f21edf5c61 100644 --- a/paimon-python/pypaimon/write/map_shared_shredding_writer.py +++ b/paimon-python/pypaimon/write/map_shared_shredding_writer.py @@ -164,9 +164,12 @@ def write_parquet(self, file_io, path, data, compression, zstd_level): try: with file_io.new_output_stream(path) as stream: with pq.ParquetWriter(stream, schema, **kwargs) as writer: - for batch in data.to_batches(max_chunksize=batch_rows): - for bounded in _bounded_batches(batch): - columns = list(bounded.columns) + 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)