This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new df5736be71 [python] Support writing shared-shredding MAPs in Parquet
(#9837)
df5736be71 is described below
commit df5736be7140947d609033665545dbe9424883a8
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:
