feat(query): dedup hybrid search by primary key on MemWAL tables (#4268)

Implementing a PK-based dedup to support WAL in hybrid search.
Basically, the existing code just calls a vector query + FTS query and
dedups locally using the row ID. Since row ID is unsupported in WAL
(currently), we update the algorithm to use the tables PK (which is
required in WAL, and uniquely linked to a PK on updates). To enable
PK-dedup branch, we've chosen to fallback on error rather than fire an
exploratory query. Specifically, we issue the regular hybrid query,
using row IDs, and when that fails with "WAL does not support row IDs"
we fallback to the PK-based variant, and then cache the fallback for
future calls from the same client. The alternative is that we chec if
WAL is enabled before issuing the query. This would require a
`get_lsm_write_spec` call, slowing down all the existing non-WAL paths
unnecessarily.

The ultimate goal is to get row ID supported in WAL, but that is pending
stable row ID GA and subsequent integration into the enterprise stack.
Once that lands, we will rip out the PK-based dedup that is only for
WAL-enabled tables and use the existing row ID code. This is the main
reason for building this adaptively. I had spec-ed out a refactor so
that rather than duplicating code between python and rust we build this
in rust and have python call that through PyO3. However, that made the
reverse path (ie. ) more difficult. So I think the tradeoff of
maintaining duplicate code in the short term for the ease of reverting
this in the medium term (hopefully 2-3 months).

---------

Co-authored-by: Claude Opus 5 <noreply@anthropic.com>

python/python/lancedb/_lancedb.pyi

@@ -514,6 +514,9 @@ 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: ...
     async def flush_lsm(self) -> None: ...
     async def compact_lsm(self) -> None: ...

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()

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,
@@ -2207,6 +2208,12 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
         self._phrase_query = None
         self._lower_bound = None
         self._upper_bound = None
+        # 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
 
     def _validate_query(self, query, vector=None, text=None):
         if query is not None and (vector is not None or text is not None):
@@ -2228,6 +2235,16 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
 
         return vector_query, text_query
 
+    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=
+        "all")` also sets for its own use — on a MemWAL table only a caller who
+        asked is refused, and the reranker still falls back to the primary key.
+        """
+        self._caller_requested_row_id = with_row_id
+        return super().with_row_id(with_row_id)
+
     def phrase_query(self, phrase_query: bool = True) -> LanceHybridQueryBuilder:
         """Set whether to use phrase query.
 
@@ -2248,16 +2265,30 @@ 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:
+        # 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
+        # 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:
-            fts_future = executor.submit(
-                self._fts_query.with_row_id(True).to_arrow, timeout=timeout
-            )
-            vector_future = executor.submit(
-                self._vector_query.with_row_id(True).to_arrow, timeout=timeout
-            )
+            fts_future = executor.submit(fts_query.to_arrow, timeout=timeout)
+            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,
@@ -2272,6 +2303,8 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
         return self._finish_hybrid_results(results)
 
     def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
+        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():
@@ -2445,7 +2478,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
         self._norm = normalize
         self._reranker = reranker
         if reranker.score == "all":
-            self.with_row_id(True)
+            # 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
 
@@ -2719,13 +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)
-        if self._columns:
-            self._vector_query.select(self._columns)
-            self._fts_query.select(self._columns)
+        # WAL-PK-FUSION: without the fallback, select `self._columns` as is.
+        self._pk_fusion = None
+        columns = self._columns
+        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:
+        # 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:
@@ -3866,6 +3913,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
         self._inner = inner
         self._norm = "score"
         self._reranker = RRFReranker()
+        # 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"
@@ -3897,7 +3947,7 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
 
         return self
 
-    def _create_child_queries(
+    async def _create_child_queries(
         self,
     ) -> Tuple["AsyncFTSQuery", "AsyncVectorQuery", int, int]:
         """Build the sub-queries that make up this hybrid query.
@@ -3906,7 +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 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)
@@ -3924,8 +3975,27 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
             limit = DEFAULT_HYBRID_LIMIT
         offset = fts_req.offset or vec_req.offset or 0
 
-        fts_query.with_row_id()
-        vec_query.with_row_id()
+        # 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:
+            # Ask for the key columns instead, and never for `_rowid`.
+            # `select` carries the whole projection; `select_source_columns`
+            # keeps only plain column references, so rebuilding from it would
+            # drop computed ones.
+            columns = self._pk_fusion.inject(fts_req.select)
+            if self._pk_fusion.injected:
+                fts_query.select(columns)
+                vec_query.select(columns)
 
         # offset() pushes the offset down into both sub-queries, which would make
         # each of them skip its own first `offset` rows. The window has to be
@@ -3944,7 +4014,23 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
         max_batch_length: Optional[int] = None,
         timeout: Optional[timedelta] = None,
     ) -> AsyncRecordBatchReader:
-        fts_query, vec_query, limit, offset = self._create_child_queries()
+        # 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
+            ),
+            self._table,
+            fused=lambda: self._pk_fusion is not None,
+            caller_requested_row_id=self._user_requested_row_id(),
+        )
+
+    async def _run_hybrid(
+        self,
+        *,
+        max_batch_length: Optional[int] = None,
+        timeout: Optional[timedelta] = None,
+    ) -> AsyncRecordBatchReader:
+        fts_query, vec_query, limit, offset = await self._create_child_queries()
 
         req = fts_query._inner.to_query_request()
         blob_auto_row_id = False
@@ -3968,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,
@@ -3979,7 +4069,9 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
             with_row_ids=True,
             offset=offset,
         )
-        if (
+        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
             and "_rowid" in result.column_names
@@ -4028,7 +4120,7 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
         plan : str
         """  # noqa: E501
 
-        fts_query, vec_query, _, _ = self._create_child_queries()
+        fts_query, vec_query, _, _ = await self._create_child_queries()
         vector_plan = await vec_query.explain_plan(verbose)
         fts_plan = await fts_query.explain_plan(verbose)
         # Indent sub-plans under the reranker
@@ -4057,7 +4149,7 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
         -------
         plan : str
         """
-        fts_query, vec_query, _, _ = self._create_child_queries()
+        fts_query, vec_query, _, _ = await self._create_child_queries()
 
         results = ["Vector Search Query:"]
         results.append(await vec_query.analyze_plan(distributed_metrics))

python/python/lancedb/remote/table.py

@@ -1036,6 +1036,14 @@ 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()
+
+    def _note_hybrid_pk_fusion(self) -> None:
+        self._table._note_hybrid_pk_fusion()
+
     def checkpoint_lsm(self) -> None:
         """Synchronous version of
         [`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""

python/python/lancedb/table.py

@@ -2647,6 +2647,18 @@ 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``.
+
+        Learned from a refusal, never probed, so asking is free. ``False`` for
+        table types that never refuse.
+        """
+        return False
+
+    def _note_hybrid_pk_fusion(self) -> None:
+        """Remember a ``_rowid`` refusal, so later hybrid queries skip it."""
+
 
 class LanceTable(Table):
     """
@@ -5294,6 +5306,15 @@ 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()
+
+    def _note_hybrid_pk_fusion(self) -> None:
+        """See [`Table._note_hybrid_pk_fusion`][lancedb.table.Table]."""
+        self._inner.note_hybrid_pk_fusion()
+
     async def checkpoint_lsm(self) -> None:
         """Converge this table's LSM write path into its base table.
 

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()

python/src/table.rs

@@ -1637,6 +1637,21 @@ 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())
+    }
+
+    /// Remember that this table refused `_rowid`, so later hybrid queries skip
+    /// straight to the primary-key fusion.
+    pub fn note_hybrid_pk_fusion(self_: PyRef<'_, Self>) -> PyResult<()> {
+        self_.inner_ref()?.base_table().note_hybrid_pk_fusion();
+        Ok(())
+    }
+
     /// Converge the table's LSM write path into its base table.
     ///
     /// Best-effort: with writes flowing, new rows may land after the last