rust/lancedb
1
0
Fork 0
mirror of https://github.com/lancedb/lancedb.git synced 2026-11-02 00:36:53 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(query): isolate the hybrid primary-key fallback for removal

Browse Source 
The primary-key fusion for MemWAL tables is a stopgap until MemWAL supports
`_rowid`, but it was woven through the hybrid execution paths in both
languages. Gather it into one module per language so it can be removed in
one pass:

- Rust: `query/wal_fusion.rs` holds the refusal detection, key resolution,
  projection injection, surrogate stamping, `pk_columns`, the remote table's
  refusal memory, and their tests. `run_hybrid` takes an
  `Option<PkFusion>` and makes three one-line calls into it.
- Python: `lancedb/_wal_hybrid.py` holds the same pieces behind a
  `PkFusion` class and a retry wrapper; `_combine_hybrid_results` is back to
  its original signature. Tests move to `tests/test_wal_hybrid.py`, with
  new unit tests for the retry wrapper.

Every remaining touchpoint is tagged `WAL-PK-FUSION` with what it reverts
to. No behavior change.

Co-Authored-By: Claude Opus 5.5 (1M context) <noreply@anthropic.com>
This commit is contained in:
Daniel Rammer
2026-09-24 09:28:21 -05:00
co-authored by Claude Opus 5.5
parent 83cfcb84d2
commit 7480a5e428
14 changed files with 1136 additions and 909 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+1
View File
@@ -509,6 +509,7 @@ class Table:
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
async def unset_lsm_write_spec(self) -> None: ...
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
# WAL-PK-FUSION: delete both.
def hybrid_pk_fusion_learned(self) -> bool: ...
def note_hybrid_pk_fusion(self) -> None: ...
async def checkpoint_lsm(self) -> None: ...
python/python/lancedb/_wal_hybrid.py
+282
View File
@@ -0,0 +1,282 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""WAL-PK-FUSION: hybrid search's primary-key fallback for MemWAL tables.
Hybrid search joins its vector and full-text legs on `_rowid`, which a MemWAL
table cannot supply: the fresh tier has no stable row id, so the server rejects
`with_row_id` before it plans. When that happens the legs are joined on a
surrogate `_rowid` built from the table's primary key instead.
This is a stopgap until MemWAL supports `_rowid`, and everything it needs lives
here so it can be removed in one pass:
1. Delete this file and `tests/test_wal_hybrid.py`.
2. Run `grep -rn WAL-PK-FUSION python/ rust/` and follow the note at each
marker; the Rust side has its own copy in
`rust/lancedb/src/query/wal_fusion.rs`.
"""
from __future__ import annotations
from typing import TYPE_CHECKING, Awaitable, Callable, List, Optional, Tuple, TypeVar
import pyarrow as pa
from ._blob import blob_v2_projection_sources
if TYPE_CHECKING:
from .types import QueryProjection
T = TypeVar("T")
# Lance records the unenforced primary key as field metadata, which survives the
# `describe` round trip — so a remote table can resolve the key without opening
# the dataset. Keys are bytes on a pyarrow field.
_PK_MARKER = b"lance-schema:unenforced-primary-key"
_PK_POSITION_MARKER = b"lance-schema:unenforced-primary-key:position"
_TRUTHY = {"1", "true", "on", "yes", "y"}
def _pk_columns(schema: pa.Schema) -> List[str]:
"""The unenforced primary key columns of `schema`, in key order.
Mirrors Lance's ordering: fields with an explicit 1-based position first, in
position order, then fields carrying only the legacy boolean marker, in
declaration order.
"""
keyed = []
for index, field in enumerate(schema):
metadata = field.metadata or {}
position = metadata.get(_PK_POSITION_MARKER)
if position is not None:
try:
position = int(position)
except ValueError:
position = None
if position is None:
legacy = metadata.get(_PK_MARKER)
if legacy is None or legacy.decode().lower() not in _TRUTHY:
continue
position = 0
keyed.append((position == 0, position, index, field.name))
return [name for *_, name in sorted(keyed)]
# Both layers raise this same sentence verbatim -- sophon's
# `LSM_WITH_ROW_ID_UNSUPPORTED` says so explicitly, and the OSS LSM scanner
# builds it from the same words -- so this substring survives either path and
# the wrapping each adds. It is a coupling to a message rather than a code,
# because the refusal arrives as a generic 400; if the server ever rewords it
# the fallback stops firing and hybrid on a MemWAL table errors visibly again,
# which is loud rather than silently wrong.
_WAL_ROW_ID_REFUSAL = "does not support with_row_id"
WAL_ROW_ID_UNSUPPORTED = (
"hybrid search on a MemWAL table cannot return _rowid: the fresh tier has "
"no stable row id, and the ids the fusion joins on are synthesized from "
"the primary key. Set use_lsm(False) to read the base table only (results "
"will exclude un-compacted MemWAL data)"
)
def _is_wal_row_id_refusal(error: BaseException) -> bool:
"""Whether `error` is the server declining `_rowid` because the table is
MemWAL-backed, as opposed to any other failure the query might hit."""
return _WAL_ROW_ID_REFUSAL in str(error)
def _with_pk_columns(columns, pk_columns: List[str]):
"""Add the key columns a projection is missing.
Returns the projection to use and the columns added, which the caller drops
from the output again so an injected key never reaches a caller who did not
select it.
"""
if not columns:
# Every non-system column, the key among them.
return columns, []
# `PyQueryRequest.select` reports a dynamic projection as (name, sql) pairs.
if isinstance(columns, (list, tuple)) and isinstance(columns[0], (list, tuple)):
columns = dict(columns)
missing = [c for c in pk_columns if c not in columns]
if isinstance(columns, dict):
return {**columns, **{c: c for c in missing}}, missing
return list(columns) + missing, missing
def _stamp_surrogate_row_ids(
tables: List[pa.Table], pk_columns: List[str]
) -> List[pa.Table]:
"""Give each table a `_rowid` column derived from its primary key.
The fusion only ever compares that column for equality — to deduplicate, to
group by in a reranker, and to restore the original scores — so on a MemWAL
table, where the server has no stable row id to hand out, any injective
mapping of the primary key does the same job. Ids are assigned in first-seen
order across `tables`, so passing them in fusion order (vector, then FTS)
preserves "first occurrence wins" and its tie break.
Writing it under the `_rowid` name is what lets every reranker, third-party
implementations included, run unmodified. The column never reaches the
caller.
"""
if not pk_columns:
raise ValueError(
"hybrid search on a MemWAL table needs an unenforced primary key to "
"deduplicate on, and this table declares none"
)
ids = {}
stamped = []
for table in tables:
if table.num_rows == 0:
# Both legs empty: the combined schema carries no key column, and no
# row that would need one.
row_ids = pa.array([], type=pa.uint64())
else:
missing = [c for c in pk_columns if c not in table.column_names]
if missing:
raise ValueError(
"hybrid search could not deduplicate: primary key "
f"column(s) {missing} missing from a result set with "
f"{table.num_rows} rows"
)
keys = zip(*(table.column(c).to_pylist() for c in pk_columns))
row_ids = pa.array(
[ids.setdefault(key, len(ids)) for key in keys], type=pa.uint64()
)
stamped.append(table.append_column("_rowid", row_ids))
return stamped
def pk_fusion_learned(table, use_lsm: Optional[bool]) -> bool:
"""Whether a previous hybrid query on `table` was refused `_rowid`, so this
one should go straight to the primary key.
Learned from a refusal, never probed, so asking is free.
"""
if table is None or use_lsm is False:
return False
# Only an unambiguous True forks to the key: a table type that answers the
# question loosely keeps the `_rowid` path.
return table._hybrid_pk_fusion_learned() is True
class PkFusion:
"""How the legs of a hybrid query on a MemWAL table are joined: on a
surrogate `_rowid` built from the primary key."""
def __init__(self, pk_columns: List[str]):
self.pk_columns = pk_columns
# Key columns the caller did not select, added to the projection so the
# surrogate can be built and dropped again before returning.
self.injected: List[str] = []
@classmethod
def for_query(
cls,
schema: pa.Schema,
projection: QueryProjection,
*,
caller_requested_row_id: bool,
) -> PkFusion:
"""Resolve the key for a hybrid query, refusing the shapes that need a
real `_rowid`."""
if caller_requested_row_id:
raise NotImplementedError(WAL_ROW_ID_UNSUPPORTED)
# A projected blob needs a real row id to fetch through, whatever else
# has turned `_with_row_id` on — including `return_score="all"`.
if blob_v2_projection_sources(schema, projection):
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"fetching blobs needs a real _rowid, and the fresh tier has no "
"stable one. Set use_lsm(False) to read the base table only "
"(results will exclude un-compacted MemWAL data)"
)
return cls(_pk_columns(schema))
def inject(self, columns):
"""`columns` with the key columns it is missing, to project in place of
the `_rowid` a MemWAL table would refuse."""
columns, self.injected = _with_pk_columns(columns, self.pk_columns)
return columns
def stamp(
self, vector_results: pa.Table, fts_results: pa.Table
) -> Tuple[pa.Table, pa.Table]:
"""Give both legs a surrogate `_rowid` built from the key."""
# Vector first: `merge_results` concatenates in that order and keeps
# the first occurrence, so first-seen ids preserve its tie break.
vector_results, fts_results = _stamp_surrogate_row_ids(
[vector_results, fts_results], self.pk_columns
)
return vector_results, fts_results
def strip(self, results: pa.Table) -> pa.Table:
"""Drop the surrogate and any injected key columns from the reranked
results.
The surrogate is an internal join key, never an answer, and a caller
who asked for `_rowid` was refused in `for_query`.
"""
drop = [c for c in ["_rowid", *self.injected] if c in results.column_names]
return results.drop(drop) if drop else results
def _check_refusal(error: Exception, *, fused: bool, caller_requested_row_id: bool):
"""Re-raise `error` unless it is a MemWAL `_rowid` refusal worth retrying on
the primary key."""
if fused or not _is_wal_row_id_refusal(error):
raise error
# `return_score="all"` turns row ids on for the reranker, not the caller, so
# it still falls back; only a caller who asked for them is refused, and with
# the reason rather than a bare 400.
if caller_requested_row_id:
raise NotImplementedError(WAL_ROW_ID_UNSUPPORTED) from error
def with_pk_fallback(
run: Callable[[], T],
table,
*,
fused: Callable[[], bool],
caller_requested_row_id: bool,
) -> T:
"""Run a hybrid query, retrying on the primary key if the table refuses
`_rowid`.
`_rowid` is asked for first rather than paying a round trip up front to
find out: only a MemWAL table refuses, and the refusal has to be handled
anyway -- a write spec installed elsewhere can arrive between any two
queries. It costs nothing server-side: it is raised before the query is
planned. `fused` reports whether the attempt already used the key.
"""
try:
return run()
except Exception as e:
_check_refusal(
e, fused=fused(), caller_requested_row_id=caller_requested_row_id
)
# Learned, so this retry and later queries go straight to the key.
table._note_hybrid_pk_fusion()
return run()
async def with_pk_fallback_async(
run: Callable[[], Awaitable[T]],
table,
*,
fused: Callable[[], bool],
caller_requested_row_id: bool,
) -> T:
"""The async form of `with_pk_fallback`."""
try:
return await run()
except Exception as e:
_check_refusal(
e, fused=fused(), caller_requested_row_id=caller_requested_row_id
)
table._note_hybrid_pk_fusion()
return await run()
python/python/lancedb/query.py
+65 -269
View File
@@ -41,6 +41,7 @@ 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 . import _wal_hybrid # WAL-PK-FUSION: delete.
from ._blob import (
BLOB_MODE_TO_HANDLING,
FetchBlobsAsync,
@@ -2175,126 +2176,6 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder):
raise NotImplementedError("Reranking is not yet supported.")
# Lance records the unenforced primary key as field metadata, which survives the
# `describe` round trip — so a remote table can resolve the key without opening
# the dataset. Keys are bytes on a pyarrow field.
_PK_MARKER = b"lance-schema:unenforced-primary-key"
_PK_POSITION_MARKER = b"lance-schema:unenforced-primary-key:position"
_TRUTHY = {"1", "true", "on", "yes", "y"}
def _pk_columns(schema: pa.Schema) -> List[str]:
"""The unenforced primary key columns of `schema`, in key order.
Mirrors Lance's ordering: fields with an explicit 1-based position first, in
position order, then fields carrying only the legacy boolean marker, in
declaration order.
"""
keyed = []
for index, field in enumerate(schema):
metadata = field.metadata or {}
position = metadata.get(_PK_POSITION_MARKER)
if position is not None:
try:
position = int(position)
except ValueError:
position = None
if position is None:
legacy = metadata.get(_PK_MARKER)
if legacy is None or legacy.decode().lower() not in _TRUTHY:
continue
position = 0
keyed.append((position == 0, position, index, field.name))
return [name for *_, name in sorted(keyed)]
# Both layers raise this same sentence verbatim -- sophon's
# `LSM_WITH_ROW_ID_UNSUPPORTED` says so explicitly, and the OSS LSM scanner
# builds it from the same words -- so this substring survives either path and
# the wrapping each adds. It is a coupling to a message rather than a code,
# because the refusal arrives as a generic 400; if the server ever rewords it
# the fallback stops firing and hybrid on a MemWAL table errors visibly again,
# which is loud rather than silently wrong.
_WAL_ROW_ID_REFUSAL = "does not support with_row_id"
_WAL_ROW_ID_UNSUPPORTED = (
"hybrid search on a MemWAL table cannot return _rowid: the fresh tier has "
"no stable row id, and the ids the fusion joins on are synthesized from "
"the primary key. Set use_lsm(False) to read the base table only (results "
"will exclude un-compacted MemWAL data)"
)
def _is_wal_row_id_refusal(error: BaseException) -> bool:
"""Whether `error` is the server declining `_rowid` because the table is
MemWAL-backed, as opposed to any other failure the query might hit."""
return _WAL_ROW_ID_REFUSAL in str(error)
def _with_pk_columns(columns, pk_columns: List[str]):
"""Add the key columns a projection is missing.
Returns the projection to use and the columns added, which the caller drops
from the output again so an injected key never reaches a caller who did not
select it.
"""
if not columns:
# Every non-system column, the key among them.
return columns, []
# `PyQueryRequest.select` reports a dynamic projection as (name, sql) pairs.
if isinstance(columns, (list, tuple)) and isinstance(columns[0], (list, tuple)):
columns = dict(columns)
missing = [c for c in pk_columns if c not in columns]
if isinstance(columns, dict):
return {**columns, **{c: c for c in missing}}, missing
return list(columns) + missing, missing
def _stamp_surrogate_row_ids(
tables: List[pa.Table], pk_columns: List[str]
) -> List[pa.Table]:
"""Give each table a `_rowid` column derived from its primary key.
The fusion only ever compares that column for equality — to deduplicate, to
group by in a reranker, and to restore the original scores — so on a MemWAL
table, where the server has no stable row id to hand out, any injective
mapping of the primary key does the same job. Ids are assigned in first-seen
order across `tables`, so passing them in fusion order (vector, then FTS)
preserves "first occurrence wins" and its tie break.
Writing it under the `_rowid` name is what lets every reranker, third-party
implementations included, run unmodified. The column never reaches the
caller.
"""
if not pk_columns:
raise ValueError(
"hybrid search on a MemWAL table needs an unenforced primary key to "
"deduplicate on, and this table declares none"
)
ids = {}
stamped = []
for table in tables:
if table.num_rows == 0:
# Both legs empty: the combined schema carries no key column, and no
# row that would need one.
row_ids = pa.array([], type=pa.uint64())
else:
missing = [c for c in pk_columns if c not in table.column_names]
if missing:
raise ValueError(
"hybrid search could not deduplicate: primary key "
f"column(s) {missing} missing from a result set with "
f"{table.num_rows} rows"
)
keys = zip(*(table.column(c).to_pylist() for c in pk_columns))
row_ids = pa.array(
[ids.setdefault(key, len(ids)) for key in keys], type=pa.uint64()
)
stamped.append(table.append_column("_rowid", row_ids))
return stamped
class LanceHybridQueryBuilder(LanceQueryBuilder):
"""
A query builder that performs hybrid vector and full text search.
@@ -2327,11 +2208,9 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._phrase_query = None
self._lower_bound = None
self._upper_bound = None
# Set by `_create_query_builders`: the key the legs are fused on when
# `_rowid` is unavailable, and the key columns added to the projection
# to make that possible.
self._fusion_pk = None
self._injected_pk = []
# WAL-PK-FUSION: delete both, and the `with_row_id` override below.
# Set by `_create_query_builders` when the legs are fused on the key.
self._pk_fusion: Optional[_wal_hybrid.PkFusion] = None
# `_with_row_id` is also turned on by `rerank(return_score="all")` for
# the reranker's own use, so it cannot answer "did the caller ask?".
self._caller_requested_row_id = False
@@ -2356,7 +2235,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
return vector_query, text_query
def with_row_id(self, with_row_id: bool) -> Self:
def with_row_id(self, with_row_id: bool) -> Self: # WAL-PK-FUSION: delete.
"""Set whether to return row ids.
Recorded separately from `_with_row_id`, which `rerank(return_score=
@@ -2386,31 +2265,19 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
raise NotImplementedError("to_query_object not yet supported on a hybrid query")
def to_arrow(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
"""Ask for `_rowid` first, falling back to the primary key only when the
server refuses, rather than paying a round trip up front to find out.
Only a MemWAL table refuses, and the refusal has to be handled anyway
-- a write spec installed elsewhere can arrive between any two queries.
It costs nothing server-side: it is raised before the query is planned.
"""
try:
return self._run_hybrid(timeout=timeout)
except Exception as e:
if self._fusion_pk is not None or not _is_wal_row_id_refusal(e):
raise
# `return_score="all"` turns row ids on for the reranker, not the
# caller, so it still falls back; only a caller who asked for them
# is refused, and with the reason rather than a bare 400.
if self._caller_requested_row_id:
raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED) from e
# Learned, so later queries on this table skip straight to it.
self._table._note_hybrid_pk_fusion()
return self._run_hybrid(timeout=timeout)
# WAL-PK-FUSION: without the fallback, this body is `_run_hybrid`'s.
return _wal_hybrid.with_pk_fallback(
lambda: self._run_hybrid(timeout=timeout),
self._table,
fused=lambda: self._pk_fusion is not None,
caller_requested_row_id=self._caller_requested_row_id,
)
def _run_hybrid(self, *, timeout: Optional[timedelta] = None) -> pa.Table:
self._create_query_builders()
fts_query, vector_query = self._fts_query, self._vector_query
if self._fusion_pk is None:
# WAL-PK-FUSION: without the fallback, both legs always ask for row ids.
if self._pk_fusion is None:
fts_query = fts_query.with_row_id(True)
vector_query = vector_query.with_row_id(True)
with ThreadPoolExecutor() as executor:
@@ -2418,6 +2285,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
vector_future = executor.submit(vector_query.to_arrow, timeout=timeout)
fts_results = fts_future.result()
vector_results = vector_future.result()
if self._pk_fusion is not None: # WAL-PK-FUSION: delete.
vector_results, fts_results = self._pk_fusion.stamp(
vector_results, fts_results
)
results = self._combine_hybrid_results(
fts_results=fts_results,
@@ -2428,50 +2299,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
limit=self._limit,
with_row_ids=True,
offset=self._offset,
pk_columns=self._fusion_pk,
)
return self._finish_hybrid_results(results)
def _resolve_fusion_pk(self) -> Optional[List[str]]:
"""The key columns the two legs must be fused on, or None for `_rowid`.
`_rowid` until a MemWAL table refuses it, after which the legs are
joined on a surrogate built from the primary key. The refusal is what
teaches us, rather than a round trip up front that every table would
pay to learn something almost none of them need.
"""
# Only an unambiguous True forks to the key: a table type that answers
# the question loosely keeps the `_rowid` path.
if (
self._use_lsm is False
or self._table._hybrid_pk_fusion_learned() is not True
):
return None
if self._caller_requested_row_id:
raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED)
# Not `_blob_auto_row_id_enabled`: that reports whether a row id would
# be added *automatically*, so it goes quiet once anything else has
# turned `_with_row_id` on — including `return_score="all"`. A projected
# blob needs a real row id to fetch through either way.
if blob_v2_projection_sources(self._table.schema, self._columns):
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"fetching blobs needs a real _rowid, and the fresh tier has no "
"stable one. Set use_lsm(False) to read the base table only "
"(results will exclude un-compacted MemWAL data)"
)
return _pk_columns(self._table.schema)
def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
if self._fusion_pk is not None:
# The surrogate is an internal join key, never an answer, and a
# caller who asked for `_rowid` was refused in `_resolve_fusion_pk`.
if "_rowid" in results.column_names:
results = results.drop(["_rowid"])
injected = [c for c in self._injected_pk if c in results.column_names]
if injected:
results = results.drop(injected)
return results
if self._pk_fusion is not None: # WAL-PK-FUSION: delete.
return self._pk_fusion.strip(results)
if self._user_requested_row_id():
return results
if self._blob_auto_row_id_enabled():
@@ -2490,14 +2323,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
limit: int,
with_row_ids: bool,
offset: Optional[int] = None,
pk_columns: Optional[List[str]] = None,
) -> pa.Table:
if pk_columns is not None:
# Vector first: `merge_results` concatenates in that order and keeps
# the first occurrence, so first-seen ids preserve its tie break.
vector_results, fts_results = _stamp_surrogate_row_ids(
[vector_results, fts_results], pk_columns
)
if norm == "rank":
vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance")
fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score")
@@ -2652,10 +2478,9 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._norm = normalize
self._reranker = reranker
if reranker.score == "all":
# The reranker needs the join column to carry the per-leg scores
# across; the caller did not ask for row ids and must not be
# refused on a MemWAL table because of it.
# Deliberately not through `with_row_id`, which records caller intent.
# WAL-PK-FUSION: without the fallback, `self.with_row_id(True)`.
# Not through `with_row_id`, which records caller intent: the
# reranker's need for row ids must not refuse a MemWAL table.
self._with_row_id = True
return self
@@ -2930,17 +2755,24 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
sub_query_limit = self._limit + (self._offset or 0)
self._vector_query.limit(sub_query_limit)
self._fts_query.limit(sub_query_limit)
self._fusion_pk = self._resolve_fusion_pk()
# WAL-PK-FUSION: without the fallback, select `self._columns` as is.
self._pk_fusion = None
columns = self._columns
if self._fusion_pk is not None:
columns, self._injected_pk = _with_pk_columns(columns, self._fusion_pk)
if _wal_hybrid.pk_fusion_learned(self._table, self._use_lsm):
self._pk_fusion = _wal_hybrid.PkFusion.for_query(
self._table.schema,
self._columns,
caller_requested_row_id=self._caller_requested_row_id,
)
columns = self._pk_fusion.inject(columns)
if columns:
self._vector_query.select(columns)
self._fts_query.select(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 and self._fusion_pk is None:
# WAL-PK-FUSION: without the fallback, `if self._with_row_id:`.
if self._with_row_id and self._pk_fusion is None:
self._vector_query.with_row_id(True)
self._fts_query.with_row_id(True)
if self._use_lsm is not None:
@@ -4081,11 +3913,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
self._inner = inner
self._norm = "score"
self._reranker = RRFReranker()
# Set by `_create_child_queries`: the key the legs are fused on when
# `_rowid` is unavailable, and the key columns added to the projection
# to make that possible.
self._fusion_pk = None
self._injected_pk = []
# WAL-PK-FUSION: delete.
# Set by `_create_child_queries` when the legs are fused on the key.
self._pk_fusion: Optional[_wal_hybrid.PkFusion] = None
def rerank(
self, reranker: Reranker = RRFReranker(), normalize: str = "score"
@@ -4117,37 +3947,6 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
return self
async def _resolve_fusion_pk(self) -> Optional[List[str]]:
"""The key columns the two legs must be fused on, or None for `_rowid`.
`_rowid` until a MemWAL table refuses it, after which the legs are
joined on a surrogate built from the primary key. The refusal is what
teaches us, rather than a round trip up front that every table would
pay to learn something almost none of them need.
"""
if self._table is None:
return None
if self._inner.to_query_request().use_lsm is False:
return None
# Only an unambiguous True forks to the key; see the sync resolver.
if self._table._hybrid_pk_fusion_learned() is not True:
return None
if self._user_requested_row_id():
raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED)
schema = await self._table.schema()
if blob_auto_row_id_for_scan(
schema,
_query_request_projection(self._inner.to_query_request()),
with_row_id=self._with_row_id,
):
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"fetching blobs needs a real _rowid, and the fresh tier has no "
"stable one. Set use_lsm(False) to read the base table only "
"(results will exclude un-compacted MemWAL data)"
)
return _pk_columns(schema)
async def _create_child_queries(
self,
) -> Tuple["AsyncFTSQuery", "AsyncVectorQuery", int, int]:
@@ -4157,8 +3956,8 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
the plans that are reported are the plans that actually run.
Returns the two sub-queries along with the effective limit and offset of
the hybrid query itself. The key the legs are fused on is left on
`self` — see `_resolve_fusion_pk`.
the hybrid query itself. The key the legs are fused on, if any, is left
on `self._pk_fusion`.
"""
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
@@ -4176,8 +3975,16 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
limit = DEFAULT_HYBRID_LIMIT
offset = fts_req.offset or vec_req.offset or 0
self._fusion_pk = await self._resolve_fusion_pk()
if self._fusion_pk is None:
# WAL-PK-FUSION: without the fallback, keep only the two `with_row_id`
# calls.
self._pk_fusion = None
if _wal_hybrid.pk_fusion_learned(self._table, fts_req.use_lsm):
self._pk_fusion = _wal_hybrid.PkFusion.for_query(
await self._table.schema(),
_query_request_projection(fts_req),
caller_requested_row_id=self._user_requested_row_id(),
)
if self._pk_fusion is None:
fts_query.with_row_id()
vec_query.with_row_id()
else:
@@ -4185,10 +3992,8 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
# `select` carries the whole projection; `select_source_columns`
# keeps only plain column references, so rebuilding from it would
# drop computed ones.
columns, self._injected_pk = _with_pk_columns(
fts_req.select, self._fusion_pk
)
if self._injected_pk:
columns = self._pk_fusion.inject(fts_req.select)
if self._pk_fusion.injected:
fts_query.select(columns)
vec_query.select(columns)
@@ -4209,21 +4014,15 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
max_batch_length: Optional[int] = None,
timeout: Optional[timedelta] = None,
) -> AsyncRecordBatchReader:
"""See `LanceHybridQueryBuilder.to_arrow` for why the key is learned
from a refusal rather than probed for."""
try:
return await self._run_hybrid(
# WAL-PK-FUSION: without the fallback, this body is `_run_hybrid`'s.
return await _wal_hybrid.with_pk_fallback_async(
lambda: self._run_hybrid(
max_batch_length=max_batch_length, timeout=timeout
)
except Exception as e:
if self._fusion_pk is not None or not _is_wal_row_id_refusal(e):
raise
if self._user_requested_row_id():
raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED) from e
self._table._note_hybrid_pk_fusion()
return await self._run_hybrid(
max_batch_length=max_batch_length, timeout=timeout
)
),
self._table,
fused=lambda: self._pk_fusion is not None,
caller_requested_row_id=self._user_requested_row_id(),
)
async def _run_hybrid(
self,
@@ -4255,6 +4054,10 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
fts_query.to_arrow(timeout=timeout),
vec_query.to_arrow(timeout=timeout),
)
if self._pk_fusion is not None: # WAL-PK-FUSION: delete.
vector_results, fts_results = self._pk_fusion.stamp(
vector_results, fts_results
)
result = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
@@ -4265,16 +4068,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
limit=limit,
with_row_ids=True,
offset=offset,
pk_columns=self._fusion_pk,
)
if self._fusion_pk is not None:
# The surrogate is an internal join key, never an answer, and a
# caller who asked for `_rowid` was refused in `_resolve_fusion_pk`.
if "_rowid" in result.column_names:
result = result.drop(["_rowid"])
injected = [c for c in self._injected_pk if c in result.column_names]
if injected:
result = result.drop(injected)
if self._pk_fusion is not None: # WAL-PK-FUSION: delete this branch.
result = self._pk_fusion.strip(result)
elif (
not self._user_requested_row_id()
and not blob_auto_row_id
python/python/lancedb/remote/table.py
+1
View File
@@ -1036,6 +1036,7 @@ class RemoteTable(Table):
"""Read the installed LsmWriteSpec, or ``None``."""
return LOOP.run(self._table.get_lsm_write_spec())
# WAL-PK-FUSION: delete both hooks.
def _hybrid_pk_fusion_learned(self) -> bool:
# `self._table` is the async table, which owns the Rust handle.
return self._table._hybrid_pk_fusion_learned()
python/python/lancedb/table.py
+2
View File
@@ -2646,6 +2646,7 @@ class Table(ABC):
to check if the table is already using the new path style.
"""
# WAL-PK-FUSION: delete both hooks, here and on AsyncTable and RemoteTable.
def _hybrid_pk_fusion_learned(self) -> bool:
"""Whether a hybrid query here has already been refused ``_rowid``.
@@ -5304,6 +5305,7 @@ class AsyncTable:
"""
return await self._inner.get_lsm_write_spec()
# WAL-PK-FUSION: delete both hooks.
def _hybrid_pk_fusion_learned(self) -> bool:
"""See [`Table._hybrid_pk_fusion_learned`][lancedb.table.Table]."""
return self._inner.hybrid_pk_fusion_learned()
python/python/tests/test_hybrid_query.py
-156
View File
@@ -495,159 +495,3 @@ def test_normalize_scores():
assert pc.equal(result, expected), (
f"Expected {expected} but got {result} for invert={invert}"
)
def _wal_table(schema: pa.Schema) -> mock.Mock:
"""A MemWAL-backed table a previous query has already been refused
`_rowid` on, so the fusion goes straight to the primary key."""
table = mock.Mock()
table._hybrid_pk_fusion_learned.return_value = True
table.schema = schema
return table
_PK_SCHEMA = pa.schema(
[
pa.field(
"id",
pa.string(),
metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
),
pa.field("text", pa.string()),
]
)
def _wal_hybrid(**select):
builder = (
LanceHybridQueryBuilder(_wal_table(_PK_SCHEMA)).vector([0.1, 0.2]).text("puppy")
)
if "columns" in select:
builder = builder.select(select["columns"])
return builder
def test_wal_hybrid_projects_the_key_instead_of_asking_for_row_ids():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
assert builder._fusion_pk == ["id"]
assert builder._injected_pk == ["id"]
for leg in (builder._vector_query, builder._fts_query):
assert leg._columns == ["text", "id"]
# A MemWAL table rejects with_row_id before it plans.
assert leg._with_row_id is not True
def test_wal_hybrid_selecting_the_key_injects_nothing():
builder = _wal_hybrid(columns=["text", "id"])
builder._create_query_builders()
assert builder._injected_pk == []
def test_wal_hybrid_without_a_projection_leaves_it_alone():
builder = _wal_hybrid()
builder._create_query_builders()
# Select-all already returns the key.
assert builder._injected_pk == []
assert builder._vector_query._columns is None
def test_base_table_hybrid_still_joins_on_row_ids():
table = mock.Mock()
table._hybrid_pk_fusion_learned.return_value = False
builder = LanceHybridQueryBuilder(table).vector([0.1, 0.2]).text("puppy")
builder._create_query_builders()
# `to_arrow` is what turns row ids on for this path; all that matters here
# is that no key was resolved and no projection was touched.
assert builder._fusion_pk is None
assert builder._injected_pk == []
def test_wal_hybrid_with_use_lsm_false_joins_on_row_ids():
builder = _wal_hybrid().use_lsm(False)
builder._create_query_builders()
assert builder._fusion_pk is None
def test_wal_hybrid_refuses_an_explicit_with_row_id():
builder = _wal_hybrid().with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
builder._create_query_builders()
def test_wal_hybrid_return_score_all_is_not_refused():
"""`return_score="all"` turns row ids on for the reranker, not the caller."""
builder = _wal_hybrid().rerank(RRFReranker(return_score="all"))
builder._create_query_builders()
assert builder._fusion_pk == ["id"]
def test_wal_hybrid_fuses_on_the_key_and_drops_it_again():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
vector_results = pa.table(
{"text": ["a", "b"], "id": ["a", "b"], "_distance": [0.1, 0.2]}
)
fts_results = pa.table({"text": ["b", "c"], "id": ["b", "c"], "_score": [1.0, 2.0]})
combined = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
vector_results=vector_results,
norm="score",
fts_query="puppy",
reranker=RRFReranker(),
limit=10,
with_row_ids=True,
pk_columns=builder._fusion_pk,
)
results = builder._finish_hybrid_results(combined)
assert "_rowid" not in results.column_names, "the surrogate is internal"
assert "id" not in results.column_names, "the caller selected only `text`"
assert sorted(results.column("text").to_pylist()) == ["a", "b", "c"]
@pytest.mark.parametrize(
"row_id_first", [True, False], ids=["row_id_first", "rerank_first"]
)
def test_wal_hybrid_refuses_row_id_even_with_return_score_all(row_id_first: bool):
"""`return_score="all"` sets the same flag the caller's request does, so
asking for both must still be refused — in either call order."""
builder = _wal_hybrid()
reranker = RRFReranker(return_score="all")
if row_id_first:
builder = builder.with_row_id(True).rerank(reranker)
else:
builder = builder.rerank(reranker).with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
builder._create_query_builders()
def test_wal_hybrid_refuses_a_blob_projection_with_return_score_all():
"""A projected blob needs a real row id to fetch through, which the
reranker turning `_with_row_id` on does not provide."""
schema = pa.schema(
[
pa.field(
"id",
pa.string(),
metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
),
pa.field("text", pa.string()),
lancedb.blob("blob"),
]
)
builder = (
LanceHybridQueryBuilder(_wal_table(schema))
.vector([0.1, 0.2])
.text("puppy")
.select(["blob"])
.rerank(RRFReranker(return_score="all"))
)
with pytest.raises(NotImplementedError, match="blob"):
builder._create_query_builders()
python/python/tests/test_wal_hybrid.py
+219
View File
@@ -0,0 +1,219 @@
# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright The LanceDB Authors
"""WAL-PK-FUSION: delete this file with `lancedb/_wal_hybrid.py`."""
from unittest import mock
import lancedb
import pyarrow as pa
import pytest
from lancedb._wal_hybrid import WAL_ROW_ID_UNSUPPORTED, with_pk_fallback
from lancedb.query import LanceHybridQueryBuilder
from lancedb.rerankers import RRFReranker
def _wal_table(schema: pa.Schema) -> mock.Mock:
"""A MemWAL-backed table a previous query has already been refused
`_rowid` on, so the fusion goes straight to the primary key."""
table = mock.Mock()
table._hybrid_pk_fusion_learned.return_value = True
table.schema = schema
return table
_PK_SCHEMA = pa.schema(
[
pa.field(
"id",
pa.string(),
metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
),
pa.field("text", pa.string()),
]
)
def _wal_hybrid(**select):
builder = (
LanceHybridQueryBuilder(_wal_table(_PK_SCHEMA)).vector([0.1, 0.2]).text("puppy")
)
if "columns" in select:
builder = builder.select(select["columns"])
return builder
def test_wal_hybrid_projects_the_key_instead_of_asking_for_row_ids():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
assert builder._pk_fusion.pk_columns == ["id"]
assert builder._pk_fusion.injected == ["id"]
for leg in (builder._vector_query, builder._fts_query):
assert leg._columns == ["text", "id"]
# A MemWAL table rejects with_row_id before it plans.
assert leg._with_row_id is not True
def test_wal_hybrid_selecting_the_key_injects_nothing():
builder = _wal_hybrid(columns=["text", "id"])
builder._create_query_builders()
assert builder._pk_fusion.injected == []
def test_wal_hybrid_without_a_projection_leaves_it_alone():
builder = _wal_hybrid()
builder._create_query_builders()
# Select-all already returns the key.
assert builder._pk_fusion.injected == []
assert builder._vector_query._columns is None
def test_base_table_hybrid_still_joins_on_row_ids():
table = mock.Mock()
table._hybrid_pk_fusion_learned.return_value = False
builder = LanceHybridQueryBuilder(table).vector([0.1, 0.2]).text("puppy")
builder._create_query_builders()
# `to_arrow` is what turns row ids on for this path; all that matters here
# is that no key was resolved and no projection was touched.
assert builder._pk_fusion is None
assert builder._vector_query._columns is None
def test_wal_hybrid_with_use_lsm_false_joins_on_row_ids():
builder = _wal_hybrid().use_lsm(False)
builder._create_query_builders()
assert builder._pk_fusion is None
def test_wal_hybrid_refuses_an_explicit_with_row_id():
builder = _wal_hybrid().with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
builder._create_query_builders()
def test_wal_hybrid_return_score_all_is_not_refused():
"""`return_score="all"` turns row ids on for the reranker, not the caller."""
builder = _wal_hybrid().rerank(RRFReranker(return_score="all"))
builder._create_query_builders()
assert builder._pk_fusion.pk_columns == ["id"]
def test_wal_hybrid_fuses_on_the_key_and_drops_it_again():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
vector_results = pa.table(
{"text": ["a", "b"], "id": ["a", "b"], "_distance": [0.1, 0.2]}
)
fts_results = pa.table({"text": ["b", "c"], "id": ["b", "c"], "_score": [1.0, 2.0]})
vector_results, fts_results = builder._pk_fusion.stamp(vector_results, fts_results)
combined = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
vector_results=vector_results,
norm="score",
fts_query="puppy",
reranker=RRFReranker(),
limit=10,
with_row_ids=True,
)
results = builder._finish_hybrid_results(combined)
assert "_rowid" not in results.column_names, "the surrogate is internal"
assert "id" not in results.column_names, "the caller selected only `text`"
assert sorted(results.column("text").to_pylist()) == ["a", "b", "c"]
@pytest.mark.parametrize(
"row_id_first", [True, False], ids=["row_id_first", "rerank_first"]
)
def test_wal_hybrid_refuses_row_id_even_with_return_score_all(row_id_first: bool):
"""`return_score="all"` sets the same flag the caller's request does, so
asking for both must still be refused — in either call order."""
builder = _wal_hybrid()
reranker = RRFReranker(return_score="all")
if row_id_first:
builder = builder.with_row_id(True).rerank(reranker)
else:
builder = builder.rerank(reranker).with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
builder._create_query_builders()
def test_wal_hybrid_refuses_a_blob_projection_with_return_score_all():
"""A projected blob needs a real row id to fetch through, which the
reranker turning `_with_row_id` on does not provide."""
schema = pa.schema(
[
pa.field(
"id",
pa.string(),
metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
),
pa.field("text", pa.string()),
lancedb.blob("blob"),
]
)
builder = (
LanceHybridQueryBuilder(_wal_table(schema))
.vector([0.1, 0.2])
.text("puppy")
.select(["blob"])
.rerank(RRFReranker(return_score="all"))
)
with pytest.raises(NotImplementedError, match="blob"):
builder._create_query_builders()
_REFUSAL = RuntimeError(
"Bad request: the MemWAL LSM scanner does not support with_row_id"
)
def test_fallback_retries_on_the_key_and_remembers_it():
table = mock.Mock()
attempts = []
def run():
attempts.append(len(attempts))
if len(attempts) == 1:
raise _REFUSAL
return "fused"
result = with_pk_fallback(
run, table, fused=lambda: False, caller_requested_row_id=False
)
assert result == "fused"
assert len(attempts) == 2
table._note_hybrid_pk_fusion.assert_called_once()
def test_fallback_explains_the_refusal_to_a_caller_who_asked_for_row_ids():
def run():
raise _REFUSAL
with pytest.raises(NotImplementedError, match="use_lsm") as info:
with_pk_fallback(
run, mock.Mock(), fused=lambda: False, caller_requested_row_id=True
)
assert str(info.value) == WAL_ROW_ID_UNSUPPORTED
@pytest.mark.parametrize(
"error, fused",
[(ValueError("unrelated"), False), (_REFUSAL, True)],
ids=["other_error", "already_fused"],
)
def test_fallback_reraises_what_it_cannot_retry(error, fused):
def run():
raise error
table = mock.Mock()
with pytest.raises(type(error)):
with_pk_fallback(run, table, fused=lambda: fused, caller_requested_row_id=False)
table._note_hybrid_pk_fusion.assert_not_called()
python/src/table.rs
+2
View File
@@ -1605,6 +1605,8 @@ impl Table {
/// Whether a hybrid query on this table has already been refused `_rowid`.
/// Learned from a previous refusal, never probed, so this is free.
///
/// WAL-PK-FUSION: delete this and `note_hybrid_pk_fusion`.
pub fn hybrid_pk_fusion_learned(self_: PyRef<'_, Self>) -> PyResult<bool> {
Ok(self_.inner_ref()?.base_table().hybrid_pk_fusion_learned())
}
rust/lancedb/src/query.rs
+22 -162
View File
@@ -48,6 +48,7 @@ use crate::{
};
mod hybrid;
pub(crate) mod wal_fusion; // WAL-PK-FUSION: delete with the module.
pub(crate) const DEFAULT_TOP_K: usize = 10;
@@ -1176,88 +1177,6 @@ pub struct VectorQuery {
request: VectorQueryRequest,
}
/// The part of the MemWAL `with_row_id` refusal that identifies it.
///
/// Both layers raise this same sentence verbatim — sophon's
/// `LSM_WITH_ROW_ID_UNSUPPORTED` says so in as many words, and the OSS LSM
/// scanner builds it from the same words — so this substring survives either
/// path and the wrapping each adds. It is a coupling to a message rather than
/// a code, because the refusal arrives as a generic 400; if the server ever
/// rewords it the fallback stops firing and hybrid on a MemWAL table goes back
/// to erroring visibly, which is loud rather than silently wrong.
const WAL_ROW_ID_REFUSAL: &str = "does not support with_row_id";
/// Whether `error` is the server declining `_rowid` because the table is
/// MemWAL-backed, as opposed to any other failure the query might hit.
fn is_wal_row_id_refusal(error: &Error) -> bool {
error.to_string().contains(WAL_ROW_ID_REFUSAL)
}
/// Why a hybrid query on a MemWAL table cannot hand back `_rowid`, raised
/// whether the refusal was anticipated or came back from the server.
fn wal_row_id_unsupported() -> Error {
Error::NotSupported {
message: "hybrid search on a MemWAL table cannot return _rowid: the fresh tier has no \
stable row id, and the ids the fusion joins on are synthesized from the \
primary key. Set use_lsm(false) to read the base table only (results will \
exclude un-compacted MemWAL data)"
.to_string(),
}
}
/// How the hybrid fusion identifies the same row across its two legs.
enum FusionKey {
/// Lance's `_rowid`, requested from the server. Every table that can supply
/// one uses this.
RowId,
/// A surrogate derived client-side from the primary key. MemWAL tables take
/// this path: the fresh tier has no stable row id, so the server rejects
/// `with_row_id` outright.
PrimaryKey {
columns: Vec<String>,
/// Key columns the caller did not ask for, added to the projection so
/// the surrogate can be built and dropped again before returning.
injected: Vec<String>,
},
}
/// The key columns `select` does not already produce, and so would have to be
/// added to it for the surrogate to be buildable.
fn pk_columns_to_inject(select: &Select, pk_columns: &[String]) -> Vec<String> {
let produced: Vec<&String> = match select {
// Already every non-system column, the key among them.
Select::All => return Vec::new(),
Select::Columns(columns) => columns.iter().collect(),
Select::Dynamic(pairs) => pairs.iter().map(|(name, _)| name).collect(),
Select::Expr(pairs) => pairs.iter().map(|(name, _)| name).collect(),
};
pk_columns
.iter()
.filter(|pk| !produced.contains(pk))
.cloned()
.collect()
}
/// Add the key columns `select` is missing, as identity projections where the
/// selection is expression-shaped.
fn inject_pk_columns(select: &mut Select, pk_columns: &[String]) {
let missing = pk_columns_to_inject(select, pk_columns);
if missing.is_empty() {
return;
}
match select {
Select::All => {}
Select::Columns(columns) => columns.extend(missing),
Select::Dynamic(pairs) => pairs.extend(missing.into_iter().map(|pk| (pk.clone(), pk))),
Select::Expr(pairs) => {
pairs.extend(missing.into_iter().map(|pk| {
let expr = crate::expr::col(&pk);
(pk, expr)
}));
}
}
}
impl VectorQuery {
fn new(base: Query) -> Self {
Self {
@@ -1490,58 +1409,38 @@ impl VectorQuery {
self
}
/// Ask for `_rowid` first and fall back to the primary key only when the
/// server refuses, rather than paying a round trip up front to find out.
///
/// Only a MemWAL table refuses, so probing would tax every table to learn
/// something almost none of them need — and the refusal has to be handled
/// regardless, since a write spec installed elsewhere can arrive between
/// any two queries. The refusal costs nothing server-side: it is raised
/// before the query is planned.
pub async fn execute_hybrid(
&self,
options: QueryExecutionOptions,
) -> Result<SendableRecordBatchStream> {
let key = self.fusion_key().await?;
let optimistic = matches!(key, FusionKey::RowId);
match self.run_hybrid(key, options.clone()).await {
// A caller who asked for `_rowid` gets the reason, not a bare 400 —
// the same message the learned path would have raised up front.
Err(e) if optimistic && self.request.base.with_row_id && is_wal_row_id_refusal(&e) => {
Err(wal_row_id_unsupported())
}
Err(e) if optimistic && is_wal_row_id_refusal(&e) => {
// Learned, so later queries on this table skip straight to it.
self.parent.note_hybrid_pk_fusion();
let key = self.pk_fusion_key().await?;
self.run_hybrid(key, options).await
}
other => other,
}
// WAL-PK-FUSION: without the fallback, this body is `run_hybrid`'s,
// with its `pk_fusion` branches taken as `None`.
wal_fusion::with_pk_fallback(self, |pk_fusion| {
self.run_hybrid(pk_fusion, options.clone())
})
.await
}
async fn run_hybrid(
&self,
key: FusionKey,
pk_fusion: Option<wal_fusion::PkFusion>,
options: QueryExecutionOptions,
) -> Result<SendableRecordBatchStream> {
let max_batch_length = options.max_batch_length as usize;
let internal_options = options.without_output_batch_length_limit();
// clone query and specify we want to include row IDs, which can be needed for reranking
let mut fts_query = Query::new(self.parent.clone());
fts_query.request = self.request.base.clone();
let mut vector_query = self.clone();
match &key {
// The legs need the join column, which the reranking needs in turn.
FusionKey::RowId => {
// WAL-PK-FUSION: without the fallback, keep only the `None` arm.
match &pk_fusion {
None => {
fts_query = fts_query.with_row_id();
vector_query = vector_query.with_row_id();
}
// Ask for the key columns instead, and never for `_rowid`: the
// server rejects it on a MemWAL table before it plans.
FusionKey::PrimaryKey { columns, .. } => {
inject_pk_columns(&mut fts_query.request.select, columns);
inject_pk_columns(&mut vector_query.request.base.select, columns);
Some(pk_fusion) => {
pk_fusion.prepare_leg(&mut fts_query.request.select);
pk_fusion.prepare_leg(&mut vector_query.request.base.select);
}
}
@@ -1564,14 +1463,9 @@ impl VectorQuery {
let mut fts_results = concat_batches(&fts_schema, fts_results.iter())?;
let mut vec_results = concat_batches(&vec_schema, vec_results.iter())?;
// Vector first: `merge_results` concatenates in that order and keeps the
// first occurrence, so first-seen ids preserve its tie break.
if let FusionKey::PrimaryKey { columns, .. } = &key {
let mut legs = [vec_results, fts_results];
hybrid::stamp_surrogate_row_ids(&mut legs, columns)?;
let [vec, fts] = legs;
vec_results = vec;
fts_results = fts;
// WAL-PK-FUSION: delete.
if let Some(pk_fusion) = &pk_fusion {
(vec_results, fts_results) = pk_fusion.stamp(vec_results, fts_results)?;
}
if matches!(self.request.base.norm, Some(NormalizeMethod::Rank)) {
@@ -1609,53 +1503,19 @@ impl VectorQuery {
results = results.slice(0, limit);
}
match &key {
FusionKey::RowId => {
// WAL-PK-FUSION: without the fallback, keep only the `None` arm.
match &pk_fusion {
None => {
if !self.request.base.with_row_id {
results = results.drop_column(ROW_ID)?;
}
}
// The surrogate is an internal join key, never an answer: it goes
// whether or not the caller asked for `_rowid`, and `fusion_key`
// has already refused the query if they did.
FusionKey::PrimaryKey { injected, .. } => {
results = results.drop_column(ROW_ID)?;
for column in injected {
// A reranker is free to reshape its output, so only drop
// what actually came back.
if results.schema().column_with_name(column).is_some() {
results = results.drop_column(column)?;
}
}
}
Some(pk_fusion) => results = pk_fusion.strip(results)?,
}
Ok(single_batch_stream(results, max_batch_length))
}
/// Whether this hybrid query can join its legs on `_rowid`, or has to build
/// a surrogate from the primary key because the table is MemWAL-backed.
/// The key to try first: `_rowid` unless a previous query on this table has
/// already been refused it.
async fn fusion_key(&self) -> Result<FusionKey> {
if self.request.base.use_lsm == Some(false) || !self.parent.hybrid_pk_fusion_learned() {
return Ok(FusionKey::RowId);
}
self.pk_fusion_key().await
}
/// Resolve the primary key the legs will be fused on. Needs the schema,
/// which the table has already fetched and cached for its own reasons.
async fn pk_fusion_key(&self) -> Result<FusionKey> {
if self.request.base.with_row_id {
return Err(wal_row_id_unsupported());
}
let schema = self.parent.schema().await?;
let columns = crate::table::primary_key::pk_columns(&schema);
let injected = pk_columns_to_inject(&self.request.base.select, &columns);
Ok(FusionKey::PrimaryKey { columns, injected })
}
async fn inner_execute_with_options(
&self,
options: QueryExecutionOptions,
rust/lancedb/src/query/hybrid.rs
+1 -187
View File
@@ -5,13 +5,10 @@ use arrow::compute::{
kernels::numeric::{div, sub},
max, min,
};
use arrow::row::{RowConverter, SortField};
use arrow_array::{Float32Array, RecordBatch, UInt64Array, cast::downcast_array};
use arrow_array::{Float32Array, RecordBatch, cast::downcast_array};
use arrow_schema::{DataType, Field, Schema, SortOptions};
use lance::dataset::ROW_ID;
use lance_arrow::RecordBatchExt;
use lance_index::{scalar::inverted::SCORE_COL, vector::DIST_COL};
use std::collections::HashMap;
use std::sync::Arc;
use crate::error::{Error, Result};
@@ -176,90 +173,6 @@ pub fn normalize_scores(
Ok(results)
}
/// Replace each batch's join column with a dense surrogate derived from
/// `pk_columns`, in place and in the order given.
///
/// The fusion joins its two legs on [`ROW_ID`], which a MemWAL table cannot
/// supply — the fresh tier has no stable row id. It only ever compares the
/// column for equality, though (`merge_results` dedups, the rerankers group and
/// the score restore looks up), so any injective mapping of the primary key
/// serves. Ids are assigned in first-seen order across `batches`, so passing
/// them in fusion order keeps the dedup's "first occurrence wins" and its tie
/// break byte-identical to what a real row id would have produced.
///
/// Writing it under the [`ROW_ID`] name is what lets every reranker, including
/// third-party implementations that hardcode the name, run unmodified. The
/// column never reaches the caller: the hybrid query drops it before returning.
pub fn stamp_surrogate_row_ids(batches: &mut [RecordBatch], pk_columns: &[String]) -> Result<()> {
if pk_columns.is_empty() {
return Err(Error::InvalidInput {
message: "hybrid search on a MemWAL table needs an unenforced primary key to \
deduplicate on, and this table declares none"
.to_string(),
});
}
let mut converter: Option<RowConverter> = None;
let mut ids: HashMap<Vec<u8>, u64> = HashMap::new();
for batch in batches.iter_mut() {
// Both legs empty: `query_schemas` synthesizes a schema carrying only
// the score and join columns, so there is no key to read — and no row
// that would need one.
let row_ids: UInt64Array = if batch.num_rows() == 0 {
UInt64Array::from(Vec::<u64>::new())
} else {
let key_columns = pk_columns
.iter()
.map(|name| {
batch
.column_by_name(name)
.cloned()
.ok_or_else(|| Error::InvalidInput {
message: format!(
"hybrid search could not deduplicate: primary key column {} \
is missing from a result set with {} rows",
name,
batch.num_rows()
),
})
})
.collect::<Result<Vec<_>>>()?;
let converter = match converter {
Some(ref converter) => converter,
None => converter.insert(RowConverter::new(
key_columns
.iter()
.map(|column| SortField::new(column.data_type().clone()))
.collect(),
)?),
};
let rows = converter.convert_columns(&key_columns)?;
let next = &mut ids;
UInt64Array::from_iter_values(rows.iter().map(|row| {
let id = next.len() as u64;
*next.entry(row.as_ref().to_vec()).or_insert(id)
}))
};
// `query_schemas` synthesizes a schema carrying `ROW_ID` when a leg came
// back empty, so the column can already be there — and it is ours to
// define either way.
let base = match batch.schema().column_with_name(ROW_ID) {
Some(_) => batch.drop_column(ROW_ID)?,
None => batch.clone(),
};
*batch = base.try_with_column(
Field::new(ROW_ID, DataType::UInt64, false),
Arc::new(row_ids),
)?;
}
Ok(())
}
#[cfg(test)]
mod test {
use super::*;
@@ -433,103 +346,4 @@ mod test {
vec![0.0, 0.0, 0.0, 0.0, 0.0]
);
}
fn batch(ids: Vec<&str>, score: &str) -> RecordBatch {
let schema = Arc::new(Schema::new(vec![
Arc::new(Field::new("id", DataType::Utf8, false)),
Arc::new(Field::new(score, DataType::Float32, false)),
]));
let scores = Float32Array::from(vec![0.0_f32; ids.len()]);
RecordBatch::try_new(
schema,
vec![Arc::new(StringArray::from(ids)), Arc::new(scores)],
)
.unwrap()
}
fn row_ids(batch: &RecordBatch) -> Vec<u64> {
let ids: UInt64Array = downcast_array(batch.column_by_name(ROW_ID).unwrap());
ids.values().to_vec()
}
/// The same key in both legs has to land on the same id or the fusion will
/// not dedup, and ids must ascend in fusion order so the tie break holds.
#[test]
fn test_surrogate_row_ids_are_shared_across_legs_in_first_seen_order() {
let mut batches = vec![
batch(vec!["a", "b", "c"], DIST_COL),
batch(vec!["b", "d"], SCORE_COL),
];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), vec![0, 1, 2]);
assert_eq!(
row_ids(&batches[1]),
vec![1, 3],
"b keeps the id the vector leg gave it"
);
}
/// A string key is the common case and cannot be cast to the u64 the
/// rerankers read, which is the whole reason for the surrogate.
#[test]
fn test_surrogate_row_ids_repeat_within_a_leg() {
let mut batches = vec![batch(vec!["a", "a", "b"], DIST_COL)];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), vec![0, 0, 1]);
}
#[test]
fn test_surrogate_row_ids_support_a_composite_key() {
let schema = Arc::new(Schema::new(vec![
Arc::new(Field::new("tenant", DataType::Utf8, false)),
Arc::new(Field::new("id", DataType::Int32, false)),
]));
let make = |tenants: Vec<&str>, ids: Vec<i32>| {
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(StringArray::from(tenants)),
Arc::new(arrow_array::Int32Array::from(ids)),
],
)
.unwrap()
};
let mut batches = vec![make(vec!["x", "y"], vec![1, 1]), make(vec!["y"], vec![1])];
stamp_surrogate_row_ids(&mut batches, &["tenant".to_string(), "id".to_string()]).unwrap();
assert_eq!(
row_ids(&batches[0]),
vec![0, 1],
"same id, different tenant"
);
assert_eq!(row_ids(&batches[1]), vec![1]);
}
/// Both legs empty: `query_schemas` synthesizes a schema with no key column,
/// and there is no row that would need one.
#[test]
fn test_surrogate_row_ids_tolerate_the_both_empty_schema() {
let mut batches = vec![
RecordBatch::new_empty(Arc::new(empty_vec_schema())),
RecordBatch::new_empty(Arc::new(empty_fts_schema())),
];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), Vec::<u64>::new());
}
/// Rows with no key would silently fuse into one bucket, so refuse.
#[test]
fn test_surrogate_row_ids_reject_a_missing_key_column() {
let mut batches = vec![batch(vec!["a"], DIST_COL)];
let err = stamp_surrogate_row_ids(&mut batches, &["nope".to_string()]).unwrap_err();
assert!(err.to_string().contains("nope"), "{err}");
}
#[test]
fn test_surrogate_row_ids_reject_a_table_with_no_primary_key() {
let mut batches = vec![batch(vec!["a"], DIST_COL)];
let err = stamp_surrogate_row_ids(&mut batches, &[]).unwrap_err();
assert!(err.to_string().contains("unenforced primary key"), "{err}");
}
}
rust/lancedb/src/query/wal_fusion.rs
+523
View File
@@ -0,0 +1,523 @@
// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The LanceDB Authors
//! WAL-PK-FUSION: hybrid search's primary-key fallback for MemWAL tables.
//!
//! Hybrid search joins its vector and full-text legs on `_rowid`, which a
//! MemWAL table cannot supply: the fresh tier has no stable row id, so the
//! server rejects `with_row_id` before it plans. When that happens the legs are
//! joined on a surrogate `_rowid` built from the table's primary key instead.
//!
//! This is a stopgap until MemWAL supports `_rowid`, and everything it needs
//! lives here so it can be removed in one pass:
//!
//! 1. Delete this file and its `mod wal_fusion;` line in `query.rs`.
//! 2. Run `grep -rn WAL-PK-FUSION rust/ python/` and follow the note at each
//! marker; the Python side has its own copy in `lancedb/_wal_hybrid.py`.
use std::collections::HashMap;
use std::future::Future;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant};
use arrow::row::{RowConverter, SortField};
use arrow_array::{RecordBatch, UInt64Array};
use arrow_schema::{DataType, Field};
use lance::dataset::ROW_ID;
use lance_arrow::RecordBatchExt;
use lance_core::datatypes::{LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION};
use lance_core::utils::parse::str_is_truthy;
use super::{Select, VectorQuery};
use crate::error::{Error, Result};
/// The part of the MemWAL `with_row_id` refusal that identifies it.
///
/// Both layers raise this same sentence verbatim — sophon's
/// `LSM_WITH_ROW_ID_UNSUPPORTED` says so in as many words, and the OSS LSM
/// scanner builds it from the same words — so this substring survives either
/// path and the wrapping each adds. It is a coupling to a message rather than
/// a code, because the refusal arrives as a generic 400; if the server ever
/// rewords it the fallback stops firing and hybrid on a MemWAL table goes back
/// to erroring visibly, which is loud rather than silently wrong.
const WAL_ROW_ID_REFUSAL: &str = "does not support with_row_id";
/// Whether `error` is the server declining `_rowid` because the table is
/// MemWAL-backed, as opposed to any other failure the query might hit.
fn is_wal_row_id_refusal(error: &Error) -> bool {
error.to_string().contains(WAL_ROW_ID_REFUSAL)
}
/// Why a hybrid query on a MemWAL table cannot hand back `_rowid`, raised
/// whether the refusal was anticipated or came back from the server.
fn wal_row_id_unsupported() -> Error {
Error::NotSupported {
message: "hybrid search on a MemWAL table cannot return _rowid: the fresh tier has no \
stable row id, and the ids the fusion joins on are synthesized from the \
primary key. Set use_lsm(false) to read the base table only (results will \
exclude un-compacted MemWAL data)"
.to_string(),
}
}
/// Run a hybrid query, joining its legs on the primary key instead of `_rowid`
/// when the table is MemWAL-backed.
///
/// `run(None)` must join the legs on `_rowid`; `run(Some(fusion))` must build
/// them with `fusion` instead.
///
/// `_rowid` is asked for first and the primary key used only when the server
/// refuses, rather than paying a round trip up front to find out. Only a
/// MemWAL table refuses, so probing would tax every table to learn something
/// almost none of them need — and the refusal has to be handled regardless,
/// since a write spec installed elsewhere can arrive between any two queries.
/// The refusal costs nothing server-side: it is raised before the query is
/// planned.
pub(super) async fn with_pk_fallback<T, F, Fut>(query: &VectorQuery, run: F) -> Result<T>
where
F: Fn(Option<PkFusion>) -> Fut,
Fut: Future<Output = Result<T>>,
{
// A previous query on this table was already refused, so skip straight to
// the key.
if query.request.base.use_lsm != Some(false) && query.parent.hybrid_pk_fusion_learned() {
return run(Some(PkFusion::resolve(query).await?)).await;
}
match run(None).await {
Err(e) if is_wal_row_id_refusal(&e) => {
// A caller who asked for `_rowid` gets the reason, not a bare 400
// — the same message the learned path raises up front.
if query.request.base.with_row_id {
return Err(wal_row_id_unsupported());
}
// Learned, so later queries on this table skip straight to it.
query.parent.note_hybrid_pk_fusion();
run(Some(PkFusion::resolve(query).await?)).await
}
other => other,
}
}
/// How long a table remembers that it refused `_rowid`.
///
/// Nothing fetches this — the refusal seeds it — so the bound exists only to
/// stop the memory outliving a write spec someone removed through another
/// handle. Being late to notice keeps the query on the primary key, which on a
/// base table is still a valid read; the window is what stops that lasting.
const MEMORY_TTL: Duration = Duration::from_secs(30);
/// A table's memory of having refused `_rowid`, backing
/// [`BaseTable::hybrid_pk_fusion_learned`](crate::table::BaseTable::hybrid_pk_fusion_learned)
/// for table types that keep one.
#[derive(Debug, Default)]
pub struct PkFusionMemory(Mutex<Option<Instant>>);
impl PkFusionMemory {
pub fn learned(&self) -> bool {
self.0
.lock()
.unwrap()
.is_some_and(|learned| learned.elapsed() < MEMORY_TTL)
}
pub fn note(&self) {
*self.0.lock().unwrap() = Some(Instant::now());
}
/// Forget the refusal, as when the write spec is removed through this handle.
pub fn forget(&self) {
*self.0.lock().unwrap() = None;
}
}
/// How the legs of a hybrid query on a MemWAL table are joined: on a surrogate
/// `_rowid` derived client-side from the primary key.
pub(super) struct PkFusion {
columns: Vec<String>,
/// Key columns the caller did not ask for, added to the projection so the
/// surrogate can be built and dropped again before returning.
injected: Vec<String>,
}
impl PkFusion {
/// Resolve the primary key the legs will be fused on. Needs the schema,
/// which the table has already fetched and cached for its own reasons.
async fn resolve(query: &VectorQuery) -> Result<Self> {
if query.request.base.with_row_id {
return Err(wal_row_id_unsupported());
}
let schema = query.parent.schema().await?;
let columns = pk_columns(&schema);
let injected = pk_columns_to_inject(&query.request.base.select, &columns);
Ok(Self { columns, injected })
}
/// Make a leg return the key columns, in place of the `_rowid` a MemWAL
/// table would refuse.
pub(super) fn prepare_leg(&self, select: &mut Select) {
inject_pk_columns(select, &self.columns);
}
/// Give both legs a surrogate `_rowid` built from the key.
pub(super) fn stamp(
&self,
vector: RecordBatch,
fts: RecordBatch,
) -> Result<(RecordBatch, RecordBatch)> {
// Vector first: `merge_results` concatenates in that order and keeps
// the first occurrence, so first-seen ids preserve its tie break.
let mut legs = [vector, fts];
stamp_surrogate_row_ids(&mut legs, &self.columns)?;
let [vector, fts] = legs;
Ok((vector, fts))
}
/// Drop the surrogate and any injected key columns from the reranked
/// results. The surrogate is an internal join key, never an answer: it
/// goes whether or not the caller asked for `_rowid`, and `resolve` has
/// already refused the query if they did.
pub(super) fn strip(&self, results: RecordBatch) -> Result<RecordBatch> {
let mut results = results.drop_column(ROW_ID)?;
for column in &self.injected {
// A reranker is free to reshape its output, so only drop what
// actually came back.
if results.schema().column_with_name(column).is_some() {
results = results.drop_column(column)?;
}
}
Ok(results)
}
}
/// The key columns `select` does not already produce, and so would have to be
/// added to it for the surrogate to be buildable.
fn pk_columns_to_inject(select: &Select, pk_columns: &[String]) -> Vec<String> {
let produced: Vec<&String> = match select {
// Already every non-system column, the key among them.
Select::All => return Vec::new(),
Select::Columns(columns) => columns.iter().collect(),
Select::Dynamic(pairs) => pairs.iter().map(|(name, _)| name).collect(),
Select::Expr(pairs) => pairs.iter().map(|(name, _)| name).collect(),
};
pk_columns
.iter()
.filter(|pk| !produced.contains(pk))
.cloned()
.collect()
}
/// Add the key columns `select` is missing, as identity projections where the
/// selection is expression-shaped.
fn inject_pk_columns(select: &mut Select, pk_columns: &[String]) {
let missing = pk_columns_to_inject(select, pk_columns);
if missing.is_empty() {
return;
}
match select {
Select::All => {}
Select::Columns(columns) => columns.extend(missing),
Select::Dynamic(pairs) => pairs.extend(missing.into_iter().map(|pk| (pk.clone(), pk))),
Select::Expr(pairs) => {
pairs.extend(missing.into_iter().map(|pk| {
let expr = crate::expr::col(&pk);
(pk, expr)
}));
}
}
}
/// The unenforced primary key columns declared in `schema`, in key order.
///
/// Reads the field metadata directly rather than going through
/// [`lance_core::datatypes::Schema::unenforced_primary_key`], because a remote
/// table only ever has the Arrow schema `describe` returned — there is no Lance
/// schema, and no field ids, on that side of the wire.
///
/// Ordering mirrors Lance's: fields with an explicit 1-based position first, in
/// position order, then fields carrying only the legacy boolean marker, in
/// declaration order. Only top-level fields are considered, which is every
/// field a supported key dtype can live on.
fn pk_columns(schema: &arrow_schema::Schema) -> Vec<String> {
let mut keyed: Vec<(bool, u32, usize, &str)> = schema
.fields()
.iter()
.enumerate()
.filter_map(|(idx, field)| {
let metadata = field.metadata();
let position = metadata
.get(LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
.and_then(|value| value.parse::<u32>().ok())
.or_else(|| {
metadata
.get(LANCE_UNENFORCED_PRIMARY_KEY)
.filter(|value| str_is_truthy(value))
.map(|_| 0)
})?;
Some((position == 0, position, idx, field.name().as_str()))
})
.collect();
keyed.sort_unstable();
keyed
.into_iter()
.map(|(_, _, _, name)| name.to_string())
.collect()
}
/// Replace each batch's join column with a dense surrogate derived from
/// `pk_columns`, in place and in the order given.
///
/// The fusion joins its two legs on [`ROW_ID`], which a MemWAL table cannot
/// supply — the fresh tier has no stable row id. It only ever compares the
/// column for equality, though (`merge_results` dedups, the rerankers group and
/// the score restore looks up), so any injective mapping of the primary key
/// serves. Ids are assigned in first-seen order across `batches`, so passing
/// them in fusion order keeps the dedup's "first occurrence wins" and its tie
/// break byte-identical to what a real row id would have produced.
///
/// Writing it under the [`ROW_ID`] name is what lets every reranker, including
/// third-party implementations that hardcode the name, run unmodified. The
/// column never reaches the caller: the hybrid query drops it before returning.
fn stamp_surrogate_row_ids(batches: &mut [RecordBatch], pk_columns: &[String]) -> Result<()> {
if pk_columns.is_empty() {
return Err(Error::InvalidInput {
message: "hybrid search on a MemWAL table needs an unenforced primary key to \
deduplicate on, and this table declares none"
.to_string(),
});
}
let mut converter: Option<RowConverter> = None;
let mut ids: HashMap<Vec<u8>, u64> = HashMap::new();
for batch in batches.iter_mut() {
// Both legs empty: `query_schemas` synthesizes a schema carrying only
// the score and join columns, so there is no key to read — and no row
// that would need one.
let row_ids: UInt64Array = if batch.num_rows() == 0 {
UInt64Array::from(Vec::<u64>::new())
} else {
let key_columns = pk_columns
.iter()
.map(|name| {
batch
.column_by_name(name)
.cloned()
.ok_or_else(|| Error::InvalidInput {
message: format!(
"hybrid search could not deduplicate: primary key column {} \
is missing from a result set with {} rows",
name,
batch.num_rows()
),
})
})
.collect::<Result<Vec<_>>>()?;
let converter = match converter {
Some(ref converter) => converter,
None => converter.insert(RowConverter::new(
key_columns
.iter()
.map(|column| SortField::new(column.data_type().clone()))
.collect(),
)?),
};
let rows = converter.convert_columns(&key_columns)?;
let next = &mut ids;
UInt64Array::from_iter_values(rows.iter().map(|row| {
let id = next.len() as u64;
*next.entry(row.as_ref().to_vec()).or_insert(id)
}))
};
// `query_schemas` synthesizes a schema carrying `ROW_ID` when a leg came
// back empty, so the column can already be there — and it is ours to
// define either way.
let base = match batch.schema().column_with_name(ROW_ID) {
Some(_) => batch.drop_column(ROW_ID)?,
None => batch.clone(),
};
*batch = base.try_with_column(
Field::new(ROW_ID, DataType::UInt64, false),
Arc::new(row_ids),
)?;
}
Ok(())
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use arrow::array::downcast_array;
use arrow_array::{Float32Array, StringArray};
use arrow_schema::Schema;
use lance_index::{scalar::inverted::SCORE_COL, vector::DIST_COL};
use super::super::hybrid::{empty_fts_schema, empty_vec_schema};
use super::*;
fn field(name: &str, metadata: &[(&str, &str)]) -> Field {
Field::new(name, DataType::Utf8, false).with_metadata(
metadata
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect::<HashMap<_, _>>(),
)
}
#[test]
fn pk_columns_reads_the_position_marker() {
let schema = Schema::new(vec![
field("other", &[]),
field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]),
]);
assert_eq!(pk_columns(&schema), vec!["id".to_string()]);
}
/// Key order is the position, not the order the columns are declared in.
#[test]
fn pk_columns_orders_a_composite_key_by_position() {
let schema = Schema::new(vec![
field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "2")]),
field("tenant", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]),
]);
assert_eq!(
pk_columns(&schema),
vec!["tenant".to_string(), "id".to_string()]
);
}
/// Datasets written before the position marker carry only the boolean.
#[test]
fn pk_columns_accepts_the_legacy_boolean_marker() {
let schema = Schema::new(vec![field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY, "true")])]);
assert_eq!(pk_columns(&schema), vec!["id".to_string()]);
}
/// Mirrors Lance: an explicit position sorts ahead of a legacy field.
#[test]
fn pk_columns_sorts_positioned_fields_ahead_of_legacy_ones() {
let schema = Schema::new(vec![
field("legacy", &[(LANCE_UNENFORCED_PRIMARY_KEY, "yes")]),
field(
"positioned",
&[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")],
),
]);
assert_eq!(
pk_columns(&schema),
vec!["positioned".to_string(), "legacy".to_string()]
);
}
#[test]
fn pk_columns_is_empty_without_a_key() {
let schema = Schema::new(vec![
field("a", &[]),
field("b", &[(LANCE_UNENFORCED_PRIMARY_KEY, "false")]),
]);
assert!(pk_columns(&schema).is_empty());
}
fn batch(ids: Vec<&str>, score: &str) -> RecordBatch {
let schema = Arc::new(Schema::new(vec![
Arc::new(Field::new("id", DataType::Utf8, false)),
Arc::new(Field::new(score, DataType::Float32, false)),
]));
let scores = Float32Array::from(vec![0.0_f32; ids.len()]);
RecordBatch::try_new(
schema,
vec![Arc::new(StringArray::from(ids)), Arc::new(scores)],
)
.unwrap()
}
fn row_ids(batch: &RecordBatch) -> Vec<u64> {
let ids: UInt64Array = downcast_array(batch.column_by_name(ROW_ID).unwrap());
ids.values().to_vec()
}
/// The same key in both legs has to land on the same id or the fusion will
/// not dedup, and ids must ascend in fusion order so the tie break holds.
#[test]
fn test_surrogate_row_ids_are_shared_across_legs_in_first_seen_order() {
let mut batches = vec![
batch(vec!["a", "b", "c"], DIST_COL),
batch(vec!["b", "d"], SCORE_COL),
];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), vec![0, 1, 2]);
assert_eq!(
row_ids(&batches[1]),
vec![1, 3],
"b keeps the id the vector leg gave it"
);
}
/// A string key is the common case and cannot be cast to the u64 the
/// rerankers read, which is the whole reason for the surrogate.
#[test]
fn test_surrogate_row_ids_repeat_within_a_leg() {
let mut batches = vec![batch(vec!["a", "a", "b"], DIST_COL)];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), vec![0, 0, 1]);
}
#[test]
fn test_surrogate_row_ids_support_a_composite_key() {
let schema = Arc::new(Schema::new(vec![
Arc::new(Field::new("tenant", DataType::Utf8, false)),
Arc::new(Field::new("id", DataType::Int32, false)),
]));
let make = |tenants: Vec<&str>, ids: Vec<i32>| {
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(StringArray::from(tenants)),
Arc::new(arrow_array::Int32Array::from(ids)),
],
)
.unwrap()
};
let mut batches = vec![make(vec!["x", "y"], vec![1, 1]), make(vec!["y"], vec![1])];
stamp_surrogate_row_ids(&mut batches, &["tenant".to_string(), "id".to_string()]).unwrap();
assert_eq!(
row_ids(&batches[0]),
vec![0, 1],
"same id, different tenant"
);
assert_eq!(row_ids(&batches[1]), vec![1]);
}
/// Both legs empty: `query_schemas` synthesizes a schema with no key column,
/// and there is no row that would need one.
#[test]
fn test_surrogate_row_ids_tolerate_the_both_empty_schema() {
let mut batches = vec![
RecordBatch::new_empty(Arc::new(empty_vec_schema())),
RecordBatch::new_empty(Arc::new(empty_fts_schema())),
];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), Vec::<u64>::new());
}
/// Rows with no key would silently fuse into one bucket, so refuse.
#[test]
fn test_surrogate_row_ids_reject_a_missing_key_column() {
let mut batches = vec![batch(vec!["a"], DIST_COL)];
let err = stamp_surrogate_row_ids(&mut batches, &["nope".to_string()]).unwrap_err();
assert!(err.to_string().contains("nope"), "{err}");
}
#[test]
fn test_surrogate_row_ids_reject_a_table_with_no_primary_key() {
let mut batches = vec![batch(vec!["a"], DIST_COL)];
let err = stamp_surrogate_row_ids(&mut batches, &[]).unwrap_err();
assert!(err.to_string().contains("unenforced primary key"), "{err}");
}
}
rust/lancedb/src/remote/table.rs
+15 -24
View File
@@ -20,6 +20,7 @@ use crate::job::Job;
use crate::materialized_view::{
MaterializedViewDefinition, MaterializedViewInfo, RefreshMaterializedViewResult, ViewProjection,
};
use crate::query::wal_fusion::PkFusionMemory; // WAL-PK-FUSION: delete.
use crate::query::{QueryFilter, QueryRequest, Select, VectorQueryRequest};
use crate::remote::job::RemoteJob;
use crate::table::AddColumnsResult;
@@ -79,7 +80,7 @@ use std::collections::{HashMap, HashSet};
use std::io::Cursor;
use std::pin::Pin;
use std::sync::{Arc, Mutex};
use std::time::{Duration, Instant, SystemTime};
use std::time::{Duration, SystemTime};
use tokio::sync::RwLock;
const REQUEST_TIMEOUT_HEADER: HeaderName = HeaderName::from_static("x-request-timeout-ms");
@@ -91,13 +92,6 @@ const METRIC_TYPE_KEY: &str = "metric_type";
const INDEX_TYPE_KEY: &str = "index_type";
const SCHEMA_CACHE_TTL: Duration = Duration::from_secs(30);
const SCHEMA_CACHE_REFRESH_WINDOW: Duration = Duration::from_secs(5);
/// How long a hybrid query remembers that this table refused `_rowid`.
///
/// Nothing fetches this — the refusal seeds it — so the bound exists only to
/// stop the memory outliving a write spec someone removed through another
/// handle. Being late to notice keeps the query on the primary key, which on a
/// base table is still a valid read; the window is what stops that lasting.
const WAL_FUSION_MEMORY_TTL: Duration = Duration::from_secs(30);
const SCHEMA_SELECTOR_CHANGED: &str = "table selector changed while fetching schema";
fn fts_query_requires_document_granularity_support(query: &FtsQuery) -> bool {
@@ -504,7 +498,8 @@ pub struct RemoteTable<S: HttpSend = Sender> {
version: Arc<RwLock<Option<u64>>>,
location: RwLock<Option<String>>,
schema_cache: BackgroundCache<SchemaRef, Error>,
wal_pk_fusion: Mutex<Option<Instant>>,
// WAL-PK-FUSION: delete, with its initializers below.
wal_pk_fusion: PkFusionMemory,
freshness: Arc<Mutex<FreshnessState>>,
/// The branch this handle is scoped to, or `None` for the main branch.
/// Stamped onto every branch-accepting request so reads and writes resolve
@@ -673,7 +668,7 @@ impl<S: HttpSend> RemoteTable<S> {
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
wal_pk_fusion: Mutex::new(None),
wal_pk_fusion: PkFusionMemory::default(),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -707,7 +702,7 @@ impl<S: HttpSend> RemoteTable<S> {
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
wal_pk_fusion: Mutex::new(None),
wal_pk_fusion: PkFusionMemory::default(),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch,
}
@@ -1442,10 +1437,6 @@ impl<S: HttpSend> RemoteTable<S> {
self.schema_cache.invalidate();
}
fn invalidate_wal_pk_fusion(&self) {
*self.wal_pk_fusion.lock().unwrap() = None;
}
fn handle_error_invalidation(&self, error: &Error) {
let status_code = match error {
Error::Http { status_code, .. } => *status_code,
@@ -1629,7 +1620,7 @@ mod test_utils {
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
wal_pk_fusion: Mutex::new(None),
wal_pk_fusion: PkFusionMemory::default(),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -1654,7 +1645,7 @@ mod test_utils {
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
wal_pk_fusion: Mutex::new(None),
wal_pk_fusion: PkFusionMemory::default(),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -1688,7 +1679,7 @@ mod test_utils {
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
wal_pk_fusion: Mutex::new(None),
wal_pk_fusion: PkFusionMemory::default(),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -3250,7 +3241,7 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
));
let (request_id, response) = self.send(request, true).await?;
self.check_table_response(&request_id, response).await?;
self.invalidate_wal_pk_fusion();
self.wal_pk_fusion.forget(); // WAL-PK-FUSION: delete.
Ok(())
}
@@ -3317,15 +3308,13 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
Ok(Some(spec))
}
// WAL-PK-FUSION: delete both.
fn hybrid_pk_fusion_learned(&self) -> bool {
self.wal_pk_fusion
.lock()
.unwrap()
.is_some_and(|learned| learned.elapsed() < WAL_FUSION_MEMORY_TTL)
self.wal_pk_fusion.learned()
}
fn note_hybrid_pk_fusion(&self) {
*self.wal_pk_fusion.lock().unwrap() = Some(Instant::now());
self.wal_pk_fusion.note();
}
async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
@@ -13615,6 +13604,8 @@ mod tests {
branch.stats().await.unwrap();
}
// WAL-PK-FUSION: delete everything from here to the end of `mod tests`.
/// Verbatim from sophon's `LSM_WITH_ROW_ID_UNSUPPORTED`; the fallback keys
/// off this sentence, so a test that paraphrased it would prove nothing.
const WAL_ROW_ID_REFUSAL_BODY: &str = r#"{"code":13,"error":"Bad request: the MemWAL LSM scanner does not support with_row_id (the LSM scanner exposes _rowaddr, not a stable _rowid); set use_lsm(false) to read the base table only (results will exclude un-compacted MemWAL data)"}"#;
rust/lancedb/src/table.rs
+3 -1
View File
@@ -80,7 +80,7 @@ pub mod freshness;
pub mod lsm_stats;
pub mod merge;
pub mod optimize;
pub(crate) mod primary_key;
mod primary_key;
pub mod query;
pub mod refresh;
pub mod schema_evolution;
@@ -709,6 +709,8 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync {
/// Whether a hybrid query on this table has already been told it cannot
/// join its legs on `_rowid`.
///
/// WAL-PK-FUSION: delete this and `note_hybrid_pk_fusion`.
///
/// Learned, never probed: hybrid optimistically asks for `_rowid` and only
/// a MemWAL table refuses, so paying a round trip up front would tax every
/// table to discover something almost none of them need. Synchronous and
rust/lancedb/src/table/primary_key.rs
-110
View File
@@ -13,7 +13,6 @@
use arrow_schema::DataType;
use lance_core::datatypes::{LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION};
use lance_core::utils::parse::str_is_truthy;
use crate::error::{Error, Result};
use crate::table::NativeTable;
@@ -112,112 +111,3 @@ fn is_supported_pk_dtype(dtype: &DataType) -> bool {
| DataType::FixedSizeBinary(_)
)
}
/// The unenforced primary key columns declared in `schema`, in key order.
///
/// Reads the field metadata directly rather than going through
/// [`lance_core::datatypes::Schema::unenforced_primary_key`], because a remote
/// table only ever has the Arrow schema `describe` returned — there is no Lance
/// schema, and no field ids, on that side of the wire.
///
/// Ordering mirrors Lance's: fields with an explicit 1-based position first, in
/// position order, then fields carrying only the legacy boolean marker, in
/// declaration order. Only top-level fields are considered, which is every
/// field a supported key dtype can live on.
pub fn pk_columns(schema: &arrow_schema::Schema) -> Vec<String> {
let mut keyed: Vec<(bool, u32, usize, &str)> = schema
.fields()
.iter()
.enumerate()
.filter_map(|(idx, field)| {
let metadata = field.metadata();
let position = metadata
.get(LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
.and_then(|value| value.parse::<u32>().ok())
.or_else(|| {
metadata
.get(LANCE_UNENFORCED_PRIMARY_KEY)
.filter(|value| str_is_truthy(value))
.map(|_| 0)
})?;
Some((position == 0, position, idx, field.name().as_str()))
})
.collect();
keyed.sort_unstable();
keyed
.into_iter()
.map(|(_, _, _, name)| name.to_string())
.collect()
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use arrow_schema::{Field, Schema};
use super::*;
fn field(name: &str, metadata: &[(&str, &str)]) -> Field {
Field::new(name, DataType::Utf8, false).with_metadata(
metadata
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect::<HashMap<_, _>>(),
)
}
#[test]
fn pk_columns_reads_the_position_marker() {
let schema = Schema::new(vec![
field("other", &[]),
field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]),
]);
assert_eq!(pk_columns(&schema), vec!["id".to_string()]);
}
/// Key order is the position, not the order the columns are declared in.
#[test]
fn pk_columns_orders_a_composite_key_by_position() {
let schema = Schema::new(vec![
field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "2")]),
field("tenant", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]),
]);
assert_eq!(
pk_columns(&schema),
vec!["tenant".to_string(), "id".to_string()]
);
}
/// Datasets written before the position marker carry only the boolean.
#[test]
fn pk_columns_accepts_the_legacy_boolean_marker() {
let schema = Schema::new(vec![field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY, "true")])]);
assert_eq!(pk_columns(&schema), vec!["id".to_string()]);
}
/// Mirrors Lance: an explicit position sorts ahead of a legacy field.
#[test]
fn pk_columns_sorts_positioned_fields_ahead_of_legacy_ones() {
let schema = Schema::new(vec![
field("legacy", &[(LANCE_UNENFORCED_PRIMARY_KEY, "yes")]),
field(
"positioned",
&[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")],
),
]);
assert_eq!(
pk_columns(&schema),
vec!["positioned".to_string(), "legacy".to_string()]
);
}
#[test]
fn pk_columns_is_empty_without_a_key() {
let schema = Schema::new(vec![
field("a", &[]),
field("b", &[(LANCE_UNENFORCED_PRIMARY_KEY, "false")]),
]);
assert!(pk_columns(&schema).is_empty());
}
}