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 08122eef82 [Python] Delegate primary-key scan sharding to Native
(#10413)
08122eef82 is described below
commit 08122eef82fcccf942ec674aac1d788599e4123b
Author: Jingsong Lee <[email protected]>
AuthorDate: Wed Oct 7 17:48:12 2026 +0800
[Python] Delegate primary-key scan sharding to Native (#10413)
---
paimon-python/pypaimon/read/scan_distribution.py | 16 +++-
.../scanner/primary_key_table_split_generator.py | 20 +++--
paimon-python/pypaimon/read/table_scan.py | 19 +++--
.../tests/native_plan_distribution_test.py | 85 ++++++++++++++++++++--
.../pypaimon/tests/native_plan_incremental_test.py | 30 ++++++++
paimon-python/pypaimon/tests/native_plan_test.py | 30 +++++++-
.../pypaimon/tests/scan_file_shard_test.py | 64 ++++++++++++++++
7 files changed, 238 insertions(+), 26 deletions(-)
diff --git a/paimon-python/pypaimon/read/scan_distribution.py
b/paimon-python/pypaimon/read/scan_distribution.py
index dcb2d1b615..87d1815340 100644
--- a/paimon-python/pypaimon/read/scan_distribution.py
+++ b/paimon-python/pypaimon/read/scan_distribution.py
@@ -15,7 +15,7 @@
# specific language governing permissions and limitations
# under the License.
-"""Row-position selection shared by the Python and native planning adapters."""
+"""Worker sharding and row-position selection for Python scan planning."""
from typing import List, Tuple
@@ -43,6 +43,20 @@ def shard_range(total: int, index: int, count: int) ->
Tuple[int, int]:
return start, start + base + int(index < remainder)
+def java_file_name_shard(file_name: str, count: int) -> int:
+ """Java SnapshotReaderImpl: abs(String.hashCode(file_name) % count)."""
+ encoded = file_name.encode('utf-16-be', 'surrogatepass')
+ hash_code = 0
+ for offset in range(0, len(encoded), 2):
+ unit = (encoded[offset] << 8) | encoded[offset + 1]
+ hash_code = (31 * hash_code + unit) & 0xffffffff
+ if hash_code >= 0x80000000:
+ hash_code -= 0x100000000
+ # Python integers do not overflow for Integer.MIN_VALUE. abs(hash) % count
+ # therefore equals Java's abs(hash % count), including negative hashes.
+ return abs(hash_code) % count
+
+
def slice_append_splits(splits: List[Split], start: int, end: int) ->
List[Split]:
"""Select physical rows in [start, end), preserving file order and DVs.
diff --git
a/paimon-python/pypaimon/read/scanner/primary_key_table_split_generator.py
b/paimon-python/pypaimon/read/scanner/primary_key_table_split_generator.py
index 06b1fc3414..01032eba53 100644
--- a/paimon-python/pypaimon/read/scanner/primary_key_table_split_generator.py
+++ b/paimon-python/pypaimon/read/scanner/primary_key_table_split_generator.py
@@ -22,6 +22,7 @@ from pypaimon.common.options.core_options import MergeEngine
from pypaimon.manifest.schema.data_file_meta import DataFileMeta
from pypaimon.manifest.schema.manifest_entry import ManifestEntry
from pypaimon.read.interval_partition import IntervalPartition
+from pypaimon.read.scan_distribution import java_file_name_shard
from pypaimon.read.scanner.split_generator import AbstractSplitGenerator
from pypaimon.read.split import Split
@@ -49,7 +50,7 @@ class PrimaryKeyTableSplitGenerator(AbstractSplitGenerator):
"""Primary key tables do not support slice-based sharding."""
raise NotImplementedError(
"Primary key tables do not support with_slice(). "
- "Use with_shard() for bucket-based parallel processing instead."
+ "Use with_shard() for parallel processing instead."
)
def create_splits(self, file_entries: List[ManifestEntry]) -> List[Split]:
@@ -119,11 +120,14 @@ class
PrimaryKeyTableSplitGenerator(AbstractSplitGenerator):
return splits
def _filter_by_shard(self, file_entries: List[ManifestEntry]) ->
List[ManifestEntry]:
+ """Match Java SnapshotReaderImpl and MergeTreeSplitGenerator.
+
+ Raw-convertible engines distribute file names; other engines keep every
+ version of a key in its bucket. Select entries before packing splits.
"""
- Filter file entries by bucket-based sharding.
- """
- filtered_entries = []
- for entry in file_entries:
- if entry.bucket % self.number_of_para_subtasks ==
self.idx_of_this_subtask:
- filtered_entries.append(entry)
- return filtered_entries
+ by_file_name = (self.deletion_vectors_enabled
+ or self.merge_engine == MergeEngine.FIRST_ROW)
+ return [entry for entry in file_entries if (
+ java_file_name_shard(entry.file.file_name,
self.number_of_para_subtasks)
+ if by_file_name else entry.bucket % self.number_of_para_subtasks
+ ) == self.idx_of_this_subtask]
diff --git a/paimon-python/pypaimon/read/table_scan.py
b/paimon-python/pypaimon/read/table_scan.py
index a2755cf6fb..44a09ee5ac 100755
--- a/paimon-python/pypaimon/read/table_scan.py
+++ b/paimon-python/pypaimon/read/table_scan.py
@@ -186,9 +186,10 @@ class TableScan:
chunk_shuffle = fs.chunk_shuffle
if chunk_shuffle is not None:
extra_options['chunk_shuffle'] = chunk_shuffle
- if fs.idx_of_this_subtask is not None:
- extra_options['shard'] = (
- fs.idx_of_this_subtask, fs.number_of_para_subtasks)
+ if (fs.idx_of_this_subtask is not None
+ and (chunk_shuffle is not None or
self.table.is_primary_key_table)):
+ extra_options['shard'] = (
+ fs.idx_of_this_subtask, fs.number_of_para_subtasks)
if (has_distribution and not self.table.is_primary_key_table
and chunk_shuffle is None):
if fs.idx_of_this_subtask is not None:
@@ -206,10 +207,15 @@ class TableScan:
self.file_scanner.partition_key_predicate,
) if predicate is not None
])
+ # Append positions and Python sorted-index refinement still precede
+ # LIMIT. Ordinary PK shards are selected inside Rust before LIMIT.
+ defer_limit = has_distribution and (
+ not self.table.is_primary_key_table
+ or self.table.options.global_index_enabled())
plan = native_plan(
self.table,
predicate=native_predicate,
- limit=None if has_distribution else self.limit,
+ limit=None if defer_limit else self.limit,
projection=(
[field.name for field in self._read_type]
if self._read_type is not None and chunk_shuffle is None
else None),
@@ -236,10 +242,7 @@ class TableScan:
and plan.snapshot_id is not None):
snapshot =
self.table.snapshot_manager().get_snapshot_by_id(plan.snapshot_id)
splits = fs._apply_primary_key_sorted_indexes(splits, snapshot)
- if chunk_shuffle is None and has_distribution:
- if self.table.is_primary_key_table:
- splits = [s for s in splits
- if s.bucket % fs.number_of_para_subtasks ==
fs.idx_of_this_subtask]
+ if chunk_shuffle is None and defer_limit:
# A partial IndexedSplit plus a file-wide DV cardinality cannot
# reveal how many deleted rows lie inside the selected range.
# Keep every selected split and let the reader enforce LIMIT.
diff --git a/paimon-python/pypaimon/tests/native_plan_distribution_test.py
b/paimon-python/pypaimon/tests/native_plan_distribution_test.py
index cfb0b54dce..bf3b0d29bf 100644
--- a/paimon-python/pypaimon/tests/native_plan_distribution_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_distribution_test.py
@@ -20,6 +20,7 @@
import tempfile
import unittest
from contextlib import ExitStack
+from dataclasses import replace
from unittest.mock import patch
import pyarrow as pa
@@ -32,6 +33,7 @@ from pypaimon.read.native_plan import (
native_version_at_least,
native_runtime_available,
)
+from pypaimon.read.scan_distribution import java_file_name_shard
from pypaimon.table.row.generic_row import GenericRow
from pypaimon.utils.range import Range
from pypaimon.write.commit_message import CommitMessage
@@ -55,16 +57,23 @@ class _DistributionFixture:
primary_keys=primary_keys), False)
return self.catalog.get_table('default.' + name)
- def _write(self, table, rows, schema=None):
+ def _write(self, table, rows, schema=None, level=None):
builder = table.new_batch_write_builder()
writer, commit = builder.new_write(), builder.new_commit()
try:
writer.write_arrow(pa.Table.from_pylist(
rows, schema=self.schema if schema is None else schema))
- commit.commit(writer.prepare_commit())
+ messages = writer.prepare_commit()
+ if level is not None:
+ # Unique-key L1 files model materialized Java compaction
output.
+ for message in messages:
+ message.new_files = [replace(file, level=level)
+ for file in message.new_files]
+ commit.commit(messages)
finally:
writer.close()
commit.close()
+ return messages
@staticmethod
def _de_options(**extra):
@@ -105,9 +114,14 @@ class _DistributionFixture:
scan.with_slice(*slice_)
if row_ranges is not None:
scan.with_row_ranges(row_ranges)
- guard = patch.object(scan.file_scanner, 'scan',
side_effect=AssertionError(
- 'distributed native plan fell back to Python')) if native else
ExitStack()
- with guard:
+ with ExitStack() as stack:
+ if native:
+ stack.enter_context(patch.object(scan.file_scanner, 'scan',
side_effect=AssertionError(
+ 'distributed native plan fell back to Python')))
+ if table.is_primary_key_table and not
table.options.global_index_enabled():
+ stack.enter_context(patch.object(
+ scan.file_scanner, '_apply_push_down_limit',
side_effect=AssertionError(
+ 'PK shard LIMIT was applied in Python')))
plan = scan.plan()
if native:
self.assertTrue(all(
@@ -279,6 +293,67 @@ class NativePlanDistributionTest(_DistributionFixture,
unittest.TestCase):
self.assertCountEqual(covered, list(latest.values()))
self.assertEqual(len({row['k'] for row in covered}), len(covered))
+ def test_raw_primary_key_shards_distribute_files_within_one_bucket(self):
+ for engine, dv in (('first-row', False), ('deduplicate', True)):
+ with self.subTest(engine=engine, dv=dv):
+ table = self._create('raw_pk_' + engine, {
+ 'bucket': '1', 'merge-engine': engine,
+ 'deletion-vectors.enabled': str(dv).lower(),
+ 'source.split.target-size': '1b',
+ }, primary_keys=['k'])
+ rows = [{'k': k, 'v': str(k)} for k in range(12)]
+ files = []
+ for row in rows:
+ files.append(self._write(table, [row],
level=1)[0].new_files[0])
+ expected_rows, snapshot_id = rows, 12
+ if dv:
+ vectors = {}
+ for key in (3, 7):
+ vector = BitmapDeletionVector()
+ vector.delete(0)
+ vectors[files[key].file_name] = vector
+ entry =
TableDeleteByRowId(table)._write_deletion_vector_index(
+ GenericRow([], []), 0, vectors)
+ commit = table.new_batch_write_builder().new_commit()
+ try:
+ commit.commit([CommitMessage(
+ partition=(), bucket=0, new_files=[],
index_adds=[entry])])
+ finally:
+ commit.close()
+ expected_rows = [row for row in rows if row['k'] not in
(3, 7)]
+ snapshot_id = 13
+ _, expected = self._read(table, False)
+ self.assertCountEqual(expected, expected_rows)
+ covered = []
+ owned_files = set()
+ for shard in range(5):
+ normal, expected = self._read(table, False, shard=(shard,
5))
+ plan = self._assert_parity(
+ table, expected, snapshot_id, ordered=False,
shard=(shard, 5))
+ for candidate in (normal, plan):
+ self.assertTrue(all(split.bucket == 0 for split in
candidate.splits()))
+ self.assertTrue(all(
+ java_file_name_shard(file.file_name, 5) == shard
+ for split in candidate.splits() for file in
split.files))
+ files = {file.file_name for split in plan.splits() for
file in split.files}
+ self.assertFalse(owned_files & files)
+ owned_files.update(files)
+ covered.extend(expected)
+ if expected:
+ _, limited = self._read(table, True, shard=(shard, 5),
limit=1)
+ self.assertEqual(len(limited), 1)
+ self.assertIn(limited[0], expected)
+ self.assertCountEqual(covered, expected_rows)
+ self.assertEqual(len(owned_files), 12)
+ # Snapshot selection must happen before shard planning.
+ old = table.copy({'scan.snapshot-id': '3'})
+ old_rows = []
+ for shard in range(5):
+ _, expected = self._read(old, False, shard=(shard, 5))
+ self._assert_parity(old, expected, 3, ordered=False,
shard=(shard, 5))
+ old_rows.extend(expected)
+ self.assertCountEqual(old_rows, rows[:3])
+
@unittest.skipUnless(native_method_available('TableScan',
'with_row_position_slice'),
'pypaimon_rust row-position selection API required')
def test_data_evolution_slice_positions_skip_row_id_gaps(self):
diff --git a/paimon-python/pypaimon/tests/native_plan_incremental_test.py
b/paimon-python/pypaimon/tests/native_plan_incremental_test.py
index 54a32fd67a..1bffc16240 100644
--- a/paimon-python/pypaimon/tests/native_plan_incremental_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_incremental_test.py
@@ -28,6 +28,7 @@ import pytest
from pypaimon import CatalogFactory, Schema
from pypaimon.common.identifier import Identifier
from pypaimon.read.native_plan import native_method_available
+from pypaimon.read.scan_distribution import java_file_name_shard
from pypaimon.read.split import DataSplit
from pypaimon.utils.range import Range
@@ -194,6 +195,35 @@ def
test_primary_key_shards_preserve_all_selected_commits(native, catalog):
(key, version) for key in range(20) for version in ('new', 'old')]
[email protected]('engine,dv', [('first-row', False), ('deduplicate',
True)])
+def test_raw_primary_key_incremental_shards_preserve_file_events(native,
catalog, engine, dv):
+ table = _table(catalog, 'raw_pk_shards', True, {
+ 'bucket': '1', 'merge-engine': engine,
+ 'deletion-vectors.enabled': str(dv).lower(),
+ 'source.split.target-size': '1b',
+ })
+ _write(table, 100, [{'k': 0, 'v': 'outside'}])
+ for key in range(8):
+ _write(table, 200 + key * 2, [{'k': key, 'v': 'old'}])
+ _write(table, 201 + key * 2, [{'k': key, 'v': 'new'}])
+ events, file_names = [], set()
+ for shard in range(5):
+ plan, rows = _read(table, native, (100, 400), shard=(shard, 5))
+ assert plan.snapshot_id == 17
+ assert all(split.is_streaming and split.bucket == 0 for split in
plan.splits())
+ names = {file.file_name for split in plan.splits() for file in
split.files}
+ assert all(java_file_name_shard(name, 5) == shard for name in names)
+ assert not file_names.intersection(names)
+ file_names.update(names)
+ events.extend(rows)
+ if rows:
+ _, limited = _read(table, native, (100, 400), shard=(shard, 5),
limit=1)
+ assert len(limited) == 1 and limited[0] in rows
+ assert len(file_names) == 16
+ assert sorted((row['k'], row['v']) for row in events) == [
+ (key, version) for key in range(8) for version in ('new', 'old')]
+
+
def test_data_evolution_positions_intersect_incremental_row_ranges(native,
catalog):
table = _table(catalog, 'de', options={
'data-evolution.enabled': 'true', 'row-tracking.enabled': 'true',
diff --git a/paimon-python/pypaimon/tests/native_plan_test.py
b/paimon-python/pypaimon/tests/native_plan_test.py
index 7f5bd3386f..3621bfdbae 100644
--- a/paimon-python/pypaimon/tests/native_plan_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_test.py
@@ -1334,18 +1334,40 @@ class NativePlanTest(unittest.TestCase):
native.assert_not_called()
fs.scan.assert_not_called()
- def test_primary_key_shard_defers_limit_until_after_bucket_selection(self):
+ def test_primary_key_shard_and_limit_are_planned_in_rust(self):
fs = Mock(partition_key_predicate=None)
scan = _scan(True, fs)
scan.table.is_primary_key_table = True
+ scan.table.options.global_index_enabled.return_value = False
scan.limit = 1
fs.idx_of_this_subtask, fs.number_of_para_subtasks = 1, 2
- splits = [Mock(bucket=0), Mock(bucket=1)]
+ # A file-name shard may own files from any bucket; never filter the
+ # already-selected Rust splits again using a Python bucket rule.
+ splits = [Mock(bucket=0)]
+ with patch('pypaimon.read.native_plan.native_plan',
return_value=Plan(splits, 3)) as native:
+ self.assertEqual(scan.plan().splits(), splits)
+ self.assertEqual(native.call_args[1]['shard'], (1, 2))
+ self.assertEqual(native.call_args[1]['limit'], 1)
+ fs._apply_push_down_limit.assert_not_called()
+ fs.scan.assert_not_called()
+
+ def
test_primary_key_shard_defers_limit_until_python_index_refinement(self):
+ fs = Mock(partition_key_predicate=None)
+ scan = _scan(True, fs)
+ scan.table.is_primary_key_table = True
+ scan.table.options.global_index_enabled.return_value = True
+ scan.limit = 1
+ fs.idx_of_this_subtask, fs.number_of_para_subtasks = 1, 2
+ scan.predicate = Mock()
+ splits = [Mock(bucket=0), Mock(bucket=0)]
+ fs._apply_primary_key_sorted_indexes.return_value = splits[1:]
fs._apply_push_down_limit.side_effect = lambda selected: selected
with patch('pypaimon.read.native_plan.native_plan',
return_value=Plan(splits, 3)) as native:
- self.assertEqual(scan.plan().splits(), [splits[1]])
+ self.assertEqual(scan.plan().splits(), splits[1:])
+ self.assertEqual(native.call_args[1]['shard'], (1, 2))
self.assertIsNone(native.call_args[1]['limit'])
- fs._apply_push_down_limit.assert_called_once_with([splits[1]])
+ fs._apply_primary_key_sorted_indexes.assert_called_once()
+ fs._apply_push_down_limit.assert_called_once_with(splits[1:])
fs.scan.assert_not_called()
def test_native_plan_requires_split_api(self):
diff --git a/paimon-python/pypaimon/tests/scan_file_shard_test.py
b/paimon-python/pypaimon/tests/scan_file_shard_test.py
new file mode 100644
index 0000000000..471859857a
--- /dev/null
+++ b/paimon-python/pypaimon/tests/scan_file_shard_test.py
@@ -0,0 +1,64 @@
+# 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.
+
+"""Golden Java shard fixtures also guard the non-native planner."""
+
+from types import SimpleNamespace
+
+import pytest
+
+from pypaimon.common.options.core_options import CoreOptions
+from pypaimon.common.options.options import Options
+from pypaimon.read.scan_distribution import java_file_name_shard
+from pypaimon.read.scanner.primary_key_table_split_generator import
PrimaryKeyTableSplitGenerator
+
+
[email protected]('name,count,expected', [
+ ('Aa', 3, 0), ('BB', 3, 0),
+ ('a', 3, 1), ('😀', 3, 1),
+ # Java hashes these strings to Integer.MIN_VALUE and a negative value.
+ ('polygenelubricants', 3, 2), ('polygenelubricants', 7, 2),
+ ('data-0.parquet', 3, 2),
+ ('', 5, 0), ('a', 1, 0),
+])
+def test_file_name_shards_match_java_golden_values(name, count, expected):
+ assert java_file_name_shard(name, count) == expected
+
+
[email protected]('engine,dv,by_file', [
+ ('deduplicate', False, False),
+ ('partial-update', False, False),
+ ('aggregation', False, False),
+ ('first-row', False, True),
+ ('deduplicate', True, True),
+])
+def test_primary_key_shard_rule_matches_java_raw_convertibility(engine, dv,
by_file):
+ table = SimpleNamespace(options=CoreOptions(Options({
+ 'merge-engine': engine, 'deletion-vectors.enabled': str(dv).lower(),
+ })))
+ generator = PrimaryKeyTableSplitGenerator(table, 100, 1)
+ names = ['Aa', 'BB', 'a', '😀', 'polygenelubricants']
+ entries = [SimpleNamespace(bucket=0, file=SimpleNamespace(file_name=name))
for name in names]
+ expected = [['Aa', 'BB'], ['a', '😀'], ['polygenelubricants']] if by_file
else [names, [], []]
+ covered = []
+ for shard in range(3):
+ generator.with_shard(shard, 3)
+ selected = generator._filter_by_shard(entries)
+ assert [entry.file.file_name for entry in selected] == expected[shard]
+ covered.extend(selected)
+ assert len(covered) == len(entries)
+ assert {id(entry) for entry in covered} == {id(entry) for entry in entries}
