From 25176778a32a421f9b10d2dfcd1aecded545d087 Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Wed, 23 Sep 2026 12:11:28 -0500 Subject: [PATCH] refactor(query): learn the hybrid fusion key from a refusal, not a probe Asking `get_lsm_write_spec` up front to decide how to fuse taxed every hybrid query, on every table, to learn something almost no table needs -- and that endpoint is not cheap: it opens the table and does a `checkout_latest()` manifest reload so a just-set spec is visible. Ask for `_rowid` optimistically instead and fall back to the primary key when the server refuses. Only a MemWAL table refuses; the refusal is raised before the query is planned, so it costs nothing server-side; and it had to be handled regardless, because a write spec installed elsewhere can arrive between any two queries. The proactive probe was belt and braces over a mechanism that already had to exist. A base table now sends nothing beyond its two legs, covered by a test that counts request paths. A MemWAL table pays one refused attempt, then remembers -- `BaseTable::hybrid_pk_fusion_learned` / `note_hybrid_pk_fusion`, seeded by the refusal and never fetched, so reading it is free. Deletes the `lsm_enabled` accessor from `BaseTable`, `Table`, the PyO3 and Python tables, and the type stubs, along with the cache that backed it. A caller who explicitly asked for `_rowid` still gets the reason rather than a bare 400: the refusal is translated into the same message the anticipated path raises. `return_score="all"` turns row ids on for the reranker rather than the caller, so it falls back like any other query -- an e2e run caught both of these. Co-Authored-By: Claude Opus 5 --- python/python/lancedb/_lancedb.pyi | 3 +- python/python/lancedb/query.py | 109 ++++-- python/python/lancedb/remote/table.py | 9 +- python/python/lancedb/table.py | 33 +- python/python/tests/test_hybrid_query.py | 7 +- python/src/table.rs | 18 +- rust/lancedb/src/query.rs | 79 +++- rust/lancedb/src/remote/table.rs | 475 ++++++++++------------- rust/lancedb/src/table.rs | 36 +- 9 files changed, 423 insertions(+), 346 deletions(-) diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 0e0eab619..da348268a 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -422,7 +422,8 @@ 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]: ... - async def lsm_enabled(self) -> bool: ... + 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: ... diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 1a9975fe2..57b0a4b0a 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -2208,6 +2208,29 @@ def _pk_columns(schema: pa.Schema) -> List[str]: 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. @@ -2352,6 +2375,28 @@ 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._user_requested_row_id() and not self._reranker_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) + + 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: @@ -2379,21 +2424,20 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): def _resolve_fusion_pk(self) -> Optional[List[str]]: """The key columns the two legs must be fused on, or None for `_rowid`. - A MemWAL table rejects `with_row_id` before it plans, because the fresh - tier has no stable row id, so there the legs are joined on a surrogate - built from the primary key instead. + `_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 PK mode: a table type that does not - # answer the question properly keeps today's `_rowid` behavior. - if self._use_lsm is False or self._table.lsm_enabled() is not True: + # 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._user_requested_row_id() and not self._reranker_requested_row_id: - raise NotImplementedError( - "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)" - ) + raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED) if self._blob_auto_row_id_enabled(): raise NotImplementedError( "hybrid search cannot project a blob column on a MemWAL table: " @@ -4061,25 +4105,20 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase): async def _resolve_fusion_pk(self) -> Optional[List[str]]: """The key columns the two legs must be fused on, or None for `_rowid`. - A MemWAL table rejects `with_row_id` before it plans, because the fresh - tier has no stable row id, so there the legs are joined on a surrogate - built from the primary key instead. + `_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 PK mode: a table type that does not - # answer the question properly keeps today's `_rowid` behavior. - if (await self._table.lsm_enabled()) is not True: + # 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( - "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)" - ) + raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED) schema = await self._table.schema() if blob_auto_row_id_for_scan( schema, @@ -4154,6 +4193,28 @@ 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( + 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 + ) + + 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() diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index d14980058..c3ee197f2 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -1023,9 +1023,12 @@ class RemoteTable(Table): """Read the installed LsmWriteSpec, or ``None``.""" return LOOP.run(self._table.get_lsm_write_spec()) - def lsm_enabled(self) -> bool: - """Whether reads route through the MemWAL. Cached by the client.""" - return LOOP.run(self._table.lsm_enabled()) + 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 diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 4f59355dc..39c493d1c 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -2467,14 +2467,17 @@ class Table(ABC): to check if the table is already using the new path style. """ - def lsm_enabled(self) -> bool: - """Whether reads on this table route through the MemWAL. + def _hybrid_pk_fusion_learned(self) -> bool: + """Whether a hybrid query here has already been refused ``_rowid``. - ``False`` for table types that have no MemWAL write path, which is what - lets query paths ask unconditionally. + 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): """ @@ -4370,10 +4373,12 @@ class LanceTable(Table): [`AsyncTable.get_lsm_write_spec`][lancedb.AsyncTable.get_lsm_write_spec].""" return LOOP.run(self._table.get_lsm_write_spec()) - def lsm_enabled(self) -> bool: - """Whether reads route through the MemWAL. See - [`AsyncTable.lsm_enabled`][lancedb.AsyncTable.lsm_enabled].""" - return LOOP.run(self._table.lsm_enabled()) + 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 @@ -5114,13 +5119,13 @@ class AsyncTable: """ return await self._inner.get_lsm_write_spec() - async def lsm_enabled(self) -> bool: - """Whether reads on this table route through the MemWAL. + def _hybrid_pk_fusion_learned(self) -> bool: + """See [`Table._hybrid_pk_fusion_learned`][lancedb.table.Table].""" + return self._inner.hybrid_pk_fusion_learned() - Equivalent to testing `get_lsm_write_spec` for ``None``, but cached by - the remote table, so query paths can ask on every call. - """ - return await self._inner.lsm_enabled() + 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. diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index 6faf80d54..fc909cfef 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -498,9 +498,10 @@ def test_normalize_scores(): def _wal_table(schema: pa.Schema) -> mock.Mock: - """A stand-in for a MemWAL-backed table: the fusion cannot use `_rowid`.""" + """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.lsm_enabled.return_value = True + table._hybrid_pk_fusion_learned.return_value = True table.schema = schema return table @@ -554,7 +555,7 @@ def test_wal_hybrid_without_a_projection_leaves_it_alone(): def test_base_table_hybrid_still_joins_on_row_ids(): table = mock.Mock() - table.lsm_enabled.return_value = False + table._hybrid_pk_fusion_learned.return_value = False builder = LanceHybridQueryBuilder(table).vector([0.1, 0.2]).text("puppy") builder._create_query_builders() diff --git a/python/src/table.rs b/python/src/table.rs index 928970d6f..8fb2ebc39 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -1481,13 +1481,17 @@ impl Table { }) } - /// Whether reads on this table route through the MemWAL. Cached by the - /// remote table, so the hybrid query path can ask on every call. - pub fn lsm_enabled(self_: PyRef<'_, Self>) -> PyResult> { - let inner = self_.inner_ref()?.clone(); - future_into_py(self_.py(), async move { - inner.lsm_enabled().await.infer_error() - }) + /// Whether a hybrid query on this table has already been refused `_rowid`. + /// Learned from a previous refusal, never probed, so this is free. + pub fn hybrid_pk_fusion_learned(self_: PyRef<'_, Self>) -> PyResult { + 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. diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index 29fc95f68..21ee5e555 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1176,6 +1176,35 @@ 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 @@ -1474,13 +1503,43 @@ 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, + } + } + + async fn run_hybrid( + &self, + key: FusionKey, + options: QueryExecutionOptions, ) -> Result { let max_batch_length = options.max_batch_length as usize; let internal_options = options.without_output_batch_length_limit(); - let key = self.fusion_key().await?; let mut fts_query = Query::new(self.parent.clone()); fts_query.request = self.request.base.clone(); @@ -1583,18 +1642,20 @@ impl VectorQuery { /// 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.lsm_enabled().await? { + 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(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(), - }); + return Err(wal_row_id_unsupported()); } let schema = self.parent.schema().await?; let columns = crate::table::primary_key::pk_columns(&schema); diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 6587dc718..53c8d9bc8 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -88,12 +88,12 @@ 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); -/// The MemWAL routing bit changes only when a write spec is installed or -/// removed, both of which invalidate this handle's cache directly. The TTL -/// bounds how long a change made through *another* handle goes unnoticed; a -/// stale `false` is caught by the `with_row_id` rejection on the query itself. -const WAL_ENABLED_CACHE_TTL: Duration = Duration::from_secs(30); -const WAL_ENABLED_CACHE_REFRESH_WINDOW: Duration = Duration::from_secs(5); +/// How long a hybrid query remembers that this table refused `_rowid`. Nothing +/// fetches this — it is seeded by the refusal itself — so the TTL only bounds +/// how long the memory outlives a write spec someone removed, and being late to +/// notice costs a base-only read that is still correct. +const WAL_FUSION_MEMORY_TTL: Duration = Duration::from_secs(30); +const WAL_FUSION_MEMORY_REFRESH_WINDOW: Duration = Duration::from_secs(5); const SCHEMA_SELECTOR_CHANGED: &str = "table selector changed while fetching schema"; fn fts_query_requires_document_granularity_support(query: &FtsQuery) -> bool { @@ -489,7 +489,7 @@ pub struct RemoteTable { version: Arc>>, location: RwLock>, schema_cache: BackgroundCache, - wal_enabled_cache: BackgroundCache, + wal_pk_fusion: BackgroundCache, 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 @@ -654,9 +654,9 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_enabled_cache: BackgroundCache::new( - WAL_ENABLED_CACHE_TTL, - WAL_ENABLED_CACHE_REFRESH_WINDOW, + wal_pk_fusion: BackgroundCache::new( + WAL_FUSION_MEMORY_TTL, + WAL_FUSION_MEMORY_REFRESH_WINDOW, ), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -691,9 +691,9 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), - wal_enabled_cache: BackgroundCache::new( - WAL_ENABLED_CACHE_TTL, - WAL_ENABLED_CACHE_REFRESH_WINDOW, + wal_pk_fusion: BackgroundCache::new( + WAL_FUSION_MEMORY_TTL, + WAL_FUSION_MEMORY_REFRESH_WINDOW, ), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch, @@ -1429,8 +1429,8 @@ impl RemoteTable { self.schema_cache.invalidate(); } - fn invalidate_wal_enabled_cache(&self) { - self.wal_enabled_cache.invalidate(); + fn invalidate_wal_pk_fusion(&self) { + self.wal_pk_fusion.invalidate(); } fn handle_error_invalidation(&self, error: &Error) { @@ -1602,9 +1602,9 @@ 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_enabled_cache: BackgroundCache::new( - WAL_ENABLED_CACHE_TTL, - WAL_ENABLED_CACHE_REFRESH_WINDOW, + wal_pk_fusion: BackgroundCache::new( + WAL_FUSION_MEMORY_TTL, + WAL_FUSION_MEMORY_REFRESH_WINDOW, ), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -1630,9 +1630,9 @@ 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_enabled_cache: BackgroundCache::new( - WAL_ENABLED_CACHE_TTL, - WAL_ENABLED_CACHE_REFRESH_WINDOW, + wal_pk_fusion: BackgroundCache::new( + WAL_FUSION_MEMORY_TTL, + WAL_FUSION_MEMORY_REFRESH_WINDOW, ), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -1667,9 +1667,9 @@ 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_enabled_cache: BackgroundCache::new( - WAL_ENABLED_CACHE_TTL, - WAL_ENABLED_CACHE_REFRESH_WINDOW, + wal_pk_fusion: BackgroundCache::new( + WAL_FUSION_MEMORY_TTL, + WAL_FUSION_MEMORY_REFRESH_WINDOW, ), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, @@ -3102,7 +3102,7 @@ impl BaseTable for RemoteTable { .json(&body); let (request_id, response) = self.send(request, true).await?; self.check_table_response(&request_id, response).await?; - self.invalidate_wal_enabled_cache(); + self.invalidate_wal_pk_fusion(); Ok(()) } @@ -3114,7 +3114,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_enabled_cache(); + self.invalidate_wal_pk_fusion(); Ok(()) } @@ -3181,40 +3181,12 @@ impl BaseTable for RemoteTable { Ok(Some(spec)) } - /// Cached because every hybrid query asks, and the answer is a routing - /// hint rather than a correctness gate: a wrong `false` surfaces as the - /// `with_row_id` rejection on the query itself, and a wrong `true` still - /// produces a valid base-only read. - /// - /// An unreachable endpoint therefore answers `false` instead of failing. - /// That hides nothing: a server without the route has no MemWAL tables, and - /// if the table itself is gone the query that follows reports it. Only a - /// success is cached, so a transient failure is re-asked on the next call. - /// - /// Reuses `get_lsm_write_spec` rather than `BackgroundCache::get`, whose - /// fetch closure must be `'static` and so cannot borrow `self`. The cost is - /// no single-flight on a cold cache — concurrent first calls each send one - /// request — and a refresh that blocks rather than happening in the - /// background. Both are one POST against a bit that changes almost never. - async fn lsm_enabled(&self) -> Result { - if let Some(enabled) = self.wal_enabled_cache.try_get() { - return Ok(enabled); - } - match self.get_lsm_write_spec().await { - Ok(spec) => { - let enabled = spec.is_some(); - self.wal_enabled_cache.seed(enabled); - Ok(enabled) - } - Err(e) => { - log::debug!( - "lsm_enabled: reading the write spec for table {} failed, \ - treating the table as base-only: {e}", - self.identifier - ); - Ok(false) - } - } + fn hybrid_pk_fusion_learned(&self) -> bool { + self.wal_pk_fusion.try_get().unwrap_or(false) + } + + fn note_hybrid_pk_fusion(&self) { + self.wal_pk_fusion.seed(true); } async fn tags(&self) -> Result> { @@ -13109,119 +13081,15 @@ mod tests { branch.stats().await.unwrap(); } - const WAL_SPEC_BODY: &str = r#"{"lsm_write_spec":{"sharding":{"mode":"unsharded"},"maintained_indexes":[],"writer_config_defaults":{}}}"#; - const NO_WAL_SPEC_BODY: &str = r#"{"lsm_write_spec":null}"#; + /// 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)"}"#; - /// One read per TTL, not one per query: hybrid asks on every call. - #[tokio::test] - async fn test_lsm_enabled_caches_the_write_spec_read() { - let calls = Arc::new(AtomicUsize::new(0)); - let seen = calls.clone(); - let table = Table::new_with_handler("my_table", move |request| { - assert_eq!( - request.url().path(), - "/v1/table/my_table/get_lsm_write_spec/" - ); - seen.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(200) - .body(WAL_SPEC_BODY) - .unwrap() - }); + use lance::dataset::ROW_ID; - assert!(table.base_table().lsm_enabled().await.unwrap()); - assert_eq!(calls.load(Ordering::SeqCst), 1); - assert!(table.base_table().lsm_enabled().await.unwrap()); - assert_eq!(calls.load(Ordering::SeqCst), 1); - - clock::advance_by(Duration::from_secs(31)); - assert!(table.base_table().lsm_enabled().await.unwrap()); - assert_eq!(calls.load(Ordering::SeqCst), 2); - } - - /// A null spec is the server saying "base-only", not an error. - #[tokio::test] - async fn test_lsm_enabled_is_false_without_a_write_spec() { - let table = Table::new_with_handler("my_table", move |_| { - http::Response::builder() - .status(200) - .body(NO_WAL_SPEC_BODY) - .unwrap() - }); - assert!(!table.base_table().lsm_enabled().await.unwrap()); - } - - /// Installing or removing a spec on this handle must not leave the routing - /// bit stale for a TTL. - #[rstest] - #[case("set_lsm_write_spec")] - #[case("unset_lsm_write_spec")] - #[tokio::test] - async fn test_lsm_enabled_invalidated_by_write_spec_change(#[case] operation: &str) { - let calls = Arc::new(AtomicUsize::new(0)); - let seen = calls.clone(); - let table = Table::new_with_handler("my_table", move |request| { - let body = if request.url().path().ends_with("/get_lsm_write_spec/") { - seen.fetch_add(1, Ordering::SeqCst); - WAL_SPEC_BODY - } else { - "{}" - }; - http::Response::builder().status(200).body(body).unwrap() - }); - - table.base_table().lsm_enabled().await.unwrap(); - assert_eq!(calls.load(Ordering::SeqCst), 1); - - match operation { - "set_lsm_write_spec" => table - .set_lsm_write_spec(LsmWriteSpec::unsharded()) - .await - .unwrap(), - "unset_lsm_write_spec" => table.unset_lsm_write_spec().await.unwrap(), - other => panic!("unexpected operation: {other}"), - } - - table.base_table().lsm_enabled().await.unwrap(); - assert_eq!(calls.load(Ordering::SeqCst), 2); - } - - /// A server with no such route has no MemWAL tables, so `false` is the - /// right answer rather than a failed query. The failure is not cached. - #[tokio::test] - async fn test_lsm_enabled_treats_an_unreachable_endpoint_as_base_only() { - let calls = Arc::new(AtomicUsize::new(0)); - let seen = calls.clone(); - let table = Table::new_with_handler("my_table", move |_| { - seen.fetch_add(1, Ordering::SeqCst); - http::Response::builder() - .status(404) - .body("not found") - .unwrap() - }); - - assert!(!table.base_table().lsm_enabled().await.unwrap()); - assert!(!table.base_table().lsm_enabled().await.unwrap()); - assert_eq!( - calls.load(Ordering::SeqCst), - 2, - "a failed read must not be cached" - ); - } - - /// The whole point of PK mode: neither leg may ask for `_rowid` (the server - /// rejects it on a MemWAL table before it plans), both must carry the key - /// so the surrogate can be built, and the key must not leak into the answer - /// when the caller did not select it. - #[tokio::test] - async fn test_hybrid_on_a_wal_table_joins_on_the_primary_key() { - use arrow_array::Float32Array; - use lance::dataset::ROW_ID; + fn pk_schema() -> Schema { use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION; - - let vector_type = - DataType::FixedSizeList(Arc::new(Field::new("item", DataType::Float32, true)), 2); - let table_schema = Schema::new(vec![ + Schema::new(vec![ Field::new("id", DataType::Utf8, false).with_metadata( [( LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(), @@ -13231,66 +13099,85 @@ mod tests { .collect(), ), Field::new("text", DataType::Utf8, true), - Field::new("vector", vector_type.clone(), true), - ]); + ]) + } - // One leg answers on `id` "a", the other on "a" and "b": the shared key - // must fuse to a single output row. - let leg = |score_column: &str, ids: Vec<&str>| { - let schema = Arc::new(Schema::new(vec![ - Field::new("id", DataType::Utf8, false), - Field::new("text", DataType::Utf8, true), - Field::new(score_column, DataType::Float32, false), - ])); - let scores = Float32Array::from(vec![1.0_f32; ids.len()]); - let texts = StringArray::from(ids.clone()); - RecordBatch::try_new( - schema, - vec![ - Arc::new(StringArray::from(ids)), - Arc::new(texts), - Arc::new(scores), - ], - ) - .unwrap() - }; + /// One leg's worth of results, keyed so the two legs share a row. + /// `row_ids` mirrors a server answering a query that asked for `_rowid`. + fn leg(score_column: &str, ids: Vec<&str>, row_ids: bool) -> RecordBatch { + use arrow_array::{Float32Array, UInt64Array}; + let mut fields = vec![ + Field::new("id", DataType::Utf8, false), + Field::new("text", DataType::Utf8, true), + Field::new(score_column, DataType::Float32, false), + ]; + let scores = Float32Array::from(vec![1.0_f32; ids.len()]); + let texts = StringArray::from(ids.clone()); + let mut columns: Vec = vec![ + Arc::new(StringArray::from(ids.clone())), + Arc::new(texts), + Arc::new(scores), + ]; + if row_ids { + fields.push(Field::new(ROW_ID, DataType::UInt64, false)); + columns.push(Arc::new(UInt64Array::from_iter_values( + ids.iter().map(|id| id.as_bytes()[0] as u64), + ))); + } + RecordBatch::try_new(Arc::new(Schema::new(fields)), columns).unwrap() + } - let bodies = Arc::new(std::sync::Mutex::new(Vec::::new())); - let seen = bodies.clone(); - let table = - Table::new_with_handler("my_table", move |request| match request.url().path() { - "/v1/table/my_table/describe/" => http::Response::builder() - .status(200) - .body(describe_response(&table_schema).into_bytes()) - .unwrap(), - "/v1/table/my_table/get_lsm_write_spec/" => http::Response::builder() - .status(200) - .body(WAL_SPEC_BODY.as_bytes().to_vec()) - .unwrap(), - "/v1/table/my_table/query/" => { - let body = request_body_json(&request); - let is_fts = body.get("full_text_query").is_some_and(|v| !v.is_null()); - seen.lock().unwrap().push(body); - let batch = if is_fts { - leg("_score", vec!["a", "b"]) - } else { - leg("_distance", vec!["a"]) - }; - http::Response::builder() - .status(200) - .body(write_ipc_file(&batch)) - .unwrap() + /// A server that refuses `_rowid` the way a MemWAL table does, recording + /// every query body so a test can see which key each attempt asked for. + fn wal_refusing_table(bodies: Arc>>) -> Table { + Table::new_with_handler("my_table", move |request| match request.url().path() { + "/v1/table/my_table/describe/" => http::Response::builder() + .status(200) + .body(describe_response(&pk_schema()).into_bytes()) + .unwrap(), + "/v1/table/my_table/query/" => { + let body = request_body_json(&request); + let asked_for_row_id = body["with_row_id"] == serde_json::Value::Bool(true); + let is_fts = body.get("full_text_query").is_some_and(|v| !v.is_null()); + bodies.lock().unwrap().push(body); + if asked_for_row_id { + return http::Response::builder() + .status(400) + .body(WAL_ROW_ID_REFUSAL_BODY.as_bytes().to_vec()) + .unwrap(); } - path => panic!("unexpected request path: {path}"), - }); + let batch = if is_fts { + leg("_score", vec!["a", "b"], false) + } else { + leg("_distance", vec!["a"], false) + }; + http::Response::builder() + .status(200) + .body(write_ipc_file(&batch)) + .unwrap() + } + path => panic!("unexpected request path: {path}"), + }) + } - let results = table + fn hybrid(table: &Table) -> crate::query::VectorQuery { + table .query() .full_text_search(FullTextSearchQuery::new("a".to_string())) - .select(Select::columns(&["text"])) - .limit(10) .nearest_to(&[0.0, 0.0]) .unwrap() + } + + /// The optimistic path: ask for `_rowid`, and when a MemWAL table refuses, + /// fall back to the primary key rather than failing. + #[tokio::test] + async fn test_hybrid_falls_back_to_the_primary_key_when_row_id_is_refused() { + let bodies = Arc::new(std::sync::Mutex::new(Vec::new())); + let table = wal_refusing_table(bodies.clone()); + + let results = hybrid(&table) + .select(Select::columns(&["text"])) + .limit(10) .execute() .await .unwrap() @@ -13299,74 +13186,124 @@ mod tests { .unwrap(); let bodies = bodies.lock().unwrap(); - assert_eq!(bodies.len(), 2, "one request per leg"); - for body in bodies.iter() { - assert_ne!( - body["with_row_id"], - serde_json::Value::Bool(true), - "a MemWAL table rejects with_row_id before it plans: {body}" - ); + let (asked_row_id, asked_pk): (Vec<_>, Vec<_>) = bodies + .iter() + .partition(|b| b["with_row_id"] == serde_json::Value::Bool(true)); + // The legs run concurrently under `try_join!`, which abandons the + // sibling as soon as one errors — so the refused attempt is one or two + // requests, never zero, and the retry always sends both. + assert!( + !asked_row_id.is_empty(), + "`_rowid` has to be tried before the fallback means anything" + ); + assert_eq!(asked_pk.len(), 2, "both legs are retried on the key"); + for body in asked_pk { let columns = body["columns"].as_array().expect("a column projection"); assert!( columns.iter().any(|c| c == "id"), - "the key has to be projected to build the surrogate: {body}" + "the retry has to project the key: {body}" ); } let batch = &results[0]; - assert!( - batch.column_by_name(ROW_ID).is_none(), - "the surrogate is an internal join key" - ); - assert!( - batch.column_by_name("id").is_none(), - "the caller selected only `text`, so the injected key must be dropped" - ); - assert_eq!(batch.num_rows(), 2, "`a` appears in both legs and fuses"); + assert!(batch.column_by_name(ROW_ID).is_none()); + assert!(batch.column_by_name("id").is_none(), "injected key leaked"); + assert_eq!(batch.num_rows(), 2, "`a` is in both legs and fuses"); } - /// A surrogate must never be handed back as if it were a row id, so an - /// explicit `with_row_id` is refused rather than quietly answered. + /// The refusal is paid once: the second query goes straight to the key. #[tokio::test] - async fn test_hybrid_on_a_wal_table_refuses_with_row_id() { - use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION; + async fn test_hybrid_remembers_the_refusal() { + let bodies = Arc::new(std::sync::Mutex::new(Vec::new())); + let table = wal_refusing_table(bodies.clone()); - let table_schema = Schema::new(vec![ - Field::new("id", DataType::Utf8, false).with_metadata( - [( - LANCE_UNENFORCED_PRIMARY_KEY_POSITION.to_string(), - "1".to_string(), - )] - .into_iter() - .collect(), - ), - Field::new("text", DataType::Utf8, true), - ]); + for _ in 0..2 { + hybrid(&table) + .limit(10) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + } + let refused = bodies + .lock() + .unwrap() + .iter() + .filter(|b| b["with_row_id"] == serde_json::Value::Bool(true)) + .count(); + assert!( + (1..=2).contains(&refused), + "only the first query should pay the refusal, saw {refused} attempts" + ); + } + + /// A table that never refuses must never pay for the fallback existing. + #[tokio::test] + async fn test_hybrid_on_a_base_table_sends_no_extra_requests() { + let paths = Arc::new(std::sync::Mutex::new(Vec::new())); + let seen = paths.clone(); let table = Table::new_with_handler("my_table", move |request| { - let body = match request.url().path() { - "/v1/table/my_table/describe/" => describe_response(&table_schema), - "/v1/table/my_table/get_lsm_write_spec/" => WAL_SPEC_BODY.to_string(), - path => panic!("the query must be refused before it is sent, got {path}"), + seen.lock().unwrap().push(request.url().path().to_string()); + let batch = if request_body_json(&request) + .get("full_text_query") + .is_some_and(|v| !v.is_null()) + { + leg("_score", vec!["a", "b"], true) + } else { + leg("_distance", vec!["a"], true) }; - http::Response::builder().status(200).body(body).unwrap() + http::Response::builder() + .status(200) + .body(write_ipc_file(&batch)) + .unwrap() }); - let err = table - .query() - .full_text_search(FullTextSearchQuery::new("a".to_string())) - .with_row_id() - .nearest_to(&[0.0, 0.0]) + hybrid(&table) + .limit(10) + .execute() + .await .unwrap() + .try_collect::>() + .await + .unwrap(); + + let paths = paths.lock().unwrap(); + assert!( + paths.iter().all(|p| p.ends_with("/query/")), + "a table that never refuses must not be probed, got {paths:?}" + ); + assert_eq!(paths.len(), 2, "one request per leg, got {paths:?}"); + } + + /// A caller who asked for `_rowid` gets the server's refusal, not a + /// surrogate dressed up as one. + #[tokio::test] + async fn test_hybrid_does_not_fall_back_when_the_caller_asked_for_row_id() { + let bodies = Arc::new(std::sync::Mutex::new(Vec::new())); + let table = wal_refusing_table(bodies.clone()); + + let err = hybrid(&table) + .with_row_id() + .limit(10) .execute() .await .err() - .expect("the query must be refused"); + .expect("the refusal must reach the caller"); - assert!( - matches!(err, Error::NotSupported { .. }), - "expected NotSupported, got {err:?}" - ); + // Translated, not passed through: the caller gets the reason and the + // way out, which a bare 400 does not carry. + assert!(matches!(err, Error::NotSupported { .. }), "{err:?}"); assert!(err.to_string().contains("use_lsm(false)"), "{err}"); + assert!( + bodies + .lock() + .unwrap() + .iter() + .all(|b| b["with_row_id"] == serde_json::Value::Bool(true)), + "no retry should have been attempted" + ); } } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 36893759c..5ae8de9e1 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -681,16 +681,29 @@ pub trait BaseTable: std::fmt::Display + std::fmt::Debug + Send + Sync { message: "get_lsm_write_spec is not supported on this table type".into(), }) } - /// Whether reads on this table route through the MemWAL. + /// Whether a hybrid query on this table has already been told it cannot + /// join its legs on `_rowid`. /// - /// Distinct from [`Self::get_lsm_write_spec`] only in that callers who need - /// the routing decision and not the spec itself can have it cached. + /// 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 + /// free by construction — an implementation may only answer from what a + /// previous query already learned. /// - /// The default is `false`: a table type with no MemWAL write path never - /// routes through one. - async fn lsm_enabled(&self) -> Result { - Ok(false) + /// The default is `false`, which keeps a table type that never refuses on + /// the `_rowid` path forever. + fn hybrid_pk_fusion_learned(&self) -> bool { + false } + + /// Record that this table refused `_rowid`, so later hybrid queries skip + /// straight to the primary-key fusion instead of paying the refusal again. + /// + /// Implementations should expire this the way they expire other table + /// metadata: a spec can be removed, after which `_rowid` works again and + /// the only cost of being late to notice is a base-only read that is still + /// correct. + fn note_hybrid_pk_fusion(&self) {} /// Seal every bucket's active memtable into L0. /// /// The default implementation returns `NotSupported`. @@ -1924,15 +1937,6 @@ impl Table { self.inner.get_lsm_write_spec().await } - /// Whether reads on this table route through the MemWAL. - /// - /// Cheaper to ask repeatedly than [`Self::get_lsm_write_spec`], which it is - /// otherwise equivalent to: implementations may cache the answer, since it - /// changes only when a write spec is installed or removed. - pub async fn lsm_enabled(&self) -> Result { - self.inner.lsm_enabled().await - } - /// Converge this table's LSM write path into its base table. /// /// One `flush` to seal every memtable into L0, then compaction triggers