From b85cc31ba382d3fc1948d571a78be9a60b855273 Mon Sep 17 00:00:00 2001 From: QuakeWang Date: Tue, 20 Oct 2026 17:08:14 +0800 Subject: [PATCH] [python] Support named column and nested-path projections in row iterators Signed-off-by: QuakeWang --- paimon-python/pypaimon/read/table_read.py | 8 +- .../tests/named_projection_iterator_test.py | 269 ++++++++++++++++++ .../pypaimon/tests/native_iterator_test.py | 47 ++- 3 files changed, 309 insertions(+), 15 deletions(-) create mode 100644 paimon-python/pypaimon/tests/named_projection_iterator_test.py diff --git a/paimon-python/pypaimon/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 6e6b9326d769..b7670b4d4e9a 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -219,10 +219,14 @@ def to_iterator(self, splits: List[Split]) -> Iterator: Rows retain their kind separately from projected fields. Like the Python split reader, the iterator reuses its row view and visits splits in order. Close the iterator when stopping before exhaustion. + + Named columns and nested paths retain their projection order. Named + VARIANT extraction expressions require the Arrow APIs. """ - if self.output_projection is not None and self.output_projection.named: + if (self.output_projection is not None and self.output_projection.named + and has_variant_extractions(self.read_type)): raise RuntimeError( - "Projection expressions are not supported by to_iterator(); " + "Variant extraction expressions are not supported by to_iterator(); " "use to_arrow() or to_arrow_batch_reader()") self._begin_auth_read(splits) limit = self.limit diff --git a/paimon-python/pypaimon/tests/named_projection_iterator_test.py b/paimon-python/pypaimon/tests/named_projection_iterator_test.py new file mode 100644 index 000000000000..8ec4e6bd9d4d --- /dev/null +++ b/paimon-python/pypaimon/tests/named_projection_iterator_test.py @@ -0,0 +1,269 @@ +# 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. + +from contextlib import ExitStack +from unittest.mock import patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.catalog.table_query_auth import TableQueryAuthResult +from pypaimon.casting.schema_evolution import read_arrow_schema +from pypaimon.schema.data_types import AtomicType, DataField, RowType +from pypaimon.table.row.internal_row import InternalRow +from pypaimon.table.row.row_kind import RowKind + + +pytestmark = pytest.mark.python_read + + +@pytest.fixture(params=['append', 'pk']) +def projected_table(tmp_path, request): + schema = pa.schema([ + pa.field('id', pa.int64(), nullable=False), + ('v', pa.int64()), + ('profile', pa.struct([('city', pa.string())])), + ('attributes', pa.map_(pa.string(), pa.string())), + ]) + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('db', True) + catalog.create_table('db.t', Schema.from_pyarrow_schema( + schema, primary_keys=['id'] if request.param == 'pk' else None, + options={'bucket': '1' if request.param == 'pk' else '-1', + 'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + 'write.native.enabled': 'false', 'read.batch-size': '2'}), False) + table = catalog.get_table('db.t') + rows = [ + {'id': 1, 'v': 10, 'profile': {'city': 'Hangzhou'}, + 'attributes': [('source', 'web'), ('source', 'mobile')]}, + {'id': 2, 'v': 20, 'profile': None, + 'attributes': [('source', None), ('source', 'later')]}, + {'id': 3, 'v': 30, 'profile': {'city': None}, 'attributes': [('other', 'value')]}, + {'id': 4, 'v': 40, 'profile': {'city': 'Shenzhen'}, 'attributes': None}, + {'id': 5, 'v': 50, 'profile': {'city': 'Beijing'}, 'attributes': []}, + ] + writes = [rows] + expected = [(1, 'Hangzhou', 'web', 10), (2, None, None, 20), + (3, None, None, 30), (4, 'Shenzhen', None, 40), (5, 'Beijing', None, 50)] + if request.param == 'pk': + writes.append([dict(rows[0], v=11, profile={'city': 'Suzhou'}, + attributes=[('source', 'api'), ('source', 'later')])]) + expected[0] = (1, 'Suzhou', 'api', 11) + for values in writes: + builder = table.new_batch_write_builder() + writer, commit = builder.new_write(), builder.new_commit() + try: + writer.write_arrow(pa.Table.from_pylist(values, schema=schema)) + commit.commit(writer.prepare_commit()) + finally: + writer.close() + commit.close() + return table, expected + + +def _rows(read, splits): + iterator = read.to_iterator(splits) + try: + result = [] + for row in iterator: + assert isinstance(row, InternalRow) + assert row.get_row_kind() == RowKind.INSERT + result.append(tuple(row.get_field(i) for i in range(len(row)))) + return result + finally: + iterator.close() + + +@pytest.mark.parametrize('projection,positions', [ + ({'user_id': 'id'}, [0]), + ({'city': 'profile.city'}, [1]), + ({'source': "attributes['source']"}, [2]), + ({'user_id': 'id', 'city': 'profile.city', 'source': "attributes['source']"}, [0, 1, 2]), + ({'city': 'profile.city', 'source': "attributes['source']", 'user_id': 'id', + 'city_again': 'profile.city', 'id_again': 'id', 'source_again': "attributes['source']"}, + [1, 2, 0, 1, 0, 2]), +]) +def test_named_iterator_matches_arrow(projected_table, projection, positions): + table, values = projected_table + builder = table.new_read_builder().with_projection(projection) + read, splits = builder.new_read(), builder.new_scan().plan().splits() + expected = [tuple(row[i] for i in positions) for row in values] + actual = _rows(read, splits) + assert sorted(actual, key=repr) == sorted(expected, key=repr) + batches = read.to_arrow_batch_reader(splits, parallelism=1) + try: + streaming = batches.read_all() + finally: + batches.close() + for arrow in (read.to_arrow(splits, parallelism=1), streaming): + assert arrow.column_names == list(projection) + assert sorted(zip(*(column.to_pylist() for column in arrow.columns)), key=repr) == sorted(actual, key=repr) + + +def test_named_iterator_keeps_filter_source_binding(projected_table): + table, _ = projected_table + predicate = table.new_read_builder().new_predicate_builder().greater_than('v', 20) + builder = table.new_read_builder().with_projection({ + 'v': 'id', 'id': 'v', 'city': 'profile.city', + }).with_filter(predicate) + assert sorted(_rows(builder.new_read(), builder.new_scan().plan().splits())) == [ + (3, 30, None), (4, 40, 'Shenzhen'), (5, 50, 'Beijing')] + + +@pytest.mark.parametrize('limit', [0, 1, 3, None]) +def test_named_iterator_limit_and_empty_input(projected_table, limit): + table, values = projected_table + builder = table.new_read_builder().with_projection({ + 'user_id': 'id', 'city': 'profile.city', 'source': "attributes['source']", + }) + if limit is not None: + builder.with_limit(limit) + read, splits = builder.new_read(), builder.new_scan().plan().splits() + actual = _rows(read, splits) + assert len(actual) == (len(values) if limit is None else limit) + assert all(row in [value[:3] for value in values] for row in actual) + assert _rows(read, []) == [] + + +@pytest.mark.parametrize('consume', [False, True]) +def test_named_iterator_is_lazy_and_closes_reader(projected_table, consume): + table, _ = projected_table + builder = table.new_read_builder().with_projection({'city': 'profile.city'}) + read, splits = builder.new_read(), builder.new_scan().plan().splits() + create = read._create_split_read + readers = [] + + def create_split_read(*args, **kwargs): + split_read = create(*args, **kwargs) + create_reader = split_read.create_reader + + def tracked_reader(): + reader = create_reader() + readers.append(stack.enter_context(patch.object(reader, 'close', wraps=reader.close))) + return reader + + split_read.create_reader = tracked_reader + return split_read + + with ExitStack() as stack, patch.object(read, '_create_split_read', side_effect=create_split_read): + iterator = read.to_iterator(splits) + assert readers == [] + try: + if consume: + assert isinstance(next(iterator), InternalRow) + assert len(readers) == 1 + readers[0].assert_not_called() + finally: + iterator.close() + if consume: + readers[0].assert_called_once_with() + else: + assert readers == [] + + +@pytest.mark.parametrize('stop', ['exhaust', 'close', 'limit', 'zero_limit']) +def test_named_iterator_adapts_native_batches(projected_table, stop): + table, _ = projected_table + builder = table.copy({'read.native.enabled': 'true'}).new_read_builder().with_projection({ + 'city': 'profile.city', 'key': 'id', 'source': "attributes['source']", + 'copy': 'profile.city', 'rowkind': 'id', + }) + if stop in ('limit', 'zero_limit'): + builder.with_limit(1 if stop == 'limit' else 0) + read, splits = builder.new_read(), builder.new_scan().plan().splits() + schema = pa.schema([pa.field('rowkind', pa.string())] + list(read_arrow_schema(read.read_type))) + batch = pa.RecordBatch.from_pylist([ + {'rowkind': '+I', 'id': 1, 'profile': {'city': 'Hangzhou'}, 'attributes': {'source': 'web'}}, + {'rowkind': '-U', 'id': 2, 'profile': None, 'attributes': {'source': None}}, + {'rowkind': '+U', 'id': 3, 'profile': {'city': None}, 'attributes': None}, + {'rowkind': '-D', 'id': 4, 'profile': {'city': 'Beijing'}, 'attributes': {'source': 'api'}}, + ], schema=schema) + closed = [] + + def batches(): + try: + yield batch + finally: + closed.append(True) + + def converted(native_read, splits, output_schema, **kwargs): + assert native_read.include_row_kind + assert kwargs == {'parallelism': 1, 'streaming': True, 'preserve_split_order': True} + return native_read._convert_native_batches(batches(), output_schema) + + with patch.object(type(read), '_try_native_batches', autospec=True, side_effect=converted) as native, \ + patch.object(read, '_create_split_read', side_effect=AssertionError('Python fallback')): + iterator = read.to_iterator(splits) + native.assert_not_called() + try: + if stop == 'close': + row = next(iterator) + actual = [(tuple(row.get_field(i) for i in range(len(row))), row.get_row_kind())] + assert closed == [] + else: + actual = [(tuple(row.get_field(i) for i in range(len(row))), row.get_row_kind()) + for row in iterator] + finally: + iterator.close() + expected = [ + (('Hangzhou', 1, 'web', 'Hangzhou', 1), RowKind.INSERT), + ((None, 2, None, None, 2), RowKind.UPDATE_BEFORE), + ((None, 3, None, None, 3), RowKind.UPDATE_AFTER), + (('Beijing', 4, 'api', 'Beijing', 4), RowKind.DELETE), + ] + assert actual == ([] if stop == 'zero_limit' else expected if stop == 'exhaust' else expected[:1]) + assert closed == ([] if stop == 'zero_limit' else [True]) + assert read.include_row_kind is False + + +def test_named_iterator_keeps_map_key_authorization_boundary(projected_table): + table, _ = projected_table + auth = TableQueryAuthResult(filter=None, column_masking={'attributes': '{"name":"NULL"}'}) + table.catalog_environment.table_query_auth = lambda options, identifier: lambda select: auth + builder = table.new_read_builder().with_projection({'source': "attributes['source']"}) + splits = builder.new_scan().plan().splits() + with pytest.raises(NotImplementedError, match='with query authorization'): + _rows(builder.new_read(), splits) + + +def test_named_projection_keeps_torch_row_boundary(projected_table): + table, _ = projected_table + read = table.new_read_builder().with_projection({'user_id': 'id'}).new_read() + with pytest.raises(RuntimeError, match='not supported by Torch row format'): + read.to_torch([]) + + +@pytest.mark.parametrize('function', ['variant_get', 'try_variant_get']) +@pytest.mark.parametrize('nested', [False, True]) +@pytest.mark.parametrize('mixed', [False, True]) +def test_named_iterator_rejects_variant_extractions(tmp_path, function, nested, mixed): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + catalog.create_database('db', True) + payload = DataField(2 if nested else 1, 'payload', AtomicType('VARIANT')) + fields = [DataField(0, 'id', AtomicType('BIGINT')), + DataField(1, 'profile', RowType(True, [payload])) if nested else payload] + catalog.create_table('db.t', Schema(fields, options={'bucket': '-1'}), False) + source = 'profile.payload' if nested else 'payload' + projection = {'x': "%s(%s, '$.x', 'float')" % (function, source)} + if mixed: + projection['user_id'] = 'id' + read = catalog.get_table('db.t').new_read_builder().with_projection(projection).new_read() + with patch.object(read, '_try_native_batches') as native: + with pytest.raises(RuntimeError, match='Variant extraction expressions are not supported by to_iterator'): + read.to_iterator([]) + native.assert_not_called() diff --git a/paimon-python/pypaimon/tests/native_iterator_test.py b/paimon-python/pypaimon/tests/native_iterator_test.py index 4f20936b75ea..41db0c8d498f 100644 --- a/paimon-python/pypaimon/tests/native_iterator_test.py +++ b/paimon-python/pypaimon/tests/native_iterator_test.py @@ -81,7 +81,8 @@ def _rows(read, splits): @pytest.mark.parametrize('mode', ['append', 'de', 'pk']) -@pytest.mark.parametrize('projection', [None, ['v', 'id'], ['id', 'id'], []]) +@pytest.mark.parametrize('projection', [ + None, ['v', 'id'], ['id', 'id'], [], {'value': 'v', 'key': 'id', 'copy': 'id'}]) @pytest.mark.parametrize('limit', [None, 0, 1, 5]) def test_iterator_uses_native_read(tmp_path, mode, projection, limit): table, _ = _table(tmp_path, mode) @@ -138,17 +139,18 @@ def _batch(ids, kinds=None): pa.array(kinds or ['+I'] * len(ids)), pa.array(ids, type=pa.int32())], names=['rowkind', 'id']) -def _planned_read(tmp_path, limit=None): +def _planned_read(tmp_path, limit=None, projection=None): table, _ = _table(tmp_path) - builder = table.new_read_builder().with_projection(['id']) + builder = table.new_read_builder().with_projection(['id'] if projection is None else projection) if limit is not None: builder.with_limit(limit) return builder.new_read(), builder.new_scan().plan().splits() @pytest.mark.parametrize('consume', [False, True]) -def test_iterator_is_lazy_and_close_stops_native_input(tmp_path, consume): - read, splits = _planned_read(tmp_path) +@pytest.mark.parametrize('projection', [['id'], {'identifier': 'id'}]) +def test_iterator_is_lazy_and_close_stops_native_input(tmp_path, consume, projection): + read, splits = _planned_read(tmp_path, projection=projection) batches = _Batches([_batch([1, 2]), RuntimeError('must not fetch next batch')]) with patch('pypaimon.read.native_plan._prepare_native_read', return_value=lambda _: batches) as prepare, \ patch.object(read, '_create_split_read', side_effect=AssertionError('Python fallback')): @@ -165,8 +167,9 @@ def test_iterator_is_lazy_and_close_stops_native_input(tmp_path, consume): @pytest.mark.parametrize('failure', ['read', 'schema', 'kind']) -def test_iterator_closes_native_input_on_failure_without_replaying_python(tmp_path, failure): - read, splits = _planned_read(tmp_path) +@pytest.mark.parametrize('projection', [['id'], {'identifier': 'id'}]) +def test_iterator_closes_native_input_on_failure_without_replaying_python(tmp_path, failure, projection): + read, splits = _planned_read(tmp_path, projection=projection) bad = (RuntimeError('native read failed') if failure == 'read' else pa.record_batch([pa.array(['+I']), pa.array(['wrong type'])], names=['rowkind', 'id']) if failure == 'schema' else _batch([2], ['invalid'])) @@ -180,8 +183,9 @@ def test_iterator_closes_native_input_on_failure_without_replaying_python(tmp_pa assert batches.closed -def test_iterator_spends_limit_before_opening_next_split(tmp_path): - read, splits = _planned_read(tmp_path, limit=3) +@pytest.mark.parametrize('projection', [['id'], {'identifier': 'id'}]) +def test_iterator_spends_limit_before_opening_next_split(tmp_path, projection): + read, splits = _planned_read(tmp_path, limit=3, projection=projection) assert len(splits) >= 2 prepared = [] opened = [] @@ -226,8 +230,9 @@ def test_iterator_qualifies_every_split_before_starting_native(tmp_path): @pytest.mark.parametrize('limit', [0, 1, None]) -def test_iterator_preserves_all_row_kinds_and_hides_transport_column(tmp_path, limit): - read, splits = _planned_read(tmp_path, limit) +@pytest.mark.parametrize('projection', [['id'], {'identifier': 'id'}, {'rowkind': 'id'}]) +def test_iterator_preserves_all_row_kinds_and_hides_transport_column(tmp_path, limit, projection): + read, splits = _planned_read(tmp_path, limit, projection) kinds = ['+I', '-U', '+U', '-D'] batches = _Batches([_batch([1, 2, 3, 4], kinds)]) with patch('pypaimon.read.native_plan._prepare_native_read', return_value=lambda _: batches) as prepare: @@ -240,7 +245,10 @@ def test_iterator_preserves_all_row_kinds_and_hides_transport_column(tmp_path, l @pytest.mark.parametrize('mode', ['append', 'pk']) @pytest.mark.parametrize('projection', [ - ['profile.score', 'id'], ["attrs['a']", 'id'], ['profile', 'profile.score', 'id']]) + ['profile.score', 'id'], ["attrs['a']", 'id'], ['profile', 'profile.score', 'id'], + {'score': 'profile.score'}, {'value': "attrs['a']"}, + {'score': 'profile.score', 'key': 'id', 'value': "attrs['a']", 'copy': 'profile.score'}, + {'v': 'id', 'id': 'v', 'score': 'profile.score'}]) def test_iterator_nested_projection_and_filter_only_inputs(tmp_path, mode, projection): table, _ = _table(tmp_path, mode) predicate = table.new_read_builder().new_predicate_builder().greater_than('v', 10) @@ -255,6 +263,19 @@ def test_iterator_nested_projection_and_filter_only_inputs(tmp_path, mode, proje assert actual and all(len(values) == len(projection) for values, _ in actual) +@pytest.mark.parametrize('mode', ['append', 'pk']) +@pytest.mark.parametrize('first', [None, 42]) +def test_named_iterator_map_key_preserves_first_match(tmp_path, mode, first): + table, schema = _table(tmp_path, mode) + _write(table, [{'id': 4, 'p': 'c', 'attrs': [('a', first), ('a', 99)]}], schema) + builder = table.new_read_builder().with_projection({ + 'value': "attrs['a']", 'copy': "attrs['a']", 'key': 'id', + }).with_filter(table.new_read_builder().new_predicate_builder().equal('id', 4)) + read, splits = builder.new_read(), builder.new_scan().plan().splits() + with patch.object(read, '_create_split_read', side_effect=AssertionError('Python fallback')): + assert _rows(read, splits) == [((first, first, 4), RowKind.INSERT)] + + @pytest.mark.parametrize('mode', ['append', 'de', 'pk']) @pytest.mark.parametrize('projection', [['id', 'id'], ['v', 'id', 'v']]) def test_duplicate_arrow_outputs_align_metadata_by_position(tmp_path, mode, projection): @@ -341,7 +362,7 @@ def test_iterator_pk_blob_respects_payload_mode_and_limit(tmp_path, descriptors) @pytest.mark.parametrize('native_write', [False, True]) -@pytest.mark.parametrize('projection', [['value', 'key'], []]) +@pytest.mark.parametrize('projection', [['value', 'key'], [], {'amount': 'value', 'identifier': 'key'}]) def test_iterator_reads_actual_incremental_changelog_kinds(tmp_path, native_write, projection): from pypaimon.tests.native_changelog_batch_test import _table as changelog_table, _write as write