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 5209d3a1e0 [python] Use native upsert for rows with different supplied
fields (#10414)
5209d3a1e0 is described below
commit 5209d3a1e09b41342dc814acad1d48f4425bbcc6
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 19:37:28 2026 +0800
[python] Use native upsert for rows with different supplied fields (#10414)
---
.../tests/native_upsert_row_shapes_test.py | 184 +++++++++++++++++++++
paimon-python/pypaimon/write/native_update.py | 41 +++--
2 files changed, 212 insertions(+), 13 deletions(-)
diff --git a/paimon-python/pypaimon/tests/native_upsert_row_shapes_test.py
b/paimon-python/pypaimon/tests/native_upsert_row_shapes_test.py
new file mode 100644
index 0000000000..7afbf2222f
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_upsert_row_shapes_test.py
@@ -0,0 +1,184 @@
+# 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.
+
+"""Keep named-row presence through Native upsert matching and appending."""
+
+from contextlib import ExitStack
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.table.row.generic_row import GenericRow
+from pypaimon.write.native_write import native_write_available
+
+
[email protected](params=[False, pytest.param(True, marks=[
+ pytest.mark.native_plan,
+ pytest.mark.skipif(not native_write_available(), reason='native writer
required'),
+])], ids=['python', 'native'])
+def native(request):
+ return request.param
+
+
[email protected](params=[False, True], ids=['batch', 'stream'])
+def stream(request):
+ return request.param
+
+
[email protected]
+def table_factory(tmp_path, native):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+
+ def create(partitioned=False):
+ fields = [('id', pa.int32()), ('value', pa.int32()), ('score',
pa.int32())]
+ if partitioned:
+ fields.append(('p', pa.int32()))
+ schema = pa.schema(fields)
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ schema, partition_keys=['p'] if partitioned else None, options={
+ 'data-evolution.enabled': 'true', 'row-tracking.enabled':
'true',
+ 'deletion-vectors.enabled': 'true',
+ 'write.native.enabled': str(native).lower(),
+ 'scan.native-plan.enabled': str(native).lower(),
+ 'read.native.enabled': str(native).lower(),
+ }), False)
+ return catalog.get_table('default.t'), schema
+
+ return create
+
+
+def _seed(table, schema, rows):
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pylist(rows, schema=schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+
+
+def _upsert(table, native, stream, rows, columns=('value',)):
+ by_name = {field.name: field for field in table.fields}
+ rows = [GenericRow(list(row.values()), [by_name[name] for name in row])
for row in rows]
+ builder = (table.new_stream_write_builder() if stream else
table.new_batch_write_builder())
+ update = builder.new_update().with_update_type(list(columns))
+ with ExitStack() as stack:
+ if native:
+ stack.enter_context(patch(
+
'pypaimon.write.table_upsert_by_key.TableUpsertByKey._upsert_row_partition',
+ side_effect=AssertionError('named-row upsert fell back to
Python')))
+ stack.enter_context(patch(
+
'pypaimon.write.table_upsert_by_key.TableUpsertByKey._build_key_to_row_ids_map',
+ side_effect=AssertionError('upsert keys were matched in
Python')))
+ messages = (update.upsert_by_key(rows, ['id'], 101) if stream
+ else update.upsert_by_key(rows, ['id']))
+ commit = builder.new_commit()
+ try:
+ if stream:
+ commit.commit(messages, 101)
+ else:
+ commit.commit(messages)
+ finally:
+ commit.close()
+ return messages
+
+
+def _rows(table):
+ builder = table.new_read_builder()
+ plan = builder.new_scan().plan()
+ return sorted(builder.new_read().to_arrow(plan.splits()).to_pylist(),
+ key=lambda row: (row['id'], row.get('p', 0)))
+
+
+def _append_shapes(messages):
+ return {tuple(file.write_cols) for message in messages for file in
message.new_files
+ if file.write_cols and 'id' in file.write_cols}
+
+
+def test_upsert_deduplicates_across_shapes_before_validating_columns(
+ native, stream, table_factory):
+ table, schema = table_factory()
+ _seed(table, schema, [
+ {'id': 1, 'value': 10, 'score': 100},
+ {'id': 1, 'value': 11, 'score': 101},
+ {'id': 2, 'value': 20, 'score': 200},
+ ])
+ messages = _upsert(table, native, stream, [
+ {'id': 1, 'score': 999}, {'id': 4, 'value': 40},
+ {'value': 111, 'id': 1}, {'id': 4, 'score': 400},
+ {'id': 2, 'value': 222, 'score': 999},
+ ])
+ assert _rows(table) == [
+ {'id': 1, 'value': 111, 'score': 100},
+ {'id': 1, 'value': 111, 'score': 101},
+ {'id': 2, 'value': 222, 'score': 200},
+ {'id': 4, 'value': None, 'score': 400},
+ ]
+ assert _append_shapes(messages) == {('id', 'score')}
+
+
+def test_upsert_distinguishes_explicit_null_from_an_absent_field(native,
stream, table_factory):
+ table, schema = table_factory()
+ _seed(table, schema, [{'id': 1, 'value': 10, 'score': 100}])
+ messages = _upsert(table, native, stream, [
+ {'id': 1, 'value': None}, {'id': 2, 'score': 200},
+ ])
+ assert _rows(table) == [
+ {'id': 1, 'value': None, 'score': 100},
+ {'id': 2, 'value': None, 'score': 200},
+ ]
+ assert _append_shapes(messages) == {('id', 'score')}
+
+
+def test_upsert_allows_different_append_shapes_in_different_partitions(native,
stream, table_factory):
+ table, schema = table_factory(partitioned=True)
+ _seed(table, schema, [
+ {'id': 1, 'value': 10, 'score': 100, 'p': 1},
+ {'id': 1, 'value': 20, 'score': 200, 'p': 2},
+ ])
+ messages = _upsert(table, native, stream, [
+ {'id': 1, 'value': 111, 'p': 1}, {'id': 4, 'value': 40, 'p': 1},
+ {'id': 4, 'score': 400, 'p': 2}, {'id': 1, 'value': 222, 'score': 999,
'p': 2},
+ {'id': 4, 'value': 41, 'p': 1},
+ ])
+ assert _rows(table) == [
+ {'id': 1, 'value': 111, 'score': 100, 'p': 1},
+ {'id': 1, 'value': 222, 'score': 200, 'p': 2},
+ {'id': 4, 'value': 41, 'score': None, 'p': 1},
+ {'id': 4, 'value': None, 'score': 400, 'p': 2},
+ ]
+ assert _append_shapes(messages) == {('id', 'value', 'p'), ('id', 'score',
'p')}
+
+
[email protected]('input_rows,message', [
+ ([{'id': 1, 'value': 111}, {'id': 2, 'score': 222}], 'value'),
+ ([{'id': 4, 'value': 40}, {'id': 5, 'score': 500}], 'same field set'),
+])
+def test_upsert_rejects_missing_winner_columns_and_inconsistent_appends(
+ native, stream, table_factory, input_rows, message):
+ table, schema = table_factory()
+ original = [{'id': 1, 'value': 10, 'score': 100}, {'id': 2, 'value': 20,
'score': 200}]
+ _seed(table, schema, original)
+ before = table.snapshot_manager().get_latest_snapshot().id
+ with pytest.raises(ValueError, match=message):
+ _upsert(table, native, stream, input_rows)
+ assert table.snapshot_manager().get_latest_snapshot().id == before
+ assert _rows(table) == original
diff --git a/paimon-python/pypaimon/write/native_update.py
b/paimon-python/pypaimon/write/native_update.py
index 96b93109cb..9b19e49b1f 100644
--- a/paimon-python/pypaimon/write/native_update.py
+++ b/paimon-python/pypaimon/write/native_update.py
@@ -136,6 +136,31 @@ def _native_update_columns_supported(table, columns):
and field.name not in inline for field in table.fields)
+def _upsert_row_batches(rows, fields, schema):
+ """Encode consecutive row shapes without padding absent fields with NULL.
+
+ Keep source order, including duplicates across shapes. Key matching,
+ last-write-wins and matched/append column validation belong to Rust core.
+ """
+ batches = []
+ start = 0
+ while start < len(rows):
+ names = set(rows[start])
+ if not names <= set(schema.names):
+ raise ValueError('upsert row fields must be in the table schema')
+ end = start + 1
+ while end < len(rows) and set(rows[end]) == names:
+ end += 1
+ selected = [field for field in fields if field.name in names]
+ batch_schema = pa.schema([schema.field(field.name) for field in
selected])
+ batches.append(pa.RecordBatch.from_pydict({
+ field.name: [value_for_arrow(row[field.name], field) for row in
rows[start:end]]
+ for field in selected
+ }, schema=batch_schema))
+ start = end
+ return batches
+
+
def create_native_upsert(table, commit_user, data, keys, columns):
"""Prepare one core upsert from Arrow columns or named row values."""
if table.options.video_frame_fields() or not
_native_update_columns_supported(table, columns):
@@ -150,20 +175,10 @@ def create_native_upsert(table, commit_user, data, keys,
columns):
if any(not _supported_upsert_key_type(schema.field(key).type) for key in
set(keys + table.partition_keys)):
return None
if not isinstance(data, pa.Table):
- # Keep absent fields absent. Heterogeneous row field sets still need
- # the row-aware updater to distinguish absence from an explicit NULL.
if not columns or not data:
return None
- names = set(data[0])
- if not names <= set(schema.names) or any(set(values) != names for
values in data):
- return None
- fields = [field for field in fields if field.name in names]
- schema = pa.schema([schema.field(field.name) for field in fields])
- data = pa.Table.from_pydict({
- field.name: [value_for_arrow(values[field.name], field) for values
in data]
- for field in fields
- }, schema=schema)
- if (len(data.column_names) != len(set(data.column_names))
+ data = _upsert_row_batches(data, fields, schema)
+ 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)):
@@ -295,7 +310,7 @@ class NativeTableUpdateByRowId(TableUpdateByRowId):
class NativeTableUpsert:
- """Submit full Arrow rows to the core Rust upsert writer."""
+ """Submit Arrow tables or row-shape batches to the core Rust upsert
writer."""
def __init__(self, table, writer, keys, data):
self.table = table
