diff --git a/paimon-python/pypaimon/read/split_read.py b/paimon-python/pypaimon/read/split_read.py index 596c661cf715..9563abf54e0a 100644 --- a/paimon-python/pypaimon/read/split_read.py +++ b/paimon-python/pypaimon/read/split_read.py @@ -22,6 +22,7 @@ from pypaimon.common.merge_engine_dispatch import build_merge_function from pypaimon.common.options.core_options import CoreOptions, MergeEngine +from pypaimon.casting.schema_evolution import read_arrow_schema from pypaimon.common.predicate import Predicate from pypaimon.deletionvectors import ( ApplyDeletionVectorReader, @@ -1762,7 +1763,9 @@ def _create_union_reader(self, need_merge_files: List[DataFileMeta], deletion_ve if not field.type.nullable: raise ValueError(f"Field {field} is not null but can't find any file contains it.") - output_schema = PyarrowFieldParser.from_paimon_schema(all_read_fields) + # Bunch readers use runtime nullability after schema evolution. Keep + # the same nested types when assembling their columns into one batch. + output_schema = read_arrow_schema(all_read_fields) return DataEvolutionMergeReader(row_offsets, field_offsets, file_record_readers, schema=output_schema) def _create_file_reader( diff --git a/paimon-python/pypaimon/tests/contiguous_window_plan_test.py b/paimon-python/pypaimon/tests/contiguous_window_plan_test.py index 8dcd6e831557..1867aad67e1a 100644 --- a/paimon-python/pypaimon/tests/contiguous_window_plan_test.py +++ b/paimon-python/pypaimon/tests/contiguous_window_plan_test.py @@ -150,6 +150,7 @@ def authorize(select): def test_independent_payload_column_files_are_read(self): self._check_independent_column_files(native=False) + @pytest.mark.native_plan @unittest.skipUnless(native_reader_available(), "pypaimon-rust native reader API not installed") def test_native_plan_retains_independent_payload_column_files(self): diff --git a/paimon-python/pypaimon/tests/native_hybrid_search_test.py b/paimon-python/pypaimon/tests/native_hybrid_search_test.py index a5bbabfe478b..24b764ef0c65 100644 --- a/paimon-python/pypaimon/tests/native_hybrid_search_test.py +++ b/paimon-python/pypaimon/tests/native_hybrid_search_test.py @@ -202,7 +202,9 @@ def test_one_authoritative_rest_snapshot_for_all_routes_without_retry(rest_catal first = table.snapshot_manager().get_latest_snapshot() indexes._append(table, _rows(8, 4), schema=SCHEMA) if response == 'error': - reply, code = ErrorResponse('TABLE', 'indexes', 'snapshot unavailable', 503), 503 + # Use a non-retryable HTTP status to isolate repeated snapshot lookup + # by the search operation from the client's legitimate 503 retries. + reply, code = ErrorResponse('TABLE', 'indexes', 'snapshot unavailable', 500), 500 else: reply = (GetTableSnapshotResponse(TableSnapshot(first, 1, 0, 7, first.time_millis)) if response == 'first' else GetTableSnapshotResponse()) diff --git a/paimon-python/pypaimon/tests/native_read_snapshot_test.py b/paimon-python/pypaimon/tests/native_read_snapshot_test.py index 74370736d73b..acbc58b9fadf 100644 --- a/paimon-python/pypaimon/tests/native_read_snapshot_test.py +++ b/paimon-python/pypaimon/tests/native_read_snapshot_test.py @@ -193,9 +193,11 @@ def test_rest_snapshot_response_is_authoritative_without_search_retry(rest_catal from pypaimon.api.api_response import ErrorResponse, GetTableSnapshotResponse catalog, server = rest_catalog table = hybrid._table(catalog) + # A 503 is retried by the HTTP client. Use 500 so the request count + # detects repeated search-level resolution rather than transport retries. reply = GetTableSnapshotResponse() if response == 'empty' else ErrorResponse( - 'TABLE', 'indexes', 'snapshot unavailable', 503) - code = 200 if response == 'empty' else 503 + 'TABLE', 'indexes', 'snapshot unavailable', 500) + code = 200 if response == 'empty' else 500 with patch.object(server, '_table_snapshot_handle', return_value=server._mock_response(reply, code)) as latest: with _native_only(): if response == 'error': diff --git a/paimon-python/pypaimon/tests/native_upsert_layout_test.py b/paimon-python/pypaimon/tests/native_upsert_layout_test.py new file mode 100644 index 000000000000..daed72532fb9 --- /dev/null +++ b/paimon-python/pypaimon/tests/native_upsert_layout_test.py @@ -0,0 +1,196 @@ +# 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. + +"""Equivalent Arrow layouts must not send upsert matching back to Python.""" + +from contextlib import ExitStack +from unittest.mock import patch + +import pyarrow as pa +import pytest + +from pypaimon import CatalogFactory, Schema +from pypaimon.schema.data_types import PyarrowFieldParser +from pypaimon.write.table_update import TableUpdate +from pypaimon.write.table_update_by_row_id import TableUpdateByRowId +from pypaimon.write.table_upsert_by_key import TableUpsertByKey + + +pytestmark = pytest.mark.native_plan + +_LAYOUTS = [ + (pa.string(), pa.large_string(), ['old', '', None, '中文']), + (pa.list_(pa.string()), pa.list_(pa.large_string()), + [['old'], [], None, ['中文', None, '']]), + (pa.struct([('text', pa.string()), ('count', pa.int32())]), + pa.struct([('text', pa.large_string()), ('count', pa.int32())]), + [{'text': 'old', 'count': 1}, None, {'text': None, 'count': 2}, {'text': '中文', 'count': 3}]), + (pa.map_(pa.string(), pa.list_(pa.string())), + pa.map_(pa.large_string(), pa.list_(pa.large_string())), + [[('old', ['first'])], [], None, [('中文', [None, ''])]]), +] + + +def _table(catalog, key_type, value_type, native): + catalog.create_database('default', True) + schema = pa.schema([('key', key_type), ('pt', pa.string()), + ('value', value_type), ('keep', pa.int32())]) + catalog.create_table('default.t', Schema.from_pyarrow_schema(schema, partition_keys=['pt'], options={ + 'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true', + 'write.native.enabled': str(native).lower(), + }), False) + table = catalog.get_table('default.t') + return table, PyarrowFieldParser.from_paimon_schema(table.fields) + + +def _commit(builder, messages, stream=False): + commit = builder.new_commit() + try: + commit.commit(messages, 7) if stream else commit.commit(messages) + finally: + commit.close() + + +def _seed(table, schema, keys, values): + builder = table.new_batch_write_builder() + writer = builder.new_write() + try: + writer.write_arrow(pa.table({'key': keys, 'pt': ['a', 'a', 'a', 'b'], + 'value': values, 'keep': [0, 1, 2, 3]}, schema=schema)) + _commit(builder, writer.prepare_commit()) + finally: + writer.close() + + +def _native_only(native): + stack = ExitStack() + if native: + for owner, method in [(TableUpsertByKey, '_upsert_partition'), + (TableUpdateByRowId, '_load_existing_files_info'), + (TableUpdate, '_build_predicate_update_table')]: + stack.enter_context(patch.object(owner, method, side_effect=AssertionError('Python update fallback'))) + return stack + + +def _check(table, values, key, new_key='new'): + # Check both readers, including the row identity of the two updated rows. + for native in (False, True): + view = table.copy({'read.native.enabled': str(native).lower(), + 'scan.native-plan.enabled': str(native).lower()}) + reader = view.new_read_builder().with_projection(['key', 'pt', 'value', 'keep', '_ROW_ID']) + result = reader.new_read().to_arrow(reader.new_scan().plan().splits()).to_pylist() + existing = sorted([row for row in result if row['keep'] is not None], key=lambda row: row['keep']) + assert [row['value'] for row in existing] == [values[3], values[3], values[2], values[0]] + assert [row['key'] for row in existing] == [key, key, None, key] + assert [row['pt'] for row in existing] == ['a', 'a', 'a', 'b'] + assert sorted(row['_ROW_ID'] for row in existing) == [0, 1, 2, 3] + added = [row for row in result if row['keep'] is None] + assert len(added) == 1 + assert added[0]['key'] == new_key + assert added[0]['value'] == values[1] and added[0]['pt'] == 'a' + assert added[0]['_ROW_ID'] == 4 + + +@pytest.mark.parametrize('native', [False, True]) +@pytest.mark.parametrize('stream', [False, True]) +@pytest.mark.parametrize('value_type,input_type,values', _LAYOUTS) +def test_upsert_large_strings_in_keys_partitions_and_nested_values( + tmp_path, native, stream, value_type, input_type, values): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + table, schema = _table(catalog, pa.string(), value_type, native) + _seed(table, schema, ['key', 'key', None, 'key'], [values[0]] * 4) + source = pa.table({ + 'value': pa.array([values[0], values[1], values[2], values[3]], type=input_type), + 'pt': pa.array(['a'] * 4, type=pa.large_string()), + 'key': pa.array(['key', 'new', None, 'key'], type=pa.large_string()), + }) + # Chunks have separate string offsets; the duplicate key spans chunks. + source = pa.concat_tables([source.slice(0, 1), source.slice(1, 2), source.slice(3)]) + builder = table.new_stream_write_builder() if stream else table.new_batch_write_builder() + with _native_only(native): + update = builder.new_update().with_update_type(['value']) + messages = (update.upsert_by_arrow_with_key(source, ['key'], 7) if stream + else update.upsert_by_arrow_with_key(source, ['key'])) + _commit(builder, messages, stream) + _check(table, values, 'key') + + +@pytest.mark.parametrize('native', [False, True]) +def test_upsert_equivalent_binary_keys_on_rest_table(native_rest_catalog, native): + table, schema = _table(native_rest_catalog, pa.binary(), pa.string(), native) + _seed(table, schema, [b'ab', b'ab', None, b'ab'], ['old'] * 4) + source = pa.table({'key': pa.array([b'ab', b'cd', None, b'ab'], type=pa.binary(2)), + 'pt': ['a'] * 4, + 'value': pa.array(['discarded', '', None, '中文'], type=pa.large_string())}) + builder = table.new_batch_write_builder() + with _native_only(native): + messages = builder.new_update().with_update_type(['value']).upsert_by_arrow_with_key(source, ['key']) + _commit(builder, messages) + _check(table, ['old', '', None, '中文'], b'ab', b'cd') + + +@pytest.mark.parametrize('native', [False, True]) +@pytest.mark.parametrize('visible_null', [False, True]) +@pytest.mark.parametrize('operation', ['upsert', 'row_id', 'predicate']) +def test_large_strings_preserve_required_nested_fields(tmp_path, native, visible_null, operation): + catalog = CatalogFactory.create({'warehouse': str(tmp_path)}) + value_type = pa.struct([pa.field('text', pa.string(), nullable=False)]) + table, schema = _table(catalog, pa.string(), value_type, native) + _seed(table, schema, ['key', 'key', None, 'key'], [{'text': 'old'}] * 4) + read = table.new_read_builder().with_projection(['keep', '_ROW_ID']) + row_ids = {row['keep']: row['_ROW_ID'] for row in + read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist()} + source = pa.table({ + 'key': ['key', None], 'pt': ['a', 'a'], + 'value': pa.array([None, {'text': None if visible_null else '中文'}], + type=pa.struct([('text', pa.large_string())])), + }) + files = set(tmp_path.rglob('*.parquet')) + builder = table.new_batch_write_builder() + with _native_only(native): + update = builder.new_update().with_update_type(['value']) + + def prepare(): + if operation == 'upsert': + return update.upsert_by_arrow_with_key(source, ['key']) + if operation == 'row_id': + return update.update_by_arrow_with_row_id( + source.select(['value']).append_column( + '_ROW_ID', pa.array([row_ids[0], row_ids[2]], type=pa.int64()))) + predicate = update.new_predicate_builder().equal('pt', 'a') + + def assignment(matched): + # The two equal keys fan out to both original rows; NULL is + # selected independently, without changing the other partition. + return pa.array([None if row['key'] is not None else source['value'][1].as_py() + for row in matched.to_pylist()], type=pa.struct([ + pa.field('text', pa.large_string(), nullable=False)])) + + return update.update_by_predicate(predicate, {'value': assignment}, read_columns=['key']) + + if visible_null: + with pytest.raises((ValueError, pa.ArrowInvalid)): + prepare() + assert set(tmp_path.rglob('*.parquet')) == files + else: + messages = prepare() + _commit(builder, messages) + reader = table.new_read_builder() + result = reader.new_read().to_arrow(reader.new_scan().plan().splits()).sort_by('keep') + assert result['value'].to_pylist() == [ + None, {'text': 'old'} if operation == 'row_id' else None, + {'text': '中文'}, {'text': 'old'}] diff --git a/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py b/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py index 8f88e872d914..b895302e2049 100644 --- a/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py +++ b/paimon-python/pypaimon/tests/table_update_by_row_id_chunked_test.py @@ -34,6 +34,22 @@ class TableUpdateByRowIdChunkedTest(unittest.TestCase): + def test_write_normalization_accepts_legacy_record_batch_interface(self): + from pypaimon.schema.data_types import AtomicType, DataField + + batch = pa.record_batch([pa.array([1, 2], type=pa.int32())], names=['payload']) + + class LegacyRecordBatch: + # Arrow 6 exposes names through schema, not column_names. + schema = batch.schema + columns = batch.columns + + updater = self._updater() + updater.table = mock.Mock(field_dict={'payload': DataField(0, 'payload', AtomicType('INT'))}) + normalized = updater._normalize_write_data(LegacyRecordBatch()) + self.assertEqual(normalized.schema, batch.schema) + self.assertEqual(normalized.columns[0].to_pylist(), [1, 2]) + def test_streaming_overlay_rejects_sidecar_before_output(self): table = mock.Mock() table.is_primary_key_table = False @@ -194,6 +210,45 @@ def test_streaming_update_rejects_non_contiguous_original_row_ids(self): ValueError, "not contiguous at row ID 0"): list(updater._merged_batches(0, updates, ["payload"])) + def test_streaming_update_closes_source_when_validation_or_writing_fails(self): + for validation_failure in (False, True): + with self.subTest(validation_failure=validation_failure): + closed = [] + batch = pa.record_batch([pa.array([1])], names=['payload']) + + def source(): + try: + yield batch + finally: + closed.append(True) + + def write(batches): + next(batches) + raise IOError('write failed') + + updater = self._updater() + updater.table = mock.sentinel.table + updater._merged_batches = mock.Mock(return_value=source()) + updater._normalize_write_data = mock.Mock( + side_effect=ValueError('validation failed') if validation_failure else None, + return_value=batch) + writer = mock.Mock() + writer.write_batches.side_effect = write + caught = [] + with mock.patch('pypaimon.write.table_update_by_row_id._RowIdUpdateFileWriter', + return_value=writer): + try: + updater._write_group_streaming(mock.Mock(values=[]), 0, mock.sentinel.data, ['payload']) + except (ValueError, IOError) as error: + # Retain the traceback: cleanup must happen immediately, + # without waiting for exception frames to be collected. + caught.append(error) + self.assertEqual(len(caught), 1) + self.assertEqual(str(caught[0]), 'validation failed' if validation_failure else 'write failed') + self.assertEqual(closed, [True]) + writer.abort.assert_called_once_with() + writer.close.assert_called_once_with() + def test_streaming_update_rejects_row_ids_before_original_group(self): updates = pa.table({ SpecialFields.ROW_ID.name: pa.array([9], type=pa.int64()), diff --git a/paimon-python/pypaimon/write/native_update.py b/paimon-python/pypaimon/write/native_update.py index df2abf5c91b1..05285adfe855 100644 --- a/paimon-python/pypaimon/write/native_update.py +++ b/paimon-python/pypaimon/write/native_update.py @@ -21,6 +21,7 @@ import pyarrow as pa +from pypaimon.schema.arrow_schema import arrow_schemas_compatible from pypaimon.schema.data_types import MapType, PyarrowFieldParser, RowType, is_blob_file_field from pypaimon.common.options.core_options import ChangelogProducer, CoreOptions from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER @@ -212,11 +213,17 @@ def create_native_upsert(table, commit_user, data, keys, columns): from pypaimon.write.native_blob_rows import NativeBlobRows blob_rows = NativeBlobRows() data = _upsert_row_batches(data, fields, schema, table, blob_rows) - elif (len(data.column_names) != len(set(data.column_names)) - or not set(data.column_names) <= set(schema.names) - or any(data.schema.field(name).type != schema.field(name).type - for name in data.column_names)): - return None + else: + if (len(data.column_names) != len(set(data.column_names)) + or not set(data.column_names) <= set(schema.names)): + return None + # Core normalizes equivalent Arrow layouts before key encoding and + # writing. Keep value conversions on the existing Python path. + supplied = pa.schema([schema.field(name) for name in data.column_names]) + if not arrow_schemas_compatible( + data.schema, supplied, check_top_level_nullability=False, + allow_binary_compatibility=True, check_nested_nullability=False): + return None writer = (native_table.new_batch_write_builder() ._with_commit_user(commit_user) .new_update() diff --git a/paimon-python/pypaimon/write/table_update_by_row_id.py b/paimon-python/pypaimon/write/table_update_by_row_id.py index 103e4400de7b..ff721625d9fe 100644 --- a/paimon-python/pypaimon/write/table_update_by_row_id.py +++ b/paimon-python/pypaimon/write/table_update_by_row_id.py @@ -28,6 +28,7 @@ from pypaimon.manifest.schema.simple_stats import SimpleStats from pypaimon.read.split import DataSplit from pypaimon.read.table_read import TableRead +from pypaimon.schema.arrow_schema import normalize_arrow_write_nullability from pypaimon.schema.data_types import ( DataField, PyarrowFieldParser, @@ -538,7 +539,8 @@ def _merged_batches(self, first_row_id, data, column_names): def _write_group_streaming(self, partition, first_row_id, data, column_names): writer = _RowIdUpdateFileWriter( self.table, tuple(partition.values), column_names) - batches = self._merged_batches(first_row_id, data, column_names) + merged_batches = self._merged_batches(first_row_id, data, column_names) + batches = (self._normalize_write_data(batch) for batch in merged_batches) try: files = writer.write_batches(batches) self._assign_update_file_metadata(files, first_row_id, column_names, {}) @@ -551,7 +553,10 @@ def _write_group_streaming(self, partition, first_row_id, data, column_names): raise finally: batches.close() - writer.close() + try: + merged_batches.close() + finally: + writer.close() def _merge_update_with_original( self, @@ -661,6 +666,13 @@ def _merge_update_with_original( return merged_table, blob_columns + def _normalize_write_data(self, data): + # DE reads relax nested nullability. Restore and validate the declared + # write types at the writer boundary, including untouched input batches. + write_schema = PyarrowFieldParser.from_paimon_schema( + [self.table.field_dict[name] for name in data.schema.names]) + return normalize_arrow_write_nullability(data, write_schema) + @classmethod def _merge_chunked_column( cls, @@ -924,6 +936,7 @@ def _write_group( success = False try: if merged_data is not None: + merged_data = self._normalize_write_data(merged_data) file_store_write = FileStoreWrite(self.table, self.commit_user) file_store_write.disable_rolling() file_store_write.write_cols = list(merged_data.column_names)