From ef54ad17ba98b9cee916b508311000e2c24851fc Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Wed, 29 Jul 2026 13:04:03 +0800 Subject: [PATCH 1/2] feat(search): support distributed batch vector queries --- docs/src/distributed-indexing.md | 26 +- lance_ray/search.py | 319 +++++++++++++++--- tests/test_distributed_indexing.py | 19 ++ tests/test_distributed_search.py | 513 +++++++++++++++++++++++++++-- 4 files changed, 799 insertions(+), 78 deletions(-) diff --git a/docs/src/distributed-indexing.md b/docs/src/distributed-indexing.md index c2146d3c..456906bc 100755 --- a/docs/src/distributed-indexing.md +++ b/docs/src/distributed-indexing.md @@ -211,9 +211,9 @@ The function returns the Lance dataset instance (optimization is applied on stor ### Distributed Vector Search -`vector_search()` - Run vector search with Ray workers and merge the global top-k on the driver. +`vector_search()` - Run single or batch vector search with Ray workers and merge the global top-k on the driver. -The driver opens one fixed dataset version, reads vector index segment metadata once, and plans work by index segment ownership. Indexed worker tasks receive only their assigned `index_segments`, so a segment covering multiple fragments is never split across workers. Fragments not covered by an index can be included as separate flat-search fallback work unless `fast_search=True`; fallback tasks use regular fragment scans and compute vector distances in Lance-Ray. +The driver opens one fixed dataset version, reads vector index segment metadata once, and plans work by index segment ownership. Indexed worker tasks receive only their assigned `index_segments`, so a segment covering multiple fragments is never split across workers. Fragments not covered by an index can be included as separate flat-search fallback work unless `fast_search=True`; fallback tasks use regular fragment scans and compute vector distances in Lance-Ray. For a fixed-size vector column, pass a two-dimensional query `[B, D]` to search a batch. The driver merges candidates independently for each query and returns up to `k` rows per query. #### `vector_search` @@ -245,7 +245,7 @@ def vector_search( | Parameter | Type | Description | |-----------|------|-------------| | `uri` | `str` or `lance.LanceDataset`, optional | Lance dataset object, or its URI. Either `uri` OR (`namespace_impl` + `table_id`) must be provided when using URI mode. If a `LanceDataset` object is provided, namespace parameters are ignored and workers reopen the same dataset URI/version. | -| `nearest` | `dict[str, Any]` | Lance vector search options. Must include `column`, `q`, and `k`. Other Lance nearest options such as `minimum_nprobes`, `maximum_nprobes`, `refine_factor`, and distance range are forwarded to every worker. Lance-Ray raises worker-side `k` to at least `k * oversample_factor` before global merge. | +| `nearest` | `dict[str, Any]` | Lance vector search options. Must include `column`, `q`, and `k`. For fixed-size vector columns, `q` may be one vector `[D]` or a batch `[B, D]`. If `metric` is omitted and a vector index is selected, indexed and fallback workers use the index metric; without an index the default is L2. Index-search options such as `minimum_nprobes`, `maximum_nprobes`, and `refine_factor` are forwarded to indexed workers; `distance_range` is also applied to flat fallback results. Lance-Ray raises worker-side `k` to at least `k * oversample_factor` before global merge. Multivector queries remain single queries and require full index coverage; flat fallback does not implement multivector distance. | | `index_name` | `str`, optional | Vector index name to use. If provided and not found, `vector_search()` raises `ValueError` instead of silently falling back. If omitted, Lance-Ray uses the first vector index covering `nearest["column"]`; if none exists, the search uses flat fallback plans unless `fast_search=True`. | | `columns` | `list[str]` or `dict[str, str]`, optional | Projection passed to the Lance scanner. When a list is provided and `_distance` is missing, Lance-Ray appends `_distance` automatically because the driver needs it for global top-k merge. | | `filter` | `Any`, optional | Filter passed unchanged to every worker scanner. | @@ -264,7 +264,7 @@ def vector_search( #### Return Value -The function returns a `pyarrow.Table` containing the global top-k rows sorted by `_distance`. If `analyze_plan=True`, it returns a `str` containing one Lance scanner analysis section per planned shard. +For a one-dimensional query, the function returns a `pyarrow.Table` containing the global top-k rows sorted by `_distance` and `_rowid`. For a two-dimensional batch query, the first column is a non-null Int32 `query_index`; rows are grouped in input-query order, and each group contains up to `k` rows sorted by `_distance` and `_rowid`. The internal `_rowid` tie-break column is omitted unless requested. If `analyze_plan=True`, the function returns a `str` containing one Lance scanner analysis section per planned shard. ## Examples @@ -393,6 +393,24 @@ results = lr.vector_search( fast_search=False, ) +# Run two queries in one Ray scheduling round. The result is one table whose +# query_index column maps every row back to query_vectors[0] or query_vectors[1]. +query_vectors = [query_vector, another_query_vector] +batch_results = lr.vector_search( + uri="path/to/dataset.lance", + nearest={ + "column": "vector", + "q": query_vectors, + "k": 10, + "minimum_nprobes": 20, + }, + index_name="idx_ivf_flat", + columns=["id", "vector"], + num_workers=8, + oversample_factor=2, + fast_search=False, +) + # Inspect the per-shard Lance scanner plans instead of executing the search. plan = lr.vector_search( uri="path/to/dataset.lance", diff --git a/lance_ray/search.py b/lance_ray/search.py index 895be2e8..bdfbefbd 100644 --- a/lance_ray/search.py +++ b/lance_ray/search.py @@ -124,6 +124,24 @@ def _canonical_index_field_names(field_names: Any) -> set[str]: return canonical_names +def _apply_index_metric_default( + nearest: dict[str, Any], + vector_index: Any | None, +) -> dict[str, Any]: + if ( + vector_index is None + or nearest.get("metric") is not None + or nearest.get("distance_type") is not None + ): + return nearest + + details = _index_value(vector_index, "details", {}) or {} + metric = _index_value(details, "metric_type") + if metric is None: + return nearest + return {**nearest, "metric": str(metric).lower()} + + def _plan_vector_search( *, fragments: list[Any], @@ -253,6 +271,7 @@ def _execute_vector_search_plan( nearest: dict[str, Any], candidate_k: int, analyze_plan: bool, + is_batch_query: bool = False, ) -> pa.Table | _SearchPlanAnalysis: dataset = _load_worker_dataset(pickled_dataset) @@ -264,6 +283,7 @@ def _execute_vector_search_plan( nearest=nearest, candidate_k=candidate_k, analyze_plan=analyze_plan, + is_batch_query=is_batch_query, ) if not _scanner_accepts_index_segments(dataset): @@ -304,6 +324,49 @@ def _scanner_accepts_index_segments(dataset: LanceDataset) -> bool: ) +def _inspect_vector_search_query( + dataset: LanceDataset, + *, + nearest: dict[str, Any], + base_scanner_options: dict[str, Any], + include_row_id: bool, +) -> tuple[bool, pa.Schema]: + probe = dataset.scanner(columns=["_distance"], nearest=nearest) + probe_schema = probe.projected_schema + is_batch_query = ( + probe_schema.names + and probe_schema.names[0] == "query_index" + and pa.types.is_int32(probe_schema.field(0).type) + and not probe_schema.field(0).nullable + ) + + schema_options = dict(base_scanner_options) + schema_options["nearest"] = nearest + schema_options["with_row_id"] = True + result_schema = dataset.scanner(**schema_options).projected_schema + if not include_row_id: + row_id_indices = result_schema.get_all_field_indices("_rowid") + if row_id_indices: + result_schema = result_schema.remove(row_id_indices[-1]) + + return is_batch_query, result_schema + + +def _projection_includes_row_id( + columns: Optional[list[str] | dict[str, str]], + scanner_options: dict[str, Any], +) -> bool: + if scanner_options.get("with_row_id"): + return True + if columns is None: + columns = scanner_options.get("columns") + if isinstance(columns, list): + return "_rowid" in columns + if isinstance(columns, dict): + return "_rowid" in columns + return False + + def _execute_flat_fallback_vector_search_plan( dataset: LanceDataset, *, @@ -312,13 +375,15 @@ def _execute_flat_fallback_vector_search_plan( nearest: dict[str, Any], candidate_k: int, analyze_plan: bool, + is_batch_query: bool, ) -> pa.Table | _SearchPlanAnalysis: vector_column = nearest["column"] + scanner_options = dict(base_scanner_options) vector_scan_column, drop_vector_column = _prepare_fallback_scan_columns( - base_scanner_options, + scanner_options, vector_column, + is_batch_query=is_batch_query, ) - scanner_options = dict(base_scanner_options) scanner_options.pop("fast_search", None) scanner_options["fragments"] = [ dataset.get_fragment(fragment_id) for fragment_id in plan.fragment_ids @@ -336,32 +401,76 @@ def _execute_flat_fallback_vector_search_plan( table = scanner.to_table() if table.num_rows == 0: table = table.append_column("_distance", pa.array([], type=pa.float32())) + if is_batch_query: + table = _add_query_index(table, []) if drop_vector_column and vector_scan_column in table.column_names: table = table.drop_columns([vector_scan_column]) return table - distances = _compute_vector_distances( - table[vector_scan_column], - nearest["q"], - _get_nearest_metric(nearest), + valid_vectors = pc.invert(pc.is_null(table[vector_scan_column])) + table = table.filter(valid_vectors) + if table.num_rows == 0: + table = table.append_column("_distance", pa.array([], type=pa.float32())) + if is_batch_query: + table = _add_query_index(table, []) + if drop_vector_column and vector_scan_column in table.column_names: + table = table.drop_columns([vector_scan_column]) + return table + + metric = _get_nearest_metric(nearest) + vector_matrix = _vector_column_to_numpy(table[vector_scan_column], metric) + query_vectors = _query_vectors_to_numpy(nearest["q"], is_batch_query, metric) + import numpy as np + + query_results = [] + for query_index, query_vector in enumerate(query_vectors): + distances = _compute_vector_distances( + vector_matrix, + query_vector, + metric, + ) + finite_distances = np.isfinite(distances) + query_result = table.filter(pa.array(finite_distances, type=pa.bool_())) + distances = distances[finite_distances] + query_result = query_result.append_column( + "_distance", pa.array(distances, type=pa.float32()) + ) + query_result = _apply_distance_range(query_result, nearest) + query_result = _take_top_k(query_result, candidate_k) + if drop_vector_column and vector_scan_column in query_result.column_names: + query_result = query_result.drop_columns([vector_scan_column]) + if is_batch_query: + query_result = _add_query_index( + query_result, + [query_index] * query_result.num_rows, + ) + query_results.append(query_result) + + table = ( + pa.concat_tables(query_results, promote_options="default") + if is_batch_query + else query_results[0] ) - table = table.append_column("_distance", pa.array(distances, type=pa.float32())) - table = _take_top_k(table, candidate_k) - if drop_vector_column and vector_scan_column in table.column_names: - table = table.drop_columns([vector_scan_column]) return table def _prepare_fallback_scan_columns( scanner_options: dict[str, Any], vector_column: str, + *, + is_batch_query: bool, ) -> tuple[str, bool]: requested_columns = scanner_options.get("columns") if requested_columns is None: return vector_column, False if isinstance(requested_columns, list): - scan_columns = [column for column in requested_columns if column != "_distance"] + virtual_columns = {"_distance"} + if is_batch_query: + virtual_columns.add("query_index") + scan_columns = [ + column for column in requested_columns if column not in virtual_columns + ] if vector_column in scan_columns: scanner_options["columns"] = scan_columns return vector_column, False @@ -391,15 +500,37 @@ def _get_nearest_metric(nearest: dict[str, Any]) -> str: return str(metric).lower() +def _query_ndim(query: Any) -> int: + import numpy as np + + return np.asarray(query).ndim + + +def _query_vectors_to_numpy( + query: Any, + is_batch_query: bool, + metric: str, +) -> Any: + import numpy as np + + dtype = np.uint8 if metric == "hamming" else np.float32 + query_array = np.asarray(query, dtype=dtype) + expected_ndim = 2 if is_batch_query else 1 + if query_array.ndim != expected_ndim: + kind = "two-dimensional batch" if is_batch_query else "one-dimensional" + raise ValueError(f"nearest['q'] must be a {kind} vector") + return query_array if is_batch_query else query_array.reshape(1, -1) + + def _compute_vector_distances( - vector_column: pa.ChunkedArray, + matrix: Any, query: Any, metric: str, ) -> Any: import numpy as np - matrix = _vector_column_to_numpy(vector_column) - query_vector = np.asarray(query, dtype=np.float32) + dtype = np.uint8 if metric == "hamming" else np.float32 + query_vector = np.asarray(query, dtype=dtype) if query_vector.ndim != 1: raise ValueError("nearest['q'] must be a one-dimensional vector") if matrix.shape[1] != query_vector.shape[0]: @@ -409,22 +540,25 @@ def _compute_vector_distances( ) if metric in ("l2", "euclidean"): - return np.linalg.norm(matrix - query_vector, axis=1).astype(np.float32) + difference = matrix - query_vector + return np.sum(difference * difference, axis=1).astype(np.float32) if metric == "cosine": query_norm = np.linalg.norm(query_vector) row_norms = np.linalg.norm(matrix, axis=1) denom = row_norms * query_norm + similarities = np.full(matrix.shape[0], np.nan, dtype=np.float32) similarities = np.divide( matrix @ query_vector, denom, - out=np.zeros(matrix.shape[0], dtype=np.float32), + out=similarities, where=denom != 0, ) return (1.0 - similarities).astype(np.float32) if metric in ("dot", "ip", "inner_product"): - return (-(matrix @ query_vector)).astype(np.float32) + return (1.0 - matrix @ query_vector).astype(np.float32) if metric == "hamming": - return np.count_nonzero(matrix != query_vector, axis=1).astype(np.float32) + xor = np.bitwise_xor(matrix, query_vector) + return np.bitwise_count(xor).sum(axis=1).astype(np.float32) raise ValueError( "Unsupported fallback vector search metric " @@ -432,26 +566,82 @@ def _compute_vector_distances( ) -def _vector_column_to_numpy(vector_column: pa.ChunkedArray) -> Any: +def _vector_column_to_numpy(vector_column: pa.ChunkedArray, metric: str) -> Any: import numpy as np values = vector_column.combine_chunks().to_pylist() if not values: - return np.empty((0, 0), dtype=np.float32) - if any(value is None for value in values): - raise ValueError("Fallback vector search does not support null vectors") - matrix = np.asarray(values, dtype=np.float32) + dtype = np.uint8 if metric == "hamming" else np.float32 + return np.empty((0, 0), dtype=dtype) + dtype = np.uint8 if metric == "hamming" else np.float32 + matrix = np.asarray(values, dtype=dtype) if matrix.ndim != 2: raise ValueError("Fallback vector search requires a list-like vector column") return matrix def _take_top_k(table: pa.Table, k: int) -> pa.Table: - sort_indices = pc.sort_indices(table, sort_keys=[("_distance", "ascending")]) + sort_keys = [("_distance", "ascending")] + if "_rowid" in table.column_names: + sort_keys.append(("_rowid", "ascending")) + sort_indices = pc.sort_indices(table, sort_keys=sort_keys) return table.take(sort_indices.slice(0, k)) -def _merge_vector_search_results(tables: list[pa.Table], k: int) -> pa.Table: +def _apply_distance_range(table: pa.Table, nearest: dict[str, Any]) -> pa.Table: + distance_range = nearest.get("distance_range") + if distance_range is None: + return table + + lower_bound, upper_bound = distance_range + if lower_bound is not None: + table = table.filter(pc.greater_equal(table["_distance"], lower_bound)) + if upper_bound is not None: + table = table.filter(pc.less(table["_distance"], upper_bound)) + return table + + +def _add_query_index( + table: pa.Table, + query_indices: list[int], +) -> pa.Table: + return table.add_column( + 0, + pa.field("query_index", pa.int32(), nullable=False), + pa.array(query_indices, type=pa.int32()), + ) + + +def _take_top_k_per_query(table: pa.Table, k: int) -> pa.Table: + import numpy as np + + sort_keys = [("query_index", "ascending"), ("_distance", "ascending")] + if "_rowid" in table.column_names: + sort_keys.append(("_rowid", "ascending")) + sort_indices = pc.sort_indices(table, sort_keys=sort_keys) + table = table.take(sort_indices) + if table.num_rows == 0: + return table + + query_indices = table["query_index"].combine_chunks().to_numpy() + row_indices = np.arange(table.num_rows) + group_starts = np.empty(table.num_rows, dtype=np.int64) + group_starts[0] = 0 + group_starts[1:] = np.where( + query_indices[1:] != query_indices[:-1], + row_indices[1:], + 0, + ) + np.maximum.accumulate(group_starts, out=group_starts) + return table.filter(pa.array(row_indices - group_starts < k)) + + +def _merge_vector_search_results( + tables: list[pa.Table], + k: int, + *, + is_batch_query: bool = False, +) -> pa.Table: non_empty_tables = [table for table in tables if table.num_rows > 0] if not non_empty_tables: return tables[0].slice(0, 0) if tables else pa.table({}) @@ -463,6 +653,13 @@ def _merge_vector_search_results(tables: list[pa.Table], k: int) -> pa.Table: "for global top-k merge" ) + if is_batch_query: + if "query_index" not in table.column_names: + raise RuntimeError( + "Distributed batch vector search results must include a " + "'query_index' column for per-query top-k merge" + ) + return _take_top_k_per_query(table, k) return _take_top_k(table, k) @@ -544,8 +741,9 @@ def vector_search( segment coverage. Indexed worker tasks search only their assigned ``index_segments``. Unindexed fallback tasks scan their assigned fragments without ``nearest`` and compute distances locally. Workers return local - candidates and the driver sorts by ``_distance`` to produce the final top-k - table. + candidates and the driver sorts by ``_distance`` and ``_rowid`` to produce + the final top-k table. Batch queries are merged independently by + ``query_index``. Args: uri: Lance dataset object or dataset URI. In URI mode, provide either @@ -553,7 +751,9 @@ def vector_search( nearest: Lance vector search options. Must include ``column``, ``q``, and ``k``. The worker-side ``k`` is raised to at least ``k * oversample_factor`` before the driver performs the final - global top-k merge. + global top-k merge. For fixed-size vector columns, ``q`` may be a + two-dimensional batch. Batch results contain a non-null Int32 + ``query_index`` column and up to ``k`` rows per query. index_name: Optional vector index name to use. If specified and the index cannot be found, ``ValueError`` is raised. If omitted, Lance-Ray uses the first vector index covering ``nearest["column"]``. @@ -589,8 +789,10 @@ def vector_search( supplied here. Returns: - A PyArrow table containing the global top-k rows sorted by ``_distance``. - If ``analyze_plan=True``, returns a string containing per-shard Lance + A PyArrow table containing the global top-k rows sorted by + ``_distance`` and ``_rowid``. Batch results are grouped by + ``query_index`` and contain a separate top-k for each query. If + ``analyze_plan=True``, returns a string containing per-shard Lance scanner analysis instead. """ if num_workers <= 0: @@ -606,10 +808,14 @@ def vector_search( base_scanner_options = dict(scanner_options or {}) _validate_search_scanner_options(base_scanner_options) - if columns is not None: - if isinstance(columns, list) and "_distance" not in columns: - columns = [*columns, "_distance"] - base_scanner_options["columns"] = columns + include_row_id = _projection_includes_row_id(columns, base_scanner_options) + effective_columns = ( + columns if columns is not None else base_scanner_options.get("columns") + ) + if effective_columns is not None: + if isinstance(effective_columns, list) and "_distance" not in effective_columns: + effective_columns = [*effective_columns, "_distance"] + base_scanner_options["columns"] = effective_columns if filter is not None: base_scanner_options["filter"] = filter base_scanner_options["fast_search"] = fast_search @@ -651,7 +857,15 @@ def vector_search( fragments = dataset.get_fragments() if not fragments: - return pa.table({}) + is_batch_query, result_schema = _inspect_vector_search_query( + dataset, + nearest=nearest, + base_scanner_options=base_scanner_options, + include_row_id=include_row_id, + ) + if analyze_plan: + return pa.table({}) + return pa.Table.from_batches([], schema=result_schema) vector_index = _select_vector_index( dataset, @@ -663,6 +877,13 @@ def vector_search( "No vector index found for column '%s'; distributed search will use flat scan", column, ) + nearest = _apply_index_metric_default(nearest, vector_index) + is_batch_query, result_schema = _inspect_vector_search_query( + dataset, + nearest=nearest, + base_scanner_options=base_scanner_options, + include_row_id=include_row_id, + ) plans = _plan_vector_search( fragments=fragments, @@ -670,10 +891,24 @@ def vector_search( num_workers=num_workers, include_unindexed=include_unindexed and not fast_search, ) + if ( + not is_batch_query + and any(not plan.index_segments for plan in plans) + and _query_ndim(nearest["q"]) == 2 + ): + raise ValueError( + "Flat fallback vector search does not support multivector queries. " + "Build an index covering all fragments or use fast_search=True." + ) if not plans: - return pa.table({}) + if analyze_plan: + return pa.table({}) + return pa.Table.from_batches([], schema=result_schema) pickled_dataset = pickle.dumps(dataset) + worker_scanner_options = dict(base_scanner_options) + if not analyze_plan: + worker_scanner_options["with_row_id"] = True try: with get_or_create_pool( @@ -688,10 +923,11 @@ def run_plan(plan: _SearchPlan) -> pa.Table | _SearchPlanAnalysis: return _execute_vector_search_plan( plan, pickled_dataset=worker_pickled_dataset, - base_scanner_options=base_scanner_options, + base_scanner_options=worker_scanner_options, nearest=nearest, candidate_k=candidate_k, analyze_plan=analyze_plan, + is_batch_query=is_batch_query, ) results = pool.map_async(run_plan, plans, chunksize=1).get() @@ -703,4 +939,11 @@ def run_plan(plan: _SearchPlan) -> pa.Table | _SearchPlanAnalysis: if analyze_plan: return _format_analyze_plan_results(results) - return _merge_vector_search_results(results, global_k) + result = _merge_vector_search_results( + results, + global_k, + is_batch_query=is_batch_query, + ) + if not include_row_id and "_rowid" in result.column_names: + result = result.drop_columns(["_rowid"]) + return result.select(result_schema.names) diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index 7cc0686c..3117fade 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -1920,6 +1920,25 @@ def test_build_distributed_vector_index(tmp_path, index_type): assert "ANNSubIndex" in plan assert index_name in plan + queries = np.asarray([q, q], dtype=np.float32) + batch = lr.vector_search( + updated_dataset, + nearest={"column": "vector", "q": queries, "k": 5}, + index_name=index_name, + columns=["id"], + num_workers=2, + fast_search=True, + ) + + assert batch.column("query_index").to_pylist() == [0] * 5 + [1] * 5 + assert ( + batch.slice(0, 5).column("id").to_pylist() + == batch.slice(5).column("id").to_pylist() + ) + assert batch.slice(0, 5).column("_distance").to_pylist() == pytest.approx( + batch.slice(5).column("_distance").to_pylist() + ) + if index_type == "IVF_PQ": stats = updated_dataset.stats.index_stats(index_name) assert stats["indices"] diff --git a/tests/test_distributed_search.py b/tests/test_distributed_search.py index 91a6aa11..62ed730f 100644 --- a/tests/test_distributed_search.py +++ b/tests/test_distributed_search.py @@ -1,10 +1,17 @@ from types import SimpleNamespace +import lance +import lance_ray as lr +import numpy as np import pyarrow as pa +import pyarrow.compute as pc import pytest from lance_ray import pool as pool_mod from lance_ray import search as search_mod from lance_ray.search import ( + _apply_distance_range, + _apply_index_metric_default, + _compute_vector_distances, _execute_vector_search_plan, _format_analyze_plan_results, _merge_vector_search_results, @@ -50,6 +57,67 @@ def fake_loads(value): return pickled_dataset +def _vector_table(vectors, ids=None, *, value_type=None): + matrix = np.asarray(vectors) + value_type = value_type or pa.float32() + vector_array = pa.FixedSizeListArray.from_arrays( + pa.array(matrix.reshape(-1), type=value_type), + matrix.shape[1], + ) + return pa.table( + { + "id": range(len(matrix)) if ids is None else ids, + "vector": vector_array, + } + ) + + +class _FallbackDataset: + def __init__(self, table, scanner_options=None): + self.table = table + self.scanner_options = scanner_options + + def get_fragment(self, fragment_id): + return f"fragment-{fragment_id}" + + def scanner(self, **kwargs): + if self.scanner_options is not None: + self.scanner_options.update(kwargs) + return SimpleNamespace(to_table=lambda: self.table) + + +def _create_partial_index_dataset( + path, + indexed_vectors, + appended_vectors, + *, + metric="l2", +): + dataset = lance.write_dataset( + _vector_table(indexed_vectors), + path, + max_rows_per_file=2, + ) + dataset.create_index( + "vector", + "IVF_FLAT", + num_partitions=1, + name="vector_idx", + metric=metric, + ) + lance.write_dataset( + _vector_table( + appended_vectors, + ids=range( + len(indexed_vectors), len(indexed_vectors) + len(appended_vectors) + ), + ), + path, + mode="append", + ) + return lance.dataset(path) + + def test_select_vector_index_raises_for_missing_explicit_index_name(): index = _index_with_segments(("S1", [1, 2])) dataset = SimpleNamespace(describe_indices=lambda: [index]) @@ -235,27 +303,14 @@ def scanner(self, columns=None): def test_execute_fallback_vector_search_plan_computes_local_top_k(monkeypatch): scanner_options = {} - vectors = pa.FixedSizeListArray.from_arrays( - pa.array([10.0, 0.0, 1.0, 0.0, 0.0, 2.0], type=pa.float32()), - 2, + dataset = _FallbackDataset( + _vector_table([[10.0, 0.0], [1.0, 0.0], [0.0, 2.0]], ids=[1, 2, 3]), + scanner_options, ) - class FakeDataset: - def __init__(self, *args, **kwargs): - pass - - def get_fragment(self, fragment_id): - return f"fragment-{fragment_id}" - - def scanner(self, **kwargs): - scanner_options.update(kwargs) - return SimpleNamespace( - to_table=lambda: pa.table({"id": [1, 2, 3], "vector": vectors}) - ) - result = _execute_vector_search_plan( _SearchPlan(fragment_ids=[7], index_segments=[]), - pickled_dataset=_mock_pickled_dataset(monkeypatch, FakeDataset()), + pickled_dataset=_mock_pickled_dataset(monkeypatch, dataset), base_scanner_options={"columns": ["id", "_distance"], "fast_search": False}, nearest={"column": "vector", "q": [0.0, 0.0], "k": 2}, candidate_k=2, @@ -266,10 +321,118 @@ def scanner(self, **kwargs): assert scanner_options["fragments"] == ["fragment-7"] assert scanner_options["columns"] == ["id", "vector"] assert result.column("id").to_pylist() == [2, 3] - assert result.column("_distance").to_pylist() == [1.0, 2.0] + assert result.column("_distance").to_pylist() == [1.0, 4.0] assert "vector" not in result.column_names +def test_execute_batch_hamming_fallback_uses_bit_distance(monkeypatch): + conversion_calls = 0 + dataset = _FallbackDataset(_vector_table([[0], [3], [255]], value_type=pa.uint8())) + original_vector_column_to_numpy = search_mod._vector_column_to_numpy + + def count_vector_column_conversion(vector_column, metric): + nonlocal conversion_calls + conversion_calls += 1 + return original_vector_column_to_numpy(vector_column, metric) + + monkeypatch.setattr( + search_mod, + "_vector_column_to_numpy", + count_vector_column_conversion, + ) + + result = _execute_vector_search_plan( + _SearchPlan(fragment_ids=[7], index_segments=[]), + pickled_dataset=_mock_pickled_dataset(monkeypatch, dataset), + base_scanner_options={"columns": ["id", "_distance"], "fast_search": False}, + nearest={ + "column": "vector", + "q": [[0], [255]], + "k": 3, + "metric": "hamming", + }, + candidate_k=3, + analyze_plan=False, + is_batch_query=True, + ) + + assert result.column("query_index").to_pylist() == [0, 0, 0, 1, 1, 1] + assert result.column("id").to_pylist() == [0, 1, 2, 2, 1, 0] + assert result.column("_distance").to_pylist() == [0.0, 2.0, 8.0, 0.0, 6.0, 8.0] + assert conversion_calls == 1 + + +def test_apply_distance_range_uses_inclusive_lower_exclusive_upper(): + table = pa.table({"id": [0, 1, 2], "_distance": [0.5, 1.0, 4.0]}) + + result = _apply_distance_range(table, {"distance_range": (0.5, 4.0)}) + + assert result.column("id").to_pylist() == [0, 1] + assert result.column("_distance").to_pylist() == [0.5, 1.0] + + +@pytest.mark.parametrize( + ("metric", "vectors", "query", "expected"), + [ + ("l2", [[0.0, 2.0], [3.0, 4.0]], [0.0, 0.0], [4.0, 25.0]), + ( + "cosine", + [[1.0, 0.0], [1.0, 1.0]], + [1.0, 0.0], + [0.0, 0.29289323], + ), + ("dot", [[1.0, 0.0], [1.0, 1.0]], [1.0, 1.0], [0.0, -1.0]), + ( + "hamming", + [[0, 0], [255, 0], [15, 240], [1, 2]], + [0, 0], + [0.0, 8.0, 8.0, 2.0], + ), + ], +) +def test_fallback_distance_conventions(metric, vectors, query, expected): + dtype = np.uint8 if metric == "hamming" else np.float32 + distances = _compute_vector_distances( + np.asarray(vectors, dtype=dtype), query, metric + ) + + assert distances.tolist() == pytest.approx(expected) + + +def test_execute_fallback_filters_null_and_invalid_cosine_vectors(monkeypatch): + table = pa.table( + { + "id": [0, 1, 2, 3], + "vector": pa.array( + [[0.0, 0.0], None, [1.0, 0.0], [0.0, 1.0]], + type=pa.list_(pa.float32(), 2), + ), + "_rowid": [0, 1, 2, 3], + } + ) + result = _execute_vector_search_plan( + _SearchPlan(fragment_ids=[7], index_segments=[]), + pickled_dataset=_mock_pickled_dataset( + monkeypatch, + _FallbackDataset(table), + ), + base_scanner_options={"fast_search": False, "with_row_id": True}, + nearest={ + "column": "vector", + "q": [[1.0, 0.0], [0.0, 0.0]], + "k": 4, + "metric": "cosine", + }, + candidate_k=4, + analyze_plan=False, + is_batch_query=True, + ) + + assert result.column("query_index").to_pylist() == [0, 0] + assert result.column("id").to_pylist() == [2, 3] + assert result.column("_distance").to_pylist() == [0.0, 1.0] + + def test_execute_indexed_vector_search_plan_can_analyze_plan(monkeypatch): scanner_options = {} @@ -370,39 +533,93 @@ def test_format_analyze_plan_results(): def test_merge_vector_search_results_returns_global_top_k(): - left = pa.table({"id": [1, 2], "_distance": [0.4, 0.1]}) - right = pa.table({"id": [3, 4], "_distance": [0.2, 0.3]}) + left = pa.table({"id": [1, 2], "_distance": [0.4, 0.1], "_rowid": [40, 20]}) + right = pa.table({"id": [3, 4], "_distance": [0.1, 0.3], "_rowid": [10, 30]}) result = _merge_vector_search_results([left, right], k=3) - assert result.column("id").to_pylist() == [2, 3, 4] - assert result.column("_distance").to_pylist() == [0.1, 0.2, 0.3] + assert result.column("id").to_pylist() == [3, 2, 4] + assert result.column("_distance").to_pylist() == [0.1, 0.1, 0.3] -def test_merge_vector_search_results_requires_distance(): - table = pa.table({"id": [1, 2]}) +def test_merge_batch_vector_search_results_returns_top_k_per_query(): + left = pa.table( + { + "query_index": pa.array([0, 0, 1], type=pa.int32()), + "id": [1, 2, 3], + "_distance": [0.4, 0.1, 0.2], + "_rowid": [40, 10, 30], + } + ) + right = pa.table( + { + "query_index": pa.array([0, 1, 1], type=pa.int32()), + "id": [4, 5, 6], + "_distance": [0.2, 0.3, 0.1], + "_rowid": [20, 50, 60], + } + ) + + result = _merge_vector_search_results( + [left, right], + k=2, + is_batch_query=True, + ) - with pytest.raises(RuntimeError, match="_distance"): - _merge_vector_search_results([table], k=1) + assert result.column("query_index").to_pylist() == [0, 0, 1, 1] + assert result.column("id").to_pylist() == [2, 4, 6, 3] + assert result.column("_distance").to_pylist() == [0.1, 0.2, 0.1, 0.2] -def test_search_scanner_options_reject_managed_options(): - with pytest.raises(ValueError, match="nearest"): - _validate_search_scanner_options({"nearest": {"column": "vector"}}) +@pytest.mark.parametrize( + ("table", "is_batch_query", "missing_column"), + [ + (pa.table({"id": [1, 2]}), False, "_distance"), + (pa.table({"id": [1], "_distance": [0.1]}), True, "query_index"), + ], +) +def test_merge_vector_search_results_requires_managed_columns( + table, + is_batch_query, + missing_column, +): + with pytest.raises(RuntimeError, match=missing_column): + _merge_vector_search_results( + [table], + k=1, + is_batch_query=is_batch_query, + ) -def test_search_scanner_options_reject_fast_search_override(): - with pytest.raises(ValueError, match="fast_search"): - _validate_search_scanner_options({"fast_search": True}) +@pytest.mark.parametrize( + "scanner_options", + [ + {"nearest": {"column": "vector"}}, + {"fast_search": True}, + ], + ids=["nearest", "fast_search"], +) +def test_search_scanner_options_reject_managed_options(scanner_options): + managed_option = next(iter(scanner_options)) + with pytest.raises(ValueError, match=managed_option): + _validate_search_scanner_options(scanner_options) -def test_vector_search_reuses_global_pool(monkeypatch): +def test_batch_vector_search_reuses_global_pool_in_one_round(monkeypatch): events = [] class FakeAsyncResult: def get(self): events.append("get") - return [pa.table({"id": [1], "_distance": [0.1]})] + return [ + pa.table( + { + "query_index": pa.array([0, 1], type=pa.int32()), + "id": [1, 2], + "_distance": [0.1, 0.2], + } + ) + ] class FakeGlobalPool: def map_async(self, func, plans, chunksize): @@ -438,6 +655,20 @@ def get_fragments(self): lambda *args, **kwargs: object(), ) monkeypatch.setattr(search_mod, "_plan_vector_search", lambda **kwargs: [plan]) + monkeypatch.setattr( + search_mod, + "_inspect_vector_search_query", + lambda *args, **kwargs: ( + True, + pa.schema( + [ + pa.field("query_index", pa.int32(), nullable=False), + pa.field("id", pa.int64()), + pa.field("_distance", pa.float32()), + ] + ), + ), + ) monkeypatch.setattr(search_mod.pickle, "dumps", lambda dataset: b"pickled-dataset") monkeypatch.setattr(search_mod.ray, "is_initialized", lambda: False) @@ -445,13 +676,14 @@ def get_fragments(self): try: result = search_mod.vector_search( uri="dataset", - nearest={"column": "vector", "q": [0.0], "k": 1}, + nearest={"column": "vector", "q": [[0.0], [1.0]], "k": 1}, num_workers=4, ) finally: pool_mod.clear_global_pool() - assert result.column("id").to_pylist() == [1] + assert result.column("query_index").to_pylist() == [0, 1] + assert result.column("id").to_pylist() == [1, 2] assert events == [ ("map_async", [plan], 1), "get", @@ -536,6 +768,16 @@ def fake_loads(value): lambda *args, **kwargs: object(), ) monkeypatch.setattr(search_mod, "_plan_vector_search", lambda **kwargs: plans) + monkeypatch.setattr( + search_mod, + "_inspect_vector_search_query", + lambda *args, **kwargs: ( + False, + pa.schema( + [pa.field("id", pa.int64()), pa.field("_distance", pa.float32())] + ), + ), + ) monkeypatch.setattr(search_mod.pickle, "dumps", fake_dumps) monkeypatch.setattr(search_mod.pickle, "loads", fake_loads) monkeypatch.setattr(search_mod.ray, "ObjectRef", FakeObjectRef, raising=False) @@ -569,6 +811,7 @@ def fake_loads(value): "scanner", { "fast_search": True, + "with_row_id": True, "nearest": {"column": "vector", "q": [0.0], "k": 1}, "index_segments": ["S1"], }, @@ -577,9 +820,207 @@ def fake_loads(value): "scanner", { "fast_search": True, + "with_row_id": True, "nearest": {"column": "vector", "q": [0.0], "k": 1}, "index_segments": ["S2"], }, ), "get", ] + + +def test_batch_vector_search_without_index_matches_single_queries(tmp_path): + vectors = np.asarray( + [ + [0.0, 0.0], + [1.0, 0.0], + [0.0, 2.0], + [3.0, 0.0], + [0.0, 4.0], + [5.0, 0.0], + ], + dtype=np.float32, + ) + dataset = lance.write_dataset( + _vector_table(vectors), + tmp_path / "batch-flat.lance", + max_rows_per_file=2, + ) + queries = np.asarray([[0.0, 0.0], [0.0, 4.0]], dtype=np.float32) + + batch = lr.vector_search( + dataset, + nearest={"column": "vector", "q": queries, "k": 2}, + scanner_options={"columns": ["id", "_rowid"]}, + num_workers=2, + ) + + assert batch.column_names == ["query_index", "id", "_distance", "_rowid"] + assert batch.column("query_index").to_pylist() == [0, 0, 1, 1] + assert batch.num_rows == len(queries) * 2 + for query_index, query in enumerate(queries): + single = lr.vector_search( + dataset, + nearest={"column": "vector", "q": query, "k": 2}, + scanner_options={"columns": ["id", "_rowid"]}, + num_workers=2, + ) + batch_slice = batch.filter(pc.field("query_index") == query_index).drop_columns( + ["query_index"] + ) + assert batch_slice.column("id").to_pylist() == single.column("id").to_pylist() + assert batch_slice.column("_distance").to_pylist() == pytest.approx( + single.column("_distance").to_pylist() + ) + assert ( + batch_slice.column("_rowid").to_pylist() + == single.column("_rowid").to_pylist() + ) + + empty = lr.vector_search( + dataset, + nearest={"column": "vector", "q": queries, "k": 2}, + columns=["id"], + num_workers=2, + fast_search=True, + ) + assert empty.num_rows == 0 + assert empty.column_names == ["query_index", "id", "_distance"] + assert empty.schema.field("query_index").type == pa.int32() + assert not empty.schema.field("query_index").nullable + + virtual_projection = lr.vector_search( + dataset, + nearest={"column": "vector", "q": queries, "k": 1}, + columns=["query_index"], + num_workers=2, + ) + assert virtual_projection.column_names == ["query_index", "_distance"] + assert virtual_projection.column("query_index").to_pylist() == [0, 1] + + +def test_batch_vector_search_rejects_dataset_query_index_column(tmp_path): + vectors = _vector_table([[0.0, 0.0], [1.0, 0.0]])["vector"] + dataset = lance.write_dataset( + pa.table({"query_index": [7, 8], "vector": vectors}), + tmp_path / "batch-query-index.lance", + ) + + with pytest.raises(ValueError, match="column 'query_index'"): + lr.vector_search( + dataset, + nearest={"column": "vector", "q": [[0.0, 0.0]], "k": 1}, + ) + + +def test_multivector_fallback_reports_unsupported_boundary(tmp_path): + vector_type = pa.list_(pa.list_(pa.float32(), 2)) + dataset = lance.write_dataset( + pa.table( + { + "id": [0, 1], + "vector": pa.array( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[-1.0, 0.0], [0.0, -1.0]], + ], + type=vector_type, + ), + } + ), + tmp_path / "multivector-flat.lance", + ) + + with pytest.raises(ValueError, match="does not support multivector"): + lr.vector_search( + dataset, + nearest={ + "column": "vector", + "q": [[1.0, 0.0], [0.0, 1.0]], + "k": 1, + }, + ) + + +def test_batch_vector_search_with_partial_index_preserves_per_query_top_k(tmp_path): + indexed_vectors = np.asarray( + [[0.0, 0.0], [1.0, 0.0], [0.0, 2.0], [3.0, 0.0]], + dtype=np.float32, + ) + appended_vectors = np.asarray([[0.0, 4.0], [5.0, 0.0]], dtype=np.float32) + dataset = _create_partial_index_dataset( + tmp_path / "batch-partial.lance", + indexed_vectors, + appended_vectors, + ) + queries = np.asarray([[0.0, 4.0], [3.0, 0.0]], dtype=np.float32) + + batch = lr.vector_search( + dataset, + nearest={"column": "vector", "q": queries, "k": 2}, + index_name="vector_idx", + columns=["id", "_rowid"], + num_workers=2, + ) + + assert batch.column_names == ["query_index", "id", "_distance", "_rowid"] + assert batch.column("query_index").to_pylist() == [0, 0, 1, 1] + assert batch.column("id").to_pylist() == [4, 2, 3, 1] + assert batch.column("_distance").to_pylist() == [0.0, 4.0, 0.0, 4.0] + + ranged_batch = lr.vector_search( + dataset, + nearest={ + "column": "vector", + "q": queries, + "k": 2, + "distance_range": (0.5, 10.0), + }, + index_name="vector_idx", + columns=["id"], + num_workers=2, + ) + assert ranged_batch.column("query_index").to_pylist() == [0, 1, 1] + assert ranged_batch.column("id").to_pylist() == [2, 1, 5] + assert ranged_batch.column("_distance").to_pylist() == [4.0, 4.0, 4.0] + + +@pytest.mark.parametrize("metric", ["COSINE", "DOT"]) +def test_apply_index_metric_default(metric): + index = SimpleNamespace(details={"metric_type": metric}) + nearest = {"column": "vector", "q": [[1.0, 0.0]], "k": 2} + + assert _apply_index_metric_default(nearest, index)["metric"] == metric.lower() + assert ( + _apply_index_metric_default({**nearest, "metric": "l2"}, index)["metric"] + == "l2" + ) + + +def test_partial_index_uses_index_metric_by_default(tmp_path): + indexed_vectors = np.asarray( + [[1.0, 0.0], [1.0, 1.0], [0.0, 1.0], [-1.0, 0.0]], + dtype=np.float32, + ) + appended_vectors = np.asarray([[0.5, 0.5], [-1.0, -1.0]], dtype=np.float32) + dataset = _create_partial_index_dataset( + tmp_path / "batch-partial-cosine.lance", + indexed_vectors, + appended_vectors, + metric="cosine", + ) + queries = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) + + result = lr.vector_search( + dataset, + nearest={"column": "vector", "q": queries, "k": 3}, + index_name="vector_idx", + columns=["id"], + num_workers=2, + ) + + assert result.column("query_index").to_pylist() == [0, 0, 0, 1, 1, 1] + assert result.column("id").to_pylist() == [0, 1, 4, 2, 1, 4] + assert result.column("_distance").to_pylist() == pytest.approx( + [0.0, 0.29289323, 0.29289323, 0.0, 0.29289323, 0.29289323] + ) From 3565a70e2ad6be8bef614ab9d417ede12a54a7f1 Mon Sep 17 00:00:00 2001 From: Xin Sun Date: Thu, 20 Aug 2026 19:46:12 +0800 Subject: [PATCH 2/2] refactor(search): add actor-based streaming vector search Keep the existing single-query API unchanged while introducing a bounded streaming session that reuses Ray actors and preserves dataset snapshots. Co-Authored-By: Claude Sonnet 4.6 noreply@anthropic.com --- docs/src/distributed-indexing.md | 78 +- lance_ray/__init__.py | 12 +- lance_ray/search.py | 1606 +++++++++++++++++++++++----- tests/test_distributed_indexing.py | 9 +- tests/test_distributed_search.py | 972 +++++++++-------- 5 files changed, 1919 insertions(+), 758 deletions(-) diff --git a/docs/src/distributed-indexing.md b/docs/src/distributed-indexing.md index 456906bc..804ca470 100755 --- a/docs/src/distributed-indexing.md +++ b/docs/src/distributed-indexing.md @@ -211,9 +211,9 @@ The function returns the Lance dataset instance (optimization is applied on stor ### Distributed Vector Search -`vector_search()` - Run single or batch vector search with Ray workers and merge the global top-k on the driver. +`vector_search()` - Run vector search with Ray workers and merge the global top-k on the driver. -The driver opens one fixed dataset version, reads vector index segment metadata once, and plans work by index segment ownership. Indexed worker tasks receive only their assigned `index_segments`, so a segment covering multiple fragments is never split across workers. Fragments not covered by an index can be included as separate flat-search fallback work unless `fast_search=True`; fallback tasks use regular fragment scans and compute vector distances in Lance-Ray. For a fixed-size vector column, pass a two-dimensional query `[B, D]` to search a batch. The driver merges candidates independently for each query and returns up to `k` rows per query. +The driver opens one fixed dataset version, reads vector index segment metadata once, and plans work by index segment ownership. Indexed worker tasks receive only their assigned `index_segments`, so a segment covering multiple fragments is never split across workers. Fragments not covered by an index can be included as separate flat-search fallback work unless `fast_search=True`; fallback tasks use regular fragment scans and compute vector distances in Lance-Ray. #### `vector_search` @@ -245,7 +245,7 @@ def vector_search( | Parameter | Type | Description | |-----------|------|-------------| | `uri` | `str` or `lance.LanceDataset`, optional | Lance dataset object, or its URI. Either `uri` OR (`namespace_impl` + `table_id`) must be provided when using URI mode. If a `LanceDataset` object is provided, namespace parameters are ignored and workers reopen the same dataset URI/version. | -| `nearest` | `dict[str, Any]` | Lance vector search options. Must include `column`, `q`, and `k`. For fixed-size vector columns, `q` may be one vector `[D]` or a batch `[B, D]`. If `metric` is omitted and a vector index is selected, indexed and fallback workers use the index metric; without an index the default is L2. Index-search options such as `minimum_nprobes`, `maximum_nprobes`, and `refine_factor` are forwarded to indexed workers; `distance_range` is also applied to flat fallback results. Lance-Ray raises worker-side `k` to at least `k * oversample_factor` before global merge. Multivector queries remain single queries and require full index coverage; flat fallback does not implement multivector distance. | +| `nearest` | `dict[str, Any]` | Lance vector search options. Must include `column`, `q`, and `k`. Other Lance nearest options such as `minimum_nprobes`, `maximum_nprobes`, `refine_factor`, and distance range are forwarded to every worker. Lance-Ray raises worker-side `k` to at least `k * oversample_factor` before global merge. | | `index_name` | `str`, optional | Vector index name to use. If provided and not found, `vector_search()` raises `ValueError` instead of silently falling back. If omitted, Lance-Ray uses the first vector index covering `nearest["column"]`; if none exists, the search uses flat fallback plans unless `fast_search=True`. | | `columns` | `list[str]` or `dict[str, str]`, optional | Projection passed to the Lance scanner. When a list is provided and `_distance` is missing, Lance-Ray appends `_distance` automatically because the driver needs it for global top-k merge. | | `filter` | `Any`, optional | Filter passed unchanged to every worker scanner. | @@ -264,7 +264,59 @@ def vector_search( #### Return Value -For a one-dimensional query, the function returns a `pyarrow.Table` containing the global top-k rows sorted by `_distance` and `_rowid`. For a two-dimensional batch query, the first column is a non-null Int32 `query_index`; rows are grouped in input-query order, and each group contains up to `k` rows sorted by `_distance` and `_rowid`. The internal `_rowid` tie-break column is omitted unless requested. If `analyze_plan=True`, the function returns a `str` containing one Lance scanner analysis section per planned shard. +The function returns a `pyarrow.Table` containing the global top-k rows sorted by `_distance`. If `analyze_plan=True`, it returns a `str` containing one Lance scanner analysis section per planned shard. + + +#### Streaming vector search + +`open_vector_search()` creates snapshot-pinned Ray actors and accepts an iterable +of query batches. Each actor keeps its Lance dataset, session, and assigned index +segments alive for the session. Uncovered fragments are distributed across the +same actors and scanned in record batches when `fast_search=False`. + +```python +import lance_ray as lr + +with lr.open_vector_search( + uri="path/to/dataset.lance", + nearest={ + "column": "embedding", + "k": 10, + }, + columns=["id"], +) as search: + for result_batch in search.map_batches(query_batch_reader): + write_results(result_batch) +``` + +`nearest["column"]` selects the vector column to search, while `columns` +selects the fields returned for each match. + +`VectorSearchStreamingOptions` controls input rebatching and the number of +in-flight query batches. `VectorSearchActorOptions` controls actor resources, +micro-batching, scanner concurrency, cache sizes, and optional index prewarming. +Lance query options such as `nprobes`, `query_parallelism`, `approx_mode`, and +`refine_factor` can be supplied in `nearest`. + +`uri` and (`namespace_impl` + `table_id`) are alternative dataset sources. +`branch` and `version` are mutually exclusive. The session pins the resolved +snapshot for its lifetime. Passing an already checked-out `LanceDataset` +preserves its current snapshot without requiring the branch name again. + +`fast_search=False` includes fragments not covered by the selected vector index. +`fast_search=True` intentionally skips those fragments. The streaming API does +not expose a second `include_unindexed` switch. + +Each emitted table has a non-null Int64 `query_index` whose value is the query's +position across the entire input stream. + +For a fixed-size vector column, each input batch may be a NumPy array with shape +`[B, D]` or an Arrow `FixedSizeList` array; each row is one logical query and +Core executes it as a batch. For a multivector column, each logical query has +shape `[M_i, D]`. Supply an Arrow `List>` array, a sequence of +two-dimensional NumPy arrays, or `[B, M, D]` when every query has the same `M`. +Lance-Ray adds the streaming `query_index` and merges indexed and +uncovered-fragment candidates for each query. ## Examples @@ -393,24 +445,6 @@ results = lr.vector_search( fast_search=False, ) -# Run two queries in one Ray scheduling round. The result is one table whose -# query_index column maps every row back to query_vectors[0] or query_vectors[1]. -query_vectors = [query_vector, another_query_vector] -batch_results = lr.vector_search( - uri="path/to/dataset.lance", - nearest={ - "column": "vector", - "q": query_vectors, - "k": 10, - "minimum_nprobes": 20, - }, - index_name="idx_ivf_flat", - columns=["id", "vector"], - num_workers=8, - oversample_factor=2, - fast_search=False, -) - # Inspect the per-shard Lance scanner plans instead of executing the search. plan = lr.vector_search( uri="path/to/dataset.lance", diff --git a/lance_ray/__init__.py b/lance_ray/__init__.py index cea756a3..883b3e62 100644 --- a/lance_ray/__init__.py +++ b/lance_ray/__init__.py @@ -24,12 +24,22 @@ write_lance, ) from .pool import clear_global_pool, get_global_pool, init_global_pool, set_global_pool -from .search import vector_search +from .search import ( + VectorSearchActorOptions, + VectorSearchSession, + VectorSearchStreamingOptions, + open_vector_search, + vector_search, +) __all__ = [ "read_lance", "write_lance", "vector_search", + "open_vector_search", + "VectorSearchSession", + "VectorSearchStreamingOptions", + "VectorSearchActorOptions", "init_global_pool", "set_global_pool", "get_global_pool", diff --git a/lance_ray/search.py b/lance_ray/search.py index bdfbefbd..1958e6c7 100644 --- a/lance_ray/search.py +++ b/lance_ray/search.py @@ -1,10 +1,14 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright The Lance Authors +import concurrent.futures import inspect import logging import math import pickle +from collections import deque +from collections.abc import Iterable, Iterator +from dataclasses import dataclass from functools import lru_cache from typing import TYPE_CHECKING, Any, NamedTuple, Optional, Union @@ -124,34 +128,15 @@ def _canonical_index_field_names(field_names: Any) -> set[str]: return canonical_names -def _apply_index_metric_default( - nearest: dict[str, Any], - vector_index: Any | None, -) -> dict[str, Any]: - if ( - vector_index is None - or nearest.get("metric") is not None - or nearest.get("distance_type") is not None - ): - return nearest - - details = _index_value(vector_index, "details", {}) or {} - metric = _index_value(details, "metric_type") - if metric is None: - return nearest - return {**nearest, "metric": str(metric).lower()} - - -def _plan_vector_search( +def _build_vector_search_plan_units( *, fragments: list[Any], vector_index: Any | None, - num_workers: int, include_unindexed: bool, -) -> list[_SearchPlan]: +) -> tuple[list[_SearchPlanUnit], list[_SearchPlanUnit], int, int]: fragment_ids = {_get_fragment_id(fragment) for fragment in fragments} if not fragment_ids: - return [] + return [], [], 0, 0 fragment_weights: dict[int, int] = {} for fragment in fragments: @@ -192,6 +177,28 @@ def _plan_vector_search( ) ) + return ( + indexed_units, + fallback_units, + len(fragment_ids), + len(fallback_fragment_ids), + ) + + +def _plan_vector_search( + *, + fragments: list[Any], + vector_index: Any | None, + num_workers: int, + include_unindexed: bool, +) -> list[_SearchPlan]: + indexed_units, fallback_units, fragment_count, fallback_count = ( + _build_vector_search_plan_units( + fragments=fragments, + vector_index=vector_index, + include_unindexed=include_unindexed, + ) + ) plans = [ *_pack_search_plan_units(indexed_units, num_workers), *_pack_search_plan_units(fallback_units, num_workers), @@ -200,12 +207,12 @@ def _plan_vector_search( if not plans: return [] - included_fallback_count = len(fallback_fragment_ids) if include_unindexed else 0 + included_fallback_count = fallback_count if include_unindexed else 0 logger.info( "Planned distributed vector search across %d tasks, %d fragments, " "%d index segments, %d fallback fragments", len(plans), - len(fragment_ids), + fragment_count, sum(len(plan.index_segments) for plan in plans), included_fallback_count, ) @@ -271,7 +278,6 @@ def _execute_vector_search_plan( nearest: dict[str, Any], candidate_k: int, analyze_plan: bool, - is_batch_query: bool = False, ) -> pa.Table | _SearchPlanAnalysis: dataset = _load_worker_dataset(pickled_dataset) @@ -283,7 +289,6 @@ def _execute_vector_search_plan( nearest=nearest, candidate_k=candidate_k, analyze_plan=analyze_plan, - is_batch_query=is_batch_query, ) if not _scanner_accepts_index_segments(dataset): @@ -324,49 +329,6 @@ def _scanner_accepts_index_segments(dataset: LanceDataset) -> bool: ) -def _inspect_vector_search_query( - dataset: LanceDataset, - *, - nearest: dict[str, Any], - base_scanner_options: dict[str, Any], - include_row_id: bool, -) -> tuple[bool, pa.Schema]: - probe = dataset.scanner(columns=["_distance"], nearest=nearest) - probe_schema = probe.projected_schema - is_batch_query = ( - probe_schema.names - and probe_schema.names[0] == "query_index" - and pa.types.is_int32(probe_schema.field(0).type) - and not probe_schema.field(0).nullable - ) - - schema_options = dict(base_scanner_options) - schema_options["nearest"] = nearest - schema_options["with_row_id"] = True - result_schema = dataset.scanner(**schema_options).projected_schema - if not include_row_id: - row_id_indices = result_schema.get_all_field_indices("_rowid") - if row_id_indices: - result_schema = result_schema.remove(row_id_indices[-1]) - - return is_batch_query, result_schema - - -def _projection_includes_row_id( - columns: Optional[list[str] | dict[str, str]], - scanner_options: dict[str, Any], -) -> bool: - if scanner_options.get("with_row_id"): - return True - if columns is None: - columns = scanner_options.get("columns") - if isinstance(columns, list): - return "_rowid" in columns - if isinstance(columns, dict): - return "_rowid" in columns - return False - - def _execute_flat_fallback_vector_search_plan( dataset: LanceDataset, *, @@ -375,15 +337,13 @@ def _execute_flat_fallback_vector_search_plan( nearest: dict[str, Any], candidate_k: int, analyze_plan: bool, - is_batch_query: bool, ) -> pa.Table | _SearchPlanAnalysis: vector_column = nearest["column"] - scanner_options = dict(base_scanner_options) vector_scan_column, drop_vector_column = _prepare_fallback_scan_columns( - scanner_options, + base_scanner_options, vector_column, - is_batch_query=is_batch_query, ) + scanner_options = dict(base_scanner_options) scanner_options.pop("fast_search", None) scanner_options["fragments"] = [ dataset.get_fragment(fragment_id) for fragment_id in plan.fragment_ids @@ -401,56 +361,19 @@ def _execute_flat_fallback_vector_search_plan( table = scanner.to_table() if table.num_rows == 0: table = table.append_column("_distance", pa.array([], type=pa.float32())) - if is_batch_query: - table = _add_query_index(table, []) - if drop_vector_column and vector_scan_column in table.column_names: - table = table.drop_columns([vector_scan_column]) - return table - - valid_vectors = pc.invert(pc.is_null(table[vector_scan_column])) - table = table.filter(valid_vectors) - if table.num_rows == 0: - table = table.append_column("_distance", pa.array([], type=pa.float32())) - if is_batch_query: - table = _add_query_index(table, []) if drop_vector_column and vector_scan_column in table.column_names: table = table.drop_columns([vector_scan_column]) return table - metric = _get_nearest_metric(nearest) - vector_matrix = _vector_column_to_numpy(table[vector_scan_column], metric) - query_vectors = _query_vectors_to_numpy(nearest["q"], is_batch_query, metric) - import numpy as np - - query_results = [] - for query_index, query_vector in enumerate(query_vectors): - distances = _compute_vector_distances( - vector_matrix, - query_vector, - metric, - ) - finite_distances = np.isfinite(distances) - query_result = table.filter(pa.array(finite_distances, type=pa.bool_())) - distances = distances[finite_distances] - query_result = query_result.append_column( - "_distance", pa.array(distances, type=pa.float32()) - ) - query_result = _apply_distance_range(query_result, nearest) - query_result = _take_top_k(query_result, candidate_k) - if drop_vector_column and vector_scan_column in query_result.column_names: - query_result = query_result.drop_columns([vector_scan_column]) - if is_batch_query: - query_result = _add_query_index( - query_result, - [query_index] * query_result.num_rows, - ) - query_results.append(query_result) - - table = ( - pa.concat_tables(query_results, promote_options="default") - if is_batch_query - else query_results[0] + distances = _compute_vector_distances( + table[vector_scan_column], + nearest["q"], + _get_nearest_metric(nearest), ) + table = table.append_column("_distance", pa.array(distances, type=pa.float32())) + table = _take_top_k(table, candidate_k) + if drop_vector_column and vector_scan_column in table.column_names: + table = table.drop_columns([vector_scan_column]) return table @@ -458,16 +381,14 @@ def _prepare_fallback_scan_columns( scanner_options: dict[str, Any], vector_column: str, *, - is_batch_query: bool, + virtual_columns: Optional[set[str]] = None, ) -> tuple[str, bool]: requested_columns = scanner_options.get("columns") if requested_columns is None: return vector_column, False if isinstance(requested_columns, list): - virtual_columns = {"_distance"} - if is_batch_query: - virtual_columns.add("query_index") + virtual_columns = virtual_columns or {"_distance"} scan_columns = [ column for column in requested_columns if column not in virtual_columns ] @@ -500,37 +421,15 @@ def _get_nearest_metric(nearest: dict[str, Any]) -> str: return str(metric).lower() -def _query_ndim(query: Any) -> int: - import numpy as np - - return np.asarray(query).ndim - - -def _query_vectors_to_numpy( - query: Any, - is_batch_query: bool, - metric: str, -) -> Any: - import numpy as np - - dtype = np.uint8 if metric == "hamming" else np.float32 - query_array = np.asarray(query, dtype=dtype) - expected_ndim = 2 if is_batch_query else 1 - if query_array.ndim != expected_ndim: - kind = "two-dimensional batch" if is_batch_query else "one-dimensional" - raise ValueError(f"nearest['q'] must be a {kind} vector") - return query_array if is_batch_query else query_array.reshape(1, -1) - - def _compute_vector_distances( - matrix: Any, + vector_column: pa.ChunkedArray, query: Any, metric: str, ) -> Any: import numpy as np - dtype = np.uint8 if metric == "hamming" else np.float32 - query_vector = np.asarray(query, dtype=dtype) + matrix = _vector_column_to_numpy(vector_column) + query_vector = np.asarray(query, dtype=np.float32) if query_vector.ndim != 1: raise ValueError("nearest['q'] must be a one-dimensional vector") if matrix.shape[1] != query_vector.shape[0]: @@ -540,25 +439,22 @@ def _compute_vector_distances( ) if metric in ("l2", "euclidean"): - difference = matrix - query_vector - return np.sum(difference * difference, axis=1).astype(np.float32) + return np.linalg.norm(matrix - query_vector, axis=1).astype(np.float32) if metric == "cosine": query_norm = np.linalg.norm(query_vector) row_norms = np.linalg.norm(matrix, axis=1) denom = row_norms * query_norm - similarities = np.full(matrix.shape[0], np.nan, dtype=np.float32) similarities = np.divide( matrix @ query_vector, denom, - out=similarities, + out=np.zeros(matrix.shape[0], dtype=np.float32), where=denom != 0, ) return (1.0 - similarities).astype(np.float32) if metric in ("dot", "ip", "inner_product"): - return (1.0 - matrix @ query_vector).astype(np.float32) + return (-(matrix @ query_vector)).astype(np.float32) if metric == "hamming": - xor = np.bitwise_xor(matrix, query_vector) - return np.bitwise_count(xor).sum(axis=1).astype(np.float32) + return np.count_nonzero(matrix != query_vector, axis=1).astype(np.float32) raise ValueError( "Unsupported fallback vector search metric " @@ -566,81 +462,31 @@ def _compute_vector_distances( ) -def _vector_column_to_numpy(vector_column: pa.ChunkedArray, metric: str) -> Any: +def _vector_column_to_numpy(vector_column: pa.ChunkedArray) -> Any: import numpy as np values = vector_column.combine_chunks().to_pylist() if not values: - dtype = np.uint8 if metric == "hamming" else np.float32 - return np.empty((0, 0), dtype=dtype) - dtype = np.uint8 if metric == "hamming" else np.float32 - matrix = np.asarray(values, dtype=dtype) + return np.empty((0, 0), dtype=np.float32) + if any(value is None for value in values): + raise ValueError("Fallback vector search does not support null vectors") + matrix = np.asarray(values, dtype=np.float32) if matrix.ndim != 2: raise ValueError("Fallback vector search requires a list-like vector column") return matrix def _take_top_k(table: pa.Table, k: int) -> pa.Table: - sort_keys = [("_distance", "ascending")] - if "_rowid" in table.column_names: - sort_keys.append(("_rowid", "ascending")) - sort_indices = pc.sort_indices(table, sort_keys=sort_keys) + sort_indices = pc.sort_indices(table, sort_keys=[("_distance", "ascending")]) return table.take(sort_indices.slice(0, k)) -def _apply_distance_range(table: pa.Table, nearest: dict[str, Any]) -> pa.Table: - distance_range = nearest.get("distance_range") - if distance_range is None: - return table - - lower_bound, upper_bound = distance_range - if lower_bound is not None: - table = table.filter(pc.greater_equal(table["_distance"], lower_bound)) - if upper_bound is not None: - table = table.filter(pc.less(table["_distance"], upper_bound)) - return table - - -def _add_query_index( - table: pa.Table, - query_indices: list[int], -) -> pa.Table: - return table.add_column( - 0, - pa.field("query_index", pa.int32(), nullable=False), - pa.array(query_indices, type=pa.int32()), - ) - - -def _take_top_k_per_query(table: pa.Table, k: int) -> pa.Table: - import numpy as np - - sort_keys = [("query_index", "ascending"), ("_distance", "ascending")] - if "_rowid" in table.column_names: - sort_keys.append(("_rowid", "ascending")) - sort_indices = pc.sort_indices(table, sort_keys=sort_keys) - table = table.take(sort_indices) - if table.num_rows == 0: - return table - - query_indices = table["query_index"].combine_chunks().to_numpy() - row_indices = np.arange(table.num_rows) - group_starts = np.empty(table.num_rows, dtype=np.int64) - group_starts[0] = 0 - group_starts[1:] = np.where( - query_indices[1:] != query_indices[:-1], - row_indices[1:], - 0, - ) - np.maximum.accumulate(group_starts, out=group_starts) - return table.filter(pa.array(row_indices - group_starts < k)) - - def _merge_vector_search_results( tables: list[pa.Table], k: int, *, - is_batch_query: bool = False, + per_query: bool = False, + deterministic: bool = False, ) -> pa.Table: non_empty_tables = [table for table in tables if table.num_rows > 0] if not non_empty_tables: @@ -653,13 +499,15 @@ def _merge_vector_search_results( "for global top-k merge" ) - if is_batch_query: + if per_query: if "query_index" not in table.column_names: raise RuntimeError( "Distributed batch vector search results must include a " "'query_index' column for per-query top-k merge" ) return _take_top_k_per_query(table, k) + if deterministic: + return _take_top_k_deterministic(table, k) return _take_top_k(table, k) @@ -741,9 +589,8 @@ def vector_search( segment coverage. Indexed worker tasks search only their assigned ``index_segments``. Unindexed fallback tasks scan their assigned fragments without ``nearest`` and compute distances locally. Workers return local - candidates and the driver sorts by ``_distance`` and ``_rowid`` to produce - the final top-k table. Batch queries are merged independently by - ``query_index``. + candidates and the driver sorts by ``_distance`` to produce the final top-k + table. Args: uri: Lance dataset object or dataset URI. In URI mode, provide either @@ -751,9 +598,7 @@ def vector_search( nearest: Lance vector search options. Must include ``column``, ``q``, and ``k``. The worker-side ``k`` is raised to at least ``k * oversample_factor`` before the driver performs the final - global top-k merge. For fixed-size vector columns, ``q`` may be a - two-dimensional batch. Batch results contain a non-null Int32 - ``query_index`` column and up to ``k`` rows per query. + global top-k merge. index_name: Optional vector index name to use. If specified and the index cannot be found, ``ValueError`` is raised. If omitted, Lance-Ray uses the first vector index covering ``nearest["column"]``. @@ -789,10 +634,8 @@ def vector_search( supplied here. Returns: - A PyArrow table containing the global top-k rows sorted by - ``_distance`` and ``_rowid``. Batch results are grouped by - ``query_index`` and contain a separate top-k for each query. If - ``analyze_plan=True``, returns a string containing per-shard Lance + A PyArrow table containing the global top-k rows sorted by ``_distance``. + If ``analyze_plan=True``, returns a string containing per-shard Lance scanner analysis instead. """ if num_workers <= 0: @@ -808,14 +651,10 @@ def vector_search( base_scanner_options = dict(scanner_options or {}) _validate_search_scanner_options(base_scanner_options) - include_row_id = _projection_includes_row_id(columns, base_scanner_options) - effective_columns = ( - columns if columns is not None else base_scanner_options.get("columns") - ) - if effective_columns is not None: - if isinstance(effective_columns, list) and "_distance" not in effective_columns: - effective_columns = [*effective_columns, "_distance"] - base_scanner_options["columns"] = effective_columns + if columns is not None: + if isinstance(columns, list) and "_distance" not in columns: + columns = [*columns, "_distance"] + base_scanner_options["columns"] = columns if filter is not None: base_scanner_options["filter"] = filter base_scanner_options["fast_search"] = fast_search @@ -857,15 +696,7 @@ def vector_search( fragments = dataset.get_fragments() if not fragments: - is_batch_query, result_schema = _inspect_vector_search_query( - dataset, - nearest=nearest, - base_scanner_options=base_scanner_options, - include_row_id=include_row_id, - ) - if analyze_plan: - return pa.table({}) - return pa.Table.from_batches([], schema=result_schema) + return pa.table({}) vector_index = _select_vector_index( dataset, @@ -877,13 +708,6 @@ def vector_search( "No vector index found for column '%s'; distributed search will use flat scan", column, ) - nearest = _apply_index_metric_default(nearest, vector_index) - is_batch_query, result_schema = _inspect_vector_search_query( - dataset, - nearest=nearest, - base_scanner_options=base_scanner_options, - include_row_id=include_row_id, - ) plans = _plan_vector_search( fragments=fragments, @@ -891,24 +715,10 @@ def vector_search( num_workers=num_workers, include_unindexed=include_unindexed and not fast_search, ) - if ( - not is_batch_query - and any(not plan.index_segments for plan in plans) - and _query_ndim(nearest["q"]) == 2 - ): - raise ValueError( - "Flat fallback vector search does not support multivector queries. " - "Build an index covering all fragments or use fast_search=True." - ) if not plans: - if analyze_plan: - return pa.table({}) - return pa.Table.from_batches([], schema=result_schema) + return pa.table({}) pickled_dataset = pickle.dumps(dataset) - worker_scanner_options = dict(base_scanner_options) - if not analyze_plan: - worker_scanner_options["with_row_id"] = True try: with get_or_create_pool( @@ -923,11 +733,10 @@ def run_plan(plan: _SearchPlan) -> pa.Table | _SearchPlanAnalysis: return _execute_vector_search_plan( plan, pickled_dataset=worker_pickled_dataset, - base_scanner_options=worker_scanner_options, + base_scanner_options=base_scanner_options, nearest=nearest, candidate_k=candidate_k, analyze_plan=analyze_plan, - is_batch_query=is_batch_query, ) results = pool.map_async(run_plan, plans, chunksize=1).get() @@ -939,11 +748,1266 @@ def run_plan(plan: _SearchPlan) -> pa.Table | _SearchPlanAnalysis: if analyze_plan: return _format_analyze_plan_results(results) - result = _merge_vector_search_results( - results, - global_k, - is_batch_query=is_batch_query, + return _merge_vector_search_results(results, global_k) + + +def _apply_index_metric_default( + nearest: dict[str, Any], + vector_index: Any | None, +) -> dict[str, Any]: + if ( + vector_index is None + or nearest.get("metric") is not None + or nearest.get("distance_type") is not None + ): + return nearest + + details = _index_value(vector_index, "details", {}) or {} + metric = _index_value(details, "metric_type") + if metric is None: + return nearest + return {**nearest, "metric": str(metric).lower()} + + +def _inspect_vector_search_query( + dataset: LanceDataset, + *, + nearest: dict[str, Any], + base_scanner_options: dict[str, Any], + include_row_id: bool, +) -> tuple[bool, pa.Schema]: + probe = dataset.scanner(columns=["_distance"], nearest=nearest) + probe_schema = probe.projected_schema + is_batch_query = ( + probe_schema.names + and probe_schema.names[0] == "query_index" + and pa.types.is_int32(probe_schema.field(0).type) + and not probe_schema.field(0).nullable + ) + + schema_options = dict(base_scanner_options) + schema_options["nearest"] = nearest + schema_options["with_row_id"] = True + result_schema = dataset.scanner(**schema_options).projected_schema + if not include_row_id: + row_id_indices = result_schema.get_all_field_indices("_rowid") + if row_id_indices: + result_schema = result_schema.remove(row_id_indices[-1]) + + return is_batch_query, result_schema + + +def _projection_includes_row_id( + columns: Optional[list[str] | dict[str, str]], + scanner_options: dict[str, Any], +) -> bool: + if scanner_options.get("with_row_id"): + return True + if columns is None: + columns = scanner_options.get("columns") + if isinstance(columns, list): + return "_rowid" in columns + if isinstance(columns, dict): + return "_rowid" in columns + return False + + +def _compute_core_vector_distances( + matrix: Any, + query: Any, + metric: str, +) -> Any: + """Compute distances using Lance Core's current scalar conventions. + + The established ``vector_search`` fallback has different public distance + conventions, so it continues to use ``_compute_vector_distances``. + """ + import numpy as np + + dtype = np.uint8 if metric == "hamming" else np.float32 + query_vector = np.asarray(query, dtype=dtype) + if query_vector.ndim != 1: + raise ValueError("nearest['q'] must be a one-dimensional vector") + if matrix.shape[1] != query_vector.shape[0]: + raise ValueError( + "Query vector dimension does not match fallback vector column " + f"dimension: {query_vector.shape[0]} != {matrix.shape[1]}" + ) + + if metric in ("l2", "euclidean"): + difference = matrix - query_vector + return np.sum(difference * difference, axis=1).astype(np.float32) + if metric == "cosine": + query_norm = np.linalg.norm(query_vector) + row_norms = np.linalg.norm(matrix, axis=1) + denom = row_norms * query_norm + similarities = np.full(matrix.shape[0], np.nan, dtype=np.float32) + similarities = np.divide( + matrix @ query_vector, + denom, + out=similarities, + where=denom != 0, + ) + return (1.0 - similarities).astype(np.float32) + if metric in ("dot", "ip", "inner_product"): + return (1.0 - matrix @ query_vector).astype(np.float32) + if metric == "hamming": + xor = np.bitwise_xor(matrix, query_vector) + return np.bitwise_count(xor).sum(axis=1).astype(np.float32) + + raise ValueError( + "Unsupported fallback vector search metric " + f"{metric!r}. Supported metrics: l2, cosine, dot, hamming" + ) + + +def _vector_column_to_numpy_for_metric( + vector_column: pa.ChunkedArray, metric: str +) -> Any: + import numpy as np + + values = vector_column.combine_chunks().to_pylist() + if not values: + dtype = np.uint8 if metric == "hamming" else np.float32 + return np.empty((0, 0), dtype=dtype) + dtype = np.uint8 if metric == "hamming" else np.float32 + matrix = np.asarray(values, dtype=dtype) + if matrix.ndim != 2: + raise ValueError("Fallback vector search requires a list-like vector column") + return matrix + + +def _take_top_k_deterministic(table: pa.Table, k: int) -> pa.Table: + sort_keys = [("_distance", "ascending")] + if "_rowid" in table.column_names: + sort_keys.append(("_rowid", "ascending")) + sort_indices = pc.sort_indices(table, sort_keys=sort_keys) + return table.take(sort_indices.slice(0, k)) + + +def _apply_distance_range(table: pa.Table, nearest: dict[str, Any]) -> pa.Table: + distance_range = nearest.get("distance_range") + if distance_range is None: + return table + + lower_bound, upper_bound = distance_range + if lower_bound is not None: + table = table.filter(pc.greater_equal(table["_distance"], lower_bound)) + if upper_bound is not None: + table = table.filter(pc.less(table["_distance"], upper_bound)) + return table + + +def _take_top_k_per_query(table: pa.Table, k: int) -> pa.Table: + import numpy as np + + sort_keys = [("query_index", "ascending"), ("_distance", "ascending")] + if "_rowid" in table.column_names: + sort_keys.append(("_rowid", "ascending")) + sort_indices = pc.sort_indices(table, sort_keys=sort_keys) + table = table.take(sort_indices) + if table.num_rows == 0: + return table + + query_indices = table["query_index"].combine_chunks().to_numpy() + row_indices = np.arange(table.num_rows) + group_starts = np.empty(table.num_rows, dtype=np.int64) + group_starts[0] = 0 + group_starts[1:] = np.where( + query_indices[1:] != query_indices[:-1], + row_indices[1:], + 0, + ) + np.maximum.accumulate(group_starts, out=group_starts) + return table.filter(pa.array(row_indices - group_starts < k)) + + +@dataclass(frozen=True) +class VectorSearchStreamingOptions: + """Controls query batching and the bounded driver pipeline.""" + + query_batch_size: Optional[int] = None + max_in_flight_batches: int = 1 + + def __post_init__(self) -> None: + if self.query_batch_size is not None and self.query_batch_size <= 0: + raise ValueError("query_batch_size must be positive") + if self.max_in_flight_batches <= 0: + raise ValueError("max_in_flight_batches must be positive") + + +@dataclass(frozen=True) +class VectorSearchActorOptions: + """Controls Ray actors, their Lance sessions, and scanner execution.""" + + num_actors: int = 4 + ray_remote_args: Optional[dict[str, Any]] = None + max_concurrent_batches: int = 1 + max_pending_calls: Optional[int] = None + micro_batch_size: Optional[int] = None + scanner_concurrency: int = 1 + index_cache_size_bytes: Optional[int] = None + metadata_cache_size_bytes: Optional[int] = None + prewarm_index: bool = False + + def __post_init__(self) -> None: + if self.num_actors <= 0: + raise ValueError("num_actors must be positive") + if self.max_concurrent_batches <= 0: + raise ValueError("max_concurrent_batches must be positive") + if self.max_pending_calls is not None and self.max_pending_calls <= 0: + raise ValueError("max_pending_calls must be positive") + if self.micro_batch_size is not None and self.micro_batch_size <= 0: + raise ValueError("micro_batch_size must be positive") + if self.scanner_concurrency <= 0: + raise ValueError("scanner_concurrency must be positive") + if self.index_cache_size_bytes is not None and self.index_cache_size_bytes < 0: + raise ValueError("index_cache_size_bytes must be non-negative") + if ( + self.metadata_cache_size_bytes is not None + and self.metadata_cache_size_bytes < 0 + ): + raise ValueError("metadata_cache_size_bytes must be non-negative") + + +@dataclass(frozen=True) +class _DatasetSnapshot: + uri: str + version: int + serialized_manifest: bytes + storage_options: dict[str, Any] + base_store_params: Optional[dict[str, dict[str, Any]]] + block_size: Optional[int] + namespace_impl: Optional[str] + namespace_properties: Optional[dict[str, str]] + table_id: Optional[list[str]] + + +@dataclass(frozen=True) +class _ActorPlan: + indexed_fragment_ids: tuple[int, ...] + index_segments: tuple[str, ...] + fallback_fragment_ids: tuple[int, ...] + weight: int + + +def _plan_streaming_vector_search( + *, + fragments: list[Any], + vector_index: Any | None, + num_actors: int, + fast_search: bool, +) -> list[_ActorPlan]: + indexed_units, fallback_units, _, _ = _build_vector_search_plan_units( + fragments=fragments, + vector_index=vector_index, + include_unindexed=not fast_search, + ) + units = [*indexed_units, *fallback_units] + + if not units: + return [] + + actor_count = min(num_actors, len(units)) + actor_weights = [0] * actor_count + indexed_fragments = [set() for _ in range(actor_count)] + index_segments = [[] for _ in range(actor_count)] + fallback_fragments = [set() for _ in range(actor_count)] + + for unit in sorted(units, key=lambda item: item.weight, reverse=True): + actor_idx = min(range(actor_count), key=lambda idx: actor_weights[idx]) + if not unit.index_segments: + fallback_fragments[actor_idx].update(unit.fragment_ids) + else: + indexed_fragments[actor_idx].update(unit.fragment_ids) + index_segments[actor_idx].extend(unit.index_segments) + actor_weights[actor_idx] += unit.weight + + return [ + _ActorPlan( + indexed_fragment_ids=tuple(sorted(indexed_fragments[idx])), + index_segments=tuple(index_segments[idx]), + fallback_fragment_ids=tuple(sorted(fallback_fragments[idx])), + weight=actor_weights[idx], + ) + for idx in range(actor_count) + ] + + +def _open_snapshot( + snapshot: _DatasetSnapshot, + *, + index_cache_size_bytes: Optional[int], + metadata_cache_size_bytes: Optional[int], +) -> LanceDataset: + import lance + + session = lance.Session( + index_cache_size_bytes=index_cache_size_bytes, + metadata_cache_size_bytes=metadata_cache_size_bytes, + ) + namespace_kwargs = get_namespace_kwargs( + snapshot.namespace_impl, + snapshot.namespace_properties, + snapshot.table_id, + ) + kwargs: dict[str, Any] = { + "storage_options": snapshot.storage_options, + "session": session, + **namespace_kwargs, + } + if snapshot.block_size is not None: + kwargs["block_size"] = snapshot.block_size + if snapshot.base_store_params is not None: + kwargs["base_store_params"] = snapshot.base_store_params + + dataset = LanceDataset( + snapshot.uri, + version=snapshot.version, + serialized_manifest=snapshot.serialized_manifest, + **kwargs, + ) + if dataset.version != snapshot.version: + raise RuntimeError( + f"Dataset snapshot changed: expected {snapshot.version}, " + f"opened {dataset.version}" + ) + return dataset + + +def _empty_fallback_table( + scanner: Any, + *, + vector_column: str, + drop_vector_column: bool, +) -> pa.Table: + schema = getattr(scanner, "projected_schema", pa.schema([])) + fields = list(schema) + if drop_vector_column: + fields = [field for field in fields if field.name != vector_column] + fields = [ + field for field in fields if field.name not in {"query_index", "_distance"} + ] + schema = pa.schema( + [ + pa.field("query_index", pa.int32(), nullable=False), + *fields, + pa.field("_distance", pa.float32()), + ] + ) + return pa.Table.from_batches([], schema=schema) + + +def _scanner_batches(scanner: Any) -> Iterable[pa.RecordBatch]: + if hasattr(scanner, "to_batches"): + return scanner.to_batches() + return scanner.to_table().to_batches() + + +def _stream_flat_fallback( + dataset: LanceDataset, + *, + fragment_ids: tuple[int, ...], + base_scanner_options: dict[str, Any], + nearest: dict[str, Any], + candidate_k: int, +) -> pa.Table: + vector_column = nearest["column"] + scanner_options = dict(base_scanner_options) + vector_scan_column, drop_vector_column = _prepare_fallback_scan_columns( + scanner_options, + vector_column, + virtual_columns={"_distance", "query_index"}, + ) + scanner_options.pop("fast_search", None) + scanner_options["fragments"] = [ + dataset.get_fragment(fragment_id) for fragment_id in fragment_ids + ] + scanner = dataset.scanner(**scanner_options) + + metric = _get_nearest_metric(nearest) + query_vectors = _canonical_query_batch(nearest["q"], metric, copy=False) + running: Optional[pa.Table] = None + for batch in _scanner_batches(scanner): + table = pa.Table.from_batches([batch]) + if table.num_rows == 0: + continue + table = table.filter(pc.invert(pc.is_null(table[vector_scan_column]))) + if table.num_rows == 0: + continue + + vector_matrix = _vector_column_to_numpy_for_metric( + table[vector_scan_column], metric + ) + query_results = [] + for query_index, query_vector in enumerate(query_vectors): + import numpy as np + + distances = _compute_core_vector_distances( + vector_matrix, + query_vector, + metric, + ) + finite = np.isfinite(distances) + query_result = table.filter(pa.array(finite, type=pa.bool_())) + query_result = query_result.append_column( + "_distance", + pa.array(distances[finite], type=pa.float32()), + ) + query_result = _apply_distance_range(query_result, nearest) + query_result = _take_top_k_deterministic(query_result, candidate_k) + if drop_vector_column and vector_scan_column in query_result.column_names: + query_result = query_result.drop_columns([vector_scan_column]) + query_result = query_result.add_column( + 0, + pa.field("query_index", pa.int32(), nullable=False), + pa.array( + [query_index] * query_result.num_rows, + type=pa.int32(), + ), + ) + query_results.append(query_result) + + current = pa.concat_tables(query_results, promote_options="default") + running = ( + current + if running is None + else _merge_vector_search_results( + [running, current], + candidate_k, + per_query=True, + ) + ) + + if running is not None: + return running + return _empty_fallback_table( + scanner, + vector_column=vector_scan_column, + drop_vector_column=drop_vector_column, + ) + + +def _indexed_search( + dataset: LanceDataset, + *, + index_segments: tuple[str, ...], + base_scanner_options: dict[str, Any], + nearest: dict[str, Any], + candidate_k: int, +) -> pa.Table: + if not _scanner_accepts_index_segments(dataset): + raise RuntimeError( + "The installed pylance scanner does not support index_segments" + ) + scanner_options = dict(base_scanner_options) + search_nearest = dict(nearest) + search_nearest["k"] = candidate_k + scanner_options.update( + nearest=search_nearest, + index_segments=index_segments, + fast_search=True, + ) + return dataset.scanner(**scanner_options).to_table() + + +def _offset_query_index( + table: pa.Table, + offset: int, + *, + output_type: pa.DataType, +) -> pa.Table: + if "query_index" not in table.column_names: + raise RuntimeError("Batch search result is missing query_index") + values = pc.cast(table["query_index"], output_type) + if offset: + values = pc.add(values, pa.scalar(offset, output_type)) + return table.set_column( + table.schema.get_field_index("query_index"), + pa.field("query_index", output_type, nullable=False), + values, + ) + + +def _canonical_query_batch( + query: Any, + metric: str, + *, + copy: bool = True, +) -> Any: + import numpy as np + + dtype = np.uint8 if metric == "hamming" else np.float32 + if isinstance(query, pa.RecordBatch | pa.Table): + if len(query.column_names) != 1: + raise ValueError( + "Arrow query batches must contain exactly one vector column" + ) + query = query.column(0) + if isinstance(query, pa.ChunkedArray): + query = query.combine_chunks() + if isinstance(query, pa.Array): + query = query.to_pylist() + if copy: + array = np.array(query, dtype=dtype, copy=True, order="C") + else: + array = np.asarray(query, dtype=dtype, order="C") + if array.size == 0: + if array.ndim == 2: + return array + return np.empty((0, 0), dtype=dtype) + if array.ndim == 1: + array = array.reshape(1, -1) + if array.ndim != 2: + raise ValueError("Each query batch must be a two-dimensional array") + return array + + +def _streaming_is_multivector_type(data_type: pa.DataType) -> bool: + if not (pa.types.is_list(data_type) or pa.types.is_large_list(data_type)): + return False + return pa.types.is_fixed_size_list(data_type.value_type) + + +def _canonical_multivector_batch( + query: Any, + metric: str, +) -> tuple[Any, ...]: + import numpy as np + + dtype = np.uint8 if metric == "hamming" else np.float32 + if isinstance(query, pa.RecordBatch | pa.Table): + if len(query.column_names) != 1: + raise ValueError( + "Arrow query batches must contain exactly one multivector column" + ) + query = query.column(0) + if isinstance(query, pa.ChunkedArray): + query = query.combine_chunks() + if isinstance(query, pa.Array): + query = query.to_pylist() + + try: + array = np.asarray(query, dtype=dtype) + except ValueError: + array = None + + if array is not None and array.ndim <= 3: + if array.size == 0: + return () + if array.ndim == 1: + return (np.array(array.reshape(1, -1), copy=True, order="C"),) + if array.ndim == 2: + return (np.array(array, copy=True, order="C"),) + return tuple(np.array(item, copy=True, order="C") for item in array) + + queries = [] + for item in query: + item_array = np.array(item, dtype=dtype, copy=True, order="C") + if item_array.ndim == 1: + item_array = item_array.reshape(1, -1) + if item_array.ndim != 2: + raise ValueError("Each multivector query must have shape [M, D]") + queries.append(item_array) + return tuple(queries) + + +def _add_query_index(table: pa.Table, query_index: int) -> pa.Table: + return table.add_column( + 0, + pa.field("query_index", pa.int32(), nullable=False), + pa.array([query_index] * table.num_rows, type=pa.int32()), + ) + + +def _multivector_fallback_search( + dataset: LanceDataset, + *, + fragment_ids: tuple[int, ...], + base_scanner_options: dict[str, Any], + nearest: dict[str, Any], + candidate_k: int, +) -> pa.Table: + query = nearest["q"] + scanner_options = dict(base_scanner_options) + scanner_options.pop("fast_search", None) + scanner_options["fragments"] = [ + dataset.get_fragment(fragment_id) for fragment_id in fragment_ids + ] + # Core requires prefilter for nearest scans scoped to explicit fragments. + scanner_options["prefilter"] = True + + search_nearest = {**nearest, "k": candidate_k} + distance_range = search_nearest.pop("distance_range", None) + scanner_options["nearest"] = search_nearest + table = dataset.scanner(**scanner_options).to_table() + + query_count = len(query) + if query_count > 1 and table.num_rows: + # Remove this compatibility offset once the minimum Core version uses + # M - sum(MaxSim) for flat multivector distance. + distances = pc.add( + table["_distance"], + pa.scalar(float(query_count - 1), pa.float32()), + ) + table = table.set_column( + table.schema.get_field_index("_distance"), + pa.field("_distance", pa.float32()), + distances, + ) + if distance_range is not None: + table = _apply_distance_range( + table, + {"distance_range": distance_range}, + ) + return _take_top_k_deterministic(table, candidate_k) + + +@ray.remote +class _VectorSearchActor: + def __init__( + self, + snapshot: _DatasetSnapshot, + plan: _ActorPlan, + base_scanner_options: dict[str, Any], + index_name: Optional[str], + is_multivector: bool, + actor_options: VectorSearchActorOptions, + ): + self._dataset = _open_snapshot( + snapshot, + index_cache_size_bytes=actor_options.index_cache_size_bytes, + metadata_cache_size_bytes=actor_options.metadata_cache_size_bytes, + ) + self._plan = plan + self._base_scanner_options = base_scanner_options + self._index_name = index_name + self._is_multivector = is_multivector + self._actor_options = actor_options + + def ready(self) -> dict[str, Any]: + return { + "version": self._dataset.version, + "index_segments": len(self._plan.index_segments), + "fallback_fragments": len(self._plan.fallback_fragment_ids), + } + + def prewarm(self) -> dict[str, Any]: + if not self._plan.index_segments or self._index_name is None: + return {"index_segments": 0, "skipped": True} + before = self._dataset.io_stats_snapshot() + self._dataset.prewarm_index( + self._index_name, + index_segments=self._plan.index_segments, + ) + after = self._dataset.io_stats_snapshot() + session = self._dataset.session() + return { + "index_segments": len(self._plan.index_segments), + "skipped": False, + "cache_entries": self._dataset._ds.index_cache_entry_count(), + "cache_size_bytes": session.index_cache_size_bytes(), + "cache_hit_rate": self._dataset._ds.index_cache_hit_rate(), + "read_bytes": after.read_bytes - before.read_bytes, + "read_iops": after.read_iops - before.read_iops, + } + + def search( + self, + query_batch: Any, + nearest: dict[str, Any], + candidate_k: int, + ) -> pa.Table: + metric = _get_nearest_metric(nearest) + if self._is_multivector: + queries = _canonical_multivector_batch(query_batch, metric) + else: + queries = _canonical_query_batch(query_batch, metric, copy=False) + micro_batch_size = self._actor_options.micro_batch_size or len(queries) + batches = [ + queries[offset : offset + micro_batch_size] + for offset in range(0, len(queries), micro_batch_size) + ] + + if self._actor_options.scanner_concurrency == 1: + results = [ + self._search_micro_batch(batch, nearest, candidate_k) + for batch in batches + ] + else: + with concurrent.futures.ThreadPoolExecutor( + max_workers=self._actor_options.scanner_concurrency + ) as pool: + results = list( + pool.map( + lambda batch: self._search_micro_batch( + batch, + nearest, + candidate_k, + ), + batches, + ) + ) + + offset = 0 + adjusted = [] + for batch, result in zip(batches, results, strict=True): + adjusted.append(_offset_query_index(result, offset, output_type=pa.int32())) + offset += len(batch) + return pa.concat_tables(adjusted, promote_options="default") + + def _search_micro_batch( + self, + query_batch: Any, + nearest: dict[str, Any], + candidate_k: int, + ) -> pa.Table: + if self._is_multivector: + results = [ + _add_query_index( + self._search_multivector_query(query, nearest, candidate_k), + query_index, + ) + for query_index, query in enumerate(query_batch) + ] + return pa.concat_tables(results, promote_options="default") + + search_nearest = {**nearest, "q": query_batch} + tables = [] + if self._plan.index_segments: + tables.append( + _indexed_search( + self._dataset, + index_segments=self._plan.index_segments, + base_scanner_options=self._base_scanner_options, + nearest=search_nearest, + candidate_k=candidate_k, + ) + ) + if self._plan.fallback_fragment_ids: + tables.append( + _stream_flat_fallback( + self._dataset, + fragment_ids=self._plan.fallback_fragment_ids, + base_scanner_options=self._base_scanner_options, + nearest=search_nearest, + candidate_k=candidate_k, + ) + ) + return _merge_vector_search_results( + tables, + candidate_k, + per_query=True, + ) + + def _search_multivector_query( + self, + query: Any, + nearest: dict[str, Any], + candidate_k: int, + ) -> pa.Table: + search_nearest = {**nearest, "q": query} + tables = [] + if self._plan.index_segments: + tables.append( + _indexed_search( + self._dataset, + index_segments=self._plan.index_segments, + base_scanner_options=self._base_scanner_options, + nearest=search_nearest, + candidate_k=candidate_k, + ) + ) + if self._plan.fallback_fragment_ids: + tables.append( + _multivector_fallback_search( + self._dataset, + fragment_ids=self._plan.fallback_fragment_ids, + base_scanner_options=self._base_scanner_options, + nearest=search_nearest, + candidate_k=candidate_k, + ) + ) + return _merge_vector_search_results( + tables, + candidate_k, + deterministic=True, + ) + + +class VectorSearchSession: + """A snapshot-pinned, actor-backed streaming vector search session.""" + + def __init__( + self, + *, + dataset: LanceDataset, + vector_type: pa.DataType, + snapshot: _DatasetSnapshot, + nearest: dict[str, Any], + index_name: Optional[str], + plans: list[_ActorPlan], + base_scanner_options: dict[str, Any], + include_row_id: bool, + global_k: int, + candidate_k: int, + streaming_options: VectorSearchStreamingOptions, + actor_options: VectorSearchActorOptions, + ): + self._dataset = dataset + self.vector_type = vector_type + self._is_multivector = _streaming_is_multivector_type(vector_type) + self._nearest = nearest + self._base_scanner_options = base_scanner_options + self._include_row_id = include_row_id + self._global_k = global_k + self._candidate_k = candidate_k + self._streaming_options = streaming_options + self._result_names: Optional[list[str]] = None + self._closed = False + self.actor_states: list[dict[str, Any]] = [] + self.prewarm_results: list[dict[str, Any]] = [] + + remote_args = dict(actor_options.ray_remote_args or {}) + remote_args.setdefault("num_cpus", 1) + remote_args["max_concurrency"] = actor_options.max_concurrent_batches + if actor_options.max_pending_calls is not None: + remote_args["max_pending_calls"] = actor_options.max_pending_calls + actor_class = _VectorSearchActor.options(**remote_args) + self._actors = [ + actor_class.remote( + snapshot, + plan, + base_scanner_options, + index_name, + self._is_multivector, + actor_options, + ) + for plan in plans + ] + try: + if self._actors: + self.actor_states = ray.get( + [actor.ready.remote() for actor in self._actors] + ) + if actor_options.prewarm_index and self._actors: + self.prewarm_results = ray.get( + [actor.prewarm.remote() for actor in self._actors] + ) + except Exception: + self.close() + raise + + def __enter__(self) -> "VectorSearchSession": + return self + + def __exit__(self, exc_type: Any, exc: Any, traceback: Any) -> None: + self.close() + + def close(self) -> None: + if self._closed: + return + self._closed = True + for actor in self._actors: + ray.kill(actor, no_restart=True) + self._actors.clear() + + def map_batches(self, query_batches: Iterable[Any]) -> Iterator[pa.Table]: + """Search an iterable of query batches with bounded memory. + + The driver canonicalizes each input batch, places it in Ray's object + store once, broadcasts the resulting reference to all search actors, + merges their local candidates per query, and yields the completed + global top-k table. At most + ``streaming_options.max_in_flight_batches`` batches are retained. + + Args: + query_batches: Iterable of regular vector batches shaped ``[B, D]`` + or multivector batches described by :func:`open_vector_search`. + + Yields: + PyArrow tables in input-batch order. ``query_index`` is an Int64 + position in the complete stream, not an index local to the batch. + """ + if self._closed: + raise RuntimeError("VectorSearchSession is closed") + + pending: deque[tuple[int, Any, list[Any]]] = deque() + global_offset = 0 + + def complete_one() -> pa.Table: + offset, query_batch, refs = pending.popleft() + tables = ray.get(refs) if refs else [] + result = self._finish_batch( + tables, + query_batch=query_batch, + global_offset=offset, + ) + return result + + for query_batch in self._iter_batches(query_batches): + query_count = len(query_batch) + if query_count == 0: + continue + while len(pending) >= self._streaming_options.max_in_flight_batches: + yield complete_one() + + query_ref = ray.put(query_batch) + refs = [ + actor.search.remote( + query_ref, + self._nearest, + self._candidate_k, + ) + for actor in self._actors + ] + pending.append((global_offset, query_batch, refs)) + global_offset += query_count + + while pending: + yield complete_one() + + def _iter_batches(self, query_batches: Iterable[Any]) -> Iterator[Any]: + import numpy as np + + metric = _get_nearest_metric(self._nearest) + target = self._streaming_options.query_batch_size + if self._is_multivector: + buffered_queries = [] + for query_batch in query_batches: + canonical = _canonical_multivector_batch(query_batch, metric) + if target is None: + yield canonical + continue + buffered_queries.extend(canonical) + while len(buffered_queries) >= target: + yield tuple(buffered_queries[:target]) + del buffered_queries[:target] + if buffered_queries: + yield tuple(buffered_queries) + return + + buffered = [] + buffered_rows = 0 + for query_batch in query_batches: + canonical = _canonical_query_batch(query_batch, metric) + if target is None: + yield canonical + continue + offset = 0 + while offset < len(canonical): + take = min(target - buffered_rows, len(canonical) - offset) + buffered.append(canonical[offset : offset + take]) + buffered_rows += take + offset += take + if buffered_rows == target: + yield np.concatenate(buffered, axis=0) + buffered.clear() + buffered_rows = 0 + if buffered: + yield np.concatenate(buffered, axis=0) + + def _finish_batch( + self, + tables: list[pa.Table], + *, + query_batch: Any, + global_offset: int, + ) -> pa.Table: + schema_query = query_batch[0] if self._is_multivector else query_batch + _, result_schema = _inspect_vector_search_query( + self._dataset, + nearest={**self._nearest, "q": schema_query}, + base_scanner_options=self._base_scanner_options, + include_row_id=self._include_row_id, + ) + if self._is_multivector: + result_schema = pa.schema( + [ + pa.field("query_index", pa.int32(), nullable=False), + *result_schema, + ] + ) + + if tables: + result = _merge_vector_search_results( + tables, + self._global_k, + per_query=True, + ) + else: + result = pa.Table.from_batches([], schema=result_schema) + + if self._result_names is None: + self._result_names = result_schema.names + result = _offset_query_index( + result, + global_offset, + output_type=pa.int64(), + ) + if not self._include_row_id and "_rowid" in result.column_names: + result = result.drop_columns(["_rowid"]) + return result.select(self._result_names) + + +def _build_driver_dataset( + uri: str | LanceDataset | None, + *, + storage_options: Optional[dict[str, Any]], + base_store_params: Optional[dict[str, dict[str, Any]]], + block_size: Optional[int], + namespace_impl: Optional[str], + namespace_properties: Optional[dict[str, str]], + table_id: Optional[list[str]], + branch: Optional[str], + version: int | str | None, +) -> tuple[LanceDataset, _DatasetSnapshot]: + if branch is not None and version is not None: + raise ValueError("branch and version are mutually exclusive") + + merged_storage_options = dict(storage_options or {}) + if isinstance(uri, LanceDataset): + dataset = uri + dataset_uri = dataset.uri + if branch is not None: + dataset = dataset.checkout_version((branch, None)) + dataset_uri = ( + f"{dataset_uri.partition('/tree/')[0].rstrip('/')}/tree/{branch}" + ) + elif version is not None: + dataset = dataset.checkout_version(version) + if not merged_storage_options: + merged_storage_options.update(_get_dataset_storage_options(dataset)) + else: + validate_uri_or_namespace(uri, namespace_impl, table_id) + dataset_uri, merged_storage_options = resolve_namespace_table( + uri, + storage_options, + namespace_impl, + namespace_properties, + table_id, + ) + kwargs: dict[str, Any] = { + "storage_options": merged_storage_options, + **get_namespace_kwargs( + namespace_impl, + namespace_properties, + table_id, + ), + } + if block_size is not None: + kwargs["block_size"] = block_size + if base_store_params is not None: + kwargs["base_store_params"] = base_store_params + dataset = LanceDataset(dataset_uri, **kwargs) + if branch is not None: + dataset = dataset.checkout_version((branch, None)) + dataset_uri = f"{dataset_uri.rstrip('/')}/tree/{branch}" + elif version is not None: + dataset = dataset.checkout_version(version) + + snapshot = _DatasetSnapshot( + uri=dataset_uri, + version=dataset.version, + serialized_manifest=dataset._ds.serialized_manifest(), + storage_options=merged_storage_options, + base_store_params=base_store_params, + block_size=block_size, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + ) + return dataset, snapshot + + +def open_vector_search( + uri: str | LanceDataset | None = None, + *, + nearest: dict[str, Any], + index_name: Optional[str] = None, + columns: Optional[list[str] | dict[str, str]] = None, + filter: Optional[Any] = None, + storage_options: Optional[dict[str, Any]] = None, + base_store_params: Optional[dict[str, dict[str, Any]]] = None, + block_size: Optional[int] = None, + namespace_impl: Optional[str] = None, + namespace_properties: Optional[dict[str, str]] = None, + table_id: Optional[list[str]] = None, + branch: Optional[str] = None, + version: int | str | None = None, + oversample_factor: float = 1.0, + fast_search: bool = False, + scanner_options: Optional[dict[str, Any]] = None, + streaming_options: Optional[VectorSearchStreamingOptions] = None, + actor_options: Optional[VectorSearchActorOptions] = None, +) -> VectorSearchSession: + """Open a reusable distributed vector search session. + + The session pins the dataset manifest when it opens, assigns index segments + and uncovered fragments to persistent Ray actors, and reuses each actor's + Lance session and index cache across query batches. Use the returned object + as a context manager so the actors are stopped when the stream finishes. + + Queries are supplied later through :meth:`VectorSearchSession.map_batches`; + do not include ``q`` in ``nearest``. For a fixed-size vector column, each + input batch is an array with shape ``[B, D]``. For a multivector column, an + input batch is a sequence of ``[M_i, D]`` arrays, an Arrow + ``List>`` array, or a ``[B, M, D]`` array when ``M`` is + fixed. Every output table contains an Int64 ``query_index`` that identifies + the query's position across the entire input stream. + + Args: + uri: Lance dataset object or dataset URI. In URI mode, provide either + ``uri`` or namespace parameters (``namespace_impl`` + ``table_id``). + An already checked-out dataset retains its exact manifest. + nearest: Lance nearest-neighbor options without ``q``. ``column`` and + ``k`` are required. Options such as ``nprobes``, + ``query_parallelism``, ``approx_mode``, ``refine_factor``, and + ``distance_range`` are forwarded to Lance Core. + index_name: Optional vector index name. If omitted, the first vector + index covering ``nearest["column"]`` is selected. + columns: Columns or projection expressions returned for each match. + ``_distance`` is added when needed for distributed top-k merging. + filter: Filter passed to each actor's Lance scanner. + storage_options: Storage options used to open the dataset. + base_store_params: Runtime options for registered external base paths. + block_size: Optional dataset I/O block size in bytes. + namespace_impl: Namespace implementation, such as ``"dir"`` or + ``"rest"``. + namespace_properties: Properties used to connect to the namespace. + table_id: Table identifier used with namespace parameters. + branch: Branch resolved and pinned when the session opens. Mutually + exclusive with ``version``. + version: Dataset version or tag to pin. Mutually exclusive with + ``branch``. + oversample_factor: Multiplier applied to each actor's local candidate + count before the driver performs the global top-k merge. + fast_search: If true, intentionally skip fragments not covered by the + selected vector index. If false, include them through flat fallback. + scanner_options: Additional Lance scanner options. ``nearest``, + ``fragments``, ``index_segments``, ``fast_search``, ``limit``, and + ``offset`` are managed by Lance-Ray and cannot be supplied here. + streaming_options: Input rebatching and bounded in-flight pipeline + settings. + actor_options: Actor count, Ray resources, micro-batching, scanner + concurrency, cache sizes, pending-call limits, and optional index + prewarming. + + Returns: + A snapshot-pinned :class:`VectorSearchSession`. + + Example: + >>> with open_vector_search( + ... "dataset.lance", + ... nearest={"column": "vector", "k": 10, "nprobes": 8}, + ... streaming_options=VectorSearchStreamingOptions( + ... query_batch_size=512, + ... max_in_flight_batches=2, + ... ), + ... ) as search: + ... for result in search.map_batches(query_batches): + ... write_result(result) + """ + streaming_options = streaming_options or VectorSearchStreamingOptions() + actor_options = actor_options or VectorSearchActorOptions() + + if block_size is not None and block_size <= 0: + raise ValueError(f"block_size must be positive, got {block_size}") + if "q" in nearest: + raise ValueError("open_vector_search receives queries through map_batches") + if not nearest.get("column"): + raise ValueError("nearest must include 'column'") + nearest = dict(nearest) + global_k, candidate_k = _candidate_k(nearest, oversample_factor) + + base_scanner_options = dict(scanner_options or {}) + _validate_search_scanner_options(base_scanner_options) + include_row_id = _projection_includes_row_id(columns, base_scanner_options) + effective_columns = ( + columns if columns is not None else base_scanner_options.get("columns") + ) + if effective_columns is not None: + if ( + isinstance(effective_columns, list) and "query_index" in effective_columns + ) or ( + isinstance(effective_columns, dict) and "query_index" in effective_columns + ): + raise ValueError( + "query_index is managed by streaming vector search and cannot " + "be included in columns" + ) + if isinstance(effective_columns, list) and "_distance" not in effective_columns: + effective_columns = [*effective_columns, "_distance"] + base_scanner_options["columns"] = effective_columns + if filter is not None: + base_scanner_options["filter"] = filter + base_scanner_options["with_row_id"] = True + + dataset, snapshot = _build_driver_dataset( + uri, + storage_options=storage_options, + base_store_params=base_store_params, + block_size=block_size, + namespace_impl=namespace_impl, + namespace_properties=namespace_properties, + table_id=table_id, + branch=branch, + version=version, + ) + if "query_index" in dataset.schema.names: + raise ValueError( + "Batch vector search cannot use a dataset containing column 'query_index'" + ) + try: + resolved_field = resolve_arrow_field_path( + dataset.schema, + nearest["column"], + ) + except KeyError as exc: + available_columns = [field.name for field in dataset.schema] + raise ValueError( + f"Column '{nearest['column']}' not found. Available: {available_columns}" + ) from exc + resolved_column = resolved_field.path + nearest = {**nearest, "column": resolved_column} + + vector_index = _select_vector_index( + dataset, + column=resolved_column, + index_name=index_name, + ) + nearest = _apply_index_metric_default(nearest, vector_index) + resolved_index_name = ( + str(_index_value(vector_index, "name")) if vector_index is not None else None + ) + plans = _plan_streaming_vector_search( + fragments=dataset.get_fragments(), + vector_index=vector_index, + num_actors=actor_options.num_actors, + fast_search=fast_search, + ) + + return VectorSearchSession( + dataset=dataset, + vector_type=resolved_field.field.type, + snapshot=snapshot, + nearest=nearest, + index_name=resolved_index_name, + plans=plans, + base_scanner_options=base_scanner_options, + include_row_id=include_row_id, + global_k=global_k, + candidate_k=candidate_k, + streaming_options=streaming_options, + actor_options=actor_options, + ) + + +def _is_multivector_type(data_type: pa.DataType) -> bool: + if not (pa.types.is_list(data_type) or pa.types.is_large_list(data_type)): + return False + return ( + pa.types.is_fixed_size_list(data_type.value_type) + or pa.types.is_list(data_type.value_type) + or pa.types.is_large_list(data_type.value_type) ) - if not include_row_id and "_rowid" in result.column_names: - result = result.drop_columns(["_rowid"]) - return result.select(result_schema.names) diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index 3117fade..acab92bc 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -1921,14 +1921,15 @@ def test_build_distributed_vector_index(tmp_path, index_type): assert index_name in plan queries = np.asarray([q, q], dtype=np.float32) - batch = lr.vector_search( + with lr.open_vector_search( updated_dataset, - nearest={"column": "vector", "q": queries, "k": 5}, + nearest={"column": "vector", "k": 5}, index_name=index_name, columns=["id"], - num_workers=2, fast_search=True, - ) + actor_options=lr.VectorSearchActorOptions(num_actors=2), + ) as search: + [batch] = list(search.map_batches([queries])) assert batch.column("query_index").to_pylist() == [0] * 5 + [1] * 5 assert ( diff --git a/tests/test_distributed_search.py b/tests/test_distributed_search.py index 62ed730f..cceb8cdc 100644 --- a/tests/test_distributed_search.py +++ b/tests/test_distributed_search.py @@ -4,17 +4,20 @@ import lance_ray as lr import numpy as np import pyarrow as pa -import pyarrow.compute as pc import pytest from lance_ray import pool as pool_mod from lance_ray import search as search_mod from lance_ray.search import ( + VectorSearchActorOptions, + VectorSearchStreamingOptions, _apply_distance_range, - _apply_index_metric_default, - _compute_vector_distances, + _canonical_multivector_batch, + _canonical_query_batch, + _compute_core_vector_distances, _execute_vector_search_plan, _format_analyze_plan_results, _merge_vector_search_results, + _plan_streaming_vector_search, _plan_vector_search, _SearchPlan, _SearchPlanAnalysis, @@ -44,24 +47,10 @@ def _index_with_segments(*segments): ) -def _mock_pickled_dataset(monkeypatch, dataset): - search_mod._load_pickled_dataset.cache_clear() - search_mod._load_pickled_dataset_ref.cache_clear() - pickled_dataset = f"pickled-dataset-{id(dataset)}".encode() - - def fake_loads(value): - assert value == pickled_dataset - return dataset - - monkeypatch.setattr(search_mod.pickle, "loads", fake_loads) - return pickled_dataset - - -def _vector_table(vectors, ids=None, *, value_type=None): - matrix = np.asarray(vectors) - value_type = value_type or pa.float32() +def _vector_table(vectors, ids=None): + matrix = np.asarray(vectors, dtype=np.float32) vector_array = pa.FixedSizeListArray.from_arrays( - pa.array(matrix.reshape(-1), type=value_type), + pa.array(matrix.reshape(-1), type=pa.float32()), matrix.shape[1], ) return pa.table( @@ -72,50 +61,17 @@ def _vector_table(vectors, ids=None, *, value_type=None): ) -class _FallbackDataset: - def __init__(self, table, scanner_options=None): - self.table = table - self.scanner_options = scanner_options - - def get_fragment(self, fragment_id): - return f"fragment-{fragment_id}" - - def scanner(self, **kwargs): - if self.scanner_options is not None: - self.scanner_options.update(kwargs) - return SimpleNamespace(to_table=lambda: self.table) +def _mock_pickled_dataset(monkeypatch, dataset): + search_mod._load_pickled_dataset.cache_clear() + search_mod._load_pickled_dataset_ref.cache_clear() + pickled_dataset = f"pickled-dataset-{id(dataset)}".encode() + def fake_loads(value): + assert value == pickled_dataset + return dataset -def _create_partial_index_dataset( - path, - indexed_vectors, - appended_vectors, - *, - metric="l2", -): - dataset = lance.write_dataset( - _vector_table(indexed_vectors), - path, - max_rows_per_file=2, - ) - dataset.create_index( - "vector", - "IVF_FLAT", - num_partitions=1, - name="vector_idx", - metric=metric, - ) - lance.write_dataset( - _vector_table( - appended_vectors, - ids=range( - len(indexed_vectors), len(indexed_vectors) + len(appended_vectors) - ), - ), - path, - mode="append", - ) - return lance.dataset(path) + monkeypatch.setattr(search_mod.pickle, "loads", fake_loads) + return pickled_dataset def test_select_vector_index_raises_for_missing_explicit_index_name(): @@ -303,14 +259,27 @@ def scanner(self, columns=None): def test_execute_fallback_vector_search_plan_computes_local_top_k(monkeypatch): scanner_options = {} - dataset = _FallbackDataset( - _vector_table([[10.0, 0.0], [1.0, 0.0], [0.0, 2.0]], ids=[1, 2, 3]), - scanner_options, + vectors = pa.FixedSizeListArray.from_arrays( + pa.array([10.0, 0.0, 1.0, 0.0, 0.0, 2.0], type=pa.float32()), + 2, ) + class FakeDataset: + def __init__(self, *args, **kwargs): + pass + + def get_fragment(self, fragment_id): + return f"fragment-{fragment_id}" + + def scanner(self, **kwargs): + scanner_options.update(kwargs) + return SimpleNamespace( + to_table=lambda: pa.table({"id": [1, 2, 3], "vector": vectors}) + ) + result = _execute_vector_search_plan( _SearchPlan(fragment_ids=[7], index_segments=[]), - pickled_dataset=_mock_pickled_dataset(monkeypatch, dataset), + pickled_dataset=_mock_pickled_dataset(monkeypatch, FakeDataset()), base_scanner_options={"columns": ["id", "_distance"], "fast_search": False}, nearest={"column": "vector", "q": [0.0, 0.0], "k": 2}, candidate_k=2, @@ -321,118 +290,10 @@ def test_execute_fallback_vector_search_plan_computes_local_top_k(monkeypatch): assert scanner_options["fragments"] == ["fragment-7"] assert scanner_options["columns"] == ["id", "vector"] assert result.column("id").to_pylist() == [2, 3] - assert result.column("_distance").to_pylist() == [1.0, 4.0] + assert result.column("_distance").to_pylist() == [1.0, 2.0] assert "vector" not in result.column_names -def test_execute_batch_hamming_fallback_uses_bit_distance(monkeypatch): - conversion_calls = 0 - dataset = _FallbackDataset(_vector_table([[0], [3], [255]], value_type=pa.uint8())) - original_vector_column_to_numpy = search_mod._vector_column_to_numpy - - def count_vector_column_conversion(vector_column, metric): - nonlocal conversion_calls - conversion_calls += 1 - return original_vector_column_to_numpy(vector_column, metric) - - monkeypatch.setattr( - search_mod, - "_vector_column_to_numpy", - count_vector_column_conversion, - ) - - result = _execute_vector_search_plan( - _SearchPlan(fragment_ids=[7], index_segments=[]), - pickled_dataset=_mock_pickled_dataset(monkeypatch, dataset), - base_scanner_options={"columns": ["id", "_distance"], "fast_search": False}, - nearest={ - "column": "vector", - "q": [[0], [255]], - "k": 3, - "metric": "hamming", - }, - candidate_k=3, - analyze_plan=False, - is_batch_query=True, - ) - - assert result.column("query_index").to_pylist() == [0, 0, 0, 1, 1, 1] - assert result.column("id").to_pylist() == [0, 1, 2, 2, 1, 0] - assert result.column("_distance").to_pylist() == [0.0, 2.0, 8.0, 0.0, 6.0, 8.0] - assert conversion_calls == 1 - - -def test_apply_distance_range_uses_inclusive_lower_exclusive_upper(): - table = pa.table({"id": [0, 1, 2], "_distance": [0.5, 1.0, 4.0]}) - - result = _apply_distance_range(table, {"distance_range": (0.5, 4.0)}) - - assert result.column("id").to_pylist() == [0, 1] - assert result.column("_distance").to_pylist() == [0.5, 1.0] - - -@pytest.mark.parametrize( - ("metric", "vectors", "query", "expected"), - [ - ("l2", [[0.0, 2.0], [3.0, 4.0]], [0.0, 0.0], [4.0, 25.0]), - ( - "cosine", - [[1.0, 0.0], [1.0, 1.0]], - [1.0, 0.0], - [0.0, 0.29289323], - ), - ("dot", [[1.0, 0.0], [1.0, 1.0]], [1.0, 1.0], [0.0, -1.0]), - ( - "hamming", - [[0, 0], [255, 0], [15, 240], [1, 2]], - [0, 0], - [0.0, 8.0, 8.0, 2.0], - ), - ], -) -def test_fallback_distance_conventions(metric, vectors, query, expected): - dtype = np.uint8 if metric == "hamming" else np.float32 - distances = _compute_vector_distances( - np.asarray(vectors, dtype=dtype), query, metric - ) - - assert distances.tolist() == pytest.approx(expected) - - -def test_execute_fallback_filters_null_and_invalid_cosine_vectors(monkeypatch): - table = pa.table( - { - "id": [0, 1, 2, 3], - "vector": pa.array( - [[0.0, 0.0], None, [1.0, 0.0], [0.0, 1.0]], - type=pa.list_(pa.float32(), 2), - ), - "_rowid": [0, 1, 2, 3], - } - ) - result = _execute_vector_search_plan( - _SearchPlan(fragment_ids=[7], index_segments=[]), - pickled_dataset=_mock_pickled_dataset( - monkeypatch, - _FallbackDataset(table), - ), - base_scanner_options={"fast_search": False, "with_row_id": True}, - nearest={ - "column": "vector", - "q": [[1.0, 0.0], [0.0, 0.0]], - "k": 4, - "metric": "cosine", - }, - candidate_k=4, - analyze_plan=False, - is_batch_query=True, - ) - - assert result.column("query_index").to_pylist() == [0, 0] - assert result.column("id").to_pylist() == [2, 3] - assert result.column("_distance").to_pylist() == [0.0, 1.0] - - def test_execute_indexed_vector_search_plan_can_analyze_plan(monkeypatch): scanner_options = {} @@ -533,93 +394,61 @@ def test_format_analyze_plan_results(): def test_merge_vector_search_results_returns_global_top_k(): - left = pa.table({"id": [1, 2], "_distance": [0.4, 0.1], "_rowid": [40, 20]}) - right = pa.table({"id": [3, 4], "_distance": [0.1, 0.3], "_rowid": [10, 30]}) + left = pa.table({"id": [1, 2], "_distance": [0.4, 0.1]}) + right = pa.table({"id": [3, 4], "_distance": [0.2, 0.3]}) result = _merge_vector_search_results([left, right], k=3) - assert result.column("id").to_pylist() == [3, 2, 4] - assert result.column("_distance").to_pylist() == [0.1, 0.1, 0.3] + assert result.column("id").to_pylist() == [2, 3, 4] + assert result.column("_distance").to_pylist() == [0.1, 0.2, 0.3] + +def test_merge_vector_search_results_requires_distance(): + table = pa.table({"id": [1, 2]}) -def test_merge_batch_vector_search_results_returns_top_k_per_query(): + with pytest.raises(RuntimeError, match="_distance"): + _merge_vector_search_results([table], k=1) + + +def test_merge_vector_search_results_can_merge_per_query(): left = pa.table( { - "query_index": pa.array([0, 0, 1], type=pa.int32()), + "query_index": [0, 0, 1], "id": [1, 2, 3], - "_distance": [0.4, 0.1, 0.2], - "_rowid": [40, 10, 30], + "_distance": [0.4, 0.1, 0.3], } ) right = pa.table( { - "query_index": pa.array([0, 1, 1], type=pa.int32()), + "query_index": [0, 1, 1], "id": [4, 5, 6], - "_distance": [0.2, 0.3, 0.1], - "_rowid": [20, 50, 60], + "_distance": [0.2, 0.4, 0.1], } ) - result = _merge_vector_search_results( - [left, right], - k=2, - is_batch_query=True, - ) + result = _merge_vector_search_results([left, right], k=2, per_query=True) - assert result.column("query_index").to_pylist() == [0, 0, 1, 1] - assert result.column("id").to_pylist() == [2, 4, 6, 3] - assert result.column("_distance").to_pylist() == [0.1, 0.2, 0.1, 0.2] + assert result["query_index"].to_pylist() == [0, 0, 1, 1] + assert result["id"].to_pylist() == [2, 4, 6, 3] -@pytest.mark.parametrize( - ("table", "is_batch_query", "missing_column"), - [ - (pa.table({"id": [1, 2]}), False, "_distance"), - (pa.table({"id": [1], "_distance": [0.1]}), True, "query_index"), - ], -) -def test_merge_vector_search_results_requires_managed_columns( - table, - is_batch_query, - missing_column, -): - with pytest.raises(RuntimeError, match=missing_column): - _merge_vector_search_results( - [table], - k=1, - is_batch_query=is_batch_query, - ) +def test_search_scanner_options_reject_managed_options(): + with pytest.raises(ValueError, match="nearest"): + _validate_search_scanner_options({"nearest": {"column": "vector"}}) -@pytest.mark.parametrize( - "scanner_options", - [ - {"nearest": {"column": "vector"}}, - {"fast_search": True}, - ], - ids=["nearest", "fast_search"], -) -def test_search_scanner_options_reject_managed_options(scanner_options): - managed_option = next(iter(scanner_options)) - with pytest.raises(ValueError, match=managed_option): - _validate_search_scanner_options(scanner_options) +def test_search_scanner_options_reject_fast_search_override(): + with pytest.raises(ValueError, match="fast_search"): + _validate_search_scanner_options({"fast_search": True}) -def test_batch_vector_search_reuses_global_pool_in_one_round(monkeypatch): +def test_vector_search_reuses_global_pool(monkeypatch): events = [] class FakeAsyncResult: def get(self): events.append("get") - return [ - pa.table( - { - "query_index": pa.array([0, 1], type=pa.int32()), - "id": [1, 2], - "_distance": [0.1, 0.2], - } - ) - ] + return [pa.table({"id": [1], "_distance": [0.1]})] class FakeGlobalPool: def map_async(self, func, plans, chunksize): @@ -655,20 +484,6 @@ def get_fragments(self): lambda *args, **kwargs: object(), ) monkeypatch.setattr(search_mod, "_plan_vector_search", lambda **kwargs: [plan]) - monkeypatch.setattr( - search_mod, - "_inspect_vector_search_query", - lambda *args, **kwargs: ( - True, - pa.schema( - [ - pa.field("query_index", pa.int32(), nullable=False), - pa.field("id", pa.int64()), - pa.field("_distance", pa.float32()), - ] - ), - ), - ) monkeypatch.setattr(search_mod.pickle, "dumps", lambda dataset: b"pickled-dataset") monkeypatch.setattr(search_mod.ray, "is_initialized", lambda: False) @@ -676,20 +491,466 @@ def get_fragments(self): try: result = search_mod.vector_search( uri="dataset", - nearest={"column": "vector", "q": [[0.0], [1.0]], "k": 1}, + nearest={"column": "vector", "q": [0.0], "k": 1}, num_workers=4, ) finally: pool_mod.clear_global_pool() - assert result.column("query_index").to_pylist() == [0, 1] - assert result.column("id").to_pylist() == [1, 2] + assert result.column("id").to_pylist() == [1] assert events == [ ("map_async", [plan], 1), "get", ] +def test_streaming_option_defaults(): + assert VectorSearchStreamingOptions() == VectorSearchStreamingOptions( + query_batch_size=None, + max_in_flight_batches=1, + ) + assert VectorSearchActorOptions() == VectorSearchActorOptions( + num_actors=4, + ray_remote_args=None, + max_concurrent_batches=1, + max_pending_calls=None, + micro_batch_size=None, + scanner_concurrency=1, + index_cache_size_bytes=None, + metadata_cache_size_bytes=None, + prewarm_index=False, + ) + assert VectorSearchActorOptions( + index_cache_size_bytes=0, + metadata_cache_size_bytes=0, + ) + + +def test_open_vector_search_requires_explicit_k(): + with pytest.raises(ValueError, match="nearest must include 'k'"): + lr.open_vector_search(nearest={"column": "vector"}) + + +@pytest.mark.parametrize( + ("metric", "vectors", "query", "expected"), + [ + ("l2", [[0.0, 2.0], [3.0, 4.0]], [0.0, 0.0], [4.0, 25.0]), + ( + "cosine", + [[1.0, 0.0], [1.0, 1.0]], + [1.0, 0.0], + [0.0, 0.29289323], + ), + ("dot", [[1.0, 0.0], [1.0, 1.0]], [1.0, 1.0], [0.0, -1.0]), + ( + "hamming", + [[0, 0], [255, 0], [15, 240], [1, 2]], + [0, 0], + [0.0, 8.0, 8.0, 2.0], + ), + ], +) +def test_streaming_fallback_distance_matches_core(metric, vectors, query, expected): + dtype = np.uint8 if metric == "hamming" else np.float32 + + distances = _compute_core_vector_distances( + np.asarray(vectors, dtype=dtype), + query, + metric, + ) + + assert distances.tolist() == pytest.approx(expected) + + +def test_streaming_distance_range_is_lower_inclusive_upper_exclusive(): + table = pa.table({"id": [0, 1, 2], "_distance": [0.5, 1.0, 4.0]}) + + result = _apply_distance_range(table, {"distance_range": (0.5, 4.0)}) + + assert result["id"].to_pylist() == [0, 1] + + +@pytest.mark.parametrize( + "columns", + [ + ["id", "query_index"], + {"query_index": "id"}, + ], +) +def test_open_vector_search_rejects_query_index_projection(columns): + with pytest.raises(ValueError, match="query_index is managed"): + lr.open_vector_search( + nearest={"column": "vector", "k": 10}, + columns=columns, + ) + + +def test_open_vector_search_rejects_dataset_query_index_column(tmp_path): + table = _vector_table([[0.0, 0.0], [1.0, 0.0]]) + table = table.append_column("query_index", pa.array([1, 2], type=pa.int32())) + dataset = lance.write_dataset(table, tmp_path / "query-index.lance") + + with pytest.raises(ValueError, match="containing column 'query_index'"): + lr.open_vector_search( + dataset, + nearest={"column": "vector", "k": 1}, + ) + + +def test_streaming_query_batches_are_canonicalized_by_column_type(): + source = np.arange(12, dtype=np.float32).reshape(3, 4)[:, ::-1] + regular = _canonical_query_batch(source, "l2") + + assert regular.flags.c_contiguous + assert regular.flags.owndata + assert not np.shares_memory(regular, source) + + multivector = _canonical_multivector_batch( + [ + np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32), + np.asarray([[1.0, 1.0]], dtype=np.float32), + ], + "cosine", + ) + assert [query.shape for query in multivector] == [(2, 2), (1, 2)] + assert all(query.flags.c_contiguous for query in multivector) + + +def test_streaming_planner_balances_indexed_and_fallback_units(): + fragments = [ + _FakeFragment(1, 100), + _FakeFragment(2, 90), + _FakeFragment(3, 80), + ] + plans = _plan_streaming_vector_search( + fragments=fragments, + vector_index=_index_with_segments(("S1", [1]), ("S2", [2])), + num_actors=2, + fast_search=False, + ) + + assert len(plans) == 2 + assert {segment for plan in plans for segment in plan.index_segments} == { + "S1", + "S2", + } + assert { + fragment_id for plan in plans for fragment_id in plan.fallback_fragment_ids + } == {3} + assert any(plan.index_segments and plan.fallback_fragment_ids for plan in plans) + + +def test_streaming_fallback_preserves_global_query_indices(tmp_path): + dataset = lance.write_dataset( + _vector_table( + [ + [0.0, 0.0], + [1.0, 0.0], + [0.0, 2.0], + [3.0, 0.0], + [0.0, 4.0], + [5.0, 0.0], + ] + ), + tmp_path / "streaming-flat.lance", + max_rows_per_file=2, + ) + + with lr.open_vector_search( + dataset, + nearest={"column": "vector", "k": 2}, + columns=["id"], + actor_options=VectorSearchActorOptions(num_actors=2), + streaming_options=VectorSearchStreamingOptions(max_in_flight_batches=2), + ) as session: + results = list( + session.map_batches( + [ + np.asarray([[0.0, 0.0], [0.0, 4.0]], dtype=np.float32), + np.asarray([[3.0, 0.0]], dtype=np.float32), + ] + ) + ) + + assert [result["query_index"].to_pylist() for result in results] == [ + [0, 0, 1, 1], + [2, 2], + ] + assert results[0]["query_index"].type == pa.int64() + assert results[0]["id"].to_pylist() == [0, 1, 4, 2] + assert results[1]["id"].to_pylist() == [3, 1] + + +def test_streaming_fast_search_without_index_returns_empty_result(tmp_path): + dataset = lance.write_dataset( + _vector_table([[0.0, 0.0], [1.0, 0.0]]), + tmp_path / "streaming-fast.lance", + ) + + with lr.open_vector_search( + dataset, + nearest={"column": "vector", "k": 1}, + columns=["id"], + fast_search=True, + ) as session: + [result] = list( + session.map_batches( + [np.asarray([[0.0, 0.0], [1.0, 0.0]], dtype=np.float32)] + ) + ) + + assert result.num_rows == 0 + assert result.column_names == ["query_index", "id", "_distance"] + assert result["query_index"].type == pa.int64() + + +def test_streaming_partial_index_merges_fallback_results(tmp_path): + path = tmp_path / "streaming-partial.lance" + dataset = lance.write_dataset( + _vector_table([[0.0, 0.0], [1.0, 0.0], [0.0, 2.0], [3.0, 0.0]]), + path, + max_rows_per_file=2, + ) + dataset.create_index( + "vector", + "IVF_FLAT", + num_partitions=1, + name="vector_idx", + ) + lance.write_dataset( + _vector_table([[0.0, 4.0], [5.0, 0.0]], ids=[4, 5]), + path, + mode="append", + ) + + with lr.open_vector_search( + lance.dataset(path), + nearest={"column": "vector", "k": 2, "nprobes": 1}, + index_name="vector_idx", + columns=["id"], + actor_options=VectorSearchActorOptions(num_actors=2), + ) as session: + [result] = list( + session.map_batches( + [np.asarray([[0.0, 4.0], [3.0, 0.0]], dtype=np.float32)] + ) + ) + + assert result["query_index"].to_pylist() == [0, 0, 1, 1] + assert result["id"].to_pylist() == [4, 2, 3, 1] + + +def test_streaming_fallback_inherits_index_metric(tmp_path): + path = tmp_path / "streaming-partial-cosine.lance" + dataset = lance.write_dataset( + _vector_table([[1.0, 0.0], [1.0, 1.0], [0.0, 1.0], [-1.0, 0.0]]), + path, + max_rows_per_file=2, + ) + dataset.create_index( + "vector", + "IVF_FLAT", + num_partitions=1, + name="vector_idx", + metric="cosine", + ) + lance.write_dataset( + _vector_table([[0.5, 0.5], [-1.0, -1.0]], ids=[4, 5]), + path, + mode="append", + ) + + with lr.open_vector_search( + lance.dataset(path), + nearest={"column": "vector", "k": 3, "nprobes": 1}, + index_name="vector_idx", + columns=["id"], + actor_options=VectorSearchActorOptions(num_actors=2), + ) as session: + [result] = list( + session.map_batches( + [np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32)] + ) + ) + + assert result["query_index"].to_pylist() == [0, 0, 0, 1, 1, 1] + assert result["id"].to_pylist() == [0, 1, 4, 2, 1, 4] + assert result["_distance"].to_pylist() == pytest.approx( + [0.0, 0.29289323, 0.29289323, 0.0, 0.29289323, 0.29289323] + ) + + +def test_streaming_can_prewarm_owned_index_segments(tmp_path): + dataset = lance.write_dataset( + _vector_table([[0.0, 0.0], [1.0, 0.0], [0.0, 2.0], [3.0, 0.0]]), + tmp_path / "streaming-prewarm.lance", + max_rows_per_file=2, + ) + dataset.create_index( + "vector", + "IVF_FLAT", + num_partitions=1, + name="vector_idx", + ) + + with lr.open_vector_search( + lance.dataset(dataset.uri), + nearest={"column": "vector", "k": 1, "nprobes": 1}, + index_name="vector_idx", + columns=["id"], + actor_options=VectorSearchActorOptions( + num_actors=1, + prewarm_index=True, + ), + ) as session: + assert session.prewarm_results[0]["skipped"] is False + assert session.prewarm_results[0]["index_segments"] == 1 + [result] = list( + session.map_batches([np.asarray([[0.0, 0.0]], dtype=np.float32)]) + ) + + assert result["id"].to_pylist() == [0] + + +def test_streaming_inherits_checked_out_dataset_snapshot(tmp_path): + path = tmp_path / "streaming-branch.lance" + dataset = lance.write_dataset(_vector_table([[0.0, 0.0]], ids=[0]), path) + branch_dataset = dataset.create_branch("experiment") + lance.write_dataset( + _vector_table([[1.0, 0.0]], ids=[1]), + path, + mode="append", + ) + + with lr.open_vector_search( + branch_dataset, + nearest={"column": "vector", "k": 2}, + columns=["id"], + actor_options=VectorSearchActorOptions(num_actors=1), + ) as session: + [result] = list( + session.map_batches([np.asarray([[0.0, 0.0]], dtype=np.float32)]) + ) + + assert result["id"].to_pylist() == [0] + assert session.actor_states[0]["version"] == branch_dataset.version + + +def test_streaming_uri_branch_uses_branch_snapshot(tmp_path): + path = tmp_path / "streaming-uri-branch.lance" + dataset = lance.write_dataset(_vector_table([[0.0, 0.0]], ids=[0]), path) + branch_dataset = dataset.create_branch("experiment") + branch_dataset = lance.write_dataset( + _vector_table([[1.0, 0.0]], ids=[1]), + branch_dataset.uri, + mode="append", + ) + + assert lance.dataset(path).version < branch_dataset.version + + with lr.open_vector_search( + str(path), + branch="experiment", + nearest={"column": "vector", "k": 2}, + columns=["id"], + actor_options=VectorSearchActorOptions(num_actors=1), + ) as session: + [result] = list( + session.map_batches([np.asarray([[0.0, 0.0]], dtype=np.float32)]) + ) + + assert result["id"].to_pylist() == [0, 1] + assert session.actor_states[0]["version"] == branch_dataset.version + + +def test_streaming_multivector_uses_additive_maxsim_distance(tmp_path): + vector_type = pa.list_(pa.list_(pa.float32(), 2)) + dataset = lance.write_dataset( + pa.table( + { + "id": [0, 1, 2], + "vector": pa.array( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[1.0, 0.0]], + [[-1.0, 0.0], [0.0, -1.0]], + ], + type=vector_type, + ), + } + ), + tmp_path / "streaming-multivector.lance", + ) + query_batch = pa.array( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[1.0, 0.0]], + ], + type=vector_type, + ) + + with lr.open_vector_search( + dataset, + nearest={"column": "vector", "k": 2, "metric": "cosine"}, + columns=["id"], + actor_options=VectorSearchActorOptions(num_actors=1), + ) as session: + [result] = list(session.map_batches([query_batch])) + + assert result["query_index"].to_pylist() == [0, 0, 1, 1] + assert result["id"].to_pylist() == [0, 1, 0, 1] + assert result["_distance"].to_pylist() == pytest.approx([0.0, 1.0, 0.0, 0.0]) + + +def test_streaming_multivector_uses_core_indexed_search(tmp_path): + vector_type = pa.list_(pa.list_(pa.float32(), 2)) + rows = [ + [[1.0, 0.0], [0.0, 1.0]], + [[1.0, 0.0]], + [[-1.0, 0.0], [0.0, -1.0]], + ] * 20 + dataset = lance.write_dataset( + pa.table( + { + "id": range(len(rows)), + "vector": pa.array(rows, type=vector_type), + } + ), + tmp_path / "streaming-multivector-indexed.lance", + ) + dataset.create_index( + "vector", + "IVF_FLAT", + num_partitions=1, + name="multivector_idx", + metric="cosine", + ) + + with lr.open_vector_search( + lance.dataset(dataset.uri), + nearest={"column": "vector", "k": 1, "nprobes": 1}, + index_name="multivector_idx", + columns=["id"], + actor_options=VectorSearchActorOptions(num_actors=1), + ) as session: + [result] = list( + session.map_batches( + [ + pa.array( + [ + [[1.0, 0.0], [0.0, 1.0]], + [[1.0, 0.0]], + ], + type=vector_type, + ) + ] + ) + ) + + assert result["query_index"].to_pylist() == [0, 1] + assert result["_distance"].to_pylist() == pytest.approx([0.0, 0.0]) + + def test_vector_search_puts_pickled_dataset_in_ray_object_store(monkeypatch): events = [] @@ -768,16 +1029,6 @@ def fake_loads(value): lambda *args, **kwargs: object(), ) monkeypatch.setattr(search_mod, "_plan_vector_search", lambda **kwargs: plans) - monkeypatch.setattr( - search_mod, - "_inspect_vector_search_query", - lambda *args, **kwargs: ( - False, - pa.schema( - [pa.field("id", pa.int64()), pa.field("_distance", pa.float32())] - ), - ), - ) monkeypatch.setattr(search_mod.pickle, "dumps", fake_dumps) monkeypatch.setattr(search_mod.pickle, "loads", fake_loads) monkeypatch.setattr(search_mod.ray, "ObjectRef", FakeObjectRef, raising=False) @@ -811,7 +1062,6 @@ def fake_loads(value): "scanner", { "fast_search": True, - "with_row_id": True, "nearest": {"column": "vector", "q": [0.0], "k": 1}, "index_segments": ["S1"], }, @@ -820,207 +1070,9 @@ def fake_loads(value): "scanner", { "fast_search": True, - "with_row_id": True, "nearest": {"column": "vector", "q": [0.0], "k": 1}, "index_segments": ["S2"], }, ), "get", ] - - -def test_batch_vector_search_without_index_matches_single_queries(tmp_path): - vectors = np.asarray( - [ - [0.0, 0.0], - [1.0, 0.0], - [0.0, 2.0], - [3.0, 0.0], - [0.0, 4.0], - [5.0, 0.0], - ], - dtype=np.float32, - ) - dataset = lance.write_dataset( - _vector_table(vectors), - tmp_path / "batch-flat.lance", - max_rows_per_file=2, - ) - queries = np.asarray([[0.0, 0.0], [0.0, 4.0]], dtype=np.float32) - - batch = lr.vector_search( - dataset, - nearest={"column": "vector", "q": queries, "k": 2}, - scanner_options={"columns": ["id", "_rowid"]}, - num_workers=2, - ) - - assert batch.column_names == ["query_index", "id", "_distance", "_rowid"] - assert batch.column("query_index").to_pylist() == [0, 0, 1, 1] - assert batch.num_rows == len(queries) * 2 - for query_index, query in enumerate(queries): - single = lr.vector_search( - dataset, - nearest={"column": "vector", "q": query, "k": 2}, - scanner_options={"columns": ["id", "_rowid"]}, - num_workers=2, - ) - batch_slice = batch.filter(pc.field("query_index") == query_index).drop_columns( - ["query_index"] - ) - assert batch_slice.column("id").to_pylist() == single.column("id").to_pylist() - assert batch_slice.column("_distance").to_pylist() == pytest.approx( - single.column("_distance").to_pylist() - ) - assert ( - batch_slice.column("_rowid").to_pylist() - == single.column("_rowid").to_pylist() - ) - - empty = lr.vector_search( - dataset, - nearest={"column": "vector", "q": queries, "k": 2}, - columns=["id"], - num_workers=2, - fast_search=True, - ) - assert empty.num_rows == 0 - assert empty.column_names == ["query_index", "id", "_distance"] - assert empty.schema.field("query_index").type == pa.int32() - assert not empty.schema.field("query_index").nullable - - virtual_projection = lr.vector_search( - dataset, - nearest={"column": "vector", "q": queries, "k": 1}, - columns=["query_index"], - num_workers=2, - ) - assert virtual_projection.column_names == ["query_index", "_distance"] - assert virtual_projection.column("query_index").to_pylist() == [0, 1] - - -def test_batch_vector_search_rejects_dataset_query_index_column(tmp_path): - vectors = _vector_table([[0.0, 0.0], [1.0, 0.0]])["vector"] - dataset = lance.write_dataset( - pa.table({"query_index": [7, 8], "vector": vectors}), - tmp_path / "batch-query-index.lance", - ) - - with pytest.raises(ValueError, match="column 'query_index'"): - lr.vector_search( - dataset, - nearest={"column": "vector", "q": [[0.0, 0.0]], "k": 1}, - ) - - -def test_multivector_fallback_reports_unsupported_boundary(tmp_path): - vector_type = pa.list_(pa.list_(pa.float32(), 2)) - dataset = lance.write_dataset( - pa.table( - { - "id": [0, 1], - "vector": pa.array( - [ - [[1.0, 0.0], [0.0, 1.0]], - [[-1.0, 0.0], [0.0, -1.0]], - ], - type=vector_type, - ), - } - ), - tmp_path / "multivector-flat.lance", - ) - - with pytest.raises(ValueError, match="does not support multivector"): - lr.vector_search( - dataset, - nearest={ - "column": "vector", - "q": [[1.0, 0.0], [0.0, 1.0]], - "k": 1, - }, - ) - - -def test_batch_vector_search_with_partial_index_preserves_per_query_top_k(tmp_path): - indexed_vectors = np.asarray( - [[0.0, 0.0], [1.0, 0.0], [0.0, 2.0], [3.0, 0.0]], - dtype=np.float32, - ) - appended_vectors = np.asarray([[0.0, 4.0], [5.0, 0.0]], dtype=np.float32) - dataset = _create_partial_index_dataset( - tmp_path / "batch-partial.lance", - indexed_vectors, - appended_vectors, - ) - queries = np.asarray([[0.0, 4.0], [3.0, 0.0]], dtype=np.float32) - - batch = lr.vector_search( - dataset, - nearest={"column": "vector", "q": queries, "k": 2}, - index_name="vector_idx", - columns=["id", "_rowid"], - num_workers=2, - ) - - assert batch.column_names == ["query_index", "id", "_distance", "_rowid"] - assert batch.column("query_index").to_pylist() == [0, 0, 1, 1] - assert batch.column("id").to_pylist() == [4, 2, 3, 1] - assert batch.column("_distance").to_pylist() == [0.0, 4.0, 0.0, 4.0] - - ranged_batch = lr.vector_search( - dataset, - nearest={ - "column": "vector", - "q": queries, - "k": 2, - "distance_range": (0.5, 10.0), - }, - index_name="vector_idx", - columns=["id"], - num_workers=2, - ) - assert ranged_batch.column("query_index").to_pylist() == [0, 1, 1] - assert ranged_batch.column("id").to_pylist() == [2, 1, 5] - assert ranged_batch.column("_distance").to_pylist() == [4.0, 4.0, 4.0] - - -@pytest.mark.parametrize("metric", ["COSINE", "DOT"]) -def test_apply_index_metric_default(metric): - index = SimpleNamespace(details={"metric_type": metric}) - nearest = {"column": "vector", "q": [[1.0, 0.0]], "k": 2} - - assert _apply_index_metric_default(nearest, index)["metric"] == metric.lower() - assert ( - _apply_index_metric_default({**nearest, "metric": "l2"}, index)["metric"] - == "l2" - ) - - -def test_partial_index_uses_index_metric_by_default(tmp_path): - indexed_vectors = np.asarray( - [[1.0, 0.0], [1.0, 1.0], [0.0, 1.0], [-1.0, 0.0]], - dtype=np.float32, - ) - appended_vectors = np.asarray([[0.5, 0.5], [-1.0, -1.0]], dtype=np.float32) - dataset = _create_partial_index_dataset( - tmp_path / "batch-partial-cosine.lance", - indexed_vectors, - appended_vectors, - metric="cosine", - ) - queries = np.asarray([[1.0, 0.0], [0.0, 1.0]], dtype=np.float32) - - result = lr.vector_search( - dataset, - nearest={"column": "vector", "q": queries, "k": 3}, - index_name="vector_idx", - columns=["id"], - num_workers=2, - ) - - assert result.column("query_index").to_pylist() == [0, 0, 0, 1, 1, 1] - assert result.column("id").to_pylist() == [0, 1, 4, 2, 1, 4] - assert result.column("_distance").to_pylist() == pytest.approx( - [0.0, 0.29289323, 0.29289323, 0.0, 0.29289323, 0.29289323] - )