diff --git a/lib/collection/src/shards/local_shard/formula_rescore.rs b/lib/collection/src/shards/local_shard/formula_rescore.rs index 70c0f291c5..e7fc146112 100644 --- a/lib/collection/src/shards/local_shard/formula_rescore.rs +++ b/lib/collection/src/shards/local_shard/formula_rescore.rs @@ -2,6 +2,7 @@ use std::sync::Arc; use std::time::Duration; use common::counter::hardware_accumulator::HwMeasurementAcc; +use common::types::ScoreType; use segment::data_types::query_context::FormulaContext; use segment::index::query_optimization::rescore_formula::parsed_formula::ParsedFormula; use segment::types::ScoredPoint; @@ -17,6 +18,7 @@ impl LocalShard { formula: ParsedFormula, prefetches_results: Vec>, limit: usize, + score_threshold: Option, timeout: Duration, hw_measurement_acc: HwMeasurementAcc, ) -> CollectionResult> { @@ -25,6 +27,7 @@ impl LocalShard { let ctx = FormulaContext { formula, prefetches_results, + score_threshold, limit, is_stopped: stopping_guard.get_is_stopped(), }; diff --git a/lib/collection/src/shards/local_shard/query.rs b/lib/collection/src/shards/local_shard/query.rs index cb0ba59deb..fe7fb8ec21 100644 --- a/lib/collection/src/shards/local_shard/query.rs +++ b/lib/collection/src/shards/local_shard/query.rs @@ -367,8 +367,15 @@ impl LocalShard { }) } ScoringQuery::Formula(formula) => { - self.rescore_with_formula(formula, sources, limit, timeout, hw_counter_acc) - .await + self.rescore_with_formula( + formula, + sources, + limit, + score_threshold.map(OrderedFloat::into_inner), + timeout, + hw_counter_acc, + ) + .await } ScoringQuery::Sample(sample) => match sample { SampleInternal::Random => { diff --git a/lib/edge/src/query.rs b/lib/edge/src/query.rs index 1e66df8b08..e7e89b64d6 100644 --- a/lib/edge/src/query.rs +++ b/lib/edge/src/query.rs @@ -4,6 +4,7 @@ use std::sync::atomic::AtomicBool; use ahash::AHashSet; use common::counter::hardware_accumulator::HwMeasurementAcc; +use common::types::ScoreType; use ordered_float::OrderedFloat; use segment::common::operation_error::{OperationError, OperationResult}; use segment::common::reciprocal_rank_fusion::rrf_scoring; @@ -243,9 +244,13 @@ impl EdgeShard { self.search(search_request) } - ScoringQuery::Formula(formula) => { - self.rescore_with_formula(formula, sources, limit, hw_counter_acc) - } + ScoringQuery::Formula(formula) => self.rescore_with_formula( + formula, + sources, + limit, + score_threshold.map(OrderedFloat::into_inner), + hw_counter_acc, + ), ScoringQuery::Sample(sample) => match sample { SampleInternal::Random => { @@ -303,12 +308,14 @@ impl EdgeShard { formula: ParsedFormula, prefetches_results: Vec>, limit: usize, + score_threshold: Option, hw_measurement_acc: HwMeasurementAcc, ) -> OperationResult> { let ctx = FormulaContext { formula, prefetches_results, limit, + score_threshold, is_stopped: Arc::new(AtomicBool::new(false)), }; diff --git a/lib/segment/src/data_types/query_context.rs b/lib/segment/src/data_types/query_context.rs index 92c68ebad8..65d1715a91 100644 --- a/lib/segment/src/data_types/query_context.rs +++ b/lib/segment/src/data_types/query_context.rs @@ -6,6 +6,7 @@ use bitvec::prelude::BitSlice; use common::counter::hardware_accumulator::HwMeasurementAcc; use common::counter::hardware_counter::HardwareCounterCell; use common::cow::SimpleCow; +use common::types::ScoreType; use sparse::common::types::{DimId, DimWeight}; use crate::data_types::tiny_map; @@ -260,5 +261,6 @@ pub struct FormulaContext { pub formula: ParsedFormula, pub prefetches_results: Vec>, pub limit: usize, + pub score_threshold: Option, pub is_stopped: Arc, } diff --git a/lib/segment/src/segment/entry.rs b/lib/segment/src/segment/entry.rs index 962e476bc1..e5188f0183 100644 --- a/lib/segment/src/segment/entry.rs +++ b/lib/segment/src/segment/entry.rs @@ -110,6 +110,7 @@ impl NonAppendableSegmentEntry for Segment { formula, prefetches_results, limit, + score_threshold, is_stopped, } = &*ctx; @@ -117,6 +118,7 @@ impl NonAppendableSegmentEntry for Segment { formula, prefetches_results, *limit, + *score_threshold, is_stopped, hw_counter, )?; diff --git a/lib/segment/src/segment/formula_rescore.rs b/lib/segment/src/segment/formula_rescore.rs index ef1586c48e..09accb1a9c 100644 --- a/lib/segment/src/segment/formula_rescore.rs +++ b/lib/segment/src/segment/formula_rescore.rs @@ -3,8 +3,8 @@ use std::sync::atomic::{AtomicBool, Ordering}; use ahash::{AHashMap, AHashSet}; use common::counter::hardware_counter::HardwareCounterCell; use common::iterator_ext::IteratorExt; -use common::types::ScoredPointOffset; -use itertools::Itertools; +use common::types::{ScoreType, ScoredPointOffset}; +use itertools::{Either, Itertools}; use super::Segment; use crate::common::operation_error::OperationResult; @@ -18,6 +18,7 @@ impl Segment { formula: &ParsedFormula, prefetches_scores: &[Vec], limit: usize, + score_threshold: Option, is_stopped: &AtomicBool, hw_counter: &HardwareCounterCell, ) -> OperationResult> { @@ -49,7 +50,7 @@ impl Segment { // Perform rescoring let mut error = None; - let rescored = points_to_rescore + let rescored_iter = points_to_rescore .into_iter() .stop_if(is_stopped) .filter_map(|internal_id| { @@ -65,10 +66,17 @@ impl Segment { None } } - }) - // Keep only the top k results - .k_largest(limit) - .collect(); + }); + + // Handle score threshold + let rescored = match score_threshold { + Some(threshold) => { + Either::Left(rescored_iter.filter(move |point| point.score >= threshold)) + } + None => Either::Right(rescored_iter), + } + .k_largest(limit) // Keep only the top k results + .collect(); if let Some(err) = error { return Err(err); diff --git a/tests/openapi/test_query_formula.py b/tests/openapi/test_query_formula.py index 563d595b4a..57d9cbe6eb 100644 --- a/tests/openapi/test_query_formula.py +++ b/tests/openapi/test_query_formula.py @@ -107,3 +107,56 @@ def test_formula(collection_name, formula, expecting): # Assert that the response contains all points assert len(points) == len(orig_scores), "Response should contain all points" + +def test_formula_with_score_threshold(collection_name): + # Insert 4 test points with numeric payload "price" and a group tag + points = [ + {"id": 1001, "vector": [0.1, 0.1, 0.1, 0.1], "payload": {"price": 0.1, "group": "threshold_test"}}, + {"id": 1002, "vector": [0.2, 0.2, 0.2, 0.2], "payload": {"price": 0.6, "group": "threshold_test"}}, + {"id": 1003, "vector": [0.3, 0.3, 0.3, 0.3], "payload": {"price": 0.4, "group": "threshold_test"}}, + {"id": 1004, "vector": [0.4, 0.4, 0.4, 0.4], "payload": {"price": 0.9, "group": "threshold_test"}}, + ] + + response = request_with_validation( + api="/collections/{collection_name}/points", + method="PUT", + path_params={"collection_name": collection_name}, + query_params={"wait": "true"}, + body={"points": points}, + ) + assert response.ok, response.json() + + # Use a formula that sets the score equal to the payload "price" + formula = "price" + + # Set a threshold of 0.5: expect ids with price >= 0.5 (1002 and 1004) + threshold = 0.5 + expected_ids = {1002, 1004} + query = { + "prefetch": {"limit": 4}, + "query": {"formula": formula, "defaults": {"price": 0.0}}, + "filter": {"must": [{"key": "group", "match": {"value": "threshold_test"}}]}, + "with_payload": True, + "limit": 10, + "score_threshold": threshold, + } + + response = request_with_validation( + api="/collections/{collection_name}/points/query", + method="POST", + path_params={"collection_name": collection_name}, + body=query, + ) + assert response.ok, response.json() + + points_resp = response.json()["result"]["points"] + returned_ids = {p.get("id") for p in points_resp} + + # Assert returned ids match expected set + assert returned_ids == expected_ids, f"Expected ids {expected_ids}, got {returned_ids}" + + # Also assert each returned point has score >= threshold + for p in points_resp: + assert p.get("score") >= threshold - 1e-8, ( + f"Point {p.get('id')} with score {p.get('score')} is below threshold {threshold}" + )