diff --git a/docs/docs/pypaimon/reading.md b/docs/docs/pypaimon/reading.md index 21fcb63192ee..ec0d09c12e33 100644 --- a/docs/docs/pypaimon/reading.md +++ b/docs/docs/pypaimon/reading.md @@ -56,33 +56,26 @@ read_builder = read_builder.with_filter(predicate_5) See [Predicate](./reading#predicate) for all supported filters and building methods. Filter by `_ROW_ID`: see [Data Evolution](./data-evolution#filter-by-_row_id). -You can also pushdown projection by `ReadBuilder`: +Project a MAP subkey with `with_projection()`: ```python -# select f3 and f2 columns -read_builder = read_builder.with_projection(['f3', 'f2']) -``` - -For tables with nested struct columns, you can project individual sub-fields using dotted names: - -```python -# Given a table with schema: id BIGINT, info ROW, val STRING - -# Select a nested sub-field and a top-level field -read_builder = read_builder.with_projection(['info.name', 'val']) - -# The result columns are flattened with underscore-joined names: -# info_name, val +read_builder = read_builder.with_projection([ + 'id', "attributes['subkey']" +]) ``` -Nested `ROW` projections are supported for ordinary append tables and primary-key -merge reads. The reader may read the containing `ROW` and extract the requested -leaf, so selecting a leaf does not guarantee physical leaf-only I/O. +Nested `ROW` fields use dot notation, for example +`with_projection(['profile.name'])`. MAP subkeys use bracket notation; +`with_projection(["attributes['sub.key']"])` treats `sub.key` as one literal +key. Shared-shredding MAP files prune unselected keys; other layouts read the +full MAP and extract the key. Limitations: -- Data-evolution tables do not support nested projection. -- `ARRAY` and `MAP` nested paths are not supported. +- Data-evolution tables support MAP keys, but not nested ROW fields. +- Filtering a projected MAP key is not supported. +- MAP-key projection with query authorization is not supported. +- `ARRAY` paths are not supported. ### Generate Splits diff --git a/paimon-python/dev/run_mixed_tests.sh b/paimon-python/dev/run_mixed_tests.sh index 66ff2b3e5f02..620c4d989a03 100755 --- a/paimon-python/dev/run_mixed_tests.sh +++ b/paimon-python/dev/run_mixed_tests.sh @@ -1027,7 +1027,8 @@ run_shared_shredding_map_test() { 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 + if ! python -m pytest java_py_read_write_test.py::JavaPyReadWriteTest \ + -k "test_read_shared_shredding_map_written_by_java or test_read_selected_shared_shredding_keys_written_by_java" -v; then echo -e "${RED}✗ Python shared-shredding MAP read test failed${NC}" return 1 fi diff --git a/paimon-python/pypaimon/data/map_shared_shredding.py b/paimon-python/pypaimon/data/map_shared_shredding.py index ecadf807a57b..a3d6c58bac55 100644 --- a/paimon-python/pypaimon/data/map_shared_shredding.py +++ b/paimon-python/pypaimon/data/map_shared_shredding.py @@ -18,7 +18,8 @@ import json import struct -from typing import Dict +from copy import copy +from typing import Dict, List import pyarrow as pa import pyarrow.compute as pc @@ -30,9 +31,13 @@ _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_COLUMNS = b"paimon.map.shared-shredding.field-columns" +_OVERFLOW_SET = b"paimon.map.shared-shredding.overflow-set" _FIELD_MAPPING = "__field_mapping" _OVERFLOW = "__overflow" _PHYSICAL_COLUMN_PREFIX = "__col_" +_SELECTED_KEYS_PREFIX = "__PAIMON_MAP_SELECTED_KEYS:" +_SELECTED_KEYS_DELIMITER = ";" def is_shared_shredding(field: pa.Field) -> bool: @@ -65,6 +70,95 @@ def parse_shared_shredding_metadata(field: pa.Field): return name_by_id, num_columns +def parse_shared_shredding_selection_metadata(field: pa.Field): + """Return file-local placement metadata required for key pruning.""" + name_by_id, num_columns = parse_shared_shredding_metadata(field) + metadata = field.metadata or {} + field_columns_json = json.loads( + _required(metadata, _FIELD_COLUMNS).decode("utf-8")) + if not isinstance(field_columns_json, dict): + raise ValueError("Shared-shredding field columns must be an object") + try: + field_to_columns = { + int(field_id): list(columns) + for field_id, columns in field_columns_json.items() + } + except (TypeError, ValueError): + raise ValueError("Shared-shredding field columns are malformed") + if not all( + isinstance(column, int) and 0 <= column < num_columns + for columns in field_to_columns.values() + if isinstance(columns, list) + for column in columns): + raise ValueError("Shared-shredding field columns are malformed") + if not all(isinstance(columns, list) + for columns in field_to_columns.values()): + raise ValueError("Shared-shredding field columns are malformed") + + overflow_json = json.loads( + _required(metadata, _OVERFLOW_SET).decode("utf-8")) + if (not isinstance(overflow_json, list) + or not all(isinstance(field_id, int) + for field_id in overflow_json)): + raise ValueError("Shared-shredding overflow set is malformed") + return name_by_id, field_to_columns, set(overflow_json), num_columns + + +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( + description)) + return description[len(_SELECTED_KEYS_PREFIX):].split( + _SELECTED_KEYS_DELIMITER) + + +def is_map_selected_keys_field(field) -> bool: + from pypaimon.schema.data_types import RowType + + return ( + isinstance(field.type, RowType) + and field.description is not None + and field.description.startswith(_SELECTED_KEYS_PREFIX) + ) + + +def map_selected_keys_field(field, keys, value_type=None): + """Build the temporary ROW used by selected-key MAP reads.""" + from pypaimon.schema.data_types import DataField, MapType, RowType + + if not keys: + raise ValueError("Selected MAP keys must not be empty") + if len(set(keys)) != len(keys): + raise ValueError("Selected MAP keys must not contain duplicates") + for key in keys: + if not isinstance(key, str): + raise TypeError("Selected MAP keys must be strings") + if _SELECTED_KEYS_DELIMITER in key: + raise ValueError( + "Selected MAP key must not contain '{}': {}".format( + _SELECTED_KEYS_DELIMITER, key)) + if key.startswith(_SELECTED_KEYS_PREFIX): + raise ValueError( + "Selected MAP key must not start with metadata prefix: {}".format( + key)) + + if value_type is None: + if not isinstance(field.type, MapType): + raise TypeError("Selected-key projection requires a MAP field") + value_type = field.type.value + children = [] + for index, key in enumerate(keys): + child_type = copy(value_type) + child_type.nullable = True + children.append(DataField(index, key, child_type)) + return DataField( + field.id, + field.name, + RowType(field.type.nullable, children), + _SELECTED_KEYS_PREFIX + _SELECTED_KEYS_DELIMITER.join(keys), + ) + + def assemble_shared_shredding_map( column: pa.StructArray, map_type: pa.MapType, @@ -211,6 +305,172 @@ def assemble_shared_shredding_map( ) +def shared_shredding_selected_paths( + field_name: str, selected_keys: List[str], metadata) -> List[str]: + """Return the physical leaf paths needed for selected literal keys.""" + name_by_id, field_to_columns, overflow_set, _ = metadata + id_by_name = {name: field_id for field_id, name in name_by_id.items()} + columns = set() + include_overflow = False + for key in selected_keys: + field_id = id_by_name.get(key) + if field_id is None: + continue + columns.update(field_to_columns.get(field_id, ())) + include_overflow = include_overflow or field_id in overflow_set + paths = ["{}.{}".format(field_name, _FIELD_MAPPING)] + paths.extend( + "{}.{}{}".format(field_name, _PHYSICAL_COLUMN_PREFIX, index) + for index in sorted(columns) + ) + if include_overflow: + paths.append("{}.{}".format(field_name, _OVERFLOW)) + return paths + + +def assemble_shared_shredding_selected_keys( + column: pa.StructArray, + selected_keys: List[str], + value_type: pa.DataType, + metadata) -> pa.StructArray: + """Materialize selected MAP values from a pruned physical struct.""" + if not pa.types.is_struct(column.type): + raise TypeError("Shared-shredding MAP must be stored as a struct") + name_by_id, field_to_columns, overflow_set, num_columns = metadata + id_by_name = {name: field_id for field_id, name in name_by_id.items()} + + 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 = {} + overflow = None + for position, field_name in enumerate(field_names[1:], 1): + if field_name == _OVERFLOW: + overflow = column.field(position) + elif field_name.startswith(_PHYSICAL_COLUMN_PREFIX): + try: + physical_columns[int( + field_name[len(_PHYSICAL_COLUMN_PREFIX):])] = column.field( + position) + except ValueError: + raise ValueError( + "Unexpected shared-shredding physical field: {}".format( + field_name)) + else: + raise ValueError( + "Unexpected shared-shredding physical field: {}".format( + field_name)) + + mapping = column.field(0).to_pylist() + null_rows = column.is_null().to_pylist() + overflow_offsets = overflow_keys = overflow_values = overflow_nulls = None + if overflow is not None: + overflow_offsets, overflow_start, overflow_end = _normalized_offsets( + overflow) + overflow_keys = overflow.keys.slice( + overflow_start, overflow_end - overflow_start).to_pylist() + overflow_values = _restore_orc_temporal_values( + overflow.items.slice( + overflow_start, overflow_end - overflow_start), value_type) + overflow_nulls = overflow.is_null().to_pylist() + + sources = [] + source_by_column = {} + for physical_index in sorted(physical_columns): + source_by_column[physical_index] = len(sources) + sources.append(_restore_orc_temporal_values( + physical_columns[physical_index], value_type)) + overflow_source = None + if overflow_values is not None: + overflow_source = len(sources) + sources.append(overflow_values) + + source_bases = [] + value_arrays = [] + next_base = 0 + for source in sources: + source_bases.append(next_base) + value_arrays.append(source) + next_base += len(source) + value_pool = ( + pa.concat_arrays(value_arrays) + if value_arrays else pa.array([], type=value_type) + ) + + children = [] + for key in selected_keys: + field_id = id_by_name.get(key) + indices = [] + candidate_columns = ( + field_to_columns.get(field_id, ()) if field_id is not None else ()) + for row in range(len(column)): + selected = None + if not null_rows[row] and field_id is not None: + 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 in candidate_columns: + if row_mapping[physical_index] == field_id: + source = source_by_column.get(physical_index) + if source is None: + raise ValueError( + "Missing shared-shredding physical column {}".format( + physical_index)) + selected = source_bases[source] + row + break + if (selected is None + and field_id in overflow_set + and overflow_source is not None + and not overflow_nulls[row]): + for item_index in range( + overflow_offsets[row], overflow_offsets[row + 1]): + if overflow_keys[item_index] == field_id: + selected = ( + source_bases[overflow_source] + item_index) + break + indices.append(selected) + children.append(pc.take( + value_pool, pa.array(indices, type=pa.int64()))) + + fields = [pa.field(key, value_type) for key in selected_keys] + mask = column.is_null() if column.null_count else None + return pa.StructArray.from_arrays(children, fields=fields, mask=mask) + + +def assemble_normal_map_selected_keys( + column: pa.MapArray, + selected_keys: List[str], + value_type: pa.DataType) -> pa.StructArray: + """Materialize selected values when an older file stores a normal MAP.""" + if not pa.types.is_map(column.type): + raise TypeError("Selected-key MAP must be stored as a map or shared struct") + offsets, start, end = _normalized_offsets(column) + keys = column.keys.slice(start, end - start).to_pylist() + values = _restore_orc_temporal_values( + column.items.slice(start, end - start), value_type) + null_rows = column.is_null().to_pylist() + children = [] + for selected_key in selected_keys: + indices = [] + for row in range(len(column)): + selected = None + if not null_rows[row]: + for item_index in range(offsets[row], offsets[row + 1]): + if keys[item_index] == selected_key: + selected = item_index + break + indices.append(selected) + children.append(pc.take( + values, pa.array(indices, type=pa.int64()))) + fields = [pa.field(key, value_type) for key in selected_keys] + mask = column.is_null() if column.null_count else None + return pa.StructArray.from_arrays(children, fields=fields, mask=mask) + + def _restore_orc_temporal_values(column, logical_type): """Restore logical temporal types from their ORC representations.""" if column.type == logical_type: diff --git a/paimon-python/pypaimon/read/read_builder.py b/paimon-python/pypaimon/read/read_builder.py index a531d5d2c83f..d64f863b9ded 100644 --- a/paimon-python/pypaimon/read/read_builder.py +++ b/paimon-python/pypaimon/read/read_builder.py @@ -15,20 +15,36 @@ # specific language governing permissions and limitations # under the License. +import ast from typing import List, Optional from pypaimon.common.predicate import Predicate from pypaimon.common.predicate_builder import PredicateBuilder from pypaimon.read.explain import ExplainResult, ExplainSplitInfo, PruningStat from pypaimon.read.explain_render import render_predicate +from pypaimon.read.push_down_utils import predicate_field_names from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.scan_stats import ScanStats from pypaimon.read.split import Split from pypaimon.read.table_read import TableRead from pypaimon.read.table_scan import TableScan -from pypaimon.schema.data_types import DataField +from pypaimon.schema.data_types import AtomicType, DataField, MapType from pypaimon.table.special_fields import SpecialFields -from pypaimon.utils.projection import Projection, is_row_type +from pypaimon.utils.projection import MapKey, Projection, is_row_type + + +class _ReadPredicateBuilder(PredicateBuilder): + + def __init__(self, fields, unsupported_fields): + super().__init__(fields) + self._unsupported_fields = unsupported_fields + + def _get_field_index(self, field: str) -> int: + if field in self._unsupported_fields: + raise NotImplementedError( + "Filtering projected MAP keys is not supported: {}".format( + field)) + return super()._get_field_index(field) class ReadBuilder: @@ -40,7 +56,7 @@ def __init__(self, table): self.table: FileStoreTable = table self._predicate: Optional[Predicate] = None # ``_projection`` stores the user-facing name list from - # :meth:`with_projection`. When dotted names are present, + # :meth:`with_projection`. When nested selectors are present, # ``_nested_paths`` is also populated and takes precedence # in ``read_type()`` and downstream consumers. self._projection: Optional[List[str]] = None @@ -60,17 +76,17 @@ def with_projection(self, projection: List[str]) -> 'ReadBuilder': """Project to the given column names. Names containing a dot (e.g. ``"struct.subfield"``) walk into ROW - children and are translated into a nested projection. Top-level- - only callers see the same observable behaviour as before — the - dotted form is opt-in. Unknown names are silently skipped to + children. A quoted bracket selector on a top-level + ``MAP`` selects one literal key (e.g. + ``"attrs['key.with.dots']"``). Unknown names are silently skipped to preserve the pre-existing contract. - Precedence: if a dotted name matches an actual top-level field, the - top-level match wins and the name is not walked as a struct path. + An exact top-level field match takes precedence over both forms. """ self._projection = projection - if projection and any('.' in name for name in projection): - self._nested_paths = self._resolve_dotted_paths(projection) + if projection and any( + '.' in name or '[' in name for name in projection): + self._nested_paths = self._resolve_projection_paths(projection) else: self._nested_paths = None return self @@ -80,16 +96,18 @@ def with_limit(self, limit: int) -> 'ReadBuilder': return self def new_scan(self) -> TableScan: + self._validate_map_key_filter() scan = TableScan( table=self.table, predicate=self._predicate, limit=self._limit, partition_predicate=self._partition_filter, ) - scan._read_type = self.read_type() + scan._read_type = self._scan_read_type() return scan def new_read(self) -> TableRead: + self._validate_map_key_filter() return TableRead( table=self.table, predicate=self._predicate, @@ -111,7 +129,8 @@ def _nested_name_paths(self) -> Optional[List[List[str]]]: return Projection.of(self._nested_paths).to_name_paths(table_fields) def new_predicate_builder(self) -> PredicateBuilder: - return PredicateBuilder(self.read_type()) + return _ReadPredicateBuilder( + self.read_type(), self._map_key_output_names()) def explain(self, verbose: bool = False) -> ExplainResult: """Produce a structured scan plan for this builder. @@ -164,16 +183,30 @@ def read_type(self) -> List[DataField]: # Helpers # ------------------------------------------------------------------ - def _resolve_dotted_paths(self, names: List[str]) -> List[List[int]]: - """Translate dotted-name projection entries into integer paths - against the current table schema. Names without dots produce - length-1 paths. - """ + def _resolve_projection_paths(self, names: List[str]) -> List[List[int]]: + """Translate ROW paths and MAP-key selectors into internal paths.""" table_fields = self.table.fields if self.table.options.row_tracking_enabled(): table_fields = SpecialFields.row_type_with_row_tracking(table_fields) top_index = {f.name: i for i, f in enumerate(table_fields)} + def resolve_row_path(top, parts): + path = [top_index[top]] + current_field = table_fields[path[0]] + for part in parts: + if not is_row_type(current_field.type): + return None + child_fields = current_field.type.fields + child_idx = next( + (i for i, f in enumerate(child_fields) + if f.name == part), + -1) + if child_idx < 0: + return None + path.append(child_idx) + current_field = child_fields[child_idx] + return path + paths: List[List[int]] = [] for name in names: # Dot can be part of a top-level field name, not only a struct path @@ -181,32 +214,104 @@ def _resolve_dotted_paths(self, names: List[str]) -> List[List[int]]: if name in top_index: paths.append([top_index[name]]) continue + + map_selector = _map_key_selector(name, table_fields) + if map_selector is not None: + top, key = map_selector + paths.append([top_index[top], MapKey(key)]) + continue + if '.' not in name: continue + + # Preserve the original ROW-path semantics before considering a + # dotted top-level field name as the path prefix. parts = name.split('.') top = parts[0] - if top not in top_index: + if top in top_index: + path = resolve_row_path(top, parts[1:]) + if path is not None: + paths.append(path) + continue + + candidates = [ + field_name for field_name in top_index + if name.startswith(field_name + '.') + and is_row_type(table_fields[top_index[field_name]].type) + ] + if not candidates: continue - path = [top_index[top]] - current_field = table_fields[path[0]] - ok = True - for part in parts[1:]: - if not is_row_type(current_field.type): - ok = False - break - child_fields = current_field.type.fields - child_idx = next( - (i for i, f in enumerate(child_fields) if f.name == part), - -1) - if child_idx < 0: - ok = False - break - path.append(child_idx) - current_field = child_fields[child_idx] - if ok: + top = max(candidates, key=len) + parts = name[len(top) + 1:].split('.') + path = resolve_row_path(top, parts) + if path is not None: paths.append(path) return paths + def _validate_map_key_filter(self): + if self._predicate is None or not self._nested_paths: + return + unsupported = ( + predicate_field_names(self._predicate) + & self._map_key_output_names() + ) + if unsupported: + raise NotImplementedError( + "Filtering projected MAP keys is not supported: {}".format( + sorted(unsupported))) + + def _map_key_output_names(self): + if not self._nested_paths: + return set() + return { + field.name + for field, path in zip(self.read_type(), self._nested_paths) + if any(isinstance(step, MapKey) for step in path) + } + + def _scan_read_type(self): + if not self._nested_paths: + return self.read_type() + table_fields = self.table.fields + if self.table.options.row_tracking_enabled(): + table_fields = SpecialFields.row_type_with_row_tracking(table_fields) + seen = set() + fields = [] + for path in self._nested_paths: + if path[0] not in seen: + seen.add(path[0]) + fields.append(table_fields[path[0]]) + return fields + + +def _map_key_selector(name, table_fields): + candidates = [ + field for field in table_fields + if _is_string_key_map(field.type) + and name.startswith(field.name + '[') + ] + if not candidates: + return None + field = max(candidates, key=lambda candidate: len(candidate.name)) + selector = name[len(field.name):] + if not selector.endswith(']'): + return None + try: + key = ast.literal_eval(selector[1:-1]) + except (SyntaxError, ValueError): + return None + if not isinstance(key, str): + return None + return field.name, key + + +def _is_string_key_map(data_type) -> bool: + return ( + isinstance(data_type, MapType) + and isinstance(data_type.key, AtomicType) + and data_type.key.type.upper() == 'STRING' + ) + def _build_explain_result(table, scan: TableScan, plan, stats: ScanStats, predicate, projection, limit, verbose: bool) -> ExplainResult: diff --git a/paimon-python/pypaimon/read/reader/data_file_batch_reader.py b/paimon-python/pypaimon/read/reader/data_file_batch_reader.py index 2bdf38689e84..bd98bccf2a7b 100644 --- a/paimon-python/pypaimon/read/reader/data_file_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/data_file_batch_reader.py @@ -22,6 +22,11 @@ from pyarrow import RecordBatch from pypaimon.common.file_io import FileIO +from pypaimon.data.map_shared_shredding import ( + assemble_normal_map_selected_keys, + is_map_selected_keys_field, + map_selected_keys, +) from pypaimon.read.partition_info import PartitionInfo from pypaimon.read.reader.format_blob_reader import FormatBlobReader from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader @@ -160,6 +165,11 @@ def __init__(self, format_reader: RecordBatchReader, index_mapping: List[int], p self.index_mapping = index_mapping self.partition_info = partition_info self.system_primary_key = system_primary_key + source_fields = ( + file_data_fields if file_data_fields is not None else fields) + self._file_data_field_map = { + field.name: field for field in source_fields + } self.schema_map = {field.name: field for field in PyarrowFieldParser.from_paimon_schema(fields)} self.row_tracking_enabled = row_tracking_enabled self.first_row_id = first_row_id @@ -222,6 +232,30 @@ def _normalize_batch(self, record_batch: RecordBatch) -> RecordBatch: names.append(target_field.name) return pa.RecordBatch.from_arrays(arrays, names=names) + def _assemble_selected_map_keys( + self, record_batch: RecordBatch) -> RecordBatch: + columns = list(record_batch.columns) + fields = list(record_batch.schema) + changed = False + for index, field in enumerate(fields): + data_field = self._file_data_field_map.get(field.name) + if (data_field is None + or not is_map_selected_keys_field(data_field) + or not pa.types.is_map(columns[index].type)): + continue + value_type = PyarrowFieldParser.from_paimon_type( + data_field.type.fields[0].type) + columns[index] = assemble_normal_map_selected_keys( + columns[index], + map_selected_keys(data_field.description), + value_type) + fields[index] = pa.field( + field.name, columns[index].type, nullable=field.nullable) + changed = True + if not changed: + return record_batch + return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) + def _align_array_by_id(self, array, file_type, target_type): """Return *array* converted to *target_type*, matching ROW sub-fields by field id (reorder, pad missing with NULL, follow renames, cast changed @@ -282,6 +316,7 @@ def read_arrow_batch(self, start_idx=None, end_idx=None) -> Optional[RecordBatch record_batch = self.format_reader.read_arrow_batch() if record_batch is None: return None + record_batch = self._assemble_selected_map_keys(record_batch) record_batch = self._normalize_batch(record_batch) if self.partition_info is None and self.index_mapping is None: diff --git a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py index 672fdcfa8a06..f3403e260bc1 100644 --- a/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py +++ b/paimon-python/pypaimon/read/reader/format_pyarrow_reader.py @@ -32,9 +32,15 @@ from pypaimon.common.options.config import CatalogOptions from pypaimon.common.options.core_options import CoreOptions from pypaimon.data.map_shared_shredding import ( + assemble_normal_map_selected_keys, assemble_shared_shredding_map, + assemble_shared_shredding_selected_keys, + is_map_selected_keys_field, is_shared_shredding, + map_selected_keys, parse_shared_shredding_metadata, + parse_shared_shredding_selection_metadata, + shared_shredding_selected_paths, ) from pypaimon.data.variant_shredding import ( VariantSchema, @@ -406,7 +412,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) + has_logical_map = any( + isinstance(field.type, MapType) + or is_map_selected_keys_field(field) + 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) @@ -427,7 +436,10 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, options is None or options.variant_shredding_enabled()) self._variant_schema_cache: Dict[pa.DataType, VariantSchema] = {} self._shared_shredding_maps = {} + self._selected_key_maps = {} + self._selected_shared_map_paths = {} logical_maps_by_source = {} + selected_maps_by_source = {} if nested_name_paths is None: source_names = [field.name for field in read_fields] else: @@ -440,6 +452,9 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, and isinstance(logical_field.type, MapType)): logical_maps_by_source.setdefault(source_name, []).append( logical_field) + elif (source_name is not None + and is_map_selected_keys_field(logical_field)): + selected_maps_by_source[source_name] = logical_field for field in metadata_schema: logical_fields = logical_maps_by_source.get(field.name, []) if logical_fields and is_shared_shredding(field): @@ -449,6 +464,20 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, logical_field.type) self._shared_shredding_maps[logical_field.name] = ( logical_arrow_type, metadata) + selected_field = selected_maps_by_source.get(field.name) + if selected_field is not None: + keys = map_selected_keys(selected_field.description) + value_type = PyarrowFieldParser.from_paimon_type( + selected_field.type.fields[0].type) + selection_metadata = None + if is_shared_shredding(field): + selection_metadata = \ + parse_shared_shredding_selection_metadata(field) + self._selected_shared_map_paths[field.name] = \ + shared_shredding_selected_paths( + field.name, keys, selection_metadata) + self._selected_key_maps[field.name] = ( + keys, value_type, selection_metadata) self._bounded_variant_read = ( self._file_format == 'parquet' and self._has_projected_variant()) @@ -485,16 +514,27 @@ def __init__(self, file_io: FileIO, file_format: str, file_path: str, # Read projected VARIANT columns in bounded batches. self._parquet_file = None + self._orc_file = None + self._orc_source = None if (self._bounded_variant_read - or self._selected_parquet_row_groups is not None): + or self._selected_parquet_row_groups is not None + or (file_format == 'parquet' + and self._selected_shared_map_paths)): import pyarrow.parquet as pq # ParquetFile(filesystem=...) is unavailable in PyArrow 6. self._parquet_file = pq.ParquetFile( file_io.filesystem.open_input_file(file_path_for_pyarrow)) + if file_format == 'orc' and self._selected_shared_map_paths: + import pyarrow.orc as orc + self._orc_source = file_io.filesystem.open_input_file( + file_path_for_pyarrow) + self._orc_file = orc.ORCFile(self._orc_source) if self._exhausted: self._raw_batches = iter(()) elif self._parquet_file is not None: self._raw_batches = self._iter_row_group_batches() + elif self._orc_file is not None: + self._raw_batches = self._iter_orc_batches() else: reader = self.dataset.scanner( columns=self._scan_columns, @@ -574,6 +614,22 @@ def _read_parquet_row_group_batches(self, row_group, columns): batch_size=self._scan_batch_size, ) + def _iter_orc_batches(self): + columns = self._physical_read_columns() + for stripe in range(self._orc_file.nstripes): + batch = self._orc_file.read_stripe(stripe, columns=columns) + for offset in range(0, batch.num_rows, self._scan_batch_size): + current = batch.slice(offset, self._scan_batch_size) + if self._scan_filter is not None: + table = ds.dataset(pa.Table.from_batches([current])).scanner( + filter=self._scan_filter).to_table() + batches = table.to_batches() + else: + batches = [current] + for filtered in batches: + if filtered.num_rows: + yield self._select_existing_fields(filtered) + def _row_group_read_columns(self): if self._has_nested_path: existing = set(self.existing_fields) @@ -583,6 +639,18 @@ def _row_group_read_columns(self): columns.append(path[0]) else: columns = list(self.existing_fields) + return self._physical_read_columns(columns) + + def _physical_read_columns(self, columns=None): + if columns is None: + columns = list(self.existing_fields) + else: + columns = list(columns) + for field_name, paths in self._selected_shared_map_paths.items(): + columns = [name for name in columns if name != field_name] + for path in paths: + if path not in columns: + columns.append(path) if self._scan_filter is not None: file_names = set(self.dataset.schema.names) for name in self._predicate_field_names: @@ -656,6 +724,9 @@ def _post_process_batch(self, batch: RecordBatch) -> RecordBatch: if self._shared_shredding_maps: batch = self._assemble_shared_shredding_maps(batch) + if self._selected_key_maps: + batch = self._assemble_selected_key_maps(batch) + if self._variant_shredding_enabled: batch = self._assemble_shredded_variants(batch) @@ -710,6 +781,29 @@ def _assemble_shared_shredding_maps( return batch return pa.RecordBatch.from_arrays(columns, schema=pa.schema(fields)) + def _assemble_selected_key_maps( + self, batch: pa.RecordBatch) -> pa.RecordBatch: + columns = list(batch.columns) + fields = list(batch.schema) + changed = False + for index, field in enumerate(fields): + selected = self._selected_key_maps.get(field.name) + if selected is None: + continue + keys, value_type, metadata = selected + if metadata is None: + columns[index] = assemble_normal_map_selected_keys( + columns[index], keys, value_type) + else: + columns[index] = assemble_shared_shredding_selected_keys( + columns[index], keys, value_type, metadata) + fields[index] = pa.field( + field.name, columns[index].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) @@ -760,6 +854,10 @@ def close(self): if close is not None: close() self._parquet_file = None + if self._orc_source is not None: + self._orc_source.close() + self._orc_source = None + self._orc_file = None def _path_exists_in_arrow_schema(schema: pa.Schema, path: List[str]) -> bool: diff --git a/paimon-python/pypaimon/read/reader/nested_leaf_batch_reader.py b/paimon-python/pypaimon/read/reader/nested_leaf_batch_reader.py index cd849454a719..59bf7def733c 100644 --- a/paimon-python/pypaimon/read/reader/nested_leaf_batch_reader.py +++ b/paimon-python/pypaimon/read/reader/nested_leaf_batch_reader.py @@ -21,11 +21,22 @@ import pyarrow.compute as pc from pyarrow import RecordBatch +from pypaimon.data.map_shared_shredding import \ + assemble_normal_map_selected_keys from pypaimon.read.reader.field_indices import blob_field_indices, vector_field_indices from pypaimon.read.reader.iface.record_batch_reader import RecordBatchReader from pypaimon.schema.data_types import DataField, PyarrowFieldParser +def _struct_field(column, name): + struct_field = getattr(pc, "struct_field", None) + if struct_field is not None: + return struct_field(column, name) + + field_index = [field.name for field in column.type].index(name) + return column.flatten()[field_index] + + class NestedLeafBatchReader(RecordBatchReader): """Extract projected nested leaves from batches of full top-level columns. @@ -57,7 +68,11 @@ def read_arrow_batch(self) -> Optional[RecordBatch]: for i, path in enumerate(self._paths): column = batch.column(path[0]) for name in path[1:]: - column = pc.struct_field(column, name) + if pa.types.is_map(column.type): + column = assemble_normal_map_selected_keys( + column, [name], column.type.item_type).field(0) + else: + column = _struct_field(column, name) target_type = self._schema.field(i).type if column.type != target_type: column = column.cast(target_type, safe=False) diff --git a/paimon-python/pypaimon/read/split_read.py b/paimon-python/pypaimon/read/split_read.py index c5300a620fa0..d32eae769709 100644 --- a/paimon-python/pypaimon/read/split_read.py +++ b/paimon-python/pypaimon/read/split_read.py @@ -28,6 +28,11 @@ PositionMappedDeletionVector, ) from pypaimon.deletionvectors.deletion_vector import DeletionVector +from pypaimon.data.map_shared_shredding import ( + is_map_selected_keys_field, + map_selected_keys, + map_selected_keys_field, +) from pypaimon.globalindex import Range from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.read.interval_partition import IntervalPartition, SortedRun @@ -77,7 +82,7 @@ builtin_seq_comparator) from pypaimon.read.split import Split from pypaimon.read.sliced_split import SlicedSplit -from pypaimon.schema.data_types import DataField, PyarrowFieldParser +from pypaimon.schema.data_types import DataField, MapType, PyarrowFieldParser from pypaimon.table.special_fields import SpecialFields from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.utils.data_evolution_utils import retrieve_anchor_file @@ -429,6 +434,16 @@ def file_reader_supplier(self, file: DataFileMeta, for_merge_read: bool, SpecialFields.row_type_with_row_tracking(self.table.table_schema.fields) if row_tracking_enabled else self.table.table_schema.fields ) + if target_fields is not None: + selected_by_id = { + field.id: field for field in target_fields + if is_map_selected_keys_field(field) + } + if selected_by_id: + table_schema_fields = [ + selected_by_id.get(field.id, field) + for field in table_schema_fields + ] # When native shard pushdown is used, the format reader only returns rows # starting from shard_range[0], so _ROW_ID must be offset accordingly. @@ -594,8 +609,21 @@ def _get_read_data_fields(self): return self._read_data_fields_from(self._get_all_data_fields()) def _read_data_fields_from(self, all_data_fields): - read_field_ids = {field.id for field in self.read_fields} - return [f for f in all_data_fields if f.id in read_field_ids] + read_fields_by_id = {field.id: field for field in self.read_fields} + result = [] + for data_field in all_data_fields: + read_field = read_fields_by_id.get(data_field.id) + if read_field is None: + continue + if (is_map_selected_keys_field(read_field) + and isinstance(data_field.type, MapType)): + data_field = map_selected_keys_field( + data_field, + map_selected_keys(read_field.description), + value_type=data_field.type.value, + ) + result.append(data_field) + return result def _final_data_fields_from(self, all_data_fields: List[DataField]) -> List[DataField]: """The per-position target fields a batch must end up as: trimmed for @@ -1197,7 +1225,6 @@ def _create_raw_reader(self) -> RecordReader: NestedLeafBatchReader reader = NestedLeafBatchReader( reader, self.outer_extract_name_paths, self.outer_flat_read_type) - if self.limit is not None and not self._post_filter_after_inline: reader = LimitedRecordBatchReader(reader, self.limit) diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 980d8fdd9c93..c5b7af32df06 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -25,6 +25,7 @@ from pypaimon.common.predicate import Predicate from pypaimon.common.predicate_json_parser import extract_referenced_fields +from pypaimon.data.map_shared_shredding import map_selected_keys_field from pypaimon.read.push_down_utils import predicate_field_names from pypaimon.read.query_auth_split import QueryAuthSplit from pypaimon.read.reader.auth_masking_reader import ( @@ -36,7 +37,8 @@ from pypaimon.read.split_read import (DataEvolutionSplitRead, MergeFileSplitRead, RawFileSplitRead, SplitRead, deferred_blob_field_names) -from pypaimon.schema.data_types import DataField, PyarrowFieldParser +from pypaimon.schema.data_types import ( + DataField, MapType, PyarrowFieldParser, is_map_blob_type) from pypaimon.table.row.offset_row import OffsetRow ROW_KIND_COLUMN = "_row_kind" @@ -918,11 +920,18 @@ def _build_split_read(self, split: Split, read_type=None, elif self.table.options.data_evolution_enabled(): if self.nested_name_paths and any( len(p) > 1 for p in self.nested_name_paths): - raise NotImplementedError( - "Nested-field projection on data-evolution tables is " - "not yet supported") - outer_extract_name_paths = None - if read_type is None and self._needs_output_projection(): + if not self._only_map_key_nested_paths(): + raise NotImplementedError( + "ROW nested-field projection on data-evolution tables " + "is not yet supported") + scan_read_type = self._with_predicate_extra_fields( + self._widen_to_top_level_for_merge()) + outer_extract_name_paths = self.nested_name_paths + else: + outer_extract_name_paths = None + if (outer_extract_name_paths is None + and read_type is None + and self._needs_output_projection()): outer_extract_name_paths = self._output_extract_name_paths() return DataEvolutionSplitRead( table=self.table, @@ -930,7 +939,9 @@ def _build_split_read(self, split: Split, read_type=None, read_type=scan_read_type, split=split, row_tracking_enabled=True, - nested_name_paths=self.nested_name_paths, + nested_name_paths=( + None if outer_extract_name_paths + else self.nested_name_paths), outer_extract_name_paths=outer_extract_name_paths, outer_flat_read_type=( self.read_type if outer_extract_name_paths else None), @@ -1007,7 +1018,10 @@ def _table_read_fields(self) -> List[DataField]: def _widen_to_top_level_for_merge(self) -> List[DataField]: """Unique top-level fields from ``self.nested_name_paths``, in path order.""" - table_fields_by_name = {f.name: f for f in self.table.fields} + table_fields_by_name = {f.name: f for f in self._table_read_fields()} + paths_by_top = {} + for path in self.nested_name_paths or []: + paths_by_top.setdefault(path[0], []).append(path) seen = set() widened: List[DataField] = [] for path in self.nested_name_paths or []: @@ -1020,9 +1034,51 @@ def _widen_to_top_level_for_merge(self) -> List[DataField]: raise ValueError( "Nested projection top-level field %r not found in " "table schema" % (top_name,)) + paths = paths_by_top[top_name] + if (isinstance(field.type, MapType) + and all(len(path) > 1 for path in paths) + and not is_map_blob_type(field.type) + and not self._map_has_aggregator(top_name)): + keys = [] + for path in paths: + if path[1] not in keys: + keys.append(path[1]) + try: + field = map_selected_keys_field(field, keys) + except ValueError: + # Keys that cannot be encoded by the Java/Spark metadata + # contract still work through the complete-MAP fallback. + pass widened.append(field) return widened + def _only_map_key_nested_paths(self) -> bool: + table_fields = {field.name: field for field in self.table.fields} + for path in self.nested_name_paths or []: + if len(path) == 1: + continue + field = table_fields.get(path[0]) + if len(path) != 2 or field is None \ + or not isinstance(field.type, MapType): + return False + return True + + def _has_map_key_projection(self) -> bool: + table_fields = {field.name: field for field in self.table.fields} + for path in self.nested_name_paths or []: + field = table_fields.get(path[0]) + if len(path) == 2 and field is not None \ + and isinstance(field.type, MapType): + return True + return False + + def _map_has_aggregator(self, field_name: str) -> bool: + options = self.table.options.options.data + return ( + "fields.{}.aggregate-function".format(field_name) in options + or "fields.default-aggregate-function" in options + ) + def __create_reader_for_split(self, split, blob_parallelism=1, limit: Optional[int] = None): auth_result = None @@ -1044,6 +1100,9 @@ def __create_reader_for_split(self, split, blob_parallelism=1, def __authed_reader(self, split, auth_result, blob_parallelism=1, limit: Optional[int] = None): + if self._has_map_key_projection(): + raise NotImplementedError( + "MAP-key projection with query authorization is not supported") table_fields = self.table.fields read_fields = self.read_type diff --git a/paimon-python/pypaimon/tests/data_evolution_formats_test.py b/paimon-python/pypaimon/tests/data_evolution_formats_test.py index 5bfb82824be9..c9703879954d 100644 --- a/paimon-python/pypaimon/tests/data_evolution_formats_test.py +++ b/paimon-python/pypaimon/tests/data_evolution_formats_test.py @@ -34,6 +34,8 @@ from pypaimon.manifest.schema.data_file_meta import DataFileMeta from pypaimon.manifest.schema.simple_stats import SimpleStats from pypaimon.read.split_read import SplitRead +from pypaimon.schema.data_types import AtomicType +from pypaimon.schema.schema_change import SchemaChange from pypaimon.table.row.generic_row import GenericRow from pypaimon.utils.range import Range @@ -239,6 +241,115 @@ def test_row_sidecar_serves_sparse_row_id_read(self): self.assertEqual(actual.column('val').to_pylist(), ['v5']) self.assertEqual(actual.column('_ROW_ID').to_pylist(), [5]) + def test_row_sidecar_map_key_uses_full_map_fallback(self): + pa_schema = pa.schema([ + ('id', pa.int32()), + ('attributes', pa.map_(pa.string(), pa.int64())), + ]) + schema = Schema.from_pyarrow_schema(pa_schema, options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + 'data-evolution.row-sidecar.enabled': 'true', + 'file.format': 'parquet', + }) + identifier = 'default.fmt_row_sidecar_map_key' + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + + wb = table.new_batch_write_builder() + tw = wb.new_write() + tc = wb.new_commit() + tw.write_arrow(pa.Table.from_arrays([ + pa.array(list(range(100)), type=pa.int32()), + pa.array( + [[('first', i)] for i in range(100)], + type=pa_schema.field('attributes').type), + ], schema=pa_schema)) + cmts = tw.prepare_commit() + tc.commit(cmts) + tw.close() + tc.close() + + data_file = next( + nf for manifest in cmts for nf in manifest.new_files + if nf.file_name.endswith('.parquet')) + self.assertEqual(1, len(self._row_sidecar_files(data_file))) + os.remove(self._file_path(data_file)) + + rb = table.new_read_builder().with_projection([ + "attributes['first']", '_ROW_ID', + ]) + pb = rb.new_predicate_builder() + rb.with_filter(pb.equal('_ROW_ID', 5)) + actual = rb.new_read().to_arrow(rb.new_scan().plan().splits()) + + self.assertEqual( + {'attributes_first': [5], '_ROW_ID': [5]}, + actual.to_pydict(), + ) + + def test_row_sidecar_map_key_with_schema_evolution(self): + for evolution in ('rename', 'value_type'): + with self.subTest(evolution=evolution): + pa_schema = pa.schema([ + ('id', pa.int32()), + ('attributes', pa.map_(pa.string(), pa.int32())), + ]) + identifier = 'default.fmt_row_sidecar_map_' + evolution + schema = Schema.from_pyarrow_schema(pa_schema, options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + 'data-evolution.row-sidecar.enabled': 'true', + 'file.format': 'parquet', + }) + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + + wb = table.new_batch_write_builder() + tw = wb.new_write() + tc = wb.new_commit() + tw.write_arrow(pa.Table.from_arrays([ + pa.array(list(range(100)), type=pa.int32()), + pa.array( + [[('first', i)] for i in range(100)], + type=pa_schema.field('attributes').type), + ], schema=pa_schema)) + cmts = tw.prepare_commit() + tc.commit(cmts) + tw.close() + tc.close() + + data_file = next( + nf for manifest in cmts for nf in manifest.new_files + if nf.file_name.endswith('.parquet')) + self.assertEqual( + 1, len(self._row_sidecar_files(data_file))) + os.remove(self._file_path(data_file)) + + if evolution == 'rename': + change = SchemaChange.rename_column( + 'attributes', 'renamed_attributes') + projection = "renamed_attributes['first']" + else: + change = SchemaChange.update_column_type( + ['attributes', 'value'], AtomicType('BIGINT')) + projection = "attributes['first']" + self.catalog.alter_table(identifier, [change], False) + table = self.catalog.get_table(identifier) + + rb = table.new_read_builder().with_projection([ + projection, '_ROW_ID', + ]) + pb = rb.new_predicate_builder() + rb.with_filter(pb.equal('_ROW_ID', 5)) + actual = rb.new_read().to_arrow( + rb.new_scan().plan().splits()) + + self.assertEqual([5], actual.column(0).to_pylist()) + if evolution == 'value_type': + self.assertEqual( + pa.int64(), actual.schema.field(0).type) + def test_parquet_column_subset_write_and_merge_read(self): """Write disjoint column subsets as parquet, merge-read via data evolution.""" pa_schema = pa.schema([ 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 618920bd5e6b..86d1255ecf80 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 @@ -1610,6 +1610,32 @@ def test_read_shared_shredding_map_written_by_java(self): for value in result.column('metrics').to_pylist()], ) + def test_read_selected_shared_shredding_keys_written_by_java(self): + 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().with_projection([ + 'id', "metrics['hot']", "metrics['overflow']", + "metrics['missing']", + ]) + result = read_builder.new_read().to_arrow( + read_builder.new_scan().plan().splits()) + result = table_sort_by(result, 'id') + + self.assertEqual( + ['id', 'metrics_hot', 'metrics_overflow', + 'metrics_missing'], + result.column_names, + ) + self.assertEqual([10, None, None, None, 60], + result.column('metrics_hot').to_pylist()) + self.assertEqual([30, None, None, None, None], + result.column('metrics_overflow').to_pylist()) + self.assertEqual([None] * 5, + result.column('metrics_missing').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 index 11b57420c067..4fd29220326d 100644 --- a/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py +++ b/paimon-python/pypaimon/tests/format_pyarrow_shared_shredding_map_test.py @@ -29,6 +29,7 @@ import pyarrow.orc as orc import pyarrow.parquet as pq +from pypaimon.data.map_shared_shredding import map_selected_keys_field from pypaimon.read.reader.format_pyarrow_reader import FormatPyArrowReader from pypaimon.schema.data_types import ( ArrayType, @@ -48,7 +49,7 @@ def to_filesystem_path(self, path): def _metadata(compression): field_dict = json.dumps( - {"camera": 0, "state": 1, "action": 2}, + {"first": 0, "second": 1, "third": 2}, separators=(",", ":"), sort_keys=True).encode("utf-8") if compression == "none": compressed = field_dict @@ -64,6 +65,9 @@ def _metadata(compression): "paimon.map.shared-shredding.field-dict": compressed.decode("latin-1"), "paimon.map.shared-shredding.field-dict-compression": compression, "paimon.map.shared-shredding.field-dict-original-size": str(len(field_dict)), + "paimon.map.shared-shredding.field-columns": json.dumps({ + "0": [0], "1": [0, 1], "2": []}), + "paimon.map.shared-shredding.overflow-set": "[2]", "paimon.map.shared-shredding.num-columns": "2", } @@ -73,12 +77,12 @@ class SharedShreddingMapReaderTest(unittest.TestCase): def setUp(self): self.tmp = tempfile.mkdtemp() self.value_arrow_type = pa.struct([ - pa.field("record_index", pa.int64()), - pa.field("timestamp_ns", pa.int64()), + pa.field("value_x", pa.int64()), + pa.field("value_y", pa.int64()), ]) self.value_type = RowType(True, [ - DataField(1, "record_index", AtomicType("BIGINT")), - DataField(2, "timestamp_ns", AtomicType("BIGINT")), + DataField(1, "value_x", AtomicType("BIGINT")), + DataField(2, "value_y", AtomicType("BIGINT")), ]) def tearDown(self): @@ -86,20 +90,20 @@ def tearDown(self): def _write(self, compression, file_format): values0 = pa.array([ - {"record_index": 10, "timestamp_ns": 100}, - {"record_index": 40, "timestamp_ns": 400}, + {"value_x": 10, "value_y": 100}, + {"value_x": 40, "value_y": 400}, None, None, ], type=self.value_arrow_type) values1 = pa.array([ - {"record_index": 20, "timestamp_ns": 200}, + {"value_x": 20, "value_y": 200}, None, None, None, ], type=self.value_arrow_type) overflow = pa.array([ - [(2, {"record_index": 30, "timestamp_ns": 300})], - [(99, {"record_index": 50, "timestamp_ns": 500})], + [(2, {"value_x": 30, "value_y": 300})], + [(99, {"value_x": 50, "value_y": 500})], None, [], ], type=pa.map_(pa.int32(), self.value_arrow_type)) @@ -115,7 +119,7 @@ def _write(self, compression, file_format): mask=pa.array([False, False, True, False]), ) field = pa.field( - "content_refs", physical.type, metadata=_metadata(compression)) + "attributes", physical.type, metadata=_metadata(compression)) table = pa.Table.from_arrays( [pa.array([0, 1, 2, 3]), physical], schema=pa.schema([pa.field("id", pa.int64()), field])) @@ -130,11 +134,11 @@ def _write(self, compression, file_format): def test_reads_complete_map_for_all_metadata_compressions(self): expected = [ [ - ("camera", {"record_index": 10, "timestamp_ns": 100}), - ("state", {"record_index": 20, "timestamp_ns": 200}), - ("action", {"record_index": 30, "timestamp_ns": 300}), + ("first", {"value_x": 10, "value_y": 100}), + ("second", {"value_x": 20, "value_y": 200}), + ("third", {"value_x": 30, "value_y": 300}), ], - [("state", {"record_index": 40, "timestamp_ns": 400})], + [("second", {"value_x": 40, "value_y": 400})], None, [], ] @@ -142,6 +146,148 @@ def test_reads_complete_map_for_all_metadata_compressions(self): with self.subTest(compression=compression): self._assert_complete_map("parquet", compression, expected) + def test_reads_only_selected_shared_shredding_keys(self): + map_field = DataField( + 0, + "attributes", + MapType( + True, AtomicType("STRING", False), self.value_type), + ) + selected_field = map_selected_keys_field( + map_field, ["first", "third", "missing"]) + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", self._write("none", "parquet"), + [selected_field], None, batch_size=2, + ) + + self.assertEqual( + [ + "attributes.__field_mapping", + "attributes.__col_0", + "attributes.__overflow", + ], + reader._row_group_read_columns(), + ) + actual = [] + while True: + batch = reader.read_arrow_batch() + if batch is None: + break + actual.extend(batch.column(0).to_pylist()) + self.assertEqual([ + { + "first": {"value_x": 10, "value_y": 100}, + "third": {"value_x": 30, "value_y": 300}, + "missing": None, + }, + {"first": None, "third": None, "missing": None}, + None, + {"first": None, "third": None, "missing": None}, + ], actual) + + def test_reads_unknown_key_without_value_columns(self): + map_field = DataField( + 0, + "attributes", + MapType( + True, AtomicType("STRING", False), self.value_type), + ) + reader = FormatPyArrowReader( + _LocalFileIO(), + "parquet", + self._write("none", "parquet"), + [map_selected_keys_field(map_field, ["missing"])], + None, + ) + + self.assertEqual( + ["attributes.__field_mapping"], + reader._row_group_read_columns(), + ) + actual = [] + while True: + batch = reader.read_arrow_batch() + if batch is None: + break + actual.extend(batch.column(0).to_pylist()) + self.assertEqual([ + {"missing": None}, {"missing": None}, None, {"missing": None}, + ], actual) + + def test_reads_only_selected_shared_shredding_keys_from_orc(self): + path = self._write("none", "orc") + physical_schema = orc.ORCFile(path).schema + fields = list(physical_schema) + index = physical_schema.get_field_index("attributes") + fields[index] = pa.field( + "attributes", fields[index].type, metadata=_metadata("none")) + map_field = DataField( + 0, + "attributes", + MapType( + True, AtomicType("STRING", False), self.value_type), + ) + with mock.patch( + "pypaimon.read.reader.format_pyarrow_reader." + "_orc_schema_with_field_metadata", + return_value=pa.schema(fields)): + reader = FormatPyArrowReader( + _LocalFileIO(), "orc", path, + [map_selected_keys_field( + map_field, ["first", "third", "missing"])], + None, + batch_size=2, + ) + actual = [] + while True: + batch = reader.read_arrow_batch() + if batch is None: + break + actual.extend(batch.column(0).to_pylist()) + + self.assertEqual([ + { + "first": {"value_x": 10, "value_y": 100}, + "third": {"value_x": 30, "value_y": 300}, + "missing": None, + }, + {"first": None, "third": None, "missing": None}, + None, + {"first": None, "third": None, "missing": None}, + ], actual) + + def test_reads_selected_keys_from_normal_map(self): + map_type = pa.map_(pa.string(), self.value_arrow_type) + path = os.path.join(self.tmp, "normal.parquet") + pq.write_table(pa.table({ + "attributes": pa.array([ + [("first", {"value_x": 1, "value_y": 10})], + [], + None, + ], type=map_type), + }), path) + map_field = DataField( + 0, + "attributes", + MapType( + True, AtomicType("STRING", False), self.value_type), + ) + reader = FormatPyArrowReader( + _LocalFileIO(), "parquet", path, + [map_selected_keys_field(map_field, ["first", "missing"])], + None, + ) + + batch = reader.read_arrow_batch() + self.assertEqual([ + { + "first": {"value_x": 1, "value_y": 10}, + "missing": None, + }, + {"first": None, "missing": None}, + None, + ], batch.column(0).to_pylist()) + def _assert_complete_map( self, file_format, compression, expected, path=None): reader = FormatPyArrowReader( @@ -149,7 +295,7 @@ def _assert_complete_map( path or self._write(compression, file_format), [DataField( 0, - "content_refs", + "attributes", MapType( True, AtomicType("STRING", False), self.value_type), )], @@ -169,19 +315,19 @@ def test_reads_arrow_schema_metadata_from_orc(self): path = self._write("none", "orc") physical_schema = orc.ORCFile(path).schema fields = list(physical_schema) - index = physical_schema.get_field_index("content_refs") + index = physical_schema.get_field_index("attributes") fields[index] = pa.field( - "content_refs", fields[index].type, metadata=_metadata("none")) + "attributes", fields[index].type, metadata=_metadata("none")) arrow_schema = base64.b64encode( pa.schema(fields).serialize().to_pybytes()) expected = [ [ - ("camera", {"record_index": 10, "timestamp_ns": 100}), - ("state", {"record_index": 20, "timestamp_ns": 200}), - ("action", {"record_index": 30, "timestamp_ns": 300}), + ("first", {"value_x": 10, "value_y": 100}), + ("second", {"value_x": 20, "value_y": 200}), + ("third", {"value_x": 30, "value_y": 300}), ], - [("state", {"record_index": 40, "timestamp_ns": 400})], + [("second", {"value_x": 40, "value_y": 400})], None, [], ] @@ -206,34 +352,89 @@ def test_restores_time_values_from_orc(self): names=["__field_mapping", "__col_0", "__col_1", "__overflow"], ) path = os.path.join(self.tmp, "time.orc") - orc.write_table(pa.table({"content_refs": physical}), path) + orc.write_table(pa.table({"attributes": physical}), path) result = self._read_orc_shared_map(path, AtomicType("TIME(3)")) self.assertEqual(pa.map_(pa.string(), pa.time32("ms")), result.type) self.assertEqual( - [[("camera", time(0, 0, 1, 234000)), - ("action", time(0, 0, 5, 678000))]], + [[("first", time(0, 0, 1, 234000)), + ("third", time(0, 0, 5, 678000))]], result.to_pylist(), ) + def test_restores_selected_time_values_from_orc(self): + physical = pa.StructArray.from_arrays( + [ + pa.array([[0, -1]], type=pa.list_(pa.int32())), + pa.array([1234], type=pa.int32()), + pa.array([None], type=pa.int32()), + pa.array( + [[(2, 5678)]], type=pa.map_(pa.int32(), pa.int32())), + ], + names=["__field_mapping", "__col_0", "__col_1", "__overflow"], + ) + path = os.path.join(self.tmp, "selected-time.orc") + orc.write_table(pa.table({"attributes": physical}), path) + physical_field = orc.ORCFile(path).schema.field("attributes") + metadata_schema = pa.schema([pa.field( + "attributes", + physical_field.type, + metadata=_metadata("none"), + )]) + map_field = DataField( + 0, + "attributes", + MapType( + True, + AtomicType("STRING", False), + AtomicType("TIME(3)"), + ), + ) + + with mock.patch( + "pypaimon.read.reader.format_pyarrow_reader." + "_orc_schema_with_field_metadata", + return_value=metadata_schema): + reader = FormatPyArrowReader( + _LocalFileIO(), + "orc", + path, + [map_selected_keys_field( + map_field, ["first", "third"])], + None, + ) + result = reader.read_arrow_batch().column(0) + + self.assertEqual( + pa.struct([ + pa.field("first", pa.time32("ms")), + pa.field("third", pa.time32("ms")), + ]), + result.type, + ) + self.assertEqual([{ + "first": time(0, 0, 1, 234000), + "third": time(0, 0, 5, 678000), + }], result.to_pylist()) + def test_restores_timestamp_precision_from_orc(self): - camera_timestamp = datetime(2024, 1, 2, 3, 4, 5, 123000) - action_timestamp = datetime(2024, 1, 2, 3, 4, 5, 678000) + first_timestamp = datetime(2024, 1, 2, 3, 4, 5, 123000) + third_timestamp = datetime(2024, 1, 2, 3, 4, 5, 678000) physical = pa.StructArray.from_arrays( [ pa.array([[0, -1]], type=pa.list_(pa.int32())), - pa.array([camera_timestamp], type=pa.timestamp("ns")), + pa.array([first_timestamp], type=pa.timestamp("ns")), pa.array([None], type=pa.timestamp("ns")), pa.array( - [[(2, action_timestamp)]], + [[(2, third_timestamp)]], type=pa.map_(pa.int32(), pa.timestamp("ns")), ), ], names=["__field_mapping", "__col_0", "__col_1", "__overflow"], ) path = os.path.join(self.tmp, "timestamp.orc") - orc.write_table(pa.table({"content_refs": physical}), path) + orc.write_table(pa.table({"attributes": physical}), path) result = self._read_orc_shared_map( path, AtomicType("TIMESTAMP(3)")) @@ -241,16 +442,16 @@ def test_restores_timestamp_precision_from_orc(self): self.assertEqual( pa.map_(pa.string(), pa.timestamp("ms")), result.type) self.assertEqual( - [[("camera", camera_timestamp), ("action", action_timestamp)]], + [[("first", first_timestamp), ("third", third_timestamp)]], result.to_pylist(), ) def test_restores_nested_timestamp_values_from_orc(self): - camera_timestamp = datetime(2024, 1, 2, 3, 4, 5, 123000) + first_timestamp = datetime(2024, 1, 2, 3, 4, 5, 123000) history_timestamp = datetime( 2024, 1, 2, 3, 4, 5, 123456, tzinfo=timezone.utc) physical_value_type = pa.struct([ - pa.field("captured_at", pa.timestamp("ns")), + pa.field("event_time", pa.timestamp("ns")), pa.field("history", pa.list_(pa.timestamp("ns", tz="UTC"))), ]) physical = pa.StructArray.from_arrays( @@ -258,7 +459,7 @@ def test_restores_nested_timestamp_values_from_orc(self): pa.array([[0, -1]], type=pa.list_(pa.int32())), pa.array( [{ - "captured_at": camera_timestamp, + "event_time": first_timestamp, "history": [history_timestamp], }], type=physical_value_type, @@ -270,9 +471,9 @@ def test_restores_nested_timestamp_values_from_orc(self): names=["__field_mapping", "__col_0", "__col_1", "__overflow"], ) path = os.path.join(self.tmp, "nested-timestamp.orc") - orc.write_table(pa.table({"content_refs": physical}), path) + orc.write_table(pa.table({"attributes": physical}), path) logical_value_type = RowType(True, [ - DataField(1, "captured_at", AtomicType("TIMESTAMP(3)")), + DataField(1, "event_time", AtomicType("TIMESTAMP(3)")), DataField( 2, "history", @@ -284,24 +485,24 @@ def test_restores_nested_timestamp_values_from_orc(self): self.assertEqual( pa.struct([ - pa.field("captured_at", pa.timestamp("ms")), + pa.field("event_time", pa.timestamp("ms")), pa.field( "history", pa.list_(pa.timestamp("us", tz="UTC"))), ]), result.type.item_type, ) self.assertEqual( - [[("camera", { - "captured_at": camera_timestamp, + [[("first", { + "event_time": first_timestamp, "history": [history_timestamp], })]], result.to_pylist(), ) def _read_orc_shared_map(self, path, value_type): - physical_field = orc.ORCFile(path).schema.field("content_refs") + physical_field = orc.ORCFile(path).schema.field("attributes") metadata_field = pa.field( - "content_refs", physical_field.type, metadata=_metadata("none")) + "attributes", physical_field.type, metadata=_metadata("none")) arrow_schema = base64.b64encode( pa.schema([metadata_field]).serialize().to_pybytes()) metadata = mock.Mock() @@ -315,7 +516,7 @@ def _read_orc_shared_map(self, path, value_type): _LocalFileIO(), "orc", path, [DataField( 0, - "content_refs", + "attributes", MapType( True, AtomicType("STRING", False), @@ -375,19 +576,19 @@ def test_restores_map_with_nested_projection_alias(self): self.assertEqual(["a_b", "a_b__0"], batch.schema.names) self.assertEqual([7], batch.column(0).to_pylist()) self.assertEqual( - [[("camera", 10)]], batch.column(1).to_pylist()) + [[("first", 10)]], batch.column(1).to_pylist()) def test_leaves_normal_map_unchanged(self): path = os.path.join(self.tmp, "normal.parquet") pq.write_table( - pa.table({"content_refs": pa.array( - [[("camera", 1)]], type=pa.map_(pa.string(), pa.int64()))}), + pa.table({"attributes": pa.array( + [[("first", 1)]], type=pa.map_(pa.string(), pa.int64()))}), path, ) reader = FormatPyArrowReader( _LocalFileIO(), "parquet", path, [DataField( - 0, "content_refs", + 0, "attributes", MapType( True, AtomicType("STRING", False), @@ -395,7 +596,7 @@ def test_leaves_normal_map_unchanged(self): None, ) self.assertEqual( - [[("camera", 1)]], reader.read_arrow_batch().column(0).to_pylist()) + [[("first", 1)]], reader.read_arrow_batch().column(0).to_pylist()) if __name__ == "__main__": diff --git a/paimon-python/pypaimon/tests/map_selected_key_projection_test.py b/paimon-python/pypaimon/tests/map_selected_key_projection_test.py new file mode 100644 index 000000000000..63670f25da50 --- /dev/null +++ b/paimon-python/pypaimon/tests/map_selected_key_projection_test.py @@ -0,0 +1,506 @@ +# 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 os +import shutil +import tempfile +import unittest + +import pyarrow as pa + +from pypaimon import CatalogFactory, Schema +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.schema.data_types import AtomicType +from pypaimon.schema.schema_change import SchemaChange + + +class MapSelectedKeyProjectionTest(unittest.TestCase): + + def setUp(self): + self.tmp = tempfile.mkdtemp() + self.catalog = CatalogFactory.create({ + 'warehouse': os.path.join(self.tmp, 'warehouse'), + }) + self.catalog.create_database('default', False) + + def tearDown(self): + shutil.rmtree(self.tmp, ignore_errors=True) + + def test_projects_literal_map_keys(self): + table = self._write_table('normal', {}) + + result = self._read(table, [ + 'id', "attributes['first']", + "attributes['key.with.dots']", + "attributes['missing']", + ]) + + self.assertEqual( + ['id', 'attributes_first', + 'attributes_key_with_dots', 'attributes_missing'], + result.column_names, + ) + self.assertEqual({ + 'id': [1, 2, 3], + 'attributes_first': [10, None, None], + 'attributes_key_with_dots': [20, None, None], + 'attributes_missing': [None, None, None], + }, result.to_pydict()) + + def test_projects_map_key_from_data_evolution_table(self): + table = self._write_table('data_evolution', { + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }) + + result = self._read(table, ['id', "attributes['first']"]) + + self.assertEqual({ + 'id': [1, 2, 3], + 'attributes_first': [10, None, None], + }, result.to_pydict()) + + def test_row_format_uses_full_map_fallback(self): + table = self._write_table('row_format', { + 'file.format': 'row', + }) + + result = self._read(table, ["attributes['first']"]) + + self.assertEqual( + [10, None, None], + result.column('attributes_first').to_pylist(), + ) + + def test_row_format_map_key_with_schema_evolution(self): + for evolution in ('rename', 'value_type'): + with self.subTest(evolution=evolution): + name = 'row_format_' + evolution + self._write_table( + name, + {'file.format': 'row'}, + value_type=pa.int32(), + ) + identifier = 'default.' + name + if evolution == 'rename': + change = SchemaChange.rename_column( + 'attributes', 'renamed_attributes') + projection = "renamed_attributes['first']" + else: + change = SchemaChange.update_column_type( + ['attributes', 'value'], AtomicType('BIGINT')) + projection = "attributes['first']" + self.catalog.alter_table(identifier, [change], False) + + result = self._read( + self.catalog.get_table(identifier), [projection]) + + self.assertEqual( + [10, None, None], result.column(0).to_pylist()) + if evolution == 'value_type': + self.assertEqual(pa.int64(), result.schema.field(0).type) + + def test_projects_map_key_with_row_tracking_fields(self): + table = self._write_table('row_tracking_fields', { + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }) + + for field in ['_ROW_ID', '_SEQUENCE_NUMBER']: + with self.subTest(field=field): + result = self._read( + table, ["attributes['first']", field]) + + self.assertEqual( + ['attributes_first', field], result.column_names) + self.assertEqual( + [10, None, None], + result.column('attributes_first').to_pylist(), + ) + self.assertEqual( + [0, 1, 2] if field == '_ROW_ID' else [1, 1, 1], + result.column(field).to_pylist(), + ) + + def test_rejects_filter_on_projected_map_key(self): + table = self._write_table('map_key_filter', {}) + builder = table.new_read_builder().with_projection( + ['id', "attributes['first']"]) + + with self.assertRaisesRegex( + NotImplementedError, 'Filtering projected MAP keys'): + builder.new_predicate_builder().equal('attributes_first', 10) + + def test_rejects_map_key_filter_after_projection_changes(self): + table = self._write_table('changed_projection_filter', {}) + builder = table.new_read_builder().with_projection( + ["attributes['first']"]) + predicate_builder = builder.new_predicate_builder() + builder.with_projection(['id']) + + with self.assertRaisesRegex( + NotImplementedError, 'Filtering projected MAP keys'): + predicate_builder.equal('attributes_first', 10) + + def test_filters_physical_column_with_map_key_projection(self): + table = self._write_table('physical_filter', {}) + builder = table.new_read_builder() + predicate = builder.new_predicate_builder().equal('id', 1) + builder.with_projection( + ["attributes['first']"]).with_filter(predicate) + + result = builder.new_read().to_arrow( + builder.new_scan().plan().splits()) + + self.assertEqual([10], result.column(0).to_pylist()) + + def test_data_evolution_filters_unprojected_physical_column(self): + table = self._write_table('data_evolution_physical_filter', { + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }) + builder = table.new_read_builder() + predicate = builder.new_predicate_builder().equal('id', 1) + builder.with_projection( + ["attributes['first']"]).with_filter(predicate) + + result = builder.new_read().to_arrow( + builder.new_scan().plan().splits()) + + self.assertEqual([10], result.column(0).to_pylist()) + + def test_physical_filter_with_conflicting_map_key_alias(self): + table = self._write_alias_collision_table('alias_filter') + builder = table.new_read_builder() + predicate = builder.new_predicate_builder().equal( + 'attributes_first', 100) + builder.with_projection( + ["attributes['first']"]).with_filter(predicate) + + result = builder.new_read().to_arrow( + builder.new_scan().plan().splits()) + + self.assertEqual(['attributes_first__0'], result.column_names) + self.assertEqual([10], result.column(0).to_pylist()) + + def test_projects_physical_column_with_conflicting_map_key_alias(self): + table = self._write_alias_collision_table('alias_projection') + builder = table.new_read_builder() + predicate = builder.new_predicate_builder().equal( + 'attributes_first', 100) + builder.with_projection([ + "attributes['first']", + 'attributes_first', + ]).with_filter(predicate) + + result = builder.new_read().to_arrow( + builder.new_scan().plan().splits()) + + self.assertEqual( + ['attributes_first__0', 'attributes_first'], + result.column_names, + ) + self.assertEqual({ + 'attributes_first__0': [10], + 'attributes_first': [100], + }, result.to_pydict()) + + def test_rejects_conflicting_derived_map_key_filter(self): + table = self._write_alias_collision_table('alias_derived_filter') + builder = table.new_read_builder().with_projection( + ["attributes['first']"]) + + with self.assertRaisesRegex( + NotImplementedError, 'Filtering projected MAP keys'): + builder.new_predicate_builder().equal('attributes_first__0', 10) + + def test_rejects_query_auth_for_projected_map_key(self): + table = self._write_table('map_key_auth', {}) + seen = [] + auth = TableQueryAuthResult( + filter=None, + column_masking={'attributes': '{"name":"NULL"}'}, + ) + table.catalog_environment.table_query_auth = ( + lambda options, identifier: lambda select: ( + seen.append(select) or auth)) + builder = table.new_read_builder().with_projection( + ['id', "attributes['first']"]) + splits = builder.new_scan().plan().splits() + + self.assertEqual([['id', 'attributes']], seen) + with self.assertRaisesRegex( + NotImplementedError, 'with query authorization'): + builder.new_read().to_arrow(splits) + + def test_unencodable_map_key_uses_full_map_fallback(self): + table = self._write_table('special_key', {}) + + result = self._read(table, ["attributes['a;b']"]) + + self.assertEqual( + [30, None, None], + result.column('attributes_a;b').to_pylist(), + ) + + def test_projects_key_from_dotted_map_column(self): + pa_schema = pa.schema([ + ('a.b', pa.map_(pa.string(), pa.int64())), + ]) + self.catalog.create_table( + 'default.dotted_map_column', + Schema.from_pyarrow_schema( + pa_schema, + options={'bucket': '-1', 'file.format': 'parquet'}, + ), + False, + ) + table = self.catalog.get_table('default.dotted_map_column') + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(pa.Table.from_arrays([ + pa.array([[('c', 10)]], type=pa_schema.field('a.b').type), + ], schema=pa_schema)) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + + result = self._read(table, ["a.b['c']"]) + + self.assertEqual([10], result.column('a.b_c').to_pylist()) + + def test_full_map_and_selected_key_use_full_map_fallback(self): + table = self._write_table('full_and_selected', {}) + + result = self._read( + table, ['attributes', "attributes['first']"]) + + self.assertEqual( + ['attributes', 'attributes_first'], result.column_names) + self.assertEqual( + [10, None, None], + result.column('attributes_first').to_pylist(), + ) + self.assertEqual( + [[('first', 10), ('key.with.dots', 20), ('a;b', 30)], [], None], + result.column('attributes').to_pylist(), + ) + + def test_projects_map_blob_keys(self): + map_type = pa.map_(pa.string(), pa.large_binary()) + pa_schema = pa.schema([ + ('id', pa.int32()), + ('payload', map_type), + ]) + self.catalog.create_table( + 'default.map_blob', + Schema.from_pyarrow_schema( + pa_schema, + options={ + 'row-tracking.enabled': 'true', + 'data-evolution.enabled': 'true', + }, + ), + False, + ) + table = self.catalog.get_table('default.map_blob') + data = pa.Table.from_arrays([ + pa.array([1, 2, 3, 4], type=pa.int32()), + pa.array([ + [('k', b'hello'), ('v', b'world')], + [('k', None)], + [], + None, + ], type=map_type), + ], schema=pa_schema) + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(data) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + + full = self._read(table, ['payload']).column('payload').to_pylist() + selected = self._read( + table, ["payload['k']", "payload['missing']"]) + + self.assertEqual( + [b'hello', None, None, None], + selected.column('payload_k').to_pylist(), + ) + self.assertEqual( + [None, None, None, None], + selected.column('payload_missing').to_pylist(), + ) + self.assertEqual( + [None if row is None else dict(row).get('k') for row in full], + selected.column('payload_k').to_pylist(), + ) + + row_ids = self._read(table, ['id', '_ROW_ID']).to_pylist() + update_builder = table.new_batch_write_builder() + update = update_builder.new_update().with_update_type(['payload']) + messages = update.update_by_arrow_with_row_id(pa.Table.from_pydict({ + '_ROW_ID': pa.array([row_ids[0]['_ROW_ID']], type=pa.int64()), + 'payload': pa.array([[('k', b'updated')]], type=map_type), + })) + update_builder.new_commit().commit(messages) + + full = self._read(table, ['payload']).column('payload').to_pylist() + selected = self._read(table, ["payload['k']"]) + self.assertEqual( + [b'updated', None, None, None], + selected.column('payload_k').to_pylist(), + ) + self.assertEqual( + [None if row is None else dict(row).get('k') for row in full], + selected.column('payload_k').to_pylist(), + ) + + from pypaimon.table.row.blob import BlobDescriptor + descriptors = self._read( + table.copy({'blob-as-descriptor': 'true'}), + ["payload['k']"], + ).column('payload_k').to_pylist() + self.assertIsInstance( + BlobDescriptor.deserialize(descriptors[0]), BlobDescriptor) + + def test_projects_map_key_after_column_rename(self): + self._write_table('renamed', {}) + self.catalog.alter_table( + 'default.renamed', + [SchemaChange.rename_column('attributes', 'renamed_attributes')], + False, + ) + + result = self._read( + self.catalog.get_table('default.renamed'), + ['id', "renamed_attributes['first']"], + ) + + self.assertEqual({ + 'id': [1, 2, 3], + 'renamed_attributes_first': [10, None, None], + }, result.to_pydict()) + + def test_projects_map_key_across_value_type_evolution(self): + value_type = pa.struct([('nested_value', pa.int32())]) + pa_schema = pa.schema([ + ('id', pa.int32()), + ('attributes', pa.map_(pa.string(), value_type)), + ]) + self.catalog.create_table( + 'default.evolved', + Schema.from_pyarrow_schema( + pa_schema, + options={'bucket': '-1', 'file.format': 'parquet'}, + ), + False, + ) + table = self.catalog.get_table('default.evolved') + data = pa.Table.from_arrays([ + pa.array([1], type=pa.int32()), + pa.array([ + [('first', {'nested_value': 10})], + ], type=pa.map_(pa.string(), value_type)), + ], schema=pa_schema) + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(data) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + self.catalog.alter_table( + 'default.evolved', + [SchemaChange.update_column_type( + ['attributes', 'value', 'nested_value'], + AtomicType('BIGINT'), + )], + False, + ) + + result = self._read( + self.catalog.get_table('default.evolved'), + ["attributes['first']"], + ) + + self.assertEqual( + pa.struct([('nested_value', pa.int64())]), + result.schema.field(0).type, + ) + self.assertEqual( + [{'nested_value': 10}], result.column(0).to_pylist()) + + def _write_table(self, name, extra_options, value_type=pa.int64()): + pa_schema = pa.schema([ + ('id', pa.int32()), + ('attributes', pa.map_(pa.string(), value_type)), + ]) + options = {'bucket': '-1', 'file.format': 'parquet'} + options.update(extra_options) + schema = Schema.from_pyarrow_schema(pa_schema, options=options) + identifier = 'default.{}'.format(name) + self.catalog.create_table(identifier, schema, False) + table = self.catalog.get_table(identifier) + data = pa.Table.from_arrays([ + pa.array([1, 2, 3], type=pa.int32()), + pa.array([ + [('first', 10), ('key.with.dots', 20), ('a;b', 30)], + [], + None, + ], type=pa.map_(pa.string(), value_type)), + ], schema=pa_schema) + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(data) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + return self.catalog.get_table(identifier) + + def _write_alias_collision_table(self, name): + pa_schema = pa.schema([ + ('attributes', pa.map_(pa.string(), pa.int64())), + ('attributes_first', pa.int64()), + ]) + self.catalog.create_table( + 'default.' + name, + Schema.from_pyarrow_schema( + pa_schema, + options={'bucket': '-1', 'file.format': 'parquet'}, + ), + False, + ) + table = self.catalog.get_table('default.' + name) + data = pa.Table.from_arrays([ + pa.array([ + [('first', 10)], + [('first', 20)], + ], type=pa_schema.field('attributes').type), + pa.array([100, 200], type=pa.int64()), + ], schema=pa_schema) + builder = table.new_batch_write_builder() + writer = builder.new_write() + writer.write_arrow(data) + builder.new_commit().commit(writer.prepare_commit()) + writer.close() + return self.catalog.get_table('default.' + name) + + @staticmethod + def _read(table, projection): + builder = table.new_read_builder().with_projection(projection) + return builder.new_read().to_arrow(builder.new_scan().plan().splits()) + + +if __name__ == '__main__': + unittest.main() diff --git a/paimon-python/pypaimon/tests/test_nested_projection_e2e.py b/paimon-python/pypaimon/tests/test_nested_projection_e2e.py index a775a6dffd33..40e53b672b29 100644 --- a/paimon-python/pypaimon/tests/test_nested_projection_e2e.py +++ b/paimon-python/pypaimon/tests/test_nested_projection_e2e.py @@ -54,7 +54,7 @@ def setUpClass(cls): def tearDownClass(cls): shutil.rmtree(cls.tempdir, ignore_errors=True) - def _create_table(self, name: str, file_format: str = 'parquet'): + def _create_table(self, name: str, file_format: str = 'parquet', rows=None): identifier = 'default.{}'.format(name) schema = Schema.from_pyarrow_schema( self.pa_schema, @@ -64,7 +64,10 @@ def _create_table(self, name: str, file_format: str = 'parquet'): table = self.catalog.get_table(identifier) wb = table.new_batch_write_builder() w = wb.new_write() - w.write_arrow(pa.Table.from_pylist(self.rows, schema=self.pa_schema)) + w.write_arrow(pa.Table.from_pylist( + self.rows if rows is None else rows, + schema=self.pa_schema, + )) wb.new_commit().commit(w.prepare_commit()) w.close() return table @@ -84,6 +87,17 @@ def test_dotted_name_returns_just_the_leaf(self): {'mv_latest_version': 200}, {'mv_latest_version': 300}]) + def test_nested_leaf_preserves_parent_null(self): + table = self._create_table('ao_nullable_parent', rows=[ + self.rows[0], + {'id': 2, 'mv': None, 'val': 'y'}, + ]) + + rb = table.new_read_builder().with_projection(['mv.latest_version']) + got = rb.new_read().to_arrow(rb.new_scan().plan().splits()) + + self.assertEqual([100, None], got.column(0).to_pylist()) + def test_mixed_nested_and_top_level_preserves_order(self): table = self._create_table('ao_mixed_order') rb = table.new_read_builder().with_projection( @@ -117,6 +131,44 @@ def test_dotted_top_level_field_kept(self): got = rb.new_read().to_arrow(rb.new_scan().plan().splits()).to_pylist() self.assertEqual(got, [{'id': 1, 'media.left': 'hello'}]) + def test_row_path_precedes_dotted_top_level_prefix(self): + pa_schema = pa.schema([ + ('a', pa.struct([ + ('b', pa.struct([('c', pa.int64())])), + ])), + ('a.b', pa.struct([ + ('c', pa.int64()), + ('d', pa.int64()), + ])), + ('id', pa.int64()), + ]) + identifier = 'default.ao_row_path_precedence' + self.catalog.create_table( + identifier, + Schema.from_pyarrow_schema(pa_schema, options={'bucket': '-1'}), + False) + table = self.catalog.get_table(identifier) + wb = table.new_batch_write_builder() + w = wb.new_write() + w.write_arrow(pa.Table.from_arrays([ + pa.array([{'b': {'c': 1}}], type=pa_schema.field('a').type), + pa.array([{'c': 99, 'd': 88}], + type=pa_schema.field('a.b').type), + pa.array([7], type=pa.int64()), + ], schema=pa_schema)) + wb.new_commit().commit(w.prepare_commit()) + w.close() + + rb = table.new_read_builder().with_projection([ + 'a.b.c', 'a.b.d', 'id', + ]) + got = rb.new_read().to_arrow(rb.new_scan().plan().splits()) + + self.assertEqual( + {'a_b_c': [1], 'a.b_d': [88], 'id': [7]}, + got.to_pydict(), + ) + def test_unknown_dotted_name_silently_skipped(self): pa_schema = pa.schema([ ('id', pa.int64()), diff --git a/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py b/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py index c80ec9791b69..61e7396e404e 100644 --- a/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py +++ b/paimon-python/pypaimon/tests/test_read_builder_nested_projection.py @@ -45,6 +45,7 @@ def setUpClass(cls): pa.field('pk', pa.int64(), nullable=False), ('mv', struct_type), ('val', pa.string()), + ('attrs', pa.map_(pa.string(), pa.int64())), ]) schema = Schema.from_pyarrow_schema( cls.pa_schema, primary_keys=['pk'], @@ -63,7 +64,7 @@ def test_no_projection_returns_full_schema(self): rb = self.table.new_read_builder() fields = rb.read_type() names = [f.name for f in fields] - self.assertEqual(names, ['pk', 'mv', 'val']) + self.assertEqual(names, ['pk', 'mv', 'val', 'attrs']) # Without an explicit projection the read_type must NOT inject # row-tracking system columns; the raw table fields are returned # verbatim. @@ -103,6 +104,27 @@ def test_dotted_name_unknown_subfield_silently_skipped(self): names = [f.name for f in rb.read_type()] self.assertEqual(names, ['pk']) + def test_bracketed_map_selector_is_one_literal_key(self): + rb = self.table.new_read_builder().with_projection( + ["attrs['key.with.dots']", 'attrs["other"]']) + + self.assertEqual( + [['attrs', 'key.with.dots'], ['attrs', 'other']], + rb._nested_name_paths(), + ) + self.assertEqual( + ['attrs_key_with_dots', 'attrs_other'], + [field.name for field in rb.read_type()], + ) + self.assertEqual( + ['attrs'], [field.name for field in rb.new_scan()._read_type]) + + def test_dot_does_not_select_map_key(self): + rb = self.table.new_read_builder().with_projection( + ['attrs.other', 'pk']) + + self.assertEqual(['pk'], [field.name for field in rb.read_type()]) + class ReadBuilderProjectionFieldIdTest(_ReadBuilderTestBase): diff --git a/paimon-python/pypaimon/utils/projection.py b/paimon-python/pypaimon/utils/projection.py index f0f5a569e6fd..9be66dad10fd 100644 --- a/paimon-python/pypaimon/utils/projection.py +++ b/paimon-python/pypaimon/utils/projection.py @@ -21,18 +21,23 @@ columns the user wants to read. Two flavours: * :class:`TopLevelProjection` selects fields by their top-level index. -* :class:`NestedProjection` accepts paths that walk into ROW children, e.g. - ``[[1, 0], [1, 2]]`` means "the 0th and 2nd children of the field at top - level index 1". The result is flattened into top-level fields whose - names are the underscore-joined original path (``a_b`` for ``a.b``, - with a ``__N`` suffix on collisions) and whose IDs are inherited from - the leaf so schema-evolution remapping by field ID still works. +* :class:`NestedProjection` accepts paths that walk into ROW children or end + in a literal MAP key. The result is flattened into top-level fields whose + names are the underscore-joined original path (``a_b`` for ``a.b``, with a + ``__N`` suffix on collisions). """ from abc import ABC, abstractmethod -from typing import List, Optional, Sequence +from copy import copy +from typing import List, NamedTuple, Optional, Sequence -from pypaimon.schema.data_types import DataField, RowType +from pypaimon.schema.data_types import DataField, MapType, RowType + + +class MapKey(NamedTuple): + """A literal MAP key in an internal projection path.""" + + value: str class Projection(ABC): @@ -57,11 +62,11 @@ def to_top_level_indexes(self) -> List[int]: @abstractmethod def to_nested_indexes(self) -> List[List[int]]: - """Return the projection as a list of paths, one per output field.""" + """Return the internal paths, one per output field.""" @abstractmethod def to_name_paths(self, row_type) -> List[List[str]]: - """Translate integer paths to field-name paths against ``row_type``. + """Translate internal paths to field-name paths against ``row_type``. For a path ``[1, 0]`` against a row whose top-level field at index 1 is a struct ``mv_col`` with sub-fields ``[LATEST_VERSION, ...]``, @@ -153,13 +158,14 @@ def to_name_paths(self, row_type) -> List[List[str]]: class NestedProjection(Projection): - """Projection over paths that may walk into ROW children. + """Projection over paths that may walk into ROW children or MAP keys. Each path navigates from a top-level field through successive ROW - children. A path of length 1 is equivalent to a top-level selection. + children, or ends with a literal :class:`MapKey`. A path of length 1 is + equivalent to a top-level selection. """ - def __init__(self, paths: Sequence[Sequence[int]]): + def __init__(self, paths: Sequence[Sequence]): if not paths: raise ValueError("NestedProjection requires at least one path") self.paths = [list(p) for p in paths] @@ -194,6 +200,13 @@ def to_name_paths(self, row_type) -> List[List[str]]: names = [field.name] for idx in path[1:]: child_type = field.type + if isinstance(idx, MapKey): + if not isinstance(child_type, MapType): + raise ValueError( + "Map-key projection expected a MAP type but got %s " + "for field '%s'" % (child_type, field.name)) + names.append(idx.value) + break if not is_row_type(child_type): raise ValueError( "Nested projection step expected a ROW type but got %s " @@ -206,14 +219,31 @@ def to_name_paths(self, row_type) -> List[List[str]]: def project(self, row_type) -> List[DataField]: fields = _row_fields(row_type) + top_level_names = {field.name for field in fields} out: List[DataField] = [] seen_names = set() dup_count = 0 for path in self.paths: field = fields[path[0]] name_parts = [field.name] + is_map_key = False for idx in path[1:]: child_type = field.type + if isinstance(idx, MapKey): + if not isinstance(child_type, MapType): + raise ValueError( + "Map-key projection expected a MAP type but got %s " + "for field '%s'" % (child_type, field.name)) + value_type = copy(child_type.value) + value_type.nullable = True + field = DataField( + id=field.id, + name=idx.value, + type=value_type, + ) + name_parts.extend(idx.value.split('.')) + is_map_key = True + break if not is_row_type(child_type): raise ValueError( "Nested projection step expected a ROW type but got %s " @@ -223,7 +253,8 @@ def project(self, row_type) -> List[DataField]: name_parts.append(field.name) base_name = "_".join(name_parts) final_name = base_name - while final_name in seen_names: + while (final_name in seen_names + or (is_map_key and final_name in top_level_names)): final_name = "%s__%d" % (base_name, dup_count) dup_count += 1 seen_names.add(final_name)