From 1fb943c1998cd52bb0898375b15660a3aef08d85 Mon Sep 17 00:00:00 2001 From: JingsongLi Date: Thu, 15 Oct 2026 16:10:16 +0800 Subject: [PATCH] [python] Use native primary-key full-text search --- .../table/source/native_full_text_search.py | 17 +- .../table/source/native_vector_search.py | 19 +- .../table/source/primary_key_scored_result.py | 25 ++- .../tests/native_pk_full_text_search_test.py | 192 ++++++++++++++++++ .../tests/native_vector_search_test.py | 4 +- .../primary_key_index_definitions_test.py | 19 +- 6 files changed, 248 insertions(+), 28 deletions(-) create mode 100644 paimon-python/pypaimon/tests/native_pk_full_text_search_test.py diff --git a/paimon-python/pypaimon/table/source/native_full_text_search.py b/paimon-python/pypaimon/table/source/native_full_text_search.py index 97a1484352f7..3665c11ebc3d 100644 --- a/paimon-python/pypaimon/table/source/native_full_text_search.py +++ b/paimon-python/pypaimon/table/source/native_full_text_search.py @@ -18,10 +18,12 @@ """Local REST full-text search delegated to Rust core.""" from pypaimon.catalog.table_query_auth import reject_search_under_query_auth +from pypaimon.common.options.core_options import GlobalIndexSearchMode from pypaimon.globalindex.full_text.native_full_text_global_index_reader import NativeFullTextIndexOptions from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult from pypaimon.read.native_plan import _native_table, _predicate_to_native, native_method_available from pypaimon.table.file_store_table import FileStoreTable +from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult from pypaimon.write.native_commit import _rest_catalog_supported @@ -31,13 +33,17 @@ def try_native_full_text_search(builder): reject_search_under_query_auth(table) if (type(table) is not FileStoreTable or not table.options.native_read_enabled() - or not table.options.data_evolution_enabled() - or not table.options.row_tracking_enabled() - or table.is_primary_key_table or table.options.file_format() != 'parquet' or not native_method_available('Table', 'new_full_text_search_builder') or not _rest_catalog_supported(table)): return None + primary_key = builder._primary_key_full_text_definition() is not None + if not primary_key and (not table.options.data_evolution_enabled() + or not table.options.row_tracking_enabled() + or table.is_primary_key_table): + return None + if primary_key and table.options.full_text_index_search_mode() != GlobalIndexSearchMode.FAST: + raise NotImplementedError('Primary-key full-text search only supports the FAST full-text index search mode') if builder._limit is None or builder._limit <= 0: raise ValueError('Limit must be positive, set via with_limit()') # The full-text adapter accepts typed lists and booleans. Reuse its mapping @@ -58,4 +64,7 @@ def try_native_full_text_search(builder): # be represented by the strict typed predicate FFI. return None # Once execution starts, propagate errors instead of running a second search. - return DictBasedScoredIndexResult(native.execute_local().row_ids()) + result = native.execute_local() + if primary_key: + return PrimaryKeyScoredResult._from_native(table, result) + return DictBasedScoredIndexResult(result.row_ids()) diff --git a/paimon-python/pypaimon/table/source/native_vector_search.py b/paimon-python/pypaimon/table/source/native_vector_search.py index b41d8f0ca71e..b0522e4f2d0d 100644 --- a/paimon-python/pypaimon/table/source/native_vector_search.py +++ b/paimon-python/pypaimon/table/source/native_vector_search.py @@ -20,11 +20,9 @@ from pypaimon.catalog.table_query_auth import reject_search_under_query_auth from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult from pypaimon.read.native_plan import ( - _catalog_metastore, _native_table, _partition_fields, _predicate_to_native, + _catalog_metastore, _native_table, _predicate_to_native, native_method_available) -from pypaimon.read.split_serializer import deserialize_split_v1 -from pypaimon.table.source.primary_key_scored_result import ( - PrimaryKeyScoredResult, PrimaryKeySearchPosition) +from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult def try_native_vector_search(builder, queries, batch=False): @@ -75,16 +73,5 @@ def try_native_vector_search(builder, queries, batch=False): native.with_query_vector(queries) result = native.execute_local() if builder._is_primary_key_vector_search(): - return _NativePrimaryKeyScoredResult(table, result) + return PrimaryKeyScoredResult._from_native(table, result) return DictBasedScoredIndexResult(result.row_ids()) - - -class _NativePrimaryKeyScoredResult(PrimaryKeyScoredResult): - """Metadata view of the physical selections already built by Rust core.""" - - def __init__(self, table, result): - self._snapshot_id = result.snapshot_id() or 0 - self._positions = tuple(PrimaryKeySearchPosition(*position) for position in result.positions()) - self._splits = tuple(deserialize_split_v1( - split, _partition_fields(table), table.trimmed_primary_keys_fields) - for split in result.splits()) 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 c0e1c2fe8360..73c370f2979b 100644 --- a/paimon-python/pypaimon/table/source/primary_key_scored_result.py +++ b/paimon-python/pypaimon/table/source/primary_key_scored_result.py @@ -69,6 +69,20 @@ def positions(self): def splits(self): return self._splits + @classmethod + def _from_native(cls, table, result): + """Convert the selections already built by core without rebuilding them.""" + from pypaimon.read.native_plan import _partition_fields + from pypaimon.read.split_serializer import deserialize_split_v1 + + value = cls.__new__(cls) + value._snapshot_id = result.snapshot_id() or 0 + value._positions = tuple(PrimaryKeySearchPosition(*position) for position in result.positions()) + value._splits = tuple(deserialize_split_v1( + split, _partition_fields(table), table.trimmed_primary_keys_fields) + for split in result.splits()) + return value + def _materialize(source_splits, positions, snapshot_id): sources = {} @@ -121,11 +135,12 @@ def _ranges(positions): def _partition_bytes(partition): - to_bytes = getattr(partition, "to_bytes", None) - if callable(to_bytes): - return bytes(to_bytes()) - values = getattr(partition, "values", ()) - return repr(tuple(values)).encode("utf-8") + from pypaimon.table.row.generic_row import GenericRowSerializer + + # Java PrimaryKeySearchPosition compares serialized BinaryRow partitions. + # Reuse the wire encoder for both GenericRow and BinaryRow, including the + # empty partition's arity prefix and eight-byte body. + return bytes(GenericRowSerializer.to_bytes(partition)) def _position_key(position): diff --git a/paimon-python/pypaimon/tests/native_pk_full_text_search_test.py b/paimon-python/pypaimon/tests/native_pk_full_text_search_test.py new file mode 100644 index 000000000000..8617989848ce --- /dev/null +++ b/paimon-python/pypaimon/tests/native_pk_full_text_search_test.py @@ -0,0 +1,192 @@ +# 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. + +"""Java primary-key full-text files searched and read through Rust core.""" + +import shutil +from contextlib import ExitStack +from pathlib import Path +from unittest.mock import patch + +import pytest + +from pypaimon import Schema +from pypaimon.api.api_response import ErrorResponse +from pypaimon.read.native_plan import _native_table, _predicate_to_native +from pypaimon.table.source.full_text_search_builder import FullTextSearchBuilderImpl +from pypaimon.table.source.primary_key_scored_result import PrimaryKeyScoredResult +from pypaimon.tests import native_vector_search_test as vectors +from pypaimon.tests.vector_search_filter_test import match_query + +rest_catalog = vectors.rest_catalog +golden_catalog = vectors.golden_catalog +pytestmark = pytest.mark.native_plan + + +@pytest.fixture +def source(rest_catalog, golden_catalog): + catalog, server = rest_catalog + golden = golden_catalog.get_table('default.test_pk_full_text_golden') + schema = golden.table_schema + catalog.create_table('default.pk_text', Schema( + schema.fields, schema.partition_keys, schema.primary_keys, schema.options), False) + destination = Path(server.data_path) / server.warehouse / 'default' / 'pk_text' + for directory in ('bucket-0', 'index', 'manifest', 'snapshot'): + shutil.copytree(str(Path(golden.table_path) / directory), str(destination / directory)) + return catalog.get_table('default.pk_text'), server + + +def _builder(table, text='Paimon', limit=10): + return table.new_full_text_search_builder().with_query('content', match_query(text)).with_limit(limit) + + +def _native(table, text='Paimon', limit=10): + return (_native_table(table).new_full_text_search_builder() + .with_query('content', match_query(text)).with_limit(limit)) + + +def _compare(table, text='Paimon', limit=10): + classic = table.copy({'read.native.enabled': 'false', 'scan.native-plan.enabled': 'false'}) + native = table.copy({'read.native.enabled': 'true', 'scan.native-plan.enabled': 'true'}) + expected = _builder(classic, text, limit).execute_local() + with patch.object(FullTextSearchBuilderImpl, 'new_full_text_scan', + side_effect=AssertionError('Python full-text search ran')): + actual = _builder(native, text, limit).execute_local() + assert isinstance(actual, PrimaryKeyScoredResult) + assert actual.snapshot_id == expected.snapshot_id + assert [(p.partition_bytes, p.bucket, p.data_file_name, p.row_position) for p in actual.positions] == [ + (p.partition_bytes, p.bucket, p.data_file_name, p.row_position) for p in expected.positions] + assert [p.score for p in actual.positions] == pytest.approx([p.score for p in expected.positions], abs=1e-6) + assert len(actual.splits) == len(expected.splits) + for actual_split, expected_split in zip(actual.splits, expected.splits): + assert len(actual_split.files) == 1 + assert actual_split.row_ranges() == expected_split.row_ranges() + assert actual_split.scores() == pytest.approx(expected_split.scores(), abs=1e-6) + rows = vectors._rows(native, actual) + assert rows.to_pylist() == vectors._rows(classic, expected).to_pylist() + return rows + + +@pytest.mark.parametrize('snapshot', [None, 1, 5, 8]) +@pytest.mark.parametrize('text', ['Paimon', 'Flink', 'no match']) +@pytest.mark.parametrize('limit', [1, 10]) +def test_java_pk_archives_preserve_positions_scores_and_snapshot(source, snapshot, text, limit): + table, _ = source + if snapshot is not None: + table = table.copy({'scan.snapshot-id': str(snapshot)}) + rows = _compare(table, text, limit) + if snapshot is None and text == 'Paimon' and limit == 10: + assert sorted(rows.column('id').to_pylist()) == [1, 3] + + +@pytest.mark.parametrize('projection', [['id'], ['content', 'id'], []]) +def test_core_result_keeps_scored_selections_and_materializes_only_the_projection(source, projection): + table, _ = source + result = _native(table).execute_local() + assert result.snapshot_id() == table.snapshot_manager().get_latest_snapshot().id + with pytest.raises(Exception, match='physical file positions'): + result.row_ids() + positions = result.positions() + assert positions and len(positions) == len(result) + reader = result.new_read_builder().with_projection(projection) + batches = list(reader.read()) + assert sum(batch.num_rows for batch in batches) == len(result) + assert all(batch.schema.names == projection + ['__paimon_search_score'] for batch in batches) + scores = [score for batch in batches for score in batch.column('__paimon_search_score').to_pylist()] + assert scores == pytest.approx([p[4] for p in positions], abs=1e-6) + # The original result can be read again without scanning or searching again. + assert sum(batch.num_rows for batch in reader.read()) == len(result) + + +def test_captured_plan_does_not_resolve_latest_again_and_accepts_another_query(source): + table, server = source + builder = _native(table) + plan = builder.new_full_text_scan().scan() + snapshot_id = plan.snapshot_id() + expected = _native(table, 'Flink', 1).execute_local().positions() + reader = _native(table, 'Flink', 1).new_full_text_read() + unavailable = server._mock_response(ErrorResponse('TABLE', 'pk_text', 'snapshot unavailable', 503), 503) + with patch.object(server, '_table_snapshot_handle', return_value=unavailable): + result = reader.read_plan(plan) + assert result.snapshot_id() == snapshot_id + assert result.positions() == expected + assert sum(batch.num_rows for batch in result.new_read_builder().with_projection(['id']).read()) == len(result) + + +def test_search_returns_metadata_without_opening_user_parquet_files(source): + table, server = source + expected = _native(table).execute_local().positions() + files = list((Path(server.data_path) / server.warehouse / 'default' / 'pk_text' / 'bucket-0').glob('*.parquet')) + assert expected and files + with ExitStack() as restore: + for path in files: + hidden = path.with_suffix('.held') + path.rename(hidden) + restore.callback(hidden.rename, path) + result = _native(table).execute_local() + assert result.positions() == expected + with pytest.raises(Exception, match='NotFound|not found|No such file|does not exist'): + result.new_read_builder().with_projection(['id']).read() + + +@pytest.mark.parametrize('case', ['table', 'filter', 'kind', 'column']) +def test_pk_plan_cannot_be_read_under_another_context(source, case): + table, _ = source + plan = _native(table).new_full_text_scan().scan() + if case == 'table': + catalog = table.catalog_environment.catalog_loader.load() + schema = table.table_schema + catalog.create_table('default.other_text', Schema( + schema.fields, schema.partition_keys, schema.primary_keys, schema.options), False) + reader = _native(catalog.get_table('default.other_text')).new_full_text_read() + with pytest.raises(Exception, match='different table, column or filter'): + reader.read_plan(plan) + elif case == 'filter': + reader = _native(table) + reader.with_filter(_predicate_to_native(table.new_read_builder().new_predicate_builder().equal('id', 1))) + with pytest.raises(Exception, match='non-partition filters'): + reader.new_full_text_read() + elif case == 'kind': + changed = table.copy({'pk-full-text.index.columns': ''}) + reader = _native(changed).new_full_text_read() + with pytest.raises(Exception, match='different table, column or filter'): + reader.read_plan(plan) + else: + reader = (_native_table(table).new_full_text_search_builder() + .with_query('id', match_query('Paimon')).with_limit(10)) + with pytest.raises(Exception, match='character string'): + reader.new_full_text_read() + + +@pytest.mark.parametrize('mode', ['full', 'detail']) +@pytest.mark.parametrize('native', [False, True]) +def test_pk_fast_only_contract_is_unchanged(source, mode, native): + table, _ = source + selected = table.copy({'read.native.enabled': str(native).lower(), 'full-text-index.search-mode': mode}) + with pytest.raises(NotImplementedError, match='FAST'): + _builder(selected).execute_local() + + +def test_native_pk_search_failure_does_not_retry_python(source): + table, server = source + table = table.copy({'read.native.enabled': 'true'}) + unavailable = server._mock_response(ErrorResponse('TABLE', 'pk_text', 'snapshot unavailable', 503), 503) + with patch.object(server, '_table_snapshot_handle', return_value=unavailable), \ + patch.object(FullTextSearchBuilderImpl, 'new_full_text_scan', + side_effect=AssertionError('Python search retried')): + with pytest.raises(Exception, match='snapshot unavailable'): + _builder(table).execute_local() diff --git a/paimon-python/pypaimon/tests/native_vector_search_test.py b/paimon-python/pypaimon/tests/native_vector_search_test.py index 5f078bf2cbec..6ecd74e28a79 100644 --- a/paimon-python/pypaimon/tests/native_vector_search_test.py +++ b/paimon-python/pypaimon/tests/native_vector_search_test.py @@ -203,8 +203,8 @@ def configure(table): result = configure(native).execute_local() assert isinstance(result, PrimaryKeyScoredResult) assert result.snapshot_id == expected.snapshot_id - assert [(p.bucket, p.data_file_name, p.row_position) for p in result.positions] == [ - (p.bucket, p.data_file_name, p.row_position) for p in expected.positions] + assert [(p.partition_bytes, p.bucket, p.data_file_name, p.row_position) for p in result.positions] == [ + (p.partition_bytes, p.bucket, p.data_file_name, p.row_position) for p in expected.positions] assert [p.score for p in result.positions] == pytest.approx([p.score for p in expected.positions]) rows = _rows(native, result) assert sorted(rows.column('id').to_pylist()) == sorted(_rows(classic, expected).column('id').to_pylist()) diff --git a/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py b/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py index 1f22cd7d6542..7e345f709413 100644 --- a/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py +++ b/paimon-python/pypaimon/tests/primary_key_index_definitions_test.py @@ -173,7 +173,8 @@ def test_scored_result_uses_file_local_positions(self): min_sequence_number=0, max_sequence_number=0, schema_id=0, level=1, extra_files=[]) source = DataSplit([data_file], partition, 3) - partition_bytes = repr(tuple(partition.values)).encode("utf-8") + from pypaimon.table.row.generic_row import GenericRowSerializer + partition_bytes = GenericRowSerializer.to_bytes(partition) result = PrimaryKeyScoredResult(7, [source], [ PrimaryKeySearchPosition(partition_bytes, 3, "data.parquet", 1, 0.9), @@ -187,6 +188,22 @@ def test_scored_result_uses_file_local_positions(self): [(r.from_, r.to) for r in result.splits[0].row_ranges()]) self.assertEqual([0.9, 0.8, 0.7], result.splits[0].scores()) + def test_search_position_partitions_use_java_binary_row_bytes(self): + from pypaimon.table.row.binary_row import BinaryRow + from pypaimon.table.row.generic_row import GenericRowSerializer + from pypaimon.table.source.primary_key_scored_result import _partition_bytes + + self.assertEqual(bytes(12), _partition_bytes(GenericRow([], []))) + fields = [DataField(0, 'pt', AtomicType('INT'))] + for value in (None, 2, 10, -1): + generic = GenericRow([value], fields) + encoded = GenericRowSerializer.to_bytes(generic) + self.assertEqual(encoded, _partition_bytes(generic)) + self.assertEqual(encoded, _partition_bytes(BinaryRow(encoded, fields))) + # BinaryRow ordering is the Java physical-position tie breaker. + self.assertLess(_partition_bytes(GenericRow([2], fields)), + _partition_bytes(GenericRow([10], fields))) + def test_pk_vector_filter_is_applied_before_top_k(self): split = SimpleNamespace(row_ranges_by_file={ "first.parquet": (Range(1, 2),),