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 da07686c68 [python] Preserve declared primary-key order in merge reads
(#9801)
da07686c68 is described below
commit da07686c68791ad2cedc4f2cb9f721240d6c8942
Author: zhigang <[email protected]>
AuthorDate: Tue Sep 15 11:15:34 2026 +0800
[python] Preserve declared primary-key order in merge reads (#9801)
---
paimon-python/pypaimon/read/split_read.py | 13 +++----
.../pypaimon/tests/reader_primary_key_test.py | 43 ++++++++++++++++++++++
.../tests/schema_evolution_pk_read_test.py | 27 ++++++++++++++
3 files changed, 76 insertions(+), 7 deletions(-)
diff --git a/paimon-python/pypaimon/read/split_read.py
b/paimon-python/pypaimon/read/split_read.py
index c5300a620f..9073072b48 100644
--- a/paimon-python/pypaimon/read/split_read.py
+++ b/paimon-python/pypaimon/read/split_read.py
@@ -632,15 +632,14 @@ class SplitRead(ABC):
self._all_data_fields_from(file_schema.fields))
def _create_key_value_fields(self, value_field: List[DataField]):
- all_fields: List[DataField] = self.table.fields
all_data_fields = []
- for field in all_fields:
- if field.name in self.trimmed_primary_key:
- key_field_name = f"{KEY_PREFIX}{field.name}"
- key_field_id = field.id + KEY_FIELD_ID_START
- key_field = DataField(key_field_id, key_field_name, field.type)
- all_data_fields.append(key_field)
+ # Merge keys must follow the same declared order as the sorted files.
+ for field in self.table.trimmed_primary_keys_fields:
+ key_field_name = f"{KEY_PREFIX}{field.name}"
+ key_field_id = field.id + KEY_FIELD_ID_START
+ key_field = DataField(key_field_id, key_field_name, field.type)
+ all_data_fields.append(key_field)
all_data_fields.append(SpecialFields.SEQUENCE_NUMBER)
all_data_fields.append(SpecialFields.VALUE_KIND)
diff --git a/paimon-python/pypaimon/tests/reader_primary_key_test.py
b/paimon-python/pypaimon/tests/reader_primary_key_test.py
index 1ed2148099..19a4f1558c 100644
--- a/paimon-python/pypaimon/tests/reader_primary_key_test.py
+++ b/paimon-python/pypaimon/tests/reader_primary_key_test.py
@@ -97,6 +97,49 @@ class PkReaderTest(unittest.TestCase):
value_kind_field_found,
"_VALUE_KIND field should exist in the written parquet file")
+ def test_composite_primary_key_order(self):
+ cases = [
+ (['customer_id', 'order_id'], []),
+ (['order_id', 'customer_id'], []),
+ (['order_id', 'dt', 'customer_id'], ['dt']),
+ ]
+ for index, (primary_keys, partition_keys) in enumerate(cases):
+ with self.subTest(primary_keys=primary_keys):
+ fields = [('customer_id', pa.int64()), ('order_id',
pa.int64()),
+ ('amount', pa.int64())]
+ base = {'customer_id': [1, 2, 1, 2], 'order_id': [10, 10, 20,
20],
+ 'amount': [100, 200, 300, 400]}
+ update = {'customer_id': [1], 'order_id': [20], 'amount':
[350]}
+ expected = {'customer_id': [1, 1, 2, 2], 'order_id': [10, 20,
10, 20],
+ 'amount': [100, 350, 200, 400]}
+ if partition_keys:
+ fields.append(('dt', pa.string()))
+ base['dt'] = expected['dt'] = ['p1'] * 4
+ update['dt'] = ['p1']
+ arrow_schema = pa.schema(fields)
+ name = 'default.test_composite_pk_order_{}'.format(index)
+ self.catalog.create_table(name, Schema.from_pyarrow_schema(
+ arrow_schema, primary_keys=primary_keys,
partition_keys=partition_keys,
+ options={'bucket': '1'}), False)
+ table = self.catalog.get_table(name)
+ for data in (base, update):
+ builder = table.new_batch_write_builder()
+ writer, committer = builder.new_write(),
builder.new_commit()
+ try:
+ writer.write_arrow(pa.Table.from_pydict(data,
schema=arrow_schema))
+ committer.commit(writer.prepare_commit())
+ finally:
+ writer.close()
+ committer.close()
+
+ actual =
self._read_test_table(table.new_read_builder()).sort_by(
+ [('customer_id', 'ascending'), ('order_id', 'ascending')])
+ self.assertEqual(actual.to_pylist(),
+ pa.Table.from_pydict(expected,
schema=arrow_schema).to_pylist())
+ projected = self._read_test_table(
+ table.new_read_builder().with_projection(['amount']))
+ self.assertEqual(sorted(projected['amount'].to_pylist()),
[100, 200, 350, 400])
+
def test_pk_orc_reader(self):
schema = Schema.from_pyarrow_schema(self.pa_schema,
partition_keys=['dt'],
diff --git a/paimon-python/pypaimon/tests/schema_evolution_pk_read_test.py
b/paimon-python/pypaimon/tests/schema_evolution_pk_read_test.py
index 4df02bb33c..6fb0afd21d 100644
--- a/paimon-python/pypaimon/tests/schema_evolution_pk_read_test.py
+++ b/paimon-python/pypaimon/tests/schema_evolution_pk_read_test.py
@@ -197,6 +197,33 @@ class SchemaEvolutionPkReadTest(unittest.TestCase):
self.assertEqual(self._read_sorted(table), [
{'b': 'b2', 'id': 1, 'a': 'a2'}])
+ def test_composite_pk_column_position(self):
+ s0 = pa.schema([('customer_id', pa.int64()), ('order_id', pa.int64()),
+ ('amount', pa.int64())])
+ table = self._create('composite_pk_position', s0,
+ primary_keys=('customer_id', 'order_id'))
+ self._write(table, pa.Table.from_pydict({
+ 'customer_id': [1, 1, 2, 2], 'order_id': [10, 20, 10, 20],
+ 'amount': [100, 200, 300, 400]}, schema=s0))
+
+ self.catalog.alter_table(
+ 'default.composite_pk_position',
+ [SchemaChange.update_column_position(Move.first('order_id'))],
False)
+ table = self.catalog.get_table('default.composite_pk_position')
+ s1 = pa.schema([('order_id', pa.int64()), ('customer_id', pa.int64()),
+ ('amount', pa.int64())])
+ self._write(table, pa.Table.from_pydict({
+ 'order_id': [10], 'customer_id': [2], 'amount': [350]}, schema=s1))
+
+ builder = table.new_read_builder()
+ actual =
builder.new_read().to_arrow(builder.new_scan().plan().splits())
+ self.assertEqual(actual.sort_by([
+ ('customer_id', 'ascending'), ('order_id',
'ascending')]).to_pylist(), [
+ {'customer_id': 1, 'order_id': 10, 'amount': 100},
+ {'customer_id': 1, 'order_id': 20, 'amount': 200},
+ {'customer_id': 2, 'order_id': 10, 'amount': 350},
+ {'customer_id': 2, 'order_id': 20, 'amount': 400}])
+
# -- B7: first-row engine + add column -------------------------------
def test_pk_first_row_add_column(self):
