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 8ff98960eac9..524344ed54e1 100644 --- a/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java +++ b/paimon-core/src/test/java/org/apache/paimon/JavaPyE2ETest.java @@ -1567,6 +1567,59 @@ public void testJavaReadMapBlobTable() throws Exception { assertAdditionalMapBlobKeyTypes(table, "python"); } + /** Java writes shared-shredding MAP columns for Python to read. */ + @Test + @EnabledIfSystemProperty(named = "run.e2e.tests", matches = "true") + public void testJavaWriteSharedShreddingMapTable() throws Exception { + for (String format : Arrays.asList("parquet", "orc")) { + Identifier identifier = identifier("shared_shredding_map_java_test_" + format); + catalog.dropTable(identifier, true); + Schema schema = + Schema.newBuilder() + .column("id", DataTypes.INT()) + .column( + "metrics", + DataTypes.MAP(DataTypes.STRING().notNull(), DataTypes.BIGINT())) + .option(BUCKET.key(), "-1") + .option(CoreOptions.FILE_FORMAT.key(), format) + .option(CoreOptions.WRITE_ONLY.key(), "true") + .option("fields.metrics.map.storage-layout", "shared-shredding") + .option("fields.metrics.map.shared-shredding.max-columns", "2") + .build(); + catalog.createTable(identifier, schema, false); + + Map first = new LinkedHashMap<>(); + first.put(BinaryString.fromString("hot"), 10L); + first.put(BinaryString.fromString("warm"), 20L); + first.put(BinaryString.fromString("overflow"), 30L); + Map second = new LinkedHashMap<>(); + second.put(BinaryString.fromString("hot"), null); + second.put(BinaryString.fromString("new"), 40L); + + FileStoreTable table = (FileStoreTable) catalog.getTable(identifier); + BatchWriteBuilder writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.write(GenericRow.of(1, new GenericMap(first))); + write.write(GenericRow.of(2, new GenericMap(second))); + write.write(GenericRow.of(3, new GenericMap(Collections.emptyMap()))); + write.write(GenericRow.of(4, null)); + commit.commit(write.prepareCommit()); + } + + Map later = new LinkedHashMap<>(); + later.put(BinaryString.fromString("late"), 50L); + later.put(BinaryString.fromString("hot"), 60L); + table = (FileStoreTable) catalog.getTable(identifier); + writeBuilder = table.newBatchWriteBuilder(); + try (BatchTableWrite write = writeBuilder.newWrite(); + BatchTableCommit commit = writeBuilder.newCommit()) { + write.write(GenericRow.of(5, new GenericMap(later))); + commit.commit(write.prepareCommit()); + } + } + } + private Map> readMapBlobRows(FileStoreTable table) throws Exception { Map> rows = new HashMap<>(); diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh index b929c6a9cd7c..66ff2b3e5f02 100755 --- a/paimon-python/dev/run_mixed_tests.sh +++ b/paimon-python/dev/run_mixed_tests.sh @@ -111,6 +111,7 @@ run_batched_java_write_tests() { core_tests="${core_tests}+testBlobWriteAlterCompact" core_tests="${core_tests}+testJavaWriteArrayBlobTable" core_tests="${core_tests}+testJavaWriteMapBlobTable" + core_tests="${core_tests}+testJavaWriteSharedShreddingMapTable" core_tests="${core_tests}+testDataEvolutionWrite" core_tests="${core_tests}+testJavaWriteRowAppendTable" if [[ "$PYTHON_MINOR" -ge 7 ]]; then @@ -1011,6 +1012,28 @@ run_map_blob_interop_test() { echo -e "${GREEN}✓ Java MAP read test completed successfully${NC}" } +run_shared_shredding_map_test() { + echo -e "${YELLOW}=== Running shared-shredding MAP Test (Java Write → Python Read) ===${NC}" + + if ! skip_batched_java_write; then + cd "$PROJECT_ROOT" + echo "Running Maven test for JavaPyE2ETest.testJavaWriteSharedShreddingMapTable..." + if ! mvn test -Dtest=org.apache.paimon.JavaPyE2ETest#testJavaWriteSharedShreddingMapTable -pl paimon-core -q -Drun.e2e.tests=true; then + echo -e "${RED}✗ Java shared-shredding MAP write test failed${NC}" + return 1 + fi + echo -e "${GREEN}✓ Java shared-shredding MAP write test completed successfully${NC}" + fi + + cd "$PAIMON_PYTHON_DIR" + echo "Running Python shared-shredding MAP read test..." + if ! python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest::test_read_shared_shredding_map_written_by_java -v; then + echo -e "${RED}✗ Python shared-shredding MAP read test failed${NC}" + return 1 + fi + echo -e "${GREEN}✓ Python shared-shredding MAP read test completed successfully${NC}" +} + # Function to run VARIANT test (Java write, Python read) run_java_variant_write_py_read_test() { echo -e "${YELLOW}=== Running VARIANT Test (Java Write, Python Read) ===${NC}" @@ -1129,6 +1152,7 @@ main() { local blob_alter_compact_result=0 local array_blob_interop_result=0 local map_blob_interop_result=0 + local shared_shredding_map_result=0 local data_evolution_result=0 local data_evolution_deletion_vector_result=0 local data_evolution_py_write_result=0 @@ -1373,6 +1397,12 @@ main() { echo "" + if ! run_shared_shredding_map_test; then + shared_shredding_map_result=1 + fi + + echo "" + # Run data evolution test (Java write, Python read). Lance variant skips # itself on <3.8 (get_file_format_params + gated Java lance read). if ! run_data_evolution_test; then @@ -1573,6 +1603,12 @@ main() { echo -e "${RED}✗ MAP Interoperability Test (Java ↔ Python): FAILED${NC}" fi + if [[ $shared_shredding_map_result -eq 0 ]]; then + echo -e "${GREEN}✓ Shared-shredding MAP Test (Java Write → Python Read): PASSED${NC}" + else + echo -e "${RED}✗ Shared-shredding MAP Test (Java Write → Python Read): FAILED${NC}" + fi + if [[ $data_evolution_result -eq 0 ]]; then echo -e "${GREEN}✓ Data Evolution Test (Java Write, Python Read): PASSED${NC}" else @@ -1614,7 +1650,7 @@ main() { # Clean up warehouse directory after all tests cleanup_warehouse - if [[ $java_write_result -eq 0 && $python_read_result -eq 0 && $python_write_result -eq 0 && $java_read_result -eq 0 && $pk_dv_result -eq 0 && $btree_index_result -eq 0 && $btree_raw_fallback_result -eq 0 && $bitmap_index_result -eq 0 && $compressed_global_index_result -eq 0 && $compressed_text_result -eq 0 && $native_fulltext_result -eq 0 && $lumina_vector_result -eq 0 && $lumina_vector_btree_result -eq 0 && $vindex_vector_result -eq 0 && $vindex_vector_raw_fallback_result -eq 0 && $compact_conflict_result -eq 0 && $blob_compact_conflict_result -eq 0 && $blob_alter_compact_result -eq 0 && $array_blob_interop_result -eq 0 && $map_blob_interop_result -eq 0 && $data_evolution_result -eq 0 && $data_evolution_deletion_vector_result -eq 0 && $data_evolution_py_write_result -eq 0 && $java_variant_write_py_read_result -eq 0 && $py_variant_write_java_read_result -eq 0 && $vector_append_table_result -eq 0 && $vector_dedicated_java_write_result -eq 0 && $vector_dedicated_py_write_result -eq 0 && $multi_vector_dedicated_java_write_result -eq 0 && $multi_vector_dedicated_py_write_result -eq 0 && $row_format_result -eq 0 ]]; then + if [[ $java_write_result -eq 0 && $python_read_result -eq 0 && $python_write_result -eq 0 && $java_read_result -eq 0 && $pk_dv_result -eq 0 && $btree_index_result -eq 0 && $btree_raw_fallback_result -eq 0 && $bitmap_index_result -eq 0 && $compressed_global_index_result -eq 0 && $compressed_text_result -eq 0 && $native_fulltext_result -eq 0 && $lumina_vector_result -eq 0 && $lumina_vector_btree_result -eq 0 && $vindex_vector_result -eq 0 && $vindex_vector_raw_fallback_result -eq 0 && $compact_conflict_result -eq 0 && $blob_compact_conflict_result -eq 0 && $blob_alter_compact_result -eq 0 && $array_blob_interop_result -eq 0 && $map_blob_interop_result -eq 0 && $shared_shredding_map_result -eq 0 && $data_evolution_result -eq 0 && $data_evolution_deletion_vector_result -eq 0 && $data_evolution_py_write_result -eq 0 && $java_variant_write_py_read_result -eq 0 && $py_variant_write_java_read_result -eq 0 && $vector_append_table_result -eq 0 && $vector_dedicated_java_write_result -eq 0 && $vector_dedicated_py_write_result -eq 0 && $multi_vector_dedicated_java_write_result -eq 0 && $multi_vector_dedicated_py_write_result -eq 0 && $row_format_result -eq 0 ]]; then echo -e "${GREEN}🎉 All tests passed! Java-Python interoperability verified.${NC}" return 0 else diff --git a/paimon-python/pypaimon/data/map_shared_shredding.py b/paimon-python/pypaimon/data/map_shared_shredding.py new file mode 100644 index 000000000000..ecadf807a57b --- /dev/null +++ b/paimon-python/pypaimon/data/map_shared_shredding.py @@ -0,0 +1,347 @@ +# 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. + +"""Read support for Paimon's shared-shredding MAP storage layout.""" + +import json +import struct +from typing import Dict + +import pyarrow as pa +import pyarrow.compute as pc + + +_STORAGE_LAYOUT = b"paimon.map.storage-layout" +_VERSION = b"paimon.map.shared-shredding.version" +_FIELD_DICT = b"paimon.map.shared-shredding.field-dict" +_FIELD_DICT_COMPRESSION = b"paimon.map.shared-shredding.field-dict-compression" +_FIELD_DICT_ORIGINAL_SIZE = b"paimon.map.shared-shredding.field-dict-original-size" +_NUM_COLUMNS = b"paimon.map.shared-shredding.num-columns" +_FIELD_MAPPING = "__field_mapping" +_OVERFLOW = "__overflow" +_PHYSICAL_COLUMN_PREFIX = "__col_" + + +def is_shared_shredding(field: pa.Field) -> bool: + metadata = field.metadata + return metadata is not None and metadata.get(_STORAGE_LAYOUT) == b"shared-shredding" + + +def parse_shared_shredding_metadata(field: pa.Field): + metadata = field.metadata or {} + version = _required_int(metadata, _VERSION) + if version != 1: + raise ValueError( + "Unsupported shared-shredding metadata version: {}".format(version)) + + original_size = _required_int(metadata, _FIELD_DICT_ORIGINAL_SIZE) + compression = metadata.get(_FIELD_DICT_COMPRESSION, b"zstd").decode("utf-8").lower() + encoded_dict = _required(metadata, _FIELD_DICT).decode("utf-8").encode("latin-1") + field_dict = json.loads( + _decompress(encoded_dict, original_size, compression).decode("utf-8")) + if not isinstance(field_dict, dict): + raise ValueError("Shared-shredding field dictionary must be an object") + if not all( + isinstance(name, str) and isinstance(field_id, int) + for name, field_id in field_dict.items()): + raise ValueError("Shared-shredding field dictionary is malformed") + name_by_id = {field_id: name for name, field_id in field_dict.items()} + num_columns = _required_int(metadata, _NUM_COLUMNS) + if num_columns < 0: + raise ValueError("Shared-shredding column count must not be negative") + return name_by_id, num_columns + + +def assemble_shared_shredding_map( + column: pa.StructArray, + map_type: pa.MapType, + name_by_id: Dict[int, str], + num_columns: int) -> pa.MapArray: + """Restore one physical shared-shredding struct as a logical MAP.""" + if not pa.types.is_struct(column.type): + raise TypeError("Shared-shredding MAP must be stored as a struct") + + field_names = [field.name for field in column.type] + if not field_names or field_names[0] != _FIELD_MAPPING: + raise ValueError( + "Shared-shredding physical struct must start with {}".format( + _FIELD_MAPPING)) + + physical_columns = [None] * num_columns + overflow = None + for position, field_name in enumerate(field_names[1:], 1): + if field_name == _OVERFLOW: + if position != len(field_names) - 1: + raise ValueError("Shared-shredding overflow must be the last field") + overflow = column.field(position) + continue + if not field_name.startswith(_PHYSICAL_COLUMN_PREFIX): + raise ValueError( + "Unexpected shared-shredding physical field: {}".format(field_name)) + try: + physical_index = int(field_name[len(_PHYSICAL_COLUMN_PREFIX):]) + except ValueError: + raise ValueError( + "Unexpected shared-shredding physical field: {}".format(field_name)) + if physical_index < 0 or physical_index >= num_columns: + raise ValueError( + "Shared-shredding physical column {} exceeds metadata column count {}".format( + physical_index, num_columns)) + if physical_columns[physical_index] is not None: + raise ValueError( + "Duplicate shared-shredding physical column {}".format( + physical_index)) + physical_columns[physical_index] = column.field(position) + + mapping_column = column.field(0) + if not ( + pa.types.is_list(mapping_column.type) + or pa.types.is_large_list(mapping_column.type)): + raise TypeError("Shared-shredding field mapping must be an array") + mapping = mapping_column.to_pylist() + null_rows = column.is_null().to_pylist() + overflow_offsets = None + overflow_keys = None + overflow_values = None + if overflow is not None: + if not pa.types.is_map(overflow.type): + raise TypeError("Shared-shredding overflow field must be a map") + overflow_offsets, overflow_start, overflow_end = _normalized_offsets(overflow) + overflow_nulls = overflow.is_null().to_pylist() + overflow_keys = overflow.keys.slice( + overflow_start, overflow_end - overflow_start).to_pylist() + overflow_values = overflow.items.slice( + overflow_start, overflow_end - overflow_start) + + sources = list(physical_columns) + if overflow_values is not None: + sources.append(overflow_values) + selected_indices = [[] for _ in sources] + entry_sources = [] + entry_positions = [] + keys = [] + offsets = [0] + + for row in range(len(column)): + if null_rows[row]: + offsets[-1] = None + offsets.append(len(keys)) + continue + + row_mapping = mapping[row] + if row_mapping is None or len(row_mapping) != num_columns: + raise ValueError( + "Shared-shredding field mapping length must equal {}".format( + num_columns)) + for physical_index, field_id in enumerate(row_mapping): + if field_id is None: + raise ValueError( + "Shared-shredding field mapping must not contain null") + name = name_by_id.get(field_id) + if field_id < 0 or name is None: + continue + if physical_columns[physical_index] is None: + raise ValueError( + "Missing shared-shredding physical column {}".format( + physical_index)) + _append_entry( + keys, entry_sources, entry_positions, selected_indices, + name, physical_index, row) + + if overflow_offsets is not None and not overflow_nulls[row]: + overflow_source = len(sources) - 1 + for item_index in range( + overflow_offsets[row], overflow_offsets[row + 1]): + name = name_by_id.get(overflow_keys[item_index]) + if name is not None: + _append_entry( + keys, entry_sources, entry_positions, selected_indices, + name, overflow_source, item_index) + offsets.append(len(keys)) + + selected_values = [] + source_bases = [] + for source, indices in zip(sources, selected_indices): + source_bases.append(sum(len(values) for values in selected_values)) + if indices: + selected = pc.take(source, pa.array(indices, type=pa.int64())) + selected_values.append( + _restore_orc_temporal_values(selected, map_type.item_type)) + else: + selected_values.append(pa.array([], type=map_type.item_type)) + + if selected_values: + value_pool = pa.concat_arrays(selected_values) + value_indices = [ + source_bases[source] + position + for source, position in zip(entry_sources, entry_positions) + ] + values = pc.take(value_pool, pa.array(value_indices, type=pa.int64())) + else: + values = pa.array([], type=map_type.item_type) + + result = pa.MapArray.from_arrays( + pa.array(offsets, type=pa.int32()), + pa.array(keys, type=map_type.key_type), + values, + ) + entries = pa.StructArray.from_arrays( + [result.keys, result.items], + fields=[map_type.key_field, map_type.item_field], + ) + return pa.Array.from_buffers( + map_type, + len(result), + result.buffers()[:2], + null_count=result.null_count, + children=[entries], + ) + + +def _restore_orc_temporal_values(column, logical_type): + """Restore logical temporal types from their ORC representations.""" + if column.type == logical_type: + return column + if pa.types.is_time(logical_type) and pa.types.is_int32(column.type): + return column.cast(logical_type) + if (pa.types.is_timestamp(logical_type) + and pa.types.is_timestamp(column.type)): + return column.cast(logical_type) + if pa.types.is_struct(logical_type) and pa.types.is_struct(column.type): + if len(column.type) != len(logical_type): + return column + fields = list(logical_type) + children = [ + _restore_orc_temporal_values(column.field(i), field.type) + for i, field in enumerate(fields) + ] + mask = column.is_null() if column.null_count else None + return pa.StructArray.from_arrays(children, fields=fields, mask=mask) + if ((pa.types.is_list(logical_type) and pa.types.is_list(column.type)) + or (pa.types.is_large_list(logical_type) + and pa.types.is_large_list(column.type))): + offsets, start, end = _normalized_offsets(column) + offsets = _nullable_offsets(column, offsets, logical_type) + values = _restore_orc_temporal_values( + column.values.slice(start, end - start), logical_type.value_type) + result = (pa.LargeListArray.from_arrays(offsets, values) + if pa.types.is_large_list(logical_type) + else pa.ListArray.from_arrays(offsets, values)) + return pa.Array.from_buffers( + logical_type, + len(result), + result.buffers()[:2], + null_count=result.null_count, + children=[values], + ) + if pa.types.is_map(logical_type) and pa.types.is_map(column.type): + offsets, start, end = _normalized_offsets(column) + offsets = _nullable_offsets(column, offsets, logical_type) + keys = _restore_orc_temporal_values( + column.keys.slice(start, end - start), logical_type.key_type) + items = _restore_orc_temporal_values( + column.items.slice(start, end - start), logical_type.item_type) + result = pa.MapArray.from_arrays(offsets, keys, items) + entries = pa.StructArray.from_arrays( + [keys, items], + fields=[logical_type.key_field, logical_type.item_field], + ) + return pa.Array.from_buffers( + logical_type, + len(result), + result.buffers()[:2], + null_count=result.null_count, + children=[entries], + ) + return column + + +def _nullable_offsets(column, offsets, logical_type): + for index, is_null in enumerate(column.is_null().to_pylist()): + if is_null: + offsets[index] = None + offset_type = (pa.int64() if pa.types.is_large_list(logical_type) + else pa.int32()) + return pa.array(offsets, type=offset_type) + + +def _append_entry(keys, entry_sources, entry_positions, selected_indices, + name, source, source_index): + keys.append(name) + entry_sources.append(source) + entry_positions.append(len(selected_indices[source])) + selected_indices[source].append(source_index) + + +def _normalized_offsets(column): + offsets_array = getattr(column, "offsets", None) + if offsets_array is None: + offsets_array = pa.Array.from_buffers( + pa.int32(), + len(column) + 1, + [None, column.buffers()[1]], + offset=column.offset, + ) + offsets = offsets_array.to_pylist() + start = offsets[0] + normalized = [value - start for value in offsets] + return normalized, start, offsets[-1] + + +def _decompress(data: bytes, original_size: int, compression: str) -> bytes: + if original_size < 0: + raise ValueError("Shared-shredding field dictionary size must not be negative") + if compression == "none": + result = data + elif compression == "zstd": + import zstandard as zstd + result = zstd.ZstdDecompressor().decompress( + data, max_output_size=original_size) + elif compression == "lz4": + if len(data) < 8: + raise ValueError("Shared-shredding LZ4 dictionary is truncated") + compressed_size, stored_size = struct.unpack_from(" pa.Schema: + """Read Paimon's Arrow schema from ORC user metadata when necessary.""" + if any(is_shared_shredding(field) for field in fallback): + return fallback + + import pyarrow.orc as orc + source = file_io.filesystem.open_input_file( + file_io.to_filesystem_path(file_path)) + try: + metadata = orc.ORCFile(source).metadata + arrow_schema = metadata.get(b"ARROW:schema") + if arrow_schema is None: + arrow_schema = metadata.get("ARROW:schema") + if arrow_schema is None: + return fallback + if isinstance(arrow_schema, str): + arrow_schema = arrow_schema.encode("latin-1") + try: + arrow_schema = base64.b64decode(arrow_schema, validate=True) + except binascii.Error: + pass + return pa.ipc.read_schema(pa.BufferReader(arrow_schema)) + finally: + source.close() + + class FormatPyArrowReader(RecordBatchReader): """ A Format Reader that reads record batch from a Parquet or ORC file using PyArrow, @@ -288,6 +322,7 @@ class FormatPyArrowReader(RecordBatchReader): When a VARIANT column is stored in the shredded Parquet format (a struct with ``metadata``, ``value``, and ``typed_value`` fields), this reader transparently reconstructs the standard ``struct`` representation. + It also restores shared-shredding MAP columns from their physical struct layout. """ def __init__(self, file_io: FileIO, file_format: str, file_path: str, @@ -371,6 +406,10 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, self._has_nested_path = has_nested_path file_schema = self.dataset.schema + has_logical_map = any(isinstance(field.type, MapType) for field in read_fields) + metadata_schema = ( + _orc_schema_with_field_metadata(file_io, file_path, file_schema) + if file_format == 'orc' and has_logical_map else file_schema) if has_nested_path: self.existing_fields = [] self.missing_fields = [] @@ -387,16 +426,51 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, self._variant_shredding_enabled = ( options is None or options.variant_shredding_enabled()) self._variant_schema_cache: Dict[pa.DataType, VariantSchema] = {} + self._shared_shredding_maps = {} + logical_maps_by_source = {} + if nested_name_paths is None: + source_names = [field.name for field in read_fields] + else: + source_names = [ + path[0] if len(path) == 1 else None + for path in nested_name_paths + ] + for logical_field, source_name in zip(read_fields, source_names): + if (source_name is not None + and isinstance(logical_field.type, MapType)): + logical_maps_by_source.setdefault(source_name, []).append( + logical_field) + for field in metadata_schema: + logical_fields = logical_maps_by_source.get(field.name, []) + if logical_fields and is_shared_shredding(field): + metadata = parse_shared_shredding_metadata(field) + for logical_field in logical_fields: + logical_arrow_type = PyarrowFieldParser.from_paimon_type( + logical_field.type) + self._shared_shredding_maps[logical_field.name] = ( + logical_arrow_type, metadata) self._bounded_variant_read = ( self._file_format == 'parquet' and self._has_projected_variant()) + self._select_nested_after_scan = False if has_nested_path and not self._bounded_variant_read: existing_set = set(self.existing_fields) columns_dict = {} - for f, path in zip(read_fields, nested_name_paths): - if f.name in existing_set: - columns_dict[f.name] = ds.field(*path) - self._scan_columns = columns_dict + try: + for f, path in zip(read_fields, nested_name_paths): + if f.name in existing_set: + columns_dict[f.name] = ds.field(*path) + self._scan_columns = columns_dict + except TypeError: + # PyArrow 6 only accepts one field name and cannot build a + # nested FieldRef. Read the required top-level columns and + # extract their children after scanning instead. + self._scan_columns = [] + for f, path in zip(read_fields, nested_name_paths): + if (f.name in existing_set + and path[0] not in self._scan_columns): + self._scan_columns.append(path[0]) + self._select_nested_after_scan = True elif has_nested_path: self._scan_columns = None else: @@ -427,7 +501,12 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, filter=self._scan_filter, batch_size=self._scan_batch_size, ).to_reader() - self._raw_batches = self._iter_reader_batches(reader) + raw_batches = self._iter_reader_batches(reader) + if self._select_nested_after_scan: + raw_batches = ( + self._select_nested_fields(batch) + for batch in raw_batches) + self._raw_batches = raw_batches def _has_projected_variant(self) -> bool: return any( @@ -574,6 +653,9 @@ def _post_process_batch(self, batch: RecordBatch) -> RecordBatch: if self._file_format == 'orc' and self._output_schema is not None: batch = self._cast_orc_time_columns(batch) + if self._shared_shredding_maps: + batch = self._assemble_shared_shredding_maps(batch) + if self._variant_shredding_enabled: batch = self._assemble_shredded_variants(batch) @@ -609,6 +691,25 @@ def _type_for_missing(name: str) -> pa.DataType: return pa.RecordBatch.from_arrays( all_columns, schema=pa.schema(out_fields)) + def _assemble_shared_shredding_maps( + self, batch: pa.RecordBatch) -> pa.RecordBatch: + columns = list(batch.columns) + fields = list(batch.schema) + changed = False + for index, field in enumerate(fields): + shared = self._shared_shredding_maps.get(field.name) + if shared is None: + continue + map_type, (name_by_id, num_columns) = shared + columns[index] = assemble_shared_shredding_map( + columns[index], map_type, name_by_id, num_columns) + fields[index] = pa.field( + field.name, map_type, nullable=field.nullable) + changed = True + if not changed: + return batch + return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) + def _assemble_shredded_variants(self, batch: pa.RecordBatch) -> pa.RecordBatch: changed = False columns = list(batch.columns) 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 d011162b4999..618920bd5e6b 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 @@ -1582,6 +1582,34 @@ def test_read_map_blob_written_by_java(self): [expected, None, None, None], ) + def test_read_shared_shredding_map_written_by_java(self): + expected = [ + {'hot': 10, 'warm': 20, 'overflow': 30}, + {'hot': None, 'new': 40}, + {}, + None, + {'late': 50, 'hot': 60}, + ] + for file_format in ('parquet', 'orc'): + with self.subTest(file_format=file_format): + table = self.catalog.get_table( + 'default.shared_shredding_map_java_test_{}'.format( + file_format)) + read_builder = table.new_read_builder() + result = read_builder.new_read().to_arrow( + read_builder.new_scan().plan().splits()) + result = table_sort_by(result, 'id') + + self.assertTrue( + pa.types.is_map(result.schema.field('metrics').type)) + self.assertEqual([1, 2, 3, 4, 5], + result.column('id').to_pylist()) + self.assertEqual( + expected, + [None if value is None else dict(value) + for value in result.column('metrics').to_pylist()], + ) + 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/format_pyarrow_shared_shredding_map_test.py b/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py new file mode 100644 index 000000000000..11b57420c067 --- /dev/null +++ b/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py @@ -0,0 +1,402 @@ +# 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 base64 +from datetime import datetime, time, timezone +import json +import os +import shutil +import struct +import tempfile +import unittest +from unittest import mock + +import pyarrow as pa +import pyarrow.fs as pafs +import pyarrow.orc as orc +import pyarrow.parquet as pq + +from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader +from pypaimon.schema.data_types import ( + ArrayType, + AtomicType, + DataField, + MapType, + RowType, +) + + +class _LocalFileIO: + filesystem = pafs.LocalFileSystem() + + def to_filesystem_path(self, path): + return path + + +def _metadata(compression): + field_dict = json.dumps( + {"camera": 0, "state": 1, "action": 2}, + separators=(",", ":"), sort_keys=True).encode("utf-8") + if compression == "none": + compressed = field_dict + elif compression == "zstd": + compressed = bytes(pa.Codec("zstd").compress(field_dict)) + else: + payload = bytes(pa.Codec("lz4_raw").compress(field_dict)) + compressed = struct.pack("