This is an automated email from the ASF dual-hosted git repository. XiaoHongbo-Hope pushed a commit to branch release-2.1 in repository https://gitbox.apache.org/repos/asf/paimon.git
commit 8360a5614d777fe5a043eafde429bd2d48d37eb1 Author: chaoyang <[email protected]> AuthorDate: Sat Sep 19 22:04:57 2026 +0800 [python] Support batch vector search with Ray (#9950) --- .github/workflows/ci-python.yml | 1 + docs/docs/pypaimon/multimodal-search.md | 29 +- paimon-python/pypaimon/multimodal/query.py | 29 +- paimon-python/pypaimon/ray/batch_vector_search.py | 116 +++++++ paimon-python/pypaimon/ray/vector_search.py | 25 +- .../pypaimon/table/source/vector_search_read.py | 29 +- .../pypaimon/tests/ray_batch_vector_search_test.py | 335 +++++++++++++++++++++ 7 files changed, 534 insertions(+), 30 deletions(-) diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index 157724ccec..9d7638bb4b 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -382,6 +382,7 @@ jobs: python -c "import ray; print(f'Ray version: {ray.__version__}')" python -c "from packaging.version import parse; import ray; assert parse(ray.__version__) == parse('$ray_version'), f'Expected Ray $ray_version, got {ray.__version__}'" tests="pypaimon/tests/ray_data_test.py::RayDataTest pypaimon/tests/ray_range_join_test.py::RayRangeJoinTest pypaimon/tests/ray_vector_search_test.py::test_raw_search_matches_exact_top_k" + tests="$tests pypaimon/tests/ray_batch_vector_search_test.py::test_raw_batch_shares_split_tasks" if [ "$ray_version" = "2.50.1" ]; then tests="$tests pypaimon/tests/ray_update_by_row_id_test.py::RayUpdateByRowIdTest::test_empty_dataset_after_transform_is_noop" fi diff --git a/docs/docs/pypaimon/multimodal-search.md b/docs/docs/pypaimon/multimodal-search.md index 73f0280b0e..f7ed76cf01 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -102,9 +102,9 @@ matches = ( ## Distributed Vector Search -Use `execution="ray"` to execute a single vector query across Ray workers and -return an Arrow table to the driver. This supports data-evolution tables; -primary-key tables, batch vector queries, and hybrid queries are not supported +Use `execution="ray"` to execute vector queries across Ray workers and return +Arrow results to the driver. This supports single and batch vector queries on +data-evolution tables; primary-key tables and hybrid queries are not supported by this execution mode. Install the same PyPaimon and index dependencies on the driver and workers: @@ -138,8 +138,27 @@ neighbors = ( the number of in-flight Ray tasks, with a default of 4. `ray_remote_args` supplies Ray task options, including resources and retry settings; `num_returns` is managed by PyPaimon. These two arguments require `execution="ray"`. -Query vectors must contain finite values. If stored vectors produce NaN scores, -Ray execution fails because NaN cannot be ranked consistently across tasks. +Single-vector Ray queries require finite query values and fail if stored +vectors produce NaN scores. + +Batch queries use the same execution options and return one Arrow table per +input vector, in input order: + +```python +batch_neighbors = ( + docs.search_vectors([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]], column="embedding") + .select(["id", "content"]) + .limit(10) + .to_arrow(execution="ray", concurrency=4) +) +``` + +Each batch task handles all query vectors for one split: index workers reuse +an open shard across bounded query blocks, and raw workers stream each data +split once with a separate top-k for each query. Batch refinement and the shared +final row lookup run on the driver. The entire batch uses one read snapshot. +Candidate traffic and result memory grow with the number of query vectors; +split large query collections into smaller batches when necessary. The driver fixes one read snapshot and plans the query. Workers search individual index shards and, when required by the table's search mode, scan unindexed data. diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index da6d712713..99a2a4993d 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -657,9 +657,29 @@ class BatchVectorQuery(_PreFilterQuery): self._vector_options = dict(vector_options or {}) super().__init__(table, pre_filter=pre_filter) - def to_arrow(self): + def to_arrow(self, *, execution="local", concurrency=None, ray_remote_args=None): + """Return one Arrow table per query, optionally searching splits on Ray. + + Ray execution supports data-evolution tables. ``concurrency`` bounds + in-flight tasks (defaults to 4); ``ray_remote_args`` configures their + resources and retries. Batch refinement and shared result lookup run + on the driver, using the same snapshot as all workers. + """ + if execution == "local": + if concurrency is not None or ray_remote_args is not None: + raise ValueError("Ray options require execution='ray'.") + query = self._for_execution() + return query._read_batch_results(query._execute_batch_vector(query)) + if execution != "ray": + raise ValueError("execution must be 'local' or 'ray'.") + + from pypaimon.ray.batch_vector_search import _execute_batch_vector_search + query = self._for_execution() - return query._read_batch_results(query._execute_batch_vector(query)) + results = _execute_batch_vector_search( + self._batch_vector_search_builder(query), + concurrency=concurrency, ray_remote_args=ray_remote_args) + return query._read_batch_results(results) def _read_batch_results(self, results): from pypaimon.globalindex.global_index_result import GlobalIndexResult @@ -703,6 +723,9 @@ class BatchVectorQuery(_PreFilterQuery): return [table.to_pylist() for table in self.to_arrow()] def _execute_batch_vector(self, query): + return self._batch_vector_search_builder(query).execute_batch_local() + + def _batch_vector_search_builder(self, query): limit = query._limit if query._limit is not None else 10 builder = ( query._table.new_batch_vector_search_builder() @@ -713,4 +736,4 @@ class BatchVectorQuery(_PreFilterQuery): ) if query._pre_filter is not None: builder = builder.with_filter(query._pre_filter) - return builder.execute_batch_local() + return builder diff --git a/paimon-python/pypaimon/ray/batch_vector_search.py b/paimon-python/pypaimon/ray/batch_vector_search.py new file mode 100644 index 0000000000..6cf1f6593d --- /dev/null +++ b/paimon-python/pypaimon/ray/batch_vector_search.py @@ -0,0 +1,116 @@ +# 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. + +"""Ray execution of snapshot-pinned batch vector queries.""" + +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.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): + concurrency, remote_args = _execution_options(concurrency, ray_remote_args) + reader = builder.new_batch_vector_search_read() + if (type(reader) is not BatchVectorSearchReadImpl + or not reader._table.options.data_evolution_enabled()): + raise ValueError("Ray vector search supports only data-evolution tables.") + _require_ray() + plan = builder.new_vector_search_scan().scan() + return _RayBatchVectorSearchRead(reader, concurrency, remote_args).read_batch_plan(plan) + + +class _RayBatchVectorSearchRead(BatchVectorSearchReadImpl): + + def __init__(self, reader, concurrency, remote_args): + super().__init__(reader._table, reader._limit, reader._vector_column, reader._query_vectors, + reader._filter, reader._partition_filter, reader._options) + 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) + results = [None] * len(splits) + with closing(_map_tasks( + _search_batch_index_split, context, items, self._concurrency, self._remote_args)) as tasks: + for ordinal, (metric, scores) in tasks: + if metric is not None: + self._set_index_metric(metric) + # Retain plan order for duplicate IDs and global per-query selection. + results[ordinal] = [DictBasedScoredIndexResult(values) for values in scores] + return results + + def _read_raw_batch_search(self, raw_row_ranges, pre_filter, index_type=None, snapshot=None): + heaps = [[] for _ in self._query_vectors] + ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) + if not ranges: + return [_scored_result(heap) for heap in heaps] + table_read, splits = self._plan_raw_read(ranges, True, snapshot) + context = (table_read, self._vector_column, self._query_vectors, + self._limit, self._search_metric(index_type)) + with closing(_map_tasks( + _search_batch_raw_split, context, splits, self._concurrency, self._remote_args)) as tasks: + for _, results in tasks: + for heap, scores in zip(heaps, results): + for row_id, score in scores.items(): + _offer_score(heap, self._limit, row_id, score) + 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) + reader, offset_reader = scorer._open_offset_reader( + split.vector_index_files, split.row_range_start, split.row_range_end) + try: + # Bound native query-matrix scratch to at most 64 queries / 4 MiB + # (or one vector), keeping the shard open across all query blocks. + block_size = max(1, min(64, (1 << 20) // max(1, len(queries[0])))) + results = [] + for start in range(0, len(queries), block_size): + search = BatchVectorSearch(queries[start:start + block_size], limit, column.name, + include_row_ids=include, options=options) + results.extend(_scores(result) for result in offset_reader.visit_batch_vector_search(search).result()) + return scorer._index_metric, results + finally: + reader.close() + + +def _search_batch_raw_split(context, split): + from pypaimon.read.table_read import _ClosableArrowBatchReader + + table_read, column, queries, limit, metric = context + scorer = BatchVectorSearchReadImpl(table_read.table, limit, column, queries) + reader, batches = table_read._new_arrow_batch_reader([split]) + with _ClosableArrowBatchReader(reader, batches) as batch_reader: + return [_scores(result) for result in scorer._score_raw_batch_queries(batch_reader, metric)] + + +def _scores(result): + if result is None: + return {} + getter = result.score_getter() + return {row_id: getter(row_id) for row_id in result.results()} diff --git a/paimon-python/pypaimon/ray/vector_search.py b/paimon-python/pypaimon/ray/vector_search.py index dfd51f1f35..13e0fb8695 100644 --- a/paimon-python/pypaimon/ray/vector_search.py +++ b/paimon-python/pypaimon/ray/vector_search.py @@ -32,20 +32,30 @@ from pypaimon.utils.roaring_bitmap import RoaringBitmap64 def _execute_vector_search(builder, *, concurrency=None, ray_remote_args=None): """Execute the builder of an already snapshot-pinned ``VectorQuery``.""" + concurrency, remote_args = _execution_options(concurrency, ray_remote_args) + reader = builder.new_vector_search_read() + if (type(reader) is not DataEvolutionVectorRead + or not reader._table.options.data_evolution_enabled()): + raise ValueError("Ray vector search supports only data-evolution tables.") + if not all(math.isfinite(float(value)) for value in reader._query_vector): + raise ValueError("Ray vector search requires a finite query vector.") + + _require_ray() + plan = builder.new_vector_search_scan().scan() + return _RayVectorSearchRead(reader, concurrency, remote_args).read_plan(plan) + + +def _execution_options(concurrency, ray_remote_args): concurrency = 4 if concurrency is None else concurrency if isinstance(concurrency, bool) or not isinstance(concurrency, int) or concurrency < 1: raise ValueError("concurrency must be a positive integer.") remote_args = dict(ray_remote_args or {}) if "num_returns" in remote_args: raise ValueError("Vector search manages num_returns; omit it from ray_remote_args.") + return concurrency, remote_args - reader = builder.new_vector_search_read() - if (type(reader) is not DataEvolutionVectorRead - or not reader._table.options.data_evolution_enabled()): - raise ValueError("Ray vector search supports only data-evolution tables.") - if not all(math.isfinite(float(value)) for value in reader._query_vector): - raise ValueError("Ray vector search requires a finite query vector.") +def _require_ray(): try: import ray # noqa: F401 except ModuleNotFoundError as error: @@ -53,9 +63,6 @@ def _execute_vector_search(builder, *, concurrency=None, ray_remote_args=None): raise raise ImportError("Ray vector search requires pypaimon[ray].") from error - plan = builder.new_vector_search_scan().scan() - return _RayVectorSearchRead(reader, concurrency, remote_args).read_plan(plan) - class _RayVectorSearchRead(DataEvolutionVectorRead): """Replace split execution while retaining local merge/refinement semantics.""" diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 489022fe1a..6a0fbf16e1 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -755,22 +755,25 @@ class BatchVectorSearchReadImpl(AbstractVectorSearchReadImpl, def _score_raw_splits(self, table_read, splits, metric): from pypaimon.read.table_read import _ClosableArrowBatchReader - heaps = [[] for _ in self._query_vectors] reader, batches = table_read._new_arrow_batch_reader(splits) # Close the underlying iterator as well if scoring fails mid-batch. with _ClosableArrowBatchReader(reader, batches) as batch_reader: - for batch in batch_reader: - row_ids = batch.column(SpecialFields.ROW_ID.name).to_pylist() - vectors = batch.column(self._vector_column.name) - for start, query_index, scores in _iter_arrow_batch_scores( - vectors, self._query_vectors, metric - ): - heap = heaps[query_index] - block_row_ids = row_ids[start:start + len(scores)] - for row_id, score in zip(block_row_ids, scores): - if score is not None: - _offer_score(heap, self._limit, row_id, score) - del batch, row_ids, vectors + return self._score_raw_batch_queries(batch_reader, metric) + + def _score_raw_batch_queries(self, batches, metric): + heaps = [[] for _ in self._query_vectors] + for batch in batches: + row_ids = batch.column(SpecialFields.ROW_ID.name).to_pylist() + vectors = batch.column(self._vector_column.name) + for start, query_index, scores in _iter_arrow_batch_scores( + vectors, self._query_vectors, metric + ): + heap = heaps[query_index] + block_row_ids = row_ids[start:start + len(scores)] + for row_id, score in zip(block_row_ids, scores): + if score is not None: + _offer_score(heap, self._limit, row_id, score) + del batch, row_ids, vectors return [_scored_result(heap) for heap in heaps] diff --git a/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py new file mode 100644 index 0000000000..78ab13c9b2 --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_batch_vector_search_test.py @@ -0,0 +1,335 @@ +# 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 types import SimpleNamespace +from unittest.mock import Mock, patch + +import numpy as np +import pyarrow as pa +import pytest + +pytest.importorskip("ray") + +import pypaimon.multimodal as pm +from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.index.index_file_meta import IndexFileMeta +from pypaimon.multimodal.query import ScanQuery +from pypaimon.ray import batch_vector_search as search_module +from pypaimon.table.source.vector_search_read import BatchVectorSearchReadImpl +from pypaimon.table.source.vector_search_scan import DataEvolutionVectorScan +from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit +from pypaimon.tests import ray_vector_search_test as fixtures +from pypaimon.tests.ray_vector_search_test import VECTORS, add_rows, build_index, ids + + +QUERIES = [[1., 1.], [-1., 2.], [1., 1.]] +ray_cluster = fixtures.ray_cluster +table = fixtures.table + + [email protected]("metric", ["l2", "cosine", "inner_product"]) [email protected]("concurrency", [1, 2]) +def test_raw_batch_shares_split_tasks(table, ray_cluster, metric, concurrency): + add_rows(table, VECTORS[:3]) + add_rows(table, VECTORS[3:], 3) + query = table.search_vectors(QUERIES, options={"metric": metric}).select(["id"]).limit(3) + expected = [] + vectors = np.array(VECTORS) + for vector in QUERIES: + if metric == "l2": + scores = -np.sum((vectors - vector) ** 2, axis=1) + elif metric == "cosine": + scores = vectors.dot(vector) / (np.linalg.norm(vectors, axis=1) * np.linalg.norm(vector)) + else: + scores = vectors.dot(vector) + expected.append(sorted(sorted(range(len(VECTORS)), key=lambda i: (-scores[i], i))[:3])) + original = search_module._map_tasks + dispatched = [] + + def record(worker, context, items, *args): + if worker is search_module._search_batch_raw_split: + dispatched.append(len(items)) + return original(worker, context, items, *args) + + with patch.object(search_module, "_map_tasks", record): + actual = query.to_arrow(execution="ray", concurrency=concurrency) + assert [ids(result) for result in actual] == expected + assert dispatched == [2] + + [email protected]("metric", ["l2", "cosine", "inner_product"]) [email protected]("refine", [False, True]) +def test_indexed_batch_and_raw_fallback_use_persisted_metric(table, ray_cluster, metric, refine): + add_rows(table, VECTORS) + build_index(table, metric) + add_rows(table, [[1., 1.], [-1., 2.]], 6) + options = {"refine_factor": "2"} if refine else {} + query = table.search_vectors(QUERIES, options=options).select(["id"]).limit(3) + expected = query.to_list() + original = search_module._map_tasks + dispatched = [] + + def record(worker, context, items, *args): + dispatched.append((worker.__name__, len(items))) + return original(worker, context, items, *args) + + with patch.object(search_module, "_map_tasks", record): + actual = query.to_arrow(execution="ray", concurrency=2) + assert [result.to_pylist() for result in actual] == expected + assert ("_search_batch_index_split", 2) in dispatched + assert ("_search_batch_raw_split", 1) in dispatched + + +def test_native_query_blocks_and_shared_final_lookup(table, ray_cluster): + add_rows(table, VECTORS) + build_index(table) + vectors = (QUERIES * 22)[:65] + query = table.search_vectors(vectors).select(["category", "id"]).limit(2) + expected = query.to_list() + lookup = ScanQuery._read_global_index_result + calls = [] + + def read(execution, result): + calls.append(result.results().cardinality()) + return lookup(execution, result) + + with patch.object(ScanQuery, "_read_global_index_result", read): + actual = query.to_arrow(execution="ray") + assert [result.to_pylist() for result in actual] == expected + assert all(result.schema.names == ["category", "id"] for result in actual) + assert len(calls) == 1 + + [email protected]("fail_read", [False, True]) +def test_raw_worker_streams_nullable_batches_and_closes_reader(table, fail_read): + query = table.search_vectors(QUERIES) + reader = query._batch_vector_search_builder(query).new_batch_vector_search_read() + schema = pa.schema([("_ROW_ID", pa.int64()), ("embedding", pa.list_(pa.float32(), 2))]) + closed = [] + + def batches(): + try: + yield pa.record_batch([[0, 1], [[1., 1.], None]], schema=schema) + if fail_read: + raise RuntimeError("injected read failure") + yield pa.record_batch([[2], [[10., 1.]]], schema=schema) + finally: + closed.append(True) + + source = batches() + arrow = pa.RecordBatchReader.from_batches(schema, source) + table_read = SimpleNamespace(table=reader._table, _new_arrow_batch_reader=Mock(return_value=(arrow, source))) + context = (table_read, reader._vector_column, QUERIES, 1, "l2") + if fail_read: + with pytest.raises(RuntimeError, match="injected read failure"): + search_module._search_batch_raw_split(context, "split") + else: + results = search_module._search_batch_raw_split(context, "split") + assert [list(scores) for scores in results] == [[0]] * 3 + table_read._new_arrow_batch_reader.assert_called_once_with(["split"]) + assert closed == [True] + assert source.gi_frame is None + + +def test_filters_deletions_and_column_updates(table, ray_cluster): + add_rows(table, VECTORS) + build_index(table) + add_rows(table, [[1., 1.], [0., 1.]], 6) + table.update("id = 2", {"category": "no"}) + table.delete("id = 0") + query = table.search_vectors(QUERIES, pre_filter="category = 'yes'").select(["id"]).limit(3) + assert [ids(result) for result in query.to_arrow(execution="ray")] == [[4, 6]] * 3 + query = table.search_vectors(QUERIES).where("id = 4").select(["id"]).limit(1) + assert [result.to_pylist() for result in query.to_arrow(execution="ray")] == query.to_list() + assert all(result.num_rows == 0 for result in table.search_vectors( + QUERIES, pre_filter="category = 'missing'").to_arrow(execution="ray")) + + +def test_partition_filter(tmp_path, ray_cluster): + table = pm.connect(options={"warehouse": str(tmp_path)}).create_table( + "partitioned", schema=fixtures.SCHEMA, partitioned=["category"], + options={"file.format": "parquet", "vector.file.format": "parquet", + "vector-index.search-mode": "full"}) + add_rows(table, VECTORS) + query = table.search_vectors(QUERIES, pre_filter="category = 'yes'").select(["id"]).limit(10) + assert [ids(result) for result in query.to_arrow(execution="ray")] == [[0, 2, 4]] * 3 + + +def test_fast_mode_does_not_scan_unindexed_rows(table, ray_cluster): + add_rows(table, VECTORS) + build_index(table) + add_rows(table, [[1., 1.]], 6) + table.raw_table = table.raw_table.copy({"vector-index.search-mode": "fast"}) + query = table.search_vectors(QUERIES).select(["id"]).limit(2) + with patch.object(search_module, "_search_batch_raw_split", side_effect=AssertionError("raw scan")): + actual = query.to_arrow(execution="ray") + assert [result.to_pylist() for result in actual] == query.to_list() + assert all(6 not in ids(result) for result in actual) + + [email protected]("stage", ["dispatch", "lookup"]) +def test_whole_batch_pins_snapshot_and_can_be_reused(table, ray_cluster, stage): + add_rows(table, VECTORS) + query = table.search_vectors(QUERIES).select(["id", "category"]).limit(1) + expected = query.to_list() + target = search_module if stage == "dispatch" else ScanQuery + name = "_map_tasks" if stage == "dispatch" else "_read_global_index_result" + original = getattr(target, name) + changed = [] + + def commit(*args): + if not changed: + table.update("id = 0", {"category": "changed"}) + changed.append(True) + return original(*args) + + with patch.object(target, name, commit): + assert [result.to_pylist() for result in query.to_arrow(execution="ray")] == expected + assert changed == [True] + assert query.to_list() != expected + assert [result.to_pylist() for result in query.to_arrow(execution="ray")] == query.to_list() + + [email protected]("selector", ["snapshot", "tag"]) +def test_snapshot_and_retained_tag(table, ray_cluster, selector): + add_rows(table, VECTORS) + snapshot = table.raw_table.snapshot_manager().get_latest_snapshot() + table.raw_table.create_tag("saved", snapshot_id=snapshot.id) + add_rows(table, [[1., 1.]], 6) + if selector == "tag": + table.raw_table.file_io.delete(table.raw_table.snapshot_manager().get_snapshot_path(snapshot.id)) + options = {"tag_name": "saved"} + else: + options = {"snapshot_id": snapshot.id} + query = table.search_vectors(QUERIES, **options).select(["id"]).limit(2) + actual = query.to_arrow(execution="ray") + assert [result.to_pylist() for result in actual] == query.to_list() + assert all(6 not in ids(result) for result in actual) + + +def test_empty_snapshot_does_not_follow_first_commit(table, ray_cluster): + query = table.search_vectors(QUERIES).select(["id"]) + scan = DataEvolutionVectorScan.scan + + def append_before_scan(search_scan): + add_rows(table, [[1., 1.]]) + return scan(search_scan) + + with patch.object(DataEvolutionVectorScan, "scan", append_before_scan), \ + patch.object(search_module, "_map_tasks", side_effect=AssertionError("empty task")): + results = query.to_arrow(execution="ray") + assert len(results) == 3 + assert all(result.schema.names == ["id"] and result.num_rows == 0 for result in results) + assert [ids(result) for result in query.to_arrow(execution="ray")] == [[0]] * 3 + + +def test_global_candidates_and_duplicate_precedence_ignore_completion_order(table): + add_rows(table, [[0., 1.], [10., 1.], [20., 1.], [1., 1.]]) + query = table.search_vectors([[1., 1.], [20., 1.]], options={"refine_factor": "2"}).limit(1)._for_execution() + reader = query._batch_vector_search_builder(query).new_batch_vector_search_read() + distributed = search_module._RayBatchVectorSearchRead(reader, 2, {}) + index_file = IndexFileMeta("ivf-flat", "fake", 1, 2) + splits = [IndexVectorSearchSplit(0, 1, [index_file]), IndexVectorSearchSplit(0, 3, [index_file])] + + def completed(*args): + # 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. + yield 1, ("l2", [{2: 8., 3: 7.}, {0: 100., 2: 8., 3: 7.}]) + yield 0, ("l2", [{0: 10., 1: 9.}, {0: 1., 1: 9.}]) + + with patch.object(search_module, "_map_tasks", completed): + results = distributed._read_batch(splits, query._table._read_snapshot) + assert [list(result.results()) for result in results] == [[0], [2]] + assert results[0].score_getter()(0) == 0.5 + + +def test_metric_mismatch_in_empty_shard_closes_tasks(table): + query = table.search_vectors(QUERIES) + reader = query._batch_vector_search_builder(query).new_batch_vector_search_read() + distributed = search_module._RayBatchVectorSearchRead(reader, 2, {}) + closed = [] + + def completed(*args): + try: + yield 0, ("l2", [{}, {}, {}]) + yield 1, ("inner_product", [{1: 10.}] * 3) + finally: + closed.append(True) + + with patch.object(search_module, "_map_tasks", completed), \ + pytest.raises(ValueError, match="different metrics"): + distributed._search_index_splits([None, None], QUERIES, 1, [], batch=True) + assert closed == [True] + + [email protected]("fail_second_block", [False, True]) +def test_index_reader_is_reused_and_closed_between_query_blocks(table, fail_second_block): + queries = [[1., 1.]] * 65 + query = table.search_vectors(queries) + reader = query._batch_vector_search_builder(query).new_batch_vector_search_read() + native = Mock() + sizes = [] + + def search(batch): + sizes.append(len(batch.vectors)) + if len(sizes) == 2 and fail_second_block: + raise ValueError("failed query block") + return SimpleNamespace(result=lambda: [DictBasedScoredIndexResult({0: 1.})] * len(batch.vectors)) + + offset = SimpleNamespace(visit_batch_vector_search=search) + split = IndexVectorSearchSplit(0, 1, []) + context = (reader._table, reader._vector_column, queries, 1, {}) + 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)) + else: + _, results = search_module._search_batch_index_split(context, (split, None)) + assert results == [{0: 1.}] * 65 + assert sizes == [64, 1] + opened.assert_called_once() + native.close.assert_called_once() + + [email protected]("kwargs, message", [ + ({"execution": "unknown"}, "execution must"), + ({"concurrency": 2}, "require execution"), + ({"ray_remote_args": {}}, "require execution"), + ({"execution": "ray", "concurrency": 0}, "positive integer"), + ({"execution": "ray", "concurrency": True}, "positive integer"), + ({"execution": "ray", "concurrency": 1.5}, "positive integer"), + ({"execution": "ray", "ray_remote_args": {"num_returns": 2}}, "num_returns"), +]) +def test_invalid_execution_options(table, kwargs, message): + with pytest.raises(ValueError, match=message): + table.search_vectors(QUERIES).to_arrow(**kwargs) + + +def test_ray_remains_optional_for_local_execution(table): + add_rows(table, VECTORS) + query = table.search_vectors(QUERIES).select(["id"]).limit(1) + with patch.dict("sys.modules", {"ray": None}): + assert len(query.to_arrow()) == 3 + with pytest.raises(ImportError, match=r"requires pypaimon\[ray\]"): + query.to_arrow(execution="ray") + + +def test_rejects_non_data_evolution_table(table): + query = table.search_vectors(QUERIES) + with patch.object(query._table.options.__class__, "data_evolution_enabled", return_value=False), \ + pytest.raises(ValueError, match="only data-evolution tables"): + query.to_arrow(execution="ray")
