diff --git a/nodejs/__test__/query.test.ts b/nodejs/__test__/query.test.ts index c3c6c6e76..1feb1d643 100644 --- a/nodejs/__test__/query.test.ts +++ b/nodejs/__test__/query.test.ts @@ -132,6 +132,40 @@ describe("Query outputSchema", () => { }, ); + it.each(["columns", "dynamic", "all"])( + "should keep hybrid search working with scoring autoprojection disabled (%s)", + async (projection) => { + await table.createIndex("text", { config: Index.fts() }); + const query = () => { + const query = table.vectorSearch([1, 2]).fullTextSearch("foo"); + if (projection === "columns") { + query.select(["a"]); + } else if (projection === "dynamic") { + query.select({ doubled: "a * 2" }); + } + return query; + }; + const expected = await query().toArrow(); + const results = await query().disableScoringAutoprojection().toArrow(); + + const rows = (result: typeof results) => + result.toArray().map((row) => ({ + a: row.a, + doubled: row.doubled, + text: row.text, + vec: row.vec ? Array.from(row.vec) : undefined, + relevance: row._relevance_score, + })); + expect(results.numRows).toBe(3); + expect(rows(results)).toEqual(rows(expected)); + expect(results.schema.fields.map((field) => field.name)).toEqual( + expected.schema.fields + .map((field) => field.name) + .filter((name) => name !== "_distance" && name !== "_score"), + ); + }, + ); + it("should preserve the scoring option when converting to a vector query", async () => { const results = await table .query() diff --git a/rust/lancedb/src/query.rs b/rust/lancedb/src/query.rs index c946ad2fe..740a92938 100644 --- a/rust/lancedb/src/query.rs +++ b/rust/lancedb/src/query.rs @@ -1556,6 +1556,33 @@ impl VectorQuery { } } + // Hybrid normalization and reranking still need each leg's score even + // when the caller disables scoring autoprojection. Select those scores + // internally and strip the injected columns only after fusion. + let mut injected_scores = Vec::new(); + if self.request.base.disable_scoring_autoprojection { + if matches!(self.request.base.select, Select::All) { + let columns = self + .parent + .schema() + .await? + .fields() + .iter() + .map(|field| field.name().clone()) + .collect(); + fts_query.request.select = Select::Columns(columns); + vector_query.request.base.select = fts_query.request.select.clone(); + } + for (select, score) in [ + (&mut fts_query.request.select, SCORE_COL), + (&mut vector_query.request.base.select, DIST_COL), + ] { + if hybrid::add_score_column(select, score) { + injected_scores.push(score); + } + } + } + vector_query.request.base.full_text_search = None; let (fts_results, vec_results) = try_join!( fts_query.execute_with_options(internal_options.clone()), @@ -1625,6 +1652,12 @@ impl VectorQuery { Some(pk_fusion) => results = pk_fusion.strip(results)?, } + for score in injected_scores { + if results.schema().column_with_name(score).is_some() { + results = results.drop_column(score)?; + } + } + Ok(single_batch_stream(results, max_batch_length)) } @@ -3590,6 +3623,57 @@ mod tests { assert!(texts.contains("cat")); // should be close by vector search assert!(texts.contains("b")); // should be close by fts search + // Scores remain available to normalization and reranking for every + // projection form, but do not leak into the opt-in query's output. + for select in [ + Select::All, + Select::columns(&["text"]), + Select::dynamic(&[("label", "upper(text)")]), + Select::expr_projection(&[("label", col("text"))]), + ] { + for norm in [NormalizeMethod::Rank, NormalizeMethod::Score] { + for term in ["b", "z"] { + let query = table + .query() + .full_text_search(FullTextSearchQuery::new(term.to_string())) + .nearest_to(&[-10.0, -10.0]) + .unwrap() + .select(select.clone()) + .norm(norm.clone()) + .with_row_id() + .limit(2); + let expected = query + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + let expected = &expected[0]; + let columns = expected + .schema() + .fields() + .iter() + .enumerate() + .filter_map(|(index, field)| { + (!matches!(field.name().as_str(), SCORE_COL | DIST_COL)) + .then_some(index) + }) + .collect::>(); + let expected = expected.project(&columns).unwrap(); + let actual = query + .disable_scoring_autoprojection() + .execute() + .await + .unwrap() + .try_collect::>() + .await + .unwrap(); + assert_eq!(actual[0], expected); + } + } + } + // ensure that this works correctly if there are no matching FTS results let fts_query = FullTextSearchQuery::new("z".to_string()); table diff --git a/rust/lancedb/src/query/hybrid.rs b/rust/lancedb/src/query/hybrid.rs index 35c64d4a9..ba94f0ebb 100644 --- a/rust/lancedb/src/query/hybrid.rs +++ b/rust/lancedb/src/query/hybrid.rs @@ -12,6 +12,36 @@ use lance_index::{scalar::inverted::SCORE_COL, vector::DIST_COL}; use std::sync::Arc; use crate::error::{Error, Result}; +use crate::query::Select; + +/// Include a leg's score in its internal projection, recording whether it was +/// injected so it can be removed after reranking. +/// +/// The caller expands `Select::All` before adding scoring columns. +pub(super) fn add_score_column(select: &mut Select, score: &str) -> bool { + match select { + Select::All => unreachable!("hybrid score projection requires explicit columns"), + Select::Columns(columns) => { + if columns.iter().any(|name| name == score) { + return false; + } + columns.push(score.to_string()); + } + Select::Dynamic(columns) => { + if columns.iter().any(|(name, _)| name == score) { + return false; + } + columns.push((score.to_string(), score.to_string())); + } + Select::Expr(columns) => { + if columns.iter().any(|(name, _)| name == score) { + return false; + } + columns.push((score.to_string(), crate::expr::col(score))); + } + } + true +} /// Converts results's score column to a rank. ///