diff --git a/docs/src/distributed-indexing.md b/docs/src/distributed-indexing.md index c2146d3c..804ca470 100755 --- a/docs/src/distributed-indexing.md +++ b/docs/src/distributed-indexing.md @@ -266,6 +266,58 @@ def vector_search( 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 ### FTS Index (Scalar) 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 895be2e8..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,16 +128,15 @@ def _canonical_index_field_names(field_names: Any) -> set[str]: return canonical_names -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: @@ -174,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), @@ -182,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, ) @@ -355,13 +380,18 @@ def _execute_flat_fallback_vector_search_plan( def _prepare_fallback_scan_columns( scanner_options: dict[str, Any], vector_column: str, + *, + 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): - scan_columns = [column for column in requested_columns if column != "_distance"] + virtual_columns = virtual_columns or {"_distance"} + 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 @@ -451,7 +481,13 @@ def _take_top_k(table: pa.Table, k: int) -> pa.Table: return table.take(sort_indices.slice(0, k)) -def _merge_vector_search_results(tables: list[pa.Table], k: int) -> pa.Table: +def _merge_vector_search_results( + tables: list[pa.Table], + k: int, + *, + 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: return tables[0].slice(0, 0) if tables else pa.table({}) @@ -463,6 +499,15 @@ def _merge_vector_search_results(tables: list[pa.Table], k: int) -> pa.Table: "for global top-k merge" ) + 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) @@ -704,3 +749,1265 @@ def run_plan(plan: _SearchPlan) -> pa.Table | _SearchPlanAnalysis: return _format_analyze_plan_results(results) 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) + ) diff --git a/tests/test_distributed_indexing.py b/tests/test_distributed_indexing.py index 7cc0686c..acab92bc 100755 --- a/tests/test_distributed_indexing.py +++ b/tests/test_distributed_indexing.py @@ -1920,6 +1920,26 @@ 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) + with lr.open_vector_search( + updated_dataset, + nearest={"column": "vector", "k": 5}, + index_name=index_name, + columns=["id"], + 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 ( + 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..cceb8cdc 100644 --- a/tests/test_distributed_search.py +++ b/tests/test_distributed_search.py @@ -1,13 +1,23 @@ from types import SimpleNamespace +import lance +import lance_ray as lr +import numpy as np import pyarrow as pa 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, + _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, @@ -37,6 +47,20 @@ def _index_with_segments(*segments): ) +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=pa.float32()), + matrix.shape[1], + ) + return pa.table( + { + "id": range(len(matrix)) if ids is None else ids, + "vector": vector_array, + } + ) + + def _mock_pickled_dataset(monkeypatch, dataset): search_mod._load_pickled_dataset.cache_clear() search_mod._load_pickled_dataset_ref.cache_clear() @@ -386,6 +410,28 @@ def test_merge_vector_search_results_requires_distance(): _merge_vector_search_results([table], k=1) +def test_merge_vector_search_results_can_merge_per_query(): + left = pa.table( + { + "query_index": [0, 0, 1], + "id": [1, 2, 3], + "_distance": [0.4, 0.1, 0.3], + } + ) + right = pa.table( + { + "query_index": [0, 1, 1], + "id": [4, 5, 6], + "_distance": [0.2, 0.4, 0.1], + } + ) + + result = _merge_vector_search_results([left, right], k=2, per_query=True) + + assert result["query_index"].to_pylist() == [0, 0, 1, 1] + assert result["id"].to_pylist() == [2, 4, 6, 3] + + def test_search_scanner_options_reject_managed_options(): with pytest.raises(ValueError, match="nearest"): _validate_search_scanner_options({"nearest": {"column": "vector"}}) @@ -458,6 +504,453 @@ def get_fragments(self): ] +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 = []