From dc14e3c78fc6771465667d1358a09b6ccbcd286b Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 15 Oct 2026 19:36:52 +0800 Subject: [PATCH 1/3] [python] Use native scans for supplied primary-key index results --- paimon-python/pypaimon/read/native_plan.py | 38 ++- .../pypaimon/read/scanner/file_scanner.py | 13 +- paimon-python/pypaimon/read/table_scan.py | 32 +-- .../table/source/primary_key_scored_result.py | 14 + .../tests/native_index_result_scan_test.py | 266 ++++++++++++++++++ .../tests/native_plan_expanded_test.py | 6 +- .../pypaimon/tests/native_plan_test.py | 16 +- 7 files changed, 349 insertions(+), 36 deletions(-) create mode 100644 paimon-python/pypaimon/tests/native_index_result_scan_test.py diff --git a/paimon-python/pypaimon/read/native_plan.py b/paimon-python/pypaimon/read/native_plan.py index 0363fd806400..44bb03d6d1d7 100644 --- a/paimon-python/pypaimon/read/native_plan.py +++ b/paimon-python/pypaimon/read/native_plan.py @@ -452,7 +452,8 @@ def native_plan( row_position_shard: Optional[Tuple[int, int]] = None, chunk_shuffle: Optional[Tuple[int, int]] = None, shard: Optional[Tuple[int, int]] = None, - partition_predicate: Optional[Predicate] = None) -> Plan: + partition_predicate: Optional[Predicate] = None, + global_index_result=None) -> Plan: """Plan with pypaimon_rust, preserving snapshot metadata. Native conversion or planning failures are handled by TableScan, which @@ -469,15 +470,40 @@ def native_plan( builder = builder.with_partition_filter(_predicate_to_native(partition_predicate)) if row_ranges is not None: builder = builder.with_row_ranges(row_ranges) - if incremental_range is None: + from pypaimon.table.source.global_index_split_result import GlobalIndexSplitResult + supplied_pk_splits = (isinstance(global_index_result, GlobalIndexSplitResult) + and not table.options.query_auth_enabled) + if supplied_pk_splits or incremental_range is None: scan = builder.new_scan() elif incremental_mode == 'delta': - # Keep the two-argument call compatible with runtimes predating the - # explicit mode API. Non-delta modes require the new binding. scan = builder.new_incremental_scan(*incremental_range) else: - scan = builder.new_incremental_scan( - *incremental_range, incremental_mode) + scan = builder.new_incremental_scan(*incremental_range, incremental_mode) + if global_index_result is not None and not table.options.query_auth_enabled: + # A PK result already selected a snapshot and physical files. Java's + # PrimaryKeyBatchScan returns those selections before ordinary planning. + from pypaimon_rust.datafusion import GlobalIndexResult, Split as NativeSplit + from pypaimon.globalindex.vector_search_result import ScoredGlobalIndexResult + from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult + if isinstance(global_index_result, GlobalIndexSplitResult): + result = (getattr(global_index_result, '_native_result', None) + if type(global_index_result) is PrimaryKeyScoredResult else None) + # Java's PK result uses snapshot 0 for an uncommitted table, + # whereas the core search uses None. Keep the Python result's ID. + if result is not None and result.snapshot_id() is None: + result = None + if result is None: + result = GlobalIndexResult.from_splits( + global_index_result.snapshot_id, + [NativeSplit.deserialize(serialize_split_v1(split)) for split in global_index_result.splits]) + elif isinstance(global_index_result, ScoredGlobalIndexResult): + score = global_index_result.score_getter() + result = GlobalIndexResult.from_scored_row_ids({ + row_id: score(row_id) for row_id in global_index_result.results()}) + else: + result = GlobalIndexResult.from_row_ranges([ + (range_.from_, range_.to) for range_ in global_index_result.results().to_range_list()]) + scan = scan.with_global_index_result(result) if row_position_slice is not None: scan = scan.with_row_position_slice(*row_position_slice) if row_position_shard is not None: diff --git a/paimon-python/pypaimon/read/scanner/file_scanner.py b/paimon-python/pypaimon/read/scanner/file_scanner.py index b938a7f7d262..85c8061975e9 100755 --- a/paimon-python/pypaimon/read/scanner/file_scanner.py +++ b/paimon-python/pypaimon/read/scanner/file_scanner.py @@ -334,10 +334,13 @@ def scan(self) -> Plan: raise ValueError( "Primary-key scan requires a GlobalIndexSplitResult, but found %s." % type(self._global_index_result).__name__) - return Plan( - list(self._global_index_result.splits), - snapshot_id=self._global_index_result.snapshot_id, - ) + # Java ignores external raw-index results under query auth: + # masking can invalidate both their positions and their scores. + if not self.table.options.query_auth_enabled: + return Plan( + list(self._global_index_result.splits), + snapshot_id=self._global_index_result.snapshot_id, + ) # Create appropriate split generator based on table type if self.chunk_shuffle is not None: self._validate_chunk_shuffle_compat() @@ -471,7 +474,7 @@ def _create_data_evolution_split_generator(self): self._scanned_snapshot = snapshot self._scanned_snapshot_id = snapshot.id if snapshot else None - if row_ranges is None: + if row_ranges is None and not self.table.options.query_auth_enabled: global_index_plan = self._global_index_result \ if self._global_index_result is not None \ else self._eval_global_index(snapshot) diff --git a/paimon-python/pypaimon/read/table_scan.py b/paimon-python/pypaimon/read/table_scan.py index 40f0383d30fc..687d5c8690db 100755 --- a/paimon-python/pypaimon/read/table_scan.py +++ b/paimon-python/pypaimon/read/table_scan.py @@ -104,7 +104,6 @@ def _native_plan_supported(self) -> bool: def _native_plan_supported_impl(self) -> bool: """Fall back to the Python scanner for scans native can't carry: - primary-key global-index results, a primary-key table whose trimmed PK is empty (PK equals the partition key; Rust rejects this schema), unsupported time travel selectors, query auth, a missing pypaimon-rust, or a catalog / identifier Rust @@ -165,8 +164,13 @@ def _native_global_index_result_supported(self) -> bool: result = self.file_scanner._global_index_result if result is None: return True - if (self.table.is_primary_key_table - or not self.file_scanner.data_evolution): + from pypaimon.read.native_plan import native_method_available + if not native_method_available('TableScan', 'with_global_index_result'): + return False + if self.table.is_primary_key_table: + from pypaimon.table.source.global_index_split_result import GlobalIndexSplitResult + return isinstance(result, GlobalIndexSplitResult) + if not self.file_scanner.data_evolution: return False from pypaimon.globalindex.global_index_result import GlobalIndexResult return isinstance(result, GlobalIndexResult) @@ -175,11 +179,7 @@ def _native_row_ranges(self) -> Optional[List[Tuple[int, int]]]: row_ranges = getattr(self.file_scanner, '_row_ranges', None) if row_ranges is not None: return [(range_.from_, range_.to) for range_ in row_ranges] - result = self.file_scanner._global_index_result - if result is None: - return None - return [(range_.from_, range_.to) - for range_ in result.results().to_range_list()] + return None def _try_native_plan(self) -> Optional[Plan]: """Plan via pypaimon_rust, then drop partitions the predicate rejects. @@ -218,6 +218,8 @@ def _try_native_plan(self) -> Optional[Plan]: extra_options['row_position_slice'] = ( fs.start_pos_of_this_subtask, fs.end_pos_of_this_subtask) row_ranges = self._native_row_ranges() + if fs._global_index_result is not None: + extra_options['global_index_result'] = fs._global_index_result if self.partition_predicate is not None: # Explicit partition selectors are scan-only. Keep the ordinary # row predicate separate so core can apply Java's override and @@ -238,6 +240,7 @@ def _try_native_plan(self) -> Optional[Plan]: ) splits = plan.splits() if (self.table.options.merge_engine() == 'first-row' + and (fs._global_index_result is None or self.table.options.query_auth_enabled) and not fs.skip_level0 and not fs.is_streaming and any(file.level == 0 for split in splits for file in split.files) and any(not split.raw_convertible and any(file.level > 0 for file in split.files) @@ -247,7 +250,9 @@ def _try_native_plan(self) -> Optional[Plan]: # physical order need not match the PK merge comparator. return None partition_predicate = self.file_scanner.partition_key_predicate - if partition_predicate is not None and not any(isinstance(split, QueryAuthSplit) for split in splits): + if (partition_predicate is not None and not ( + self.table.is_primary_key_table and fs._global_index_result is not None) + and not any(isinstance(split, QueryAuthSplit) for split in splits)): # Authorized user predicates observe masked partition values; # Rust has already applied only the safe raw partition tests. splits = [s for s in splits @@ -260,15 +265,6 @@ def _try_native_plan(self) -> Optional[Plan]: if not (getattr(fs, 'deletion_vectors_enabled', False) and not self.table.is_primary_key_table): splits = fs._apply_push_down_limit(splits) - # Attach scores to the row ranges retained by native planning. - from pypaimon.globalindex.indexed_split import IndexedSplit, scores_for_ranges - from pypaimon.globalindex.vector_search_result import ScoredGlobalIndexResult - result = fs._global_index_result - if fs._row_ranges is None and isinstance(result, ScoredGlobalIndexResult): - splits = [IndexedSplit( - split.data_split(), split.row_ranges(), - scores_for_ranges(result.score_getter(), split.row_ranges()), - ) for split in splits] return Plan(splits, snapshot_id=plan.snapshot_id) except Exception as e: _raise_if_native_fork_safety_error(e) diff --git a/paimon-python/pypaimon/table/source/primary_key_scored_result.py b/paimon-python/pypaimon/table/source/primary_key_scored_result.py index 73c370f2979b..ea836ed48977 100644 --- a/paimon-python/pypaimon/table/source/primary_key_scored_result.py +++ b/paimon-python/pypaimon/table/source/primary_key_scored_result.py @@ -56,6 +56,7 @@ def __init__(self, snapshot_id, source_splits, positions): raise ValueError("Primary-key search result contains duplicate physical positions.") self._splits = tuple(_materialize( source_splits, positions, snapshot_id)) + self._native_result = None @property def snapshot_id(self): @@ -67,8 +68,20 @@ def positions(self): @property def splits(self): + # Exposing mutable Python splits transfers metadata ownership to the + # caller. A later native scan must observe any edits to that view. + self._native_result = None return self._splits + def __getstate__(self): + # SearchResult owns a live core table. Transport the materialized + # Python metadata. copy.copy also uses this protocol and shares the + # mutable splits, so the original must stop using its native cache too. + self._native_result = None + state = self.__dict__.copy() + state.pop('_native_result', None) + return state + @classmethod def _from_native(cls, table, result): """Convert the selections already built by core without rebuilding them.""" @@ -81,6 +94,7 @@ def _from_native(cls, table, result): value._splits = tuple(deserialize_split_v1( split, _partition_fields(table), table.trimmed_primary_keys_fields) for split in result.splits()) + value._native_result = result return value diff --git a/paimon-python/pypaimon/tests/native_index_result_scan_test.py b/paimon-python/pypaimon/tests/native_index_result_scan_test.py new file mode 100644 index 000000000000..bae7555cd4c1 --- /dev/null +++ b/paimon-python/pypaimon/tests/native_index_result_scan_test.py @@ -0,0 +1,266 @@ +# 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. + +"""Externally computed PK selections keep their snapshot and physical addresses.""" + +import copy +import pickle +from unittest.mock import patch + +import pytest + +from pypaimon import Schema +from pypaimon.read.scanner.file_scanner import FileScanner +from pypaimon.read.split_serializer import serialize_split_v1 +from pypaimon.table.source.global_index_split_result import GlobalIndexSplitResult +from pypaimon.tests import native_pk_full_text_search_test as text +from pypaimon.tests import native_query_auth_test as auth +from pypaimon.tests import native_vector_search_test as vectors + +rest_catalog = text.rest_catalog +golden_catalog = text.golden_catalog +source = text.source +guarded_table = auth.guarded_table +pk_vector_table = vectors.pk_vector_table +vector_table = vectors.vector_table +pytestmark = pytest.mark.native_plan + + +class _SplitResult(GlobalIndexSplitResult): + def __init__(self, plan): + self._plan = plan + + @property + def snapshot_id(self): + return self._plan.snapshot_id + + @property + def splits(self): + return self._plan.splits() + + +@pytest.mark.parametrize('native_read', ['true', 'false']) +@pytest.mark.parametrize('snapshot', [1, 5, 8]) +@pytest.mark.parametrize('query', ['Paimon', 'no match']) +def test_pk_supplied_results_use_native_scan_and_preserve_metadata(source, native_read, snapshot, query): + table, _ = source + captured = table.copy({'scan.snapshot-id': str(snapshot), 'read.native.enabled': 'false'}) + result = text._builder(captured, query).execute_local() + latest = table.copy({'scan.native-plan.enabled': 'true', 'read.native.enabled': native_read}) + builder = latest.new_read_builder().with_projection(['id']) + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = builder.new_scan().with_global_index_result(result).plan() + assert plan.snapshot_id == result.snapshot_id + assert [serialize_split_v1(s) for s in plan.splits()] == [serialize_split_v1(s) for s in result.splits] + expected = captured.new_read_builder().with_projection(['id']).new_read().to_arrow(list(result.splits)) + actual = builder.new_read().to_arrow(plan.splits()) + assert actual.to_pylist() == expected.to_pylist() + + +@pytest.mark.parametrize('hint', ['filter', 'limit', 'shard']) +def test_supplied_pk_splits_precede_ordinary_planning_hints(source, hint): + table, _ = source + result = text._builder(table.copy({'read.native.enabled': 'false'})).execute_local() + builder = table.copy({'scan.native-plan.enabled': 'true'}).new_read_builder() + if hint == 'filter': + builder.with_filter(builder.new_predicate_builder().greater_than('id', 100)) + elif hint == 'limit': + builder.with_limit(0) + scan = builder.new_scan().with_global_index_result(result) + if hint == 'shard': + scan.with_shard(1, 2) + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = scan.plan() + assert plan.snapshot_id == result.snapshot_id + assert [serialize_split_v1(s) for s in plan.splits()] == [serialize_split_v1(s) for s in result.splits] + if hint in ('filter', 'limit'): + assert builder.new_read().to_arrow(plan.splits()).num_rows == 0 + + +def test_core_search_result_passes_directly_to_scan_and_scored_splits_round_trip(source): + from pypaimon_rust.datafusion import GlobalIndexResult, Split + table, _ = source + native = text._native_table(table) + result = text._native(table).execute_local() + scan = native.new_read_builder().with_projection(['id']).new_scan().with_global_index_result(result) + for plan in [scan.plan(), scan.plan()]: + assert plan.snapshot_id() == result.snapshot_id() + assert [s.serialize() for s in plan.splits()] == result.splits() + splits = [pickle.loads(pickle.dumps(Split.deserialize(s))) for s in result.splits()] + carried = GlobalIndexResult.from_splits(result.snapshot_id(), iter(splits)) + plan = native.new_read_builder().new_scan().with_global_index_result(carried).plan() + assert [s.serialize() for s in plan.splits()] == result.splits() + reader = native.new_read_builder().with_projection(['id']).new_read() + assert sum(batch.num_rows for batch in reader.read(plan.splits())) == 2 + + +@pytest.mark.parametrize('family', ['full-text', 'vector']) +def test_native_search_result_hands_core_selection_directly_to_scan(source, pk_vector_table, family): + table = source[0] if family == 'full-text' else pk_vector_table + table = table.copy({'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'}) + search = text._builder(table) if family == 'full-text' else vectors._single(table, [1, 0, 0, 0], 2) + result = search.execute_local() + expected = sorted(position.row_position for position in result.positions) + with patch('pypaimon.read.native_plan.serialize_split_v1', + side_effect=AssertionError('Native result round-tripped through Python')), \ + patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = table.new_read_builder().new_scan().with_global_index_result(result).plan() + assert plan.snapshot_id == result.snapshot_id + assert sorted(position for split in plan.splits() for r in split.row_ranges() + for position in range(r.from_, r.to + 1)) == expected + + +def test_exposed_python_splits_and_pickled_results_use_current_metadata(source): + table, _ = source + table = table.copy({'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'}) + result = text._builder(table).execute_local() + transported = pickle.loads(pickle.dumps(result)) + builder = table.new_read_builder().with_projection(['id']) + plan = builder.new_scan().with_global_index_result(transported).plan() + assert sorted(builder.new_read().to_arrow(plan.splits()).column('id').to_pylist()) == [1, 3] + # Accessing splits exposes mutable lists. Subsequent scans must respect + # the edited selections, including a deliberately empty one. + for split in result.splits: + split.row_ranges().clear() + split.scores().clear() + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = builder.new_scan().with_global_index_result(result).plan() + assert plan.snapshot_id == result.snapshot_id + assert all(split.row_count == 0 and split.scores() == [] for split in plan.splits()) + assert builder.new_read().to_arrow(plan.splits()).num_rows == 0 + + +def test_empty_native_pk_search_preserves_java_snapshot_zero(source, rest_catalog): + catalog, _ = rest_catalog + schema = source[0].table_schema + catalog.create_table('default.empty_text', Schema( + schema.fields, schema.partition_keys, schema.primary_keys, schema.options), False) + table = catalog.get_table('default.empty_text').copy({ + 'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'}) + result = text._builder(table).execute_local() + assert result.snapshot_id == 0 + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = table.new_read_builder().new_scan().with_global_index_result(result).plan() + assert plan.snapshot_id == 0 + assert plan.splits() == [] + + +def test_shallow_copy_split_edits_are_observed_when_scanning_the_original(source): + table = source[0].copy({'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'}) + result = text._builder(table).execute_local() + alias = copy.copy(result) + for split in alias.splits: + split.row_ranges().clear() + split.scores().clear() + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = table.new_read_builder().new_scan().with_global_index_result(result).plan() + assert plan.snapshot_id == result.snapshot_id + assert all(split.row_count == 0 for split in plan.splits()) + + +@pytest.mark.parametrize('native_read', ['true', 'false']) +def test_pk_vector_results_use_the_same_native_scan(pk_vector_table, native_read): + source = pk_vector_table.copy({'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false', + 'vector-index.search-mode': 'fast'}) + result = vectors._single(source, [1, 0, 0, 0]).execute_local() + expected = vectors._rows(source, result).column('id').to_pylist() + table = pk_vector_table.copy({'read.native.enabled': native_read, 'scan.native-plan.enabled': 'true'}) + builder = table.new_read_builder().with_projection(['id']) + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = builder.new_scan().with_global_index_result(result).plan() + assert [serialize_split_v1(s) for s in plan.splits()] == [serialize_split_v1(s) for s in result.splits] + assert sorted(builder.new_read().to_arrow(plan.splits()).column('id').to_pylist()) == sorted(expected) + + +@pytest.mark.parametrize('native_read', ['true', 'false']) +def test_nonindexed_pk_splits_are_preserved_too(source, native_read): + table, _ = source + original = table.copy({'scan.native-plan.enabled': 'false', 'read.native.enabled': 'false'}) + original_builder = original.new_read_builder() + result = _SplitResult(original_builder.new_scan().plan()) + expected = original_builder.new_read().to_arrow(result.splits).to_pylist() + builder = table.copy({'scan.native-plan.enabled': 'true', 'read.native.enabled': native_read}).new_read_builder() + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = builder.new_scan().with_global_index_result(result).plan() + assert [serialize_split_v1(s) for s in plan.splits()] == [serialize_split_v1(s) for s in result.splits] + assert builder.new_read().to_arrow(plan.splits()).to_pylist() == expected + + +@pytest.mark.parametrize('selection', ['all', 'slice', 'shard', 'empty']) +def test_de_scored_results_align_scores_in_core(vector_table, selection): + from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult + table = vector_table.copy({'scan.native-plan.enabled': 'true', 'read.native.enabled': 'true'}) + result = DictBasedScoredIndexResult({4: 0.2, 0: 0.9, 3: 0.7} if selection != 'empty' else {}) + scan = table.new_read_builder().new_scan().with_global_index_result(result) + if selection == 'slice': + scan.with_slice(3, 5) + elif selection == 'shard': + scan.with_shard(1, 2) + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = scan.plan() + expected = {0: 0.9, 3: 0.7, 4: 0.2} + if selection in ('slice', 'shard'): + expected.pop(0) + elif selection == 'empty': + expected = {} + actual = {} + for split in plan.splits(): + for range_ in split.row_ranges(): + for row_id in range(range_.from_, range_.to + 1): + actual[row_id] = split.get_score(row_id) + assert actual == pytest.approx(expected) + + +def test_de_supplied_result_preserves_explicit_incremental_window(vector_table): + from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult + from pypaimon.read.native_plan import native_plan + table = vector_table.copy({'read.native.enabled': 'true'}) + start = table.snapshot_manager().get_latest_snapshot().id + vectors._write(table, [{'id': 7, 'embedding': [1, 0], 'pt': 1}]) + end = table.snapshot_manager().get_latest_snapshot().id + plan = native_plan(table, incremental_range=(start, end), + global_index_result=DictBasedScoredIndexResult({0: 0.1, 6: 0.9})) + assert plan.snapshot_id == end + assert len(plan.splits()) == 1 + assert plan.splits()[0].scores() == pytest.approx([0.9]) + assert vectors._rows(table, DictBasedScoredIndexResult({6: 0.9})).column('id').to_pylist() == [7] + builder = table.new_read_builder().with_projection(['id']) + assert builder.new_read().to_arrow(plan.splits()).to_pylist() == [{'id': 7}] + + +@pytest.mark.parametrize('guarded_table', [ + ('primary-key', {'bucket': '1'}), + ('data-evolution', {'row-tracking.enabled': 'true', 'data-evolution.enabled': 'true'}), +], indirect=True) +@pytest.mark.parametrize('native', [False, True]) +def test_authorized_scans_ignore_external_raw_index_selections(guarded_table, native): + from pypaimon.api.api_response import AuthTableQueryResponse + from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult + from pypaimon.read.plan import Plan + table, server = guarded_table + server.table_query_auth_store[table.identifier.get_full_name()] = AuthTableQueryResponse( + filter=[auth._filter('id', 'GREATER_THAN', [1])]) + result = _SplitResult(Plan([], 1)) if table.is_primary_key_table else DictBasedScoredIndexResult({0: 0.9}) + table = table.copy({'scan.native-plan.enabled': str(native).lower(), 'read.native.enabled': str(native).lower()}) + builder = table.new_read_builder().with_projection(['id']) + scan = builder.new_scan().with_global_index_result(result) + if native: + with patch.object(FileScanner, 'scan', side_effect=AssertionError('Python scan ran')): + plan = scan.plan() + else: + plan = scan.plan() + assert sorted(builder.new_read().to_arrow(plan.splits()).column('id').to_pylist()) == [2, 3, 4] diff --git a/paimon-python/pypaimon/tests/native_plan_expanded_test.py b/paimon-python/pypaimon/tests/native_plan_expanded_test.py index f2bd2b7b7fb2..af4874ae10ed 100644 --- a/paimon-python/pypaimon/tests/native_plan_expanded_test.py +++ b/paimon-python/pypaimon/tests/native_plan_expanded_test.py @@ -184,7 +184,11 @@ def test_scored_de_preserves_score_mapping_and_selection(catalog, native, dv): plan, actual = read(table, native, predicate, shard, result=result) for split in plan.splits(): ids = [i for r in split.row_ranges() for i in range(r.from_, r.to + 1)] - assert split.scores() == [scores[i] for i in ids] + # Java IndexedSplit stores float[], including when a Native + # plan crosses the split wire format back into Python. + expected_scores = [pa.scalar(scores[i], type=pa.float32()).as_py() + if native else scores[i] for i in ids] + assert split.scores() == expected_scores python_rows = read(table, False, predicate, shard, result=result)[1] assert actual == python_rows assert all(row['id'] in scores and (not dv or row['id'] != 1) for row in actual) diff --git a/paimon-python/pypaimon/tests/native_plan_test.py b/paimon-python/pypaimon/tests/native_plan_test.py index d2a522189e1b..f9b629d4e43c 100644 --- a/paimon-python/pypaimon/tests/native_plan_test.py +++ b/paimon-python/pypaimon/tests/native_plan_test.py @@ -191,6 +191,7 @@ def setUp(self): fake_df.PaimonCatalog = type( 'PaimonCatalog', (), {'get_table': lambda self, name: None}) fake_df.Split = type('Split', (), {'serialize': lambda self: b''}) + fake_df.TableScan = type('TableScan', (), {'with_global_index_result': lambda self, result: self}) fake_mod = ModuleType('pypaimon_rust') fake_mod.datafusion = fake_df patcher = patch.dict( @@ -381,7 +382,7 @@ def test_plan_forwards_filter_limit_partition_and_time_travel(self): row_ranges=None, ) - def test_plan_forwards_global_index_row_ranges(self): + def test_plan_forwards_global_index_result(self): fs = Mock(partition_key_predicate=None) scan = _scan(native_enabled=True, file_scanner=fs) fs.data_evolution = True @@ -397,7 +398,8 @@ def test_plan_forwards_global_index_row_ranges(self): predicate=None, limit=None, projection=None, - row_ranges=[(1, 2), (5, 5)], + row_ranges=None, + global_index_result=fs._global_index_result, ) fs.scan.assert_not_called() self.assertEqual(plan.splits(), [split]) @@ -416,12 +418,13 @@ def test_empty_global_index_result_does_not_fall_back(self): predicate=None, limit=None, projection=None, - row_ranges=[], + row_ranges=None, + global_index_result=fs._global_index_result, ) fs.scan.assert_not_called() self.assertEqual(plan.splits(), []) - def test_scored_global_index_result_uses_native_ranges(self): + def test_scored_global_index_result_retains_core_scores(self): from pypaimon.globalindex.indexed_split import IndexedSplit from pypaimon.read.split import DataSplit @@ -430,10 +433,11 @@ def test_scored_global_index_result_uses_native_ranges(self): fs.data_evolution = True bitmap = GlobalIndexResult.from_range(Range(1, 1)).results() fs._global_index_result = ScoredGlobalIndexResult.create(bitmap, lambda _: 0.75) - split = IndexedSplit(DataSplit([], None, 0), [Range(1, 1)]) + split = IndexedSplit(DataSplit([], None, 0), [Range(1, 1)], [0.75]) with patch('pypaimon.read.native_plan.native_plan', return_value=Plan([split], 9)) as np: plan = scan.plan() - self.assertEqual(np.call_args[1]['row_ranges'], [(1, 1)]) + self.assertIs(np.call_args[1]['global_index_result'], fs._global_index_result) + self.assertIsNone(np.call_args[1]['row_ranges']) self.assertEqual(plan.splits()[0].scores(), [0.75]) self.assertEqual(plan.snapshot_id, 9) fs.scan.assert_not_called() From ff213959ae6356e4b05440381df876b3482f7942 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 15 Oct 2026 21:06:31 +0800 Subject: [PATCH 2/3] test(python): use shared native GlobalIndexResult for searches and scans --- paimon-python/pypaimon/read/native_plan.py | 1 + .../table/source/primary_key_scored_result.py | 4 +- .../tests/native_index_result_scan_test.py | 47 ++++++++++++++++++- 3 files changed, 49 insertions(+), 3 deletions(-) diff --git a/paimon-python/pypaimon/read/native_plan.py b/paimon-python/pypaimon/read/native_plan.py index 44bb03d6d1d7..b3c8267598be 100644 --- a/paimon-python/pypaimon/read/native_plan.py +++ b/paimon-python/pypaimon/read/native_plan.py @@ -486,6 +486,7 @@ def native_plan( from pypaimon.globalindex.vector_search_result import ScoredGlobalIndexResult from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult if isinstance(global_index_result, GlobalIndexSplitResult): + # Searches and scans share the same core GlobalIndexResult. result = (getattr(global_index_result, '_native_result', None) if type(global_index_result) is PrimaryKeyScoredResult else None) # Java's PK result uses snapshot 0 for an uncommitted table, diff --git a/paimon-python/pypaimon/table/source/primary_key_scored_result.py b/paimon-python/pypaimon/table/source/primary_key_scored_result.py index ea836ed48977..6b80fcad60b8 100644 --- a/paimon-python/pypaimon/table/source/primary_key_scored_result.py +++ b/paimon-python/pypaimon/table/source/primary_key_scored_result.py @@ -74,7 +74,7 @@ def splits(self): return self._splits def __getstate__(self): - # SearchResult owns a live core table. Transport the materialized + # GlobalIndexResult owns a live core table. Transport the materialized # Python metadata. copy.copy also uses this protocol and shares the # mutable splits, so the original must stop using its native cache too. self._native_result = None @@ -84,7 +84,7 @@ def __getstate__(self): @classmethod def _from_native(cls, table, result): - """Convert the selections already built by core without rebuilding them.""" + """Wrap core GlobalIndexResult; retain it for direct native scan handoff.""" from pypaimon.read.native_plan import _partition_fields from pypaimon.read.split_serializer import deserialize_split_v1 diff --git a/paimon-python/pypaimon/tests/native_index_result_scan_test.py b/paimon-python/pypaimon/tests/native_index_result_scan_test.py index bae7555cd4c1..5b7065103a42 100644 --- a/paimon-python/pypaimon/tests/native_index_result_scan_test.py +++ b/paimon-python/pypaimon/tests/native_index_result_scan_test.py @@ -25,8 +25,10 @@ from pypaimon import Schema from pypaimon.read.scanner.file_scanner import FileScanner -from pypaimon.read.split_serializer import serialize_split_v1 +from pypaimon.read.native_plan import _partition_fields +from pypaimon.read.split_serializer import deserialize_split_v1, serialize_split_v1 from pypaimon.table.source.global_index_split_result import GlobalIndexSplitResult +from pypaimon.tests import native_hybrid_search_test as hybrid from pypaimon.tests import native_pk_full_text_search_test as text from pypaimon.tests import native_query_auth_test as auth from pypaimon.tests import native_vector_search_test as vectors @@ -96,6 +98,7 @@ def test_core_search_result_passes_directly_to_scan_and_scored_splits_round_trip table, _ = source native = text._native_table(table) result = text._native(table).execute_local() + assert isinstance(result, GlobalIndexResult) scan = native.new_read_builder().with_projection(['id']).new_scan().with_global_index_result(result) for plan in [scan.plan(), scan.plan()]: assert plan.snapshot_id() == result.snapshot_id() @@ -114,6 +117,8 @@ def test_native_search_result_hands_core_selection_directly_to_scan(source, pk_v table = table.copy({'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'}) search = text._builder(table) if family == 'full-text' else vectors._single(table, [1, 0, 0, 0], 2) result = search.execute_local() + from pypaimon_rust.datafusion import GlobalIndexResult + assert isinstance(result._native_result, GlobalIndexResult) expected = sorted(position.row_position for position in result.positions) with patch('pypaimon.read.native_plan.serialize_split_v1', side_effect=AssertionError('Native result round-tripped through Python')), \ @@ -264,3 +269,43 @@ def test_authorized_scans_ignore_external_raw_index_selections(guarded_table, na else: plan = scan.plan() assert sorted(builder.new_read().to_arrow(plan.splits()).column('id').to_pylist()) == [2, 3, 4] + + +@pytest.mark.parametrize('family', ['vector', 'batch-vector', 'full-text', 'hybrid']) +def test_all_core_searches_return_one_result_for_read_and_scan(rest_catalog, family): + import pyarrow as pa + from pypaimon_rust.datafusion import GlobalIndexResult + table = hybrid._table(rest_catalog[0]) + native = text._native_table(table) + if family == 'vector': + results = [(native.new_vector_search_builder().with_vector_column('embedding') + .with_query_vector([0., 0.]).with_limit(3).execute_local())] + elif family == 'batch-vector': + results = (native.new_batch_vector_search_builder().with_vector_column('embedding') + .with_query_vectors([[0., 0.], [7., 0.]]).with_limit(3).execute_batch_local()) + assert len(results) == 2 + elif family == 'full-text': + results = [text._native(table, 'paimon', 3).execute_local()] + else: + results = [(native.new_hybrid_search_builder().with_limit(3) + .add_vector_route('embedding', [0., 0.], 3, 1., {}) + .add_full_text_route('content', hybrid.match_query('paimon'), 3, 1., {}) + .execute_local())] + for result in results: + assert isinstance(result, GlobalIndexResult) + assert result.snapshot_id() == table.snapshot_manager().get_latest_snapshot().id + expected = result.row_ids() + ranked = pa.Table.from_batches(result.new_read_builder().with_projection(['id']).read()) + assert ranked.column_names == ['id', '__paimon_search_score'] + assert sorted(ranked.column('__paimon_search_score').to_pylist()) == pytest.approx(sorted(expected.values())) + builder = native.new_read_builder().with_projection(['id']) + plan = builder.new_scan().with_global_index_result(result).plan() + rows = pa.Table.from_batches(builder.new_read().read(plan.splits())) + assert sorted(rows.column('id').to_pylist()) == sorted(ranked.column('id').to_pylist()) + carried_scores = {} + for selected in plan.splits(): + split = deserialize_split_v1(selected.serialize(), _partition_fields(table), []) + for range_ in split.row_ranges(): + for row_id in range(range_.from_, range_.to + 1): + carried_scores[row_id] = split.get_score(row_id) + assert carried_scores == pytest.approx(expected) From 3856839dbd6e8bbfec29413ce60c18f09347ce59 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 15 Oct 2026 23:08:44 +0800 Subject: [PATCH 3/3] test(python): set explicit query authorization in index planning regression --- .../pypaimon/tests/global_index_test.py | 29 +++++++++++++------ 1 file changed, 20 insertions(+), 9 deletions(-) diff --git a/paimon-python/pypaimon/tests/global_index_test.py b/paimon-python/pypaimon/tests/global_index_test.py index 044d08ae6c6b..c07833a70f30 100644 --- a/paimon-python/pypaimon/tests/global_index_test.py +++ b/paimon-python/pypaimon/tests/global_index_test.py @@ -375,15 +375,26 @@ def test_split_planning_merges_indexed_and_unindexed_ranges(self): scanner.open_file_cost = 1 scanner._deletion_files_map = unittest.mock.MagicMock(return_value={}) - with unittest.mock.patch( - "pypaimon.read.scanner.file_scanner.DataEvolutionSplitGenerator" - ) as split_generator: - scanner._create_data_evolution_split_generator() - - self.assertEqual( - [Range(1, 1), Range(10, 10 ** 12)], - split_generator.call_args[0][4], - ) + for query_auth in (False, True): + with self.subTest(query_auth=query_auth): + scanner.table.options = CoreOptions(Options({ + "query-auth.enabled": str(query_auth).lower(), + })) + scanner._eval_global_index.reset_mock() + with unittest.mock.patch( + "pypaimon.read.scanner.file_scanner.DataEvolutionSplitGenerator" + ) as split_generator: + scanner._create_data_evolution_split_generator() + + self.assertEqual( + None if query_auth else [Range(1, 1), Range(10, 10 ** 12)], + split_generator.call_args[0][4], + ) + if query_auth: + scanner._eval_global_index.assert_not_called() + else: + scanner._eval_global_index.assert_called_once_with( + scanner.manifest_scanner.return_value[1]) def test_eval_global_index_keeps_none_as_full_scan(self): from pypaimon.read.scanner.file_scanner import FileScanner