This is an automated email from the ASF dual-hosted git repository.
JingsongLi pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/paimon.git
The following commit(s) were added to refs/heads/master by this push:
new 89a7ec3f7b [python] Use native upserts for equivalent Arrow input
layouts (#10624)
89a7ec3f7b is described below
commit 89a7ec3f7ba2e2c7aaa851f7cf05b5514b24d0c4
Author: Jingsong Lee <[email protected]>
AuthorDate: Sat Oct 17 17:34:21 2026 +0800
[python] Use native upserts for equivalent Arrow input layouts (#10624)
---
paimon-python/pypaimon/read/split_read.py | 5 +-
.../pypaimon/tests/contiguous_window_plan_test.py | 1 +
.../pypaimon/tests/native_hybrid_search_test.py | 4 +-
.../pypaimon/tests/native_read_snapshot_test.py | 6 +-
.../pypaimon/tests/native_upsert_layout_test.py | 196 +++++++++++++++++++++
.../tests/table_update_by_row_id_chunked_test.py | 55 ++++++
paimon-python/pypaimon/write/native_update.py | 17 +-
.../pypaimon/write/table_update_by_row_id.py | 17 +-
8 files changed, 290 insertions(+), 11 deletions(-)
diff --git a/paimon-python/pypaimon/read/split_read.py
b/paimon-python/pypaimon/read/split_read.py
index 596c661cf7..9563abf54e 100644
--- a/paimon-python/pypaimon/read/split_read.py
+++ b/paimon-python/pypaimon/read/split_read.py
@@ -22,6 +22,7 @@ from typing import Callable, Dict, List, Optional, Tuple
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 @@ class DataEvolutionSplitRead(SplitRead):
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 8dcd6e8315..1867aad67e 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 @@ class ContiguousWindowPlanTest(unittest.TestCase):
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 a5bbabfe47..24b764ef0c 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 74370736d7..acbc58b9fa 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 0000000000..daed72532f
--- /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
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
[email protected]('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')
+
+
[email protected]('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')
+
+
[email protected]('native', [False, True])
[email protected]('visible_null', [False, True])
[email protected]('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 8f88e872d9..b895302e20 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 @@ from pypaimon.write.writer.single_file_writer import
SingleFileWriter
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 @@ class TableUpdateByRowIdChunkedTest(unittest.TestCase):
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 df2abf5c91..05285adfe8 100644
--- a/paimon-python/pypaimon/write/native_update.py
+++ b/paimon-python/pypaimon/write/native_update.py
@@ -21,6 +21,7 @@ import json
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 103e4400de..ff721625d9 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.data_file_meta import
DataFileMeta
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 @@ class TableUpdateByRowId:
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 @@ class TableUpdateByRowId:
raise
finally:
batches.close()
- writer.close()
+ try:
+ merged_batches.close()
+ finally:
+ writer.close()
def _merge_update_with_original(
self,
@@ -661,6 +666,13 @@ class TableUpdateByRowId:
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 @@ class TableUpdateByRowId:
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)
