use std::collections::BTreeSet;
use std::error::Error;
use std::fmt::{self, Display, Formatter};
use lodestone_types::{
DenseVectorValue, MAX_SPARSE_VECTOR_DIMENSION, MAX_VECTOR_BYTES_PER_QUERY, MAX_VECTOR_BYTES_PER_REQUEST,
MAX_VECTOR_SET_BYTES_PER_QUERY, MAX_VECTOR_SET_QUERY_VECTORS, SparseVectorDefinition, SparseVectorValue,
SparseVectorWeightType, VectorSetValue, is_canonical_field_path,
};
use serde::{Deserialize, Deserializer, Serialize};
use crate::Query;
pub const RETRIEVAL_FORMAT_VERSION: u16 = 1;
pub const MAX_RETRIEVAL_BRANCHES: usize = 16;
pub const MAX_RETRIEVAL_BRANCH_NAME_BYTES: usize = 64;
pub const MAX_RETRIEVAL_WINDOW: usize = 10_000;
pub const MAX_RETRIEVAL_CANDIDATES: usize = 10_000;
pub const MIN_RETRIEVAL_RRF_CONSTANT: u16 = 1;
pub const MAX_RETRIEVAL_RRF_CONSTANT: u16 = 1_000;
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(rename_all = "snake_case")]
pub enum EngineDenseRetrievalMode {
Auto,
Ann,
Exact,
}
#[derive(Clone, Copy, Debug, Deserialize, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum EngineDenseRetrievalThreshold {
MinimumScore {
#[serde(deserialize_with = "deserialize_f64")]
value: f64,
},
MaximumDistance {
#[serde(deserialize_with = "deserialize_f64")]
value: f64,
},
}
impl EngineDenseRetrievalThreshold {
fn validate(self) -> Result<(), RetrievalPlanError> {
match self {
Self::MinimumScore { value } if value.is_finite() => Ok(()),
Self::MaximumDistance { value } if value.is_finite() && value >= 0.0 => Ok(()),
Self::MinimumScore { .. } => Err(RetrievalPlanError::InvalidMinimumScore),
Self::MaximumDistance { .. } => Err(RetrievalPlanError::InvalidMaximumDistance),
}
}
}
#[derive(Clone, Debug, Deserialize, PartialEq, Serialize)]
#[serde(tag = "kind", rename_all = "snake_case", deny_unknown_fields)]
pub enum EngineRetrievalBranch {
Lexical {
name: String,
query: Query,
#[serde(default, skip_serializing_if = "Option::is_none")]
filter: Option<Query>,
window: usize,
#[serde(deserialize_with = "deserialize_f64")]
weight: f64,
},
Dense {
name: String,
field: String,
query_vector: DenseVectorValue,
#[serde(default, skip_serializing_if = "Option::is_none")]
filter: Option<Query>,
window: usize,
#[serde(default, skip_serializing_if = "Option::is_none")]
candidates: Option<usize>,
#[serde(deserialize_with = "deserialize_f64")]
weight: f64,
mode: EngineDenseRetrievalMode,
#[serde(default, skip_serializing_if = "Option::is_none")]
threshold: Option<EngineDenseRetrievalThreshold>,
},
Sparse {
name: String,
field: String,
#[serde(deserialize_with = "deserialize_sparse_vector")]
query_vector: SparseVectorValue,
#[serde(default, skip_serializing_if = "Option::is_none")]
filter: Option<Query>,
window: usize,
#[serde(deserialize_with = "deserialize_f64")]
weight: f64,
},
LateInteraction {
name: String,
field: String,
query_vectors: VectorSetValue,
#[serde(default, skip_serializing_if = "Option::is_none")]
filter: Option<Query>,
window: usize,
candidates: usize,
#[serde(deserialize_with = "deserialize_f64")]
weight: f64,
},
}
impl EngineRetrievalBranch {
#[must_use]
pub fn name(&self) -> &str {
match self {
Self::Lexical { name, .. }
| Self::Dense { name, .. }
| Self::Sparse { name, .. }
| Self::LateInteraction { name, .. } => name,
}
}
#[must_use]
pub const fn window(&self) -> usize {
match self {
Self::Lexical { window, .. }
| Self::Dense { window, .. }
| Self::Sparse { window, .. }
| Self::LateInteraction { window, .. } => *window,
}
}
#[must_use]
pub const fn weight(&self) -> f64 {
match self {
Self::Lexical { weight, .. }
| Self::Dense { weight, .. }
| Self::Sparse { weight, .. }
| Self::LateInteraction { weight, .. } => *weight,
}
}
#[must_use]
pub const fn filter(&self) -> Option<&Query> {
match self {
Self::Lexical { filter, .. }
| Self::Dense { filter, .. }
| Self::Sparse { filter, .. }
| Self::LateInteraction { filter, .. } => filter.as_ref(),
}
}
fn validate(&self, branch_ordinal: usize) -> Result<(), RetrievalPlanError> {
validate_branch_name(self.name(), branch_ordinal)?;
if !(1..=MAX_RETRIEVAL_WINDOW).contains(&self.window()) {
return Err(RetrievalPlanError::InvalidBranchWindow {
branch_ordinal,
observed: self.window(),
});
}
if !self.weight().is_finite() || self.weight() <= 0.0 {
return Err(RetrievalPlanError::InvalidBranchWeight(branch_ordinal));
}
if let Self::Dense {
field,
query_vector,
window,
candidates,
mode,
threshold,
..
} = self
{
if !is_canonical_field_path(field) {
return Err(RetrievalPlanError::InvalidDenseField(branch_ordinal));
}
query_vector
.validate_bounded(MAX_VECTOR_BYTES_PER_QUERY)
.map_err(|_| RetrievalPlanError::InvalidDenseValue(branch_ordinal))?;
if candidates.is_some_and(|value| value < *window || value > MAX_RETRIEVAL_CANDIDATES) {
return Err(RetrievalPlanError::InvalidDenseCandidates {
branch_ordinal,
observed: candidates.unwrap_or_default(),
window: *window,
});
}
if *mode == EngineDenseRetrievalMode::Exact && candidates.is_some() {
return Err(RetrievalPlanError::ExactCandidates(branch_ordinal));
}
if let Some(threshold) = threshold {
threshold.validate()?;
}
}
if let Self::Sparse {
field, query_vector, ..
} = self
{
if !is_canonical_field_path(field) {
return Err(RetrievalPlanError::InvalidSparseField(branch_ordinal));
}
query_vector
.validate_bounded(
SparseVectorDefinition {
maximum_dimension: MAX_SPARSE_VECTOR_DIMENSION,
weight_type: SparseVectorWeightType::F32,
},
MAX_VECTOR_BYTES_PER_QUERY,
)
.map_err(|_| RetrievalPlanError::InvalidSparseValue(branch_ordinal))?;
}
if let Self::LateInteraction {
field,
query_vectors,
window,
candidates,
..
} = self
{
if !is_canonical_field_path(field) {
return Err(RetrievalPlanError::InvalidLateInteractionField(branch_ordinal));
}
if *candidates < *window || *candidates > MAX_RETRIEVAL_CANDIDATES {
return Err(RetrievalPlanError::InvalidLateInteractionCandidates {
branch_ordinal,
observed: *candidates,
window: *window,
});
}
validate_query_vector_set(query_vectors)
.map_err(|_| RetrievalPlanError::InvalidLateInteractionValue(branch_ordinal))?;
}
Ok(())
}
}
#[derive(Clone, Copy, Debug, Deserialize, Eq, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct EngineWeightedRrfFusion {
pub rank_constant: u16,
}
#[derive(Clone, Debug, PartialEq, Serialize)]
#[serde(deny_unknown_fields)]
pub struct EngineRetrievalPlan {
pub format_version: u16,
#[serde(default, skip_serializing_if = "Option::is_none")]
pub global_filter: Option<Query>,
pub branches: Vec<EngineRetrievalBranch>,
pub fusion: EngineWeightedRrfFusion,
pub result_window: usize,
}
impl<'de> Deserialize<'de> for EngineRetrievalPlan {
fn deserialize<D>(deserializer: D) -> Result<Self, D::Error>
where
D: Deserializer<'de>,
{
#[derive(Deserialize)]
#[serde(deny_unknown_fields)]
struct Wire {
format_version: u16,
#[serde(default)]
global_filter: Option<Query>,
branches: Vec<serde_json::Value>,
fusion: EngineWeightedRrfFusion,
result_window: usize,
}
let wire = Wire::deserialize(deserializer)?;
let branches = wire
.branches
.into_iter()
.map(|branch| serde_json::from_value(branch).map_err(serde::de::Error::custom))
.collect::<Result<Vec<_>, D::Error>>()?;
Ok(Self {
format_version: wire.format_version,
global_filter: wire.global_filter,
branches,
fusion: wire.fusion,
result_window: wire.result_window,
})
}
}
impl EngineRetrievalPlan {
pub fn validate(&self) -> Result<(), RetrievalPlanError> {
if self.format_version != RETRIEVAL_FORMAT_VERSION {
return Err(RetrievalPlanError::UnsupportedVersion(self.format_version));
}
if self.branches.is_empty() || self.branches.len() > MAX_RETRIEVAL_BRANCHES {
return Err(RetrievalPlanError::InvalidBranchCount(self.branches.len()));
}
if !(MIN_RETRIEVAL_RRF_CONSTANT..=MAX_RETRIEVAL_RRF_CONSTANT).contains(&self.fusion.rank_constant) {
return Err(RetrievalPlanError::InvalidRrfConstant(self.fusion.rank_constant));
}
if !(1..=MAX_RETRIEVAL_WINDOW).contains(&self.result_window) {
return Err(RetrievalPlanError::InvalidResultWindow(self.result_window));
}
let mut names = BTreeSet::new();
let mut branch_work = 0usize;
let mut vector_bytes = 0usize;
for (branch_ordinal, branch) in self.branches.iter().enumerate() {
branch.validate(branch_ordinal)?;
if !names.insert(branch.name()) {
return Err(RetrievalPlanError::DuplicateBranchName(branch_ordinal));
}
branch_work = branch_work
.checked_add(branch.window())
.ok_or(RetrievalPlanError::BranchWorkOverflow)?;
let branch_vector_bytes = match branch {
EngineRetrievalBranch::Dense { query_vector, .. } => query_vector
.validate_bounded(MAX_VECTOR_BYTES_PER_QUERY)
.map_err(|_| RetrievalPlanError::InvalidDenseValue(branch_ordinal))?,
EngineRetrievalBranch::Sparse { query_vector, .. } => query_vector
.validate_bounded(
SparseVectorDefinition {
maximum_dimension: MAX_SPARSE_VECTOR_DIMENSION,
weight_type: SparseVectorWeightType::F32,
},
MAX_VECTOR_BYTES_PER_QUERY,
)
.map_err(|_| RetrievalPlanError::InvalidSparseValue(branch_ordinal))?,
EngineRetrievalBranch::LateInteraction { query_vectors, .. } => {
validate_query_vector_set(query_vectors)
.map_err(|_| RetrievalPlanError::InvalidLateInteractionValue(branch_ordinal))?
}
EngineRetrievalBranch::Lexical { .. } => 0,
};
vector_bytes = vector_bytes
.checked_add(branch_vector_bytes)
.ok_or(RetrievalPlanError::VectorRequestBytesExceeded)?;
}
if branch_work > MAX_RETRIEVAL_BRANCHES * MAX_RETRIEVAL_WINDOW {
return Err(RetrievalPlanError::BranchWorkOverflow);
}
if vector_bytes > MAX_VECTOR_BYTES_PER_REQUEST {
return Err(RetrievalPlanError::VectorRequestBytesExceeded);
}
Ok(())
}
pub fn dense_branches(&self) -> impl Iterator<Item = &EngineRetrievalBranch> {
self.branches
.iter()
.filter(|branch| matches!(branch, EngineRetrievalBranch::Dense { .. }))
}
}
fn validate_query_vector_set(value: &VectorSetValue) -> Result<usize, ()> {
if value.vectors.is_empty() || value.vectors.len() > MAX_VECTOR_SET_QUERY_VECTORS {
return Err(());
}
let mut bytes = 0_usize;
let mut dimensions = None;
for vector in &value.vectors {
let vector_bytes = vector.validate_bounded(MAX_VECTOR_BYTES_PER_QUERY).map_err(|_| ())?;
let vector_dimensions = vector.element_count();
if dimensions.is_some_and(|expected| vector_dimensions != Some(expected)) {
return Err(());
}
if dimensions.is_none() {
dimensions = vector_dimensions;
}
bytes = bytes.checked_add(vector_bytes).ok_or(())?;
}
(bytes <= MAX_VECTOR_SET_BYTES_PER_QUERY).then_some(bytes).ok_or(())
}
fn deserialize_f64<'de, D>(deserializer: D) -> Result<f64, D::Error>
where
D: Deserializer<'de>,
{
serde_json::Number::deserialize(deserializer)
.map_err(|error| serde::de::Error::custom(format!("invalid retrieval number: {error}")))?
.to_string()
.parse::<f64>()
.map_err(serde::de::Error::custom)
}
fn deserialize_sparse_vector<'de, D>(deserializer: D) -> Result<SparseVectorValue, D::Error>
where
D: Deserializer<'de>,
{
SparseVectorValue::deserialize(deserializer)
.map_err(|error| serde::de::Error::custom(format!("invalid sparse vector: {error}")))
}
fn validate_branch_name(name: &str, branch_ordinal: usize) -> Result<(), RetrievalPlanError> {
if name.is_empty()
|| name.len() > MAX_RETRIEVAL_BRANCH_NAME_BYTES
|| !name
.bytes()
.all(|byte| byte.is_ascii_alphanumeric() || matches!(byte, b'_' | b'-'))
{
return Err(RetrievalPlanError::InvalidBranchName(branch_ordinal));
}
Ok(())
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub enum RetrievalPlanError {
UnsupportedVersion(u16),
InvalidBranchCount(usize),
InvalidBranchName(usize),
DuplicateBranchName(usize),
InvalidBranchWindow {
branch_ordinal: usize,
observed: usize,
},
InvalidBranchWeight(usize),
InvalidDenseField(usize),
InvalidDenseValue(usize),
InvalidSparseField(usize),
InvalidSparseValue(usize),
InvalidLateInteractionField(usize),
InvalidLateInteractionValue(usize),
InvalidDenseCandidates {
branch_ordinal: usize,
observed: usize,
window: usize,
},
ExactCandidates(usize),
InvalidLateInteractionCandidates {
branch_ordinal: usize,
observed: usize,
window: usize,
},
InvalidMinimumScore,
InvalidMaximumDistance,
InvalidRrfConstant(u16),
InvalidResultWindow(usize),
BranchWorkOverflow,
VectorRequestBytesExceeded,
}
impl Display for RetrievalPlanError {
fn fmt(&self, formatter: &mut Formatter<'_>) -> fmt::Result {
match self {
Self::UnsupportedVersion(version) => write!(formatter, "unsupported retrieval format version {version}"),
Self::InvalidBranchCount(observed) => write!(
formatter,
"retrieval branch count {observed} is outside 1..={MAX_RETRIEVAL_BRANCHES}"
),
Self::InvalidBranchName(branch) => write!(formatter, "retrieval branch {branch} name is invalid"),
Self::DuplicateBranchName(branch) => {
write!(formatter, "retrieval branch {branch} repeats an earlier branch name")
}
Self::InvalidBranchWindow {
branch_ordinal,
observed,
} => write!(
formatter,
"retrieval branch {branch_ordinal} window {observed} is outside 1..={MAX_RETRIEVAL_WINDOW}"
),
Self::InvalidBranchWeight(branch) => {
write!(formatter, "retrieval branch {branch} weight is not positive and finite")
}
Self::InvalidDenseField(branch) => {
write!(formatter, "dense retrieval branch {branch} field is not canonical")
}
Self::InvalidDenseValue(branch) => {
write!(
formatter,
"dense retrieval branch {branch} has an invalid or oversized vector value"
)
}
Self::InvalidSparseField(branch) => {
write!(formatter, "sparse retrieval branch {branch} field is not canonical")
}
Self::InvalidSparseValue(branch) => {
write!(
formatter,
"sparse retrieval branch {branch} has an invalid or oversized vector value"
)
}
Self::InvalidLateInteractionField(branch) => {
write!(
formatter,
"late-interaction retrieval branch {branch} field is not canonical"
)
}
Self::InvalidLateInteractionValue(branch) => {
write!(
formatter,
"late-interaction retrieval branch {branch} has an invalid or oversized vector set"
)
}
Self::InvalidDenseCandidates {
branch_ordinal,
observed,
window,
} => write!(
formatter,
"dense retrieval branch {branch_ordinal} candidates {observed} are below window {window} or above {MAX_RETRIEVAL_CANDIDATES}"
),
Self::ExactCandidates(branch) => {
write!(
formatter,
"exact dense retrieval branch {branch} cannot specify ANN candidate effort"
)
}
Self::InvalidLateInteractionCandidates {
branch_ordinal,
observed,
window,
} => write!(
formatter,
"late-interaction retrieval branch {branch_ordinal} candidates {observed} are below window {window} or above {MAX_RETRIEVAL_CANDIDATES}"
),
Self::InvalidMinimumScore => formatter.write_str("retrieval minimum score is not finite"),
Self::InvalidMaximumDistance => {
formatter.write_str("retrieval maximum distance is not finite and nonnegative")
}
Self::InvalidRrfConstant(observed) => write!(
formatter,
"retrieval RRF constant {observed} is outside {MIN_RETRIEVAL_RRF_CONSTANT}..={MAX_RETRIEVAL_RRF_CONSTANT}"
),
Self::InvalidResultWindow(observed) => write!(
formatter,
"retrieval result window {observed} is outside 1..={MAX_RETRIEVAL_WINDOW}"
),
Self::BranchWorkOverflow => formatter.write_str("retrieval aggregate branch work exceeds its hard bound"),
Self::VectorRequestBytesExceeded => {
formatter.write_str("retrieval aggregate query vector bytes exceed the request ceiling")
}
}
}
}
impl Error for RetrievalPlanError {}
#[cfg(test)]
mod tests {
use super::*;
fn dense(name: &str, mode: EngineDenseRetrievalMode) -> EngineRetrievalBranch {
EngineRetrievalBranch::Dense {
name: name.to_owned(),
field: "embedding".to_owned(),
query_vector: DenseVectorValue::Numeric { values: vec![1.0, 0.0] },
filter: None,
window: 10,
candidates: (mode != EngineDenseRetrievalMode::Exact).then_some(64),
weight: 1.0,
mode,
threshold: None,
}
}
fn sparse(name: &str) -> EngineRetrievalBranch {
EngineRetrievalBranch::Sparse {
name: name.to_owned(),
field: "terms".to_owned(),
query_vector: SparseVectorValue {
entries: vec![
lodestone_types::SparseVectorEntry {
dimension: 1,
weight: 2.0,
},
lodestone_types::SparseVectorEntry {
dimension: 8,
weight: -0.5,
},
],
},
filter: None,
window: 10,
weight: 0.5,
}
}
fn late_interaction(name: &str) -> EngineRetrievalBranch {
EngineRetrievalBranch::LateInteraction {
name: name.to_owned(),
field: "tokens".to_owned(),
query_vectors: VectorSetValue {
vectors: vec![
DenseVectorValue::Numeric { values: vec![1.0, 0.0] },
DenseVectorValue::Numeric { values: vec![0.0, 1.0] },
],
},
filter: None,
window: 10,
candidates: 64,
weight: 0.75,
}
}
#[test]
fn retrieval_plan_accepts_named_lexical_dense_sparse_and_late_interaction_branches() -> Result<(), Box<dyn Error>> {
let sparse_branch = sparse("terms");
let sparse_bytes = serde_json::to_vec(&sparse_branch)?;
assert_eq!(
serde_json::from_slice::<EngineRetrievalBranch>(&sparse_bytes)?,
sparse_branch
);
let branch_vector = vec![
EngineRetrievalBranch::Lexical {
name: "text".to_owned(),
query: Query::MatchAll,
filter: None,
window: 20,
weight: 0.75,
},
dense("semantic", EngineDenseRetrievalMode::Ann),
sparse("terms"),
late_interaction("tokens"),
];
let plan = EngineRetrievalPlan {
format_version: RETRIEVAL_FORMAT_VERSION,
global_filter: Some(Query::Term {
field: "kind".to_owned(),
value: serde_json::json!("article"),
}),
branches: branch_vector,
fusion: EngineWeightedRrfFusion { rank_constant: 60 },
result_window: 20,
};
plan.validate()?;
let encoded = serde_json::to_vec(&plan)?;
let decoded: EngineRetrievalPlan = serde_json::from_slice(&encoded)?;
assert_eq!(decoded, plan);
Ok(())
}
#[test]
fn retrieval_plan_refuses_duplicate_names_nonfinite_values_and_exact_effort() {
let mut plan = EngineRetrievalPlan {
format_version: RETRIEVAL_FORMAT_VERSION,
global_filter: None,
branches: vec![
dense("same", EngineDenseRetrievalMode::Ann),
dense("same", EngineDenseRetrievalMode::Auto),
],
fusion: EngineWeightedRrfFusion { rank_constant: 60 },
result_window: 10,
};
assert_eq!(plan.validate(), Err(RetrievalPlanError::DuplicateBranchName(1)));
plan.branches = vec![dense("exact", EngineDenseRetrievalMode::Exact)];
if let EngineRetrievalBranch::Dense { candidates, .. } = &mut plan.branches[0] {
*candidates = Some(20);
}
assert_eq!(plan.validate(), Err(RetrievalPlanError::ExactCandidates(0)));
if let EngineRetrievalBranch::Dense {
candidates,
query_vector,
..
} = &mut plan.branches[0]
{
*candidates = None;
*query_vector = DenseVectorValue::Numeric { values: vec![f64::NAN] };
}
assert_eq!(plan.validate(), Err(RetrievalPlanError::InvalidDenseValue(0)));
plan.branches = vec![sparse("sparse")];
if let EngineRetrievalBranch::Sparse { query_vector, .. } = &mut plan.branches[0] {
query_vector.entries[1].dimension = 1;
}
assert_eq!(plan.validate(), Err(RetrievalPlanError::InvalidSparseValue(0)));
plan.branches = vec![late_interaction("tokens")];
if let EngineRetrievalBranch::LateInteraction {
query_vectors,
candidates,
..
} = &mut plan.branches[0]
{
*candidates = 1;
query_vectors.vectors[1] = DenseVectorValue::Numeric { values: vec![1.0] };
}
assert_eq!(
plan.validate(),
Err(RetrievalPlanError::InvalidLateInteractionCandidates {
branch_ordinal: 0,
observed: 1,
window: 10,
})
);
if let EngineRetrievalBranch::LateInteraction {
query_vectors,
candidates,
..
} = &mut plan.branches[0]
{
*candidates = 64;
query_vectors.vectors[1] = DenseVectorValue::Numeric { values: vec![1.0] };
}
assert_eq!(plan.validate(), Err(RetrievalPlanError::InvalidLateInteractionValue(0)));
}
}