[python] Use native scans for supplied primary-key index results

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,41 @@ 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):
+            # 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,
+            # 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:

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)

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)

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,11 +68,23 @@ 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):
+        # 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
+        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."""
+        """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
 
@@ -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
 
 

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