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

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

Browse Source 
Hybrid search fuses its two legs on `_rowid`, so both sub-queries set
`with_row_id` -- which a MemWAL table rejects before it plans, because the
fresh tier has no stable row id. Every hybrid shape and every reranker
therefore 400s on a WAL-backed table.

Join on the primary key instead. Every WAL-managed table declares an
unenforced primary key and the WAL union already dedups each leg by it, so
the two legs share a key space by construction -- and because they are
separate round trips that can land on different snapshots, joining on the
key is strictly more correct than joining on row ids, which would emit the
same logical row twice at two versions.

The key is dictionary-encoded into a dense u64 written under the `_rowid`
name. The fusion only ever compares that column for equality -- dedup,
grouping, score restoration -- so any injective mapping serves, and writing
it under that name is what lets every reranker, third-party implementations
included, run unmodified. It never reaches the caller.

Routing is decided by `lsm_enabled()`, added to `BaseTable` with a `false`
default and implemented on the remote table over the existing
`get_lsm_write_spec` endpoint, cached beside the schema. No server change.

Co-Authored-By: Claude Opus 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_01UozJQ9kQxntWarST7hC971
This commit is contained in:
Daniel Rammer
2026-09-16 15:47:06 -05:00
co-authored by Claude Opus 5
parent bc4497b21a
commit 959645f1c7
11 changed files with 1163 additions and 26 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/python/lancedb/_lancedb.pyi
+1
View File
@@ -422,6 +422,7 @@ class Table:
async def set_lsm_write_spec(self, spec: LsmWriteSpec) -> None: ...
async def unset_lsm_write_spec(self) -> None: ...
async def get_lsm_write_spec(self) -> Optional[LsmWriteSpec]: ...
async def lsm_enabled(self) -> bool: ...
async def checkpoint_lsm(self) -> None: ...
async def flush_lsm(self) -> None: ...
async def compact_lsm(self) -> None: ...
python/python/lancedb/query.py
+238 -18
View File
@@ -2175,6 +2175,103 @@ class LanceEmptyQueryBuilder(LanceQueryBuilder):
raise NotImplementedError("Reranking is not yet supported.")
# Lance records the unenforced primary key as field metadata, which survives the
# `describe` round trip — so a remote table can resolve the key without opening
# the dataset. Keys are bytes on a pyarrow field.
_PK_MARKER = b"lance-schema:unenforced-primary-key"
_PK_POSITION_MARKER = b"lance-schema:unenforced-primary-key:position"
_TRUTHY = {"1", "true", "on", "yes", "y"}
def _pk_columns(schema: pa.Schema) -> List[str]:
"""The unenforced primary key columns of `schema`, in key order.
Mirrors Lance's ordering: fields with an explicit 1-based position first, in
position order, then fields carrying only the legacy boolean marker, in
declaration order.
"""
keyed = []
for index, field in enumerate(schema):
metadata = field.metadata or {}
position = metadata.get(_PK_POSITION_MARKER)
if position is not None:
try:
position = int(position)
except ValueError:
position = None
if position is None:
legacy = metadata.get(_PK_MARKER)
if legacy is None or legacy.decode().lower() not in _TRUTHY:
continue
position = 0
keyed.append((position == 0, position, index, field.name))
return [name for *_, name in sorted(keyed)]
def _with_pk_columns(columns, pk_columns: List[str]):
"""Add the key columns a projection is missing.
Returns the projection to use and the columns added, which the caller drops
from the output again so an injected key never reaches a caller who did not
select it.
"""
if not columns:
# Every non-system column, the key among them.
return columns, []
# `PyQueryRequest.select` reports a dynamic projection as (name, sql) pairs.
if isinstance(columns, (list, tuple)) and isinstance(columns[0], (list, tuple)):
columns = dict(columns)
missing = [c for c in pk_columns if c not in columns]
if isinstance(columns, dict):
return {**columns, **{c: c for c in missing}}, missing
return list(columns) + missing, missing
def _stamp_surrogate_row_ids(
tables: List[pa.Table], pk_columns: List[str]
) -> List[pa.Table]:
"""Give each table a `_rowid` column derived from its primary key.
The fusion only ever compares that column for equality — to deduplicate, to
group by in a reranker, and to restore the original scores — so on a MemWAL
table, where the server has no stable row id to hand out, any injective
mapping of the primary key does the same job. Ids are assigned in first-seen
order across `tables`, so passing them in fusion order (vector, then FTS)
preserves "first occurrence wins" and its tie break.
Writing it under the `_rowid` name is what lets every reranker, third-party
implementations included, run unmodified. The column never reaches the
caller.
"""
if not pk_columns:
raise ValueError(
"hybrid search on a MemWAL table needs an unenforced primary key to "
"deduplicate on, and this table declares none"
)
ids = {}
stamped = []
for table in tables:
if table.num_rows == 0:
# Both legs empty: the combined schema carries no key column, and no
# row that would need one.
row_ids = pa.array([], type=pa.uint64())
else:
missing = [c for c in pk_columns if c not in table.column_names]
if missing:
raise ValueError(
"hybrid search could not deduplicate: primary key "
f"column(s) {missing} missing from a result set with "
f"{table.num_rows} rows"
)
keys = zip(*(table.column(c).to_pylist() for c in pk_columns))
row_ids = pa.array(
[ids.setdefault(key, len(ids)) for key in keys], type=pa.uint64()
)
stamped.append(table.append_column("_rowid", row_ids))
return stamped
class LanceHybridQueryBuilder(LanceQueryBuilder):
"""
A query builder that performs hybrid vector and full text search.
@@ -2207,6 +2304,13 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._phrase_query = None
self._lower_bound = None
self._upper_bound = None
# Set by `_create_query_builders`: the key the legs are fused on when
# `_rowid` is unavailable, and the key columns added to the projection
# to make that possible.
self._fusion_pk = None
self._injected_pk = []
# `rerank(return_score="all")` turns `_with_row_id` on for its own use.
self._reranker_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):
@@ -2249,13 +2353,13 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
def to_arrow(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:
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()
@@ -2268,10 +2372,47 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
limit=self._limit,
with_row_ids=True,
offset=self._offset,
pk_columns=self._fusion_pk,
)
return self._finish_hybrid_results(results)
def _resolve_fusion_pk(self) -> Optional[List[str]]:
"""The key columns the two legs must be fused on, or None for `_rowid`.
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.
"""
# 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:
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)"
)
if self._blob_auto_row_id_enabled():
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"fetching blobs needs a real _rowid, and the fresh tier has no "
"stable one. Set use_lsm(False) to read the base table only "
"(results will exclude un-compacted MemWAL data)"
)
return _pk_columns(self._table.schema)
def _finish_hybrid_results(self, results: pa.Table) -> pa.Table:
if self._fusion_pk is not None:
# The surrogate is an internal join key, never an answer, and a
# caller who asked for `_rowid` was refused in `_resolve_fusion_pk`.
if "_rowid" in results.column_names:
results = results.drop(["_rowid"])
injected = [c for c in self._injected_pk if c in results.column_names]
if injected:
results = results.drop(injected)
return results
if self._user_requested_row_id():
return results
if self._blob_auto_row_id_enabled():
@@ -2290,7 +2431,14 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
limit: int,
with_row_ids: bool,
offset: Optional[int] = None,
pk_columns: Optional[List[str]] = None,
) -> pa.Table:
if pk_columns is not None:
# Vector first: `merge_results` concatenates in that order and keeps
# the first occurrence, so first-seen ids preserve its tie break.
vector_results, fts_results = _stamp_surrogate_row_ids(
[vector_results, fts_results], pk_columns
)
if norm == "rank":
vector_results = LanceHybridQueryBuilder._rank(vector_results, "_distance")
fts_results = LanceHybridQueryBuilder._rank(fts_results, "_score")
@@ -2445,6 +2593,10 @@ class LanceHybridQueryBuilder(LanceQueryBuilder):
self._norm = normalize
self._reranker = reranker
if reranker.score == "all":
# The reranker needs the join column to carry the per-leg scores
# across; the caller did not ask for row ids and must not be
# refused on a MemWAL table because of it.
self._reranker_requested_row_id = True
self.with_row_id(True)
return self
@@ -2719,13 +2871,17 @@ 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)
self._fusion_pk = self._resolve_fusion_pk()
columns = self._columns
if self._fusion_pk is not None:
columns, self._injected_pk = _with_pk_columns(columns, self._fusion_pk)
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:
if self._with_row_id and self._fusion_pk 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 +4022,11 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
self._inner = inner
self._norm = "score"
self._reranker = RRFReranker()
# Set by `_create_child_queries`: the key the legs are fused on when
# `_rowid` is unavailable, and the key columns added to the projection
# to make that possible.
self._fusion_pk = None
self._injected_pk = []
def rerank(
self, reranker: Reranker = RRFReranker(), normalize: str = "score"
@@ -3897,7 +4058,43 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
return self
def _create_child_queries(
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.
"""
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:
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)"
)
schema = await self._table.schema()
if blob_auto_row_id_for_scan(
schema,
_query_request_projection(self._inner.to_query_request()),
with_row_id=self._with_row_id,
):
raise NotImplementedError(
"hybrid search cannot project a blob column on a MemWAL table: "
"fetching blobs needs a real _rowid, and the fresh tier has no "
"stable one. Set use_lsm(False) to read the base table only "
"(results will exclude un-compacted MemWAL data)"
)
return _pk_columns(schema)
async def _create_child_queries(
self,
) -> Tuple["AsyncFTSQuery", "AsyncVectorQuery", int, int]:
"""Build the sub-queries that make up this hybrid query.
@@ -3906,7 +4103,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 is left on
`self` — see `_resolve_fusion_pk`.
"""
fts_query = AsyncFTSQuery(self._inner.to_fts_query(), self._table)
vec_query = AsyncVectorQuery(self._inner.to_vector_query(), self._table)
@@ -3924,8 +4122,21 @@ 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()
self._fusion_pk = await self._resolve_fusion_pk()
if self._fusion_pk 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._injected_pk = _with_pk_columns(
fts_req.select, self._fusion_pk
)
if self._injected_pk:
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 +4155,7 @@ 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()
fts_query, vec_query, limit, offset = await self._create_child_queries()
req = fts_query._inner.to_query_request()
blob_auto_row_id = False
@@ -3978,8 +4189,17 @@ class AsyncHybridQuery(AsyncStandardQuery, AsyncVectorQueryBase):
limit=limit,
with_row_ids=True,
offset=offset,
pk_columns=self._fusion_pk,
)
if (
if self._fusion_pk is not None:
# The surrogate is an internal join key, never an answer, and a
# caller who asked for `_rowid` was refused in `_resolve_fusion_pk`.
if "_rowid" in result.column_names:
result = result.drop(["_rowid"])
injected = [c for c in self._injected_pk if c in result.column_names]
if injected:
result = result.drop(injected)
elif (
not self._user_requested_row_id()
and not blob_auto_row_id
and "_rowid" in result.column_names
@@ -4028,7 +4248,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 +4277,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
+4
View File
@@ -1023,6 +1023,10 @@ 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 checkpoint_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
python/python/lancedb/table.py
+21
View File
@@ -2467,6 +2467,14 @@ 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.
``False`` for table types that have no MemWAL write path, which is what
lets query paths ask unconditionally.
"""
return False
class LanceTable(Table):
"""
@@ -4362,6 +4370,11 @@ 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 checkpoint_lsm(self) -> None:
"""Synchronous version of
[`AsyncTable.checkpoint_lsm`][lancedb.AsyncTable.checkpoint_lsm]."""
@@ -5101,6 +5114,14 @@ 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.
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()
async def checkpoint_lsm(self) -> None:
"""Converge this table's LSM write path into its base table.
python/python/tests/test_hybrid_query.py
+112
View File
@@ -495,3 +495,115 @@ def test_normalize_scores():
assert pc.equal(result, expected), (
f"Expected {expected} but got {result} for invert={invert}"
)
def _wal_table(schema: pa.Schema) -> mock.Mock:
"""A stand-in for a MemWAL-backed table: the fusion cannot use `_rowid`."""
table = mock.Mock()
table.lsm_enabled.return_value = True
table.schema = schema
return table
_PK_SCHEMA = pa.schema(
[
pa.field(
"id",
pa.string(),
metadata={b"lance-schema:unenforced-primary-key:position": b"1"},
),
pa.field("text", pa.string()),
]
)
def _wal_hybrid(**select):
builder = (
LanceHybridQueryBuilder(_wal_table(_PK_SCHEMA)).vector([0.1, 0.2]).text("puppy")
)
if "columns" in select:
builder = builder.select(select["columns"])
return builder
def test_wal_hybrid_projects_the_key_instead_of_asking_for_row_ids():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
assert builder._fusion_pk == ["id"]
assert builder._injected_pk == ["id"]
for leg in (builder._vector_query, builder._fts_query):
assert leg._columns == ["text", "id"]
# A MemWAL table rejects with_row_id before it plans.
assert leg._with_row_id is not True
def test_wal_hybrid_selecting_the_key_injects_nothing():
builder = _wal_hybrid(columns=["text", "id"])
builder._create_query_builders()
assert builder._injected_pk == []
def test_wal_hybrid_without_a_projection_leaves_it_alone():
builder = _wal_hybrid()
builder._create_query_builders()
# Select-all already returns the key.
assert builder._injected_pk == []
assert builder._vector_query._columns is None
def test_base_table_hybrid_still_joins_on_row_ids():
table = mock.Mock()
table.lsm_enabled.return_value = False
builder = LanceHybridQueryBuilder(table).vector([0.1, 0.2]).text("puppy")
builder._create_query_builders()
# `to_arrow` is what turns row ids on for this path; all that matters here
# is that no key was resolved and no projection was touched.
assert builder._fusion_pk is None
assert builder._injected_pk == []
def test_wal_hybrid_with_use_lsm_false_joins_on_row_ids():
builder = _wal_hybrid().use_lsm(False)
builder._create_query_builders()
assert builder._fusion_pk is None
def test_wal_hybrid_refuses_an_explicit_with_row_id():
builder = _wal_hybrid().with_row_id(True)
with pytest.raises(NotImplementedError, match="use_lsm"):
builder._create_query_builders()
def test_wal_hybrid_return_score_all_is_not_refused():
"""`return_score="all"` turns row ids on for the reranker, not the caller."""
builder = _wal_hybrid().rerank(RRFReranker(return_score="all"))
builder._create_query_builders()
assert builder._fusion_pk == ["id"]
def test_wal_hybrid_fuses_on_the_key_and_drops_it_again():
builder = _wal_hybrid(columns=["text"])
builder._create_query_builders()
vector_results = pa.table(
{"text": ["a", "b"], "id": ["a", "b"], "_distance": [0.1, 0.2]}
)
fts_results = pa.table({"text": ["b", "c"], "id": ["b", "c"], "_score": [1.0, 2.0]})
combined = LanceHybridQueryBuilder._combine_hybrid_results(
fts_results=fts_results,
vector_results=vector_results,
norm="score",
fts_query="puppy",
reranker=RRFReranker(),
limit=10,
with_row_ids=True,
pk_columns=builder._fusion_pk,
)
results = builder._finish_hybrid_results(combined)
assert "_rowid" not in results.column_names, "the surrogate is internal"
assert "id" not in results.column_names, "the caller selected only `text`"
assert sorted(results.column("text").to_pylist()) == ["a", "b", "c"]
python/src/table.rs
+9
View File
@@ -1481,6 +1481,15 @@ 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()
})
}
/// Converge the table's LSM write path into its base table.
///
/// Best-effort: with writes flowing, new rows may land after the last
rust/lancedb/src/query.rs
+126 -6
View File
@@ -1176,6 +1176,72 @@ pub struct VectorQuery {
request: VectorQueryRequest,
}
/// How the hybrid fusion identifies the same row across its two legs.
enum FusionKey {
/// Lance's `_rowid`, requested from the server. Every table that can supply
/// one uses this.
RowId,
/// A surrogate derived client-side from the primary key. MemWAL tables take
/// this path: the fresh tier has no stable row id, so the server rejects
/// `with_row_id` outright.
PrimaryKey {
columns: Vec<String>,
/// Key columns the caller did not ask for, added to the projection so
/// the surrogate can be built and dropped again before returning.
injected: Vec<String>,
},
}
/// The key columns `select` does not already produce, and so would have to be
/// added to it for the surrogate to be buildable.
fn pk_columns_to_inject(select: &Select, pk_columns: &[String]) -> Vec<String> {
let produced: Vec<&String> = match select {
// Already every non-system column, the key among them.
Select::All => return Vec::new(),
Select::Columns(columns) => columns.iter().collect(),
Select::Dynamic(pairs) => pairs.iter().map(|(name, _)| name).collect(),
Select::Expr(pairs) => pairs.iter().map(|(name, _)| name).collect(),
};
pk_columns
.iter()
.filter(|pk| !produced.contains(pk))
.cloned()
.collect()
}
/// Add the key columns `select` is missing, as identity projections where the
/// selection is expression-shaped.
fn inject_pk_columns(select: &mut Select, pk_columns: &[String]) {
let missing = pk_columns_to_inject(select, pk_columns);
if missing.is_empty() {
return;
}
match select {
Select::All => {}
Select::Columns(columns) => columns.extend(missing),
Select::Dynamic(pairs) => pairs.extend(missing.into_iter().map(|pk| (pk.clone(), pk))),
Select::Expr(pairs) => {
pairs.extend(missing.into_iter().map(|pk| {
let expr = crate::expr::col(&pk);
(pk, expr)
}));
}
}
}
/// Drop `columns` from `batch`, ignoring any that are not present.
fn drop_columns(batch: RecordBatch, columns: &[String]) -> Result<RecordBatch> {
let keep: Vec<usize> = batch
.schema()
.fields()
.iter()
.enumerate()
.filter(|(_, field)| !columns.iter().any(|name| name == field.name()))
.map(|(idx, _)| idx)
.collect();
Ok(batch.project(&keep)?)
}
impl VectorQuery {
fn new(base: Query) -> Self {
Self {
@@ -1414,12 +1480,24 @@ impl VectorQuery {
) -> Result<SendableRecordBatchStream> {
let max_batch_length = options.max_batch_length as usize;
let internal_options = options.without_output_batch_length_limit();
// clone query and specify we want to include row IDs, which can be needed for reranking
let key = self.fusion_key().await?;
let mut fts_query = Query::new(self.parent.clone());
fts_query.request = self.request.base.clone();
fts_query = fts_query.with_row_id();
let mut vector_query = self.clone().with_row_id();
let mut vector_query = self.clone();
match &key {
// The legs need the join column, which the reranking needs in turn.
FusionKey::RowId => {
fts_query = fts_query.with_row_id();
vector_query = vector_query.with_row_id();
}
// Ask for the key columns instead, and never for `_rowid`: the
// server rejects it on a MemWAL table before it plans.
FusionKey::PrimaryKey { columns, .. } => {
inject_pk_columns(&mut fts_query.request.select, columns);
inject_pk_columns(&mut vector_query.request.base.select, columns);
}
}
vector_query.request.base.full_text_search = None;
let (fts_results, vec_results) = try_join!(
@@ -1440,6 +1518,16 @@ impl VectorQuery {
let mut fts_results = concat_batches(&fts_schema, fts_results.iter())?;
let mut vec_results = concat_batches(&vec_schema, vec_results.iter())?;
// Vector first: `merge_results` concatenates in that order and keeps the
// first occurrence, so first-seen ids preserve its tie break.
if let FusionKey::PrimaryKey { columns, .. } = &key {
let mut legs = [vec_results, fts_results];
hybrid::stamp_surrogate_row_ids(&mut legs, columns)?;
let [vec, fts] = legs;
vec_results = vec;
fts_results = fts;
}
if matches!(self.request.base.norm, Some(NormalizeMethod::Rank)) {
vec_results = hybrid::rank(vec_results, DIST_COL, None)?;
fts_results = hybrid::rank(fts_results, SCORE_COL, None)?;
@@ -1475,13 +1563,45 @@ impl VectorQuery {
results = results.slice(0, limit);
}
if !self.request.base.with_row_id {
results = results.drop_column(ROW_ID)?;
match &key {
FusionKey::RowId => {
if !self.request.base.with_row_id {
results = results.drop_column(ROW_ID)?;
}
}
// The surrogate is an internal join key, never an answer: it goes
// whether or not the caller asked for `_rowid`, and `fusion_key`
// has already refused the query if they did.
FusionKey::PrimaryKey { injected, .. } => {
results = results.drop_column(ROW_ID)?;
results = drop_columns(results, injected)?;
}
}
Ok(single_batch_stream(results, max_batch_length))
}
/// Whether this hybrid query can join its legs on `_rowid`, or has to build
/// a surrogate from the primary key because the table is MemWAL-backed.
async fn fusion_key(&self) -> Result<FusionKey> {
if self.request.base.use_lsm == Some(false) || !self.parent.lsm_enabled().await? {
return Ok(FusionKey::RowId);
}
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(),
});
}
let schema = self.parent.schema().await?;
let columns = crate::table::primary_key::pk_columns(&schema);
let injected = pk_columns_to_inject(&self.request.base.select, &columns);
Ok(FusionKey::PrimaryKey { columns, injected })
}
async fn inner_execute_with_options(
&self,
options: QueryExecutionOptions,
rust/lancedb/src/query/hybrid.rs
+192 -1
View File
@@ -5,10 +5,13 @@ use arrow::compute::{
kernels::numeric::{div, sub},
max, min,
};
use arrow_array::{Float32Array, RecordBatch, cast::downcast_array};
use arrow::row::{RowConverter, SortField};
use arrow_array::{Float32Array, RecordBatch, UInt64Array, cast::downcast_array};
use arrow_schema::{DataType, Field, Schema, SortOptions};
use lance::dataset::ROW_ID;
use lance_arrow::RecordBatchExt;
use lance_index::{scalar::inverted::SCORE_COL, vector::DIST_COL};
use std::collections::HashMap;
use std::sync::Arc;
use crate::error::{Error, Result};
@@ -173,6 +176,83 @@ pub fn normalize_scores(
Ok(results)
}
/// Replace each batch's join column with a dense surrogate derived from
/// `pk_columns`, in place and in the order given.
///
/// The fusion joins its two legs on [`ROW_ID`], which a MemWAL table cannot
/// supply — the fresh tier has no stable row id. It only ever compares the
/// column for equality, though (`merge_results` dedups, the rerankers group and
/// the score restore looks up), so any injective mapping of the primary key
/// serves. Ids are assigned in first-seen order across `batches`, so passing
/// them in fusion order keeps the dedup's "first occurrence wins" and its tie
/// break byte-identical to what a real row id would have produced.
///
/// Writing it under the [`ROW_ID`] name is what lets every reranker, including
/// third-party implementations that hardcode the name, run unmodified. The
/// column never reaches the caller: the hybrid query drops it before returning.
pub fn stamp_surrogate_row_ids(batches: &mut [RecordBatch], pk_columns: &[String]) -> Result<()> {
if pk_columns.is_empty() {
return Err(Error::InvalidInput {
message: "hybrid search on a MemWAL table needs an unenforced primary key to \
deduplicate on, and this table declares none"
.to_string(),
});
}
let mut converter: Option<RowConverter> = None;
let mut ids: HashMap<Vec<u8>, u64> = HashMap::new();
for batch in batches.iter_mut() {
// Both legs empty: `query_schemas` synthesizes a schema carrying only
// the score and join columns, so there is no key to read — and no row
// that would need one.
let row_ids: UInt64Array = if batch.num_rows() == 0 {
UInt64Array::from(Vec::<u64>::new())
} else {
let key_columns = pk_columns
.iter()
.map(|name| {
batch
.column_by_name(name)
.cloned()
.ok_or_else(|| Error::InvalidInput {
message: format!(
"hybrid search could not deduplicate: primary key column {} \
is missing from a result set with {} rows",
name,
batch.num_rows()
),
})
})
.collect::<Result<Vec<_>>>()?;
let converter = match converter {
Some(ref converter) => converter,
None => converter.insert(RowConverter::new(
key_columns
.iter()
.map(|column| SortField::new(column.data_type().clone()))
.collect(),
)?),
};
let rows = converter.convert_columns(&key_columns)?;
let next = &mut ids;
UInt64Array::from_iter_values(rows.iter().map(|row| {
let id = next.len() as u64;
*next.entry(row.as_ref().to_vec()).or_insert(id)
}))
};
*batch = batch.try_with_column(
Field::new(ROW_ID, DataType::UInt64, false),
Arc::new(row_ids),
)?;
}
Ok(())
}
#[cfg(test)]
mod test {
use super::*;
@@ -346,4 +426,115 @@ mod test {
vec![0.0, 0.0, 0.0, 0.0, 0.0]
);
}
fn batch(ids: Vec<&str>, score: &str) -> RecordBatch {
let schema = Arc::new(Schema::new(vec![
Arc::new(Field::new("id", DataType::Utf8, false)),
Arc::new(Field::new(score, DataType::Float32, false)),
]));
let scores = Float32Array::from(vec![0.0_f32; ids.len()]);
RecordBatch::try_new(
schema,
vec![Arc::new(StringArray::from(ids)), Arc::new(scores)],
)
.unwrap()
}
fn row_ids(batch: &RecordBatch) -> Vec<u64> {
let ids: UInt64Array = downcast_array(batch.column_by_name(ROW_ID).unwrap());
ids.values().to_vec()
}
/// The same key in both legs has to land on the same id or the fusion will
/// not dedup, and ids must ascend in fusion order so the tie break holds.
#[test]
fn test_surrogate_row_ids_are_shared_across_legs_in_first_seen_order() {
let mut batches = vec![
batch(vec!["a", "b", "c"], DIST_COL),
batch(vec!["b", "d"], SCORE_COL),
];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), vec![0, 1, 2]);
assert_eq!(
row_ids(&batches[1]),
vec![1, 3],
"b keeps the id the vector leg gave it"
);
}
/// A string key is the common case and cannot be cast to the u64 the
/// rerankers read, which is the whole reason for the surrogate.
#[test]
fn test_surrogate_row_ids_repeat_within_a_leg() {
let mut batches = vec![batch(vec!["a", "a", "b"], DIST_COL)];
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), vec![0, 0, 1]);
}
#[test]
fn test_surrogate_row_ids_support_a_composite_key() {
let schema = Arc::new(Schema::new(vec![
Arc::new(Field::new("tenant", DataType::Utf8, false)),
Arc::new(Field::new("id", DataType::Int32, false)),
]));
let make = |tenants: Vec<&str>, ids: Vec<i32>| {
RecordBatch::try_new(
schema.clone(),
vec![
Arc::new(StringArray::from(tenants)),
Arc::new(arrow_array::Int32Array::from(ids)),
],
)
.unwrap()
};
let mut batches = vec![make(vec!["x", "y"], vec![1, 1]), make(vec!["y"], vec![1])];
stamp_surrogate_row_ids(&mut batches, &["tenant".to_string(), "id".to_string()]).unwrap();
assert_eq!(
row_ids(&batches[0]),
vec![0, 1],
"same id, different tenant"
);
assert_eq!(row_ids(&batches[1]), vec![1]);
}
/// Both legs empty: `query_schemas` synthesizes a schema with no key column,
/// and there is no row that would need one.
#[test]
fn test_surrogate_row_ids_tolerate_the_both_empty_schema() {
let mut batches = vec![
RecordBatch::new_empty(Arc::new(empty_vec_schema())),
RecordBatch::new_empty(Arc::new(empty_fts_schema())),
];
// The synthesized schemas already carry ROW_ID, so stamping would
// duplicate the field name; drop it first the way the caller does.
for b in batches.iter_mut() {
let keep: Vec<_> = b
.schema()
.fields()
.iter()
.filter(|f| f.name() != ROW_ID)
.map(|f| b.schema().index_of(f.name()).unwrap())
.collect();
*b = b.project(&keep).unwrap();
}
stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap();
assert_eq!(row_ids(&batches[0]), Vec::<u64>::new());
}
/// Rows with no key would silently fuse into one bucket, so refuse.
#[test]
fn test_surrogate_row_ids_reject_a_missing_key_column() {
let mut batches = vec![batch(vec!["a"], DIST_COL)];
let err = stamp_surrogate_row_ids(&mut batches, &["nope".to_string()]).unwrap_err();
assert!(err.to_string().contains("nope"), "{err}");
}
#[test]
fn test_surrogate_row_ids_reject_a_table_with_no_primary_key() {
let mut batches = vec![batch(vec!["a"], DIST_COL)];
let err = stamp_surrogate_row_ids(&mut batches, &[]).unwrap_err();
assert!(err.to_string().contains("unenforced primary key"), "{err}");
}
}
rust/lancedb/src/remote/table.rs
+330
View File
@@ -88,6 +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);
const SCHEMA_SELECTOR_CHANGED: &str = "table selector changed while fetching schema";
fn fts_query_requires_document_granularity_support(query: &FtsQuery) -> bool {
@@ -483,6 +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>,
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
@@ -647,6 +654,10 @@ 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,
),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -680,6 +691,10 @@ 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,
),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch,
}
@@ -1414,6 +1429,10 @@ impl<S: HttpSend> RemoteTable<S> {
self.schema_cache.invalidate();
}
fn invalidate_wal_enabled_cache(&self) {
self.wal_enabled_cache.invalidate();
}
fn handle_error_invalidation(&self, error: &Error) {
let status_code = match error {
Error::Http { status_code, .. } => *status_code,
@@ -1583,6 +1602,10 @@ 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,
),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -1607,6 +1630,10 @@ 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,
),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -1640,6 +1667,10 @@ 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,
),
freshness: Arc::new(Mutex::new(FreshnessState::default())),
branch: None,
}
@@ -3071,6 +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();
Ok(())
}
@@ -3082,6 +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();
Ok(())
}
@@ -3148,6 +3181,42 @@ 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)
}
}
}
async fn tags(&self) -> Result<Box<dyn Tags + '_>> {
Ok(Box::new(RemoteTags { inner: self }))
}
@@ -13039,4 +13108,265 @@ mod tests {
.unwrap();
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}"#;
/// 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()
});
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;
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![
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),
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()
};
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()
}
path => panic!("unexpected request path: {path}"),
});
let results = table
.query()
.full_text_search(FullTextSearchQuery::new("a".to_string()))
.select(Select::columns(&["text"]))
.limit(10)
.nearest_to(&[0.0, 0.0])
.unwrap()
.execute()
.await
.unwrap()
.try_collect::<Vec<_>>()
.await
.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 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}"
);
}
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");
}
/// 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.
#[tokio::test]
async fn test_hybrid_on_a_wal_table_refuses_with_row_id() {
use lance_core::datatypes::LANCE_UNENFORCED_PRIMARY_KEY_POSITION;
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),
]);
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}"),
};
http::Response::builder().status(200).body(body).unwrap()
});
let err = table
.query()
.full_text_search(FullTextSearchQuery::new("a".to_string()))
.with_row_id()
.nearest_to(&[0.0, 0.0])
.unwrap()
.execute()
.await
.err()
.expect("the query must be refused");
assert!(
matches!(err, Error::NotSupported { .. }),
"expected NotSupported, got {err:?}"
);
assert!(err.to_string().contains("use_lsm(false)"), "{err}");
}
}
rust/lancedb/src/table.rs
+20 -1
View File
@@ -78,7 +78,7 @@ pub mod delete;
pub mod lsm_stats;
pub mod merge;
pub mod optimize;
mod primary_key;
pub(crate) mod primary_key;
pub mod query;
pub mod refresh;
pub mod schema_evolution;
@@ -681,6 +681,16 @@ 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.
///
/// 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.
///
/// 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)
}
/// Seal every bucket's active memtable into L0.
///
/// The default implementation returns `NotSupported`.
@@ -1914,6 +1924,15 @@ 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
rust/lancedb/src/table/primary_key.rs
+110
View File
@@ -13,6 +13,7 @@
use arrow_schema::DataType;
use lance_core::datatypes::{LANCE_UNENFORCED_PRIMARY_KEY, LANCE_UNENFORCED_PRIMARY_KEY_POSITION};
use lance_core::utils::parse::str_is_truthy;
use crate::error::{Error, Result};
use crate::table::NativeTable;
@@ -111,3 +112,112 @@ fn is_supported_pk_dtype(dtype: &DataType) -> bool {
| DataType::FixedSizeBinary(_)
)
}
/// The unenforced primary key columns declared in `schema`, in key order.
///
/// Reads the field metadata directly rather than going through
/// [`lance_core::datatypes::Schema::unenforced_primary_key`], because a remote
/// table only ever has the Arrow schema `describe` returned — there is no Lance
/// schema, and no field ids, on that side of the wire.
///
/// Ordering mirrors Lance's: fields with an explicit 1-based position first, in
/// position order, then fields carrying only the legacy boolean marker, in
/// declaration order. Only top-level fields are considered, which is every
/// field a supported key dtype can live on.
pub fn pk_columns(schema: &arrow_schema::Schema) -> Vec<String> {
let mut keyed: Vec<(bool, u32, usize, &str)> = schema
.fields()
.iter()
.enumerate()
.filter_map(|(idx, field)| {
let metadata = field.metadata();
let position = metadata
.get(LANCE_UNENFORCED_PRIMARY_KEY_POSITION)
.and_then(|value| value.parse::<u32>().ok())
.or_else(|| {
metadata
.get(LANCE_UNENFORCED_PRIMARY_KEY)
.filter(|value| str_is_truthy(value))
.map(|_| 0)
})?;
Some((position == 0, position, idx, field.name().as_str()))
})
.collect();
keyed.sort_unstable();
keyed
.into_iter()
.map(|(_, _, _, name)| name.to_string())
.collect()
}
#[cfg(test)]
mod tests {
use std::collections::HashMap;
use arrow_schema::{Field, Schema};
use super::*;
fn field(name: &str, metadata: &[(&str, &str)]) -> Field {
Field::new(name, DataType::Utf8, false).with_metadata(
metadata
.iter()
.map(|(k, v)| (k.to_string(), v.to_string()))
.collect::<HashMap<_, _>>(),
)
}
#[test]
fn pk_columns_reads_the_position_marker() {
let schema = Schema::new(vec![
field("other", &[]),
field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]),
]);
assert_eq!(pk_columns(&schema), vec!["id".to_string()]);
}
/// Key order is the position, not the order the columns are declared in.
#[test]
fn pk_columns_orders_a_composite_key_by_position() {
let schema = Schema::new(vec![
field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "2")]),
field("tenant", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]),
]);
assert_eq!(
pk_columns(&schema),
vec!["tenant".to_string(), "id".to_string()]
);
}
/// Datasets written before the position marker carry only the boolean.
#[test]
fn pk_columns_accepts_the_legacy_boolean_marker() {
let schema = Schema::new(vec![field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY, "true")])]);
assert_eq!(pk_columns(&schema), vec!["id".to_string()]);
}
/// Mirrors Lance: an explicit position sorts ahead of a legacy field.
#[test]
fn pk_columns_sorts_positioned_fields_ahead_of_legacy_ones() {
let schema = Schema::new(vec![
field("legacy", &[(LANCE_UNENFORCED_PRIMARY_KEY, "yes")]),
field(
"positioned",
&[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")],
),
]);
assert_eq!(
pk_columns(&schema),
vec!["positioned".to_string(), "legacy".to_string()]
);
}
#[test]
fn pk_columns_is_empty_without_a_key() {
let schema = Schema::new(vec![
field("a", &[]),
field("b", &[(LANCE_UNENFORCED_PRIMARY_KEY, "false")]),
]);
assert!(pk_columns(&schema).is_empty());
}
}