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 c0398c9d7e830f0218eb47bc0b5f7c276689b14d Author: chaoyang <[email protected]> AuthorDate: Fri Sep 18 12:13:39 2026 +0800 [python] Support distributed vector search with Ray (#9930) --- .github/workflows/ci-python.yml | 2 +- docs/docs/pypaimon/multimodal-search.md | 64 ++++ .../lumina/lumina_vector_global_index_reader.py | 2 +- .../vindex/vindex_vector_global_index_reader.py | 2 +- paimon-python/pypaimon/multimodal/query.py | 27 +- paimon-python/pypaimon/ray/vector_search.py | 182 +++++++++ .../pypaimon/table/source/vector_search_read.py | 17 +- .../pypaimon/tests/ray_vector_search_test.py | 405 +++++++++++++++++++++ .../tests/vector_metric_consistency_test.py | 23 +- paimon-python/pypaimon/tests/vindex_input_test.py | 11 +- 10 files changed, 721 insertions(+), 14 deletions(-) diff --git a/.github/workflows/ci-python.yml b/.github/workflows/ci-python.yml index f7e127f074..157724ccec 100644 --- a/.github/workflows/ci-python.yml +++ b/.github/workflows/ci-python.yml @@ -381,7 +381,7 @@ jobs: python -m pip install --no-cache-dir -q ray==$ray_version 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" + 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" 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 034162016e..9ec110a42c 100644 --- a/docs/docs/pypaimon/multimodal-search.md +++ b/docs/docs/pypaimon/multimodal-search.md @@ -100,6 +100,70 @@ 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 +by this execution mode. + +Install the same PyPaimon and index dependencies on the driver and workers: + +```shell +python -m pip install 'pypaimon[ray,vindex]' +``` + +```python +import ray + +ray.init() # Use address="auto" to connect to an existing cluster. + +neighbors = ( + docs.search( + [0.1, 0.2, 0.3], + column="embedding", + pre_filter="category = 'lake'", + ) + .select(["id", "content"]) + .limit(10) + .to_arrow( + execution="ray", + concurrency=4, + ray_remote_args={"num_cpus": 1}, + ) +) +``` + +`execution` defaults to `"local"`. `concurrency` is a positive integer limiting +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. + +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. +They return candidate row IDs and scores. The driver retains the local search +rules for global candidate selection, refinement, filtering, and result lookup. +In particular, `pre_filter` filters candidates before ranking; `where()` filters +the selected rows and can return fewer than the requested number of results. + +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 +partial results. Snapshot pinning does not prevent file expiration; the files +needed by the query must remain available for its duration. + +All workers must be able to access the table's storage. Local filesystem paths +are suitable for a local Ray cluster; multiple nodes require shared storage. +Each index task searches one shard, while its native index I/O settings still +apply. Raw-scan parallelism is limited by the number of planned read splits, +controlled by the table's `source.split.target-size` option. Refinement and final +row lookup run on the driver. Candidate traffic grows +with the number of index shards and the configured refinement budget, so Ray +execution is most useful when shard search or raw scanning outweighs scheduling +and transfer costs. Small queries can be faster locally. + ## Search Hybrid Use `search_hybrid` to combine vector and full-text routes, then rerank the diff --git a/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py b/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py index 2abba80617..0ddd1e3325 100644 --- a/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py +++ b/paimon-python/pypaimon/globalindex/lumina/lumina_vector_global_index_reader.py @@ -174,7 +174,7 @@ class LuminaVectorGlobalIndexReader(GlobalIndexReader): self._searcher = LuminaSearcher(self._options) self._searcher.open_stream(stream, self._io_meta.file_size) self._stream = stream - except Exception: + except BaseException: stream.close() raise diff --git a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py index 3a7efb7d83..c3fc51f61a 100644 --- a/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py +++ b/paimon-python/pypaimon/globalindex/vindex/vindex_vector_global_index_reader.py @@ -223,7 +223,7 @@ class VindexVectorGlobalIndexReader(GlobalIndexReader): self._reader = reader self._search_params_type = SearchParams self._stream = stream - except Exception: + except BaseException: try: if reader is not None: reader.close() diff --git a/paimon-python/pypaimon/multimodal/query.py b/paimon-python/pypaimon/multimodal/query.py index 9cfea57e6f..60c3452dc5 100644 --- a/paimon-python/pypaimon/multimodal/query.py +++ b/paimon-python/pypaimon/multimodal/query.py @@ -532,7 +532,32 @@ class VectorQuery(_PreFilterQuery): super().__init__( table, result_factory=self._execute_vector, pre_filter=pre_filter) + def to_arrow(self, *, execution="local", concurrency=None, ray_remote_args=None): + """Execute a vector query locally or on Ray workers, returning an Arrow table. + + Ray execution supports data-evolution tables. ``concurrency`` bounds + in-flight tasks (defaults to 4); ``ray_remote_args`` configures their + resources and retries. Refinement and result lookup run on the driver. + """ + if execution == "local": + if concurrency is not None or ray_remote_args is not None: + raise ValueError("Ray options require execution='ray'.") + return super().to_arrow() + if execution != "ray": + raise ValueError("execution must be 'local' or 'ray'.") + + from pypaimon.ray.vector_search import _execute_vector_search + + query = self._for_execution() + result = _execute_vector_search( + self._vector_search_builder(query), + concurrency=concurrency, ray_remote_args=ray_remote_args) + return query._read_global_index_result(result) + def _execute_vector(self, query): + return self._vector_search_builder(query).execute_local() + + def _vector_search_builder(self, query): limit = query._limit if query._limit is not None else 10 builder = ( query._table.new_vector_search_builder() @@ -543,7 +568,7 @@ class VectorQuery(_PreFilterQuery): ) if query._pre_filter is not None: builder = builder.with_filter(query._pre_filter) - return builder.execute_local() + return builder class TextQuery(_PreFilterQuery): diff --git a/paimon-python/pypaimon/ray/vector_search.py b/paimon-python/pypaimon/ray/vector_search.py new file mode 100644 index 0000000000..348df0f93c --- /dev/null +++ b/paimon-python/pypaimon/ray/vector_search.py @@ -0,0 +1,182 @@ +# 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, single-vector queries.""" + +from contextlib import closing +import math + +from pypaimon.globalindex.vector_search_result import DictBasedScoredIndexResult +from pypaimon.table.source.vector_search_read import ( + DataEvolutionVectorRead, + _filtered_raw_row_ranges, + _offer_score, + _scored_result, +) +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 = 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.") + + 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.") + + try: + import ray # noqa: F401 + except ModuleNotFoundError as error: + if error.name != "ray": + 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.""" + + def __init__(self, reader, concurrency, remote_args): + super().__init__( + reader._table, reader._limit, reader._vector_column, reader._query_vector, + reader._filter, reader._partition_filter, reader._options) + self._concurrency = concurrency + self._remote_args = remote_args + + def _score_raw_batches(self, batches, query_vector, metric, score_candidates=None, reject_nan=True): + # Refinement also runs through this scorer, before its local top-k. + return super()._score_raw_batches(batches, query_vector, metric, score_candidates, reject_nan=True) + + 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) + results = [None] * len(splits) + with closing(_map_tasks( + _search_index_split, context, items, self._concurrency, self._remote_args)) as tasks: + for i, (metric, scores) in tasks: + if metric is not None: + self._set_index_metric(metric) + # Keep plan order, including duplicate-row precedence, regardless + # of worker completion order. Do not refine or truncate per worker. + results[i] = DictBasedScoredIndexResult(scores) + return results + + def _read_raw_search(self, raw_row_ranges, pre_filter, query_vector, + index_type=None, include_filter=True, + score_candidates=None, snapshot=None): + if score_candidates is not None: + # Refinement uses the globally selected indexed candidates. + return super()._read_raw_search( + raw_row_ranges, pre_filter, query_vector, index_type, + include_filter, score_candidates, snapshot) + + ranges = _filtered_raw_row_ranges(raw_row_ranges, pre_filter) + if not ranges: + return DictBasedScoredIndexResult({}) + table_read, splits = self._plan_raw_read(ranges, include_filter, snapshot) + context = (table_read, self._vector_column, query_vector, + self._limit, self._search_metric(index_type)) + heap = [] + # Keep each planned data split intact: its files may be vertical column + # groups that must be read together to reconstruct a logical row. + with closing(_map_tasks( + _search_raw_split, context, splits, self._concurrency, self._remote_args)) as tasks: + for _, scores in tasks: + for row_id, score in scores.items(): + _offer_score(heap, self._limit, row_id, score) + 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) + result = reader._eval( + split.row_range_start, split.row_range_end, + split.vector_index_files, query, limit, include).result() + # The reader closes its native resources before returning a completed result. + # Return metric even for zero hits, so inconsistent indexes cannot be hidden + # by a scalar filter and raw fallback uses the persisted index metric. + return reader._index_metric, _scores(result) + + +def _search_raw_split(context, split): + from pypaimon.read.table_read import _ClosableArrowBatchReader + + table_read, column, query, limit, metric = context + scorer = DataEvolutionVectorRead(table_read.table, limit, column, query) + reader, batches = table_read._new_arrow_batch_reader([split]) + with _ClosableArrowBatchReader(reader, batches) as batch_reader: + return _scores(scorer._score_raw_batches(batch_reader, query, metric, reject_nan=True)) + + +def _scores(result): + if result is None: + return {} + getter = result.score_getter() + scores = {row_id: getter(row_id) for row_id in result.results()} + if any(math.isnan(score) for score in scores.values()): + raise ValueError("Ray vector search cannot rank NaN scores.") + return scores + + +def _map_tasks(worker, context, items, concurrency, remote_args): + """Yield completed tasks with their plan ordinal, bounding in-flight work.""" + import ray + + if not items: + return + # Share only serializable planning state, never open readers or closures. + context_ref = ray.put(context) + remote = ray.remote(worker).options(**remote_args) + remaining = iter(enumerate(items)) + pending = {} + try: + while True: + while len(pending) < concurrency: + item = next(remaining, None) + if item is None: + break + ordinal, split = item + pending[remote.remote(context_ref, split)] = ordinal + if not pending: + break + ready, _ = ray.wait(list(pending), num_returns=1) + ref = ready[0] + result = ray.get(ref) + yield pending.pop(ref), result + finally: + for ref in pending: + try: + ray.cancel(ref) + except Exception: + # Preserve the search failure if the cluster is already gone. + pass diff --git a/paimon-python/pypaimon/table/source/vector_search_read.py b/paimon-python/pypaimon/table/source/vector_search_read.py index 178a7ea907..489022fe1a 100644 --- a/paimon-python/pypaimon/table/source/vector_search_read.py +++ b/paimon-python/pypaimon/table/source/vector_search_read.py @@ -273,7 +273,7 @@ class AbstractVectorSearchReadImpl: try: self._record_index_metric(reader, vector_index_files[0].index_type, metric_lock) return reader, OffsetGlobalIndexReader(reader, row_range_start, row_range_end) - except Exception: + except BaseException: reader.close() raise @@ -314,11 +314,16 @@ class AbstractVectorSearchReadImpl: if table is None or table.num_rows == 0: return DictBasedScoredIndexResult({}) + return self._score_raw_batches( + table.to_batches(max_chunksize=_score_block_size(query_vector)), + query_vector, self._search_metric(index_type), score_candidates) + + def _score_raw_batches(self, batches, query_vector, metric, score_candidates=None, reject_nan=False): + """Score bounded Arrow batches, shared by local and distributed reads.""" + import math + top_k_heap = [] - metric = self._search_metric(index_type) - block_size = _score_block_size(query_vector) - for start in range(0, table.num_rows, block_size): - block = table.slice(start, block_size) + for block in batches: row_ids = block.column(SpecialFields.ROW_ID.name).to_pylist() vectors = block.column(self._vector_column.name) if score_candidates is not None: @@ -332,6 +337,8 @@ class AbstractVectorSearchReadImpl: row_ids, _iter_arrow_scores(vectors, query_vector, metric) ): if score is not None: + if reject_nan and math.isnan(score): + raise ValueError("Ray vector search cannot rank NaN scores.") _offer_score(top_k_heap, self._limit, row_id, score) return _scored_result(top_k_heap) diff --git a/paimon-python/pypaimon/tests/ray_vector_search_test.py b/paimon-python/pypaimon/tests/ray_vector_search_test.py new file mode 100644 index 0000000000..06a7446863 --- /dev/null +++ b/paimon-python/pypaimon/tests/ray_vector_search_test.py @@ -0,0 +1,405 @@ +# 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. + +import os +from contextlib import closing +from unittest.mock import patch + +import numpy as np +import pyarrow as pa +import pytest + +ray = pytest.importorskip("ray") + +import pypaimon.multimodal as pm +from pypaimon.multimodal.query import ScanQuery +from pypaimon.ray import vector_search as search_module +from pypaimon.snapshot.time_travel_util import TimeTravelUtil +from pypaimon.table.source.vector_search_scan import DataEvolutionVectorScan + + +SCHEMA = pa.schema([ + ("id", pa.int64()), ("category", pa.string()), ("embedding", pa.list_(pa.float32(), 2))]) +VECTORS = [[0., 1.], [2., 1.], [1., 3.], [-2., 2.], [4., 1.], [1., -3.]] + + [email protected](scope="module") +def ray_cluster(): + started = not ray.is_initialized() + if started: + ray.init(address="local", num_cpus=2, include_dashboard=False, + object_store_memory=100 * 1024 * 1024) + yield + if started: + ray.shutdown() + + [email protected] +def table(tmp_path): + return pm.connect(options={"warehouse": str(tmp_path)}).create_table( + "vectors", schema=SCHEMA, + options={"file.format": "parquet", "vector.file.format": "parquet", + "vector-index.search-mode": "full", "read.batch-size": "2", + "source.split.target-size": "1 b", "source.split.open-file-cost": "1 b"}) + + +def add_rows(table, vectors, start=0): + ids = list(range(start, start + len(vectors))) + table.add(pa.table({"id": ids, "category": ["yes" if i % 2 == 0 else "no" for i in ids], + "embedding": vectors}, schema=SCHEMA)) + + +def build_index(table, metric="l2"): + pytest.importorskip("paimon_vindex") + # Build before any deletes. The index builder currently requires DVs off; + # subsequent reads/updates use the table's original DV-enabled options. + table.raw_table.copy({"deletion-vectors.enabled": "false"}).create_global_index( + "embedding", "ivf-flat", options={ + "global-index.row-count-per-shard": "3", + "ivf-flat.nlist": "1", "ivf-flat.distance.metric": metric}) + + +def ids(result): + return sorted(result.column("id").to_pylist()) + + [email protected]("metric", ["l2", "cosine", "inner_product"]) [email protected]("concurrency", [1, 2]) +def test_raw_search_matches_exact_top_k(table, ray_cluster, metric, concurrency): + add_rows(table, VECTORS[:3]) + add_rows(table, VECTORS[3:], 3) + query_vector = [1., 1.] + query = table.search(query_vector, options={"metric": metric}).select(["id"]).limit(3) + vectors = np.array(VECTORS) + if metric == "l2": + scores = -np.sum((vectors - query_vector) ** 2, axis=1) + elif metric == "cosine": + scores = vectors.dot(query_vector) / (np.linalg.norm(vectors, axis=1) * np.sqrt(2.)) + else: + scores = vectors.dot(query_vector) + expected = sorted(sorted(range(len(VECTORS)), key=lambda i: (-scores[i], i))[:3]) + assert ids(query.to_arrow()) == expected + original = search_module._map_tasks + dispatched = [] + + def record(worker, context, items, *args): + dispatched.append(len(items)) + return original(worker, context, items, *args) + + with patch.object(search_module, "_map_tasks", record): + assert ids(query.to_arrow(execution="ray", concurrency=concurrency)) == expected + assert dispatched == [2] + + [email protected]("metric", ["l2", "cosine", "inner_product"]) [email protected]("refine", [False, True]) +def test_indexed_search_and_raw_fallback_use_persisted_metric(table, ray_cluster, metric, refine): + add_rows(table, VECTORS) + build_index(table, metric) + add_rows(table, [[1., 1.], [20., 1.]], 6) + # No metric override: raw workers must inherit it from the actual indexes. + options = {"ivf.nprobe": "1"} + if refine: + options["refine_factor"] = "2" + query = table.search([1., 1.], options=options).select(["id"]).limit(3) + expected = ids(query.to_arrow()) + observed = [] + original = search_module._map_tasks + + def record(worker, context, items, *args): + observed.append((worker.__name__, len(items))) + return original(worker, context, items, *args) + + with patch.object(search_module, "_map_tasks", record): + assert ids(query.to_arrow(execution="ray", concurrency=2)) == expected + assert ("_search_index_split", 2) in observed + assert any(name == "_search_raw_split" and count > 0 for name, count in observed) + + +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") + # The scalar filter has no index, so fallback overlaps vector-index coverage. + query = table.search([1., 1.], pre_filter="category = 'yes'").select(["id"]).limit(3) + assert ids(query.to_arrow()) == [4, 6] + assert ids(query.to_arrow(execution="ray", concurrency=2)) == [4, 6] + # Post-filtering keeps its existing meaning: it need not fill the requested K. + query = table.search([1., 1.]).where("id = 4").select(["id"]).limit(1) + assert query.to_arrow().num_rows == 0 + assert query.to_arrow(execution="ray").num_rows == 0 + assert table.search([1., 1.], pre_filter="category = 'missing'").to_arrow(execution="ray").num_rows == 0 + + +def test_partition_filter(tmp_path, ray_cluster): + table = pm.connect(options={"warehouse": str(tmp_path)}).create_table( + "partitioned", schema=SCHEMA, partitioned=["category"], + options={"file.format": "parquet", "vector.file.format": "parquet", + "vector-index.search-mode": "full"}) + add_rows(table, VECTORS) + query = table.search([1., 1.], column="embedding", pre_filter="category = 'yes'").select(["id"]).limit(10) + assert ids(query.to_arrow()) == [0, 2, 4] + assert ids(query.to_arrow(execution="ray")) == [0, 2, 4] + + +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([1., 1.]).select(["id"]).limit(2) + with patch.object(search_module, "_search_raw_split", side_effect=AssertionError("raw scan")): + result = query.to_arrow(execution="ray") + assert ids(result) == ids(query.to_arrow()) + assert 6 not in ids(result) + + [email protected]("change", ["update", "delete"]) +def test_query_pins_worker_reads_and_lookup_but_can_be_reused(table, ray_cluster, change): + add_rows(table, [[1., 1.], [10., 1.]]) + query = table.search([1., 1.]).select(["id", "category"]).limit(1) + expected = query.to_arrow().to_pylist() + original = search_module._map_tasks + changed = [] + + def change_before_dispatch(*args): + if not changed: + if change == "delete": + table.delete("id = 0") + else: + table.update("id = 0", {"category": "changed"}) + changed.append(True) + return original(*args) + + with patch.object(search_module, "_map_tasks", change_before_dispatch): + assert query.to_arrow(execution="ray").to_pylist() == expected + latest = query.to_arrow().to_pylist() + assert latest != expected + assert query.to_arrow(execution="ray").to_pylist() == latest + + +def test_empty_snapshot_does_not_follow_first_commit(table, ray_cluster): + query = table.search([1., 1.]).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")): + result = query.to_arrow(execution="ray") + assert result.schema.names == ["id"] + assert result.num_rows == 0 + assert ids(query.to_arrow(execution="ray")) == [0] + + [email protected]("selector", ["snapshot", "tag"]) +def test_snapshot_and_retained_tag_metadata_reach_workers(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)) + query = table.search([1., 1.], tag_name="saved") + else: + query = table.search([1., 1.], snapshot_id=snapshot.id) + query.select(["id"]).limit(2) + assert ids(query.to_arrow(execution="ray")) == ids(query.to_arrow()) == [0, 1] + + +def test_lookup_is_pinned_after_worker_completion(table, ray_cluster): + add_rows(table, [[1., 1.], [10., 1.]]) + original = ScanQuery._read_global_index_result + + def delete_before_lookup(query, result): + table.delete("id = 0") + return original(query, result) + + with patch.object(ScanQuery, "_read_global_index_result", delete_before_lookup): + assert ids(table.search([1., 1.]).select(["id"]).limit(1).to_arrow(execution="ray")) == [0] + + +def test_warmed_read_context_and_distinct_worker_processes(table, ray_cluster): + add_rows(table, VECTORS[:3]) + add_rows(table, VECTORS[3:], 3) + query = table.search([1., 1.]).select(["id"]).limit(2) + query.to_arrow() # Warm filesystem and metadata state before serialization. + pinned = query._for_execution()._table + + def inspect_context(table, ordinal): + import time + time.sleep(0.2) + return os.getpid(), TimeTravelUtil.resolve_snapshot(table).id + + results = list(search_module._map_tasks(inspect_context, pinned, [0, 1], 2, {"num_cpus": 1})) + assert len({result[0] for _, result in results}) == 2 + assert all(pid != os.getpid() and snapshot == 2 for _, (pid, snapshot) in results) + assert ids(query.to_arrow(execution="ray")) == [0, 1] + + +def test_task_retry_keeps_original_snapshot(table, ray_cluster, tmp_path): + add_rows(table, [[1., 1.]]) + pinned = table.search([1., 1.])._for_execution()._table + add_rows(table, [[2., 1.]], 1) + marker = str(tmp_path / "retry-marker") + + def retry_once(table, marker): + try: + fd = os.open(marker, os.O_CREAT | os.O_EXCL | os.O_WRONLY) + except FileExistsError: + builder = table.new_read_builder().with_projection(["id"]) + data = builder.new_read().to_arrow(builder.new_scan().plan().splits()) + return TimeTravelUtil.resolve_snapshot(table).id, data.column("id").to_pylist() + os.close(fd) + raise RuntimeError("retry this read") + + assert list(search_module._map_tasks( + retry_once, pinned, [marker], 1, {"max_retries": 1, "retry_exceptions": True})) == [(0, (1, [0]))] + + +def test_failure_cancels_outstanding_tasks(ray_cluster): + def fail(context, item): + if item == 0: + raise ValueError("injected search failure") + import time + time.sleep(30) + + with patch.object(ray, "cancel", wraps=ray.cancel) as cancel: + with pytest.raises(ValueError, match="injected search failure"): + list(search_module._map_tasks(fail, None, [0, 1, 2], 2, {})) + assert cancel.call_count == 2 + + +def test_map_tasks_bounds_submissions_and_cancels_on_early_close(ray_cluster): + def echo(context, item): + return item + + wait = ray.wait + counts = [] + + def record_wait(refs, **kwargs): + counts.append(len(refs)) + return wait(refs, **kwargs) + + with patch.object(ray, "wait", record_wait), patch.object(ray, "cancel", wraps=ray.cancel) as cancel: + with closing(search_module._map_tasks(echo, None, list(range(10)), 2, {})) as tasks: + next(tasks) + assert counts == [2] + assert cancel.call_count == 1 + + +def test_refinement_uses_global_candidates_despite_reverse_completion_order(table): + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit + from pypaimon.index.index_file_meta import IndexFileMeta + + add_rows(table, [[0., 1.], [10., 1.], [20., 1.], [1., 1.]]) + query = table.search([1., 1.], options={"refine_factor": "2"}).limit(1)._for_execution() + reader = query._vector_search_builder(query).new_vector_search_read() + distributed = search_module._RayVectorSearchRead(reader, 2, {}) + index_file = IndexFileMeta("ivf-flat", "fake", 1, 2, None) + splits = [IndexVectorSearchSplit(0, 1, [index_file]), IndexVectorSearchSplit(2, 3, [index_file])] + + def completed(*args): + # Row 3 is the true nearest, but is outside the GLOBAL approximate top-2. + # Refining independently per shard would incorrectly bring it back in. + yield 1, ("l2", {2: 8., 3: 7.}) + yield 0, ("l2", {0: 10., 1: 9.}) + + with patch.object(search_module, "_map_tasks", completed): + result = distributed._read_indexed(splits, [1., 1.], query._table._read_snapshot) + assert list(result.results()) == [0] + assert result.score_getter()(0) == 0.5 # L2 score is 1 / (1 + squared distance). + + +def test_duplicate_scores_keep_plan_order(table): + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit + from pypaimon.index.index_file_meta import IndexFileMeta + + query = table.search([1., 1.]).limit(2)._for_execution() + reader = query._vector_search_builder(query).new_vector_search_read() + distributed = search_module._RayVectorSearchRead(reader, 2, {}) + index_file = IndexFileMeta("ivf-flat", "fake", 1, 2) + splits = [IndexVectorSearchSplit(0, 1, [index_file]), IndexVectorSearchSplit(0, 3, [index_file])] + + def completed(*args): + yield 1, ("l2", {0: 100.}) + yield 0, ("l2", {0: 1., 1: 2.}) + + with patch.object(search_module, "_map_tasks", completed): + result = distributed._read_indexed(splits, [1., 1.], None) + assert result.score_getter()(0) == 1. + assert list(result.top_k(1).results()) == [1] + + +def test_metric_mismatch_fails_even_when_one_shard_has_no_hits(table): + from pypaimon.table.source.vector_search_split import IndexVectorSearchSplit + + query = table.search([1., 1.]) + reader = query._vector_search_builder(query).new_vector_search_read() + distributed = search_module._RayVectorSearchRead(reader, 2, {}) + closed = [] + + def completed(*args): + try: + yield 0, ("l2", {}) + yield 1, ("inner_product", {1: 10.}) + finally: + closed.append(True) + + with patch.object(search_module, "_map_tasks", completed), \ + pytest.raises(ValueError, match="different metrics"): + distributed._search_index_splits( + [IndexVectorSearchSplit(0, 1, []), IndexVectorSearchSplit(2, 3, [])], [1., 1.], 1, []) + assert closed == [True] + + [email protected]("kwargs, message", [ + ({"execution": "unknown"}, "execution must"), + ({"concurrency": 2}, "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([1., 1.]).to_arrow(**kwargs) + + +def test_rejects_non_data_evolution_table(table): + query = table.search([1., 1.]) + 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") + + [email protected]("value", [float("nan"), float("inf"), -float("inf")]) +def test_rejects_nonfinite_query_before_dispatch(table, value): + with pytest.raises(ValueError, match="finite query vector"): + table.search([value, 1.]).to_arrow(execution="ray") + + +def test_nan_scores_fail_before_worker_top_k(table, ray_cluster): + add_rows(table, [[1., 1.], [float("nan"), 1.], [10., 1.]]) + # NaN may be discarded by a local heap. Validate before that truncation, + # otherwise a different split layout can silently change the answer. + with pytest.raises(ValueError, match="cannot rank NaN"): + table.search([1., 1.]).limit(1).to_arrow(execution="ray") diff --git a/paimon-python/pypaimon/tests/vector_metric_consistency_test.py b/paimon-python/pypaimon/tests/vector_metric_consistency_test.py index 8b761efada..bdb29735a3 100644 --- a/paimon-python/pypaimon/tests/vector_metric_consistency_test.py +++ b/paimon-python/pypaimon/tests/vector_metric_consistency_test.py @@ -178,22 +178,41 @@ class VectorMetricResolutionTest(unittest.TestCase): def test_reader_closes_on_metadata_and_metric_errors(self): entry = _entry(None, field_id=1, index_type='ivf-flat', file_name='vectors.index', row_range_start=0, row_range_end=1) - for failure in ('metadata', 'query', 'shard'): + for failure in ('metadata', 'query', 'shard', 'interrupt'): with self.subTest(failure=failure): native = Mock(spec=['vector_metric', 'close']) native.vector_metric.return_value = 'inner_product' reader = self._reader({'metric': 'l2'} if failure == 'query' else {}) if failure == 'metadata': native.vector_metric.side_effect = RuntimeError('invalid index metadata') + if failure == 'interrupt': + native.vector_metric.side_effect = KeyboardInterrupt() if failure == 'shard': previous = Mock(spec=['vector_metric']) previous.vector_metric.return_value = 'cosine' reader._record_index_metric(previous, 'ivf-flat') with patch('pypaimon.table.source.vector_search_read._create_vector_reader', return_value=native): - with self.assertRaises((RuntimeError, ValueError)): + with self.assertRaises((RuntimeError, ValueError, KeyboardInterrupt)): reader._open_offset_reader([entry.index_file], 0, 1) native.close.assert_called_once_with() + def test_lumina_open_closes_stream_on_interrupt(self): + from pypaimon.globalindex.global_index_meta import GlobalIndexIOMeta + from pypaimon.globalindex.lumina.lumina_vector_global_index_reader import LuminaVectorGlobalIndexReader + + stream, native, file_io = Mock(), Mock(), Mock() + file_io.new_input_stream.return_value = stream + native.open_stream.side_effect = KeyboardInterrupt() + meta = GlobalIndexIOMeta( + file_name='vectors.index', file_size=1, + metadata=b'{"index.dimension":"1","distance.metric":"l2"}') + with patch.dict('sys.modules', {'lumina_data': Mock(LuminaSearcher=Mock(return_value=native))}): + with self.assertRaises(KeyboardInterrupt): + with LuminaVectorGlobalIndexReader(file_io, '/index', [meta]) as reader: + reader._ensure_loaded() + stream.close.assert_called_once_with() + native.close.assert_called_once_with() + def test_metric_does_not_leak_between_read_calls(self): _install_raw_vector_read_builder(self.table, 'embedding', {0: [2.0], 1: [3.0]}) split = RawVectorSearchSplit([Range(0, 1)], [], 'ivf-flat') diff --git a/paimon-python/pypaimon/tests/vindex_input_test.py b/paimon-python/pypaimon/tests/vindex_input_test.py index 892231462d..8e5935425b 100644 --- a/paimon-python/pypaimon/tests/vindex_input_test.py +++ b/paimon-python/pypaimon/tests/vindex_input_test.py @@ -182,7 +182,12 @@ class VindexInputTest(unittest.TestCase): def test_reader_releases_workers_when_native_close_fails(self): self._check_reader_cleanup("close") - def _check_reader_cleanup(self, phase): + def test_reader_releases_workers_when_initialization_is_interrupted(self): + for phase in ("constructor", "initialize"): + with self.subTest(phase=phase): + self._check_reader_cleanup(phase, KeyboardInterrupt) + + def _check_reader_cleanup(self, phase, error_type=OSError): stream = mock.Mock(spec=["read_at", "close"]) stream.read_at.return_value = b"x" io_ = mock.Mock() @@ -199,7 +204,7 @@ class VindexInputTest(unittest.TestCase): raise error return native - error = OSError("native reader failed") + error = error_type("native reader failed") if phase == "initialize": native.optimize_for_search.side_effect = error if phase == "close": @@ -210,7 +215,7 @@ class VindexInputTest(unittest.TestCase): reader = VindexVectorGlobalIndexReader( io_, "s3://bucket", [GlobalIndexIOMeta(file_name="index", file_size=2)]) if phase in ("constructor", "initialize"): - with self.assertRaises(OSError) as raised: + with self.assertRaises(error_type) as raised: reader._ensure_loaded() self.assertIs(error, raised.exception) else:
