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

refactor(query): learn the hybrid fusion key from a refusal, not a probe

Browse Source 
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 <noreply@anthropic.com>
This commit is contained in:
Daniel Rammer
2026-09-23 12:11:28 -05:00
co-authored by Claude Opus 5
parent 959645f1c7
commit 25176778a3
9 changed files with 423 additions and 346 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+2 -1
View File
@@ -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: ...
python/python/lancedb/query.py
+85 -24
View File
@@ -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()
python/python/lancedb/remote/table.py
+6 -3
View File
@@ -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
python/python/lancedb/table.py
+19 -14
View File
@@ -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.
python/python/tests/test_hybrid_query.py
+4 -3
View File
@@ -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()
python/src/table.rs
+11 -7
View File
@@ -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<Bound<'_, PyAny>> {
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<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.
rust/lancedb/src/query.rs
+70 -9
View File
@@ -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<SendableRecordBatchStream> {
let key = self.fusion_key().await?;
let optimistic = matches!(key, FusionKey::RowId);
match self.run_hybrid(key, options.clone()).await {
// A caller who asked for `_rowid` gets the reason, not a bare 400 —
// the same message the learned path would have raised up front.
Err(e) if optimistic && self.request.base.with_row_id && is_wal_row_id_refusal(&e) => {
Err(wal_row_id_unsupported())
}
Err(e) if optimistic && is_wal_row_id_refusal(&e) => {
// Learned, so later queries on this table skip straight to it.
self.parent.note_hybrid_pk_fusion();
let key = self.pk_fusion_key().await?;
self.run_hybrid(key, options).await
}
other => other,
}
}
async fn run_hybrid(
&self,
key: FusionKey,
options: QueryExecutionOptions,
) -> Result<SendableRecordBatchStream> {
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<FusionKey> {
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<FusionKey> {
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);
rust/lancedb/src/remote/table.rs
+206 -269
View File
@@ -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<S: HttpSend = Sender> {
version: Arc<RwLock<Option<u64>>>,
location: RwLock<Option<String>>,
schema_cache: BackgroundCache<SchemaRef, Error>,
wal_enabled_cache: BackgroundCache<bool, Error>,
wal_pk_fusion: BackgroundCache<bool, Error>,
freshness: Arc<Mutex<FreshnessState>>,
/// The branch this handle is scoped to, or `None` for the main branch.
/// Stamped onto every branch-accepting request so reads and writes resolve
@@ -654,9 +654,9 @@ impl<S: HttpSend> RemoteTable<S> {
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
wal_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<S: HttpSend> RemoteTable<S> {
version: Arc::new(RwLock::new(None)),
location: RwLock::new(None),
schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW),
wal_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<S: HttpSend> RemoteTable<S> {
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<S: HttpSend> BaseTable for RemoteTable<S> {
.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<S: HttpSend> BaseTable for RemoteTable<S> {
));
let (request_id, response) = self.send(request, true).await?;
self.check_table_response(&request_id, response).await?;
self.invalidate_wal_enabled_cache();
self.invalidate_wal_pk_fusion();
Ok(())
}
@@ -3181,40 +3181,12 @@ impl<S: HttpSend> BaseTable for RemoteTable<S> {
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<bool> {
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<Box<dyn Tags + '_>> {
@@ -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<arrow_array::ArrayRef> = 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::<serde_json::Value>::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<std::sync::Mutex<Vec<serde_json::Value>>>) -> 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::<Vec<_>>()
.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::<Vec<_>>()
.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"
);
}
}
rust/lancedb/src/table.rs
+20 -16
View File
@@ -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<bool> {
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<bool> {
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