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 c9a36a81d7 [python] Use native scans for supplied primary-key index
results (#10591)
c9a36a81d7 is described below
commit c9a36a81d7ba761f544be5f07c9defcea6d4ffc2
Author: Jingsong Lee <[email protected]>
AuthorDate: Thu Oct 15 23:22:17 2026 +0800
[python] Use native scans for supplied primary-key index results (#10591)
---
paimon-python/pypaimon/read/native_plan.py | 39 ++-
.../pypaimon/read/scanner/file_scanner.py | 13 +-
paimon-python/pypaimon/read/table_scan.py | 32 +--
.../table/source/primary_key_scored_result.py | 16 +-
paimon-python/pypaimon/tests/global_index_test.py | 29 +-
.../tests/native_index_result_scan_test.py | 311 +++++++++++++++++++++
.../pypaimon/tests/native_plan_expanded_test.py | 6 +-
paimon-python/pypaimon/tests/native_plan_test.py | 16 +-
8 files changed, 416 insertions(+), 46 deletions(-)
diff --git a/paimon-python/pypaimon/read/native_plan.py
b/paimon-python/pypaimon/read/native_plan.py
index 0363fd8064..b3c8267598 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,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:
diff --git a/paimon-python/pypaimon/read/scanner/file_scanner.py
b/paimon-python/pypaimon/read/scanner/file_scanner.py
index b938a7f7d2..85c8061975 100755
--- a/paimon-python/pypaimon/read/scanner/file_scanner.py
+++ b/paimon-python/pypaimon/read/scanner/file_scanner.py
@@ -334,10 +334,13 @@ class FileScanner:
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 @@ class FileScanner:
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 40f0383d30..687d5c8690 100755
--- a/paimon-python/pypaimon/read/table_scan.py
+++ b/paimon-python/pypaimon/read/table_scan.py
@@ -104,7 +104,6 @@ class TableScan:
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 @@ class TableScan:
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 @@ class TableScan:
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 @@ class TableScan:
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 @@ class TableScan:
)
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 @@ class TableScan:
# 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 @@ class TableScan:
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 73c370f297..6b80fcad60 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 @@ class PrimaryKeyScoredResult(GlobalIndexSplitResult):
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 @@ class PrimaryKeyScoredResult(GlobalIndexSplitResult):
@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 @@ class PrimaryKeyScoredResult(GlobalIndexSplitResult):
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/global_index_test.py
b/paimon-python/pypaimon/tests/global_index_test.py
index 044d08ae6c..c07833a70f 100644
--- a/paimon-python/pypaimon/tests/global_index_test.py
+++ b/paimon-python/pypaimon/tests/global_index_test.py
@@ -375,15 +375,26 @@ class GlobalIndexScalarFallbackTest(unittest.TestCase):
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
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 0000000000..5b7065103a
--- /dev/null
+++ b/paimon-python/pypaimon/tests/native_index_result_scan_test.py
@@ -0,0 +1,311 @@
+# 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.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
+
+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()
+
+
[email protected]('native_read', ['true', 'false'])
[email protected]('snapshot', [1, 5, 8])
[email protected]('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()
+
+
[email protected]('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()
+ 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()
+ 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
+
+
[email protected]('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()
+ 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')), \
+ 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())
+
+
[email protected]('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)
+
+
[email protected]('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
+
+
[email protected]('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}]
+
+
[email protected]('guarded_table', [
+ ('primary-key', {'bucket': '1'}),
+ ('data-evolution', {'row-tracking.enabled': 'true',
'data-evolution.enabled': 'true'}),
+], indirect=True)
[email protected]('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]
+
+
[email protected]('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)
diff --git a/paimon-python/pypaimon/tests/native_plan_expanded_test.py
b/paimon-python/pypaimon/tests/native_plan_expanded_test.py
index f2bd2b7b7f..af4874ae10 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 d2a522189e..f9b629d4e4 100644
--- a/paimon-python/pypaimon/tests/native_plan_test.py
+++ b/paimon-python/pypaimon/tests/native_plan_test.py
@@ -191,6 +191,7 @@ class NativePlanTest(unittest.TestCase):
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 @@ class NativePlanTest(unittest.TestCase):
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 @@ class NativePlanTest(unittest.TestCase):
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 @@ class NativePlanTest(unittest.TestCase):
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 @@ class NativePlanTest(unittest.TestCase):
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()
