diff --git a/docs/src/python/python.md b/docs/src/python/python.md index 84188aecb..af3f9d598 100644 --- a/docs/src/python/python.md +++ b/docs/src/python/python.md @@ -508,6 +508,8 @@ and combined with [BooleanQuery][lancedb.query.BooleanQuery]. ::: lancedb.query.MultiMatchQuery +::: lancedb.query.CombinedFieldsQuery + ::: lancedb.query.BooleanQuery ::: lancedb.query.FullTextOperator diff --git a/docs/web/search/full-text-search.mdx b/docs/web/search/full-text-search.mdx index 5e0d7d8d7..858492576 100644 --- a/docs/web/search/full-text-search.mdx +++ b/docs/web/search/full-text-search.mdx @@ -748,6 +748,50 @@ const multiMatchBoostingResults = await table.query() ``` +### Search Columns as One Field {#combined-fields-search} + +`MultiMatchQuery` scores each column on its own and keeps the best match. +`CombinedFieldsQuery` treats the columns as one field (BM25F), so with the `AND` +operator a row matches when every term appears in at least one of the columns. +The query below finds `car` in `text` and `merrily` in `text2`, returning rows 1 +and 3. The same query with `MultiMatchQuery` returns no rows, because neither +column contains both words on its own. + +Every column must use the same tokenizer configuration. The two indexes in this +example are built with the same `FTS` settings. Boosts work as they do for +`MultiMatchQuery`: one weight per column, each at least 1. + + +```python Python icon="python" +from lancedb.query import CombinedFieldsQuery, FullTextOperator + +# Every term must appear in at least one of the two columns +combined_results = ( + table.search( + CombinedFieldsQuery("car merrily", ["text", "text2"], operator=FullTextOperator.AND), + ) + .select(["id", "text", "text2", "_score"]) + .limit(100) + .to_arrow().to_pandas() +) + +# Weight matches in text2 twice as much as matches in text +combined_boosting_results = ( + table.search( + CombinedFieldsQuery("car merrily", ["text", "text2"], boosts=[1.0, 2.0]), + ) + .select(["id", "text", "text2", "_score"]) + .limit(100) + .to_arrow().to_pandas() +) +``` + + +The typed `CombinedFieldsQuery` convenience class and this example are specific +to the Python SDK. Rust applications can construct the core query type exposed +by `lancedb::index::scalar`. TypeScript and Java do not yet provide typed +convenience entry points for this query. + ### Boolean Queries {#boolean-queries} LanceDB supports boolean logic in full-text search, allowing you to combine multiple queries using `and` and `or` operators. This is useful when you want to match documents that satisfy multiple conditions (intersection) or at least one of several conditions (union). diff --git a/python/python/lancedb/query.py b/python/python/lancedb/query.py index 625a17524..c9ca258e7 100644 --- a/python/python/lancedb/query.py +++ b/python/python/lancedb/query.py @@ -445,6 +445,7 @@ class FullTextQueryType(str, Enum): MATCH_PHRASE = "match_phrase" BOOST = "boost" MULTI_MATCH = "multi_match" + COMBINED_FIELDS = "combined_fields" BOOLEAN = "boolean" @@ -669,6 +670,42 @@ class MultiMatchQuery(FullTextQuery): return FullTextQueryType.MULTI_MATCH +@pydantic.dataclasses.dataclass +class CombinedFieldsQuery(FullTextQuery): + """ + Combined-fields (BM25F) query for full-text search. + + Unlike `MultiMatchQuery`, which scores each column on its own and keeps + the best one, this query treats the columns as a single virtual field, so + term statistics are blended across them. A query term can then match in + any of the columns, e.g. a first name in one column and a last name in + another. + + Parameters + ---------- + query : str + The query string to match against. + columns : list[str] + The columns combined into the virtual field. Every column must have the + same tokenizer configuration, and at least one must have an FTS index. + boosts : list[float], optional + The weight of each column, aligned with `columns`. Each weight must be + at least 1. If not provided, every column has a weight of 1. + operator : FullTextOperator, default OR + The operator to use for combining the query terms. + If `AND`, every term must appear in at least one of the columns. + If `OR`, at least one term must match. + """ + + query: str + columns: list[str] + boosts: Optional[list[float]] = pydantic.Field(None, kw_only=True) + operator: FullTextOperator = pydantic.Field(FullTextOperator.OR, kw_only=True) + + def query_type(self) -> FullTextQueryType: + return FullTextQueryType.COMBINED_FIELDS + + @pydantic.dataclasses.dataclass class BooleanQuery(FullTextQuery): """ diff --git a/python/python/tests/test_fts.py b/python/python/tests/test_fts.py index a02c4db28..486c1ce62 100644 --- a/python/python/tests/test_fts.py +++ b/python/python/tests/test_fts.py @@ -26,7 +26,9 @@ from lancedb.db import DBConnection from lancedb.index import FTS from lancedb.query import ( BoostQuery, + CombinedFieldsQuery, DocumentGranularity, + FullTextOperator, MatchQuery, MultiMatchQuery, PhraseQuery, @@ -1367,6 +1369,103 @@ def test_match_query_fuzziness_search(tmp_db, options, expected): assert sorted(row["text"] for row in results) == sorted(expected) +@pytest.mark.parametrize( + "operator,combined_ids,multi_match_ids", + [ + # Combined fields treat the columns as one field, so AND also matches + # row 0, where "john" is in the title and "smith" is in the body. + (FullTextOperator.AND, {0, 1}, {1}), + (FullTextOperator.OR, {0, 1, 2}, {0, 1, 2}), + ], +) +def test_combined_fields_query(mem_db, operator, combined_ids, multi_match_ids): + table = mem_db.create_table( + "people", + pa.table( + { + "id": [0, 1, 2], + "title": ["john", "john smith", "john"], + "body": ["smith", "foo", "alice"], + } + ), + ) + table.create_index("title", config=FTS()) + table.create_index("body", config=FTS()) + + def ids(query): + return {row["id"] for row in table.search(query).select(["id"]).to_list()} + + combined = CombinedFieldsQuery("john smith", ["title", "body"], operator=operator) + multi_match = MultiMatchQuery("john smith", ["title", "body"], operator=operator) + assert ids(combined) == combined_ids + assert ids(multi_match) == multi_match_ids + + +def test_combined_fields_query_boosts(mem_db): + table = mem_db.create_table( + "docs", pa.table({"title": ["hello"], "body": ["world"]}) + ) + table.create_index("title", config=FTS()) + table.create_index("body", config=FTS()) + + results = table.search( + CombinedFieldsQuery("hello", ["title", "body"], boosts=[1.5, 1.0]) + ).to_list() + assert len(results) == 1 + + # Each weight must be at least 1, and there must be one per column. + with pytest.raises(ValueError, match="boost"): + table.search( + CombinedFieldsQuery("hello", ["title", "body"], boosts=[0.5, 1.0]) + ).to_list() + with pytest.raises(ValueError, match="boosts"): + table.search( + CombinedFieldsQuery("hello", ["title", "body"], boosts=[1.0]) + ).to_list() + + +@pytest.mark.asyncio +async def test_combined_fields_query_round_trip(mem_db_async): + table = await mem_db_async.create_table( + "docs", pa.table({"title": ["hello"], "body": ["world"]}) + ) + query = CombinedFieldsQuery( + "hello world", + ["title", "body"], + boosts=[2.0, 1.0], + operator=FullTextOperator.AND, + ) + + query_obj = table.query().nearest_to_text(query).to_query_object() + + assert query_obj.full_text_query.query == query + + +@pytest.mark.asyncio +async def test_combined_fields_query_nested_round_trip(mem_db_async): + table = await mem_db_async.create_table( + "docs", pa.table({"title": ["hello"], "body": ["world"]}) + ) + query = BooleanQuery( + [ + ( + Occur.MUST, + CombinedFieldsQuery( + "hello world", + ["title", "body"], + boosts=[2.0, 1.0], + operator=FullTextOperator.AND, + ), + ), + (Occur.MUST_NOT, MatchQuery("draft", "body")), + ] + ) + + query_obj = table.query().nearest_to_text(query).to_query_object() + + assert query_obj.full_text_query.query == query + + def test_fts_query_to_json(): """Test that FTS query to_json() produces valid JSON strings with exact format.""" @@ -1504,6 +1603,17 @@ def test_fts_query_to_json(): ) assert json_str == expected + # Test CombinedFieldsQuery + combined_fields = CombinedFieldsQuery( + "python", ["tags", "title"], boosts=[2.0, 1.0], operator=FullTextOperator.AND + ) + json_str = combined_fields.to_json() + expected = ( + '{"combined_fields":{"query":"python","columns":["tags","title"],' + '"boost":[2.0,1.0],"operator":"And"}}' + ) + assert json_str == expected + # Test complex nested BooleanQuery inner1 = BooleanQuery( [ diff --git a/python/src/query.rs b/python/src/query.rs index c2bd6afb3..8687b3756 100644 --- a/python/src/query.rs +++ b/python/src/query.rs @@ -17,8 +17,8 @@ use arrow::pyarrow::FromPyArrow; use arrow::pyarrow::IntoPyArrow; use arrow::pyarrow::ToPyArrow; use lancedb::index::scalar::{ - BooleanQuery, BoostQuery, DocumentGranularity, FtsQuery, FullTextSearchQuery, MatchQuery, - MultiMatchQuery, Occur, Operator, PhraseQuery, + BooleanQuery, BoostQuery, CombinedFieldsQuery, DocumentGranularity, FtsQuery, + FullTextSearchQuery, MatchQuery, MultiMatchQuery, Occur, Operator, PhraseQuery, }; use lancedb::query::AnalyzePlanDistributedMetrics; use lancedb::query::QueryBase; @@ -148,6 +148,26 @@ impl<'a, 'py> FromPyObject<'a, 'py> for PyLanceDB { Ok(Self(q.with_operator(op).into())) } + "CombinedFieldsQuery" => { + let query = ob.getattr("query")?.extract()?; + let columns = ob.getattr("columns")?.extract()?; + let boosts: Option> = ob.getattr("boosts")?.extract()?; + let operator: String = ob.getattr("operator")?.extract()?; + + let q = CombinedFieldsQuery::try_new(query, columns) + .map_err(|e| PyValueError::new_err(format!("Invalid query: {}", e)))?; + let q = if let Some(boosts) = boosts { + q.try_with_boosts(boosts) + .map_err(|e| PyValueError::new_err(format!("Invalid boosts: {}", e)))? + } else { + q + }; + + let op = Operator::try_from(operator.as_str()) + .map_err(|e| PyValueError::new_err(format!("Invalid operator: {}", e)))?; + + Ok(Self(q.with_operator(op).into())) + } "BooleanQuery" => { let queries: Vec<(String, Self)> = ob.getattr("queries")?.extract()?; let mut sub_queries = Vec::with_capacity(queries.len()); @@ -178,9 +198,6 @@ impl<'py> IntoPyObject<'py> for PyLanceDB { .expect("Failed to import namespace"); match self.0 { - FtsQuery::CombinedFields(_) => Err(PyValueError::new_err( - "Combined-fields FTS queries are not yet supported in Python", - )), FtsQuery::Match(query) => { let kwargs = PyDict::new(py); kwargs.set_item("boost", query.boost)?; @@ -236,6 +253,18 @@ impl<'py> IntoPyObject<'py> for PyLanceDB { .getattr(intern!(py, "MultiMatchQuery"))? .call((first.terms.clone(), columns), Some(&kwargs)) } + FtsQuery::CombinedFields(query) => { + let (columns, boosts): (Vec<_>, Vec<_>) = query + .weighted_columns() + .map(|(column, boost)| (column.to_string(), boost)) + .unzip(); + let kwargs = PyDict::new(py); + kwargs.set_item("boosts", boosts)?; + kwargs.set_item::<_, &str>("operator", query.operator().into())?; + namespace + .getattr(intern!(py, "CombinedFieldsQuery"))? + .call((query.terms(), columns), Some(&kwargs)) + } FtsQuery::Boolean(query) => { let mut queries: Vec<(&str, Bound<'py, PyAny>)> = Vec::with_capacity( query.should.len() + query.must.len() + query.must_not.len(),