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 12bb2ae313 [python] Bound full-text candidates while preserving score
ties (#10583)
12bb2ae313 is described below
commit 12bb2ae313d592703a8df893ef840586be48e4ce
Author: chaoyang <[email protected]>
AuthorDate: Fri Oct 16 21:56:48 2026 +0800
[python] Bound full-text candidates while preserving score ties (#10583)
---
.../native_full_text_global_index_reader.py | 37 +++--
paimon-python/pypaimon/ray/full_text_search.py | 6 +-
.../pypaimon/table/source/full_text_read.py | 36 +++--
.../pypaimon/tests/full_text_top_k_test.py | 152 +++++++++++++++++++++
.../pypaimon/tests/vector_search_filter_test.py | 20 +--
5 files changed, 216 insertions(+), 35 deletions(-)
diff --git
a/paimon-python/pypaimon/globalindex/full_text/native_full_text_global_index_reader.py
b/paimon-python/pypaimon/globalindex/full_text/native_full_text_global_index_reader.py
index 629dee6454..5d621e0efe 100644
---
a/paimon-python/pypaimon/globalindex/full_text/native_full_text_global_index_reader.py
+++
b/paimon-python/pypaimon/globalindex/full_text/native_full_text_global_index_reader.py
@@ -86,11 +86,12 @@ class PaimonFullTextInput:
class NativeFullTextGlobalIndexReader(GlobalIndexReader):
"""Full-text global index reader using paimon-full-text."""
- def __init__(self, file_io, index_path: str, io_metas:
List[GlobalIndexIOMeta]):
+ def __init__(self, file_io, index_path: str, io_metas:
List[GlobalIndexIOMeta], max_candidates=None):
assert len(io_metas) == 1, "Expected exactly one index file per shard"
self._file_io = file_io
self._index_path = index_path
self._io_meta = io_metas[0]
+ self._max_candidates = max_candidates
self._searcher = None
self._native_reader = None
self._stream = None
@@ -107,13 +108,8 @@ class NativeFullTextGlobalIndexReader(GlobalIndexReader):
filter_bytes = (
None if include_row_ids is None else include_row_ids.serialize()
)
- row_ids, scores = self._native_reader.search(
- full_text_search.query,
- limit=limit,
- filter_bytes=filter_bytes)
- id_to_scores = dict(zip(row_ids, scores))
- return _completed_future(
- DictBasedScoredIndexResult(id_to_scores).top_k(limit))
+ return _completed_future(_search_top_k(
+ self._native_reader, full_text_search.query, limit, filter_bytes,
self._max_candidates))
def _ensure_loaded(self):
if self._searcher is not None:
@@ -203,6 +199,31 @@ class NativeFullTextGlobalIndexReader(GlobalIndexReader):
self._stream = None
+def _search_top_k(reader, query, limit, filter_bytes=None,
max_candidates=None):
+ # The native collector breaks score ties by document ordinal, whereas
+ # Python ranks ties by row ID. Read beyond the boundary until every tied
+ # candidate is available, without fetching the whole corpus otherwise.
+ maximum = 2147483647 if max_candidates is None else min(max_candidates,
2147483647)
+ if maximum == 0:
+ return DictBasedScoredIndexResult({})
+ fetch = min(limit + 1, maximum)
+ probes = 0
+ while True:
+ row_ids, scores = reader.search(query, limit=fetch,
filter_bytes=filter_bytes)
+ selected = DictBasedScoredIndexResult(dict(zip(row_ids,
scores))).top_k(limit)
+ getter = selected.score_getter()
+ if (len(row_ids) < fetch or len(row_ids) <= limit or fetch == maximum
+ or min(scores) < min(getter(row_id) for row_id in
selected.results())):
+ # Do not retain a score getter closing over the larger candidate
map.
+ return DictBasedScoredIndexResult({row_id: getter(row_id) for
row_id in selected.results()})
+ probes += 1
+ # Every probe reruns the native query. After two small probes, fetch
+ # the known corpus bound once instead of repeatedly scanning a large
+ # tie group. Never use INT_MAX as a fallback without a known bound:
+ # the native wrapper allocates output buffers of the requested size.
+ fetch = maximum if max_candidates is not None and probes >= 2 else
min(fetch * 2, maximum)
+
+
def _read_fully(stream, length: int) -> bytes:
buf = bytearray()
remaining = length
diff --git a/paimon-python/pypaimon/ray/full_text_search.py
b/paimon-python/pypaimon/ray/full_text_search.py
index 29114a0c36..57d945ef4c 100644
--- a/paimon-python/pypaimon/ray/full_text_search.py
+++ b/paimon-python/pypaimon/ray/full_text_search.py
@@ -25,7 +25,7 @@ from pypaimon.globalindex.global_index_result import
GlobalIndexResult
from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
from pypaimon.ray.vector_search import _execution_options, _map_tasks,
_require_ray, _scores, _shared_scalar_filters
from pypaimon.table.source.full_text_read import (
- DataEvolutionFullTextRead, _index_type, _raw_row_ranges,
_split_search_splits,
+ DataEvolutionFullTextRead, _index_candidate_limit, _index_type,
_raw_row_ranges, _split_search_splits,
)
from pypaimon.table.source import global_index_live_row_filter
from pypaimon.utils.range import Range
@@ -68,9 +68,11 @@ class _RayFullTextRead(DataEvolutionFullTextRead):
def _eval_column_query(self, splits_by_column, live_rows):
items = splits_by_column.get(self._text_columns[0].name, [])
merged = {}
+ reader = self._worker_reader()
+ reader._limit = _index_candidate_limit(items, self._limit)
with _shared_scalar_filters(items, self._filter, self._remote_args,
_prepare_scalar_filter) as prepare,
closing(_map_tasks(
- _search_index, self._worker_reader(),
items,
+ _search_index, reader, items,
self._concurrency, self._remote_args,
True, prepare)) as tasks:
for _, scores in tasks:
# Duplicate row IDs keep the first planned shard's score.
diff --git a/paimon-python/pypaimon/table/source/full_text_read.py
b/paimon-python/pypaimon/table/source/full_text_read.py
index 21cc5a5cf9..def8fc754e 100644
--- a/paimon-python/pypaimon/table/source/full_text_read.py
+++ b/paimon-python/pypaimon/table/source/full_text_read.py
@@ -175,13 +175,15 @@ class DataEvolutionFullTextRead(FullTextRead):
splits = splits_by_column.get(column, [])
if not splits:
return ScoredGlobalIndexResult.create_empty()
+ reader = copy(self)
+ reader._limit = _index_candidate_limit(splits, self._limit)
futures = []
for split in splits:
include_row_ids = global_index_live_row_filter.for_range(
live_rows, split.row_range_start, split.row_range_end)
if include_row_ids is not None and include_row_ids.is_empty():
continue
- futures.append(self._eval(
+ futures.append(reader._eval(
split.row_range_start, split.row_range_end,
split.full_text_index_files,
include_row_ids,
@@ -223,13 +225,13 @@ class DataEvolutionFullTextRead(FullTextRead):
reader = _create_full_text_reader(
index_type, file_io, index_path,
- index_io_meta_list
+ index_io_meta_list, row_range_end - row_range_start + 1
)
full_text_search = FullTextSearch(
field_name=self._text_columns[0].name,
query=self._query,
- limit=_candidate_limit(row_range_start, row_range_end),
+ limit=min(self._limit, row_range_end - row_range_start + 1),
)
if include_row_ids is not None:
full_text_search =
full_text_search.with_include_row_ids(include_row_ids)
@@ -251,7 +253,6 @@ class DataEvolutionFullTextRead(FullTextRead):
raise ValueError(f"Unsupported full-text index type:
'{index_type}'")
row_range_start = raw_row_ranges[0].from_
- row_range_end = raw_row_ranges[-1].to
search_kwargs = {}
if include_row_ids is not None:
raw_rows = GlobalIndexResult.from_ranges(raw_row_ranges).results()
@@ -277,7 +278,8 @@ class DataEvolutionFullTextRead(FullTextRead):
index_bytes,
row_range_start,
self._query,
- _candidate_limit(row_range_start, row_range_end),
+ min(self._limit, sum(row_range.count() for row_range in
raw_row_ranges)),
+ max_candidates=table.num_rows,
**search_kwargs,
).top_k(self._limit)
@@ -329,12 +331,16 @@ class DataEvolutionFullTextRead(FullTextRead):
writer.close()
-def _candidate_limit(row_range_start: int, row_range_end: int) -> int:
- size = row_range_end - row_range_start + 1
- return min(size, 2147483647)
+def _index_candidate_limit(splits, limit):
+ ranges = sorted((split.row_range_start, split.row_range_end) for split in
splits)
+ if any(start <= previous_end for (_, previous_end), (start, _) in
zip(ranges, ranges[1:])):
+ # Overlapping shards may assign different scores to the same row.
+ # Preserve first-shard precedence before truncating the global result.
+ return min(max(end - start + 1 for start, end in ranges), 2147483647)
+ return limit
-def _create_full_text_reader(index_type, file_io, index_path,
index_io_meta_list):
+def _create_full_text_reader(index_type, file_io, index_path,
index_io_meta_list, max_candidates=None):
"""Create a global index reader for full-text search."""
from pypaimon.globalindex.full_text.native_full_text_global_index_reader
import (
FULL_TEXT_IDENTIFIER,
@@ -343,7 +349,7 @@ def _create_full_text_reader(index_type, file_io,
index_path, index_io_meta_list
from
pypaimon.globalindex.full_text.native_full_text_global_index_reader import (
NativeFullTextGlobalIndexReader,
)
- return NativeFullTextGlobalIndexReader(file_io, index_path,
index_io_meta_list)
+ return NativeFullTextGlobalIndexReader(file_io, index_path,
index_io_meta_list, max_candidates)
raise ValueError(f"Unsupported full-text index type: '{index_type}'")
@@ -383,10 +389,10 @@ def _index_type(index_splits):
return None
-def _search_raw_full_text(index_bytes, row_range_start, query, limit,
include_row_ids=None):
+def _search_raw_full_text(index_bytes, row_range_start, query, limit,
include_row_ids=None, max_candidates=None):
from paimon_ftindex import FullTextIndexReader
from pypaimon.globalindex.full_text.native_full_text_global_index_reader
import (
- PaimonFullTextInput,
+ PaimonFullTextInput, _search_top_k,
)
reader = FullTextIndexReader(PaimonFullTextInput(BytesIO(index_bytes)))
@@ -395,11 +401,11 @@ def _search_raw_full_text(index_bytes, row_range_start,
query, limit, include_ro
if include_row_ids is not None:
relative_ids =
GlobalIndexResult.create(include_row_ids).offset(-row_range_start).results()
search_kwargs["filter_bytes"] = relative_ids.serialize()
- row_ids, scores = reader.search(query, limit=limit, **search_kwargs)
+ result = _search_top_k(reader, query, limit,
max_candidates=max_candidates, **search_kwargs)
return DictBasedScoredIndexResult(
{
- row_range_start + row_id: score
- for row_id, score in zip(row_ids, scores)
+ row_range_start + row_id: result.score_getter()(row_id)
+ for row_id in result.results()
})
finally:
reader.close()
diff --git a/paimon-python/pypaimon/tests/full_text_top_k_test.py
b/paimon-python/pypaimon/tests/full_text_top_k_test.py
new file mode 100644
index 0000000000..04b0854a93
--- /dev/null
+++ b/paimon-python/pypaimon/tests/full_text_top_k_test.py
@@ -0,0 +1,152 @@
+# 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.
+
+from io import BytesIO
+from unittest.mock import patch
+
+import pytest
+
+from pypaimon.globalindex.full_text.native_full_text_global_index_reader
import (
+ NativeFullTextIndexOptions, PaimonFullTextInput, _search_top_k,
+)
+from pypaimon.globalindex.global_index_reader import _completed_future
+from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
+from pypaimon.table.source.full_text_read import DataEvolutionFullTextRead,
_search_raw_full_text
+from pypaimon.table.source.full_text_search_split import
IndexFullTextSearchSplit
+from pypaimon.tests.vector_search_filter_test import _StubTable, _entry,
_field, _bitmap, match_query
+
+
+def _index(rows):
+ native = pytest.importorskip("paimon_ftindex")
+ writer =
native.FullTextIndexWriter(NativeFullTextIndexOptions().to_native_options())
+ try:
+ for row_id, text in rows:
+ writer.add_document(row_id, text)
+ output = BytesIO()
+ writer.write(output)
+ return output.getvalue()
+ finally:
+ writer.close()
+
+
[email protected]("filtered", [False, True])
+def test_top_k_preserves_row_id_ties_and_raw_offsets(filtered):
+ data = _index([(row_id, "paimon") for row_id in [100, 50, 30, 1, 2, 3]])
+ include = _bitmap(150, 130, 102) if filtered else None
+ result = _search_raw_full_text(data, 100, match_query("paimon"), 1,
include, max_candidates=6)
+ assert list(result.results()) == ([102] if filtered else [101])
+
+
+def test_top_k_does_not_fetch_the_whole_index_without_boundary_ties():
+ native = pytest.importorskip("paimon_ftindex")
+ data = _index([(i, "paimon " + "lake " * i) for i in range(128)])
+ reader = native.FullTextIndexReader(PaimonFullTextInput(BytesIO(data)))
+ calls = []
+
+ class RecordingReader:
+ def search(self, query, **kwargs):
+ calls.append(kwargs["limit"])
+ return reader.search(query, **kwargs)
+
+ try:
+ ids, scores = reader.search(match_query("paimon"), limit=128)
+ expected = DictBasedScoredIndexResult(dict(zip(ids, scores))).top_k(3)
+ actual = _search_top_k(RecordingReader(), match_query("paimon"), 3,
max_candidates=128)
+ assert list(actual.results()) == list(expected.results())
+ assert {i: actual.score_getter()(i) for i in actual.results()} == {
+ i: expected.score_getter()(i) for i in expected.results()}
+ assert calls[0] == 4
+ assert max(calls) < 128
+ finally:
+ reader.close()
+
+
[email protected]("filtered", [False, True])
+def test_large_boundary_ties_cap_native_query_repetitions(filtered):
+ native = pytest.importorskip("paimon_ftindex")
+ count = 1024
+ data = _index([(i, "paimon document" + str(i)) for i in
reversed(range(count))])
+ reader = native.FullTextIndexReader(PaimonFullTextInput(BytesIO(data)))
+ include = _bitmap(*range(0, count, 3)).serialize() if filtered else None
+ calls = []
+
+ class RecordingReader:
+ def search(self, query, **kwargs):
+ calls.append(kwargs["limit"])
+ assert kwargs["filter_bytes"] == include
+ return reader.search(query, **kwargs)
+
+ try:
+ ids, scores = reader.search(match_query("paimon"), limit=count,
filter_bytes=include)
+ expected = DictBasedScoredIndexResult(dict(zip(ids, scores))).top_k(10)
+ actual = _search_top_k(RecordingReader(), match_query("paimon"), 10,
include, count)
+ assert {i: actual.score_getter()(i) for i in actual.results()} == {
+ i: expected.score_getter()(i) for i in expected.results()}
+ assert calls == [11, 22, count]
+ finally:
+ reader.close()
+
+
[email protected]("bound", [None, 128])
+def test_small_boundary_ties_finish_before_the_corpus_bound(bound):
+ native = pytest.importorskip("paimon_ftindex")
+ data = _index([(i, "paimon" if i < 3 else "paimon " + "lake " * i)
+ for i in reversed(range(128))])
+ reader = native.FullTextIndexReader(PaimonFullTextInput(BytesIO(data)))
+ calls = []
+
+ class RecordingReader:
+ def search(self, query, **kwargs):
+ calls.append(kwargs["limit"])
+ return reader.search(query, **kwargs)
+
+ try:
+ actual = _search_top_k(RecordingReader(), match_query("paimon"), 1,
max_candidates=bound)
+ assert list(actual.results()) == [0]
+ assert calls == [2, 4]
+ finally:
+ reader.close()
+
+
[email protected]("overlap", [False, True])
+def
test_per_shard_top_k_preserves_global_ranking_and_duplicate_precedence(overlap):
+ field = _field(0, "text", "STRING")
+ starts = [0, 0] if overlap else [0, 3]
+ entries = [_entry(None, 0, "full-text", str(i), start, start + 2)
+ for i, start in enumerate(starts)]
+ splits = [IndexFullTextSearchSplit("text", start, start + 2,
[entry.index_file])
+ for start, entry in zip(starts, entries)]
+ scores = [{0: 1., 1: 3., 2: 2.}, {0: 10., 1: 4., 2: 2.}]
+ limits = []
+
+ class Reader:
+ def __init__(self, index):
+ self.index = index
+
+ def visit_full_text_search(self, search):
+ limits.append(search.limit)
+ return
_completed_future(DictBasedScoredIndexResult(scores[self.index]).top_k(search.limit))
+
+ def close(self):
+ pass
+
+ with patch("pypaimon.table.source.full_text_read._create_full_text_reader",
+ side_effect=lambda kind, io, path, metas, bound:
Reader(int(metas[0].file_name))):
+ read = DataEvolutionFullTextRead(_StubTable([field], entries), 1,
field, match_query("paimon"))
+ result = read._eval_column_query({"text": splits}, None)
+ assert list(result.results()) == ([1] if overlap else [3])
+ assert limits == ([3, 3] if overlap else [1, 1])
diff --git a/paimon-python/pypaimon/tests/vector_search_filter_test.py
b/paimon-python/pypaimon/tests/vector_search_filter_test.py
index b64625dc27..6b68f94dc0 100644
--- a/paimon-python/pypaimon/tests/vector_search_filter_test.py
+++ b/paimon-python/pypaimon/tests/vector_search_filter_test.py
@@ -655,7 +655,7 @@ class NativeFullTextIndexOptionsTest(unittest.TestCase):
self.assertEqual([7], sorted(list(result.results())))
self.assertEqual(match_query("paimon"), calls[0][0])
- self.assertEqual(10, calls[0][1])
+ self.assertEqual(11, calls[0][1]) # One lookahead hit detects score
ties at top-k.
self.assertIsNone(calls[0][2])
def test_reader_delegates_boolean_query_json_to_native(self):
@@ -806,7 +806,7 @@ class FullTextSearchBuilderDslTest(unittest.TestCase):
_patch_snapshot(self, [entry])
captured_searches = []
- def _fake_create(index_type, file_io, index_path, index_io_meta_list):
+ def _fake_create(index_type, file_io, index_path, index_io_meta_list,
max_candidates=None):
class _FakeReader:
def visit_full_text_search(self_inner, fts):
captured_searches.append(fts)
@@ -837,7 +837,7 @@ class FullTextSearchBuilderDslTest(unittest.TestCase):
self.assertEqual(10, captured_searches[0].limit)
self.assertEqual(match_query("paimon", "And"),
captured_searches[0].query)
- def
test_full_text_builder_sends_leaf_candidate_limit_for_compound_queries(self):
+ def test_full_text_builder_sends_top_k_for_compound_queries(self):
from pypaimon.table.source.full_text_search_builder import (
FullTextSearchBuilderImpl,
)
@@ -853,7 +853,7 @@ class FullTextSearchBuilderDslTest(unittest.TestCase):
_patch_snapshot(self, [entry])
captured_limits = []
- def _fake_create(index_type, file_io, index_path, index_io_meta_list):
+ def _fake_create(index_type, file_io, index_path, index_io_meta_list,
max_candidates=None):
class _FakeReader:
def visit_full_text_search(self_inner, fts):
captured_limits.append(fts.limit)
@@ -882,7 +882,7 @@ class FullTextSearchBuilderDslTest(unittest.TestCase):
.read([split])
)
- self.assertEqual([100], captured_limits)
+ self.assertEqual([3], captured_limits)
def test_full_text_scan_ignores_other_index_types_on_same_column(self):
from pypaimon.table.source.full_text_scan import
DataEvolutionFullTextScan
@@ -3252,7 +3252,7 @@ class FullTextSearchManySplitsTest(unittest.TestCase):
captured_io_metas = []
def _fake_create(index_type, file_io, index_path,
- index_io_meta_list):
+ index_io_meta_list, max_candidates=None):
captured_io_metas.append(list(index_io_meta_list))
class _FakeReader:
@@ -3304,7 +3304,7 @@ class FullTextSearchManySplitsTest(unittest.TestCase):
partition_filter = object()
captured = []
- def _fake_create(index_type, file_io, index_path, index_io_meta_list):
+ def _fake_create(index_type, file_io, index_path, index_io_meta_list,
max_candidates=None):
class _FakeReader:
def visit_full_text_search(self_inner, fts):
captured.append(fts.include_row_ids)
@@ -3370,7 +3370,7 @@ class FullTextSearchManySplitsTest(unittest.TestCase):
)
raw_split = RawFullTextSearchSplit([Range(5, 9)])
- def _fake_create(index_type, file_io, index_path, index_io_meta_list):
+ def _fake_create(index_type, file_io, index_path, index_io_meta_list,
max_candidates=None):
class _FakeReader:
def visit_full_text_search(self_inner, fts):
return _completed_future(DictBasedScoredIndexResult({2:
1.0}))
@@ -3402,7 +3402,7 @@ class FullTextSearchManySplitsTest(unittest.TestCase):
self.assertEqual(([5, 6, 8], ["raw paimon", None, "raw full-text"], 5),
build_raw_index.call_args.args)
search_raw.assert_called_once_with(
- b"raw-index", 5, match_query("paimon"), 5)
+ b"raw-index", 5, match_query("paimon"), 5, max_candidates=3)
self.assertEqual([2, 8], sorted(list(result.results())))
self.assertEqual(1.0, result.score_getter()(2))
self.assertEqual(9.0, result.score_getter()(8))
@@ -3428,7 +3428,7 @@ class FullTextSearchManySplitsTest(unittest.TestCase):
_patch_snapshot(self, entries)
def _fake_create(index_type, file_io, index_path,
- index_io_meta_list):
+ index_io_meta_list, max_candidates=None):
row_id = index_io_meta_list[0].file_name
row_id = int(row_id.split("-")[1].split(".")[0])
