Expose boost query (#250)

master
alex-au-922 2024-04-25 05:57:16 +08:00 committed by GitHub
parent ed7374c7bd
commit c74990aeb8
No known key found for this signature in database
GPG Key ID: B5690EEEBB952194
3 changed files with 119 additions and 10 deletions

View File

@ -178,4 +178,13 @@ impl Query {
inner: Box::new(dismax_query),
})
}
#[staticmethod]
#[pyo3(signature = (query, boost))]
pub(crate) fn boost_query(query: Query, boost: f32) -> PyResult<Query> {
let inner = tv::query::BoostQuery::new(query.inner, boost);
Ok(Query {
inner: Box::new(inner),
})
}
}

View File

@ -213,6 +213,10 @@ class Query:
def disjunction_max_query(subqueries: Sequence[Query], tie_breaker: Optional[float] = None) -> Query:
pass
@staticmethod
def boost_query(query: Query, boost: float) -> Query:
pass
class Order(Enum):
Asc = 1

View File

@ -899,3 +899,99 @@ class TestQuery(object):
with pytest.raises(TypeError, match = r"'str' object cannot be converted to 'Query'"):
query = Query.disjunction_max_query([query1, "not a query"], tie_breaker=0.5)
def test_boost_query(self, ram_index):
index = ram_index
query1 = Query.term_query(index.schema, "title", "sea")
boosted_query = Query.boost_query(query1, 2.0)
# Normal boost query
assert (
repr(boosted_query)
== """Query(Boost(query=TermQuery(Term(field=0, type=Str, "sea")), boost=2))"""
)
query2 = Query.fuzzy_term_query(index.schema, "title", "ice")
combined_query = Query.boolean_query([
(Occur.Should, boosted_query),
(Occur.Should, query2)
])
boosted_query = Query.boost_query(combined_query, 2.0)
# Boosted boolean query
assert (
repr(boosted_query)
== """Query(Boost(query=BooleanQuery { subqueries: [(Should, Boost(query=TermQuery(Term(field=0, type=Str, "sea")), boost=2)), (Should, FuzzyTermQuery { term: Term(field=0, type=Str, "ice"), distance: 1, transposition_cost_one: true, prefix: false })] }, boost=2))"""
)
boosted_query = Query.boost_query(query1, 0.1)
# Check for decimal boost values
assert(
repr(boosted_query)
== """Query(Boost(query=TermQuery(Term(field=0, type=Str, "sea")), boost=0.1))"""
)
boosted_query = Query.boost_query(query1, 0.0)
# Check for zero boost values
assert(
repr(boosted_query)
== """Query(Boost(query=TermQuery(Term(field=0, type=Str, "sea")), boost=0))"""
)
result = index.searcher().search(boosted_query, 10)
for _score, _ in result.hits:
# the score should be 0.0
assert _score == pytest.approx(0.0)
boosted_query = Query.boost_query(
Query.boost_query(
query1, 0.1
), 0.1
)
# Check for nested boost queries
assert(
repr(boosted_query)
== """Query(Boost(query=Boost(query=TermQuery(Term(field=0, type=Str, "sea")), boost=0.1), boost=0.1))"""
)
result = index.searcher().search(boosted_query, 10)
for _score, _ in result.hits:
# the score should be very small, due to
# the unknown score of BM25, we can only check for the relative difference
assert _score == pytest.approx(0.01, rel = 1)
boosted_query = Query.boost_query(
query1, -0.1
)
# Check for negative boost values
assert(
repr(boosted_query)
== """Query(Boost(query=TermQuery(Term(field=0, type=Str, "sea")), boost=-0.1))"""
)
result = index.searcher().search(boosted_query, 10)
# Even with a negative boost, the query should still match the document
assert len(result.hits) == 1
titles = set()
for _score, doc_address in result.hits:
# the score should be negative
assert _score < 0
titles.update(index.searcher().doc(doc_address)["title"])
assert titles == {"The Old Man and the Sea"}
# wrong query type
with pytest.raises(TypeError, match = r"'int' object cannot be converted to 'Query'"):
Query.boost_query(1, 0.1)
# wrong boost type
with pytest.raises(TypeError, match = r"argument 'boost': must be real number, not str"):
Query.boost_query(query1, "0.1")
# no boost type error
with pytest.raises(TypeError, match = r"Query.boost_query\(\) missing 1 required positional argument: 'boost'"):
Query.boost_query(query1)