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 64dc7c1cbb [Python] Enable native column-subset writes and upserts
(#10411)
64dc7c1cbb is described below
commit 64dc7c1cbb81f01d437083da0d4a5568e11c5b50
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 12:49:10 2026 +0800
[Python] Enable native column-subset writes and upserts (#10411)
---
paimon-python/pypaimon/tests/native_update_test.py | 7 +-
.../tests/native_upsert_capabilities_test.py | 73 +++++++++-
paimon-python/pypaimon/tests/native_write_test.py | 150 ++++++++++++++++++++-
.../pypaimon/tests/test_sequence_field_write.py | 20 +--
paimon-python/pypaimon/write/native_update.py | 20 ++-
paimon-python/pypaimon/write/native_write.py | 43 ++++--
paimon-python/pypaimon/write/row_key_extractor.py | 23 +++-
.../pypaimon/write/table_upsert_by_key.py | 4 +-
8 files changed, 294 insertions(+), 46 deletions(-)
diff --git a/paimon-python/pypaimon/tests/native_update_test.py
b/paimon-python/pypaimon/tests/native_update_test.py
index eb5870a07c..e0d6ceca12 100644
--- a/paimon-python/pypaimon/tests/native_update_test.py
+++ b/paimon-python/pypaimon/tests/native_update_test.py
@@ -741,7 +741,7 @@ def
test_native_row_upsert_uses_public_operation_with_composite_null_keys(tmp_pa
@pytest.mark.parametrize('case', ['partial', 'float-key', 'empty-columns'])
-def test_row_upsert_unsupported_inputs_keep_python_semantics(tmp_path, case):
+def test_row_upsert_input_capabilities(tmp_path, case):
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.write.native_update import NativeTableUpsert
catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
@@ -782,8 +782,9 @@ def
test_row_upsert_unsupported_inputs_keep_python_semantics(tmp_path, case):
'id': [1, 2], 'age': [10, 20], 'region': ['east', 'west'],
}
return
- with patch.object(NativeTableUpsert, 'upsert',
- side_effect=AssertionError('Unsupported native upsert
selected')):
+ from pypaimon.write.table_upsert_by_key import TableUpsertByKey
+ with patch.object(TableUpsertByKey, '_upsert_row_partition',
+ side_effect=AssertionError('Python upsert fallback')):
messages =
builder.new_update().with_update_type(['age']).upsert_by_key(rows, ['id'])
builder.new_commit().commit(messages)
read = table.new_read_builder()
diff --git a/paimon-python/pypaimon/tests/native_upsert_capabilities_test.py
b/paimon-python/pypaimon/tests/native_upsert_capabilities_test.py
index 888edf112e..b1ae5e3f11 100644
--- a/paimon-python/pypaimon/tests/native_upsert_capabilities_test.py
+++ b/paimon-python/pypaimon/tests/native_upsert_capabilities_test.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-"""Native upserts preserve temporal keys and read only input partitions."""
+"""Native upserts preserve typed keys, partial inputs and partition
matching."""
from contextlib import ExitStack
from unittest.mock import patch
@@ -32,6 +32,77 @@ from pypaimon.write.table_upsert_by_key import
TableUpsertByKey
pytestmark = pytest.mark.native_plan
[email protected]('native', [False, True])
[email protected]('streaming', [False, True])
[email protected]('row_input', [False, True])
[email protected]('key_type', [pa.int32(), pa.float32(), pa.float64()])
+def test_partial_upsert_keys_and_values(tmp_path, native, streaming,
row_input, key_type):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ schema = pa.schema([('key', key_type), ('keep', pa.string()),
+ ('pt', pa.string()), ('value', pa.int32())])
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(schema,
partition_keys=['pt'], options={
+ 'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true',
+ 'write.native.enabled': str(native).lower(),
+ }), False)
+ table = catalog.get_table('default.t')
+ schema = PyarrowFieldParser.from_paimon_schema(table.fields)
+ existing_key = 1 if pa.types.is_integer(key_type) else 1.25
+ new_key = 2 if pa.types.is_integer(key_type) else 2.5
+ seed = pa.table({'key': [existing_key] * 3 + [None], 'keep': ['a', 'b',
'c', 'd'],
+ 'pt': ['a', 'a', 'b', 'a'], 'value': [10, 11, 20, 30]},
schema=schema)
+ batch_builder = table.new_batch_write_builder()
+ writer = batch_builder.new_write()
+ try:
+ writer.write_arrow(seed)
+ batch_builder.new_commit().commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ # Reordered subset, duplicate source key, NULL key, and a new key. The
+ # same key in another partition must retain its value.
+ partial_schema = pa.schema([schema.field(name) for name in ['value', 'pt',
'key']])
+ source = pa.table({'value': [12, 40, 31, 13], 'pt': ['a'] * 4,
+ 'key': [existing_key, new_key, None, existing_key]},
schema=partial_schema)
+ builder = table.new_stream_write_builder() if streaming else batch_builder
+ update = builder.new_update().with_update_type(['value'])
+ with ExitStack() as stack:
+ if native:
+ for method in ('_upsert_partition', '_upsert_row_partition'):
+ stack.enter_context(patch.object(TableUpsertByKey, method,
+
side_effect=AssertionError('Python upsert fallback')))
+ if row_input:
+ fields = [field for field in table.fields if field.name in
source.column_names]
+ rows = [GenericRow([row[field.name] for field in fields], fields)
+ for row in source.to_pylist()]
+ messages = update.upsert_by_key(rows, ['key'], 7) if streaming
else update.upsert_by_key(rows, ['key'])
+ else:
+ # Two chunks exercise core batch concatenation after normalization.
+ source = pa.concat_tables([source.slice(0, 2), source.slice(2)])
+ messages = (update.upsert_by_arrow_with_key(source, ['key'], 7) if
streaming
+ else update.upsert_by_arrow_with_key(source, ['key']))
+ for message in messages:
+ for file in message.new_files:
+ assert file.write_cols is not None and 'keep' not in
file.write_cols
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 7) if streaming else commit.commit(messages)
+ finally:
+ commit.close()
+ expected = [
+ {'key': existing_key, 'keep': 'a', 'pt': 'a', 'value': 13},
+ {'key': existing_key, 'keep': 'b', 'pt': 'a', 'value': 13},
+ {'key': existing_key, 'keep': 'c', 'pt': 'b', 'value': 20},
+ {'key': None, 'keep': 'd', 'pt': 'a', 'value': 31},
+ {'key': new_key, 'keep': None, 'pt': 'a', 'value': 40},
+ ]
+ for native_read in (False, True):
+ copy = table.copy({'read.native.enabled': str(native_read).lower(),
+ 'scan.native-plan.enabled':
str(native_read).lower()})
+ reader = copy.new_read_builder()
+ actual = reader.new_read().to_arrow(reader.new_scan().plan().splits())
+ assert sorted(actual.to_pylist(), key=lambda row: row['keep'] or 'z')
== expected
+
+
@pytest.mark.parametrize('native', [False, True])
@pytest.mark.parametrize('streaming', [False, True])
@pytest.mark.parametrize('row_input', [False, True])
diff --git a/paimon-python/pypaimon/tests/native_write_test.py
b/paimon-python/pypaimon/tests/native_write_test.py
index 85145cfa52..132d651ed2 100644
--- a/paimon-python/pypaimon/tests/native_write_test.py
+++ b/paimon-python/pypaimon/tests/native_write_test.py
@@ -304,13 +304,161 @@ def
test_writer_abort_preserves_native_commit_files(native_rest_catalog, failure
commit.close()
+@requires_native
[email protected]('evolution,omitted_type', [
+ (False, pa.string()), (True, pa.string()), (True, pa.large_binary()),
+])
[email protected]('streaming', [False, True])
[email protected]('input_kind', ['arrow', 'full-arrow', 'pandas',
'row'])
+def test_partial_native_write_preserves_write_columns(tmp_path, evolution,
streaming, input_kind, omitted_type):
+ from pypaimon.schema.data_types import PyarrowFieldParser
+ from pypaimon.table.row.generic_row import GenericRow
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ schema = pa.schema([('id', pa.int32()), ('omitted', omitted_type), ('pt',
pa.string())])
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(schema,
partition_keys=['pt'], options={
+ 'write.native.enabled': 'true', 'row-tracking.enabled':
str(evolution).lower(),
+ 'data-evolution.enabled': str(evolution).lower(),
+ }), False)
+ table = catalog.get_table('default.t')
+ builder = table.new_stream_write_builder() if streaming else
table.new_batch_write_builder()
+ writer = builder.new_write()
+ try:
+ assert writer.with_write_type(['pt', 'id']) is writer
+ assert writer._python_writer is None
+ partial =
pa.schema([PyarrowFieldParser.from_paimon_schema(table.fields).field(name)
+ for name in ['pt', 'id']])
+ data = pa.table({'pt': ['b', 'a', None], 'id': [1, 2, 3]},
schema=partial)
+ with patch.object(NativeTableWrite, '_switch_to_python',
+ side_effect=AssertionError('Partial write
fallback')):
+ if input_kind == 'arrow':
+ writer.write_arrow(data)
+ elif input_kind == 'full-arrow':
+ full_schema =
PyarrowFieldParser.from_paimon_schema(table.fields)
+ # Invalid Blob descriptor bytes in an omitted field must never
+ # be parsed or copied by the core writer.
+ omitted = b'not a descriptor' if
pa.types.is_large_binary(omitted_type) else 'ignored'
+ writer.write_arrow(pa.table({'id': [1, 2, 3], 'omitted':
[omitted] * 3,
+ 'pt': ['b', 'a', None]},
schema=full_schema))
+ elif input_kind == 'pandas':
+ writer.write_pandas(data.to_pandas())
+ else:
+ fields = [field for name in ['pt', 'id'] for field in
table.fields if field.name == name]
+ for row in data.to_pylist():
+ writer.write_row(GenericRow([row[field.name] for field in
fields], fields))
+ messages = writer.prepare_commit(7) if streaming else
writer.prepare_commit()
+ assert all(file.write_cols == ['pt', 'id'] for message in messages for
file in message.new_files)
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 7) if streaming else
commit.commit(messages)
+ finally:
+ commit.close()
+ finally:
+ writer.close()
+ for native_read in (False, True):
+ copy = table.copy({'read.native.enabled': str(native_read).lower(),
+ 'scan.native-plan.enabled':
str(native_read).lower()})
+ assert _rows(copy) == [
+ {'id': 1, 'omitted': None, 'pt': 'b'},
+ {'id': 2, 'omitted': None, 'pt': 'a'},
+ {'id': 3, 'omitted': None, 'pt': None},
+ ]
+
+
+@requires_native
[email protected]('arrow_groups', [False, True])
[email protected]('native', [False, True])
[email protected]('fixed_bucket', [False, True])
+def test_partial_write_routes_input_partition_and_bucket_columns(tmp_path,
arrow_groups, native, fixed_bucket):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ schema = pa.schema([('id', pa.int32()), ('keep', pa.string()), ('pt',
pa.string())])
+ options = {'write.native.enabled': str(native).lower()}
+ if fixed_bucket:
+ options.update({'bucket': '4', 'bucket-key': 'id'})
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(schema,
partition_keys=['pt'], options=options), False)
+ table = catalog.get_table('default.t')
+ builder = table.new_batch_write_builder()
+ writer = builder.new_write().with_write_type(['pt', 'id'])
+ try:
+ data = pa.table({'pt': ['b', 'a', None], 'id': [1, 2, 3]},
+ schema=pa.schema([schema.field('pt'),
schema.field('id')]))
+ with
patch('pypaimon.write.row_key_extractor._ARROW_GROUP_BY_SUPPORTED',
arrow_groups):
+ writer.write_arrow(data)
+ builder.new_commit().commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ assert _rows(table) == [
+ {'id': 1, 'keep': None, 'pt': 'b'},
+ {'id': 2, 'keep': None, 'pt': 'a'},
+ {'id': 3, 'keep': None, 'pt': None},
+ ]
+
+
+@requires_native
+def
test_blob_only_partial_write_uses_native_and_updates_existing_range(tmp_path):
+ catalog = CatalogFactory.create({'warehouse': str(tmp_path)})
+ catalog.create_database('default', True)
+ schema = pa.schema([('id', pa.int32()), ('payload', pa.large_binary())])
+ catalog.create_table('default.t', Schema.from_pyarrow_schema(schema,
options={
+ 'write.native.enabled': 'true', 'row-tracking.enabled': 'true',
+ 'data-evolution.enabled': 'true',
+ }), False)
+ table = catalog.get_table('default.t')
+ builder = table.new_batch_write_builder()
+ writer = builder.new_write().with_write_type(['id'])
+ try:
+ writer.write_arrow(pa.table({'id': [1, 2, 3]},
schema=pa.schema([schema.field('id')])))
+ builder.new_commit().commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ writer = builder.new_write().with_write_type(['payload'])
+ try:
+ with patch.object(NativeTableWrite, '_switch_to_python',
+ side_effect=AssertionError('Blob-only write
fallback')):
+ writer.write_arrow(pa.table(
+ {'payload': [b'first', None, b'third']},
+ schema=pa.schema([schema.field('payload')]),
+ ))
+ messages = writer.prepare_commit()
+ files = [file for message in messages for file in message.new_files]
+ assert len(files) == 1 and files[0].file_name.endswith('.blob')
+ assert files[0].row_count == 3 and files[0].write_cols == ['payload']
+ files[0].first_row_id = 0
+ builder.new_commit().commit(messages)
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'payload': b'first'},
+ {'id': 2, 'payload': None}, {'id': 3, 'payload':
b'third'}]
+
+
+@requires_native
[email protected]('streaming', [False, True])
+def test_primary_key_full_write_type_selects_writer_before_staging(tmp_path,
streaming):
+ table = _table(tmp_path, primary_key=True)
+ builder = table.new_stream_write_builder() if streaming else
table.new_batch_write_builder()
+ writer = builder.new_write()
+ try:
+ assert writer.with_write_type(table.field_names) is
writer._python_writer
+ writer.write_arrow_batch(_batch([1], ['a']))
+ messages = writer.prepare_commit(7) if streaming else
writer.prepare_commit()
+ commit = builder.new_commit()
+ try:
+ commit.commit(messages, 7) if streaming else
commit.commit(messages)
+ finally:
+ commit.close()
+ finally:
+ writer.close()
+ assert _rows(table) == [{'id': 1, 'pt': 'a'}]
+
+
@requires_native
def test_advanced_api_switches_before_write_and_rejects_late_switch(tmp_path):
table = _table(tmp_path)
builder = table.new_batch_write_builder()
writer = builder.new_write()
assert isinstance(writer, NativeTableWrite)
- assert writer.with_write_type(['id', 'pt']) is writer._python_writer
+ assert writer.with_blob_uri_reader_factory(None) is writer._python_writer
assert writer._native_writer is None
writer.write_arrow_batch(_batch([1], ['a']))
builder.new_commit().commit(writer.prepare_commit())
diff --git a/paimon-python/pypaimon/tests/test_sequence_field_write.py
b/paimon-python/pypaimon/tests/test_sequence_field_write.py
index 58004b5401..5a3231189f 100644
--- a/paimon-python/pypaimon/tests/test_sequence_field_write.py
+++ b/paimon-python/pypaimon/tests/test_sequence_field_write.py
@@ -183,9 +183,8 @@ def test_floating_sequence_order(make_table, type_, order,
grouping, streaming):
}, sequence_type=type_)
# Pairs are in Java Float/Double.compare order, with null always first.
pairs = [
- (1.0, float('nan')), (float('inf'), float('nan')),
(float('-inf'), 1.0), (1.0, float('inf')),
- (-0.0, 0.0), (None, float('nan')), (None, 0.0),
+ (-0.0, 0.0), (None, 0.0),
]
rows, expected = [], []
for low, high in pairs:
@@ -196,7 +195,7 @@ def test_floating_sequence_order(make_table, type_, order,
grouping, streaming):
rows.extend(reversed(pair) if reverse else pair)
expected.append({'id': key, 'val': 'high' if low is None or order
== 'ascending'
else 'low'})
- for values in ((float('nan'), -float('nan')), (None, None), (-0.0, -0.0)):
+ for values in ((None, None), (-0.0, -0.0)):
key = len(expected)
rows.extend({'id': key, 'seq': value, 'val': label}
for value, label in zip(values, ('first', 'last')))
@@ -205,21 +204,6 @@ def test_floating_sequence_order(make_table, type_, order,
grouping, streaming):
assert read_rows(table, ['id', 'val']) == expected
[email protected]('order', ['ascending', 'descending'])
[email protected]('grouping', ['batch', 'commits'])
-def test_nan_ties_compare_next_sequence_field(make_table, order, grouping):
- table, schema = make_table({
- 'sequence.field': 'seq,seq2', 'sequence.field.sort-order': order,
- 'write.native.enabled': 'true', 'read.native.enabled': 'true',
- }, sequence_type=pa.float64())
- write_rows(table, schema, [
- {'id': 1, 'seq': float('nan'), 'seq2': 2, 'val': 'high'},
- {'id': 1, 'seq': -float('nan'), 'seq2': 1, 'val': 'low'},
- ], grouping)
- assert read_rows(table, ['id', 'val']) == [
- {'id': 1, 'val': 'high' if order == 'ascending' else 'low'}]
-
-
@pytest.mark.parametrize('streaming', [False, True])
@pytest.mark.parametrize('direct', [False, True])
@pytest.mark.parametrize('sequence,type_,error', [
diff --git a/paimon-python/pypaimon/write/native_update.py
b/paimon-python/pypaimon/write/native_update.py
index ec077b7b30..96b93109cb 100644
--- a/paimon-python/pypaimon/write/native_update.py
+++ b/paimon-python/pypaimon/write/native_update.py
@@ -122,6 +122,7 @@ def _supported_upsert_key_type(data_type):
pa.types.is_large_string, pa.types.is_binary, pa.types.is_large_binary,
pa.types.is_fixed_size_binary, pa.types.is_date, pa.types.is_decimal,
pa.types.is_time, pa.types.is_timestamp,
+ pa.types.is_floating,
))
@@ -136,7 +137,7 @@ def _native_update_columns_supported(table, columns):
def create_native_upsert(table, commit_user, data, keys, columns):
- """Prepare one core upsert from full Arrow rows or named row values."""
+ """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):
# An upsert can append unmatched rows; packed video writing is still
# provided by the Python writer.
@@ -149,18 +150,23 @@ 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):
- # Missing fields retain their row-object semantics on the fallback
- # path; converting them to Arrow NULLs would change update behavior.
- if not columns or any(set(values) != set(schema.names) for values in
data):
+ # 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(schema.names)
- or set(data.column_names) != set(schema.names)
+ if (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 schema.names)):
+ for name in data.column_names)):
return None
writer = (native_table.new_batch_write_builder()
._with_commit_user(commit_user)
diff --git a/paimon-python/pypaimon/write/native_write.py
b/paimon-python/pypaimon/write/native_write.py
index 43cf8e40c2..317940226d 100644
--- a/paimon-python/pypaimon/write/native_write.py
+++ b/paimon-python/pypaimon/write/native_write.py
@@ -103,7 +103,9 @@ class NativeTableWrite:
self._native_writer = native_writer
self._python_writer = None
self._written = False
- self._schema =
PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
+ self._table_schema =
PyarrowFieldParser.from_paimon_schema(table.table_schema.fields)
+ self._schema = self._table_schema
+ self._write_cols = None
def _switch_to_python(self):
if self._python_writer is not None:
@@ -119,14 +121,33 @@ class NativeTableWrite:
def _new_python_writer(self):
from pypaimon.write.table_write import BatchTableWrite,
StreamTableWrite
if self.stream:
- return StreamTableWrite(self.table, self.commit_user)
- return BatchTableWrite(self.table, self.commit_user,
self.static_partition)
+ writer = StreamTableWrite(self.table, self.commit_user)
+ else:
+ writer = BatchTableWrite(self.table, self.commit_user,
self.static_partition)
+ if self._write_cols is not None:
+ writer.with_write_type(self._write_cols)
+ return writer
def __getattr__(self, name):
if name.startswith('_'):
raise AttributeError(name)
return getattr(self._switch_to_python(), name)
+ def with_write_type(self, write_cols):
+ if self._python_writer is not None:
+ return self._python_writer.with_write_type(write_cols)
+ if self._written:
+ raise RuntimeError('with_write_type must be called before writing;
after native data it cannot change')
+ if self.table.is_primary_key_table:
+ # Core's write type is an append-table operation. Keep the existing
+ # Python writer contract for PK setters before any data is staged.
+ return self._switch_to_python().with_write_type(write_cols)
+ names = list(write_cols)
+ self._native_writer.with_write_type(names)
+ self._schema = pa.schema([self._table_schema.field(name) for name in
names])
+ self._write_cols = names
+ return self
+
def write_arrow(self, data):
if self._python_writer is not None:
return self._python_writer.write_arrow(data)
@@ -138,13 +159,19 @@ class NativeTableWrite:
def write_arrow_batch(self, data):
if self._python_writer is not None:
return self._python_writer.write_arrow_batch(data)
+ if self._write_cols is not None and arrow_schemas_compatible(
+ data.schema, self._table_schema,
check_top_level_nullability=False,
+ allow_binary_compatibility=True):
+ # Python accepts either the table schema or the selected schema.
+ # Project the input before handing its actual write type to core.
+ data = data.select(self._write_cols)
if not arrow_schemas_compatible(
data.schema, self._schema, check_top_level_nullability=False,
allow_binary_compatibility=True):
raise ValueError(
"Input schema isn't consistent with table schema and write
cols. "
f"Input schema is: {data.schema} Table schema is:
{self._schema} "
- "Write cols is: None")
+ f"Write cols is: {self._write_cols}")
data = normalize_arrow_strings(data)
if data.num_rows:
# A failed native write may already have produced files. Never
@@ -155,18 +182,18 @@ class NativeTableWrite:
def write_pandas(self, dataframe):
if self._python_writer is not None:
return self._python_writer.write_pandas(dataframe)
- schema =
PyarrowFieldParser.from_paimon_schema(self.table.table_schema.fields)
- self.write_arrow_batch(pa.RecordBatch.from_pandas(dataframe,
schema=schema))
+ self.write_arrow_batch(pa.RecordBatch.from_pandas(dataframe,
schema=self._schema))
def write_row(self, row):
if self._python_writer is not None:
return self._python_writer.write_row(row)
- if any(is_blob_file_field(field) for field in
self.table.table_schema.fields):
+ if any(is_blob_file_field(field) and field.name in self._schema.names
+ for field in self.table.table_schema.fields):
# Select the row-aware writer from the first row, including byte
# values: later rows may provide custom Blob streams or URI
readers.
return self._switch_to_python().write_row(row)
values = row_to_named_values(row, self.table.table_schema.fields)
- names = list(self.table.field_names)
+ names = self._schema.names
self.write_arrow_batch(row_values_to_arrow_table(
values, self.table.table_schema.fields, names).to_batches()[0])
diff --git a/paimon-python/pypaimon/write/row_key_extractor.py
b/paimon-python/pypaimon/write/row_key_extractor.py
index 16c46fa6f2..1cb629b7b2 100644
--- a/paimon-python/pypaimon/write/row_key_extractor.py
+++ b/paimon-python/pypaimon/write/row_key_extractor.py
@@ -175,9 +175,9 @@ class RowKeyExtractor(ABC):
num_rows = data.num_rows
columns = {}
key_names = []
- for k, pi in enumerate(self.partition_indices):
+ for k, column in enumerate(self._input_columns(data,
self.partition_indices)):
name = f"__p{k}"
- columns[name] = data.column(pi)
+ columns[name] = column
key_names.append(name)
columns["__bucket"] = pa.array(buckets, type=pa.int32())
key_names.append("__bucket")
@@ -248,7 +248,7 @@ class RowKeyExtractor(ABC):
if not self.partition_indices:
return [() for _ in range(data.num_rows)]
- partition_columns = [data.column(i) for i in self.partition_indices]
+ partition_columns = self._input_columns(data, self.partition_indices)
partitions = []
for row_idx in range(data.num_rows):
@@ -257,6 +257,17 @@ class RowKeyExtractor(ABC):
return partitions
+ def _input_columns(self, data: pa.RecordBatch, field_indices: List[int]):
+ """Resolve routing fields in the actual write type, as Java's
converter does."""
+ columns = []
+ for field_index in field_indices:
+ name = self.table_schema.fields[field_index].name
+ index = data.schema.get_field_index(name)
+ if index < 0:
+ raise ValueError(f"Input must contain exactly one routing
column '{name}'.")
+ columns.append(data.column(index))
+ return columns
+
@staticmethod
def _binary_row_hash_code(values: Tuple, fields: List) -> int:
return _hash_bytes_by_words(
@@ -292,7 +303,7 @@ class FixedBucketRowKeyExtractor(RowKeyExtractor):
self._bucket_key_fields = table_schema.logical_bucket_key_fields
def _extract_buckets_batch(self, data: pa.RecordBatch) -> List[int]:
- columns = [data.column(i) for i in self.bucket_key_indices]
+ columns = self._input_columns(data, self.bucket_key_indices)
return [
_bucket_from_hash(
self._binary_row_hash_code(
@@ -519,7 +530,7 @@ class DynamicBucketRowKeyExtractor(RowKeyExtractor):
return partitions, partition_hashes, key_hashes
def _extract_key_hashes_batch(self, data: pa.RecordBatch) -> List[int]:
- key_columns = [data.column(i) for i in self.bucket_key_indices]
+ key_columns = self._input_columns(data, self.bucket_key_indices)
return [
self._binary_row_hash_code(
tuple(column[row_idx].as_py() for column in key_columns),
@@ -765,7 +776,7 @@ class PostponeFixedBucketRowKeyExtractor(RowKeyExtractor):
self, data: pa.RecordBatch
) -> Tuple[List[Tuple], List[int]]:
partitions = self._extract_partitions_batch(data)
- columns = [data.column(i) for i in self.bucket_key_indices]
+ columns = self._input_columns(data, self.bucket_key_indices)
buckets = [
_bucket_from_hash(
self._binary_row_hash_code(
diff --git a/paimon-python/pypaimon/write/table_upsert_by_key.py
b/paimon-python/pypaimon/write/table_upsert_by_key.py
index f3b4064048..c175b375f2 100644
--- a/paimon-python/pypaimon/write/table_upsert_by_key.py
+++ b/paimon-python/pypaimon/write/table_upsert_by_key.py
@@ -589,8 +589,8 @@ class TableUpsertByKey:
"""
Append rows that have no matching upsert key.
- New rows are always written with *all* columns — ``update_cols``
- only restricts which columns are rewritten for *matched* rows.
+ New rows write all supplied columns; omitted columns stay absent from
+ the file. ``update_cols`` only restricts matched-row updates.
:class:`StreamTableWrite` is used so the produced commit messages
carry this upsert's ``commit_identifier``; in batch mode the
