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 847a717d77 [python] Expand Native upserts, Dictionary MERGE and
Parquet vector updates (#10629)
847a717d77 is described below
commit 847a717d77968a1c9bd334c3d329cdbdd69754f7
Author: Jingsong Lee <[email protected]>
AuthorDate: Sat Oct 17 21:19:27 2026 +0800
[python] Expand Native upserts, Dictionary MERGE and Parquet vector updates
(#10629)
---
.../tests/native_merge_dictionary_keys_test.py | 161 ++++++++++++
.../tests/native_parquet_vector_update_test.py | 235 +++++++++++++++++
.../tests/native_upsert_key_normalization_test.py | 232 +++++++++++++++++
.../tests/native_upsert_selected_values_test.py | 283 +++++++++++++++++++++
paimon-python/pypaimon/write/native_merge_into.py | 9 +-
paimon-python/pypaimon/write/native_update.py | 41 ++-
paimon-python/pypaimon/write/table_update.py | 4 +-
7 files changed, 952 insertions(+), 13 deletions(-)
diff --git a/paimon-python/pypaimon/tests/native_merge_dictionary_keys_test.py
b/paimon-python/pypaimon/tests/native_merge_dictionary_keys_test.py
new file mode 100644
index 0000000000..9065db1ae1
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_merge_dictionary_keys_test.py
@@ -0,0 +1,161 @@
+# 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.
+
+"""Native MERGE must join dictionary values, never dictionary indices."""
+
+from contextlib import ExitStack
+from unittest.mock import patch
+
+import pyarrow as pa
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.table.data_evolution_merge_into import WhenMatched,
WhenNotMatched, lit, source_col
+
+
+pytestmark = pytest.mark.native_plan
+_KEYS = [(pa.int32(), 1, 2, 3), (pa.string(), '中文', 'keep', 'new'),
+ (pa.binary(), b'key', b'keep', b'new')]
+
+
+def _dictionary(values, data_type):
+ return pa.array(values, type=data_type).dictionary_encode()
+
+
+def _table(tmp_path, native, key_type, key, other):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('db', True)
+ schema = pa.schema([('key', key_type), ('pt', pa.string()), ('value',
pa.int32()), ('keep', pa.int32())])
+ catalog.create_table('db.t', Schema.from_pyarrow_schema(schema,
partition_keys=['pt'], options={
+ 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true',
+ 'deletion-vectors.enabled': 'true', 'write.native.enabled':
str(native).lower(),
+ }), False)
+ table = catalog.get_table('db.t')
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.table({'key': [key, key, other, None], 'pt':
['a'] * 4,
+ 'value': [0] * 4, 'keep': [0, 1, 2, 3]},
schema=schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ return table
+
+
+def _source(key_type, values):
+ source = pa.table({'match_key': _dictionary([key for key, _ in values],
key_type),
+ 'match_pt': _dictionary(['a'] * len(values),
pa.string()),
+ 'value': pa.array([value for _, value in values],
type=pa.int32())})
+ # Retain source chunks while using the shared dictionaries required by
+ # PyArrow's reference join. Rust's independent dictionaries have core
tests.
+ return pa.concat_tables([source.slice(index, 1) for index in
range(len(values))])
+
+
+def _merge(table, source, native, stream, matched, inserted=()):
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ with ExitStack() as stack:
+ if native:
+
stack.enter_context(patch('pypaimon.table.data_evolution_merge_into._build_tables',
+ side_effect=AssertionError('Python MERGE
join')))
+ update = builder.new_update()
+ kwargs = dict(on={'key': 'match_key', 'pt': 'match_pt'},
when_matched=matched, when_not_matched=inserted)
+ messages = update.merge_into(source, commit_identifier=53, **kwargs)
if stream else update.merge_into(
+ source, **kwargs)
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 53) if stream else commit.commit(messages)
+ finally:
+ commit.close()
+
+
+def _rows(table):
+ for native in (False, True):
+ view = table.copy({'read.native.enabled': str(native).lower(),
+ 'scan.native-plan.enabled': str(native).lower()})
+ builder = view.new_read_builder().with_projection(['key', 'pt',
'value', 'keep', '_ROW_ID'])
+ yield
builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist()
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
[email protected]('key_type,key,other,new', _KEYS)
+def test_dictionary_merge_update_insert_and_null_key_semantics(tmp_path,
native, stream, key_type, key, other, new):
+ table = _table(tmp_path, native, key_type, key, other)
+ before = {row['keep']: row['_ROW_ID'] for row in next(_rows(table))}
+ source = _source(key_type, [(key, 42), (new, 7), (None, 9)])
+ _merge(table, source, native, stream,
+ [WhenMatched.update({'value': source_col('value')},
condition='s.value > t.value')],
+ [WhenNotMatched({'key': source_col('match_key'), 'pt':
source_col('match_pt'),
+ 'value': source_col('value'), 'keep': lit(None)})])
+ for rows in _rows(table):
+ existing = sorted([row for row in rows if row['keep'] is not None],
key=lambda row: row['keep'])
+ assert [row['value'] for row in existing] == [42, 42, 0, 0]
+ assert {row['keep']: row['_ROW_ID'] for row in existing} == before
+ inserted = [row for row in rows if row['keep'] is None]
+ assert len(inserted) == 2
+ assert {(row['key'], row['value'], row['pt']) for row in inserted} ==
{(new, 7, 'a'), (None, 9, 'a')}
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('key_type,key,other,new', _KEYS)
+def test_dictionary_cardinality_uses_values_across_source_chunks(
+ tmp_path, native, key_type, key, other, new):
+ table = _table(tmp_path, native, key_type, key, other)
+ source = _source(key_type, [(key, 1), (new, 2), (key, 3)])
+ before = next(_rows(table))
+ files = set(tmp_path.rglob('*.parquet'))
+ with pytest.raises(ValueError, match='multiple source rows'):
+ _merge(table, source, native, False,
+ [WhenMatched.update({'value': source_col('value')},
condition='s.value = 99')])
+ assert set(tmp_path.rglob('*.parquet')) == files
+ for rows in _rows(table):
+ assert sorted(rows, key=lambda row: row['keep']) == sorted(before,
key=lambda row: row['keep'])
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('key_type,key,other,new', _KEYS)
+def test_dictionary_unconditional_delete_accepts_duplicate_source_matches(
+ tmp_path, native, key_type, key, other, new):
+ table = _table(tmp_path, native, key_type, key, other)
+ source = _source(key_type, [(key, 1), (new, 2), (key, 3)])
+ _merge(table, source, native, False, [WhenMatched.delete()])
+ for rows in _rows(table):
+ assert sorted(row['keep'] for row in rows) == [2, 3]
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
+def test_pandas_categorical_merge_does_not_require_python_join(tmp_path,
native):
+ pd = pytest.importorskip('pandas')
+ table = _table(tmp_path, native, pa.string(), '中文', 'keep')
+ source = pd.DataFrame({'match_key': pd.Categorical(['中文', 'new']),
+ 'match_pt': pd.Categorical(['a', 'a']), 'value':
[42, 7]})
+ _merge(table, source, native, False, [WhenMatched.update({'value':
source_col('value')})])
+ for rows in _rows(table):
+ assert [row['value'] for row in sorted(rows, key=lambda row:
row['keep'])] == [42, 42, 0, 0]
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
+def
test_dictionary_null_values_and_null_indices_do_not_match_in_merge(tmp_path,
native, stream):
+ table = _table(tmp_path, native, pa.string(), 'key', 'other')
+ keys = pa.DictionaryArray.from_arrays(pa.array([0, 1, None],
type=pa.int8()), pa.array([None, 'key']))
+ source = pa.table({'match_key': keys, 'match_pt': ['a'] * 3,
+ 'value': pa.array([9, 42, 11], type=pa.int32())})
+ _merge(table, source, native, stream, [WhenMatched.update({'value':
source_col('value')})])
+ for rows in _rows(table):
+ assert [row['value'] for row in sorted(rows, key=lambda row:
row['keep'])] == [42, 42, 0, 0]
diff --git a/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py
b/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py
new file mode 100644
index 0000000000..90a6c1b4f8
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_parquet_vector_update_test.py
@@ -0,0 +1,235 @@
+# 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.
+
+"""Parquet vector sidecars must not force row-level updates back to Python."""
+
+from contextlib import ExitStack
+from unittest.mock import patch
+
+import pyarrow as pa
+import pyarrow.compute as pc
+import pytest
+
+from pypaimon import CatalogFactory, Schema
+from pypaimon.read.table_read import TableRead
+from pypaimon.schema.data_types import PyarrowFieldParser
+from pypaimon.table.data_evolution_merge_into import WhenMatched,
WhenNotMatched, source_col
+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
+
+
+def _table(tmp_path, native):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('db', True)
+ schema = pa.schema([('id', pa.int32()), ('pt', pa.string()), ('score',
pa.int32()),
+ ('embedding', pa.list_(pa.float32(), 2)), ('payload',
pa.large_binary())])
+ catalog.create_table('db.t', Schema.from_pyarrow_schema(schema,
partition_keys=['pt'], options={
+ 'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true',
'deletion-vectors.enabled': 'true',
+ 'file.format': 'parquet', 'vector.file.format': 'parquet',
+ 'write.native.enabled': str(native).lower(), 'read.batch-size': '1',
+ }), False)
+ table = catalog.get_table('db.t')
+ schema = PyarrowFieldParser.from_paimon_schema(table.fields)
+ rows = [dict(id=i, pt='a' if i <= 2 else 'b', score=i,
+ embedding=[float(i), float(i + 1)], payload=('blob-%d' %
i).encode()) for i in range(1, 5)]
+ 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))
+ messages = writer.prepare_commit()
+ assert any(file.file_name.endswith('.vector.parquet') for message in
messages for file in message.new_files)
+ assert any(file.file_name.endswith('.blob') for message in messages
for file in message.new_files)
+ commit.commit(messages)
+ finally:
+ writer.close()
+ commit.close()
+ return table, schema
+
+
+def _rows(table):
+ for native in (False, True):
+ view = table.copy({'read.native.enabled': str(native).lower(),
+ 'scan.native-plan.enabled': str(native).lower()})
+ builder = view.new_read_builder().with_projection(table.field_names +
['_ROW_ID'])
+ rows =
builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist()
+ yield sorted(rows, key=lambda row: row['id'])
+
+
+def _native_only(native):
+ stack = ExitStack()
+ if native:
+ for owner, method in [(TableUpdateByRowId,
'_load_existing_files_info'),
+ (TableUpsertByKey, '_upsert_partition'),
+ (TableUpdate, '_build_predicate_update_table')]:
+ stack.enter_context(patch.object(owner, method,
side_effect=AssertionError('Python update fallback')))
+
stack.enter_context(patch('pypaimon.table.data_evolution_merge_into._build_tables',
+ side_effect=AssertionError('Python merge
fallback')))
+ return stack
+
+
+def _commit(builder, messages, stream):
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 71) if stream else commit.commit(messages)
+ finally:
+ commit.close()
+
+
+def _replacement(schema, columns, ids):
+ values = {'score': [42] * len(ids), 'embedding': [[9.0, 8.0]] * len(ids),
'payload': [b'new'] * len(ids)}
+ return pa.table({name: pa.array(values[name],
type=schema.field(name).type) for name in columns})
+
+
+def _update(table, schema, native, stream, operation, columns, before):
+ ids = {row['id']: row['_ROW_ID'] for row in before}
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ update = builder.new_update().with_update_type(columns)
+ data = _replacement(schema, columns, [1, 3])
+ with _native_only(native):
+ if operation in ('row_id', 'incremental'):
+ data = data.append_column('_ROW_ID', pa.array([ids[1], ids[3]],
type=pa.int64()))
+ if operation == 'row_id':
+ messages = (update.update_by_arrow_with_row_id(data, 71) if
stream
+ else update.update_by_arrow_with_row_id(data))
+ else:
+ updater = update.new_update_by_row_id(71) if stream else
update.new_update_by_row_id()
+ messages = updater.update_columns(data, columns)
+ elif operation == 'upsert':
+ data = data.append_column('id', pa.array([1, 3],
type=pa.int32())).append_column('pt', ['a', 'b'])
+ messages = (update.upsert_by_arrow_with_key(data, ['id'], 71) if
stream
+ else update.upsert_by_arrow_with_key(data, ['id']))
+ elif operation == 'predicate':
+ predicate = update.new_predicate_builder().is_in('id', [1, 3])
+ values = {'score': 42, 'embedding': [9.0, 8.0], 'payload': b'new'}
+ assignments = {name: values[name] for name in columns}
+ messages = (update.update_by_predicate(predicate, assignments, 71)
if stream
+ else update.update_by_predicate(predicate,
assignments))
+ else:
+ data = data.append_column('id', pa.array([1, 3],
type=pa.int32())).append_column('pt', ['a', 'b'])
+ kwargs = dict(on=['id', 'pt'], when_matched=[WhenMatched.update(
+ {name: source_col(name) for name in columns})])
+ messages = update.merge_into(data, commit_identifier=71, **kwargs)
if stream else update.merge_into(
+ data, **kwargs)
+ _commit(builder, messages, stream)
+ return messages
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
[email protected]('operation', ['row_id', 'incremental', 'upsert',
'predicate', 'merge'])
[email protected]('columns', [['score'], ['embedding'], ['score',
'embedding', 'payload']])
+def test_parquet_vector_and_blob_partial_updates_use_native_core(
+ tmp_path, native, stream, operation, columns):
+ table, schema = _table(tmp_path, native)
+ before = next(_rows(table))
+ messages = _update(table, schema, native, stream, operation, columns,
before)
+ files = [file for message in messages for file in message.new_files]
+ if native:
+ assert any(file.file_name.endswith('.vector.parquet') for file in
files) == ('embedding' in columns)
+ assert any(file.file_name.endswith('.blob') for file in files) ==
('payload' in columns)
+ for file in files:
+ if file.file_name.endswith('.vector.parquet'):
+ assert file.write_cols == ['embedding']
+ assert file.first_row_id is not None
+ for rows in _rows(table):
+ expected = []
+ for original in before:
+ row = dict(original)
+ if row['id'] in (1, 3):
+ for name, value in [('score', 42), ('embedding', [9.0, 8.0]),
('payload', b'new')]:
+ if name in columns:
+ row[name] = value
+ expected.append(row)
+ assert rows == expected
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
[email protected]('native_read', [False, True], ids=['python-read',
'native-read'])
+def
test_shard_normal_updates_preserve_parquet_vector_and_blob_sidecars(tmp_path,
native, stream, native_read):
+ table, schema = _table(tmp_path, native)
+ table = table.copy({'read.native.enabled': str(native_read).lower(),
'read.parallelism': '4'})
+ before = next(_rows(table))
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ update = builder.new_update().with_update_type(['score'])
+ update.with_read_projection(['id', 'score'])
+
+ def reversed_groups(self, readers):
+ # A valid completion order for parallel readers must not change which
+ # row receives an update. Read actual files, but finish the last group
first.
+ for reader in reversed(readers):
+ try:
+ yield from reader
+ finally:
+ reader.close()
+
+ with ExitStack() as stack:
+ stack.enter_context(patch.object(TableRead,
'_native_batches_parallel_streaming', reversed_groups))
+ if native:
+
stack.enter_context(patch('pypaimon.write.table_update.ShardTableUpdator.__init__',
+ side_effect=AssertionError('Python shard
fallback')))
+ shard = update.new_shard_updator(0, 1)
+ reader = shard.arrow_reader()
+ try:
+ for batch in reader:
+ scores = pc.add(batch.column('score'), pa.scalar(10,
type=pa.int32()))
+
shard.update_by_arrow_batch(pa.RecordBatch.from_arrays([scores],
names=['score']))
+ finally:
+ reader.close()
+ messages = shard.prepare_commit()
+ _commit(builder, messages, stream)
+ for rows in _rows(table):
+ assert rows == [dict(row, score=row['score'] + 10) for row in before]
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
+def test_merge_reads_a_parquet_vector_source_in_native_core(tmp_path, native):
+ table, schema = _table(tmp_path, native)
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_table('db.source', Schema.from_pyarrow_schema(schema,
partition_keys=['pt'], options={
+ 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true',
+ 'file.format': 'parquet', 'vector.file.format': 'parquet',
+ }), False)
+ source = catalog.get_table('db.source')
+ seed = source.new_batch_write_builder()
+ writer = seed.new_write()
+ try:
+ writer.write_arrow(pa.Table.from_pylist([
+ dict(id=1, pt='a', score=42, embedding=[9.0, 8.0], payload=b'new'),
+ dict(id=5, pt='b', score=50, embedding=[5.0, 6.0],
payload=b'insert'),
+ ], schema=schema))
+ seed.new_commit().commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ before = next(_rows(table))
+ builder = table.new_batch_write_builder()
+ with _native_only(native):
+ messages = builder.new_update().merge_into(
+ source, on=['id', 'pt'],
+ when_matched=[WhenMatched.update({name: source_col(name) for name
in ('score', 'embedding', 'payload')})],
+ when_not_matched=[WhenNotMatched('*')])
+ _commit(builder, messages, False)
+ for rows in _rows(table):
+ assert len(rows) == 5
+ assert rows[0] == dict(before[0], score=42, embedding=[9.0, 8.0],
payload=b'new')
+ assert rows[1:4] == before[1:]
+ assert {k: v for k, v in rows[4].items() if k != '_ROW_ID'} == dict(
+ id=5, pt='b', score=50, embedding=[5.0, 6.0], payload=b'insert')
diff --git
a/paimon-python/pypaimon/tests/native_upsert_key_normalization_test.py
b/paimon-python/pypaimon/tests/native_upsert_key_normalization_test.py
new file mode 100644
index 0000000000..004c653269
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_upsert_key_normalization_test.py
@@ -0,0 +1,232 @@
+# 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.
+
+"""Logical upsert keys must not depend on Arrow widths or dictionaries."""
+
+import datetime
+from contextlib import ExitStack
+from decimal import Decimal
+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_upsert_by_key import TableUpsertByKey
+
+
+pytestmark = pytest.mark.native_plan
+
+_KEYS = [
+ (pa.int8(), pa.int64(), 1, 2),
+ (pa.int16(), pa.uint32(), 1, 2),
+ (pa.int32(), pa.int64(), 1, 2),
+ (pa.int64(), pa.uint64(), 1, 2),
+ (pa.decimal128(10, 2), pa.decimal128(15, 4), Decimal('1.25'),
Decimal('-2.50')),
+ (pa.date32(), pa.date64(), datetime.date(1969, 12, 31),
datetime.date(1970, 1, 2)),
+ (pa.time32('ms'), pa.time64('us'), datetime.time(1, 2, 3, 400000),
datetime.time(4, 5, 6)),
+ (pa.timestamp('ms'), pa.timestamp('us'), datetime.datetime(1969, 12, 31,
23, 59, 59, 999000),
+ datetime.datetime(1970, 1, 2)),
+ (pa.timestamp('us'), pa.timestamp('ms'), datetime.datetime(1969, 12, 31,
23, 59, 59, 999000),
+ datetime.datetime(1970, 1, 2)),
+ (pa.timestamp('us', 'UTC'), pa.timestamp('ms', 'America/New_York'),
+ datetime.datetime(2024, 7, 1, tzinfo=datetime.timezone.utc),
+ datetime.datetime(2024, 1, 1, tzinfo=datetime.timezone.utc)),
+ (pa.string(), pa.dictionary(pa.int8(), pa.string()), '中文', 'other'),
+ (pa.int32(), pa.dictionary(pa.int16(), pa.int64()), 1, 2),
+ (pa.binary(), pa.dictionary(pa.int8(), pa.binary()), b'key', b'other'),
+]
+
+
+def _array(values, data_type):
+ if pa.types.is_dictionary(data_type):
+ return pa.array(values,
type=data_type.value_type).dictionary_encode().cast(data_type)
+ return pa.array(values, type=data_type)
+
+
+def _table(tmp_path, native, key_type, key, other, partitioned=False):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ schema = pa.schema([('key', key_type), ('pt', pa.string()), ('value',
pa.int32()), ('keep', pa.int32())])
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(
+ schema, partition_keys=['pt'] if partitioned else None, options={
+ 'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true',
+ 'write.native.enabled': str(native).lower(),
+ }), False)
+ table = catalog.get_table('default.t')
+ schema = PyarrowFieldParser.from_paimon_schema(table.fields)
+ builder = table.new_batch_write_builder()
+ writer, commit = builder.new_write(), builder.new_commit()
+ try:
+ writer.write_arrow(pa.table({'key': [key, key, other, None], 'pt':
['a', 'a', 'b', 'a'],
+ 'value': [0] * 4, 'keep': [0, 1, 2, 3]},
schema=schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ return table, schema
+
+
+def _upsert(table, source, native, stream, columns=('value',)):
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ with ExitStack() as stack:
+ if native:
+ stack.enter_context(patch.object(
+ TableUpsertByKey, '_upsert_partition',
side_effect=AssertionError('Python key matching')))
+ update = builder.new_update()
+ if columns is not None:
+ update.with_update_type(list(columns))
+ messages = (update.upsert_by_arrow_with_key(source, ['key'], 37) if
stream
+ else update.upsert_by_arrow_with_key(source, ['key']))
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 37) if stream else commit.commit(messages)
+ finally:
+ commit.close()
+
+
+def _rows(table):
+ for native in (False, True):
+ view = table.copy({'read.native.enabled': str(native).lower(),
+ 'scan.native-plan.enabled': str(native).lower()})
+ builder = view.new_read_builder().with_projection(['key', 'pt',
'value', 'keep', '_ROW_ID'])
+ rows =
builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist()
+ yield sorted(rows, key=lambda row: row['keep'])
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
[email protected]('key_type,input_type,key,other', _KEYS)
+def test_upsert_normalizes_match_keys_without_changing_row_identity(
+ tmp_path, native, stream, key_type, input_type, key, other):
+ table, _ = _table(tmp_path, native, key_type, key, other)
+ before = next(_rows(table))
+ # Different dictionaries and the last duplicate occupy separate chunks.
+ chunks = [pa.table({'key': _array([value], input_type), 'value':
pa.array([number], type=pa.int32())})
+ for value, number in [(key, 1), (None, 9), (key, 42)]]
+ _upsert(table, pa.concat_tables(chunks), native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [42, 42, 0, 9]
+ assert [{k: v for k, v in row.items() if k != 'value'} for row in
rows] == [
+ {k: v for k, v in row.items() if k != 'value'} for row in before]
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
+def
test_partition_key_dictionaries_and_integer_widths_share_one_identity(tmp_path,
native, stream):
+ table, _ = _table(tmp_path, native, pa.int32(), 1, 2, partitioned=True)
+ source = pa.table({'key': pa.array([1, 2, 1], type=pa.int64()),
+ 'pt': _array(['a', 'b', 'a'], pa.dictionary(pa.int8(),
pa.string())),
+ 'value': pa.array([1, 7, 42], type=pa.int32())})
+ before = next(_rows(table))
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [42, 42, 7, 0]
+ assert [row['_ROW_ID'] for row in rows] == [row['_ROW_ID'] for row in
before]
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('key_type,input_type,key,other,new_key', [
+ (pa.int32(), pa.int64(), 1, 2, 3),
+ (pa.decimal128(10, 2), pa.decimal128(15, 4), Decimal('1.00'),
Decimal('2.00'), Decimal('3.00')),
+ (pa.timestamp('ms'), pa.timestamp('us'), datetime.datetime(2024, 1, 1),
+ datetime.datetime(2024, 1, 2), datetime.datetime(2024, 1, 3)),
+])
+def test_new_rows_retain_original_key_write_type_validation(
+ tmp_path, native, key_type, input_type, key, other, new_key):
+ table, _ = _table(tmp_path, native, key_type, key, other)
+ before = next(_rows(table))
+ source = pa.table({'key': _array([new_key], input_type), 'value':
pa.array([42], type=pa.int32())})
+ with pytest.raises((ValueError, TypeError, NotImplementedError)):
+ _upsert(table, source, native, False)
+ for rows in _rows(table):
+ assert rows == before
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('key_type,input_type,key,other,invalid', [
+ (pa.int8(), pa.int64(), 1, 2, 256),
+ (pa.int64(), pa.uint64(), 1, 2, 2 ** 63),
+ (pa.decimal128(10, 2), pa.decimal128(15, 4), Decimal('1.25'),
Decimal('2.50'), Decimal('1.2501')),
+ (pa.timestamp('ms'), pa.timestamp('us'), datetime.datetime(2024, 1, 1),
+ datetime.datetime(2024, 1, 2), datetime.datetime(2024, 1, 1, 0, 0, 0, 1)),
+])
+def test_lossy_key_conversion_cannot_update_an_existing_row(
+ tmp_path, native, key_type, input_type, key, other, invalid):
+ table, _ = _table(tmp_path, native, key_type, key, other)
+ before = next(_rows(table))
+ source = pa.table({'key': _array([invalid], input_type), 'value':
pa.array([42], type=pa.int32())})
+ with pytest.raises((ValueError, TypeError, NotImplementedError)):
+ _upsert(table, source, native, False)
+ for rows in _rows(table):
+ assert rows == before
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
[email protected]('key_type,key,other', [
+ (pa.int32(), 1, 2), (pa.string(), '中文', 'other'), (pa.binary(), b'key',
b'other'),
+])
+def test_dictionary_keys_can_also_be_written_by_full_row_updates(tmp_path,
native, stream, key_type, key, other):
+ table, _ = _table(tmp_path, native, key_type, key, other)
+ before = next(_rows(table))
+ source = pa.table({'key': _dictionary_array([key], key_type), 'pt':
['updated'],
+ 'value': pa.array([42], type=pa.int32()), 'keep':
pa.array([99], type=pa.int32())})
+ _upsert(table, source, native, stream, columns=None)
+ for rows in _rows(table):
+ by_id = {row['_ROW_ID']: row for row in rows}
+ for original in before:
+ expected = dict(original)
+ if original['key'] == key:
+ expected.update(pt='updated', value=42, keep=99)
+ assert by_id[original['_ROW_ID']] == expected
+
+
+def _dictionary_array(values, value_type):
+ return pa.array(values, type=value_type).dictionary_encode()
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
+def test_untyped_null_keys_match_only_existing_null_key_rows(tmp_path, native):
+ table, _ = _table(tmp_path, native, pa.int32(), 1, 2)
+ source = pa.table({'key': pa.nulls(1), 'value': pa.array([9],
type=pa.int32())})
+ _upsert(table, source, native, False)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [0, 0, 0, 9]
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
+def
test_dictionary_null_values_and_null_indices_share_upsert_identity(tmp_path,
native, stream):
+ table, _ = _table(tmp_path, native, pa.string(), 'key', 'other')
+ keys = pa.DictionaryArray.from_arrays(pa.array([0, 1, None],
type=pa.int8()), pa.array([None, 'key']))
+ source = pa.table({'key': keys, 'value': pa.array([9, 42, 11],
type=pa.int32())})
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [42, 42, 0, 11]
+
+
[email protected]('native', [False, True], ids=['python', 'native'])
[email protected]('stream', [False, True], ids=['batch', 'stream'])
+def
test_decimal_zero_and_null_keys_rescale_across_the_full_scale_range(tmp_path,
native, stream):
+ table, _ = _table(tmp_path, native, pa.decimal128(38, 38), Decimal('0'),
Decimal('1e-38'))
+ source = pa.table({'key': pa.array([Decimal('0'), None],
type=pa.decimal128(38, -2)),
+ 'value': pa.array([42, 9], type=pa.int32())})
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [42, 42, 0, 9]
diff --git a/paimon-python/pypaimon/tests/native_upsert_selected_values_test.py
b/paimon-python/pypaimon/tests/native_upsert_selected_values_test.py
new file mode 100644
index 0000000000..9a6f6df9a6
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_upsert_selected_values_test.py
@@ -0,0 +1,283 @@
+# 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.
+
+"""Upsert converts only values selected by matching and last-write-wins."""
+
+import datetime
+from contextlib import ExitStack
+from decimal import Decimal
+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_upsert_by_key import TableUpsertByKey
+
+
+pytestmark = pytest.mark.native_plan
+
+
+def _table(tmp_path, native, value_type=pa.int32()):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ schema = pa.schema([('key', pa.int32()), ('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 _seed(table, schema, value):
+ builder = table.new_batch_write_builder()
+ writer = builder.new_write()
+ commit = builder.new_commit()
+ try:
+ writer.write_arrow(pa.table({'key': [1, 1, 1], 'pt': ['a', 'a', 'b'],
+ 'value': [value] * 3, 'keep': [10, 11,
12]}, schema=schema))
+ commit.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ commit.close()
+ read = table.new_read_builder().with_projection(['keep', '_ROW_ID'])
+ rows =
read.new_read().to_arrow(read.new_scan().plan().splits()).to_pylist()
+ return [row['_ROW_ID'] for row in sorted(rows, key=lambda row:
row['keep'])]
+
+
+def _upsert(table, source, native, stream, columns=('value',)):
+ builder = table.new_stream_write_builder() if stream else
table.new_batch_write_builder()
+ with ExitStack() as stack:
+ if native:
+ stack.enter_context(patch.object(
+ TableUpsertByKey, '_upsert_partition',
+ side_effect=AssertionError('upsert fell back to Python')))
+ update = builder.new_update().with_update_type(list(columns))
+ messages = (update.upsert_by_arrow_with_key(source, ['key'], 17) if
stream
+ else update.upsert_by_arrow_with_key(source, ['key']))
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 17) if stream else commit.commit(messages)
+ finally:
+ commit.close()
+ return messages
+
+
+def _rows(table, native_modes=(False, True)):
+ for native in native_modes:
+ view = table.copy({'read.native.enabled': str(native).lower(),
+ 'scan.native-plan.enabled': str(native).lower()})
+ builder = view.new_read_builder().with_projection(['key', 'pt',
'value', 'keep', '_ROW_ID'])
+ rows =
builder.new_read().to_arrow(builder.new_scan().plan().splits()).to_pylist()
+ yield sorted(rows, key=lambda row: row['keep'])
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
[email protected]('value_type,input_type,initial,value,expected', [
+ (pa.int32(), pa.string(), 0, '42', 42),
+ (pa.int32(), pa.float64(), 0, -42.9, -42),
+ (pa.int64(), pa.int32(), 0, 42, 42),
+ (pa.int32(), pa.null(), 0, None, None),
+ (pa.decimal128(8, 2), pa.string(), Decimal('0.00'), '42.25',
Decimal('42.25')),
+ (pa.list_(pa.int32()), pa.list_(pa.string()), [], ['42', None, '-7'], [42,
None, -7]),
+ (pa.struct([('n', pa.int32())]), pa.struct([('n', pa.string())]),
+ {'n': 0}, {'n': '42'}, {'n': 42}),
+ (pa.binary(), pa.string(), b'old', '中文', '中文'.encode()),
+ (pa.binary(), pa.null(), b'old', None, None),
+ (pa.large_binary(), pa.string(), b'old', 'new', b'new'),
+ (pa.large_binary(), pa.null(), b'old', None, None),
+ (pa.int32(), pa.binary(), 0, b'42', 42),
+ (pa.list_(pa.int32()), pa.list_(pa.binary()), [], [b'42', None], [42,
None]),
+ (pa.struct([('data', pa.binary()), ('n', pa.int32())]),
+ pa.struct([('data', pa.binary()), ('n', pa.string())]),
+ {'data': b'old', 'n': 0}, {'data': b'new', 'n': '42'}, {'data': b'new',
'n': 42}),
+ (pa.struct([('data', pa.binary()), ('n', pa.string())]),
+ pa.struct([('data', pa.binary()), ('n', pa.int32())]),
+ {'data': b'old', 'n': 'old'}, {'data': b'new', 'n': 42}, {'data': b'new',
'n': '42'}),
+ (pa.map_(pa.string(), pa.binary()), pa.struct([('a', pa.string())]),
+ [('a', b'old')], {'a': 'new'}, [('a', b'new')]),
+ (pa.map_(pa.string(), pa.large_binary()), pa.map_(pa.string(),
pa.string()),
+ [('a', b'old')], [('a', 'new')], [('a', b'new')]),
+ (pa.list_(pa.int32()), pa.struct([('dummy', pa.int32())]), [], None, None),
+])
+def test_matched_values_use_row_id_update_casts(
+ tmp_path, native, stream, value_type, input_type, initial, value,
expected):
+ table, schema = _table(tmp_path, native, value_type)
+ row_ids = _seed(table, schema, initial)
+ source = pa.table({'key': pa.array([1], type=pa.int32()), 'pt': ['a'],
+ 'value': pa.array([value], type=input_type)})
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [expected, expected, initial]
+ assert [row['keep'] for row in rows] == [10, 11, 12]
+ assert [row['_ROW_ID'] for row in rows] == row_ids
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
+def test_matched_rows_ignore_unwritten_payload_columns(tmp_path, native,
stream):
+ table, schema = _table(tmp_path, native)
+ _seed(table, schema, 0)
+ source = pa.table({'key': pa.array([1], type=pa.int32()), 'pt': ['a'],
+ 'value': pa.array([42], type=pa.int32()), 'keep': ['not
an integer']})
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [42, 42, 0]
+ assert [row['keep'] for row in rows] == [10, 11, 12]
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
[email protected]('encoded_payload', [False, True], ids=['map',
'dictionary-map'])
[email protected]('target_key,input_key,initial,value,expected', [
+ (pa.int32(), pa.int64(), 1, 2, 2),
+ (pa.string(), pa.dictionary(pa.int8(), pa.string()), 'old', 'new', 'new'),
+ (pa.decimal128(10, 2), pa.decimal128(8, 3), Decimal('1.00'),
Decimal('1.235'), Decimal('1.24')),
+ (pa.time32('ms'), pa.time64('us'), datetime.time(1), datetime.time(1, 2,
3, 123456),
+ datetime.time(1, 2, 3, 123000)),
+])
+def test_matched_blob_map_keys_follow_java_serialization(
+ tmp_path, native, stream, encoded_payload, target_key, input_key,
initial, value, expected):
+ table, schema = _table(tmp_path, native, pa.map_(target_key,
pa.large_binary()))
+ row_ids = _seed(table, schema, [(initial, b'old')])
+ keys = (pa.array([value], type=input_key.value_type).dictionary_encode()
+ if pa.types.is_dictionary(input_key) else pa.array([value],
type=input_key))
+ payload = pa.MapArray.from_arrays(
+ pa.array([0, 1], type=pa.int32()), keys, pa.array([b'new'],
type=pa.large_binary()))
+ if encoded_payload:
+ payload = pa.DictionaryArray.from_arrays(pa.array([0],
type=pa.int8()), payload)
+ source = pa.table({'key': pa.array([1], type=pa.int32()), 'pt': ['a'],
'value': payload})
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [[(expected, b'new')]] * 2 +
[[(initial, b'old')]]
+ assert [row['_ROW_ID'] for row in rows] == row_ids
+ # An unmatched outer key is an ordinary append, whose Arrow types remain
strict.
+ source = source.set_column(0, 'key', pa.array([2], type=pa.int32()))
+ with pytest.raises((ValueError, TypeError, NotImplementedError)):
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['_ROW_ID'] for row in rows] == row_ids
+
+
[email protected]('native', [False, True])
[email protected]('key_type,input_type,initial,values', [
+ (pa.time32('ms'), pa.time64('us'), datetime.time(1),
+ [datetime.time(1, 2, 3, 123001), datetime.time(1, 2, 3, 123999)]),
+ (pa.decimal128(10, 2), pa.decimal128(8, 3), Decimal('1.00'),
[Decimal('1.231'), Decimal('1.234')]),
+])
+def test_blob_map_duplicate_keys_are_rejected_after_java_normalization(
+ tmp_path, native, key_type, input_type, initial, values):
+ table, schema = _table(tmp_path, native, pa.map_(key_type,
pa.large_binary()))
+ row_ids = _seed(table, schema, [(initial, b'old')])
+ source = pa.table({'key': pa.array([1], type=pa.int32()), 'pt': ['a'],
+ 'value': pa.array([[(value, b'new') for value in
values]],
+ type=pa.map_(input_type,
pa.large_binary()))})
+ with pytest.raises(ValueError, match='keys must be unique'):
+ _upsert(table, source, native, False)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [[(initial, b'old')]] * 3
+ assert [row['_ROW_ID'] for row in rows] == row_ids
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
+def test_discarded_duplicate_values_are_not_cast(tmp_path, native, stream):
+ table, schema = _table(tmp_path, native)
+ _seed(table, schema, 0)
+ first = pa.table({'key': pa.array([1], type=pa.int32()), 'pt': ['a'],
'value': ['invalid']})
+ last = pa.table({'key': pa.array([1], type=pa.int32()), 'pt': ['a'],
'value': ['42']})
+ # The winning key is in a separate source chunk.
+ _upsert(table, pa.concat_tables([first, last]), native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [42, 42, 0]
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
+def
test_discarded_duplicates_do_not_validate_required_nested_children(tmp_path,
native, stream):
+ value_type = pa.struct([pa.field('n', pa.int32(), nullable=False),
pa.field('label', pa.string())])
+ table, schema = _table(tmp_path, native, value_type)
+ _seed(table, schema, {'n': 0, 'label': 'old'})
+ source = pa.table({'key': pa.array([1, 1], type=pa.int32()), 'pt': ['a',
'a'],
+ 'value': pa.array([{'n': None, 'label': 'discarded'},
{'n': 42, 'label': 'final'}],
+ type=pa.struct([('n', pa.int32()),
('label', pa.string())]))})
+ _upsert(table, source, native, stream)
+ # Python's format reader cannot restore required ROW child metadata. Read
+ # the files from both update implementations through the native reader.
+ for rows in _rows(table, native_modes=(True,)):
+ assert [row['value'] for row in rows] == [
+ {'n': 42, 'label': 'final'}, {'n': 42, 'label': 'final'}, {'n': 0,
'label': 'old'}]
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
+def test_invalid_surviving_update_value_propagates_without_retry(tmp_path,
native, stream):
+ table, schema = _table(tmp_path, native)
+ _seed(table, schema, 0)
+ source = pa.table({'key': pa.array([1, 1], type=pa.int32()), 'pt': ['a',
'a'],
+ 'value': ['42', 'invalid']})
+ with pytest.raises((ValueError, TypeError)):
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [0, 0, 0]
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
+def test_new_rows_keep_ordinary_write_type_validation(tmp_path, native,
stream):
+ table, schema = _table(tmp_path, native)
+ row_ids = _seed(table, schema, 0)
+ source = pa.table({'key': pa.array([2], type=pa.int32()), 'pt': ['a'],
'value': ['42']})
+ with pytest.raises((ValueError, TypeError, NotImplementedError)):
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [0, 0, 0]
+ assert [row['_ROW_ID'] for row in rows] == row_ids
+
+
[email protected]('native', [False, True])
[email protected]('stream', [False, True])
[email protected]('target_type,source_type,initial,value,expected', [
+ (pa.binary(), pa.large_binary(), b'old', b'\xff\0new', b'\xff\0new'),
+ (pa.large_binary(), pa.binary(), b'old', b'\xff\0new', b'\xff\0new'),
+ (pa.binary(), pa.large_binary(), b'old', None, None),
+ (pa.large_binary(), pa.binary(), b'old', None, None),
+ (pa.map_(pa.string(), pa.binary()),
+ pa.map_(pa.field('k', pa.string(), nullable=False), pa.field('v',
pa.large_binary())),
+ [('a', b'old')], [('a', b'new')], [('a', b'new')]),
+ (pa.map_(pa.string(), pa.binary()), pa.struct([('a', pa.large_binary())]),
+ [('a', b'old')], {'a': b'new'}, [('a', b'new')]),
+ (pa.struct([('a', pa.binary())]), pa.map_(pa.string(), pa.large_binary()),
+ {'a': b'old'}, [('a', b'new')], {'a': b'new'}),
+ (pa.map_(pa.string(), pa.binary()),
+ pa.list_(pa.list_(pa.large_binary())),
+ [('a', b'old')], [[b'a', b'new']], [('a', b'new')]),
+])
+def test_matched_binary_values_use_existing_row_id_casts(
+ tmp_path, native, stream, target_type, source_type, initial, value,
expected):
+ table, schema = _table(tmp_path, native, target_type)
+ row_ids = _seed(table, schema, initial)
+ source = pa.table({'key': pa.array([1], type=pa.int32()), 'pt': ['a'],
+ 'value': pa.array([value], type=source_type)})
+ _upsert(table, source, native, stream)
+ for rows in _rows(table):
+ assert [row['value'] for row in rows] == [expected, expected, initial]
+ assert [row['_ROW_ID'] for row in rows] == row_ids
diff --git a/paimon-python/pypaimon/write/native_merge_into.py
b/paimon-python/pypaimon/write/native_merge_into.py
index 7999c00d72..331804fc6c 100644
--- a/paimon-python/pypaimon/write/native_merge_into.py
+++ b/paimon-python/pypaimon/write/native_merge_into.py
@@ -24,7 +24,8 @@ from pypaimon.common.options.core_options import CoreOptions
from pypaimon.snapshot.snapshot import BATCH_COMMIT_IDENTIFIER
from pypaimon.write.native_commit import create_native_write_table,
from_native_commit_messages
from pypaimon.write.native_update import (
- _native_row_id_table, _native_update_columns_supported,
_supported_upsert_key_type,
+ _dictionary_value_type, _native_row_id_table,
_native_update_columns_supported,
+ _native_update_formats_supported, _supported_upsert_key_type,
)
from pypaimon.write.native_write import _native_partition_types_supported
@@ -49,8 +50,8 @@ def create_native_merge_into(table, source, on, matched,
not_matched, commit_use
from pypaimon.read.native_plan import _native_blob_view_supported
from pypaimon.table.file_store_table import FileStoreTable
if (type(source) is not FileStoreTable
- or source.options.file_format() != 'parquet'
- or source.options.video_frame_fields() or
source.options.with_vector_format()
+ or not _native_update_formats_supported(source)
+ or source.options.video_frame_fields()
or (source.is_primary_key_table and not
source.trimmed_primary_keys)):
return None
from pypaimon.read.table_scan import _validate_scan_mode
@@ -79,7 +80,7 @@ def create_native_merge_into(table, source, on, matched,
not_matched, commit_use
return None
for target, origin in zip(target_keys, source_keys):
target_type = pa.int64() if target == '_ROW_ID' else
context.full_pa_schema.field(target).type
- source_type = target_type if self_merge else
source.schema.field(origin).type
+ source_type = target_type if self_merge else
_dictionary_value_type(source.schema.field(origin).type)
if source_type != target_type or not
_supported_upsert_key_type(target_type):
return None
diff --git a/paimon-python/pypaimon/write/native_update.py
b/paimon-python/pypaimon/write/native_update.py
index 05285adfe8..dbf0b50d17 100644
--- a/paimon-python/pypaimon/write/native_update.py
+++ b/paimon-python/pypaimon/write/native_update.py
@@ -35,12 +35,16 @@ from pypaimon.write.table_update_by_row_id import
TableUpdateByRowId
from pypaimon.write.row_utils import _contains_blob_value, value_for_arrow
+def _native_update_formats_supported(table):
+ return (table.options.file_format() == 'parquet'
+ and (not table.options.with_vector_format() or
table.options.vector_file_format() == 'parquet'))
+
+
def _native_update_table(table):
"""Resolve common native update prerequisites before preparing changes."""
if (type(table) is not FileStoreTable
or not table.options.native_write_enabled()
- or table.options.file_format() != 'parquet'
- or table.options.with_vector_format()
+ or not _native_update_formats_supported(table)
or table.options.changelog_producer() != ChangelogProducer.NONE
or any(table.options.options.contains_key(key) for key in
SCAN_KEYS)
or not native_write_available()):
@@ -154,6 +158,27 @@ def _supported_upsert_key_type(data_type):
))
+def _dictionary_value_type(data_type):
+ while pa.types.is_dictionary(data_type):
+ data_type = data_type.value_type
+ return data_type
+
+
+def _native_upsert_key_type_compatible(source, target):
+ """Qualify layouts; exact values and conversion bounds belong to core."""
+ source = _dictionary_value_type(source)
+ if pa.types.is_null(source):
+ return True
+ for family in (pa.types.is_integer, pa.types.is_decimal, pa.types.is_date,
pa.types.is_time):
+ if family(source) and family(target):
+ return True
+ if pa.types.is_timestamp(source) and pa.types.is_timestamp(target):
+ return (source.tz is None) == (target.tz is None)
+ return arrow_schemas_compatible(
+ pa.schema([pa.field('key', source)]), pa.schema([pa.field('key',
target)]),
+ check_top_level_nullability=False, allow_binary_compatibility=True)
+
+
def _native_update_columns_supported(table, columns):
"""Core writes Parquet columns and Java-compatible sparse Blob deltas."""
if type(table) is not FileStoreTable:
@@ -217,12 +242,12 @@ def create_native_upsert(table, commit_user, data, keys,
columns):
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):
+ # Keys need a shared representation before matching. Core selects
+ # surviving payloads before applying row-ID update or append rules;
+ # validating all payload types here would reject valid matched updates.
+ key_names = list(dict.fromkeys(keys + table.partition_keys))
+ if not
all(_native_upsert_key_type_compatible(data.schema.field(name).type,
schema.field(name).type)
+ for name in key_names):
return None
writer = (native_table.new_batch_write_builder()
._with_commit_user(commit_user)
diff --git a/paimon-python/pypaimon/write/table_update.py
b/paimon-python/pypaimon/write/table_update.py
index 82e5fc7fc1..32a4562afb 100644
--- a/paimon-python/pypaimon/write/table_update.py
+++ b/paimon-python/pypaimon/write/table_update.py
@@ -957,7 +957,9 @@ class ShardTableUpdator:
def arrow_reader(self) -> pyarrow.ipc.RecordBatchReader:
read_builder = self.table.new_read_builder()
read_builder.with_projection(self.projection)
- return read_builder.new_read().to_arrow_batch_reader(self.splits)
+ # Updates consume row ranges in split order. Parallel native readers
+ # return batches in completion order, which would assign values to
other rows.
+ return read_builder.new_read().to_arrow_batch_reader(self.splits,
parallelism=1)
def prepare_commit(self) -> List[CommitMessage]:
commit_messages = []
