mirror of
https://github.com/lancedb/lancedb.git
synced 2026-10-30 15:26:52 +00:00
feat(python): add CombinedFieldsQuery for BM25F full text search
Add a typed Python `CombinedFieldsQuery` for BM25F search across multiple columns. The query mirrors Lance's combined-fields semantics with ordered per-column boosts and AND/OR term matching. Convert it in both directions at the PyO3 boundary, including nested Boolean queries, and document the Python typed convenience entry. Closes #4652.
This commit is contained in:
@@ -508,6 +508,8 @@ and combined with [BooleanQuery][lancedb.query.BooleanQuery].
|
||||
|
||||
::: lancedb.query.MultiMatchQuery
|
||||
|
||||
::: lancedb.query.CombinedFieldsQuery
|
||||
|
||||
::: lancedb.query.BooleanQuery
|
||||
|
||||
::: lancedb.query.FullTextOperator
|
||||
|
||||
@@ -748,6 +748,50 @@ const multiMatchBoostingResults = await table.query()
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
### 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.
|
||||
|
||||
<CodeGroup>
|
||||
```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()
|
||||
)
|
||||
```
|
||||
</CodeGroup>
|
||||
|
||||
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).
|
||||
|
||||
@@ -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):
|
||||
"""
|
||||
|
||||
@@ -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(
|
||||
[
|
||||
|
||||
+34
-5
@@ -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<FtsQuery> {
|
||||
|
||||
Ok(Self(q.with_operator(op).into()))
|
||||
}
|
||||
"CombinedFieldsQuery" => {
|
||||
let query = ob.getattr("query")?.extract()?;
|
||||
let columns = ob.getattr("columns")?.extract()?;
|
||||
let boosts: Option<Vec<f32>> = 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<FtsQuery> {
|
||||
.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<FtsQuery> {
|
||||
.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(),
|
||||
|
||||
Reference in New Issue
Block a user