mirror of
https://github.com/lancedb/lancedb.git
synced 2026-08-18 12:08:35 +00:00
Heng Ge
f655f62e09
## What
MemWAL LSM **read** support. When a table has an LSM write spec
(`set_lsm_write_spec`), `merge_insert` upserts live in the MemWAL
active/frozen memtables and flushed SSTables until an external
compaction merges them into the base table, so a normal scan returns
**stale** data. This routes reads through Lance's `LsmScanner` so
queries also surface that in-flight data, deduplicated by primary key
(newest generation wins).
## How
- Adds a **`use_lsm: Option<bool>`** query flag, symmetric with the
`merge_insert` flag:
- **unset** — auto-route through the LSM scanner when the table carries
a write spec
- **`use_lsm(true)`** — force the LSM path; error if there is no spec
- **`use_lsm(false)`** — read the base table only (the escape hatch)
- Plain scan, single-column full-text search, and single-vector ANN all
run through one `LsmScanner` (assembled from on-disk shard manifests
plus the cached writer's in-memory memtables), so a `where` predicate is
honored as a **prefilter** uniformly — including for vector search.
- **Compaction-aware snapshots:** an SSTable generation is dropped only
once it is both compacted into the base table and covered by the arm's
base-index catch-up (`index_catchup`); plain scans use the compaction
watermark alone.
- Query shapes the scanner cannot honor hard-error with guidance to set
`use_lsm(false)`: hybrid, multi/binary vectors, `with_row_id`,
reranking, `order_by`, dynamic/Substrait projection or filters,
`distance_range`, `use_index(false)`, postfilter, take-by-row-id/offset,
reads from a time-traveled version, and an unmaintained or ambiguous
(multiple) FTS/vector index. Namespace-pushdown queries fall back to
local execution when a spec is present; WAL-only writers are handled.
- Exposed across the Rust core and the Python (`use_lsm`) and TypeScript
(`useLsm`) bindings, including `TakeQuery`.
Rebased from Lance `7.2.0-beta.3` to `10.0.0-beta.3`.
4322 lines
154 KiB
Python
4322 lines
154 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
||
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
|
||
|
||
from __future__ import annotations
|
||
|
||
import asyncio
|
||
from abc import ABC, abstractmethod
|
||
from concurrent.futures import ThreadPoolExecutor
|
||
from datetime import timedelta
|
||
from enum import Enum
|
||
from typing import (
|
||
TYPE_CHECKING,
|
||
Any,
|
||
Dict,
|
||
List,
|
||
Literal,
|
||
Optional,
|
||
Protocol,
|
||
Tuple,
|
||
Type,
|
||
TypeVar,
|
||
Union,
|
||
runtime_checkable,
|
||
)
|
||
|
||
import deprecation
|
||
import numpy as np
|
||
import pyarrow as pa
|
||
import pyarrow.compute as pc
|
||
import pydantic
|
||
from typing_extensions import Annotated
|
||
|
||
from lancedb._lancedb import fts_query_to_json
|
||
from lancedb.background_loop import LOOP
|
||
from lancedb.pydantic import PYDANTIC_VERSION
|
||
|
||
from . import __version__
|
||
from .arrow import AsyncRecordBatchReader
|
||
from .dependencies import pandas as pd
|
||
from .expr import Expr
|
||
from .rerankers.base import Reranker
|
||
from .rerankers.rrf import RRFReranker
|
||
from .rerankers.util import check_reranker_result
|
||
from .schema import is_blob_like_field, schema_has_blob_field
|
||
from .util import flatten_columns
|
||
from ._blob import (
|
||
BLOB_MODE_TO_HANDLING,
|
||
FetchBlobsAsync,
|
||
FetchBlobsSync,
|
||
blob_auto_row_id_for_scan,
|
||
blob_v2_projection_sources,
|
||
finalize_blob_query_table,
|
||
replace_v2_blob_columns_with_bytes,
|
||
replace_v2_blob_columns_with_bytes_sync,
|
||
supports_blob_auto_row_id,
|
||
validate_blob_mode,
|
||
)
|
||
from .types import BlobMode, QueryProjection
|
||
|
||
if TYPE_CHECKING:
|
||
import sys
|
||
|
||
import PIL
|
||
import polars as pl
|
||
|
||
from ._lancedb import FTSQuery as LanceFTSQuery
|
||
from ._lancedb import HybridQuery as LanceHybridQuery
|
||
from ._lancedb import PyQueryRequest
|
||
from ._lancedb import Query as LanceQuery
|
||
from ._lancedb import TakeQuery as LanceTakeQuery
|
||
from ._lancedb import VectorQuery as LanceVectorQuery
|
||
from .common import VEC
|
||
from .pydantic import LanceModel
|
||
from .table import AsyncTable, Table
|
||
|
||
if sys.version_info >= (3, 11):
|
||
from typing import Self
|
||
else:
|
||
from typing_extensions import Self
|
||
|
||
T = TypeVar("T", bound="LanceModel")
|
||
AnalyzePlanDistributedMetrics = Literal["aggregate", "per_worker", "full"]
|
||
|
||
|
||
@runtime_checkable
|
||
class _LanceScanner(Protocol):
|
||
projected_schema: pa.Schema | None
|
||
schema: pa.Schema | None
|
||
|
||
def to_pandas(self, blob_mode: BlobMode | None = ..., **kwargs) -> pd.DataFrame: ...
|
||
|
||
def to_pyarrow(self): ...
|
||
|
||
def to_table(self) -> pa.Table: ...
|
||
|
||
def to_reader(self): ...
|
||
|
||
|
||
def _blob_mode_requires_native_pandas(blob_mode: BlobMode, schema: pa.Schema) -> bool:
|
||
return blob_mode in BLOB_MODE_TO_HANDLING and schema_has_blob_field(schema)
|
||
|
||
|
||
def _unsupported_blob_pandas_error(reason: str) -> RuntimeError:
|
||
return RuntimeError(
|
||
"blob columns require Lance native scanner conversion for query "
|
||
f"to_pandas(), but {reason}. Use a plain scan query or remove blob "
|
||
"columns from the projection."
|
||
)
|
||
|
||
|
||
def _query_is_plain_scan(query: Query) -> bool:
|
||
return (
|
||
query.vector is None
|
||
and query.full_text_query is None
|
||
and not query.postfilter
|
||
and not query.order_by
|
||
)
|
||
|
||
|
||
def _filter_to_sql(filter: Optional[Union[str, Expr]]) -> Optional[str]:
|
||
if filter is None:
|
||
return None
|
||
if isinstance(filter, Expr):
|
||
return filter.to_sql()
|
||
return filter
|
||
|
||
|
||
def _combine_where(
|
||
existing: Optional[Union[str, Expr]], new: Union[str, Expr]
|
||
) -> Union[str, Expr]:
|
||
"""Combine a new filter with an existing one using a logical AND.
|
||
|
||
Calling ``where`` more than once composes the filters with AND instead of
|
||
replacing the previous filter. Two :class:`~lancedb.expr.Expr` filters are
|
||
combined as an expression; otherwise both filters are lowered to SQL strings
|
||
and combined as SQL.
|
||
"""
|
||
if existing is None:
|
||
return new
|
||
existing_is_expr = isinstance(existing, Expr)
|
||
new_is_expr = isinstance(new, Expr)
|
||
if existing_is_expr and new_is_expr:
|
||
return existing & new
|
||
existing_sql = existing.to_sql() if existing_is_expr else existing
|
||
new_sql = new.to_sql() if new_is_expr else new
|
||
return f"({existing_sql}) AND ({new_sql})"
|
||
|
||
|
||
def _projection_to_scanner_kwargs(columns: QueryProjection) -> Dict[str, Any]:
|
||
if columns is None:
|
||
return {}
|
||
if isinstance(columns, list):
|
||
if all(isinstance(column, str) for column in columns):
|
||
return {"columns": columns}
|
||
if all(isinstance(column, tuple) and len(column) == 2 for column in columns):
|
||
return {
|
||
"columns": {
|
||
name: expr.to_sql() if isinstance(expr, Expr) else expr
|
||
for name, expr in columns
|
||
}
|
||
}
|
||
# Let Lance raise the detailed projection validation error.
|
||
return {"columns": columns}
|
||
|
||
projection = {}
|
||
for name, expr in columns.items():
|
||
if isinstance(expr, Expr):
|
||
expr = expr.to_sql()
|
||
projection[name] = expr
|
||
return {"columns": projection}
|
||
|
||
|
||
def _scanner_kwargs_for_query(
|
||
query: Query,
|
||
blob_mode: BlobMode,
|
||
dataset: Optional[Any] = None,
|
||
*,
|
||
with_row_id: Optional[bool] = None,
|
||
) -> Dict[str, Any]:
|
||
fragments = _scanner_fragments_for_query(query, dataset)
|
||
kwargs = {
|
||
**_projection_to_scanner_kwargs(query.columns),
|
||
"filter": _filter_to_sql(query.filter),
|
||
"limit": query.limit,
|
||
"offset": query.offset,
|
||
"with_row_id": with_row_id if with_row_id is not None else query.with_row_id,
|
||
"with_row_address": query.with_row_address,
|
||
"fast_search": query.fast_search,
|
||
"blob_handling": BLOB_MODE_TO_HANDLING[blob_mode],
|
||
"fragments": fragments,
|
||
}
|
||
return {key: value for key, value in kwargs.items() if value is not None}
|
||
|
||
|
||
def _scanner_fragments_for_query(query: Query, dataset: Optional[Any]) -> Optional[Any]:
|
||
if query.fragments is not None and query.fragment_ids is not None:
|
||
raise ValueError("fragments and fragment_ids cannot both be set")
|
||
if query.fragments is not None:
|
||
return query.fragments
|
||
if query.fragment_ids is None:
|
||
return None
|
||
if dataset is None:
|
||
raise ValueError("fragment_ids require a Lance dataset")
|
||
|
||
requested = set(query.fragment_ids)
|
||
fragments = [
|
||
fragment
|
||
for fragment in dataset.get_fragments()
|
||
if fragment.fragment_id in requested
|
||
]
|
||
found = {fragment.fragment_id for fragment in fragments}
|
||
missing = requested - found
|
||
if missing:
|
||
missing_ids = ", ".join(str(fragment_id) for fragment_id in sorted(missing))
|
||
raise ValueError(f"fragment_ids not found in dataset: {missing_ids}")
|
||
return fragments
|
||
|
||
|
||
def _ensure_lazy_blob_frame(
|
||
df: "pd.DataFrame", schema: pa.Schema, blob_mode: BlobMode
|
||
) -> "pd.DataFrame":
|
||
if blob_mode != "lazy" or not schema_has_blob_field(schema) or len(df) == 0:
|
||
return df
|
||
|
||
for field in schema:
|
||
if not is_blob_like_field(field) or field.name not in df.columns:
|
||
continue
|
||
value = df[field.name].iloc[0]
|
||
if value is not None and not hasattr(value, "readall"):
|
||
raise _unsupported_blob_pandas_error(
|
||
"the Lance scanner did not return lazy blob files"
|
||
)
|
||
return df
|
||
|
||
|
||
def _scanner_to_table(scanner: _LanceScanner) -> pa.Table:
|
||
if hasattr(scanner, "to_pyarrow"):
|
||
reader = scanner.to_pyarrow()
|
||
return reader.read_all()
|
||
if hasattr(scanner, "to_table"):
|
||
return scanner.to_table()
|
||
reader = scanner.to_reader()
|
||
return reader.read_all()
|
||
|
||
|
||
def _scanner_to_pandas(
|
||
scanner: _LanceScanner, blob_mode: BlobMode, **kwargs
|
||
) -> pd.DataFrame:
|
||
schema = getattr(scanner, "projected_schema", None)
|
||
if schema is None:
|
||
schema = getattr(scanner, "schema", None)
|
||
if schema is None:
|
||
schema = getattr(scanner, "dataset_schema", None)
|
||
if callable(schema):
|
||
schema = schema()
|
||
if hasattr(scanner, "to_pandas"):
|
||
try:
|
||
df = scanner.to_pandas(blob_mode=blob_mode, **kwargs)
|
||
except TypeError as err:
|
||
message = str(err)
|
||
if "blob_mode" not in message and "unexpected keyword" not in message:
|
||
raise
|
||
df = scanner.to_pandas(**kwargs)
|
||
if schema is not None:
|
||
return _ensure_lazy_blob_frame(df, schema, blob_mode)
|
||
return df
|
||
|
||
tbl = _scanner_to_table(scanner)
|
||
if blob_mode == "lazy" and schema_has_blob_field(tbl.schema):
|
||
raise _unsupported_blob_pandas_error(
|
||
"the Lance scanner does not expose to_pandas"
|
||
)
|
||
return tbl.to_pandas(**kwargs)
|
||
|
||
|
||
def _finish_plain_scan_pandas(
|
||
scanner: _LanceScanner,
|
||
*,
|
||
blob_mode: BlobMode,
|
||
blob_sources: dict[str, str],
|
||
fetch_blobs: FetchBlobsSync,
|
||
strip_auto_row_id: bool,
|
||
flatten: Optional[Union[int, bool]],
|
||
**kwargs,
|
||
) -> pd.DataFrame:
|
||
if blob_sources:
|
||
tbl = _scanner_to_table(scanner)
|
||
tbl = replace_v2_blob_columns_with_bytes_sync(tbl, blob_sources, fetch_blobs)
|
||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||
tbl = tbl.drop_columns(["_rowid"])
|
||
if flatten is not None:
|
||
tbl = flatten_columns(tbl, flatten)
|
||
return tbl.to_pandas(**kwargs)
|
||
if flatten is not None:
|
||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||
tbl = tbl.drop_columns(["_rowid"])
|
||
return tbl.to_pandas(**kwargs)
|
||
df = _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||
if strip_auto_row_id and "_rowid" in df.columns:
|
||
return df.drop(columns=["_rowid"])
|
||
return df
|
||
|
||
|
||
async def _finish_plain_scan_pandas_async(
|
||
scanner: _LanceScanner,
|
||
*,
|
||
blob_mode: BlobMode,
|
||
blob_sources: dict[str, str],
|
||
fetch_blobs: FetchBlobsAsync,
|
||
strip_auto_row_id: bool,
|
||
flatten: Optional[Union[int, bool]],
|
||
**kwargs,
|
||
) -> pd.DataFrame:
|
||
if blob_sources:
|
||
tbl = _scanner_to_table(scanner)
|
||
tbl = await replace_v2_blob_columns_with_bytes(tbl, blob_sources, fetch_blobs)
|
||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||
tbl = tbl.drop_columns(["_rowid"])
|
||
if flatten is not None:
|
||
tbl = flatten_columns(tbl, flatten)
|
||
return tbl.to_pandas(**kwargs)
|
||
if flatten is not None:
|
||
tbl = flatten_columns(_scanner_to_table(scanner), flatten)
|
||
if strip_auto_row_id and "_rowid" in tbl.column_names:
|
||
tbl = tbl.drop_columns(["_rowid"])
|
||
return tbl.to_pandas(**kwargs)
|
||
df = _scanner_to_pandas(scanner, blob_mode, **kwargs)
|
||
if strip_auto_row_id and "_rowid" in df.columns:
|
||
return df.drop(columns=["_rowid"])
|
||
return df
|
||
|
||
|
||
# Pydantic validation function for vector queries
|
||
def ensure_vector_query(
|
||
val: Any,
|
||
) -> Union[List[float], List[List[float]], pa.Array, List[pa.Array]]:
|
||
if isinstance(val, list):
|
||
if len(val) == 0:
|
||
raise ValueError("Vector query must be a non-empty list")
|
||
sample = val[0]
|
||
else:
|
||
if isinstance(val, float):
|
||
raise ValueError(
|
||
"Vector query must be a list of floats or a list of lists of floats"
|
||
)
|
||
sample = val
|
||
if isinstance(sample, pa.Array):
|
||
# val is array or list of array
|
||
return val
|
||
if isinstance(sample, list):
|
||
if len(sample) == 0:
|
||
raise ValueError("Vector query must be a non-empty list")
|
||
if isinstance(sample[0], float):
|
||
# val is list of list of floats
|
||
return val
|
||
if isinstance(sample, float):
|
||
# val is a list of floats
|
||
return val
|
||
|
||
|
||
class ColumnOrdering(pydantic.BaseModel):
|
||
column_name: str
|
||
ascending: bool = True
|
||
nulls_first: bool = False
|
||
|
||
|
||
class FullTextQueryType(str, Enum):
|
||
MATCH = "match"
|
||
MATCH_PHRASE = "match_phrase"
|
||
BOOST = "boost"
|
||
MULTI_MATCH = "multi_match"
|
||
BOOLEAN = "boolean"
|
||
|
||
|
||
class FullTextOperator(str, Enum):
|
||
AND = "AND"
|
||
OR = "OR"
|
||
|
||
|
||
class Occur(str, Enum):
|
||
SHOULD = "SHOULD"
|
||
MUST = "MUST"
|
||
MUST_NOT = "MUST_NOT"
|
||
|
||
|
||
@pydantic.dataclasses.dataclass
|
||
class FullTextQuery(ABC):
|
||
@abstractmethod
|
||
def query_type(self) -> FullTextQueryType:
|
||
"""
|
||
Get the query type of the query.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
The type of the query.
|
||
"""
|
||
pass
|
||
|
||
def to_json(self) -> str:
|
||
"""
|
||
Convert the query to a JSON string.
|
||
|
||
Returns
|
||
-------
|
||
str
|
||
A JSON string representation of the query.
|
||
|
||
Examples
|
||
--------
|
||
>>> from lancedb.query import MatchQuery
|
||
>>> query = MatchQuery("puppy", "text", fuzziness=2)
|
||
>>> query.to_json()
|
||
'{"match":{"column":"text","terms":"puppy","boost":1.0,"fuzziness":2,"max_expansions":50,"operator":"Or","prefix_length":0}}'
|
||
"""
|
||
return fts_query_to_json(self)
|
||
|
||
def __and__(self, other: "FullTextQuery") -> "FullTextQuery":
|
||
"""
|
||
Combine two queries with a logical AND operation.
|
||
|
||
Parameters
|
||
----------
|
||
other : FullTextQuery
|
||
The other query to combine with.
|
||
|
||
Returns
|
||
-------
|
||
FullTextQuery
|
||
A new query that combines both queries with AND.
|
||
"""
|
||
return BooleanQuery([(Occur.MUST, self), (Occur.MUST, other)])
|
||
|
||
def __or__(self, other: "FullTextQuery") -> "FullTextQuery":
|
||
"""
|
||
Combine two queries with a logical OR operation.
|
||
|
||
Parameters
|
||
----------
|
||
other : FullTextQuery
|
||
The other query to combine with.
|
||
|
||
Returns
|
||
-------
|
||
FullTextQuery
|
||
A new query that combines both queries with OR.
|
||
"""
|
||
return BooleanQuery([(Occur.SHOULD, self), (Occur.SHOULD, other)])
|
||
|
||
|
||
@pydantic.dataclasses.dataclass
|
||
class MatchQuery(FullTextQuery):
|
||
"""
|
||
Match query for full-text search.
|
||
|
||
Parameters
|
||
----------
|
||
query : str
|
||
The query string to match against.
|
||
column : str
|
||
The name of the column to match against.
|
||
boost : float, default 1.0
|
||
The boost factor for the query.
|
||
The score of each matching document is multiplied by this value.
|
||
fuzziness : int, optional
|
||
The maximum edit distance for each term in the match query.
|
||
Defaults to 0 (exact match).
|
||
If None, fuzziness is applied automatically by the rules:
|
||
- 0 for terms with length <= 2
|
||
- 1 for terms with length <= 5
|
||
- 2 for terms with length > 5
|
||
max_expansions : int, optional
|
||
The maximum number of terms to consider for fuzzy matching.
|
||
Defaults to 50.
|
||
operator : FullTextOperator, default OR
|
||
The operator to use for combining the query results.
|
||
Can be either `AND` or `OR`.
|
||
If `AND`, all terms in the query must match.
|
||
If `OR`, at least one term in the query must match.
|
||
prefix_length : int, optional
|
||
The number of beginning characters being unchanged for fuzzy matching.
|
||
This is useful to achieve prefix matching.
|
||
"""
|
||
|
||
query: str
|
||
column: str
|
||
boost: float = pydantic.Field(1.0, kw_only=True)
|
||
fuzziness: int = pydantic.Field(0, kw_only=True)
|
||
max_expansions: int = pydantic.Field(50, kw_only=True)
|
||
operator: FullTextOperator = pydantic.Field(FullTextOperator.OR, kw_only=True)
|
||
prefix_length: int = pydantic.Field(0, kw_only=True)
|
||
|
||
def query_type(self) -> FullTextQueryType:
|
||
return FullTextQueryType.MATCH
|
||
|
||
|
||
@pydantic.dataclasses.dataclass
|
||
class PhraseQuery(FullTextQuery):
|
||
"""
|
||
Phrase query for full-text search.
|
||
|
||
Parameters
|
||
----------
|
||
query : str
|
||
The query string to match against.
|
||
column : str
|
||
The name of the column to match against.
|
||
"""
|
||
|
||
query: str
|
||
column: str
|
||
slop: int = pydantic.Field(0, kw_only=True)
|
||
|
||
def query_type(self) -> FullTextQueryType:
|
||
return FullTextQueryType.MATCH_PHRASE
|
||
|
||
|
||
@pydantic.dataclasses.dataclass
|
||
class BoostQuery(FullTextQuery):
|
||
"""
|
||
Boost query for full-text search.
|
||
|
||
Parameters
|
||
----------
|
||
positive : dict
|
||
The positive query object.
|
||
negative : dict
|
||
The negative query object.
|
||
negative_boost : float, default 0.5
|
||
The boost factor for the negative query.
|
||
"""
|
||
|
||
positive: FullTextQuery
|
||
negative: FullTextQuery
|
||
negative_boost: float = pydantic.Field(0.5, kw_only=True)
|
||
|
||
def query_type(self) -> FullTextQueryType:
|
||
return FullTextQueryType.BOOST
|
||
|
||
|
||
@pydantic.dataclasses.dataclass
|
||
class MultiMatchQuery(FullTextQuery):
|
||
"""
|
||
Multi-match query for full-text search.
|
||
|
||
Parameters
|
||
----------
|
||
query : str | list[Query]
|
||
If a string, the query string to match against.
|
||
columns : list[str]
|
||
The list of columns to match against.
|
||
boosts : list[float], optional
|
||
The list of boost factors for each column. If not provided,
|
||
all columns will have the same boost factor.
|
||
operator : FullTextOperator, default OR
|
||
The operator to use for combining the query results.
|
||
Can be either `AND` or `OR`.
|
||
It would be applied to all columns individually.
|
||
For example, if the operator is `AND`,
|
||
then the query "hello world" is equal to
|
||
`match("hello AND world", column1) OR match("hello AND world", column2)`.
|
||
"""
|
||
|
||
query: str
|
||
columns: list[str]
|
||
boosts: Optional[list[float]] = pydantic.Field(None, kw_only=True)
|
||
operator: FullTextOperator = pydantic.Field(FullTextOperator.OR, kw_only=True)
|
||
|
||
def query_type(self) -> FullTextQueryType:
|
||
return FullTextQueryType.MULTI_MATCH
|
||
|
||
|
||
@pydantic.dataclasses.dataclass
|
||
class BooleanQuery(FullTextQuery):
|
||
"""
|
||
Boolean query for full-text search.
|
||
|
||
Parameters
|
||
----------
|
||
queries : list[tuple(Occur, FullTextQuery)]
|
||
The list of queries with their occurrence requirements.
|
||
Each tuple contains an Occur value (MUST, SHOULD, or MUST_NOT)
|
||
and a FullTextQuery to apply.
|
||
"""
|
||
|
||
queries: list[tuple[Occur, FullTextQuery]]
|
||
|
||
def query_type(self) -> FullTextQueryType:
|
||
return FullTextQueryType.BOOLEAN
|
||
|
||
|
||
class FullTextSearchQuery(pydantic.BaseModel):
|
||
"""A LanceDB Full Text Search Query
|
||
|
||
Attributes
|
||
----------
|
||
columns: List[str]
|
||
The columns to search
|
||
|
||
If None, then the table should select the column automatically.
|
||
query: str | FullTextQuery
|
||
If a string, it is treated as a MatchQuery.
|
||
If a FullTextQuery object, it is used directly.
|
||
"""
|
||
|
||
columns: Optional[List[str]] = None
|
||
query: Union[str, FullTextQuery]
|
||
|
||
|
||
class Query(pydantic.BaseModel):
|
||
"""A LanceDB Query
|
||
|
||
Queries are constructed by the `Table.search` method. This class is a
|
||
python representation of the query. Normally you will not need to interact
|
||
with this class directly. You can build up a query and execute it using
|
||
collection methods such as `to_batches()`, `to_arrow()`, `to_pandas()`,
|
||
etc.
|
||
|
||
However, you can use the `to_query()` method to get the underlying query object.
|
||
This can be useful for serializing a query or using it in a different context.
|
||
|
||
Attributes
|
||
----------
|
||
filter : Optional[str]
|
||
sql filter to refine the query with
|
||
limit : Optional[int]
|
||
The limit on the number of results to return. If this is a vector or FTS query,
|
||
then this is required. If this is a plain SQL query, then this is optional.
|
||
offset: Optional[int]
|
||
The offset to start fetching results from
|
||
|
||
This is ignored for vector / FTS search (will be None).
|
||
columns : Optional[Union[List[str], Dict[str, str]]]
|
||
which columns to return in the results
|
||
|
||
This can be a list of column names or a dictionary. If it is a dictionary,
|
||
then the keys are the column names and the values are sql expressions to
|
||
use to calculate the result.
|
||
|
||
If this is None then all columns are returned. This can be expensive.
|
||
with_row_id : Optional[bool]
|
||
if True then include the row id in the results
|
||
vector : Optional[Union[List[float], List[List[float]], pa.Array, List[pa.Array]]]
|
||
the vector to search for, if this a vector search or hybrid search. It will
|
||
be None for full text search and plain SQL filtering.
|
||
vector_column : Optional[str]
|
||
the name of the vector column to use for vector search
|
||
|
||
If this is None then a default vector column will be used.
|
||
distance_type : Optional[str]
|
||
the distance type to use for vector search
|
||
|
||
This can be l2 (default), cosine and dot. See [metric definitions][search] for
|
||
more details.
|
||
|
||
If this is not a vector search this will be None.
|
||
postfilter : bool
|
||
if True then apply the filter after vector / FTS search. This is ignored for
|
||
plain SQL filtering.
|
||
nprobes : Optional[int]
|
||
The number of IVF partitions to search. If this is None then a default
|
||
number of partitions will be used.
|
||
|
||
- A higher number makes search more accurate but also slower.
|
||
|
||
- See discussion in [Querying an ANN Index][querying-an-ann-index] for
|
||
tuning advice.
|
||
|
||
Will be None if this is not a vector search.
|
||
refine_factor : Optional[int]
|
||
Refine the results by reading extra elements and re-ranking them in memory.
|
||
|
||
- A higher number makes search more accurate but also slower.
|
||
|
||
- See discussion in [Querying an ANN Index][querying-an-ann-index] for
|
||
tuning advice.
|
||
|
||
Will be None if this is not a vector search.
|
||
lower_bound : Optional[float]
|
||
The lower bound for distance search
|
||
|
||
Only results with a distance greater than or equal to this value
|
||
will be returned.
|
||
|
||
This will only be set on vector search.
|
||
upper_bound : Optional[float]
|
||
The upper bound for distance search
|
||
|
||
Only results with a distance less than or equal to this value
|
||
will be returned.
|
||
|
||
This will only be set on vector search.
|
||
ef : Optional[int]
|
||
The size of the nearest neighbor list maintained during HNSW search
|
||
|
||
This will only be set on vector search.
|
||
full_text_query : Optional[Union[str, dict]]
|
||
The full text search query
|
||
|
||
This can be a string or a dictionary. A dictionary will be used to search
|
||
multiple columns. The keys are the column names and the values are the
|
||
search queries.
|
||
|
||
This will only be set on FTS or hybrid queries.
|
||
fast_search: Optional[bool]
|
||
Skip a flat search of unindexed data. This will improve
|
||
search performance but search results will not include unindexed data.
|
||
|
||
The default is False
|
||
"""
|
||
|
||
# The name of the vector column to use for vector search.
|
||
vector_column: Optional[str] = None
|
||
|
||
# vector to search for
|
||
#
|
||
# Note: today this will be floats on the sync path and pa.Array on the async
|
||
# path though in the future we should unify this to pa.Array everywhere
|
||
vector: Annotated[
|
||
Optional[Union[List[float], List[List[float]], pa.Array, List[pa.Array]]],
|
||
ensure_vector_query,
|
||
] = None
|
||
|
||
# sql filter or type-safe Expr to refine the query with
|
||
filter: Optional[Union[str, Expr]] = None
|
||
|
||
# if True then apply the filter after vector search
|
||
postfilter: Optional[bool] = None
|
||
|
||
# full text search query
|
||
full_text_query: Optional[FullTextSearchQuery] = None
|
||
|
||
# top k results to return
|
||
limit: Optional[int] = None
|
||
|
||
# distance type to use for vector search
|
||
distance_type: Optional[str] = None
|
||
|
||
# which columns to return in the results (dict values may be str or Expr)
|
||
columns: QueryProjection = None
|
||
|
||
# minimum number of IVF partitions to search
|
||
#
|
||
# If None then a default value (20) will be used.
|
||
minimum_nprobes: Optional[int] = None
|
||
|
||
# maximum number of IVF partitions to search
|
||
#
|
||
# If None then a default value (20) will be used.
|
||
#
|
||
# If 0 then no limit will be applied and all partitions could be searched
|
||
# if needed to satisfy the limit.
|
||
maximum_nprobes: Optional[int] = None
|
||
|
||
# lower bound for distance search
|
||
lower_bound: Optional[float] = None
|
||
|
||
# upper bound for distance search
|
||
upper_bound: Optional[float] = None
|
||
|
||
# multiplier for the number of results to inspect for reranking
|
||
refine_factor: Optional[int] = None
|
||
|
||
# if true, include the row id in the results
|
||
with_row_id: Optional[bool] = None
|
||
|
||
# if true, include the row address in the results
|
||
with_row_address: Optional[bool] = None
|
||
|
||
# Lance fragments or fragment ids to scan on scanner-backed plain queries
|
||
fragments: Optional[Any] = None
|
||
fragment_ids: Optional[List[int]] = None
|
||
|
||
# offset to start fetching results from
|
||
offset: Optional[int] = None
|
||
|
||
# if true, will only search the indexed data
|
||
fast_search: Optional[bool] = None
|
||
|
||
# MemWAL LSM read routing: None auto-routes when the table carries a write
|
||
# spec, True forces the LSM scanner (errors without a spec), False reads the
|
||
# base table only
|
||
use_lsm: Optional[bool] = None
|
||
|
||
# size of the nearest neighbor list maintained during HNSW search
|
||
ef: Optional[int] = None
|
||
|
||
# Bypass the vector index and use a brute force search
|
||
bypass_vector_index: Optional[bool] = None
|
||
|
||
order_by: Optional[List[ColumnOrdering]] = None
|
||
|
||
@classmethod
|
||
def from_inner(cls, req: PyQueryRequest) -> Self:
|
||
query = cls()
|
||
query.limit = req.limit
|
||
query.offset = req.offset
|
||
query.filter = req.filter
|
||
query.full_text_query = req.full_text_search
|
||
query.columns = req.select
|
||
query.with_row_id = req.with_row_id
|
||
# use_lsm is a genuine tri-state (None / True / False); preserve it as-is
|
||
# so a round-tripped query keeps an explicit False.
|
||
query.use_lsm = req.use_lsm
|
||
query.vector_column = req.column
|
||
query.vector = req.query_vector
|
||
query.distance_type = req.distance_type
|
||
query.minimum_nprobes = req.minimum_nprobes
|
||
query.maximum_nprobes = req.maximum_nprobes
|
||
query.lower_bound = req.lower_bound
|
||
query.upper_bound = req.upper_bound
|
||
query.ef = req.ef
|
||
query.refine_factor = req.refine_factor
|
||
query.bypass_vector_index = req.bypass_vector_index
|
||
query.postfilter = req.postfilter
|
||
if req.order_by is not None:
|
||
query.order_by = [ColumnOrdering(**o) for o in req.order_by]
|
||
if req.full_text_search is not None:
|
||
query.full_text_query = FullTextSearchQuery(
|
||
columns=None,
|
||
query=req.full_text_search,
|
||
)
|
||
return query
|
||
|
||
# This tells pydantic to allow custom types (needed for the `vector` query since
|
||
# pa.Array wouln't be allowed otherwise)
|
||
if PYDANTIC_VERSION.major < 2: # Pydantic 1.x compat
|
||
|
||
class Config:
|
||
arbitrary_types_allowed = True
|
||
else:
|
||
model_config = {"arbitrary_types_allowed": True}
|
||
|
||
|
||
class LanceQueryBuilder(ABC):
|
||
"""An abstract query builder. Subclasses are defined for vector search,
|
||
full text search, hybrid, and plain SQL filtering.
|
||
"""
|
||
|
||
@classmethod
|
||
def create(
|
||
cls,
|
||
table: "Table",
|
||
query: Optional[Union[np.ndarray, str, "PIL.Image.Image", Tuple]],
|
||
query_type: str,
|
||
vector_column_name: str,
|
||
ordering_field_name: Optional[str] = None,
|
||
fts_columns: Optional[Union[str, List[str]]] = None,
|
||
fast_search: bool = None,
|
||
) -> Self:
|
||
"""
|
||
Create a query builder based on the given query and query type.
|
||
|
||
Parameters
|
||
----------
|
||
table: Table
|
||
The table to query.
|
||
query: Optional[Union[np.ndarray, str, "PIL.Image.Image", Tuple]]
|
||
The query to use. If None, an empty query builder is returned
|
||
which performs simple SQL filtering.
|
||
query_type: str
|
||
The type of query to perform. One of "vector", "fts", "hybrid", or "auto".
|
||
If "auto", the query type is inferred based on the query.
|
||
vector_column_name: str
|
||
The name of the vector column to use for vector search.
|
||
ordering_field_name: Optional[str]
|
||
.. deprecated:: 0.27.0
|
||
Use ``order_by()`` method instead.
|
||
fts_columns: Optional[Union[str, List[str]]]
|
||
The columns to search in for full text search.
|
||
fast_search: bool
|
||
Skip flat search of unindexed data.
|
||
"""
|
||
if ordering_field_name is not None:
|
||
import warnings
|
||
|
||
warnings.warn(
|
||
"ordering_field_name is deprecated, use .order_by() method instead.",
|
||
DeprecationWarning,
|
||
stacklevel=2,
|
||
)
|
||
# Check hybrid search first as it supports empty query pattern
|
||
if query_type == "hybrid":
|
||
# hybrid fts and vector query
|
||
return LanceHybridQueryBuilder(
|
||
table, query, vector_column_name, fts_columns=fts_columns
|
||
)
|
||
|
||
if query is None:
|
||
return LanceEmptyQueryBuilder(table)
|
||
|
||
# remember the string query for reranking purpose
|
||
str_query = query if isinstance(query, str) else None
|
||
|
||
# convert "auto" query_type to "vector", "fts"
|
||
# or "hybrid" and convert the query to vector if needed
|
||
query, query_type = cls._resolve_query(
|
||
table, query, query_type, vector_column_name
|
||
)
|
||
|
||
if query_type == "hybrid":
|
||
return LanceHybridQueryBuilder(
|
||
table, query, vector_column_name, fts_columns=fts_columns
|
||
)
|
||
|
||
if isinstance(query, (str, FullTextQuery)):
|
||
# fts
|
||
return LanceFtsQueryBuilder(
|
||
table,
|
||
query,
|
||
ordering_field_name=ordering_field_name,
|
||
fts_columns=fts_columns,
|
||
fast_search=fast_search,
|
||
)
|
||
|
||
if isinstance(query, list):
|
||
query = np.array(query, dtype=np.float32)
|
||
elif isinstance(query, np.ndarray):
|
||
query = query.astype(np.float32)
|
||
else:
|
||
raise TypeError(f"Unsupported query type: {type(query)}")
|
||
|
||
return LanceVectorQueryBuilder(
|
||
table, query, vector_column_name, str_query, fast_search
|
||
)
|
||
|
||
@classmethod
|
||
def _resolve_query(cls, table, query, query_type, vector_column_name):
|
||
# If query_type is fts, then query must be a string.
|
||
# otherwise raise TypeError
|
||
if query_type == "fts":
|
||
if not isinstance(query, (str, FullTextQuery)):
|
||
raise TypeError(
|
||
f"'fts' query must be a string or FullTextQuery: {type(query)}"
|
||
)
|
||
return query, query_type
|
||
elif query_type == "vector":
|
||
query = cls._query_to_vector(table, query, vector_column_name)
|
||
return query, query_type
|
||
elif query_type == "auto":
|
||
if isinstance(query, (list, np.ndarray)):
|
||
return query, "vector"
|
||
else:
|
||
conf = table.embedding_functions.get(vector_column_name)
|
||
if conf is not None:
|
||
query = conf.function.compute_query_embeddings_with_retry(query)[0]
|
||
return query, "vector"
|
||
else:
|
||
return query, "fts"
|
||
else:
|
||
raise ValueError(
|
||
f"Invalid query_type, must be 'vector', 'fts', or 'auto': {query_type}"
|
||
)
|
||
|
||
@classmethod
|
||
def _query_to_vector(cls, table, query, vector_column_name):
|
||
if isinstance(query, (list, np.ndarray)):
|
||
return query
|
||
conf = table.embedding_functions.get(vector_column_name)
|
||
if conf is not None:
|
||
return conf.function.compute_query_embeddings_with_retry(query)[0]
|
||
else:
|
||
msg = f"No embedding function for {vector_column_name}"
|
||
raise ValueError(msg)
|
||
|
||
def __init__(self, table: "Table"):
|
||
self._table = table
|
||
self._limit = None
|
||
self._offset = None
|
||
self._columns = None
|
||
self._where = None
|
||
self._postfilter = None
|
||
self._with_row_id = None
|
||
self._with_row_address = None
|
||
self._fragments = None
|
||
self._fragment_ids = None
|
||
self._use_lsm = None
|
||
self._vector = None
|
||
self._text = None
|
||
self._ef = None
|
||
self._bypass_vector_index = None
|
||
self._order_by = None
|
||
|
||
@deprecation.deprecated(
|
||
deprecated_in="0.3.1",
|
||
removed_in="0.4.0",
|
||
current_version=__version__,
|
||
details="Use to_pandas() instead",
|
||
)
|
||
def to_df(self) -> "pd.DataFrame":
|
||
"""
|
||
*Deprecated alias for `to_pandas()`. Please use `to_pandas()` instead.*
|
||
|
||
Execute the query and return the results as a pandas DataFrame.
|
||
In addition to the selected columns, LanceDB also returns a vector
|
||
and also the "_distance" column which is the distance between the query
|
||
vector and the returned vector.
|
||
"""
|
||
return self.to_pandas()
|
||
|
||
def to_pandas(
|
||
self,
|
||
flatten: Optional[Union[int, bool]] = None,
|
||
*,
|
||
blob_mode: BlobMode = "lazy",
|
||
timeout: Optional[timedelta] = None,
|
||
**kwargs,
|
||
) -> "pd.DataFrame":
|
||
"""
|
||
Execute the query and return the results as a pandas DataFrame.
|
||
In addition to the selected columns, LanceDB also returns a vector
|
||
and also the "_distance" column which is the distance between the query
|
||
vector and the returned vector.
|
||
|
||
Parameters
|
||
----------
|
||
flatten: Optional[Union[int, bool]]
|
||
If flatten is True, flatten all nested columns.
|
||
If flatten is an integer, flatten the nested columns up to the
|
||
specified depth.
|
||
If unspecified, do not flatten the nested columns.
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
blob_mode: str, default "lazy"
|
||
Controls how blob columns are returned for plain scan queries.
|
||
Vector, FTS, hybrid, and other non-native query shapes keep the
|
||
existing Arrow conversion path and only support blob descriptions.
|
||
**kwargs
|
||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||
optional flattening.
|
||
"""
|
||
validate_blob_mode(blob_mode)
|
||
output_schema = getattr(self, "output_schema", None)
|
||
if output_schema is not None:
|
||
schema = output_schema()
|
||
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
||
native_error = None
|
||
if (flatten is None or blob_mode == "descriptions") and timeout is None:
|
||
try:
|
||
df = self._plain_scan_to_pandas(
|
||
blob_mode, flatten=flatten, **kwargs
|
||
)
|
||
if df is not None:
|
||
return df
|
||
except Exception as err:
|
||
native_error = err
|
||
reason = (
|
||
"this query shape cannot use Lance native pandas conversion"
|
||
if native_error is None
|
||
else str(native_error)
|
||
)
|
||
raise _unsupported_blob_pandas_error(reason) from native_error
|
||
|
||
tbl = flatten_columns(self.to_arrow(timeout=timeout), flatten)
|
||
if _blob_mode_requires_native_pandas(blob_mode, tbl.schema):
|
||
raise _unsupported_blob_pandas_error(
|
||
"this query shape cannot use Lance native pandas conversion"
|
||
)
|
||
return tbl.to_pandas(**kwargs)
|
||
|
||
@abstractmethod
|
||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||
"""
|
||
Execute the query and return the results as an
|
||
[Apache Arrow Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table).
|
||
|
||
In addition to the selected columns, LanceDB also returns a vector
|
||
and also the "_distance" column which is the distance between the query
|
||
vector and the returned vectors.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
@abstractmethod
|
||
def to_batches(
|
||
self,
|
||
/,
|
||
batch_size: Optional[int] = None,
|
||
*,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> pa.RecordBatchReader:
|
||
"""
|
||
Execute the query and return the results as a pyarrow
|
||
[RecordBatchReader](https://arrow.apache.org/docs/python/generated/pyarrow.RecordBatchReader.html)
|
||
|
||
For v2 blob projections, ``to_batches`` keeps the auto ``_rowid``
|
||
column visible so batch consumers can call ``fetch_blobs``. Use
|
||
``to_arrow``, ``to_list``, or ``to_pandas`` if you want LanceDB to hide
|
||
auto row ids in the final collected result.
|
||
|
||
Parameters
|
||
----------
|
||
batch_size: int
|
||
The maximum number of selected records in a RecordBatch object.
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def to_list(self, *, timeout: Optional[timedelta] = None) -> List[dict]:
|
||
"""
|
||
Execute the query and return the results as a list of dictionaries.
|
||
|
||
Each list entry is a dictionary with the selected column names as keys,
|
||
or all table columns if `select` is not called. The vector and the "_distance"
|
||
fields are returned whether or not they're explicitly selected.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
"""
|
||
return self.to_arrow(timeout=timeout).to_pylist()
|
||
|
||
def to_pydantic(
|
||
self, model: type[T], *, timeout: Optional[timedelta] = None
|
||
) -> list[T]:
|
||
"""Return the table as a list of pydantic models.
|
||
|
||
Parameters
|
||
----------
|
||
model: Type[LanceModel]
|
||
The pydantic model to use.
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
|
||
Returns
|
||
-------
|
||
List[LanceModel]
|
||
"""
|
||
return [model(**row) for row in self.to_arrow(timeout=timeout).to_pylist()]
|
||
|
||
def to_polars(self, *, timeout: Optional[timedelta] = None) -> "pl.DataFrame":
|
||
"""
|
||
Execute the query and return the results as a Polars DataFrame.
|
||
In addition to the selected columns, LanceDB also returns a vector
|
||
and also the "_distance" column which is the distance between the query
|
||
vector and the returned vector.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
"""
|
||
import polars as pl
|
||
|
||
return pl.from_arrow(self.to_arrow(timeout=timeout))
|
||
|
||
def limit(self, limit: Union[int, None]) -> Self:
|
||
"""Set the maximum number of results to return.
|
||
|
||
Parameters
|
||
----------
|
||
limit: int
|
||
The maximum number of results to return.
|
||
The default query limit is 10 results.
|
||
For ANN/KNN queries, you must specify a limit.
|
||
For plain searches, all records are returned if limit not set.
|
||
*WARNING* if you have a large dataset, setting
|
||
the limit to a large number, e.g. the table size,
|
||
can potentially result in reading a
|
||
large amount of data into memory and cause
|
||
out of memory issues.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
if limit is None or limit <= 0:
|
||
if isinstance(self, LanceVectorQueryBuilder):
|
||
raise ValueError("Limit is required for ANN/KNN queries")
|
||
else:
|
||
self._limit = None
|
||
else:
|
||
self._limit = limit
|
||
return self
|
||
|
||
def offset(self, offset: int) -> Self:
|
||
"""Set the offset for the results.
|
||
|
||
Parameters
|
||
----------
|
||
offset: int
|
||
The offset to start fetching results from.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
if offset is None or offset <= 0:
|
||
self._offset = 0
|
||
else:
|
||
self._offset = offset
|
||
return self
|
||
|
||
def select(self, columns: Union[list[str], dict[str, Union[str, Expr]]]) -> Self:
|
||
"""Set the columns to return.
|
||
|
||
Parameters
|
||
----------
|
||
columns: list of str, or dict of str to str or Expr
|
||
List of column names to be fetched.
|
||
Or a dictionary of column names to SQL expressions or
|
||
:class:`~lancedb.expr.Expr` objects.
|
||
All columns are fetched if None or unspecified.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
if isinstance(columns, list) or isinstance(columns, dict):
|
||
self._columns = columns
|
||
else:
|
||
raise ValueError("columns must be a list or a dictionary")
|
||
return self
|
||
|
||
def where(self, where: Union[str, Expr], prefilter: bool = True) -> Self:
|
||
"""Set the where clause.
|
||
|
||
Parameters
|
||
----------
|
||
where: str or :class:`~lancedb.expr.Expr`
|
||
The filter condition. Can be a SQL string or a type-safe
|
||
:class:`~lancedb.expr.Expr` built with :func:`~lancedb.expr.col`
|
||
and :func:`~lancedb.expr.lit`.
|
||
prefilter: bool, default True
|
||
If True, apply the filter before vector search, otherwise the
|
||
filter is applied on the result of vector search.
|
||
This feature is **EXPERIMENTAL** and may be removed and modified
|
||
without warning in the future.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
|
||
Notes
|
||
-----
|
||
Calling this multiple times combines the filters with a logical AND
|
||
rather than replacing the previous filter.
|
||
"""
|
||
self._where = _combine_where(self._where, where)
|
||
self._postfilter = not prefilter
|
||
return self
|
||
|
||
def with_row_id(self, with_row_id: bool) -> Self:
|
||
"""Set whether to return row ids.
|
||
|
||
Parameters
|
||
----------
|
||
with_row_id: bool
|
||
If True, return _rowid column in the results.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._with_row_id = with_row_id
|
||
return self
|
||
|
||
def _user_requested_row_id(self) -> bool:
|
||
return self._with_row_id is True
|
||
|
||
def _blob_auto_row_id_enabled(self) -> bool:
|
||
if not supports_blob_auto_row_id(self._table):
|
||
return False
|
||
return blob_auto_row_id_for_scan(
|
||
self._table,
|
||
self._table.schema,
|
||
self._columns,
|
||
with_row_id=self._with_row_id,
|
||
)
|
||
|
||
def _scan_needs_row_id(self) -> bool:
|
||
return self._user_requested_row_id() or self._blob_auto_row_id_enabled()
|
||
|
||
def _query_for_scan(self) -> Query:
|
||
query = self.to_query_object()
|
||
if self._scan_needs_row_id():
|
||
query.with_row_id = True
|
||
return query
|
||
|
||
def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table:
|
||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||
blob_paths = (
|
||
blob_v2_projection_sources(self._table.schema, self._columns).keys()
|
||
if blob_auto_row_id
|
||
else ()
|
||
)
|
||
return finalize_blob_query_table(
|
||
tbl,
|
||
user_requested_row_id=self._user_requested_row_id(),
|
||
blob_auto_row_id=blob_auto_row_id,
|
||
blob_paths=blob_paths,
|
||
)
|
||
|
||
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||
"""Set whether to return row addresses.
|
||
|
||
Parameters
|
||
----------
|
||
with_row_address: bool, default True
|
||
If True, return the _rowaddr column in the results.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._with_row_address = with_row_address
|
||
return self
|
||
|
||
def with_fragments(self, fragments: Any) -> Self:
|
||
"""Set the Lance fragments to scan for plain scanner-backed queries."""
|
||
self._fragments = fragments
|
||
return self
|
||
|
||
def fragment_ids(self, fragment_ids: List[int]) -> Self:
|
||
"""Set the Lance fragment ids to scan for plain scanner-backed queries."""
|
||
self._fragment_ids = fragment_ids
|
||
return self
|
||
|
||
def use_lsm(self, enable: bool) -> Self:
|
||
"""Control MemWAL LSM read routing for this query.
|
||
|
||
By default (unset), a query against a table with an LSM write spec is
|
||
routed through the LSM scanner so it also returns data written via the
|
||
``merge_insert`` LSM path that has not yet been compacted into the base
|
||
table (active/frozen memtables + flushed generations); a table without a
|
||
spec reads the base table.
|
||
|
||
Parameters
|
||
----------
|
||
enable : bool
|
||
``True`` forces the LSM scanner and errors if the table has no LSM
|
||
write spec. ``False`` bypasses the MemWAL and reads the base table
|
||
only, even when a spec is present.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._use_lsm = enable
|
||
return self
|
||
|
||
def explain_plan(self, verbose: Optional[bool] = False) -> str:
|
||
"""Return the execution plan for this query.
|
||
|
||
Examples
|
||
--------
|
||
>>> import lancedb
|
||
>>> db = lancedb.connect("./.lancedb")
|
||
>>> table = db.create_table("my_table", [{"vector": [99.0, 99]}])
|
||
>>> query = [100, 100]
|
||
>>> plan = table.search(query).explain_plan(True)
|
||
>>> print(plan) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
|
||
GlobalLimitExec: skip=0, fetch=10
|
||
FilterExec: _distance@2 IS NOT NULL
|
||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||
KNNVectorDistance: metric=l2
|
||
LanceRead: uri=..., projection=[vector], ...
|
||
|
||
Parameters
|
||
----------
|
||
verbose : bool, default False
|
||
Use a verbose output format.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
""" # noqa: E501
|
||
return self._table._explain_plan(self.to_query_object(), verbose=verbose)
|
||
|
||
def order_by(self, ordering: Optional[List[ColumnOrdering]]) -> Self:
|
||
"""
|
||
Set the ordering for the results.
|
||
|
||
Parameters
|
||
----------
|
||
ordering: Optional[List[ColumnOrdering]]
|
||
The ordering to use for the results. If None, then the default ordering
|
||
will be used.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._order_by = ordering
|
||
return self
|
||
|
||
def analyze_plan(
|
||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||
) -> str:
|
||
"""
|
||
Run the query and return its execution plan with runtime metrics.
|
||
|
||
This returns detailed metrics for each step, such as elapsed time,
|
||
rows processed, bytes read, and I/O stats. It is useful for debugging
|
||
and performance tuning.
|
||
|
||
Examples
|
||
--------
|
||
>>> import lancedb
|
||
>>> db = lancedb.connect("./.lancedb")
|
||
>>> table = db.create_table("my_table", [{"vector": [99.0, 99]}])
|
||
>>> query = [100, 100]
|
||
>>> plan = table.search(query).analyze_plan()
|
||
>>> print(plan) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||
AnalyzeExec verbose=true, elapsed=..., metrics=...
|
||
TracedExec, elapsed=..., metrics=...
|
||
ProjectionExec: elapsed=..., expr=[...],
|
||
metrics=[output_rows=..., elapsed_compute=..., output_bytes=...]
|
||
GlobalLimitExec: elapsed=..., skip=0, fetch=10,
|
||
metrics=[output_rows=..., elapsed_compute=..., output_bytes=...]
|
||
FilterExec: elapsed=..., _distance@2 IS NOT NULL, metrics=[...]
|
||
SortExec: elapsed=..., TopK(fetch=10), expr=[...],
|
||
preserve_partitioning=[...],
|
||
metrics=[output_rows=..., elapsed_compute=...,
|
||
output_bytes=..., row_replacements=...]
|
||
KNNVectorDistance: elapsed=..., metric=l2,
|
||
metrics=[output_rows=..., elapsed_compute=...,
|
||
output_bytes=..., output_batches=...]
|
||
LanceRead: elapsed=..., uri=..., projection=[vector],
|
||
num_fragments=..., range_before=None, range_after=None,
|
||
row_id=true, row_addr=false,
|
||
full_filter=--, refine_filter=--,
|
||
metrics=[output_rows=..., elapsed_compute=..., output_bytes=...,
|
||
fragments_scanned=..., ranges_scanned=1, rows_scanned=1,
|
||
bytes_read=..., iops=..., requests=..., task_wait_time=...]
|
||
|
||
Parameters
|
||
----------
|
||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||
Defaults to "aggregate".
|
||
How distributed worker metrics are displayed for remote query plans.
|
||
"aggregate" preserves the legacy summary, "per_worker" shows each
|
||
worker separately, and "full" includes both.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
The physical query execution plan with runtime metrics.
|
||
"""
|
||
return self._table._analyze_plan(
|
||
self.to_query_object(), distributed_metrics=distributed_metrics
|
||
)
|
||
|
||
def vector(self, vector: Union[np.ndarray, list]) -> Self:
|
||
"""Set the vector to search for.
|
||
|
||
Parameters
|
||
----------
|
||
vector: np.ndarray or list
|
||
The vector to search for.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def text(self, text: str | FullTextQuery) -> Self:
|
||
"""Set the text to search for.
|
||
|
||
Parameters
|
||
----------
|
||
text: str | FullTextQuery
|
||
If a string, it is treated as a MatchQuery.
|
||
If a FullTextQuery object, it is used directly.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
@abstractmethod
|
||
def rerank(self, reranker: Reranker) -> Self:
|
||
"""Rerank the results using the specified reranker.
|
||
|
||
Parameters
|
||
----------
|
||
reranker: Reranker
|
||
The reranker to use.
|
||
|
||
Returns
|
||
-------
|
||
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
def _plain_scan_to_pandas(
|
||
self,
|
||
blob_mode: BlobMode,
|
||
flatten: Optional[Union[int, bool]] = None,
|
||
**kwargs,
|
||
) -> Optional["pd.DataFrame"]:
|
||
query = self.to_query_object()
|
||
if not _query_is_plain_scan(query):
|
||
return None
|
||
|
||
dataset = self._table.to_lance()
|
||
blob_auto_row_id = self._blob_auto_row_id_enabled()
|
||
blob_sources = (
|
||
blob_v2_projection_sources(self._table.schema, query.columns)
|
||
if blob_mode == "bytes"
|
||
else {}
|
||
)
|
||
scanner = dataset.scanner(
|
||
**_scanner_kwargs_for_query(
|
||
query,
|
||
"descriptions" if blob_sources else blob_mode,
|
||
dataset,
|
||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
||
)
|
||
)
|
||
return _finish_plain_scan_pandas(
|
||
scanner,
|
||
blob_mode=blob_mode,
|
||
blob_sources=blob_sources,
|
||
fetch_blobs=self._table.fetch_blobs,
|
||
strip_auto_row_id=blob_auto_row_id,
|
||
flatten=flatten,
|
||
**kwargs,
|
||
)
|
||
|
||
@abstractmethod
|
||
def to_query_object(self) -> Query:
|
||
"""Return a serializable representation of the query
|
||
|
||
Returns
|
||
-------
|
||
Query
|
||
The serializable representation of the query
|
||
"""
|
||
raise NotImplementedError
|
||
|
||
|
||
class LanceVectorQueryBuilder(LanceQueryBuilder):
|
||
"""
|
||
Examples
|
||
--------
|
||
>>> import lancedb
|
||
>>> data = [{"vector": [1.1, 1.2], "b": 2},
|
||
... {"vector": [0.5, 1.3], "b": 4},
|
||
... {"vector": [0.4, 0.4], "b": 6},
|
||
... {"vector": [0.4, 0.4], "b": 10}]
|
||
>>> db = lancedb.connect("./.lancedb")
|
||
>>> table = db.create_table("my_table", data=data)
|
||
>>> (table.search([0.4, 0.4])
|
||
... .distance_type("cosine")
|
||
... .where("b < 10")
|
||
... .select(["b", "vector"])
|
||
... .limit(2)
|
||
... .to_pandas())
|
||
b vector _distance
|
||
0 6 [0.4, 0.4] 0.000000
|
||
1 2 [1.1, 1.2] 0.000944
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
table: "Table",
|
||
query: Union[np.ndarray, list, "PIL.Image.Image"],
|
||
vector_column: str,
|
||
str_query: Optional[str] = None,
|
||
fast_search: bool = None,
|
||
):
|
||
super().__init__(table)
|
||
self._query = query
|
||
self._distance_type = None
|
||
self._minimum_nprobes = None
|
||
self._maximum_nprobes = None
|
||
self._lower_bound = None
|
||
self._upper_bound = None
|
||
self._refine_factor = None
|
||
self._vector_column = vector_column
|
||
self._postfilter = None
|
||
self._reranker = None
|
||
self._str_query = str_query
|
||
self._fast_search = fast_search
|
||
|
||
def metric(self, metric: Literal["l2", "cosine", "dot"]) -> LanceVectorQueryBuilder:
|
||
"""Set the distance metric to use.
|
||
|
||
This is an alias for distance_type() and may be deprecated in the future.
|
||
Please use distance_type() instead.
|
||
|
||
Parameters
|
||
----------
|
||
metric: "l2" or "cosine" or "dot"
|
||
The distance metric to use. By default "l2" is used.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
return self.distance_type(metric)
|
||
|
||
def distance_type(
|
||
self, distance_type: Literal["l2", "cosine", "dot"]
|
||
) -> "LanceVectorQueryBuilder":
|
||
"""Set the distance metric to use.
|
||
|
||
When performing a vector search we try and find the "nearest" vectors according
|
||
to some kind of distance metric. This parameter controls which distance metric
|
||
to use.
|
||
|
||
Note: if there is a vector index then the distance type used MUST match the
|
||
distance type used to train the vector index. If this is not done then the
|
||
results will be invalid.
|
||
|
||
Parameters
|
||
----------
|
||
distance_type: "l2" or "cosine" or "dot"
|
||
The distance metric to use. By default "l2" is used.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._distance_type = distance_type.lower()
|
||
return self
|
||
|
||
def nprobes(self, nprobes: int) -> LanceVectorQueryBuilder:
|
||
"""Set the number of probes to use.
|
||
|
||
Higher values will yield better recall (more likely to find vectors if
|
||
they exist) at the expense of latency.
|
||
|
||
See discussion in [Querying an ANN Index][querying-an-ann-index] for
|
||
tuning advice.
|
||
|
||
This method sets both the minimum and maximum number of probes to the same
|
||
value. See `minimum_nprobes` and `maximum_nprobes` for more fine-grained
|
||
control.
|
||
|
||
Parameters
|
||
----------
|
||
nprobes: int
|
||
The number of probes to use.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._minimum_nprobes = nprobes
|
||
self._maximum_nprobes = nprobes
|
||
return self
|
||
|
||
def minimum_nprobes(self, minimum_nprobes: int) -> LanceVectorQueryBuilder:
|
||
"""Set the minimum number of probes to use.
|
||
|
||
See `nprobes` for more details.
|
||
|
||
These partitions will be searched on every vector query and will increase recall
|
||
at the expense of latency.
|
||
"""
|
||
self._minimum_nprobes = minimum_nprobes
|
||
return self
|
||
|
||
def maximum_nprobes(self, maximum_nprobes: int) -> LanceVectorQueryBuilder:
|
||
"""Set the maximum number of probes to use.
|
||
|
||
See `nprobes` for more details.
|
||
|
||
If this value is greater than `minimum_nprobes` then the excess partitions
|
||
will be searched only if we have not found enough results.
|
||
|
||
This can be useful when there is a narrow filter to allow these queries to
|
||
spend more time searching and avoid potential false negatives.
|
||
|
||
If this value is 0 then no limit will be applied and all partitions could be
|
||
searched if needed to satisfy the limit.
|
||
"""
|
||
self._maximum_nprobes = maximum_nprobes
|
||
return self
|
||
|
||
def distance_range(
|
||
self, lower_bound: Optional[float] = None, upper_bound: Optional[float] = None
|
||
) -> LanceVectorQueryBuilder:
|
||
"""Set the distance range to use.
|
||
|
||
Only rows with distances within range [lower_bound, upper_bound)
|
||
will be returned.
|
||
|
||
Parameters
|
||
----------
|
||
lower_bound: Optional[float]
|
||
The lower bound of the distance range.
|
||
upper_bound: Optional[float]
|
||
The upper bound of the distance range.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._lower_bound = lower_bound
|
||
self._upper_bound = upper_bound
|
||
return self
|
||
|
||
def ef(self, ef: int) -> LanceVectorQueryBuilder:
|
||
"""Set the number of candidates to consider during search.
|
||
|
||
Higher values will yield better recall (more likely to find vectors if
|
||
they exist) at the expense of latency.
|
||
|
||
This only applies to the HNSW-related index.
|
||
The default value is 1.5 * limit.
|
||
|
||
Parameters
|
||
----------
|
||
ef: int
|
||
The number of candidates to consider during search.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._ef = ef
|
||
return self
|
||
|
||
def refine_factor(self, refine_factor: int) -> LanceVectorQueryBuilder:
|
||
"""Set the refine factor to use, increasing the number of vectors sampled.
|
||
|
||
As an example, a refine factor of 2 will sample 2x as many vectors as
|
||
requested, re-ranks them, and returns the top half most relevant results.
|
||
|
||
See discussion in [Querying an ANN Index][querying-an-ann-index] for
|
||
tuning advice.
|
||
|
||
Parameters
|
||
----------
|
||
refine_factor: int
|
||
The refine factor to use.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._refine_factor = refine_factor
|
||
return self
|
||
|
||
def output_schema(self) -> pa.Schema:
|
||
"""
|
||
Return the output schema for the query
|
||
|
||
This does not execute the query.
|
||
"""
|
||
return self._table._output_schema(self.to_query_object())
|
||
|
||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||
"""
|
||
Execute the query and return the results as an
|
||
[Apache Arrow Table](https://arrow.apache.org/docs/python/generated/pyarrow.Table.html#pyarrow.Table).
|
||
|
||
In addition to the selected columns, LanceDB also returns a vector
|
||
and also the "_distance" column which is the distance between the query
|
||
vector and the returned vectors.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
"""
|
||
return self._finalize_blob_query_table(
|
||
self.to_batches(timeout=timeout).read_all()
|
||
)
|
||
|
||
def to_query_object(self) -> Query:
|
||
"""
|
||
Build a Query object
|
||
|
||
This can be used to serialize a query
|
||
"""
|
||
vector = self._query if isinstance(self._query, list) else self._query.tolist()
|
||
if isinstance(vector[0], np.ndarray):
|
||
vector = [v.tolist() for v in vector]
|
||
return Query(
|
||
vector=vector,
|
||
filter=self._where,
|
||
postfilter=self._postfilter,
|
||
limit=self._limit,
|
||
distance_type=self._distance_type,
|
||
columns=self._columns,
|
||
minimum_nprobes=self._minimum_nprobes,
|
||
maximum_nprobes=self._maximum_nprobes,
|
||
lower_bound=self._lower_bound,
|
||
upper_bound=self._upper_bound,
|
||
refine_factor=self._refine_factor,
|
||
vector_column=self._vector_column,
|
||
with_row_id=self._with_row_id,
|
||
with_row_address=self._with_row_address,
|
||
fragments=self._fragments,
|
||
fragment_ids=self._fragment_ids,
|
||
use_lsm=self._use_lsm,
|
||
offset=self._offset,
|
||
fast_search=self._fast_search,
|
||
ef=self._ef,
|
||
bypass_vector_index=self._bypass_vector_index,
|
||
order_by=self._order_by,
|
||
)
|
||
|
||
def to_batches(
|
||
self,
|
||
/,
|
||
batch_size: Optional[int] = None,
|
||
*,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> pa.RecordBatchReader:
|
||
"""
|
||
Execute the query and return the result as a RecordBatchReader object.
|
||
|
||
Parameters
|
||
----------
|
||
batch_size: int
|
||
The maximum number of selected records in a RecordBatch object.
|
||
timeout: timedelta, default None
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
|
||
Returns
|
||
-------
|
||
pa.RecordBatchReader
|
||
"""
|
||
vector = self._query if isinstance(self._query, list) else self._query.tolist()
|
||
if isinstance(vector[0], np.ndarray):
|
||
vector = [v.tolist() for v in vector]
|
||
query = self._query_for_scan()
|
||
result_set = self._table._execute_query(
|
||
query, batch_size=batch_size, timeout=timeout
|
||
)
|
||
if self._reranker is not None:
|
||
rs_table = result_set.read_all()
|
||
result_set = self._reranker.rerank_vector(self._str_query, rs_table)
|
||
check_reranker_result(result_set)
|
||
# convert result_set back to RecordBatchReader
|
||
result_set = pa.RecordBatchReader.from_batches(
|
||
result_set.schema, result_set.to_batches()
|
||
)
|
||
|
||
return result_set
|
||
|
||
def where(
|
||
self, where: Union[str, Expr], prefilter: bool = None
|
||
) -> LanceVectorQueryBuilder:
|
||
"""Set the where clause.
|
||
|
||
Parameters
|
||
----------
|
||
where: str or :class:`~lancedb.expr.Expr`
|
||
The filter condition. Can be a SQL string or a type-safe
|
||
:class:`~lancedb.expr.Expr` built with :func:`~lancedb.expr.col`
|
||
and :func:`~lancedb.expr.lit`.
|
||
prefilter: bool, default True
|
||
If True, apply the filter before vector search, otherwise the
|
||
filter is applied on the result of vector search.
|
||
|
||
Returns
|
||
-------
|
||
LanceQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
|
||
Notes
|
||
-----
|
||
Calling this multiple times combines the filters with a logical AND
|
||
rather than replacing the previous filter.
|
||
"""
|
||
self._where = _combine_where(self._where, where)
|
||
if prefilter is not None:
|
||
self._postfilter = not prefilter
|
||
return self
|
||
|
||
def rerank(
|
||
self, reranker: Reranker, query_string: Optional[str] = None
|
||
) -> LanceVectorQueryBuilder:
|
||
"""Rerank the results using the specified reranker.
|
||
|
||
Parameters
|
||
----------
|
||
reranker: Reranker
|
||
The reranker to use.
|
||
|
||
query_string: Optional[str]
|
||
The query to use for reranking. This needs to be specified explicitly here
|
||
as the query used for vector search may already be vectorized and the
|
||
reranker requires a string query.
|
||
This is only required if the query used for vector search is not a string.
|
||
Note: This doesn't yet support the case where the query is multimodal or a
|
||
list of vectors.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._reranker = reranker
|
||
if self._str_query is None and query_string is None:
|
||
raise ValueError(
|
||
"""
|
||
The query used for vector search is not a string.
|
||
In this case, the reranker query needs to be specified explicitly.
|
||
"""
|
||
)
|
||
if query_string is not None and not isinstance(query_string, str):
|
||
raise ValueError("Reranking currently only supports string queries")
|
||
self._str_query = query_string if query_string is not None else self._str_query
|
||
if reranker.score == "all":
|
||
self.with_row_id(True)
|
||
return self
|
||
|
||
def bypass_vector_index(self) -> LanceVectorQueryBuilder:
|
||
"""
|
||
If this is called then any vector index is skipped
|
||
|
||
An exhaustive (flat) search will be performed. The query vector will
|
||
be compared to every vector in the table. At high scales this can be
|
||
expensive. However, this is often still useful. For example, skipping
|
||
the vector index can give you ground truth results which you can use to
|
||
calculate your recall to select an appropriate value for nprobes.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceVectorQueryBuilder object.
|
||
"""
|
||
self._bypass_vector_index = True
|
||
return self
|
||
|
||
def fast_search(self) -> LanceVectorQueryBuilder:
|
||
"""
|
||
Skip a flat search of unindexed data. This will improve
|
||
search performance but search results will not include unindexed data.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceVectorQueryBuilder object.
|
||
"""
|
||
self._fast_search = True
|
||
return self
|
||
|
||
|
||
class LanceFtsQueryBuilder(LanceQueryBuilder):
|
||
"""A builder for full text search for LanceDB."""
|
||
|
||
def __init__(
|
||
self,
|
||
table: "Table",
|
||
query: str | FullTextQuery,
|
||
ordering_field_name: Optional[str] = None,
|
||
fts_columns: Optional[Union[str, List[str]]] = None,
|
||
fast_search: bool = None,
|
||
):
|
||
super().__init__(table)
|
||
self._query = query
|
||
self._phrase_query = False
|
||
# Deprecated compatibility parameter. Native FTS ordering is now
|
||
# configured through order_by(); LanceQueryBuilder.create emits the warning.
|
||
_ = ordering_field_name
|
||
self._reranker = None
|
||
self._fast_search = fast_search
|
||
if isinstance(fts_columns, str):
|
||
fts_columns = [fts_columns]
|
||
self._fts_columns = fts_columns
|
||
|
||
def phrase_query(self, phrase_query: bool = True) -> LanceFtsQueryBuilder:
|
||
"""Set whether to use phrase query.
|
||
|
||
Parameters
|
||
----------
|
||
phrase_query: bool, default True
|
||
If True, then an unquoted string query will be wrapped in quotes.
|
||
|
||
Returns
|
||
-------
|
||
LanceFtsQueryBuilder
|
||
The LanceFtsQueryBuilder object.
|
||
"""
|
||
self._phrase_query = phrase_query
|
||
return self
|
||
|
||
def _query_with_phrase_semantics(self) -> str | FullTextQuery:
|
||
query = self._query
|
||
if not self._phrase_query:
|
||
return query
|
||
if isinstance(query, str):
|
||
if not query.startswith('"') or not query.endswith('"'):
|
||
return f'"{query}"'
|
||
return query
|
||
if isinstance(query, PhraseQuery):
|
||
return query
|
||
raise TypeError(
|
||
"phrase_query() requires a string or PhraseQuery, "
|
||
f"got {type(query).__name__}"
|
||
)
|
||
|
||
def fast_search(self) -> LanceFtsQueryBuilder:
|
||
"""
|
||
Skip a flat search of unindexed data. This will improve
|
||
search performance but search results will not include unindexed data.
|
||
|
||
Returns
|
||
-------
|
||
LanceFtsQueryBuilder
|
||
The LanceFtsQueryBuilder object.
|
||
"""
|
||
self._fast_search = True
|
||
return self
|
||
|
||
def to_query_object(self) -> Query:
|
||
return Query(
|
||
columns=self._columns,
|
||
filter=self._where,
|
||
limit=self._limit,
|
||
postfilter=self._postfilter,
|
||
with_row_id=self._with_row_id,
|
||
with_row_address=self._with_row_address,
|
||
fragments=self._fragments,
|
||
fragment_ids=self._fragment_ids,
|
||
use_lsm=self._use_lsm,
|
||
full_text_query=FullTextSearchQuery(
|
||
query=self._query_with_phrase_semantics(), columns=self._fts_columns
|
||
),
|
||
offset=self._offset,
|
||
fast_search=self._fast_search,
|
||
order_by=self._order_by,
|
||
)
|
||
|
||
def output_schema(self) -> pa.Schema:
|
||
"""
|
||
Return the output schema for the query
|
||
|
||
This does not execute the query.
|
||
"""
|
||
return self._table._output_schema(self.to_query_object())
|
||
|
||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||
self._table._ensure_no_legacy_fts_index()
|
||
|
||
query = self._query_for_scan()
|
||
results = self._table._execute_query(query, timeout=timeout)
|
||
results = results.read_all()
|
||
if self._reranker is not None:
|
||
results = self._reranker.rerank_fts(self._query, results)
|
||
check_reranker_result(results)
|
||
return self._finalize_blob_query_table(results)
|
||
|
||
def to_batches(
|
||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||
):
|
||
raise NotImplementedError("to_batches on an FTS query")
|
||
|
||
def rerank(self, reranker: Reranker) -> LanceFtsQueryBuilder:
|
||
"""Rerank the results using the specified reranker.
|
||
|
||
Parameters
|
||
----------
|
||
reranker: Reranker
|
||
The reranker to use.
|
||
|
||
Returns
|
||
-------
|
||
LanceFtsQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._reranker = reranker
|
||
if reranker.score == "all":
|
||
self.with_row_id(True)
|
||
return self
|
||
|
||
|
||
class LanceEmptyQueryBuilder(LanceQueryBuilder):
|
||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||
return self._finalize_blob_query_table(
|
||
self.to_batches(timeout=timeout).read_all()
|
||
)
|
||
|
||
def to_query_object(self) -> Query:
|
||
return Query(
|
||
columns=self._columns,
|
||
filter=self._where,
|
||
limit=self._limit,
|
||
with_row_id=self._with_row_id,
|
||
with_row_address=self._with_row_address,
|
||
fragments=self._fragments,
|
||
fragment_ids=self._fragment_ids,
|
||
use_lsm=self._use_lsm,
|
||
offset=self._offset,
|
||
order_by=self._order_by,
|
||
)
|
||
|
||
def output_schema(self) -> pa.Schema:
|
||
query = self.to_query_object()
|
||
return self._table._output_schema(query)
|
||
|
||
def to_batches(
|
||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||
) -> pa.RecordBatchReader:
|
||
query = self._query_for_scan()
|
||
return self._table._execute_query(query, batch_size=batch_size, timeout=timeout)
|
||
|
||
def rerank(self, reranker: Reranker) -> LanceEmptyQueryBuilder:
|
||
"""Rerank the results using the specified reranker.
|
||
|
||
Parameters
|
||
----------
|
||
reranker: Reranker
|
||
The reranker to use.
|
||
|
||
Returns
|
||
-------
|
||
LanceEmptyQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
raise NotImplementedError("Reranking is not yet supported.")
|
||
|
||
|
||
class LanceHybridQueryBuilder(LanceQueryBuilder):
|
||
"""
|
||
A query builder that performs hybrid vector and full text search.
|
||
Results are combined and reranked based on the specified reranker.
|
||
By default, the results are reranked using the RRFReranker, which
|
||
uses reciprocal rank fusion score for reranking.
|
||
|
||
To make the vector and fts results comparable, the scores are normalized.
|
||
Instead of normalizing scores, the `normalize` parameter can be set to "rank"
|
||
in the `rerank` method to convert the scores to ranks and then normalize them.
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
table: "Table",
|
||
query: Optional[Union[str, FullTextQuery]] = None,
|
||
vector_column: Optional[str] = None,
|
||
fts_columns: Optional[Union[str, List[str]]] = None,
|
||
):
|
||
super().__init__(table)
|
||
self._query = query
|
||
self._vector_column = vector_column
|
||
self._fts_columns = fts_columns
|
||
self._norm = None
|
||
self._reranker = None
|
||
self._minimum_nprobes = None
|
||
self._maximum_nprobes = None
|
||
self._refine_factor = None
|
||
self._distance_type = None
|
||
self._phrase_query = None
|
||
self._lower_bound = None
|
||
self._upper_bound = None
|
||
|
||
def _validate_query(self, query, vector=None, text=None):
|
||
if query is not None and (vector is not None or text is not None):
|
||
raise ValueError(
|
||
"You can either provide a string query in search() method "
|
||
"or set `vector()` and `text()` explicitly for hybrid search."
|
||
"But not both."
|
||
)
|
||
|
||
vector_query = vector if vector is not None else query
|
||
if not isinstance(vector_query, (str, list, np.ndarray)):
|
||
raise ValueError("Vector query must be either a string or a vector")
|
||
|
||
text_query = text or query
|
||
if text_query is None:
|
||
raise ValueError("Text query must be provided for hybrid search.")
|
||
if not isinstance(text_query, (str, FullTextQuery)):
|
||
raise ValueError("Text query must be a string or FullTextQuery")
|
||
|
||
return vector_query, text_query
|
||
|
||
def phrase_query(self, phrase_query: bool = True) -> LanceHybridQueryBuilder:
|
||
"""Set whether to use phrase query.
|
||
|
||
Parameters
|
||
----------
|
||
phrase_query: bool, default True
|
||
If True, then an unquoted string query will be wrapped in quotes.
|
||
|
||
Returns
|
||
-------
|
||
LanceHybridQueryBuilder
|
||
The LanceHybridQueryBuilder object.
|
||
"""
|
||
self._phrase_query = phrase_query
|
||
return self
|
||
|
||
def to_query_object(self) -> Query:
|
||
raise NotImplementedError("to_query_object not yet supported on a hybrid query")
|
||
|
||
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
|
||
self._create_query_builders()
|
||
with ThreadPoolExecutor() as executor:
|
||
fts_future = executor.submit(
|
||
self._fts_query.with_row_id(True).to_arrow, timeout=timeout
|
||
)
|
||
vector_future = executor.submit(
|
||
self._vector_query.with_row_id(True).to_arrow, timeout=timeout
|
||
)
|
||
fts_results = fts_future.result()
|
||
vector_results = vector_future.result()
|
||
|
||
results = self._combine_hybrid_results(
|
||
fts_results=fts_results,
|
||
vector_results=vector_results,
|
||
norm=self._norm,
|
||
fts_query=self._fts_query._query,
|
||
reranker=self._reranker,
|
||
limit=self._limit,
|
||
with_row_ids=True,
|
||
)
|
||
return self._finish_hybrid_results(results)
|
||
|
||
def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
|
||
if self._user_requested_row_id():
|
||
return results
|
||
if self._blob_auto_row_id_enabled():
|
||
return self._finalize_blob_query_table(results)
|
||
if "_rowid" in results.column_names:
|
||
return results.drop(["_rowid"])
|
||
return results
|
||
|
||
@staticmethod
|
||
def _combine_hybrid_results(
|
||
fts_results: pa.Table,
|
||
vector_results: pa.Table,
|
||
norm: str,
|
||
fts_query: str,
|
||
reranker,
|
||
limit: int,
|
||
with_row_ids: bool,
|
||
) -> pa.Table:
|
||
if norm == "rank":
|
||
vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance")
|
||
fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score")
|
||
|
||
# If both result sets are empty (e.g. after hard filtering),
|
||
# return early to avoid errors in reranking or score restoration.
|
||
if vector_results.num_rows == 0 and fts_results.num_rows == 0:
|
||
# Build a minimal empty table with the _relevance_score column
|
||
combined_schema = pa.unify_schemas(
|
||
[vector_results.schema, fts_results.schema],
|
||
)
|
||
empty = pa.table(
|
||
{
|
||
col: pa.array([], type=combined_schema.field(col).type)
|
||
for col in combined_schema.names
|
||
}
|
||
)
|
||
empty = empty.append_column(
|
||
"_relevance_score", pa.array([], type=pa.float32())
|
||
)
|
||
if not with_row_ids and "_rowid" in empty.column_names:
|
||
empty = empty.drop(["_rowid"])
|
||
return empty
|
||
|
||
original_distances = None
|
||
original_scores = None
|
||
original_distance_row_ids = None
|
||
original_score_row_ids = None
|
||
# normalize the scores to be between 0 and 1, 0 being most relevant
|
||
# We check whether the results (vector and FTS) are empty, because when
|
||
# they are, they often are missing the _rowid column, which causes an error
|
||
if vector_results.num_rows > 0:
|
||
distance_i = vector_results.column_names.index("_distance")
|
||
original_distances = vector_results.column(distance_i)
|
||
original_distance_row_ids = vector_results.column("_rowid")
|
||
vector_results = vector_results.set_column(
|
||
distance_i,
|
||
vector_results.field(distance_i),
|
||
LanceHybridQueryBuilder._normalize_scores(original_distances),
|
||
)
|
||
|
||
# In fts higher scores represent relevance. Not inverting them here as
|
||
# rerankers might need to preserve this score to support `return_score="all"`
|
||
if fts_results.num_rows > 0:
|
||
score_i = fts_results.column_names.index("_score")
|
||
original_scores = fts_results.column(score_i)
|
||
original_score_row_ids = fts_results.column("_rowid")
|
||
fts_results = fts_results.set_column(
|
||
score_i,
|
||
fts_results.field(score_i),
|
||
LanceHybridQueryBuilder._normalize_scores(original_scores),
|
||
)
|
||
|
||
results = reranker.rerank_hybrid(fts_query, vector_results, fts_results)
|
||
|
||
check_reranker_result(results)
|
||
|
||
if "_distance" in results.column_names and original_distances is not None:
|
||
# restore the original distances
|
||
indices = pc.index_in(
|
||
results["_rowid"], original_distance_row_ids, skip_nulls=True
|
||
)
|
||
original_distances = pc.take(original_distances, indices)
|
||
distance_i = results.column_names.index("_distance")
|
||
results = results.set_column(distance_i, "_distance", original_distances)
|
||
|
||
if "_score" in results.column_names and original_scores is not None:
|
||
# restore the original scores
|
||
indices = pc.index_in(
|
||
results["_rowid"], original_score_row_ids, skip_nulls=True
|
||
)
|
||
original_scores = pc.take(original_scores, indices)
|
||
score_i = results.column_names.index("_score")
|
||
results = results.set_column(score_i, "_score", original_scores)
|
||
|
||
results = results.slice(length=limit)
|
||
|
||
if not with_row_ids:
|
||
results = results.drop(["_rowid"])
|
||
|
||
return results
|
||
|
||
def to_batches(
|
||
self, /, batch_size: Optional[int] = None, timeout: Optional[timedelta] = None
|
||
):
|
||
raise NotImplementedError("to_batches not yet supported on a hybrid query")
|
||
|
||
@staticmethod
|
||
def _rank(results: pa.Table, column: str, ascending: bool = True):
|
||
if len(results) == 0:
|
||
return results
|
||
# Get the _score column from results
|
||
scores = results.column(column).to_numpy()
|
||
sort_indices = np.argsort(scores)
|
||
if not ascending:
|
||
sort_indices = sort_indices[::-1]
|
||
ranks = np.empty_like(sort_indices)
|
||
ranks[sort_indices] = np.arange(len(scores)) + 1
|
||
# replace the _score column with the ranks
|
||
_score_idx = results.column_names.index(column)
|
||
results = results.set_column(
|
||
_score_idx, column, pa.array(ranks, type=pa.float32())
|
||
)
|
||
return results
|
||
|
||
@staticmethod
|
||
def _normalize_scores(scores: pa.Array, invert=False) -> pa.Array:
|
||
if len(scores) == 0:
|
||
return scores
|
||
# normalize the scores by subtracting the min and dividing by the max
|
||
min, max = pc.min_max(scores).values()
|
||
rng = pc.subtract(max, min)
|
||
|
||
if not pc.equal(rng, pa.scalar(0.0)).as_py():
|
||
scores = pc.divide(pc.subtract(scores, min), rng)
|
||
elif not pc.equal(max, pa.scalar(0.0)).as_py():
|
||
# If rng is 0, then we at least want the scores to be 0
|
||
scores = pc.subtract(scores, min)
|
||
|
||
if invert:
|
||
scores = pc.subtract(1, scores)
|
||
|
||
return scores
|
||
|
||
def rerank(
|
||
self,
|
||
reranker: Reranker = RRFReranker(),
|
||
normalize: str = "score",
|
||
) -> LanceHybridQueryBuilder:
|
||
"""
|
||
Rerank the hybrid search results using the specified reranker. The reranker
|
||
must be an instance of Reranker class.
|
||
|
||
Parameters
|
||
----------
|
||
reranker: Reranker, default RRFReranker()
|
||
The reranker to use. Must be an instance of Reranker class.
|
||
normalize: str, default "score"
|
||
The method to normalize the scores. Can be "rank" or "score". If "rank",
|
||
the scores are converted to ranks and then normalized. If "score", the
|
||
scores are normalized directly.
|
||
Returns
|
||
-------
|
||
LanceHybridQueryBuilder
|
||
The LanceHybridQueryBuilder object.
|
||
"""
|
||
if normalize not in ["rank", "score"]:
|
||
raise ValueError("normalize must be 'rank' or 'score'.")
|
||
if reranker and not isinstance(reranker, Reranker):
|
||
raise ValueError("reranker must be an instance of Reranker class.")
|
||
|
||
self._norm = normalize
|
||
self._reranker = reranker
|
||
if reranker.score == "all":
|
||
self.with_row_id(True)
|
||
|
||
return self
|
||
|
||
def nprobes(self, nprobes: int) -> LanceHybridQueryBuilder:
|
||
"""
|
||
Set the number of probes to use for vector search.
|
||
|
||
Higher values will yield better recall (more likely to find vectors if
|
||
they exist) at the expense of latency.
|
||
|
||
Parameters
|
||
----------
|
||
nprobes: int
|
||
The number of probes to use.
|
||
|
||
Returns
|
||
-------
|
||
LanceHybridQueryBuilder
|
||
The LanceHybridQueryBuilder object.
|
||
"""
|
||
self._minimum_nprobes = nprobes
|
||
self._maximum_nprobes = nprobes
|
||
return self
|
||
|
||
def minimum_nprobes(self, minimum_nprobes: int) -> LanceHybridQueryBuilder:
|
||
"""Set the minimum number of probes to use.
|
||
|
||
See `nprobes` for more details.
|
||
"""
|
||
self._minimum_nprobes = minimum_nprobes
|
||
return self
|
||
|
||
def maximum_nprobes(self, maximum_nprobes: int) -> LanceHybridQueryBuilder:
|
||
"""Set the maximum number of probes to use.
|
||
|
||
See `nprobes` for more details.
|
||
"""
|
||
self._maximum_nprobes = maximum_nprobes
|
||
return self
|
||
|
||
def distance_range(
|
||
self, lower_bound: Optional[float] = None, upper_bound: Optional[float] = None
|
||
) -> LanceHybridQueryBuilder:
|
||
"""
|
||
Set the distance range to use.
|
||
|
||
Only rows with distances within range [lower_bound, upper_bound)
|
||
will be returned.
|
||
|
||
Parameters
|
||
----------
|
||
lower_bound: Optional[float]
|
||
The lower bound of the distance range.
|
||
upper_bound: Optional[float]
|
||
The upper bound of the distance range.
|
||
|
||
Returns
|
||
-------
|
||
LanceHybridQueryBuilder
|
||
The LanceHybridQueryBuilder object.
|
||
"""
|
||
self._lower_bound = lower_bound
|
||
self._upper_bound = upper_bound
|
||
return self
|
||
|
||
def ef(self, ef: int) -> LanceHybridQueryBuilder:
|
||
"""
|
||
Set the number of candidates to consider during search.
|
||
|
||
Higher values will yield better recall (more likely to find vectors if
|
||
they exist) at the expense of latency.
|
||
|
||
This only applies to the HNSW-related index.
|
||
The default value is 1.5 * limit.
|
||
|
||
Parameters
|
||
----------
|
||
ef: int
|
||
The number of candidates to consider during search.
|
||
|
||
Returns
|
||
-------
|
||
LanceHybridQueryBuilder
|
||
The LanceHybridQueryBuilder object.
|
||
"""
|
||
self._ef = ef
|
||
return self
|
||
|
||
def metric(self, metric: Literal["l2", "cosine", "dot"]) -> LanceHybridQueryBuilder:
|
||
"""Set the distance metric to use.
|
||
|
||
This is an alias for distance_type() and may be deprecated in the future.
|
||
Please use distance_type() instead.
|
||
|
||
Parameters
|
||
----------
|
||
metric: "l2" or "cosine" or "dot"
|
||
The distance metric to use. By default "l2" is used.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
return self.distance_type(metric)
|
||
|
||
def distance_type(
|
||
self, distance_type: Literal["l2", "cosine", "dot"]
|
||
) -> "LanceHybridQueryBuilder":
|
||
"""Set the distance metric to use.
|
||
|
||
When performing a vector search we try and find the "nearest" vectors according
|
||
to some kind of distance metric. This parameter controls which distance metric
|
||
to use.
|
||
|
||
Note: if there is a vector index then the distance type used MUST match the
|
||
distance type used to train the vector index. If this is not done then the
|
||
results will be invalid.
|
||
|
||
Parameters
|
||
----------
|
||
distance_type: "l2" or "cosine" or "dot"
|
||
The distance metric to use. By default "l2" is used.
|
||
|
||
Returns
|
||
-------
|
||
LanceVectorQueryBuilder
|
||
The LanceQueryBuilder object.
|
||
"""
|
||
self._distance_type = distance_type.lower()
|
||
return self
|
||
|
||
def refine_factor(self, refine_factor: int) -> LanceHybridQueryBuilder:
|
||
"""
|
||
Refine the vector search results by reading extra elements and
|
||
re-ranking them in memory.
|
||
|
||
Parameters
|
||
----------
|
||
refine_factor: int
|
||
The refine factor to use.
|
||
|
||
Returns
|
||
-------
|
||
LanceHybridQueryBuilder
|
||
The LanceHybridQueryBuilder object.
|
||
"""
|
||
self._refine_factor = refine_factor
|
||
return self
|
||
|
||
def vector(self, vector: Union[np.ndarray, list]) -> LanceHybridQueryBuilder:
|
||
self._vector = vector
|
||
return self
|
||
|
||
def text(self, text: str | FullTextQuery) -> LanceHybridQueryBuilder:
|
||
self._text = text
|
||
return self
|
||
|
||
def bypass_vector_index(self) -> LanceHybridQueryBuilder:
|
||
"""
|
||
If this is called then any vector index is skipped
|
||
|
||
An exhaustive (flat) search will be performed. The query vector will
|
||
be compared to every vector in the table. At high scales this can be
|
||
expensive. However, this is often still useful. For example, skipping
|
||
the vector index can give you ground truth results which you can use to
|
||
calculate your recall to select an appropriate value for nprobes.
|
||
|
||
Returns
|
||
-------
|
||
LanceHybridQueryBuilder
|
||
The LanceHybridQueryBuilder object.
|
||
"""
|
||
self._bypass_vector_index = True
|
||
return self
|
||
|
||
def explain_plan(self, verbose: Optional[bool] = False) -> str:
|
||
"""Return the execution plan for this query.
|
||
|
||
Examples
|
||
--------
|
||
>>> import lancedb
|
||
>>> db = lancedb.connect("./.lancedb")
|
||
>>> table = db.create_table("my_table", [{"vector": [99.0, 99]}])
|
||
>>> query = [100, 100]
|
||
>>> plan = table.search(query).explain_plan(True)
|
||
>>> print(plan) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
|
||
GlobalLimitExec: skip=0, fetch=10
|
||
FilterExec: _distance@2 IS NOT NULL
|
||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||
KNNVectorDistance: metric=l2
|
||
LanceRead: uri=..., projection=[vector], ...
|
||
|
||
Parameters
|
||
----------
|
||
verbose : bool, default False
|
||
Use a verbose output format.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
""" # noqa: E501
|
||
self._create_query_builders()
|
||
|
||
reranker_label = str(self._reranker) if self._reranker else "No reranker"
|
||
vector_plan = self._table._explain_plan(
|
||
self._vector_query.to_query_object(), verbose=verbose
|
||
)
|
||
fts_plan = self._table._explain_plan(
|
||
self._fts_query.to_query_object(), verbose=verbose
|
||
)
|
||
# Indent sub-plans under the reranker
|
||
indented_vector = "\n".join(" " + line for line in vector_plan.splitlines())
|
||
indented_fts = "\n".join(" " + line for line in fts_plan.splitlines())
|
||
return f"{reranker_label}\n {indented_vector}\n {indented_fts}"
|
||
|
||
def analyze_plan(
|
||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||
) -> str:
|
||
"""Execute the query and display with runtime metrics.
|
||
|
||
Parameters
|
||
----------
|
||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||
Defaults to "aggregate".
|
||
How distributed worker metrics are displayed for remote query plans.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
"""
|
||
self._create_query_builders()
|
||
|
||
results = ["Vector Search Plan:"]
|
||
results.append(
|
||
self._table._analyze_plan(
|
||
self._vector_query.to_query_object(),
|
||
distributed_metrics=distributed_metrics,
|
||
)
|
||
)
|
||
results.append("FTS Search Plan:")
|
||
results.append(
|
||
self._table._analyze_plan(
|
||
self._fts_query.to_query_object(),
|
||
distributed_metrics=distributed_metrics,
|
||
)
|
||
)
|
||
return "\n".join(results)
|
||
|
||
def _create_query_builders(self):
|
||
"""Set up and configure the vector and FTS query builders."""
|
||
vector_query, fts_query = self._validate_query(
|
||
self._query, self._vector, self._text
|
||
)
|
||
self._fts_query = LanceFtsQueryBuilder(
|
||
self._table, fts_query, fts_columns=self._fts_columns
|
||
)
|
||
vector_query = self._query_to_vector(
|
||
self._table, vector_query, self._vector_column
|
||
)
|
||
self._vector_query = LanceVectorQueryBuilder(
|
||
self._table, vector_query, self._vector_column
|
||
)
|
||
|
||
# Apply common configurations
|
||
if self._limit:
|
||
self._vector_query.limit(self._limit)
|
||
self._fts_query.limit(self._limit)
|
||
if self._columns:
|
||
self._vector_query.select(self._columns)
|
||
self._fts_query.select(self._columns)
|
||
if self._where:
|
||
self._vector_query.where(self._where, not self._postfilter)
|
||
self._fts_query.where(self._where, not self._postfilter)
|
||
if self._with_row_id:
|
||
self._vector_query.with_row_id(True)
|
||
self._fts_query.with_row_id(True)
|
||
if self._use_lsm is not None:
|
||
self._vector_query.use_lsm(self._use_lsm)
|
||
self._fts_query.use_lsm(self._use_lsm)
|
||
if self._phrase_query:
|
||
self._fts_query.phrase_query(True)
|
||
if self._distance_type:
|
||
self._vector_query.metric(self._distance_type)
|
||
if self._minimum_nprobes:
|
||
self._vector_query.minimum_nprobes(self._minimum_nprobes)
|
||
if self._maximum_nprobes is not None:
|
||
self._vector_query.maximum_nprobes(self._maximum_nprobes)
|
||
if self._refine_factor:
|
||
self._vector_query.refine_factor(self._refine_factor)
|
||
if self._ef:
|
||
self._vector_query.ef(self._ef)
|
||
if self._bypass_vector_index:
|
||
self._vector_query.bypass_vector_index()
|
||
if self._lower_bound is not None or self._upper_bound is not None:
|
||
self._vector_query.distance_range(
|
||
lower_bound=self._lower_bound, upper_bound=self._upper_bound
|
||
)
|
||
|
||
if self._reranker is None:
|
||
self._reranker = RRFReranker()
|
||
|
||
|
||
class AsyncQueryBase(object):
|
||
"""
|
||
Base class for all async queries (take, scan, vector, fts, hybrid)
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
inner: Union[LanceQuery, LanceVectorQuery, LanceTakeQuery],
|
||
table: Optional["AsyncTable"] = None,
|
||
):
|
||
"""
|
||
Construct an AsyncQueryBase
|
||
|
||
This method is not intended to be called directly. Instead, use the
|
||
[AsyncTable.query][lancedb.table.AsyncTable.query] method to create a query.
|
||
"""
|
||
self._inner = inner
|
||
self._table = table
|
||
self._with_row_address = None
|
||
self._fragments = None
|
||
self._fragment_ids = None
|
||
self._with_row_id = None
|
||
self._blob_auto_row_id = False
|
||
self._blob_paths: tuple[str, ...] = ()
|
||
|
||
def to_query_object(self) -> Query:
|
||
"""
|
||
Convert the query into a query object
|
||
|
||
This is currently experimental but can be useful as the query object is pure
|
||
python and more easily serializable.
|
||
"""
|
||
query = Query.from_inner(self._inner.to_query_request())
|
||
query.with_row_id = self._user_requested_row_id()
|
||
query.with_row_address = self._with_row_address
|
||
query.fragments = self._fragments
|
||
query.fragment_ids = self._fragment_ids
|
||
return query
|
||
|
||
def _user_requested_row_id(self) -> bool:
|
||
return self._with_row_id is True
|
||
|
||
def _blob_auto_row_id_enabled(self) -> bool:
|
||
return self._blob_auto_row_id
|
||
|
||
def _finalize_blob_query_table(self, tbl: pa.Table) -> pa.Table:
|
||
return finalize_blob_query_table(
|
||
tbl,
|
||
user_requested_row_id=self._user_requested_row_id(),
|
||
blob_auto_row_id=self._blob_auto_row_id_enabled(),
|
||
blob_paths=self._blob_paths,
|
||
)
|
||
|
||
async def _maybe_add_blob_row_id(self) -> None:
|
||
if self._table is None or not supports_blob_auto_row_id(self._table):
|
||
self._blob_auto_row_id = False
|
||
self._blob_paths = ()
|
||
return
|
||
|
||
req = self._inner.to_query_request()
|
||
schema = await self._table.schema()
|
||
self._blob_auto_row_id = blob_auto_row_id_for_scan(
|
||
self._table,
|
||
schema,
|
||
req.select,
|
||
with_row_id=self._with_row_id,
|
||
)
|
||
if not self._blob_auto_row_id:
|
||
self._blob_paths = ()
|
||
return
|
||
self._blob_paths = tuple(blob_v2_projection_sources(schema, req.select).keys())
|
||
self._inner.with_row_id()
|
||
|
||
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
|
||
"""
|
||
Return only the specified columns.
|
||
|
||
By default a query will return all columns from the table. However, this can
|
||
have a very significant impact on latency. LanceDb stores data in a columnar
|
||
fashion. This
|
||
means we can finely tune our I/O to select exactly the columns we need.
|
||
|
||
As a best practice you should always limit queries to the columns that you need.
|
||
If you pass in a list of column names then only those columns will be
|
||
returned.
|
||
|
||
You can also use this method to create new "dynamic" columns based on your
|
||
existing columns. For example, you may not care about "a" or "b" but instead
|
||
simply want "a + b". This is often seen in the SELECT clause of an SQL query
|
||
(e.g. `SELECT a+b FROM my_table`).
|
||
|
||
To create dynamic columns you can pass in a dict[str, str]. A column will be
|
||
returned for each entry in the map. The key provides the name of the column.
|
||
The value is an SQL string used to specify how the column is calculated.
|
||
|
||
For example, an SQL query might state `SELECT a + b AS combined, c`. The
|
||
equivalent input to this method would be `{"combined": "a + b", "c": "c"}`.
|
||
|
||
Columns will always be returned in the order given, even if that order is
|
||
different than the order used when adding the data.
|
||
"""
|
||
if isinstance(columns, list) and all(isinstance(c, str) for c in columns):
|
||
self._inner.select_columns(columns)
|
||
elif isinstance(columns, dict) and all(isinstance(k, str) for k in columns):
|
||
if any(isinstance(v, Expr) for v in columns.values()):
|
||
# At least one value is an Expr — use the type-safe path.
|
||
from .expr import _coerce
|
||
|
||
pairs = [(k, _coerce(v)._inner) for k, v in columns.items()]
|
||
self._inner.select_expr(pairs)
|
||
elif all(isinstance(v, str) for v in columns.values()):
|
||
self._inner.select(list(columns.items()))
|
||
else:
|
||
raise TypeError(
|
||
"dict values must be str or Expr, got "
|
||
+ str({k: type(v) for k, v in columns.items()})
|
||
)
|
||
else:
|
||
raise TypeError("columns must be a list of column names or a dict")
|
||
return self
|
||
|
||
def with_row_id(self) -> Self:
|
||
"""
|
||
Include the _rowid column in the results.
|
||
"""
|
||
self._with_row_id = True
|
||
self._inner.with_row_id()
|
||
return self
|
||
|
||
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||
"""
|
||
Include the _rowaddr column in scanner-backed plain query results.
|
||
"""
|
||
self._with_row_address = with_row_address
|
||
return self
|
||
|
||
def with_fragments(self, fragments: Any) -> Self:
|
||
"""
|
||
Restrict scanner-backed plain query results to the given Lance fragments.
|
||
"""
|
||
self._fragments = fragments
|
||
return self
|
||
|
||
def fragment_ids(self, fragment_ids: List[int]) -> Self:
|
||
"""
|
||
Restrict scanner-backed plain query results to the given Lance fragment ids.
|
||
"""
|
||
self._fragment_ids = fragment_ids
|
||
return self
|
||
|
||
async def to_batches(
|
||
self,
|
||
*,
|
||
max_batch_length: Optional[int] = None,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> AsyncRecordBatchReader:
|
||
"""
|
||
Execute the query and return the results as an Apache Arrow RecordBatchReader.
|
||
|
||
Parameters
|
||
----------
|
||
|
||
max_batch_length: Optional[int]
|
||
The maximum number of selected records in a single RecordBatch object.
|
||
If not specified, a default batch length is used.
|
||
It is possible for batches to be smaller than the provided length if the
|
||
underlying data is stored in smaller chunks.
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
"""
|
||
await self._maybe_add_blob_row_id()
|
||
return AsyncRecordBatchReader(
|
||
await self._inner.execute(
|
||
max_batch_length=max_batch_length, timeout=timeout
|
||
)
|
||
)
|
||
|
||
async def output_schema(self) -> pa.Schema:
|
||
"""
|
||
Return the output schema for the query
|
||
|
||
This does not execute the query.
|
||
"""
|
||
return await self._inner.output_schema()
|
||
|
||
async def to_arrow(self, timeout: Optional[timedelta] = None) -> pa.Table:
|
||
"""
|
||
Execute the query and collect the results into an Apache Arrow Table.
|
||
|
||
This method will collect all results into memory before returning. If
|
||
you expect a large number of results, you may want to use
|
||
[to_batches][lancedb.query.AsyncQueryBase.to_batches]
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
"""
|
||
batch_iter = await self.to_batches(timeout=timeout)
|
||
return self._finalize_blob_query_table(
|
||
pa.Table.from_batches(await batch_iter.read_all(), schema=batch_iter.schema)
|
||
)
|
||
|
||
async def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
||
"""
|
||
Execute the query and return the results as a list of dictionaries.
|
||
|
||
Each list entry is a dictionary with the selected column names as keys,
|
||
or all table columns if `select` is not called. The vector and the "_distance"
|
||
fields are returned whether or not they're explicitly selected.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
"""
|
||
return (await self.to_arrow(timeout=timeout)).to_pylist()
|
||
|
||
async def to_pandas(
|
||
self,
|
||
flatten: Optional[Union[int, bool]] = None,
|
||
timeout: Optional[timedelta] = None,
|
||
*,
|
||
blob_mode: BlobMode = "lazy",
|
||
**kwargs,
|
||
) -> "pd.DataFrame":
|
||
"""
|
||
Execute the query and collect the results into a pandas DataFrame.
|
||
|
||
This method will collect all results into memory before returning. If you
|
||
expect a large number of results, you may want to use
|
||
[to_batches][lancedb.query.AsyncQueryBase.to_batches] and convert each batch to
|
||
pandas separately.
|
||
|
||
Examples
|
||
--------
|
||
|
||
>>> import asyncio
|
||
>>> from lancedb import connect_async
|
||
>>> async def doctest_example():
|
||
... conn = await connect_async("./.lancedb")
|
||
... table = await conn.create_table("my_table", data=[{"a": 1, "b": 2}])
|
||
... async for batch in await table.query().to_batches():
|
||
... batch_df = batch.to_pandas()
|
||
>>> asyncio.run(doctest_example())
|
||
|
||
Parameters
|
||
----------
|
||
flatten: Optional[Union[int, bool]]
|
||
If flatten is True, flatten all nested columns.
|
||
If flatten is an integer, flatten the nested columns up to the
|
||
specified depth.
|
||
If unspecified, do not flatten the nested columns.
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
blob_mode: str, default "lazy"
|
||
Controls how blob columns are returned for plain scan queries.
|
||
Vector, FTS, hybrid, and other non-native query shapes keep the
|
||
existing Arrow conversion path and only support blob descriptions.
|
||
**kwargs
|
||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||
optional flattening.
|
||
"""
|
||
validate_blob_mode(blob_mode)
|
||
if hasattr(self._inner, "output_schema"):
|
||
schema = await self.output_schema()
|
||
if _blob_mode_requires_native_pandas(blob_mode, schema):
|
||
native_error = None
|
||
if (flatten is None or blob_mode == "descriptions") and timeout is None:
|
||
try:
|
||
df = await self._plain_scan_to_pandas(
|
||
blob_mode, flatten=flatten, **kwargs
|
||
)
|
||
if df is not None:
|
||
return df
|
||
except Exception as err:
|
||
native_error = err
|
||
reason = (
|
||
"this query shape cannot use Lance native pandas conversion"
|
||
if native_error is None
|
||
else str(native_error)
|
||
)
|
||
raise _unsupported_blob_pandas_error(reason) from native_error
|
||
|
||
tbl = flatten_columns(await self.to_arrow(timeout=timeout), flatten)
|
||
if _blob_mode_requires_native_pandas(blob_mode, tbl.schema):
|
||
raise _unsupported_blob_pandas_error(
|
||
"this query shape cannot use Lance native pandas conversion"
|
||
)
|
||
return tbl.to_pandas(**kwargs)
|
||
|
||
async def _plain_scan_to_pandas(
|
||
self,
|
||
blob_mode: BlobMode,
|
||
flatten: Optional[Union[int, bool]] = None,
|
||
**kwargs,
|
||
) -> Optional["pd.DataFrame"]:
|
||
if self._table is None:
|
||
return None
|
||
|
||
query = self.to_query_object()
|
||
if not _query_is_plain_scan(query):
|
||
return None
|
||
|
||
schema = await self._table.schema()
|
||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||
self._table,
|
||
schema,
|
||
query.columns,
|
||
with_row_id=self._with_row_id,
|
||
)
|
||
blob_sources = (
|
||
blob_v2_projection_sources(schema, query.columns)
|
||
if blob_mode == "bytes"
|
||
else {}
|
||
)
|
||
dataset = await self._table._to_lance()
|
||
scanner = dataset.scanner(
|
||
**_scanner_kwargs_for_query(
|
||
query,
|
||
"descriptions" if blob_sources else blob_mode,
|
||
dataset,
|
||
with_row_id=query.with_row_id or blob_auto_row_id or bool(blob_sources),
|
||
)
|
||
)
|
||
return await _finish_plain_scan_pandas_async(
|
||
scanner,
|
||
blob_mode=blob_mode,
|
||
blob_sources=blob_sources,
|
||
fetch_blobs=self._table.fetch_blobs,
|
||
strip_auto_row_id=blob_auto_row_id,
|
||
flatten=flatten,
|
||
**kwargs,
|
||
)
|
||
|
||
async def to_polars(
|
||
self,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> "pl.DataFrame":
|
||
"""
|
||
Execute the query and collect the results into a Polars DataFrame.
|
||
|
||
This method will collect all results into memory before returning. If you
|
||
expect a large number of results, you may want to use
|
||
[to_batches][lancedb.query.AsyncQueryBase.to_batches] and convert each batch to
|
||
polars separately.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
|
||
Examples
|
||
--------
|
||
|
||
>>> import asyncio
|
||
>>> import polars as pl
|
||
>>> from lancedb import connect_async
|
||
>>> async def doctest_example():
|
||
... conn = await connect_async("./.lancedb")
|
||
... table = await conn.create_table("my_table", data=[{"a": 1, "b": 2}])
|
||
... async for batch in await table.query().to_batches():
|
||
... batch_df = pl.from_arrow(batch)
|
||
>>> asyncio.run(doctest_example())
|
||
"""
|
||
import polars as pl
|
||
|
||
return pl.from_arrow(await self.to_arrow(timeout=timeout))
|
||
|
||
async def to_pydantic(
|
||
self, model: Type[LanceModel], *, timeout: Optional[timedelta] = None
|
||
) -> List[LanceModel]:
|
||
"""
|
||
Convert results to a list of pydantic models.
|
||
|
||
Parameters
|
||
----------
|
||
model : Type[LanceModel]
|
||
The pydantic model to use.
|
||
timeout : timedelta, optional
|
||
The maximum time to wait for the query to complete.
|
||
If None, wait indefinitely.
|
||
|
||
Returns
|
||
-------
|
||
list[LanceModel]
|
||
"""
|
||
return [
|
||
model(**row) for row in (await self.to_arrow(timeout=timeout)).to_pylist()
|
||
]
|
||
|
||
async def explain_plan(self, verbose: Optional[bool] = False):
|
||
"""Return the execution plan for this query.
|
||
|
||
Examples
|
||
--------
|
||
>>> import asyncio
|
||
>>> from lancedb import connect_async
|
||
>>> async def doctest_example():
|
||
... conn = await connect_async("./.lancedb")
|
||
... table = await conn.create_table("my_table", [{"vector": [99.0, 99.0]}])
|
||
... plan = await table.query().nearest_to([1.0, 2.0]).explain_plan(True)
|
||
... print(plan)
|
||
>>> asyncio.run(doctest_example()) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
|
||
GlobalLimitExec: skip=0, fetch=10
|
||
FilterExec: _distance@2 IS NOT NULL
|
||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||
KNNVectorDistance: metric=l2
|
||
LanceRead: uri=..., projection=[vector], ...
|
||
<BLANKLINE>
|
||
|
||
Parameters
|
||
----------
|
||
verbose : bool, default False
|
||
Use a verbose output format.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
""" # noqa: E501
|
||
return await self._inner.explain_plan(verbose)
|
||
|
||
async def analyze_plan(
|
||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||
) -> str:
|
||
"""Execute the query and display with runtime metrics.
|
||
|
||
Parameters
|
||
----------
|
||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||
Defaults to "aggregate".
|
||
How distributed worker metrics are displayed for remote query plans.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
"""
|
||
return await self._inner.analyze_plan(distributed_metrics)
|
||
|
||
|
||
class AsyncStandardQuery(AsyncQueryBase):
|
||
"""
|
||
Base class for "standard" async queries (all but take currently)
|
||
"""
|
||
|
||
def __init__(
|
||
self,
|
||
inner: Union[LanceQuery, LanceVectorQuery],
|
||
table: Optional["AsyncTable"] = None,
|
||
):
|
||
"""
|
||
Construct an AsyncStandardQuery
|
||
|
||
This method is not intended to be called directly. Instead, use the
|
||
[AsyncTable.query][lancedb.table.AsyncTable.query] method to create a query.
|
||
"""
|
||
super().__init__(inner, table)
|
||
|
||
def where(self, predicate: Union[str, Expr]) -> Self:
|
||
"""
|
||
Only return rows matching the given predicate
|
||
|
||
The predicate can be a SQL string or a type-safe
|
||
:class:`~lancedb.expr.Expr` built with :func:`~lancedb.expr.col`
|
||
and :func:`~lancedb.expr.lit`.
|
||
|
||
Examples
|
||
--------
|
||
|
||
>>> predicate = "x > 10"
|
||
>>> predicate = "y > 0 AND y < 100"
|
||
>>> predicate = "x > 5 OR y = 'test'"
|
||
|
||
Filtering performance can often be improved by creating a scalar index
|
||
on the filter column(s).
|
||
|
||
Calling this multiple times combines the filters with a logical AND
|
||
rather than replacing the previous filter.
|
||
"""
|
||
if isinstance(predicate, Expr):
|
||
self._inner.where_expr(predicate._inner)
|
||
else:
|
||
self._inner.where(predicate)
|
||
return self
|
||
|
||
def limit(self, limit: int) -> Self:
|
||
"""
|
||
Set the maximum number of results to return.
|
||
|
||
By default, a plain search has no limit. If this method is not
|
||
called then every valid row from the table will be returned.
|
||
"""
|
||
self._inner.limit(limit)
|
||
return self
|
||
|
||
def offset(self, offset: int) -> Self:
|
||
"""
|
||
Set the offset for the results.
|
||
|
||
Parameters
|
||
----------
|
||
offset: int
|
||
The offset to start fetching results from.
|
||
"""
|
||
self._inner.offset(offset)
|
||
return self
|
||
|
||
def order_by(self, ordering: Optional[List[ColumnOrdering]]) -> Self:
|
||
"""
|
||
Set the ordering for the results.
|
||
|
||
Parameters
|
||
----------
|
||
ordering: Optional[List[ColumnOrdering]]
|
||
The ordering to use for the results. If None, then the default ordering
|
||
will be used.
|
||
"""
|
||
if ordering is None:
|
||
self._inner.order_by(None)
|
||
else:
|
||
self._inner.order_by(
|
||
[
|
||
o.model_dump() if hasattr(o, "model_dump") else o.dict()
|
||
for o in ordering
|
||
]
|
||
)
|
||
return self
|
||
|
||
def fast_search(self) -> Self:
|
||
"""
|
||
Skip searching un-indexed data.
|
||
|
||
This can make queries faster, but will miss any data that has not been
|
||
indexed.
|
||
|
||
!!! tip
|
||
You can add new data into an existing index by calling
|
||
[AsyncTable.optimize][lancedb.table.AsyncTable.optimize].
|
||
"""
|
||
self._inner.fast_search()
|
||
return self
|
||
|
||
def use_lsm(self, enable: bool) -> Self:
|
||
"""
|
||
Control MemWAL LSM read routing for this query.
|
||
|
||
By default (unset), a query against a table with an LSM write spec (see
|
||
[AsyncTable.set_lsm_write_spec][lancedb.table.AsyncTable.set_lsm_write_spec])
|
||
is routed through the LSM scanner so it also returns data written via the
|
||
``merge_insert`` LSM path that has not yet been compacted into the base
|
||
table (the active/frozen in-memory memtables and the flushed generations),
|
||
deduplicated by primary key; a table without a spec reads the base table.
|
||
|
||
Parameters
|
||
----------
|
||
enable : bool
|
||
``True`` forces the LSM scanner and errors if the table has no LSM
|
||
write spec. ``False`` bypasses the MemWAL and reads the base table
|
||
only, even when a spec is present.
|
||
"""
|
||
self._inner.use_lsm(enable)
|
||
return self
|
||
|
||
def postfilter(self) -> Self:
|
||
"""
|
||
If this is called then filtering will happen after the search instead of
|
||
before.
|
||
By default filtering will be performed before the search. This is how
|
||
filtering is typically understood to work. This prefilter step does add some
|
||
additional latency. Creating a scalar index on the filter column(s) can
|
||
often improve this latency. However, sometimes a filter is too complex or
|
||
scalar indices cannot be applied to the column. In these cases postfiltering
|
||
can be used instead of prefiltering to improve latency.
|
||
Post filtering applies the filter to the results of the search. This
|
||
means we only run the filter on a much smaller set of data. However, it can
|
||
cause the query to return fewer than `limit` results (or even no results) if
|
||
none of the nearest results match the filter.
|
||
Post filtering happens during the "refine stage" (described in more detail in
|
||
@see {@link VectorQuery#refineFactor}). This means that setting a higher refine
|
||
factor can often help restore some of the results lost by post filtering.
|
||
"""
|
||
self._inner.postfilter()
|
||
return self
|
||
|
||
|
||
class AsyncQuery(AsyncStandardQuery):
|
||
def __init__(self, inner: LanceQuery, table: Optional["AsyncTable"] = None):
|
||
"""
|
||
Construct an AsyncQuery
|
||
|
||
This method is not intended to be called directly. Instead, use the
|
||
[AsyncTable.query][lancedb.table.AsyncTable.query] method to create a query.
|
||
"""
|
||
super().__init__(inner, table)
|
||
self._inner = inner
|
||
|
||
@classmethod
|
||
def _query_vec_to_array(self, vec: Union[VEC, Tuple]):
|
||
if isinstance(vec, list):
|
||
return pa.array(vec)
|
||
if isinstance(vec, np.ndarray):
|
||
return pa.array(vec)
|
||
if isinstance(vec, pa.Array):
|
||
return vec
|
||
if isinstance(vec, pa.ChunkedArray):
|
||
return vec.combine_chunks()
|
||
if isinstance(vec, tuple):
|
||
return pa.array(vec)
|
||
# We've checked everything we formally support in our typings
|
||
# but, as a fallback, let pyarrow try and convert it anyway.
|
||
# This can allow for some more exotic things like iterables
|
||
return pa.array(vec)
|
||
|
||
def nearest_to(
|
||
self,
|
||
query_vector: Union[VEC, Tuple, List[VEC]],
|
||
) -> AsyncVectorQuery:
|
||
"""
|
||
Find the nearest vectors to the given query vector.
|
||
|
||
This converts the query from a plain query to a vector query.
|
||
|
||
This method will attempt to convert the input to the query vector
|
||
expected by the embedding model. If the input cannot be converted
|
||
then an error will be thrown.
|
||
|
||
By default, there is no embedding model, and the input should be
|
||
something that can be converted to a pyarrow array of floats. This
|
||
includes lists, numpy arrays, and tuples.
|
||
|
||
If there is only one vector column (a column whose data type is a
|
||
fixed size list of floats) then the column does not need to be specified.
|
||
If there is more than one vector column you must use
|
||
[AsyncVectorQuery.column][lancedb.query.AsyncVectorQuery.column] to specify
|
||
which column you would like to compare with.
|
||
|
||
If no index has been created on the vector column then a vector query
|
||
will perform a distance comparison between the query vector and every
|
||
vector in the database and then sort the results. This is sometimes
|
||
called a "flat search"
|
||
|
||
For small databases, with tens of thousands of vectors or less, this can
|
||
be reasonably fast. In larger databases you should create a vector index
|
||
on the column. If there is a vector index then an "approximate" nearest
|
||
neighbor search (frequently called an ANN search) will be performed. This
|
||
search is much faster, but the results will be approximate.
|
||
|
||
The query can be further parameterized using the returned builder. There
|
||
are various ANN search parameters that will let you fine tune your recall
|
||
accuracy vs search latency.
|
||
|
||
Vector searches always have a [limit][]. If `limit` has not been called then
|
||
a default `limit` of 10 will be used.
|
||
|
||
Typically, a single vector is passed in as the query. However, you can also
|
||
pass in multiple vectors. When multiple vectors are passed in, if the vector
|
||
column is with multivector type, then the vectors will be treated as a single
|
||
query. Or the vectors will be treated as multiple queries, this can be useful
|
||
if you want to find the nearest vectors to multiple query vectors.
|
||
This is not expected to be faster than making multiple queries concurrently;
|
||
it is just a convenience method. If multiple vectors are passed in then
|
||
an additional column `query_index` will be added to the results. This column
|
||
will contain the index of the query vector that the result is nearest to.
|
||
"""
|
||
if query_vector is None:
|
||
raise ValueError("query_vector can not be None")
|
||
|
||
if (
|
||
isinstance(query_vector, (list, np.ndarray, pa.Array))
|
||
and len(query_vector) > 0
|
||
and isinstance(query_vector[0], (list, np.ndarray, pa.Array))
|
||
):
|
||
# multiple have been passed
|
||
query_vectors = [AsyncQuery._query_vec_to_array(v) for v in query_vector]
|
||
new_self = self._inner.nearest_to(query_vectors[0])
|
||
for v in query_vectors[1:]:
|
||
new_self.add_query_vector(v)
|
||
return AsyncVectorQuery(new_self, self._table)
|
||
else:
|
||
return AsyncVectorQuery(
|
||
self._inner.nearest_to(AsyncQuery._query_vec_to_array(query_vector)),
|
||
self._table,
|
||
)
|
||
|
||
def nearest_to_text(
|
||
self, query: str | FullTextQuery, columns: Union[str, List[str], None] = None
|
||
) -> AsyncFTSQuery:
|
||
"""
|
||
Find the documents that are most relevant to the given text query.
|
||
|
||
This method will perform a full text search on the table and return
|
||
the most relevant documents. The relevance is determined by BM25.
|
||
|
||
The columns to search must be with native FTS index
|
||
(Tantivy-based can't work with this method).
|
||
|
||
By default, all indexed columns are searched,
|
||
now only one column can be searched at a time.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The text query to search for.
|
||
columns: str or list of str, default None
|
||
The columns to search in. If None, all indexed columns are searched.
|
||
For now only one column can be searched at a time.
|
||
"""
|
||
if isinstance(columns, str):
|
||
columns = [columns]
|
||
if columns is None:
|
||
columns = []
|
||
|
||
if isinstance(query, str):
|
||
return AsyncFTSQuery(
|
||
self._inner.nearest_to_text({"query": query, "columns": columns}),
|
||
self._table,
|
||
)
|
||
# FullTextQuery object
|
||
return AsyncFTSQuery(self._inner.nearest_to_text({"query": query}), self._table)
|
||
|
||
|
||
class AsyncFTSQuery(AsyncStandardQuery):
|
||
"""A query for full text search for LanceDB."""
|
||
|
||
def __init__(self, inner: LanceFTSQuery, table: Optional["AsyncTable"] = None):
|
||
super().__init__(inner, table)
|
||
self._inner = inner
|
||
self._reranker = None
|
||
|
||
def get_query(self) -> str:
|
||
return self._inner.get_query()
|
||
|
||
def rerank(
|
||
self,
|
||
reranker: Reranker = RRFReranker(),
|
||
) -> AsyncFTSQuery:
|
||
if reranker and not isinstance(reranker, Reranker):
|
||
raise ValueError("reranker must be an instance of Reranker class.")
|
||
|
||
self._reranker = reranker
|
||
|
||
return self
|
||
|
||
def nearest_to(
|
||
self,
|
||
query_vector: Union[VEC, Tuple, List[VEC]],
|
||
) -> AsyncHybridQuery:
|
||
"""
|
||
In addition doing text search on the LanceDB Table, also
|
||
find the nearest vectors to the given query vector.
|
||
|
||
This converts the query from a FTS Query to a Hybrid query. Results
|
||
from the vector search will be combined with results from the FTS query.
|
||
|
||
This method will attempt to convert the input to the query vector
|
||
expected by the embedding model. If the input cannot be converted
|
||
then an error will be thrown.
|
||
|
||
By default, there is no embedding model, and the input should be
|
||
something that can be converted to a pyarrow array of floats. This
|
||
includes lists, numpy arrays, and tuples.
|
||
|
||
If there is only one vector column (a column whose data type is a
|
||
fixed size list of floats) then the column does not need to be specified.
|
||
If there is more than one vector column you must use
|
||
[AsyncVectorQuery.column][lancedb.query.AsyncVectorQuery.column] to specify
|
||
which column you would like to compare with.
|
||
|
||
If no index has been created on the vector column then a vector query
|
||
will perform a distance comparison between the query vector and every
|
||
vector in the database and then sort the results. This is sometimes
|
||
called a "flat search"
|
||
|
||
For small databases, with tens of thousands of vectors or less, this can
|
||
be reasonably fast. In larger databases you should create a vector index
|
||
on the column. If there is a vector index then an "approximate" nearest
|
||
neighbor search (frequently called an ANN search) will be performed. This
|
||
search is much faster, but the results will be approximate.
|
||
|
||
The query can be further parameterized using the returned builder. There
|
||
are various ANN search parameters that will let you fine tune your recall
|
||
accuracy vs search latency.
|
||
|
||
Hybrid searches always have a [limit][]. If `limit` has not been called then
|
||
a default `limit` of 10 will be used.
|
||
|
||
Typically, a single vector is passed in as the query. However, you can also
|
||
pass in multiple vectors. This can be useful if you want to find the nearest
|
||
vectors to multiple query vectors. This is not expected to be faster than
|
||
making multiple queries concurrently; it is just a convenience method.
|
||
If multiple vectors are passed in then an additional column `query_index`
|
||
will be added to the results. This column will contain the index of the
|
||
query vector that the result is nearest to.
|
||
"""
|
||
if query_vector is None:
|
||
raise ValueError("query_vector can not be None")
|
||
|
||
if (
|
||
isinstance(query_vector, list)
|
||
and len(query_vector) > 0
|
||
and not isinstance(query_vector[0], (float, int))
|
||
):
|
||
# multiple have been passed
|
||
query_vectors = [AsyncQuery._query_vec_to_array(v) for v in query_vector]
|
||
new_self = self._inner.nearest_to(query_vectors[0])
|
||
for v in query_vectors[1:]:
|
||
new_self.add_query_vector(v)
|
||
return AsyncHybridQuery(new_self, self._table)
|
||
else:
|
||
return AsyncHybridQuery(
|
||
self._inner.nearest_to(AsyncQuery._query_vec_to_array(query_vector)),
|
||
self._table,
|
||
)
|
||
|
||
async def to_batches(
|
||
self,
|
||
*,
|
||
max_batch_length: Optional[int] = None,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> AsyncRecordBatchReader:
|
||
reader = await super().to_batches(timeout=timeout)
|
||
results = pa.Table.from_batches(await reader.read_all(), reader.schema)
|
||
if self._reranker:
|
||
results = self._reranker.rerank_fts(self.get_query(), results)
|
||
return AsyncRecordBatchReader(results, max_batch_length=max_batch_length)
|
||
|
||
|
||
class AsyncVectorQueryBase:
|
||
def column(self, column: str) -> Self:
|
||
"""
|
||
Set the vector column to query
|
||
|
||
This controls which column is compared to the query vector supplied in
|
||
the call to [AsyncQuery.nearest_to][lancedb.query.AsyncQuery.nearest_to].
|
||
|
||
This parameter must be specified if the table has more than one column
|
||
whose data type is a fixed-size-list of floats.
|
||
"""
|
||
self._inner.column(column)
|
||
return self
|
||
|
||
def nprobes(self, nprobes: int) -> Self:
|
||
"""
|
||
Set the number of partitions to search (probe)
|
||
|
||
This argument is only used when the vector column has an IVF-based index.
|
||
If there is no index then this value is ignored.
|
||
|
||
The IVF stage of IVF PQ divides the input into partitions (clusters) of
|
||
related values.
|
||
|
||
The partition whose centroids are closest to the query vector will be
|
||
exhaustiely searched to find matches. This parameter controls how many
|
||
partitions should be searched.
|
||
|
||
Increasing this value will increase the recall of your query but will
|
||
also increase the latency of your query. The default value is 20. This
|
||
default is good for many cases but the best value to use will depend on
|
||
your data and the recall that you need to achieve.
|
||
|
||
For best results we recommend tuning this parameter with a benchmark against
|
||
your actual data to find the smallest possible value that will still give
|
||
you the desired recall.
|
||
"""
|
||
self._inner.nprobes(nprobes)
|
||
return self
|
||
|
||
def minimum_nprobes(self, minimum_nprobes: int) -> Self:
|
||
"""Set the minimum number of probes to use.
|
||
|
||
See `nprobes` for more details.
|
||
|
||
These partitions will be searched on every indexed vector query and will
|
||
increase recall at the expense of latency.
|
||
"""
|
||
self._inner.minimum_nprobes(minimum_nprobes)
|
||
return self
|
||
|
||
def maximum_nprobes(self, maximum_nprobes: int) -> Self:
|
||
"""Set the maximum number of probes to use.
|
||
|
||
See `nprobes` for more details.
|
||
|
||
If this value is greater than `minimum_nprobes` then the excess partitions
|
||
will be searched only if we have not found enough results.
|
||
|
||
This can be useful when there is a narrow filter to allow these queries to
|
||
spend more time searching and avoid potential false negatives.
|
||
|
||
If this value is 0 then no limit will be applied and all partitions could be
|
||
searched if needed to satisfy the limit.
|
||
"""
|
||
self._inner.maximum_nprobes(maximum_nprobes)
|
||
return self
|
||
|
||
def distance_range(
|
||
self, lower_bound: Optional[float] = None, upper_bound: Optional[float] = None
|
||
) -> Self:
|
||
"""Set the distance range to use.
|
||
|
||
Only rows with distances within range [lower_bound, upper_bound)
|
||
will be returned.
|
||
|
||
Parameters
|
||
----------
|
||
lower_bound: Optional[float]
|
||
The lower bound of the distance range.
|
||
upper_bound: Optional[float]
|
||
The upper bound of the distance range.
|
||
|
||
Returns
|
||
-------
|
||
AsyncVectorQuery
|
||
The AsyncVectorQuery object.
|
||
"""
|
||
self._inner.distance_range(lower_bound, upper_bound)
|
||
return self
|
||
|
||
def ef(self, ef: int) -> Self:
|
||
"""
|
||
Set the number of candidates to consider during search
|
||
|
||
This argument is only used when the vector column has an HNSW index.
|
||
If there is no index then this value is ignored.
|
||
|
||
Increasing this value will increase the recall of your query but will also
|
||
increase the latency of your query. The default value is 1.5 * limit. This
|
||
default is good for many cases but the best value to use will depend on your
|
||
data and the recall that you need to achieve.
|
||
"""
|
||
self._inner.ef(ef)
|
||
return self
|
||
|
||
def refine_factor(self, refine_factor: int) -> Self:
|
||
"""
|
||
A multiplier to control how many additional rows are taken during the refine
|
||
step
|
||
|
||
This argument is only used when the vector column has an IVF PQ index.
|
||
If there is no index then this value is ignored.
|
||
|
||
An IVF PQ index stores compressed (quantized) values. They query vector is
|
||
compared against these values and, since they are compressed, the comparison is
|
||
inaccurate.
|
||
|
||
This parameter can be used to refine the results. It can improve both improve
|
||
recall and correct the ordering of the nearest results.
|
||
|
||
To refine results LanceDb will first perform an ANN search to find the nearest
|
||
`limit` * `refine_factor` results. In other words, if `refine_factor` is 3 and
|
||
`limit` is the default (10) then the first 30 results will be selected. LanceDb
|
||
then fetches the full, uncompressed, values for these 30 results. The results
|
||
are then reordered by the true distance and only the nearest 10 are kept.
|
||
|
||
Note: there is a difference between calling this method with a value of 1 and
|
||
never calling this method at all. Calling this method with any value will have
|
||
an impact on your search latency. When you call this method with a
|
||
`refine_factor` of 1 then LanceDb still needs to fetch the full, uncompressed,
|
||
values so that it can potentially reorder the results.
|
||
|
||
Note: if this method is NOT called then the distances returned in the _distance
|
||
column will be approximate distances based on the comparison of the quantized
|
||
query vector and the quantized result vectors. This can be considerably
|
||
different than the true distance between the query vector and the actual
|
||
uncompressed vector.
|
||
"""
|
||
self._inner.refine_factor(refine_factor)
|
||
return self
|
||
|
||
def distance_type(self, distance_type: str) -> Self:
|
||
"""
|
||
Set the distance metric to use
|
||
|
||
When performing a vector search we try and find the "nearest" vectors according
|
||
to some kind of distance metric. This parameter controls which distance metric
|
||
to use. See @see {@link IvfPqOptions.distanceType} for more details on the
|
||
different distance metrics available.
|
||
|
||
Note: if there is a vector index then the distance type used MUST match the
|
||
distance type used to train the vector index. If this is not done then the
|
||
results will be invalid.
|
||
|
||
By default "l2" is used.
|
||
"""
|
||
self._inner.distance_type(distance_type)
|
||
return self
|
||
|
||
def bypass_vector_index(self) -> Self:
|
||
"""
|
||
If this is called then any vector index is skipped
|
||
|
||
An exhaustive (flat) search will be performed. The query vector will
|
||
be compared to every vector in the table. At high scales this can be
|
||
expensive. However, this is often still useful. For example, skipping
|
||
the vector index can give you ground truth results which you can use to
|
||
calculate your recall to select an appropriate value for nprobes.
|
||
"""
|
||
self._inner.bypass_vector_index()
|
||
return self
|
||
|
||
|
||
class AsyncVectorQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||
def __init__(self, inner: LanceVectorQuery, table: Optional["AsyncTable"] = None):
|
||
"""
|
||
Construct an AsyncVectorQuery
|
||
|
||
This method is not intended to be called directly. Instead, create
|
||
a query first with [AsyncTable.query][lancedb.table.AsyncTable.query] and then
|
||
use [AsyncQuery.nearest_to][lancedb.query.AsyncQuery.nearest_to]] to convert to
|
||
a vector query. Or you can use
|
||
[AsyncTable.vector_search][lancedb.table.AsyncTable.vector_search]
|
||
"""
|
||
super().__init__(inner, table)
|
||
self._inner = inner
|
||
self._reranker = None
|
||
self._query_string = None
|
||
|
||
def rerank(
|
||
self, reranker: Reranker = RRFReranker(), query_string: Optional[str] = None
|
||
) -> AsyncHybridQuery:
|
||
if reranker and not isinstance(reranker, Reranker):
|
||
raise ValueError("reranker must be an instance of Reranker class.")
|
||
|
||
self._reranker = reranker
|
||
|
||
if not self._query_string and not query_string:
|
||
raise ValueError("query_string must be provided to rerank the results.")
|
||
|
||
self._query_string = query_string
|
||
|
||
return self
|
||
|
||
def nearest_to_text(
|
||
self, query: str | FullTextQuery, columns: Union[str, List[str], None] = None
|
||
) -> AsyncHybridQuery:
|
||
"""
|
||
Find the documents that are most relevant to the given text query,
|
||
in addition to vector search.
|
||
|
||
This converts the vector query into a hybrid query.
|
||
|
||
This search will perform a full text search on the table and return
|
||
the most relevant documents, combined with the vector query results.
|
||
The text relevance is determined by BM25.
|
||
|
||
The columns to search must be with native FTS index
|
||
(Tantivy-based can't work with this method).
|
||
|
||
By default, all indexed columns are searched,
|
||
now only one column can be searched at a time.
|
||
|
||
Parameters
|
||
----------
|
||
query: str
|
||
The text query to search for.
|
||
columns: str or list of str, default None
|
||
The columns to search in. If None, all indexed columns are searched.
|
||
For now only one column can be searched at a time.
|
||
"""
|
||
if isinstance(columns, str):
|
||
columns = [columns]
|
||
if columns is None:
|
||
columns = []
|
||
|
||
if isinstance(query, str):
|
||
return AsyncHybridQuery(
|
||
self._inner.nearest_to_text({"query": query, "columns": columns}),
|
||
self._table,
|
||
)
|
||
# FullTextQuery object
|
||
return AsyncHybridQuery(
|
||
self._inner.nearest_to_text({"query": query}), self._table
|
||
)
|
||
|
||
async def to_batches(
|
||
self,
|
||
*,
|
||
max_batch_length: Optional[int] = None,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> AsyncRecordBatchReader:
|
||
reader = await super().to_batches(timeout=timeout)
|
||
results = pa.Table.from_batches(await reader.read_all(), reader.schema)
|
||
if self._reranker:
|
||
results = self._reranker.rerank_vector(self._query_string, results)
|
||
return AsyncRecordBatchReader(results, max_batch_length=max_batch_length)
|
||
|
||
|
||
class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
|
||
"""
|
||
A query builder that performs hybrid vector and full text search.
|
||
Results are combined and reranked based on the specified reranker.
|
||
By default, the results are reranked using the RRFReranker, which
|
||
uses reciprocal rank fusion score for reranking.
|
||
|
||
To make the vector and fts results comparable, the scores are normalized.
|
||
Instead of normalizing scores, the `normalize` parameter can be set to "rank"
|
||
in the `rerank` method to convert the scores to ranks and then normalize them.
|
||
"""
|
||
|
||
def __init__(self, inner: LanceHybridQuery, table: Optional["AsyncTable"] = None):
|
||
super().__init__(inner, table)
|
||
self._inner = inner
|
||
self._norm = "score"
|
||
self._reranker = RRFReranker()
|
||
|
||
def rerank(
|
||
self, reranker: Reranker = RRFReranker(), normalize: str = "score"
|
||
) -> AsyncHybridQuery:
|
||
"""
|
||
Rerank the hybrid search results using the specified reranker. The reranker
|
||
must be an instance of Reranker class.
|
||
|
||
Parameters
|
||
----------
|
||
reranker: Reranker, default RRFReranker()
|
||
The reranker to use. Must be an instance of Reranker class.
|
||
normalize: str, default "score"
|
||
The method to normalize the scores. Can be "rank" or "score". If "rank",
|
||
the scores are converted to ranks and then normalized. If "score", the
|
||
scores are normalized directly.
|
||
Returns
|
||
-------
|
||
AsyncHybridQuery
|
||
The AsyncHybridQuery object.
|
||
"""
|
||
if normalize not in ["rank", "score"]:
|
||
raise ValueError("normalize must be 'rank' or 'score'.")
|
||
if reranker and not isinstance(reranker, Reranker):
|
||
raise ValueError("reranker must be an instance of Reranker class.")
|
||
|
||
self._norm = normalize
|
||
self._reranker = reranker
|
||
|
||
return self
|
||
|
||
async def to_batches(
|
||
self,
|
||
*,
|
||
max_batch_length: Optional[int] = None,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> AsyncRecordBatchReader:
|
||
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
|
||
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
|
||
|
||
req = fts_query._inner.to_query_request()
|
||
blob_auto_row_id = False
|
||
blob_paths: tuple[str, ...] = ()
|
||
if self._table is not None and supports_blob_auto_row_id(self._table):
|
||
schema = await self._table.schema()
|
||
blob_auto_row_id = blob_auto_row_id_for_scan(
|
||
self._table,
|
||
schema,
|
||
req.select,
|
||
with_row_id=self._with_row_id,
|
||
)
|
||
if blob_auto_row_id:
|
||
blob_paths = tuple(
|
||
blob_v2_projection_sources(schema, req.select).keys()
|
||
)
|
||
self._blob_auto_row_id = blob_auto_row_id
|
||
self._blob_paths = blob_paths
|
||
|
||
fts_query.with_row_id()
|
||
vec_query.with_row_id()
|
||
|
||
fts_results, vector_results = await asyncio.gather(
|
||
fts_query.to_arrow(timeout=timeout),
|
||
vec_query.to_arrow(timeout=timeout),
|
||
)
|
||
|
||
result = LanceHybridQueryBuilder._combine_hybrid_results(
|
||
fts_results=fts_results,
|
||
vector_results=vector_results,
|
||
norm=self._norm,
|
||
fts_query=fts_query.get_query(),
|
||
reranker=self._reranker,
|
||
limit=self._inner.get_limit(),
|
||
with_row_ids=True,
|
||
)
|
||
if (
|
||
not self._user_requested_row_id()
|
||
and not blob_auto_row_id
|
||
and "_rowid" in result.column_names
|
||
):
|
||
result = result.drop(["_rowid"])
|
||
|
||
return AsyncRecordBatchReader(result, max_batch_length=max_batch_length)
|
||
|
||
async def explain_plan(self, verbose: Optional[bool] = False):
|
||
"""Return the execution plan for this query.
|
||
|
||
The output includes both the vector and FTS search plans.
|
||
|
||
Examples
|
||
--------
|
||
>>> import asyncio
|
||
>>> from lancedb import connect_async
|
||
>>> from lancedb.index import FTS
|
||
>>> async def doctest_example():
|
||
... conn = await connect_async("./.lancedb")
|
||
... table = await conn.create_table("my_table", [{"vector": [99.0, 99.0], "text": "hello world"}])
|
||
... await table.create_index("text", config=FTS(with_position=False))
|
||
... plan = await table.query().nearest_to([1.0, 2.0]).nearest_to_text("hello").explain_plan(True)
|
||
... print(plan)
|
||
>>> asyncio.run(doctest_example()) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||
RRFReranker(K=60)
|
||
ProjectionExec: expr=[vector@0 as vector, text@3 as text, _distance@2 as _distance]
|
||
LanceRead: uri=..., projection=[text], source=stream(_rowid)
|
||
GlobalLimitExec: skip=0, fetch=10
|
||
FilterExec: _distance@2 IS NOT NULL
|
||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||
KNNVectorDistance: metric=l2
|
||
LanceRead: uri=..., projection=[vector], ...
|
||
ProjectionExec: expr=[vector@2 as vector, text@3 as text, _score@1 as _score]
|
||
LanceRead: uri=..., projection=[vector, text], source=stream(_rowid)
|
||
GlobalLimitExec: skip=0, fetch=10
|
||
MatchQuery: column=text, query=[hello]
|
||
|
||
Parameters
|
||
----------
|
||
verbose : bool, default False
|
||
Use a verbose output format.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
""" # noqa: E501
|
||
|
||
vector_plan = await self._inner.to_vector_query().explain_plan(verbose)
|
||
fts_plan = await self._inner.to_fts_query().explain_plan(verbose)
|
||
# Indent sub-plans under the reranker
|
||
indented_vector = "\n".join(" " + line for line in vector_plan.splitlines())
|
||
indented_fts = "\n".join(" " + line for line in fts_plan.splitlines())
|
||
return f"{self._reranker}\n {indented_vector}\n {indented_fts}"
|
||
|
||
async def analyze_plan(
|
||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||
) -> str:
|
||
"""
|
||
Execute the query and return the physical execution plan with runtime metrics.
|
||
|
||
This runs both the vector and FTS (full-text search) queries and returns
|
||
detailed metrics for each step of execution—such as rows processed,
|
||
elapsed time, I/O stats, and more. It’s useful for debugging and
|
||
performance analysis.
|
||
|
||
Parameters
|
||
----------
|
||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||
Defaults to "aggregate".
|
||
How distributed worker metrics are displayed for remote query plans.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
"""
|
||
results = ["Vector Search Query:"]
|
||
results.append(
|
||
await self._inner.to_vector_query().analyze_plan(distributed_metrics)
|
||
)
|
||
results.append("FTS Search Query:")
|
||
results.append(
|
||
await self._inner.to_fts_query().analyze_plan(distributed_metrics)
|
||
)
|
||
|
||
return "\n".join(results)
|
||
|
||
|
||
class AsyncTakeQuery(AsyncQueryBase):
|
||
"""
|
||
Builder for parameterizing and executing take queries.
|
||
"""
|
||
|
||
def __init__(self, inner: LanceTakeQuery, table: Optional["AsyncTable"] = None):
|
||
super().__init__(inner, table)
|
||
|
||
def use_lsm(self, enable: bool) -> "AsyncTakeQuery":
|
||
"""Control MemWAL LSM read routing for this take query.
|
||
|
||
``False`` bypasses the MemWAL and reads the base table only — the escape
|
||
hatch, since take-by-row-id/offset is not supported on the LSM scanner.
|
||
"""
|
||
self._inner.use_lsm(enable)
|
||
return self
|
||
|
||
async def _plain_scan_to_pandas(
|
||
self,
|
||
blob_mode: BlobMode,
|
||
flatten: Optional[Union[int, bool]] = None,
|
||
**kwargs,
|
||
) -> Optional["pd.DataFrame"]:
|
||
return None
|
||
|
||
|
||
class BaseQueryBuilder(object):
|
||
"""
|
||
Wraps AsyncQueryBase and provides a synchronous interface
|
||
"""
|
||
|
||
def __init__(self, inner: AsyncQueryBase):
|
||
self._inner = inner
|
||
|
||
def to_query_object(self) -> Query:
|
||
return self._inner.to_query_object()
|
||
|
||
def select(self, columns: Union[List[str], dict[str, str]]) -> Self:
|
||
"""
|
||
Return only the specified columns.
|
||
|
||
By default a query will return all columns from the table. However, this can
|
||
have a very significant impact on latency. LanceDb stores data in a columnar
|
||
fashion. This
|
||
means we can finely tune our I/O to select exactly the columns we need.
|
||
|
||
As a best practice you should always limit queries to the columns that you need.
|
||
If you pass in a list of column names then only those columns will be
|
||
returned.
|
||
|
||
You can also use this method to create new "dynamic" columns based on your
|
||
existing columns. For example, you may not care about "a" or "b" but instead
|
||
simply want "a + b". This is often seen in the SELECT clause of an SQL query
|
||
(e.g. `SELECT a+b FROM my_table`).
|
||
|
||
To create dynamic columns you can pass in a dict[str, str]. A column will be
|
||
returned for each entry in the map. The key provides the name of the column.
|
||
The value is an SQL string used to specify how the column is calculated.
|
||
|
||
For example, an SQL query might state `SELECT a + b AS combined, c`. The
|
||
equivalent input to this method would be `{"combined": "a + b", "c": "c"}`.
|
||
|
||
Columns will always be returned in the order given, even if that order is
|
||
different than the order used when adding the data.
|
||
"""
|
||
self._inner.select(columns)
|
||
return self
|
||
|
||
def with_row_id(self) -> Self:
|
||
"""
|
||
Include the _rowid column in the results.
|
||
"""
|
||
self._inner.with_row_id()
|
||
return self
|
||
|
||
def use_lsm(self, enable: bool) -> Self:
|
||
"""
|
||
Control MemWAL LSM read routing for this query.
|
||
|
||
``False`` bypasses the MemWAL and reads the base table only, the escape
|
||
hatch for shapes the LSM scanner cannot honor (e.g. take-by-row-id).
|
||
"""
|
||
self._inner.use_lsm(enable)
|
||
return self
|
||
|
||
def with_row_address(self, with_row_address: bool = True) -> Self:
|
||
"""
|
||
Include the _rowaddr column in scanner-backed plain query results.
|
||
"""
|
||
self._inner.with_row_address(with_row_address)
|
||
return self
|
||
|
||
def with_fragments(self, fragments: Any) -> Self:
|
||
"""
|
||
Restrict scanner-backed plain query results to the given Lance fragments.
|
||
"""
|
||
self._inner.with_fragments(fragments)
|
||
return self
|
||
|
||
def fragment_ids(self, fragment_ids: List[int]) -> Self:
|
||
"""
|
||
Restrict scanner-backed plain query results to the given Lance fragment ids.
|
||
"""
|
||
self._inner.fragment_ids(fragment_ids)
|
||
return self
|
||
|
||
def output_schema(self) -> pa.Schema:
|
||
"""
|
||
Return the output schema for the query
|
||
|
||
This does not execute the query.
|
||
"""
|
||
return LOOP.run(self._inner.output_schema())
|
||
|
||
def to_batches(
|
||
self,
|
||
*,
|
||
max_batch_length: Optional[int] = None,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> pa.RecordBatchReader:
|
||
"""
|
||
Execute the query and return the results as an Apache Arrow RecordBatchReader.
|
||
|
||
Parameters
|
||
----------
|
||
|
||
max_batch_length: Optional[int]
|
||
The maximum number of selected records in a single RecordBatch object.
|
||
If not specified, a default batch length is used.
|
||
It is possible for batches to be smaller than the provided length if the
|
||
underlying data is stored in smaller chunks.
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
"""
|
||
async_reader = LOOP.run(
|
||
self._inner.to_batches(max_batch_length=max_batch_length, timeout=timeout)
|
||
)
|
||
|
||
def iter_sync():
|
||
try:
|
||
while True:
|
||
yield LOOP.run(async_reader.__anext__())
|
||
except StopAsyncIteration:
|
||
return
|
||
|
||
return pa.RecordBatchReader.from_batches(async_reader.schema, iter_sync())
|
||
|
||
def to_arrow(self, timeout: Optional[timedelta] = None) -> pa.Table:
|
||
"""
|
||
Execute the query and collect the results into an Apache Arrow Table.
|
||
|
||
This method will collect all results into memory before returning. If
|
||
you expect a large number of results, you may want to use
|
||
[to_batches][lancedb.query.AsyncQueryBase.to_batches]
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
"""
|
||
return LOOP.run(self._inner.to_arrow(timeout))
|
||
|
||
def to_list(self, timeout: Optional[timedelta] = None) -> List[dict]:
|
||
"""
|
||
Execute the query and return the results as a list of dictionaries.
|
||
|
||
Each list entry is a dictionary with the selected column names as keys,
|
||
or all table columns if `select` is not called. The vector and the "_distance"
|
||
fields are returned whether or not they're explicitly selected.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
"""
|
||
return LOOP.run(self._inner.to_list(timeout))
|
||
|
||
def to_pandas(
|
||
self,
|
||
flatten: Optional[Union[int, bool]] = None,
|
||
timeout: Optional[timedelta] = None,
|
||
*,
|
||
blob_mode: BlobMode = "lazy",
|
||
**kwargs,
|
||
) -> "pd.DataFrame":
|
||
"""
|
||
Execute the query and collect the results into a pandas DataFrame.
|
||
|
||
This method will collect all results into memory before returning. If you
|
||
expect a large number of results, you may want to use
|
||
[to_batches][lancedb.query.AsyncQueryBase.to_batches] and convert each batch to
|
||
pandas separately.
|
||
|
||
Examples
|
||
--------
|
||
|
||
>>> import asyncio
|
||
>>> from lancedb import connect_async
|
||
>>> async def doctest_example():
|
||
... conn = await connect_async("./.lancedb")
|
||
... table = await conn.create_table("my_table", data=[{"a": 1, "b": 2}])
|
||
... async for batch in await table.query().to_batches():
|
||
... batch_df = batch.to_pandas()
|
||
>>> asyncio.run(doctest_example())
|
||
|
||
Parameters
|
||
----------
|
||
flatten: Optional[Union[int, bool]]
|
||
If flatten is True, flatten all nested columns.
|
||
If flatten is an integer, flatten the nested columns up to the
|
||
specified depth.
|
||
If unspecified, do not flatten the nested columns.
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
blob_mode: str, default "lazy"
|
||
Controls how blob columns are returned for plain scan queries.
|
||
**kwargs
|
||
Forwarded to pyarrow.Table.to_pandas after query execution and
|
||
optional flattening.
|
||
"""
|
||
return LOOP.run(
|
||
self._inner.to_pandas(flatten, timeout, blob_mode=blob_mode, **kwargs)
|
||
)
|
||
|
||
def to_polars(
|
||
self,
|
||
timeout: Optional[timedelta] = None,
|
||
) -> "pl.DataFrame":
|
||
"""
|
||
Execute the query and collect the results into a Polars DataFrame.
|
||
|
||
This method will collect all results into memory before returning. If you
|
||
expect a large number of results, you may want to use
|
||
[to_batches][lancedb.query.AsyncQueryBase.to_batches] and convert each batch to
|
||
polars separately.
|
||
|
||
Parameters
|
||
----------
|
||
timeout: Optional[timedelta]
|
||
The maximum time to wait for the query to complete.
|
||
If not specified, no timeout is applied. If the query does not
|
||
complete within the specified time, an error will be raised.
|
||
|
||
Examples
|
||
--------
|
||
|
||
>>> import asyncio
|
||
>>> import polars as pl
|
||
>>> from lancedb import connect_async
|
||
>>> async def doctest_example():
|
||
... conn = await connect_async("./.lancedb")
|
||
... table = await conn.create_table("my_table", data=[{"a": 1, "b": 2}])
|
||
... async for batch in await table.query().to_batches():
|
||
... batch_df = pl.from_arrow(batch)
|
||
>>> asyncio.run(doctest_example())
|
||
"""
|
||
return LOOP.run(self._inner.to_polars(timeout))
|
||
|
||
def explain_plan(self, verbose: Optional[bool] = False):
|
||
"""Return the execution plan for this query.
|
||
|
||
Examples
|
||
--------
|
||
>>> import asyncio
|
||
>>> from lancedb import connect_async
|
||
>>> async def doctest_example():
|
||
... conn = await connect_async("./.lancedb")
|
||
... table = await conn.create_table("my_table", [{"vector": [99.0, 99.0]}])
|
||
... plan = await table.query().nearest_to([1.0, 2.0]).explain_plan(True)
|
||
... print(plan)
|
||
>>> asyncio.run(doctest_example()) # doctest: +ELLIPSIS, +NORMALIZE_WHITESPACE
|
||
ProjectionExec: expr=[vector@0 as vector, _distance@2 as _distance]
|
||
GlobalLimitExec: skip=0, fetch=10
|
||
FilterExec: _distance@2 IS NOT NULL
|
||
SortExec: TopK(fetch=10), expr=[_distance@2 ASC NULLS LAST, _rowid@1 ASC NULLS LAST], preserve_partitioning=[false]
|
||
KNNVectorDistance: metric=l2
|
||
LanceRead: uri=..., projection=[vector], ...
|
||
<BLANKLINE>
|
||
|
||
Parameters
|
||
----------
|
||
verbose : bool, default False
|
||
Use a verbose output format.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
""" # noqa: E501
|
||
return LOOP.run(self._inner.explain_plan(verbose))
|
||
|
||
def analyze_plan(
|
||
self, distributed_metrics: AnalyzePlanDistributedMetrics = "aggregate"
|
||
) -> str:
|
||
"""Execute the query and display with runtime metrics.
|
||
|
||
Parameters
|
||
----------
|
||
distributed_metrics : Literal["aggregate", "per_worker", "full"]
|
||
Defaults to "aggregate".
|
||
How distributed worker metrics are displayed for remote query plans.
|
||
|
||
Returns
|
||
-------
|
||
plan : str
|
||
"""
|
||
return LOOP.run(self._inner.analyze_plan(distributed_metrics))
|
||
|
||
|
||
class LanceTakeQueryBuilder(BaseQueryBuilder):
|
||
"""
|
||
Builder for parameterizing and executing take queries.
|
||
"""
|
||
|
||
def __init__(self, inner: AsyncTakeQuery):
|
||
super().__init__(inner)
|