You are viewing a plain text version of this content. The canonical link for it is here.
Posted to commits@paimon.apache.org by ho...@apache.org on 2026/09/20 03:51:25 UTC
(paimon) 11/20: [python] Support selected-key projection for shared-shredding MAP (#9789)
This is an automated email from the ASF dual-hosted git repository.
XiaoHongbo-Hope pushed a commit to branch release-2.1
in repository https://gitbox.apache.org/repos/asf/paimon.git
commit d716a5e160243c867939bf509f69de622fe16f42
Author: XiaoHongbo <xi...@alibaba-inc.com>
AuthorDate: Tue Sep 15 14:33:40 2026 +0800
[python] Support selected-key projection for shared-shredding MAP (#9789)
---
docs/docs/pypaimon/reading.md | 33 +-
paimon-python/dev/run_mixed_tests.sh | 3 +-
.../pypaimon/data/map_shared_shredding.py | 262 ++++++++++-
paimon-python/pypaimon/read/read_builder.py | 175 +++++--
.../pypaimon/read/reader/data_file_batch_reader.py | 35 ++
.../pypaimon/read/reader/format_pyarrow_reader.py | 102 ++++-
.../read/reader/nested_leaf_batch_reader.py | 17 +-
paimon-python/pypaimon/read/split_read.py | 35 +-
paimon-python/pypaimon/read/table_read.py | 75 ++-
.../pypaimon/tests/data_evolution_formats_test.py | 111 +++++
.../pypaimon/tests/e2e/java_py_read_write_test.py | 26 ++
.../format_pyarrow_shared_shredding_map_test.py | 295 ++++++++++--
.../tests/map_selected_key_projection_test.py | 506 +++++++++++++++++++++
.../pypaimon/tests/test_nested_projection_e2e.py | 56 ++-
.../tests/test_read_builder_nested_projection.py | 24 +-
paimon-python/pypaimon/utils/projection.py | 59 ++-
16 files changed, 1678 insertions(+), 136 deletions(-)
diff --git a/docs/docs/pypaimon/reading.md b/docs/docs/pypaimon/reading.md
index 21fcb63192..ec0d09c12e 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<name STRING, age INT>, 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<ROW>` 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<ROW>` 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 66ff2b3e5f..620c4d989a 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 ecadf807a5..a3d6c58bac 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 = 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_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 a531d5d2c8..d64f863b9d 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 @@ class ReadBuilder:
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 @@ class 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<STRING, ...>`` 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 @@ class 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 @@ class ReadBuilder:
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 @@ class ReadBuilder:
# 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 @@ class ReadBuilder:
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 2bdf38689e..bd98bccf2a 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 @@ import pyarrow.compute as pc
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 @@ class DataFileBatchReader(RecordBatchReader):
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 @@ class DataFileBatchReader(RecordBatchReader):
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 @@ class DataFileBatchReader(RecordBatchReader):
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 672fdcfa8a..f3403e260b 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.file_io import FileIO
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
# 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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 @@ class FormatPyArrowReader(RecordBatchReader):
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 cd849454a7..59bf7def73 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 as pa
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 @@ class NestedLeafBatchReader(RecordBatchReader):
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 9073072b48..2d35830ecc 100644
--- a/paimon-python/pypaimon/read/split_read.py
+++ b/paimon-python/pypaimon/read/split_read.py
@@ -28,6 +28,11 @@ from pypaimon.deletionvectors import (
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 @@ from pypaimon.read.reader.sort_merge_reader import (SortMergeReaderWithMinHeap,
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 @@ class SplitRead(ABC):
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 @@ class SplitRead(ABC):
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
@@ -1196,7 +1224,6 @@ class DataEvolutionSplitRead(SplitRead):
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 980d8fdd9c..c5b7af32df 100644
--- a/paimon-python/pypaimon/read/table_read.py
+++ b/paimon-python/pypaimon/read/table_read.py
@@ -25,6 +25,7 @@ import pyarrow
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 import Split
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 @@ class TableRead:
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 5bfb82824b..c970387995 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 import CatalogFactory, Schema
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 @@ class DataEvolutionFormatsTest(unittest.TestCase):
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 618920bd5e..86d1255ecf 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 @@ class JavaPyReadWriteTest(unittest.TestCase):
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 11b57420c0..4fd2922032 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.fs as pafs
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 @@ class _LocalFileIO:
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
path or self._write(compression, file_format),
[DataField(
0,
- "content_refs",
+ "attributes",
MapType(
True, AtomicType("STRING", False), self.value_type),
)],
@@ -169,19 +315,19 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
_LocalFileIO(), "orc", path,
[DataField(
0,
- "content_refs",
+ "attributes",
MapType(
True,
AtomicType("STRING", False),
@@ -375,19 +576,19 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 @@ class SharedShreddingMapReaderTest(unittest.TestCase):
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 0000000000..63670f25da
--- /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 a775a6dffd..40e53b672b 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 @@ class _AppendOnlyNestedBase(unittest.TestCase):
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 @@ class _AppendOnlyNestedBase(unittest.TestCase):
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 @@ class AppendOnlyNestedParquetTest(_AppendOnlyNestedBase):
{'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 @@ class AppendOnlyNestedParquetTest(_AppendOnlyNestedBase):
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 c80ec9791b..61e7396e40 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 @@ class _ReadBuilderTestBase(unittest.TestCase):
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 @@ class ReadBuilderProjectionStateTest(_ReadBuilderTestBase):
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 @@ class ReadBuilderProjectionStateTest(_ReadBuilderTestBase):
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 f0f5a569e6..9be66dad10 100644
--- a/paimon-python/pypaimon/utils/projection.py
+++ b/paimon-python/pypaimon/utils/projection.py
@@ -21,18 +21,23 @@ A projection maps a source row type to a flat list of ``DataField``: the
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 @@ class Projection(ABC):
@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 @@ class TopLevelProjection(Projection):
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 @@ class NestedProjection(Projection):
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 @@ class NestedProjection(Projection):
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 @@ class NestedProjection(Projection):
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)