diff --git a/python/python/lancedb/_lancedb.pyi b/python/python/lancedb/_lancedb.pyi index 0f7b110ac..0e0eab619 100644 --- a/python/python/lancedb/_lancedb.pyi +++ b/python/python/lancedb/_lancedb.pyi @@ -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: ... diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index cef89e63c..1a9975fe2 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -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)) diff --git a/python/python/lancedb/remote/table.py b/python/python/lancedb/remote/table.py index 89b958165..d14980058 100644 --- a/python/python/lancedb/remote/table.py +++ b/python/python/lancedb/remote/table.py @@ -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].""" diff --git a/python/python/lancedb/table.py b/python/python/lancedb/table.py index 72362acad..4f59355dc 100644 --- a/python/python/lancedb/table.py +++ b/python/python/lancedb/table.py @@ -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. diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index 61b8001cf..6faf80d54 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -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"] diff --git a/python/src/table.rs b/python/src/table.rs index 784d29136..928970d6f 100644 --- a/python/src/table.rs +++ b/python/src/table.rs @@ -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> { + 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 diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index ab1bdfc4a..29fc95f68 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -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, + /// Key columns the caller did not ask for, added to the projection so + /// the surrogate can be built and dropped again before returning. + injected: Vec, + }, +} + +/// The key columns `select` does not already produce, and so would have to be +/// added to it for the surrogate to be buildable. +fn pk_columns_to_inject(select: &Select, pk_columns: &[String]) -> Vec { + let produced: Vec<&String> = match select { + // Already every non-system column, the key among them. + Select::All => return Vec::new(), + Select::Columns(columns) => columns.iter().collect(), + Select::Dynamic(pairs) => pairs.iter().map(|(name, _)| name).collect(), + Select::Expr(pairs) => pairs.iter().map(|(name, _)| name).collect(), + }; + pk_columns + .iter() + .filter(|pk| !produced.contains(pk)) + .cloned() + .collect() +} + +/// Add the key columns `select` is missing, as identity projections where the +/// selection is expression-shaped. +fn inject_pk_columns(select: &mut Select, pk_columns: &[String]) { + let missing = pk_columns_to_inject(select, pk_columns); + if missing.is_empty() { + return; + } + match select { + Select::All => {} + Select::Columns(columns) => columns.extend(missing), + Select::Dynamic(pairs) => pairs.extend(missing.into_iter().map(|pk| (pk.clone(), pk))), + Select::Expr(pairs) => { + pairs.extend(missing.into_iter().map(|pk| { + let expr = crate::expr::col(&pk); + (pk, expr) + })); + } + } +} + +/// Drop `columns` from `batch`, ignoring any that are not present. +fn drop_columns(batch: RecordBatch, columns: &[String]) -> Result { + let keep: Vec = 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 { 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 { + 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, diff --git a/rust/lancedb/src/query/hybrid.rs b/rust/lancedb/src/query/hybrid.rs index 35c64d4a9..9c86202df 100644 --- a/rust/lancedb/src/query/hybrid.rs +++ b/rust/lancedb/src/query/hybrid.rs @@ -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 = None; + let mut ids: HashMap, u64> = HashMap::new(); + + for batch in batches.iter_mut() { + // Both legs empty: `query_schemas` synthesizes a schema carrying only + // the score and join columns, so there is no key to read — and no row + // that would need one. + let row_ids: UInt64Array = if batch.num_rows() == 0 { + UInt64Array::from(Vec::::new()) + } else { + let key_columns = pk_columns + .iter() + .map(|name| { + batch + .column_by_name(name) + .cloned() + .ok_or_else(|| Error::InvalidInput { + message: format!( + "hybrid search could not deduplicate: primary key column {} \ + is missing from a result set with {} rows", + name, + batch.num_rows() + ), + }) + }) + .collect::>>()?; + + let converter = match converter { + Some(ref converter) => converter, + None => converter.insert(RowConverter::new( + key_columns + .iter() + .map(|column| SortField::new(column.data_type().clone())) + .collect(), + )?), + }; + + let rows = converter.convert_columns(&key_columns)?; + let next = &mut ids; + UInt64Array::from_iter_values(rows.iter().map(|row| { + let id = next.len() as u64; + *next.entry(row.as_ref().to_vec()).or_insert(id) + })) + }; + + *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 { + let ids: UInt64Array = downcast_array(batch.column_by_name(ROW_ID).unwrap()); + ids.values().to_vec() + } + + /// The same key in both legs has to land on the same id or the fusion will + /// not dedup, and ids must ascend in fusion order so the tie break holds. + #[test] + fn test_surrogate_row_ids_are_shared_across_legs_in_first_seen_order() { + let mut batches = vec![ + batch(vec!["a", "b", "c"], DIST_COL), + batch(vec!["b", "d"], SCORE_COL), + ]; + stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap(); + + assert_eq!(row_ids(&batches[0]), vec![0, 1, 2]); + assert_eq!( + row_ids(&batches[1]), + vec![1, 3], + "b keeps the id the vector leg gave it" + ); + } + + /// A string key is the common case and cannot be cast to the u64 the + /// rerankers read, which is the whole reason for the surrogate. + #[test] + fn test_surrogate_row_ids_repeat_within_a_leg() { + let mut batches = vec![batch(vec!["a", "a", "b"], DIST_COL)]; + stamp_surrogate_row_ids(&mut batches, &["id".to_string()]).unwrap(); + assert_eq!(row_ids(&batches[0]), vec![0, 0, 1]); + } + + #[test] + fn test_surrogate_row_ids_support_a_composite_key() { + let schema = Arc::new(Schema::new(vec![ + Arc::new(Field::new("tenant", DataType::Utf8, false)), + Arc::new(Field::new("id", DataType::Int32, false)), + ])); + let make = |tenants: Vec<&str>, ids: Vec| { + RecordBatch::try_new( + schema.clone(), + vec![ + Arc::new(StringArray::from(tenants)), + Arc::new(arrow_array::Int32Array::from(ids)), + ], + ) + .unwrap() + }; + let mut batches = vec![make(vec!["x", "y"], vec![1, 1]), make(vec!["y"], vec![1])]; + stamp_surrogate_row_ids(&mut batches, &["tenant".to_string(), "id".to_string()]).unwrap(); + + assert_eq!( + row_ids(&batches[0]), + vec![0, 1], + "same id, different tenant" + ); + assert_eq!(row_ids(&batches[1]), vec![1]); + } + + /// Both legs empty: `query_schemas` synthesizes a schema with no key column, + /// and there is no row that would need one. + #[test] + fn test_surrogate_row_ids_tolerate_the_both_empty_schema() { + let mut batches = vec![ + RecordBatch::new_empty(Arc::new(empty_vec_schema())), + RecordBatch::new_empty(Arc::new(empty_fts_schema())), + ]; + // 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::::new()); + } + + /// Rows with no key would silently fuse into one bucket, so refuse. + #[test] + fn test_surrogate_row_ids_reject_a_missing_key_column() { + let mut batches = vec![batch(vec!["a"], DIST_COL)]; + let err = stamp_surrogate_row_ids(&mut batches, &["nope".to_string()]).unwrap_err(); + assert!(err.to_string().contains("nope"), "{err}"); + } + + #[test] + fn test_surrogate_row_ids_reject_a_table_with_no_primary_key() { + let mut batches = vec![batch(vec!["a"], DIST_COL)]; + let err = stamp_surrogate_row_ids(&mut batches, &[]).unwrap_err(); + assert!(err.to_string().contains("unenforced primary key"), "{err}"); + } } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 5c32e3cdd..6587dc718 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -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 { version: Arc>>, location: RwLock>, schema_cache: BackgroundCache, + wal_enabled_cache: BackgroundCache, freshness: Arc>, /// The branch this handle is scoped to, or `None` for the main branch. /// Stamped onto every branch-accepting request so reads and writes resolve @@ -647,6 +654,10 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + wal_enabled_cache: BackgroundCache::new( + WAL_ENABLED_CACHE_TTL, + WAL_ENABLED_CACHE_REFRESH_WINDOW, + ), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch: None, } @@ -680,6 +691,10 @@ impl RemoteTable { version: Arc::new(RwLock::new(None)), location: RwLock::new(None), schema_cache: BackgroundCache::new(SCHEMA_CACHE_TTL, SCHEMA_CACHE_REFRESH_WINDOW), + wal_enabled_cache: BackgroundCache::new( + WAL_ENABLED_CACHE_TTL, + WAL_ENABLED_CACHE_REFRESH_WINDOW, + ), freshness: Arc::new(Mutex::new(FreshnessState::default())), branch, } @@ -1414,6 +1429,10 @@ impl RemoteTable { 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 BaseTable for RemoteTable { .json(&body); let (request_id, response) = self.send(request, true).await?; self.check_table_response(&request_id, response).await?; + self.invalidate_wal_enabled_cache(); Ok(()) } @@ -3082,6 +3114,7 @@ impl BaseTable for RemoteTable { )); let (request_id, response) = self.send(request, true).await?; self.check_table_response(&request_id, response).await?; + self.invalidate_wal_enabled_cache(); Ok(()) } @@ -3148,6 +3181,42 @@ impl BaseTable for RemoteTable { Ok(Some(spec)) } + /// Cached because every hybrid query asks, and the answer is a routing + /// hint rather than a correctness gate: a wrong `false` surfaces as the + /// `with_row_id` rejection on the query itself, and a wrong `true` still + /// produces a valid base-only read. + /// + /// An unreachable endpoint therefore answers `false` instead of failing. + /// That hides nothing: a server without the route has no MemWAL tables, and + /// if the table itself is gone the query that follows reports it. Only a + /// success is cached, so a transient failure is re-asked on the next call. + /// + /// Reuses `get_lsm_write_spec` rather than `BackgroundCache::get`, whose + /// fetch closure must be `'static` and so cannot borrow `self`. The cost is + /// no single-flight on a cold cache — concurrent first calls each send one + /// request — and a refresh that blocks rather than happening in the + /// background. Both are one POST against a bit that changes almost never. + async fn lsm_enabled(&self) -> Result { + if let Some(enabled) = self.wal_enabled_cache.try_get() { + return Ok(enabled); + } + match self.get_lsm_write_spec().await { + Ok(spec) => { + let enabled = spec.is_some(); + self.wal_enabled_cache.seed(enabled); + Ok(enabled) + } + Err(e) => { + log::debug!( + "lsm_enabled: reading the write spec for table {} failed, \ + treating the table as base-only: {e}", + self.identifier + ); + Ok(false) + } + } + } + async fn tags(&self) -> Result> { 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::::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::>() + .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}"); + } } diff --git a/rust/lancedb/src/table.rs b/rust/lancedb/src/table.rs index 7f139c4cb..36893759c 100644 --- a/rust/lancedb/src/table.rs +++ b/rust/lancedb/src/table.rs @@ -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 { + 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 { + 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 diff --git a/rust/lancedb/src/table/primary_key.rs b/rust/lancedb/src/table/primary_key.rs index 8a7b48efa..f8932bfa1 100644 --- a/rust/lancedb/src/table/primary_key.rs +++ b/rust/lancedb/src/table/primary_key.rs @@ -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 { + let mut keyed: Vec<(bool, u32, usize, &str)> = schema + .fields() + .iter() + .enumerate() + .filter_map(|(idx, field)| { + let metadata = field.metadata(); + let position = metadata + .get(LANCE_UNENFORCED_PRIMARY_KEY_POSITION) + .and_then(|value| value.parse::().ok()) + .or_else(|| { + metadata + .get(LANCE_UNENFORCED_PRIMARY_KEY) + .filter(|value| str_is_truthy(value)) + .map(|_| 0) + })?; + Some((position == 0, position, idx, field.name().as_str())) + }) + .collect(); + keyed.sort_unstable(); + keyed + .into_iter() + .map(|(_, _, _, name)| name.to_string()) + .collect() +} + +#[cfg(test)] +mod tests { + use std::collections::HashMap; + + use arrow_schema::{Field, Schema}; + + use super::*; + + fn field(name: &str, metadata: &[(&str, &str)]) -> Field { + Field::new(name, DataType::Utf8, false).with_metadata( + metadata + .iter() + .map(|(k, v)| (k.to_string(), v.to_string())) + .collect::>(), + ) + } + + #[test] + fn pk_columns_reads_the_position_marker() { + let schema = Schema::new(vec![ + field("other", &[]), + field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]), + ]); + assert_eq!(pk_columns(&schema), vec!["id".to_string()]); + } + + /// Key order is the position, not the order the columns are declared in. + #[test] + fn pk_columns_orders_a_composite_key_by_position() { + let schema = Schema::new(vec![ + field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "2")]), + field("tenant", &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")]), + ]); + assert_eq!( + pk_columns(&schema), + vec!["tenant".to_string(), "id".to_string()] + ); + } + + /// Datasets written before the position marker carry only the boolean. + #[test] + fn pk_columns_accepts_the_legacy_boolean_marker() { + let schema = Schema::new(vec![field("id", &[(LANCE_UNENFORCED_PRIMARY_KEY, "true")])]); + assert_eq!(pk_columns(&schema), vec!["id".to_string()]); + } + + /// Mirrors Lance: an explicit position sorts ahead of a legacy field. + #[test] + fn pk_columns_sorts_positioned_fields_ahead_of_legacy_ones() { + let schema = Schema::new(vec![ + field("legacy", &[(LANCE_UNENFORCED_PRIMARY_KEY, "yes")]), + field( + "positioned", + &[(LANCE_UNENFORCED_PRIMARY_KEY_POSITION, "1")], + ), + ]); + assert_eq!( + pk_columns(&schema), + vec!["positioned".to_string(), "legacy".to_string()] + ); + } + + #[test] + fn pk_columns_is_empty_without_a_key() { + let schema = Schema::new(vec![ + field("a", &[]), + field("b", &[(LANCE_UNENFORCED_PRIMARY_KEY, "false")]), + ]); + assert!(pk_columns(&schema).is_empty()); + } +}