From 83cfcb84d20f3a937eeff919e2921f08f85d775f Mon Sep 17 00:00:00 2001 From: Daniel Rammer Date: Wed, 23 Sep 2026 14:20:56 -0500 Subject: [PATCH] fix(query): restore the empty-result and real-row-id contracts Three defects from review, each reproduced before it was touched. A hybrid query matching nothing errored instead of returning nothing. When a leg comes back with no batches, `query_schemas` synthesizes a schema that already carries `_rowid`, and stamping then tried to append a second column of that name. The column is ours to define, so replace it rather than append. The unit test covering that case stripped `_rowid` by hand first, claiming to do "what the caller does" -- the caller does no such thing, so the test asserted a shape `run_hybrid` never produces. It now runs the synthesized schemas untouched, and a mock-server test covers a no-match query end to end. An explicit `with_row_id` combined with `return_score="all"` was accepted and then silently stripped, in either call order, because both set the same `_with_row_id` flag and the guard inferred caller intent from it. Record the caller's own request instead of inferring it: `with_row_id` on the hybrid builder sets `_caller_requested_row_id`, and `rerank` sets `_with_row_id` directly without going through it. A projected Blob v2 column slipped the same way. `blob_auto_row_id_for_scan` reports whether a row id would be added *automatically*, so it answers False once anything else has turned `_with_row_id` on -- including that reranker -- while the blob still needs a real row id to fetch through. Check the projected sources directly. Co-Authored-By: Claude Opus 5 --- python/python/lancedb/query.py | 29 +++++++++++---- python/python/tests/test_hybrid_query.py | 43 ++++++++++++++++++++++ rust/lancedb/src/query/hybrid.rs | 21 +++++------ rust/lancedb/src/remote/table.rs | 46 ++++++++++++++++++++++++ 4 files changed, 119 insertions(+), 20 deletions(-) diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 57b0a4b0a..f12cc9c25 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -2332,8 +2332,9 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): # 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 + # `_with_row_id` is also turned on by `rerank(return_score="all")` for + # the reranker's own use, so it cannot answer "did the caller ask?". + self._caller_requested_row_id = False def _validate_query(self, query, vector=None, text=None): if query is not None and (vector is not None or text is not None): @@ -2355,6 +2356,16 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): return vector_query, text_query + def with_row_id(self, with_row_id: bool) -> Self: + """Set whether to return row ids. + + Recorded separately from `_with_row_id`, which `rerank(return_score= + "all")` also sets for its own use — on a MemWAL table only a caller who + asked is refused, and the reranker still falls back to the primary key. + """ + self._caller_requested_row_id = with_row_id + return super().with_row_id(with_row_id) + def phrase_query(self, phrase_query: bool = True) -> LanceHybridQueryBuilder: """Set whether to use phrase query. @@ -2390,7 +2401,7 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): # `return_score="all"` turns row ids on for the reranker, not the # caller, so it still falls back; only a caller who asked for them # is refused, and with the reason rather than a bare 400. - if self._user_requested_row_id() and not self._reranker_requested_row_id: + if self._caller_requested_row_id: raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED) from e # Learned, so later queries on this table skip straight to it. self._table._note_hybrid_pk_fusion() @@ -2436,9 +2447,13 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): or self._table._hybrid_pk_fusion_learned() is not True ): return None - if self._user_requested_row_id() and not self._reranker_requested_row_id: + if self._caller_requested_row_id: raise NotImplementedError(_WAL_ROW_ID_UNSUPPORTED) - if self._blob_auto_row_id_enabled(): + # Not `_blob_auto_row_id_enabled`: that reports whether a row id would + # be added *automatically*, so it goes quiet once anything else has + # turned `_with_row_id` on — including `return_score="all"`. A projected + # blob needs a real row id to fetch through either way. + if blob_v2_projection_sources(self._table.schema, self._columns): 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 " @@ -2640,8 +2655,8 @@ class LanceHybridQueryBuilder(LanceQueryBuilder): # 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) + # Deliberately not through `with_row_id`, which records caller intent. + self._with_row_id = True return self diff --git a/python/python/tests/test_hybrid_query.py b/python/python/tests/test_hybrid_query.py index fc909cfef..7d86d37f0 100644 --- a/python/python/tests/test_hybrid_query.py +++ b/python/python/tests/test_hybrid_query.py @@ -608,3 +608,46 @@ def test_wal_hybrid_fuses_on_the_key_and_drops_it_again(): assert "_rowid" not in results.column_names, "the surrogate is internal" assert "id" not in results.column_names, "the caller selected only `text`" assert sorted(results.column("text").to_pylist()) == ["a", "b", "c"] + + +@pytest.mark.parametrize( + "row_id_first", [True, False], ids=["row_id_first", "rerank_first"] +) +def test_wal_hybrid_refuses_row_id_even_with_return_score_all(row_id_first: bool): + """`return_score="all"` sets the same flag the caller's request does, so + asking for both must still be refused — in either call order.""" + builder = _wal_hybrid() + reranker = RRFReranker(return_score="all") + if row_id_first: + builder = builder.with_row_id(True).rerank(reranker) + else: + builder = builder.rerank(reranker).with_row_id(True) + + with pytest.raises(NotImplementedError, match="use_lsm"): + builder._create_query_builders() + + +def test_wal_hybrid_refuses_a_blob_projection_with_return_score_all(): + """A projected blob needs a real row id to fetch through, which the + reranker turning `_with_row_id` on does not provide.""" + schema = pa.schema( + [ + pa.field( + "id", + pa.string(), + metadata={b"lance-schema:unenforced-primary-key:position": b"1"}, + ), + pa.field("text", pa.string()), + lancedb.blob("blob"), + ] + ) + builder = ( + LanceHybridQueryBuilder(_wal_table(schema)) + .vector([0.1, 0.2]) + .text("puppy") + .select(["blob"]) + .rerank(RRFReranker(return_score="all")) + ) + + with pytest.raises(NotImplementedError, match="blob"): + builder._create_query_builders() diff --git a/rust/lancedb/src/query/hybrid.rs b/rust/lancedb/src/query/hybrid.rs index 9c86202df..e656daa0d 100644 --- a/rust/lancedb/src/query/hybrid.rs +++ b/rust/lancedb/src/query/hybrid.rs @@ -244,7 +244,14 @@ pub fn stamp_surrogate_row_ids(batches: &mut [RecordBatch], pk_columns: &[String })) }; - *batch = batch.try_with_column( + // `query_schemas` synthesizes a schema carrying `ROW_ID` when a leg came + // back empty, so the column can already be there — and it is ours to + // define either way. + let base = match batch.schema().column_with_name(ROW_ID) { + Some(_) => batch.drop_column(ROW_ID)?, + None => batch.clone(), + }; + *batch = base.try_with_column( Field::new(ROW_ID, DataType::UInt64, false), Arc::new(row_ids), )?; @@ -507,18 +514,6 @@ mod test { 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()); } diff --git a/rust/lancedb/src/remote/table.rs b/rust/lancedb/src/remote/table.rs index 15f452f23..f4147c5ce 100644 --- a/rust/lancedb/src/remote/table.rs +++ b/rust/lancedb/src/remote/table.rs @@ -13745,6 +13745,52 @@ mod tests { assert_eq!(batch.num_rows(), 2, "`a` is in both legs and fuses"); } + /// A hybrid query that matches nothing must return nothing, not fail. + /// When both legs come back empty, `query_schemas` synthesizes a schema + /// that already carries `_rowid`, which the stamping has to tolerate. + #[tokio::test] + async fn test_hybrid_with_no_matches_returns_an_empty_result() { + 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(&pk_schema()).into_bytes()) + .unwrap(), + "/v1/table/my_table/query/" => { + let body = request_body_json(&request); + if body["with_row_id"] == serde_json::Value::Bool(true) { + return http::Response::builder() + .status(400) + .body(WAL_ROW_ID_REFUSAL_BODY.as_bytes().to_vec()) + .unwrap(); + } + // A leg that matched nothing: schema, no batches. + let schema = + Arc::new(Schema::new(vec![Field::new("id", DataType::Utf8, false)])); + let mut body = Vec::new(); + { + let mut writer = + arrow_ipc::writer::FileWriter::try_new(&mut body, &schema).unwrap(); + writer.finish().unwrap(); + } + http::Response::builder().status(200).body(body).unwrap() + } + path => panic!("unexpected request path: {path}"), + } + }); + + let results = hybrid(&table) + .limit(10) + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + + assert_eq!(results.iter().map(|b| b.num_rows()).sum::(), 0); + } + /// The refusal is paid once: the second query goes straight to the key. #[tokio::test] async fn test_hybrid_remembers_the_refusal() {