use super::super::rerank_test_support as support;
use std::num::NonZeroU64;
use serde_json::{Value, json};
use support::{Fixture, TestResult, patients, request, values};
const FOUR: NonZeroU64 = NonZeroU64::MIN.saturating_add(3);
const PIPELINE: &str =
"patients | similar(.a; record(\"p0\"); dot) | limit(20) | rerank(.b; record(\"p0\"); dot)";
fn fixture() -> Result<Fixture, Box<dyn std::error::Error>> {
let mut fixture = Fixture::new(None, &patients())?;
fixture.vectors(
"a",
&[
(0, [1.0, 0.0]),
(1, [4.0, 0.0]),
(2, [3.0, 0.0]),
(3, [2.0, 0.0]),
],
)?;
fixture.vectors("b", &[(0, [1.0, 0.0]), (1, [2.0, 0.0]), (3, [5.0, 0.0])])?;
Ok(fixture)
}
fn ids(answer: &Value) -> Result<Vec<&str>, Box<dyn std::error::Error>> {
answer
.get("rows")
.and_then(Value::as_array)
.ok_or("missing rows")?
.iter()
.map(|row| {
row.get("origin")
.and_then(|origin| origin.get("id"))
.and_then(Value::as_str)
.ok_or_else(|| "missing origin id".into())
})
.collect()
}
#[test]
fn rerank_after_similar_and_limit_orders_by_the_second_field() -> TestResult {
let fixture = fixture()?;
let answer = fixture.run(&request(PIPELINE, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p3", "p1", "p0", "p2"]);
assert_eq!(
answer.get("unscored").ok_or("missing unscored count")?,
&json!(1)
);
assert_eq!(answer.get("next").ok_or("missing next")?, &Value::Null);
let rows = answer
.get("rows")
.and_then(Value::as_array)
.ok_or("missing rows")?;
let scores = rows
.iter()
.map(|row| {
row.get("origin")
.and_then(|origin| origin.get("score"))
.ok_or("missing score")
})
.collect::<Result<Vec<_>, _>>()?;
assert_eq!(
scores,
[&json!(5.0), &json!(2.0), &json!(1.0), &Value::Null]
);
Ok(())
}
#[test]
fn jq_score_binding_is_the_rerank_score_including_null() -> TestResult {
let fixture = fixture()?;
let answer = fixture.run(&request(&format!("{PIPELINE} | '$score'"), FOUR), &[])?;
assert_eq!(
values(&answer)?,
[json!(5.0), json!(2.0), json!(1.0), Value::Null]
);
assert_eq!(ids(&answer)?, ["p3", "p1", "p0", "p2"]);
assert_eq!(answer.get("unscored"), Some(&json!(1)));
Ok(())
}
#[test]
fn rerank_drains_only_records_before_its_upstream_limit() -> TestResult {
let fixture = fixture()?;
let query =
"patients | similar(.a; record(\"p0\"); dot) | limit(2) | rerank(.b; record(\"p0\"); dot)";
let answer = fixture.run(&request(query, FOUR), &[])?;
assert_eq!(ids(&answer)?, ["p1", "p2"]);
assert_eq!(answer.get("unscored"), Some(&json!(1)));
assert_eq!(answer.get("next"), Some(&Value::Null));
Ok(())
}
#[test]
fn rerank_refuses_past_the_scan_limit_and_names_it() -> TestResult {
let mut fixture = fixture()?;
fixture.scan(NonZeroU64::MIN.saturating_add(1));
let error = fixture
.run(
&request("patients | rerank(.b; record(\"p0\"); dot)", FOUR),
&[],
)
.err()
.ok_or("partial rerank escaped scan limit")?;
assert_eq!(error.code, "query_scan_limit");
assert!(error.message.contains('2'), "{}", error.message);
Ok(())
}
#[test]
fn rerank_refuses_an_output_page_cut_instead_of_returning_next() -> TestResult {
let fixture = fixture()?;
let error = fixture
.run(&request(PIPELINE, NonZeroU64::MIN), &[])
.err()
.ok_or("rerank returned a partial page")?;
assert_eq!(error.code, "query_output_limit");
Ok(())
}
#[test]
fn rerank_cannot_read_a_denied_question_record() -> TestResult {
let fixture = fixture()?;
let error = fixture
.run(
&request("patients | rerank(.b; record(\"p0\"); dot)", FOUR),
&[("patients", "p0")],
)
.err()
.ok_or("denied question record was read")?;
assert_eq!(error.code, "query_search_record_missing");
Ok(())
}
#[test]
fn rerank_l2_orders_the_nearest_vector_first() -> TestResult {
let fixture = fixture()?;
let answer = fixture.run(
&request("patients | rerank(.b; record(\"p0\"); l2)", FOUR),
&[],
)?;
assert_eq!(ids(&answer)?, ["p0", "p1", "p3", "p2"]);
assert_eq!(answer.get("next"), Some(&Value::Null));
Ok(())
}
#[test]
fn rerank_requires_the_query_worker() -> TestResult {
let mut fixture = fixture()?;
fixture.no_worker();
let error = fixture
.run(&request(PIPELINE, FOUR), &[])
.err()
.ok_or("missing query worker was accepted")?;
assert_eq!(error.code, "query_engine_unavailable");
Ok(())
}