This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 478a7201c451fa8bc96f81185457f18c062f8c18 Author: XiaoHongbo <[email protected]> AuthorDate: Wed Sep 16 15:22:02 2026 +0800 [python] Support writing shared-shredding MAPs in Parquet (#9837) --- .../test/java/org/apache/paimon/JavaPyE2ETest.java | 61 +++ paimon-python/dev/run_mixed_tests.sh | 15 + .../pypaimon/common/options/core_options.py | 44 ++ .../pypaimon/data/map_shared_shredding.py | 51 +- .../pypaimon/tests/e2e/java_py_read_write_test.py | 46 ++ .../tests/map_shared_shredding_write_test.py | 385 ++++++++++++++ .../pypaimon/write/map_shared_shredding_writer.py | 573 +++++++++++++++++++++ .../pypaimon/write/writer/data_vector_writer.py | 8 +- paimon-python/pypaimon/write/writer/data_writer.py | 34 +- .../write/writer/dedicated_format_writer.py | 8 +- 10 files changed, 1212 insertions(+), 13 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 524344ed54..8a59c43326 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,67 @@ public class JavaPyE2ETest { } } + /** 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<Integer, Map<String, Long>> rows = new HashMap<>(); + List<Split> splits = new ArrayList<>(table.newSnapshotReader().read().dataSplits()); + try (org.apache.paimon.reader.RecordReader<InternalRow> reader = + table.newRead().createReader(splits)) { + 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; + } + InternalMap map = row.getMap(1); + InternalArray keys = map.keyArray(); + InternalArray values = map.valueArray(); + Map<String, Long> 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<Integer, Map<Integer, byte[]>> readMapBlobRows(FileStoreTable table) throws Exception { Map<Integer, Map<Integer, byte[]>> rows = new HashMap<>(); diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh index 620c4d989a..54fef2a4e1 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 6586e051a5..2fe0d039fd 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] = ( @@ -1894,6 +1898,46 @@ class CoreOptions: .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 a3d6c58bac..13c15df75f 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 @@ _FIELD_DICT_ORIGINAL_SIZE = b"paimon.map.shared-shredding.field-dict-original-si _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("<ii", len(payload), len(data)) + payload + raise ValueError( + "Unsupported shared-shredding dictionary compression: {}".format( + compression)) + + def _required(metadata, key): try: return metadata[key] 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 86d1255ecf..d850206277 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 @@ -1636,6 +1636,52 @@ class JavaPyReadWriteTest(unittest.TestCase): self.assertEqual([None] * 5, result.column('metrics_missing').to_pylist()) + 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], + 'metrics': [ + [('hot', 10), ('warm', 20), ('overflow', 30)], + [('hot', None), ('new', 40)], + [], + 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' + self.catalog.drop_table(table_name, True) + self.catalog.create_table( + table_name, + Schema.from_pyarrow_schema(schema, options={ + 'bucket': '-1', + 'file.format': 'parquet', + '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, + ) + table = self.catalog.get_table(table_name) + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(rows) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + def test_write_map_blob_for_java(self): map_blob_type = pa.map_(pa.int32(), pa.large_binary()) boolean_map_blob_type = pa.map_(pa.bool_(), pa.large_binary()) diff --git a/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py new file mode 100644 index 0000000000..eba20b4270 --- /dev/null +++ b/paimon-python/pypaimon/tests/map_shared_shredding_write_test.py @@ -0,0 +1,385 @@ +# 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. + +import shutil +import subprocess +import sys +import tempfile +import unittest +from unittest.mock import patch + +import pyarrow as pa +import pyarrow.parquet as pq + +from pypaimon import CatalogFactory, Schema +from pypaimon.data.map_shared_shredding import ( + is_shared_shredding, + parse_shared_shredding_metadata, +) + + +class MapSharedShreddingWriteTest(unittest.TestCase): + + def setUp(self): + self.temp_dir = tempfile.mkdtemp() + self.catalog = CatalogFactory.create({"warehouse": self.temp_dir}) + self.catalog.create_database("default", True) + self.arrow_schema = pa.schema([ + pa.field("id", pa.int32()), + pa.field("metrics", pa.map_(pa.string(), pa.int64())), + ]) + + def tearDown(self): + shutil.rmtree(self.temp_dir, ignore_errors=True) + + def test_write_and_read_parquet(self): + expected = [ + [("hot", 10), ("warm", 20), ("overflow", 30)], + [("hot", None), ("new", 40)], + [], + None, + ] + data = pa.Table.from_pydict({ + "id": [1, 2, 3, 4], + "metrics": expected, + }, schema=self.arrow_schema) + + table = self._create_table("parquet", max_columns=2) + messages = self._write(table, data) + + physical_field = pq.read_schema( + messages[0].new_files[0].file_path).field("metrics") + self.assertTrue(pa.types.is_struct(physical_field.type)) + self.assertTrue(is_shared_shredding(physical_field)) + name_by_id, num_columns = \ + parse_shared_shredding_metadata(physical_field) + self.assertEqual(2, num_columns) + self.assertEqual( + {"hot", "warm", "overflow", "new"}, + set(name_by_id.values()), + ) + + read_builder = table.new_read_builder() + result = read_builder.new_read().to_arrow( + read_builder.new_scan().plan().splits()) + self.assertEqual(expected, result.column("metrics").to_pylist()) + + selected = table.new_read_builder().with_projection([ + "id", "metrics['hot']", "metrics['overflow']", + ]) + result = selected.new_read().to_arrow( + selected.new_scan().plan().splits()) + self.assertEqual( + [10, None, None, None], + result.column("metrics_hot").to_pylist(), + ) + self.assertEqual( + [30, None, None, None], + result.column("metrics_overflow").to_pylist(), + ) + + def test_reject_orc(self): + with self.assertRaisesRegex( + ValueError, + "PyPaimon MAP shared-shredding writes only support parquet"): + writer = self._create_table( + "orc", max_columns=2).new_batch_write_builder().new_write() + writer.write_arrow(pa.Table.from_pydict({ + "id": [1], + "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_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): + 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_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_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 + + 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, + extra_options={ + "data-evolution.enabled": "true", + "row-tracking.enabled": "true", + "target-file-row-num": "2", + }, + ) + data = pa.Table.from_pydict({ + "id": [1, 2, 3, 4], + "metrics": [ + [("a", 1)], + [("a", 2)], + [("a", 3), ("b", 4)], + [("b", 5)], + ], + }, schema=self.arrow_schema) + + messages = self._write(table, data) + files = [file for message in messages for file in message.new_files] + self.assertEqual(2, len(files)) + counts = [] + for file in files: + field = pq.read_schema(file.file_path).field("metrics") + counts.append(parse_shared_shredding_metadata(field)[1]) + self.assertEqual([4, 1], counts) + + read_builder = table.new_read_builder() + result = read_builder.new_read().to_arrow( + read_builder.new_scan().plan().splits()) + self.assertEqual(4, result.num_rows) + + def _create_table( + self, file_format, max_columns, extra_options=None): + options = { + "file.format": file_format, + "fields.metrics.map.storage-layout": "shared-shredding", + "fields.metrics.map.shared-shredding.max-columns": str(max_columns), + } + options.update(extra_options or {}) + name = "default.map_write_{}_{}".format( + file_format, len(self.catalog.list_tables("default"))) + self.catalog.create_table( + name, + Schema.from_pyarrow_schema(self.arrow_schema, options=options), + False, + ) + return self.catalog.get_table(name) + + @staticmethod + def _write(table, data): + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(data) + messages = writer.prepare_commit() + builder.new_commit().commit(messages) + writer.close() + return messages + +if __name__ == "__main__": + unittest.main() diff --git a/paimon-python/pypaimon/write/map_shared_shredding_writer.py b/paimon-python/pypaimon/write/map_shared_shredding_writer.py new file mode 100644 index 0000000000..07f21edf5c --- /dev/null +++ b/paimon-python/pypaimon/write/map_shared_shredding_writer.py @@ -0,0 +1,573 @@ +# 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. + +"""Java-compatible shared-shredding MAP conversion for data files.""" + +import math +from array import array +from collections import deque + +import pyarrow as pa + +from pypaimon.data.map_shared_shredding import ( + _normalized_offsets, + shared_shredding_metadata, +) +from pypaimon.schema.data_types import ( + ArrayType, + AtomicType, + MapType, + MultisetType, + RowType, + VectorType, +) +from pypaimon.table.bucket_mode import BucketMode + + +_FIELD_MAPPING = "__field_mapping" +_OVERFLOW = "__overflow" +_PHYSICAL_COLUMN_PREFIX = "__col_" +_FIELD_ID_BASE = 2147483647 // 4 +_FIELD_ID_DEPTH_LIMIT = 1 << 10 +_CONVERSION_BYTES = 8 * 1024 * 1024 + + +class MapSharedShreddingWriter: + """Converts configured logical MAP columns before each file write.""" + + 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() + self._fields = {} + self._max_columns = {} + self._policies = {} + 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( + "Unsupported MAP storage layout for '{}': {}".format( + name, layout)) + if layout != "shared-shredding": + continue + self._validate_field(field) + self._fields[name] = field + self._max_columns[name] = \ + options.map_shared_shredding_max_columns(name) + self._policies[name] = \ + options.map_shared_shredding_column_placement_policy(name) + self._recent_widths[name] = deque(maxlen=20) + + if configured.difference(field_by_name): + name = sorted(configured.difference(field_by_name))[0] + raise ValueError( + "Column '{}' is configured with map.storage-layout but does " + "not exist in table schema.".format(name)) + if not self._fields: + return + + self._validate_format("file.format", file_format) + self._validate_format("changelog.file.format", changelog_format) + self._validate_compression("file.compression", options.file_compression()) + 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) + for field in fields): + raise ValueError( + "MAP shared-shredding cannot be used with VARIANT fields.") + if any(_contains_type(field.type, lambda t: isinstance(t, MultisetType)) + for field in fields): + raise ValueError( + "MAP shared-shredding cannot be used with MULTISET fields.") + + def is_active(self): + 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.""" + import pyarrow.parquet as pq + + fields = list(data.schema) + completed = {} + converters = {} + field_index = {field.name: index for index, field in enumerate(fields)} + + for name, logical_field in self._fields.items(): + index = field_index.get(name) + if index is None: + continue + input_field = fields[index] + if not pa.types.is_map(input_field.type): + raise TypeError( + "Shared-shredding column '{}' must be a MAP.".format(name)) + num_columns = self._next_num_columns(name) + converter = _MapFieldConverter( + num_columns, + self._policies[name], + input_field.type.item_type, + logical_field.type.value, + ) + # 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, + converter.field_to_columns, + converter.overflow_set, + num_columns, + converter.max_row_width, + self._options.file_compression(), + )) + metadata.update(_field_id_metadata(logical_field.id)) + fields[index] = pa.field( + name, + converter.physical_type, + nullable=input_field.nullable, + metadata=metadata, + ) + completed[name] = converter.max_row_width + 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 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 + except Exception: + file_io.delete_quietly(path) + raise + return completed + + def file_completed(self, completed): + for name, max_row_width in completed.items(): + self._recent_widths[name].append(max_row_width) + + def _next_num_columns(self, name): + widths = self._recent_widths[name] + max_columns = self._max_columns[name] + if not widths: + return max_columns + ordered = sorted(widths) + rank = max(1, min(int(math.ceil(0.90 * len(ordered))), len(ordered))) + percentile = ordered[rank - 1] + maximum = ordered[-1] + if (maximum - percentile <= 4 + or maximum <= int(math.ceil(percentile * 1.25))): + width = maximum + else: + width = percentile + return max(1, min(width, max_columns)) + + def _configured_fields(self): + prefix = "fields." + suffix = ".map.storage-layout" + result = set() + for key in self._options.options.to_map(): + if key.startswith(prefix) and key.endswith(suffix): + result.add(key[len(prefix):-len(suffix)]) + return result + + @staticmethod + def _validate_field(field): + if not isinstance(field.type, MapType): + raise ValueError( + "Column '{}' is configured with map.storage-layout=" + "shared-shredding but its type is not MAP.".format(field.name)) + key_type = field.type.key + if (not isinstance(key_type, AtomicType) + or not key_type.type.upper().startswith( + ("STRING", "VARCHAR"))): + raise ValueError( + "Shared-shredding column '{}' must use STRING keys.".format( + field.name)) + if key_type.nullable: + raise ValueError( + "Shared-shredding MAP keys cannot be nullable for column '{}'." + .format(field.name)) + if _contains_type(field.type.value, _is_blob): + raise ValueError( + "MAP shared-shredding cannot contain BLOB fields.") + if _contains_type(field.type.value, lambda t: isinstance(t, VectorType)): + raise ValueError( + "MAP shared-shredding cannot contain VECTOR fields.") + + @staticmethod + def _validate_format(option, file_format): + if file_format and file_format.lower() != "parquet": + raise ValueError( + "PyPaimon MAP shared-shredding writes only support parquet, " + "but {} is {}.".format(option, file_format)) + + @staticmethod + def _validate_compression(option, compression): + if compression and compression.lower() not in ("none", "lz4", "zstd"): + raise ValueError( + "MAP shared-shredding only supports none/lz4/zstd compression, " + "but {} is {}.".format(option, compression)) + + +class _MapFieldConverter: + + def __init__(self, num_columns, policy, item_type, logical_item_type): + self.num_columns = num_columns + self.policy = policy + self.name_to_id = {} + self.field_to_columns = {} + self.overflow_set = set() + self.max_row_width = 0 + self._resident = [-1] * num_columns + self._last_used = [0] * num_columns + self._clock = 0 + self.physical_type = _physical_struct_type( + num_columns, item_type, logical_item_type) + + def convert(self, column): + 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 _place(self, keys): + field_ids = [] + 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) + + mapping, overflow = self._allocate(field_ids) + for column_id, field_id in enumerate(mapping): + if field_id >= 0: + self.field_to_columns.setdefault(field_id, set()).add(column_id) + self.overflow_set.update(overflow) + self.max_row_width = max(self.max_row_width, len(field_ids)) + return field_ids, mapping, overflow + + 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 _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 + 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", + 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, + True, + 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_vector_writer.py b/paimon-python/pypaimon/write/writer/data_vector_writer.py index 9294f67d92..4825b22af1 100644 --- a/paimon-python/pypaimon/write/writer/data_vector_writer.py +++ b/paimon-python/pypaimon/write/writer/data_vector_writer.py @@ -261,11 +261,13 @@ class DataVectorWriter(DataWriter): if data.num_rows == 0: return None + 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: @@ -290,7 +292,7 @@ class DataVectorWriter(DataWriter): 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 +313,8 @@ class DataVectorWriter(DataWriter): 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 24a67186c3..46db0578ce 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.manifest.schema.simple_stats import SimpleStats 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,13 @@ class DataWriter(ABC): # 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, + ) # 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,6 +267,7 @@ class DataWriter(ABC): extra_files = [] row_sidecar_path = None changelog_meta = None + shared_shredding_stats = {} if self._variant_shredding: data = self._apply_variant_shredding(data) @@ -270,7 +279,7 @@ class DataWriter(ABC): # 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: @@ -300,7 +309,7 @@ class DataWriter(ABC): # 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 +320,23 @@ class DataWriter(ABC): 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,10 +379,18 @@ class DataWriter(ABC): 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) + 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. @@ -414,8 +433,7 @@ class DataWriter(ABC): 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 c894182fbe..f79369d1cd 100644 --- a/paimon-python/pypaimon/write/writer/dedicated_format_writer.py +++ b/paimon-python/pypaimon/write/writer/dedicated_format_writer.py @@ -703,12 +703,14 @@ class DedicatedFormatWriter(DataWriter): if data.num_rows == 0: return None + 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: @@ -728,7 +730,9 @@ class DedicatedFormatWriter(DataWriter): 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, 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:
