# 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 . 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,
    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 query.take_offsets 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 _query_request_projection(req: "PyQueryRequest") -> QueryProjection:
    if req.select_source_columns is not None:
        return req.select_source_columns
    return req.select


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 DocumentGranularity(str, Enum):
    """The unit treated as one full-text-search document."""

    ROW = "row"
    LIST_ELEMENT = "list_element"


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.
    document_granularity : DocumentGranularity, optional
        Explicitly select row or deepest-list-element documents. If omitted,
        the indexed granularity is inferred. When both granularities are indexed
        for the field, this must be specified. With no index, row granularity is used.
    """

    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)
    document_granularity: Optional[DocumentGranularity] = pydantic.Field(
        None, 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.
    slop : int, default 0
        The maximum number of intervening positions permitted in the phrase.
    document_granularity : DocumentGranularity, optional
        Explicitly select row or deepest-list-element documents. If omitted,
        the indexed granularity is inferred. When both granularities are indexed
        for the field, this must be specified. With no index, row granularity is used.
    """

    query: str
    column: str
    slop: int = pydantic.Field(0, kw_only=True)
    document_granularity: Optional[DocumentGranularity] = pydantic.Field(
        None, 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](https://lancedb.com/docs/search/vector-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](https://lancedb.com/docs/indexing/)
          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](https://lancedb.com/docs/indexing/)
          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

    # Dataset offsets whose duplicate occurrences must be restored after lookup.
    # This is populated when a take query is converted to this serializable form.
    take_offsets: Optional[List[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.take_offsets = req.take_offsets
        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)
    model_config = pydantic.ConfigDict(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:
        return blob_auto_row_id_for_scan(
            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](https://lancedb.com/docs/indexing/)
        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](https://lancedb.com/docs/indexing/)
        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,
            offset=self._offset,
        )
        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,
        offset: Optional[int] = None,
    ) -> 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(offset=offset or 0, 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:
            # The final offset/limit window is sliced out of the combined,
            # reranked results, so each sub-query must fetch enough rows to
            # cover the skipped prefix as well as the window itself.
            sub_query_limit = self._limit + (self._offset or 0)
            self._vector_query.limit(sub_query_limit)
            self._fts_query.limit(sub_query_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 is not None:
            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:
            self._blob_auto_row_id = False
            self._blob_paths = ()
            return

        req = self._inner.to_query_request()
        schema = await self._table.schema()
        projection = _query_request_projection(req)
        self._blob_auto_row_id = blob_auto_row_id_for_scan(
            schema,
            projection,
            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, projection).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(
            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() 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][lancedb.query.AsyncVectorQuery.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. Flat
        searches share one table scan across the query vectors, avoiding the scan
        and memory amplification of making multiple queries concurrently. 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][lancedb.query.AsyncHybridQuery.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. Flat searches share one table scan across
        the query vectors instead of issuing concurrent full scans.
        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:
            schema = await self._table.schema()
            projection = _query_request_projection(req)
            blob_auto_row_id = blob_auto_row_id_for_scan(
                schema,
                projection,
                with_row_id=self._with_row_id,
            )
            if blob_auto_row_id:
                blob_paths = tuple(
                    blob_v2_projection_sources(schema, projection).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)
