diff --git a/docs/src/distributed-indexing.md b/docs/src/distributed-indexing.md index ba54e9bd..f413e810 100755 --- a/docs/src/distributed-indexing.md +++ b/docs/src/distributed-indexing.md @@ -203,9 +203,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` @@ -237,7 +237,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. | @@ -256,7 +256,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 @@ -385,6 +385,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 81d9854d..f7dbe1f0 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -1756,6 +1756,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] + )