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 e93ab384de [python] Evaluate Ray vector prefilters on workers (#10446)
e93ab384de is described below
commit e93ab384dee11d690752507a1c78c6b9aa98306b
Author: chaoyang <[email protected]>
AuthorDate: Fri Oct 9 21:57:18 2026 +0800
[python] Evaluate Ray vector prefilters on workers (#10446)
---
docs/docs/pypaimon/multimodal-search.md | 9 ++
paimon-python/pypaimon/ray/batch_vector_search.py | 32 +++--
paimon-python/pypaimon/ray/vector_search.py | 95 +++++++++++---
.../table/source/global_index_live_row_filter.py | 26 ++--
.../pypaimon/table/source/vector_search_read.py | 25 ++--
.../tests/global_index_live_row_filter_test.py | 49 +++++++
.../pypaimon/tests/ray_batch_vector_search_test.py | 10 +-
.../tests/ray_vector_filter_exactness_test.py | 143 +++++++++++++++++++--
.../pypaimon/tests/ray_vector_search_test.py | 6 +-
9 files changed, 330 insertions(+), 65 deletions(-)
diff --git a/docs/docs/pypaimon/multimodal-search.md
b/docs/docs/pypaimon/multimodal-search.md
index 47b4f8698f..b317e75180 100644
--- a/docs/docs/pypaimon/multimodal-search.md
+++ b/docs/docs/pypaimon/multimodal-search.md
@@ -248,6 +248,15 @@ filtering and final row lookup.
In particular, `pre_filter` filters candidates before ranking; `where()`
filters
the selected rows and can return fewer than the requested number of results.
+For indexed vector searches, workers evaluate scalar predicates and build
+live-row filters at the query snapshot. Tasks using the same scalar index files
+share one Ray scalar-index evaluation; its result is passed directly to
dependent
+tasks without materializing the bitmap on the driver. Candidate verification
+reads only filter columns and row IDs in the task's shard range, without
scanning
+the scalar indexes again. The final shard filter is reused across a batch's
query
+vectors. This lets expensive verification scale with search concurrency;
partially
+overlapping scalar index inputs can still cause repeated index reads.
+
Index search, raw scans, refinement, lookup, and task retries use the same read
snapshot. The existing `snapshot_id` and `tag_name` arguments to `search()`
also
work with Ray execution. A failed task fails the query rather than returning
diff --git a/paimon-python/pypaimon/ray/batch_vector_search.py
b/paimon-python/pypaimon/ray/batch_vector_search.py
index b7b4303a6f..1a226020cb 100644
--- a/paimon-python/pypaimon/ray/batch_vector_search.py
+++ b/paimon-python/pypaimon/ray/batch_vector_search.py
@@ -21,11 +21,12 @@ from contextlib import closing
from pypaimon.globalindex.batch_vector_search import BatchVectorSearch
from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
-from pypaimon.ray.vector_search import _execution_options, _map_tasks,
_require_ray
+from pypaimon.ray.vector_search import (
+ _execution_options, _map_tasks, _require_ray, _shared_scalar_filters,
_worker_pre_filter,
+)
from pypaimon.table.source.vector_search_read import (
BatchVectorSearchReadImpl, _filtered_raw_row_ranges, _offer_score,
_scored_result,
)
-from pypaimon.utils.roaring_bitmap import RoaringBitmap64
def _execute_batch_vector_search(builder, *, concurrency=None,
ray_remote_args=None):
@@ -47,12 +48,14 @@ class _RayBatchVectorSearchRead(BatchVectorSearchReadImpl):
self._concurrency = concurrency
self._remote_args = remote_args
- def _search_index_splits(self, splits, query, search_limit, pre_filters,
batch=False):
- items = [(split, None if not pre_filters or pre_filters[i] is None
- else pre_filters[i].serialize()) for i, split in
enumerate(splits)]
- context = (self._table, self._vector_column, query, search_limit,
self._options)
- with closing(_map_tasks(
- _search_batch_index_split, context, items, self._concurrency,
self._remote_args, True)) as tasks:
+ def _search_index_splits(self, splits, query, search_limit, snapshot,
batch=False):
+ # Send predicates and metadata, leaving each worker to build and
consume
+ # its own filter. Batch queries reuse the same filter for every vector.
+ context = (self._table, self._vector_column, query, search_limit,
self._options,
+ self._filter, self._partition_filter, snapshot)
+ with _shared_scalar_filters(splits, self._filter, self._remote_args)
as prepare, closing(_map_tasks(
+ _search_batch_index_split, context, splits, self._concurrency,
+ self._remote_args, True, prepare)) as tasks:
for _, (metric, scores) in tasks:
if metric is not None:
self._set_index_metric(metric)
@@ -97,11 +100,14 @@ class _RayBatchVectorSearchRead(BatchVectorSearchReadImpl):
return [_scored_result(heap) for heap in heaps]
-def _search_batch_index_split(context, item):
- table, column, queries, limit, options = context
- split, include_bytes = item
- scorer = BatchVectorSearchReadImpl(table, limit, column, queries,
options=options)
- include = None if include_bytes is None else
RoaringBitmap64.deserialize(include_bytes)
+def _search_batch_index_split(context, split, scalar_filter=None):
+ table, column, queries, limit, options, filter_, partition_filter,
snapshot = context
+ # Explicit shard indexes are evaluated below. Candidate verification must
+ # not rediscover and scan table-wide scalar indexes in every worker.
+ table = table.copy({"global-index.enabled": "false"})
+ scorer = BatchVectorSearchReadImpl(
+ table, limit, column, queries, filter_, partition_filter, options)
+ include = _worker_pre_filter(scorer, split, snapshot, scalar_filter)
reader, offset_reader = scorer._open_offset_reader(
split.vector_index_files, split.row_range_start, split.row_range_end)
try:
diff --git a/paimon-python/pypaimon/ray/vector_search.py
b/paimon-python/pypaimon/ray/vector_search.py
index 9e99425827..0b592ec887 100644
--- a/paimon-python/pypaimon/ray/vector_search.py
+++ b/paimon-python/pypaimon/ray/vector_search.py
@@ -17,9 +17,11 @@
"""Ray execution of snapshot-pinned, single-vector queries."""
-from contextlib import closing
+from collections import Counter
+from contextlib import closing, contextmanager
import math
+from pypaimon.globalindex.global_index_result import GlobalIndexResult
from pypaimon.globalindex.vector_search_result import
DictBasedScoredIndexResult
from pypaimon.table.source.vector_search_read import (
DataEvolutionVectorRead,
@@ -27,6 +29,8 @@ from pypaimon.table.source.vector_search_read import (
_offer_score,
_scored_result,
)
+from pypaimon.table.source import global_index_live_row_filter
+from pypaimon.utils.range import Range
from pypaimon.utils.roaring_bitmap import RoaringBitmap64
@@ -74,14 +78,13 @@ class _RayVectorSearchRead(DataEvolutionVectorRead):
self._concurrency = concurrency
self._remote_args = remote_args
- def _search_index_splits(self, splits, query, search_limit, pre_filters,
batch=False):
- # Filters are planned once on the driver, at the query snapshot. Each
- # worker receives only its own include-row bitmap, not a table-wide
set.
- items = [(split, None if not pre_filters or pre_filters[i] is None
- else pre_filters[i].serialize()) for i, split in
enumerate(splits)]
- context = (self._table, self._vector_column, query, search_limit,
self._options)
- with closing(_map_tasks(
- _search_index_split, context, items, self._concurrency,
self._remote_args, True)) as tasks:
+ def _search_index_splits(self, splits, query, search_limit, snapshot,
batch=False):
+ # Send predicates and metadata, leaving each worker to build and
consume
+ # its own filter. Batch queries reuse the same filter for every vector.
+ context = (self._table, self._vector_column, query, search_limit,
self._options,
+ self._filter, self._partition_filter, snapshot)
+ with _shared_scalar_filters(splits, self._filter, self._remote_args)
as prepare, closing(_map_tasks(
+ _search_index_split, context, splits, self._concurrency,
self._remote_args, True, prepare)) as tasks:
for _, (metric, scores) in tasks:
if metric is not None:
self._set_index_metric(metric)
@@ -114,11 +117,14 @@ class _RayVectorSearchRead(DataEvolutionVectorRead):
return _scored_result(heap)
-def _search_index_split(context, item):
- table, column, query, limit, options = context
- split, include_bytes = item
- reader = DataEvolutionVectorRead(table, limit, column, query,
options=options)
- include = None if include_bytes is None else
RoaringBitmap64.deserialize(include_bytes)
+def _search_index_split(context, split, scalar_filter=None):
+ table, column, query, limit, options, filter_, partition_filter, snapshot
= context
+ # Explicit shard indexes are evaluated below. Candidate verification must
+ # not rediscover and scan table-wide scalar indexes in every worker.
+ table = table.copy({"global-index.enabled": "false"})
+ reader = DataEvolutionVectorRead(
+ table, limit, column, query, filter_, partition_filter, options)
+ include = _worker_pre_filter(reader, split, snapshot, scalar_filter)
result = reader._eval(
split.row_range_start, split.row_range_end,
split.vector_index_files, query, limit, include).result()
@@ -128,6 +134,62 @@ def _search_index_split(context, item):
return reader._index_metric, _scores(result)
+def _worker_pre_filter(reader, split, snapshot, scalar_filter):
+ live_rows = global_index_live_row_filter.live_rows(
+ reader._table, reader._partition_filter, snapshot,
+ row_ranges=[Range(split.row_range_start, split.row_range_end)])
+ if scalar_filter is None:
+ matched = reader._scalar_matched_rows([split], snapshot)
+ else:
+ bitmap, exact = scalar_filter
+ result = None if bitmap is None else GlobalIndexResult.create(
+ RoaringBitmap64.deserialize(bitmap), is_exact=exact)
+ matched = reader._refine_scalar_result([split], result, snapshot)
+ filters = reader._combine_pre_filters([split], live_rows, matched)
+ return filters[0] if filters else None
+
+
+def _prepare_scalar_filter(context, split):
+ table, column, query, limit, options, filter_, partition_filter, snapshot
= context
+ reader = DataEvolutionVectorRead(table, limit, column, query, filter_,
partition_filter, options)
+ result = reader._scalar_index_result([split], snapshot)
+ return (None, False) if result is None else (result.results().serialize(),
result.is_exact())
+
+
+@contextmanager
+def _shared_scalar_filters(splits, predicate, remote_args):
+ """Evaluate identical scalar-index inputs once, without fetching masks on
the driver."""
+ import ray
+
+ def key(split):
+ return tuple(sorted(f.file_name for f in split.scalar_index_files))
+
+ counts = Counter(key(split) for split in splits) if predicate is not None
else {}
+ shared = {files for files, count in counts.items() if files and count > 1}
+ pending = {}
+ remote = ray.remote(_prepare_scalar_filter).options(**remote_args) if
shared else None
+
+ def prepare(context_ref, split):
+ files = key(split) if shared else None
+ if files not in shared:
+ return split, None
+ if files not in pending:
+ pending[files] = remote.remote(context_ref, split)
+ # A top-level ObjectRef dependency is resolved before allocating the
+ # search worker, so waiting searches cannot starve their filter tasks.
+ return split, pending[files]
+
+ try:
+ yield prepare
+ finally:
+ for ref in pending.values():
+ try:
+ ray.cancel(ref)
+ except Exception:
+ # Preserve the search failure if the cluster is already gone.
+ pass
+
+
def _search_raw_split(context, split):
from pypaimon.read.table_read import _ClosableArrowBatchReader
@@ -148,7 +210,7 @@ def _scores(result):
return scores
-def _map_tasks(worker, context, items, concurrency, remote_args,
ordered=False):
+def _map_tasks(worker, context, items, concurrency, remote_args,
ordered=False, prepare_item=None):
"""Bound in-flight work and, optionally, completed results awaiting plan
order."""
import ray
@@ -168,7 +230,8 @@ def _map_tasks(worker, context, items, concurrency,
remote_args, ordered=False):
if item is None:
break
ordinal, split = item
- pending[remote.remote(context_ref, split)] = ordinal
+ args = (split,) if prepare_item is None else
prepare_item(context_ref, split)
+ pending[remote.remote(context_ref, *args)] = ordinal
if not pending:
break
ready, _ = ray.wait(list(pending), num_returns=1)
diff --git
a/paimon-python/pypaimon/table/source/global_index_live_row_filter.py
b/paimon-python/pypaimon/table/source/global_index_live_row_filter.py
index fe59df49c2..7e8565f9e2 100644
--- a/paimon-python/pypaimon/table/source/global_index_live_row_filter.py
+++ b/paimon-python/pypaimon/table/source/global_index_live_row_filter.py
@@ -20,17 +20,19 @@
from typing import Optional
from pypaimon.deletionvectors.deletion_vector import DeletionVector
+from pypaimon.globalindex.indexed_split import IndexedSplit
from pypaimon.read.query_auth_split import QueryAuthSplit
from pypaimon.read.split import DataSplit
from pypaimon.utils.range import Range
from pypaimon.utils.roaring_bitmap import RoaringBitmap64
-def live_rows(table, partition_filter=None, snapshot=None) ->
Optional[RoaringBitmap64]:
+def live_rows(table, partition_filter=None, snapshot=None, row_ranges=None) ->
Optional[RoaringBitmap64]:
"""Return live global row ids at ``snapshot`` for deletion-vector tables.
``None`` means no live-row filter is needed. This keeps tables without
deletion vectors on the old zero-overhead path.
+ ``row_ranges`` restricts both file planning and the returned live-row
bitmap.
"""
options = getattr(table, "options", None)
deletion_vectors_enabled = getattr(options, "deletion_vectors_enabled",
None)
@@ -44,16 +46,20 @@ def live_rows(table, partition_filter=None, snapshot=None)
-> Optional[RoaringBi
if partition_filter is not None:
read_builder = read_builder.with_partition_filter(partition_filter)
- data_splits = [
- (split.split if isinstance(split, QueryAuthSplit) else split)
- for split in read_builder.new_scan().plan().splits()
- ]
- data_splits = [s for s in data_splits if isinstance(s, DataSplit)]
+ scan = read_builder.new_scan()
+ if row_ranges is not None:
+ scan = scan.with_row_ranges(row_ranges)
+ data_splits = []
+ for split in scan.plan().splits():
+ while isinstance(split, (QueryAuthSplit, IndexedSplit)):
+ split = split.split if isinstance(split, QueryAuthSplit) else
split.data_split()
+ if isinstance(split, DataSplit):
+ data_splits.append(split)
rows = RoaringBitmap64()
# Phase 1: union every file's row-id range.
for split in data_splits:
- _add_row_ranges(rows, split)
+ _add_row_ranges(rows, split, row_ranges)
# Phase 2: subtract each DV. Ranges are all unioned first (no re-add), and
# peak memory stays at one DV at a time.
for split in data_splits:
@@ -81,11 +87,13 @@ def for_range(live_row_ids: Optional[RoaringBitmap64],
return None if include.cardinality() == row_range.count() else include
-def _add_row_ranges(rows: RoaringBitmap64, split: DataSplit) -> None:
+def _add_row_ranges(rows: RoaringBitmap64, split: DataSplit, row_ranges=None)
-> None:
for data_file in split.files:
row_id_range = data_file.row_id_range()
if row_id_range is not None:
- rows.add_range(row_id_range.from_, row_id_range.to)
+ ranges = [row_id_range] if row_ranges is None else
Range.and_([row_id_range], row_ranges)
+ for r in ranges:
+ rows.add_range(r.from_, r.to)
def _subtract_deleted_rows(table, rows: RoaringBitmap64, split: DataSplit) ->
None:
diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py
b/paimon-python/pypaimon/table/source/vector_search_read.py
index 0df1fec983..f90d0a4d93 100644
--- a/paimon-python/pypaimon/table/source/vector_search_read.py
+++ b/paimon-python/pypaimon/table/source/vector_search_read.py
@@ -133,6 +133,10 @@ class AbstractVectorSearchReadImpl:
live_rows = global_index_live_row_filter.live_rows(
self._table, self._partition_filter, snapshot)
matched_rows = self._scalar_matched_rows(splits, snapshot)
+ return self._combine_pre_filters(splits, live_rows, matched_rows)
+
+ @staticmethod
+ def _combine_pre_filters(splits, live_rows, matched_rows):
if live_rows is None and matched_rows is None:
return []
@@ -158,6 +162,10 @@ class AbstractVectorSearchReadImpl:
if self._filter is None:
return None
+ return self._refine_scalar_result(splits,
self._scalar_index_result(splits, snapshot), snapshot)
+
+ def _scalar_index_result(self, splits, snapshot=None):
+ """Evaluate scalar indexes without reading candidate rows."""
# Collect scalar index files across splits, deduplicated by file name.
seen = set()
scalar_files = []
@@ -169,7 +177,7 @@ class AbstractVectorSearchReadImpl:
scalar_files.append(index_file)
if not scalar_files:
- return RoaringBitmap64()
+ return GlobalIndexResult.create_empty()
from pypaimon.globalindex.data_evolution_global_index_scanner import
DataEvolutionGlobalIndexScanner
scanner = DataEvolutionGlobalIndexScanner.create(
@@ -179,11 +187,13 @@ class AbstractVectorSearchReadImpl:
snapshot=snapshot,
)
if scanner is None:
- return RoaringBitmap64()
+ return GlobalIndexResult.create_empty()
try:
- result = scanner.scan(self._filter)
+ return scanner.scan(self._filter)
finally:
scanner.close()
+
+ def _refine_scalar_result(self, splits, result, snapshot=None):
if result is not None and result.is_exact():
return result.results()
if not self._table.options.global_index_filter_refine_from_data():
@@ -489,7 +499,8 @@ class AbstractVectorSearchReadImpl:
future.add_done_callback(lambda _: reader.close())
return future
- def _search_index_splits(self, splits, query, search_limit, pre_filters,
batch=False):
+ def _search_index_splits(self, splits, query, search_limit, snapshot,
batch=False):
+ pre_filters = self._pre_filters(splits, snapshot)
# Native readers finish their search before returning a completed
Future.
# Schedule the entire open/search/close operation, not just
Future.result().
option = CoreOptions.GLOBAL_INDEX_THREAD_NUM
@@ -681,10 +692,9 @@ class
DataEvolutionVectorRead(AbstractVectorSearchReadImpl, VectorSearchRead):
def _read_indexed(self, splits, query_vector, snapshot):
index_type = _vector_index_type(splits)
search_limit = self._indexed_search_limit(index_type)
- pre_filters = self._pre_filters(splits, snapshot)
merged_scores = {}
with closing(self._search_index_splits(
- splits, query_vector, search_limit, pre_filters)) as results:
+ splits, query_vector, search_limit, snapshot)) as results:
for split_result in results:
_merge_index_scores(merged_scores, split_result)
del split_result
@@ -717,11 +727,10 @@ class
BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl,
# passing that split's pre-filter. Each future returns n per-query
results.
index_type = _vector_index_type(index_splits)
search_limit = self._indexed_search_limit(index_type)
- pre_filters = self._pre_filters(index_splits, snapshot)
# Merge each query vector's indexed results across index splits.
merged_scores = [{} for _ in range(n)]
with closing(self._search_index_splits(
- index_splits, self._query_vectors, search_limit, pre_filters,
batch=True)) as results:
+ index_splits, self._query_vectors, search_limit, snapshot,
batch=True)) as results:
for split_results in results:
for i in range(n):
_merge_index_scores(merged_scores[i], split_results[i])
diff --git a/paimon-python/pypaimon/tests/global_index_live_row_filter_test.py
b/paimon-python/pypaimon/tests/global_index_live_row_filter_test.py
new file mode 100644
index 0000000000..4348cfed71
--- /dev/null
+++ b/paimon-python/pypaimon/tests/global_index_live_row_filter_test.py
@@ -0,0 +1,49 @@
+# 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 unittest.mock import patch
+
+import pyarrow as pa
+
+import pypaimon.multimodal as pm
+from pypaimon.deletionvectors.deletion_vector import DeletionVector
+from pypaimon.table.source.global_index_live_row_filter import live_rows
+from pypaimon.utils.range import Range
+
+
+def test_scoped_live_rows_only_load_relevant_deletions(tmp_path):
+ table = pm.connect(options={"warehouse": str(tmp_path)}).create_table(
+ "rows", schema=pa.schema([("id", pa.int64())]),
options={"file.format": "parquet"})
+ table.add(pa.table({"id": [0, 1, 2]}))
+ table.add(pa.table({"id": [3, 4, 5]}))
+ snapshot = table.raw_table.snapshot_manager().get_latest_snapshot()
+ table.delete("id = 0 OR id = 3")
+ with patch.object(DeletionVector, "read", wraps=DeletionVector.read) as
read:
+ assert list(live_rows(table.raw_table)) == [1, 2, 4, 5]
+ all_reads = read.call_count
+ read.reset_mock()
+ # These bounds lie inside a file range and produce IndexedSplit
wrappers.
+ assert list(live_rows(table.raw_table, row_ranges=[Range(0, 1)])) ==
[1]
+ assert 0 < read.call_count < all_reads
+ read.reset_mock()
+ assert list(live_rows(table.raw_table, row_ranges=[Range(3, 4)])) ==
[4]
+ assert 0 < read.call_count < all_reads
+ read.reset_mock()
+ assert list(live_rows(table.raw_table, snapshot=snapshot,
row_ranges=[Range(0, 1)])) == [0, 1]
+ read.assert_not_called()
+ assert list(live_rows(table.raw_table, row_ranges=[])) == []
+ read.assert_not_called()
diff --git a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
index a28c82f90b..05a8e1d6a8 100644
--- a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
+++ b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py
@@ -252,7 +252,7 @@ def test_global_candidates_and_duplicate_precedence(table):
return
# First query must exclude exact nearest row 3; second query has its
own candidates.
# Duplicate row 0 must retain the earlier split's score for the second
query.
- assert args[-1] is True
+ assert args[2] is True
yield 0, ("l2", [{0: 10., 1: 9.}, {0: 1., 1: 9.}])
yield 1, ("l2", [{2: 8., 3: 7.}, {0: 100., 2: 8., 3: 7.}])
@@ -277,7 +277,7 @@ def test_metric_mismatch_in_empty_shard_closes_tasks(table):
with patch.object(search_module, "_map_tasks", completed), \
pytest.raises(ValueError, match="different metrics"):
- list(distributed._search_index_splits([None, None], QUERIES, 1, [],
batch=True))
+ list(distributed._search_index_splits([None, None], QUERIES, 1, None,
batch=True))
assert closed == [True]
@@ -297,13 +297,13 @@ def
test_index_reader_is_reused_and_closed_between_query_blocks(table, fail_seco
offset = SimpleNamespace(visit_batch_vector_search=search)
split = IndexVectorSearchSplit(0, 1, [])
- context = (reader._table, reader._vector_column, queries, 1, {})
+ context = (reader._table, reader._vector_column, queries, 1, {}, None,
None, None)
with patch.object(BatchVectorSearchReadImpl, "_open_offset_reader",
return_value=(native, offset)) as opened:
if fail_second_block:
with pytest.raises(ValueError, match="failed query block"):
- search_module._search_batch_index_split(context, (split, None))
+ search_module._search_batch_index_split(context, split)
else:
- _, results = search_module._search_batch_index_split(context,
(split, None))
+ _, results = search_module._search_batch_index_split(context,
split)
assert results == [{0: 1.}] * 65
assert sizes == [64, 1]
opened.assert_called_once()
diff --git a/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py
b/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py
index eae220840a..983ba96ef8 100644
--- a/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py
+++ b/paimon-python/pypaimon/tests/ray_vector_filter_exactness_test.py
@@ -38,23 +38,144 @@ def
test_ray_applies_exact_row_filter_before_worker_top_k(table, ray_cluster, ba
scalar_index(table)
table.raw_table = table.raw_table.copy({
"vector-index.search-mode": mode,
"global-index.filter.refine-from-data": str(refine).lower()})
- original = AbstractVectorSearchReadImpl._matching_candidate_rows
- arrow_read = TableRead._new_arrow_batch_reader
+ from pypaimon.ray import batch_vector_search, vector_search
+ from pypaimon.table.source import global_index_live_row_filter
+
+ module = batch_vector_search if batch else vector_search
+ original_map = module._map_tasks
calls = []
- def no_vectors(read, *args, **kwargs):
- assert "embedding" not in [field.name for field in read.read_type]
- return arrow_read(read, *args, **kwargs)
+ def record(worker, context, items, *args):
+ def checked(context, split, scalar_filter=None):
+ # These patches run in the actual Ray worker, not on the driver.
+ original = AbstractVectorSearchReadImpl._matching_candidate_rows
+ arrow_read = TableRead._new_arrow_batch_reader
+ observed = []
+
+ def no_vectors(read, *args, **kwargs):
+ assert "embedding" not in [field.name for field in
read.read_type]
+ return arrow_read(read, *args, **kwargs)
+
+ def verify(reader, candidates, snapshot):
+ observed.append(list(candidates))
+ assert snapshot is not None
+ assert all(split.row_range_start <= row_id <=
split.row_range_end for row_id in candidates)
+ from pypaimon.globalindex.data_evolution_global_index_scanner
import DataEvolutionGlobalIndexScanner
+ with patch.object(TableRead, "_new_arrow_batch_reader",
no_vectors), \
+ patch.object(DataEvolutionGlobalIndexScanner, "create",
+
wraps=DataEvolutionGlobalIndexScanner.create) as scan:
+ result = original(reader, candidates, snapshot)
+ scan.assert_not_called()
+ return result
- def verify(reader, candidates, snapshot):
- calls.append(list(candidates))
- assert snapshot is not None
- with patch.object(TableRead, "_new_arrow_batch_reader", no_vectors):
- return original(reader, candidates, snapshot)
+ with patch.object(AbstractVectorSearchReadImpl,
"_matching_candidate_rows", verify):
+ result = worker(context, split, scalar_filter)
+ return result, observed
- with patch.object(AbstractVectorSearchReadImpl,
"_matching_candidate_rows", verify):
+ for ordinal, (result, observed) in original_map(checked, context,
items, *args):
+ calls.extend(observed)
+ yield ordinal, result
+
+ with patch.object(module, "_map_tasks", record), \
+ patch.object(AbstractVectorSearchReadImpl, "_pre_filters",
+ side_effect=AssertionError("driver must not build
filters")), \
+ patch.object(global_index_live_row_filter, "live_rows",
+ side_effect=AssertionError("driver must not read
deletion vectors")):
result = query(table, "name LIKE '%zeta%'",
batch).to_arrow(execution="ray", concurrency=2)
actual = [value.to_pylist() for value in result] if batch else
result.to_pylist()
expected = [{"id": 1}] if refine else []
assert actual == ([expected, expected] if batch else expected)
+ # The whole batch shares a single verification read in its index task.
assert calls == ([[0, 1, 2]] if refine else [])
+
+
[email protected]("batch", [False, True])
[email protected]("kind", ["btree", "bitmap"])
+def
test_worker_filters_overlapping_scalar_shards_at_planned_snapshot(tmp_path,
ray_cluster, batch, kind):
+ import pyarrow as pa
+ import pypaimon.multimodal as pm
+ from pypaimon.ray import batch_vector_search, vector_search
+
+ pytest.importorskip("paimon_vindex")
+ schema = pa.schema([("id", pa.int64()), ("name", pa.string()),
+ ("embedding", pa.list_(pa.float32(), 2))])
+ table = pm.connect(options={"warehouse": str(tmp_path)}).create_table(
+ "vectors", schema=schema, options={
+ "file.format": "parquet", "vector.file.format": "parquet",
+ "global-index.filter.refine-from-data": "true"})
+ for start in range(0, 12, 3):
+ row_ids = list(range(start, start + 3))
+ table.add(pa.table({"id": row_ids, "name": ["keep" if i % 2 == 0 else
"drop" for i in row_ids],
+ "embedding": [[float(i), 1.] for i in row_ids]},
schema=schema))
+ table.raw_table.create_global_index("embedding", "ivf-flat", options={
+ "global-index.row-count-per-shard": "3", "ivf-flat.nlist": "1"})
+ # Scalar and vector shard boundaries intentionally differ.
+ table.raw_table.copy({"deletion-vectors.enabled":
"false"}).create_global_index(
+ "name", kind, options={"global-index.row-count-per-shard": "5"})
+ table.delete("id = 2")
+ search = table.search_vectors([[0., 1.]] * 2, pre_filter="name LIKE
'%keep%'") if batch else table.search(
+ [0., 1.], pre_filter="name LIKE '%keep%'")
+ search = search.select(["id"]).limit(3)
+ expected = [{"id": 0}, {"id": 4}, {"id": 6}]
+ expected = [expected, expected] if batch else expected
+ assert search.to_list() == expected
+ module = batch_vector_search if batch else vector_search
+ original_map = module._map_tasks
+ dispatched = []
+
+ def commit_before_workers(worker, context, splits, *args):
+ if not dispatched:
+ table.delete("id = 0")
+ dispatched.append(True)
+ return original_map(worker, context, splits, *args)
+
+ with patch.object(module, "_map_tasks", commit_before_workers), \
+ patch.object(AbstractVectorSearchReadImpl, "_pre_filters",
+ side_effect=AssertionError("driver filter")):
+ result = search.to_list(execution="ray", concurrency=2)
+ assert result == expected
+ assert dispatched == [True]
+ assert search.to_list(execution="ray", concurrency=2) == search.to_list()
+ assert search.to_list() != expected
+
+
[email protected]("batch", [False, True])
[email protected]("mode", ["exact", "candidate", "unavailable"])
+def test_shared_scalar_index_is_evaluated_once_on_ray(table, tmp_path,
ray_cluster, batch, mode):
+ import os
+ from pathlib import Path
+
+ import pyarrow as pa
+ from pypaimon.ray import vector_search
+
+ schema = pa.schema([("id", pa.int64()), ("name", pa.string()),
+ ("embedding", pa.list_(pa.float32(), 2))])
+ table.add(pa.table({"id": [3, 4, 5], "name": ["delta zeta", "epsilon",
"zeta"],
+ "embedding": [[0.5, 1.], [3., 1.], [4., 1.]]},
schema=schema))
+ table.raw_table.create_global_index("embedding", "ivf-flat", options={
+ "ivf-flat.nlist": "1", "ivf-flat.distance.metric": "l2"})
+ scalar_index(table, "bitmap" if mode == "exact" else "btree")
+ options = {"global-index.filter.refine-from-data": "true"}
+ if mode == "unavailable":
+ options["btree-index.fallback-scan-max-size"] = "0 b"
+ table.raw_table = table.raw_table.copy(options)
+ marker = str(tmp_path / "scalar-evaluations")
+ original = vector_search._prepare_scalar_filter
+
+ def record(context, split):
+ with open(marker, "a") as output:
+ output.write(str(os.getpid()) + "\n")
+ return original(context, split)
+
+ with patch.object(vector_search, "_prepare_scalar_filter", record), \
+ patch.object(AbstractVectorSearchReadImpl, "_scalar_index_result",
+ side_effect=AssertionError("driver scalar
evaluation")):
+ # Each task requires the whole two-CPU cluster. Pending searches must
+ # not occupy resources while waiting for their shared filter
dependency.
+ result = query(table, "name LIKE '%zeta%'", batch).to_list(
+ execution="ray", concurrency=2, ray_remote_args={"num_cpus": 2})
+ expected = [{"id": 3}]
+ assert result == ([expected, expected] if batch else expected)
+ worker_pids = Path(marker).read_text().splitlines()
+ assert len(worker_pids) == 1
+ assert int(worker_pids[0]) != os.getpid()
diff --git a/paimon-python/pypaimon/tests/ray_vector_search_test.py
b/paimon-python/pypaimon/tests/ray_vector_search_test.py
index ca1b361521..fe560bccdd 100644
--- a/paimon-python/pypaimon/tests/ray_vector_search_test.py
+++ b/paimon-python/pypaimon/tests/ray_vector_search_test.py
@@ -327,7 +327,7 @@ def test_refinement_uses_global_candidates(table):
return
# Row 3 is the true nearest, but is outside the GLOBAL approximate
top-2.
# Refining independently per shard would incorrectly bring it back in.
- assert args[-1] is True # Index tasks request plan-ordered results.
+ assert args[2] is True # Index tasks request plan-ordered results.
yield 0, ("l2", {0: 10., 1: 9.})
yield 1, ("l2", {2: 8., 3: 7.})
@@ -348,7 +348,7 @@ def test_duplicate_scores_keep_plan_order(table):
splits = [IndexVectorSearchSplit(0, 1, [index_file]),
IndexVectorSearchSplit(0, 3, [index_file])]
def completed(*args):
- assert args[-1] is True
+ assert args[5] is True
yield 0, ("l2", {0: 1., 1: 2.})
yield 1, ("l2", {0: 100.})
@@ -376,7 +376,7 @@ def
test_metric_mismatch_fails_even_when_one_shard_has_no_hits(table):
with patch.object(search_module, "_map_tasks", completed), \
pytest.raises(ValueError, match="different metrics"):
list(distributed._search_index_splits(
- [IndexVectorSearchSplit(0, 1, []), IndexVectorSearchSplit(2, 3,
[])], [1., 1.], 1, []))
+ [IndexVectorSearchSplit(0, 1, []), IndexVectorSearchSplit(2, 3,
[])], [1., 1.], 1, None))
assert closed == [True]
