fix: handle score_threshold in Formula queries (#8097)

* handle `score_threshold` in Formula queries

* AI: add openapi test

prompt: Upload 4 points with a numeric payload, use the payload as the
score, and set a score threshold. We can assert which ids should be in
the result and which shouldn't

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<Vec<ScoredPoint>>,
         limit: usize,
+        score_threshold: Option<ScoreType>,
         timeout: Duration,
         hw_measurement_acc: HwMeasurementAcc,
     ) -> CollectionResult<Vec<ScoredPoint>> {
@@ -25,6 +27,7 @@ impl LocalShard {
         let ctx = FormulaContext {
             formula,
             prefetches_results,
+            score_threshold,
             limit,
             is_stopped: stopping_guard.get_is_stopped(),
         };

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 => {

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<Vec<ScoredPoint>>,
         limit: usize,
+        score_threshold: Option<ScoreType>,
         hw_measurement_acc: HwMeasurementAcc,
     ) -> OperationResult<Vec<ScoredPoint>> {
         let ctx = FormulaContext {
             formula,
             prefetches_results,
             limit,
+            score_threshold,
             is_stopped: Arc::new(AtomicBool::new(false)),
         };
 

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<Vec<ScoredPoint>>,
     pub limit: usize,
+    pub score_threshold: Option<ScoreType>,
     pub is_stopped: Arc<AtomicBool>,
 }

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,
         )?;

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<ScoredPoint>],
         limit: usize,
+        score_threshold: Option<ScoreType>,
         is_stopped: &AtomicBool,
         hw_counter: &HardwareCounterCell,
     ) -> OperationResult<Vec<ScoredPointOffset>> {
@@ -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);

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}"
+        )