diff --git a/paimon-python/pypaimon/read/read_dispatch.py b/paimon-python/pypaimon/read/read_dispatch.py index 195c0826ab50..f56c651b7835 100644 --- a/paimon-python/pypaimon/read/read_dispatch.py +++ b/paimon-python/pypaimon/read/read_dispatch.py @@ -35,9 +35,12 @@ def to_iterator(self, splits: List[Split]) -> Iterator: Python split reader, the iterator reuses its row view and visits splits in order. Close the iterator when stopping before exhaustion. """ - if self.output_projection is not None and self.output_projection.named: + from pypaimon.read.variant_read_type import has_variant_extractions + + 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/read/table_read.py b/paimon-python/pypaimon/read/table_read.py index 0108f7204621..38ee402ce583 100644 --- a/paimon-python/pypaimon/read/table_read.py +++ b/paimon-python/pypaimon/read/table_read.py @@ -106,6 +106,9 @@ 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. """ from pypaimon.read.read_dispatch import to_iterator return to_iterator(self, splits) 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..b78f27e71725 --- /dev/null +++ b/paimon-python/pypaimon/tests/named_projection_iterator_test.py @@ -0,0 +1,270 @@ +# 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.schema.arrow_schema 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: + # Core applies LIMIT before returning batches to Python. + yield batch.slice(0, 1) if stop == 'limit' else batch + finally: + closed.append(True) + + def converted(native_read, splits, output_schema, **kwargs): + assert native_read.include_row_kind + assert kwargs == {'parallelism': 1, 'ordered': 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 ca396fda3902..fe484a9e6ed7 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) @@ -178,9 +179,9 @@ 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) splits = builder.new_scan().plan().splits() @@ -192,8 +193,9 @@ def _planned_read(tmp_path, limit=None): @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.execution.native.read.table_read.native_read', return_value=batches) as prepare, \ patch.object(read, '_create_split_read', side_effect=AssertionError('Python fallback')): @@ -210,8 +212,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'])) @@ -225,8 +228,9 @@ def test_iterator_closes_native_input_on_failure_without_replaying_python(tmp_pa assert batches.closed -def test_iterator_delegates_ordered_global_limit_in_one_native_read(tmp_path): - read, splits = _planned_read(tmp_path, limit=3) +@pytest.mark.parametrize('projection', [['id'], {'identifier': 'id'}]) +def test_iterator_delegates_ordered_global_limit_in_one_native_read(tmp_path, projection): + read, splits = _planned_read(tmp_path, limit=3, projection=projection) assert len(splits) >= 2 batches = _Batches([_batch([1, 2]), _batch([3])]) with patch('pypaimon.execution.native.read.table_read.native_read', return_value=batches) as native, \ @@ -264,8 +268,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'] count = 4 if limit is None else limit batches = _Batches([_batch([1, 2, 3, 4][:count], kinds[:count])]) @@ -278,7 +283,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) @@ -293,6 +301,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): @@ -379,7 +400,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