From 7480a5e428cec56f8a84bbbfd1d4c04dcc46f1b3 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Thu, 24 Sep 2026 09:28:21 -0500 Subject: [PATCH] refactor(query): isolate the hybrid primary-key fallback for removal 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` 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) --- python/python/lancedb/_lancedb.pyi | 1 + python/python/lancedb/_wal_hybrid.py | 282 ++++++++++++ python/python/lancedb/query.py | 334 +++------------ python/python/lancedb/remote/table.py | 1 + python/python/lancedb/table.py | 2 + python/python/tests/test_hybrid_query.py | 156 ------- python/python/tests/test_wal_hybrid.py | 219 ++++++++++ python/src/table.rs | 2 + rust/lancedb/src/query.rs | 184 +------- rust/lancedb/src/query/hybrid.rs | 188 +------- rust/lancedb/src/query/wal_fusion.rs | 523 +++++++++++++++++++++++ rust/lancedb/src/remote/table.rs | 39 +- rust/lancedb/src/table.rs | 4 +- rust/lancedb/src/table/primary_key.rs | 110 ----- 14 files changed, 1136 insertions(+), 909 deletions(-) create mode 100644 python/python/lancedb/_wal_hybrid.py create mode 100644 python/python/tests/test_wal_hybrid.py create mode 100644 rust/lancedb/src/query/wal_fusion.rs diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index bfb1b2b17..8595b823b 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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: ... diff --git a/python/python/lancedb/_wal_hybrid.py b/python/python/lancedb/_wal_hybrid.py new file mode 100644 index 000000000..67597950d --- /dev/null +++ b/python/python/lancedb/_wal_hybrid.py @@ -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() diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index f12cc9c25..c42344ad3 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -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 diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 797b18d9b..44d16211c 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -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() diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 23bb05605..493569829 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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() diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index 7d86d37f0..61b8001cf 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -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() diff --git a/python/python/tests/test_wal_hybrid.py b/python/python/tests/test_wal_hybrid.py new file mode 100644 index 000000000..11a9755e0 --- /dev/null +++ b/python/python/tests/test_wal_hybrid.py @@ -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() diff --git a/python/src/table.rs b/python/src/table.rs index e9aa0e906..b7763d0c6 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -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 { Ok(self_.inner_ref()?.base_table().hybrid_pk_fusion_learned()) } diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 4eb7d30e7..c8b19f013 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -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, - /// 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, - }, -} - -/// 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 { - 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 { - 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, options: QueryExecutionOptions, ) -> Result { 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 { - 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 { - 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, diff --git a/rust/lancedb/src/query/hybrid.rs b/rust/lancedb/src/query/hybrid.rs index e656daa0d..35c64d4a9 100644 --- a/rust/lancedb/src/query/hybrid.rs +++ b/rust/lancedb/src/query/hybrid.rs @@ -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 = None; - let mut ids: HashMap, 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::::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::>>()?; - - 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 { - 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| { - 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::::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}"); - } } diff --git a/rust/lancedb/src/query/wal_fusion.rs b/rust/lancedb/src/query/wal_fusion.rs new file mode 100644 index 000000000..b438f6150 --- /dev/null +++ b/rust/lancedb/src/query/wal_fusion.rs @@ -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(query: &VectorQuery, run: F) -> Result +where + F: Fn(Option) -> Fut, + Fut: Future>, +{ + // 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>); + +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, + /// 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, +} + +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 { + 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 { + 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 { + 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 { + 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::().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 = None; + let mut ids: HashMap, 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::::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::>>()?; + + 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::>(), + ) + } + + #[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 { + 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| { + 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::::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}"); + } +} diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index f4147c5ce..d6abfea9a 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -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 { version: Arc>>, location: RwLock>, schema_cache: BackgroundCache, - wal_pk_fusion: Mutex>, + // WAL-PK-FUSION: delete, with its initializers below. + wal_pk_fusion: PkFusionMemory, freshness: Arc>, /// 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 RemoteTable { 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 RemoteTable { 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 RemoteTable { 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 BaseTable for RemoteTable { )); 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 BaseTable for RemoteTable { 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> { @@ -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)"}"#; diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 5ab805745..dcb3e1578 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -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 diff --git a/rust/lancedb/src/table/primary_key.rs b/rust/lancedb/src/table/primary_key.rs index f8932bfa1..8a7b48efa 100644 --- a/rust/lancedb/src/table/primary_key.rs +++ b/rust/lancedb/src/table/primary_key.rs @@ -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 { - 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::().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::>(), - ) - } - - #[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()); - } -}