[python] Use native upserts for equivalent Arrow input layouts

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(

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):

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())

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':

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'}]

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()),

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()