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 6ae5a9e337 [python] Support distributed vector search with Ray (#9930)
6ae5a9e337 is described below
commit 6ae5a9e337653363d4c440b4cfa7084f3e93228e
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:
