use crate::actions::helpers;
use crate::actions::vector_index::{
hash_attr, parse_attr_defs, parse_vector_defs, scalar_type_str,
};
use crate::actions::vector_lifecycle::VectorIndexLifecycle;
use crate::errors::{DynoxideError, Result};
use crate::expressions;
use crate::expressions::PathElement;
use crate::expressions::condition::{CompOp, ConditionExpr, Operand};
use crate::storage_backend::{StorageBackend, VectorCandidateRow};
use crate::types::{AttributeValue, Item, VectorIndex};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
#[derive(Debug, Default, Deserialize)]
struct SearchVectorsRequestRaw {
#[serde(rename = "TableName", default)]
table_name: Option<String>,
#[serde(rename = "IndexName", default)]
index_name: Option<String>,
#[serde(rename = "SearchVector", default)]
search_vector: Option<Vec<AttributeValue>>,
#[serde(rename = "TopK", default)]
top_k: Option<i32>,
#[serde(rename = "SearchConditionExpression", default)]
search_condition_expression: Option<String>,
#[serde(rename = "ProjectionExpression", default)]
projection_expression: Option<String>,
#[serde(rename = "ExpressionAttributeNames", default)]
expression_attribute_names: Option<HashMap<String, String>>,
#[serde(rename = "ExpressionAttributeValues", default)]
expression_attribute_values: Option<HashMap<String, AttributeValue>>,
#[serde(rename = "ReturnConsumedCapacity", default)]
return_consumed_capacity: Option<String>,
}
#[derive(Debug, Default)]
pub struct SearchVectorsRequest {
pub table_name: String,
pub index_name: String,
pub search_vector: Vec<AttributeValue>,
pub top_k: i64,
pub search_condition_expression: Option<String>,
pub projection_expression: Option<String>,
pub expression_attribute_names: Option<HashMap<String, String>>,
pub expression_attribute_values: Option<HashMap<String, AttributeValue>>,
pub return_consumed_capacity: Option<String>,
}
impl<'de> serde::Deserialize<'de> for SearchVectorsRequest {
fn deserialize<D: serde::Deserializer<'de>>(
deserializer: D,
) -> std::result::Result<Self, D::Error> {
let raw = SearchVectorsRequestRaw::deserialize(deserializer)?;
use crate::validation::{
TableNameContext, format_validation_errors, table_name_constraint_errors,
};
let mut errors = Vec::new();
errors.extend(table_name_constraint_errors(
raw.table_name.as_deref(),
TableNameContext::ReadWrite,
));
match raw.index_name.as_deref() {
None => errors.push(
"Value null at 'indexName' failed to satisfy constraint: \
Member must not be null"
.to_string(),
),
Some(name) if name.len() < 3 => errors.push(
"Value at 'IndexName' failed to satisfy constraint: \
Member must have length greater than or equal to 3"
.to_string(),
),
Some(name) if name.len() > 255 => errors.push(
"Value at 'IndexName' failed to satisfy constraint: \
Member must have length less than or equal to 255"
.to_string(),
),
Some(_) => {}
}
match raw.search_vector.as_deref() {
None => errors.push(
"Value null at 'searchVector' failed to satisfy constraint: \
Member must not be null"
.to_string(),
),
Some([]) => errors.push(
"Value at 'SearchVector' failed to satisfy constraint: \
Member must have length greater than or equal to 1"
.to_string(),
),
Some(_) => {}
}
match raw.top_k {
None => errors.push(
"Value null at 'topK' failed to satisfy constraint: \
Member must not be null"
.to_string(),
),
Some(v) if v < 1 => errors.push(format!(
"Value '{v}' at 'topK' failed to satisfy constraint: \
Member must have value greater than or equal to 1"
)),
Some(_) => {}
}
if let Some(msg) = crate::validation::return_consumed_capacity_rejection(
raw.return_consumed_capacity.as_deref(),
) {
errors.push(msg);
}
if let Some(msg) = format_validation_errors(&errors) {
return Err(serde::de::Error::custom(format!("VALIDATION:{msg}")));
}
Ok(SearchVectorsRequest {
table_name: raw.table_name.unwrap_or_default(),
index_name: raw.index_name.unwrap_or_default(),
search_vector: raw.search_vector.unwrap_or_default(),
top_k: raw.top_k.map(i64::from).unwrap_or_default(),
search_condition_expression: raw.search_condition_expression,
projection_expression: raw.projection_expression,
expression_attribute_names: raw.expression_attribute_names,
expression_attribute_values: raw.expression_attribute_values,
return_consumed_capacity: raw.return_consumed_capacity,
})
}
}
#[derive(Debug, Default, Serialize)]
pub struct SearchVectorsResponse {
#[serde(rename = "SearchResults")]
pub search_results: Vec<SearchResult>,
#[serde(rename = "ConsumedCapacity", skip_serializing_if = "Option::is_none")]
pub consumed_capacity: Option<crate::types::VectorSearchCapacity>,
}
fn search_scanned_bytes(candidates: &[VectorCandidateRow]) -> f64 {
let scanned: i64 = candidates.iter().map(|row| row.entry_bytes).sum();
crate::types::vector_request_bytes(scanned.max(0) as usize)
}
#[derive(Debug, Serialize)]
pub struct SearchResult {
#[serde(rename = "Item")]
pub item: Item,
#[serde(rename = "Score")]
pub score: f64,
}
fn invalid_search_vector() -> DynoxideError {
DynoxideError::ValidationException(
"Search vector contains invalid values. All values in the search vector must be a \
32-bit floating-point number attribute"
.to_string(),
)
}
fn invalid_comparator() -> DynoxideError {
DynoxideError::ValidationException(
"Invalid SearchConditionExpression: Invalid comparator used in SearchConditionExpression"
.to_string(),
)
}
struct EqualityPair {
attr: String,
nested: bool,
value: AttributeValue,
}
fn collect_equalities(
expr: &ConditionExpr,
tracker: &expressions::TrackedExpressionAttributes,
pairs: &mut Vec<EqualityPair>,
attrs: &mut Vec<String>,
) -> Result<()> {
match expr {
ConditionExpr::And(left, right) => {
collect_equalities(left, tracker, pairs, attrs)?;
collect_equalities(right, tracker, pairs, attrs)?;
Ok(())
}
ConditionExpr::Comparison { left, op, right } => {
if *op != CompOp::Eq {
return Err(invalid_comparator());
}
let (path, value_ref) = match (left, right) {
(Operand::Path(p), Operand::ValueRef(v))
| (Operand::ValueRef(v), Operand::Path(p)) => (p, v),
_ => return Err(invalid_comparator()),
};
let top = match path.first() {
Some(PathElement::Attribute(name)) => tracker.resolve_name(name).map_err(|e| {
DynoxideError::ValidationException(format!(
"Invalid SearchConditionExpression: {e}"
))
})?,
_ => return Err(invalid_comparator()),
};
let value = tracker
.resolve_value(value_ref)
.map_err(|e| {
DynoxideError::ValidationException(format!(
"Invalid SearchConditionExpression: {e}"
))
})?
.clone();
attrs.push(top.clone());
pairs.push(EqualityPair {
attr: top,
nested: path.len() > 1,
value,
});
Ok(())
}
_ => Err(invalid_comparator()),
}
}
fn score_candidate(distance_function: &str, query: &[f32], stored: &[f32]) -> f64 {
let score = match distance_function {
"EUCLIDEAN" => {
let mut sum = 0.0f64;
for (q, s) in query.iter().zip(stored) {
let d = f64::from(*q) - f64::from(*s);
sum += d * d;
}
sum.sqrt()
}
"DOT_PRODUCT" => {
let mut dot = 0.0f64;
for (q, s) in query.iter().zip(stored) {
dot += f64::from(*q) * f64::from(*s);
}
dot
}
_ => {
let mut dot = 0.0f64;
let mut qq = 0.0f64;
let mut ss = 0.0f64;
for (q, s) in query.iter().zip(stored) {
dot += f64::from(*q) * f64::from(*s);
qq += f64::from(*q) * f64::from(*q);
ss += f64::from(*s) * f64::from(*s);
}
if qq == 0.0 || ss == 0.0 {
return 1.0;
}
if query == stored {
return 0.0;
}
let c = 1.0 - dot / (qq.sqrt() * ss.sqrt());
c.clamp(0.0, 2.0)
}
};
score.clamp(-f64::from(f32::MAX), f64::from(f32::MAX))
}
fn higher_is_better(distance_function: &str) -> bool {
distance_function == "DOT_PRODUCT"
}
fn parse_query_vector(search_vector: &[AttributeValue], vix: &VectorIndex) -> Result<Vec<f32>> {
let mut out = Vec::with_capacity(search_vector.len());
for elem in search_vector {
let AttributeValue::N(n) = elem else {
return Err(invalid_search_vector());
};
match n.parse::<f32>() {
Ok(v) if v.is_finite() => out.push(v),
_ => return Err(invalid_search_vector()),
}
}
if out.len() != vix.dimensions as usize {
return Err(DynoxideError::ValidationException(format!(
"Input search vector dimension {} does not match vector index dimension {}",
out.len(),
vix.dimensions
)));
}
Ok(out)
}
pub async fn execute<S: StorageBackend>(
storage: &S,
request: SearchVectorsRequest,
lifecycle: &VectorIndexLifecycle,
) -> Result<SearchVectorsResponse> {
crate::validation::validate_table_name(&request.table_name)?;
if !(1..=100).contains(&request.top_k) {
return Err(DynoxideError::ValidationException(format!(
"Provided TopK value '{}' is out of valid range. The value must be between 1 and \
100 inclusive",
request.top_k
)));
}
{
let mut expr = Vec::new();
if request.search_condition_expression.is_some() {
expr.push("SearchConditionExpression");
}
if request.projection_expression.is_some() {
expr.push("ProjectionExpression");
}
let no_raw_eav: Option<serde_json::Value> = None;
let ctx = helpers::ExpressionParamContext {
non_expression_params: Vec::new(),
expression_params: expr,
all_expression_param_names: vec!["SearchConditionExpression"],
expression_attribute_names: &request.expression_attribute_names,
expression_attribute_values: &request.expression_attribute_values,
expression_attribute_values_raw: &no_raw_eav,
};
helpers::validate_expression_params(&ctx)?;
}
let tracker = expressions::TrackedExpressionAttributes::new(
&request.expression_attribute_names,
&request.expression_attribute_values,
);
let mut pairs: Vec<EqualityPair> = Vec::new();
let mut referenced_attrs: Vec<String> = Vec::new();
let condition = match request.search_condition_expression.as_deref() {
Some("") => {
return Err(DynoxideError::ValidationException(
"Invalid SearchConditionExpression: The expression can not be empty;".to_string(),
));
}
Some(expr_str) => {
let parsed = expressions::condition::parse(expr_str).map_err(|e| {
DynoxideError::ValidationException(format!(
"Invalid SearchConditionExpression: {e}"
))
})?;
expressions::condition::validate_name_refs(
&parsed,
&request.expression_attribute_names,
)
.map_err(|e| {
DynoxideError::ValidationException(format!(
"Invalid SearchConditionExpression: {e}"
))
})?;
collect_equalities(&parsed, &tracker, &mut pairs, &mut referenced_attrs)?;
Some(parsed)
}
None => None,
};
let projection = match request.projection_expression.as_deref() {
Some("") => {
return Err(DynoxideError::ValidationException(
"Invalid ProjectionExpression: The expression can not be empty;".to_string(),
));
}
Some(proj_expr) => Some(
expressions::projection::parse(proj_expr)
.map_err(DynoxideError::ValidationException)?,
),
None => None,
};
let meta = helpers::require_table_for_item_op(storage, &request.table_name).await?;
let vixs = parse_vector_defs(&meta)?;
let Some(vix) = vixs.iter().find(|v| v.index_name == request.index_name) else {
return Err(DynoxideError::ValidationException(format!(
"The table does not have the specified index: {}",
request.index_name
)));
};
if !lifecycle
.phase_of_index(storage, &request.table_name, &request.index_name)
.is_searchable()
{
return Err(DynoxideError::ValidationException(format!(
"Cannot search backfilling vector index: {}",
request.index_name
)));
}
let query_vec = parse_query_vector(&request.search_vector, vix)?;
for attr in &referenced_attrs {
let in_schema = vix
.search_schema
.as_ref()
.is_some_and(|schema| schema.iter().any(|e| e.attribute_name == *attr));
if !in_schema {
return Err(DynoxideError::ValidationException(format!(
"SearchConditionExpression must not contain any attributes that is not in \
SearchSchema. Invalid attribute: {attr}"
)));
}
}
if pairs.iter().any(|p| p.nested) {
return Err(DynoxideError::ValidationException(
"Invalid SearchConditionExpression: SearchConditionExpression cannot have \
conditions on nested attributes"
.to_string(),
));
}
{
let mut seen: Vec<&str> = Vec::with_capacity(pairs.len());
for pair in &pairs {
if seen.contains(&pair.attr.as_str()) {
return Err(DynoxideError::ValidationException(
"Invalid SearchConditionExpression: SearchConditionExpression must only \
contain one condition per attribute"
.to_string(),
));
}
seen.push(&pair.attr);
}
}
let index_hash_attr = hash_attr(vix);
if let Some(hash) = index_hash_attr {
let provided = pairs.iter().any(|p| p.attr == hash && !p.nested);
if !provided {
return Err(DynoxideError::ValidationException(
"SearchConditionExpression must be provided when SearchSchema has a HASH key"
.to_string(),
));
}
}
if let Some(ref cond) = condition {
tracker.track_condition_expr(cond);
}
if let Some(ref proj) = projection {
tracker.track_projection_expr(proj);
expressions::projection::validate(proj, &tracker)
.map_err(DynoxideError::ValidationException)?;
}
let mut hash_scope: Option<String> = None;
let mut unscopable_hash = false;
if let Some(hash) = index_hash_attr {
let attr_defs = parse_attr_defs(&meta)?;
for pair in pairs.iter().filter(|p| p.attr == hash) {
if let Some(def) = attr_defs.iter().find(|d| d.attribute_name == hash) {
let declared = scalar_type_str(&def.attribute_type);
let actual = pair.value.type_name();
if actual != declared {
return Err(DynoxideError::ValidationException(format!(
"Type of '{hash}' attribute in SearchConditionExpression ({actual}) \
does not match type in search schema ({declared})"
)));
}
}
if matches!(&pair.value, AttributeValue::S(s) if s.is_empty()) {
return Err(DynoxideError::ValidationException(format!(
"One or more parameter values are not valid. The AttributeValue for a key \
attribute cannot contain an empty string value. Key: {hash}"
)));
}
match pair.value.to_key_string() {
Some(key) => hash_scope = Some(key),
None => unscopable_hash = true,
}
}
}
let candidates: Vec<VectorCandidateRow> = if unscopable_hash {
Vec::new()
} else {
storage
.query_vector_candidates(
&request.table_name,
&request.index_name,
hash_scope.as_deref(),
)
.await?
};
let filter_pairs: Vec<&EqualityPair> = pairs
.iter()
.filter(|p| Some(p.attr.as_str()) != index_hash_attr)
.collect();
let mut scored: Vec<(f64, &VectorCandidateRow)> = Vec::with_capacity(candidates.len());
for row in &candidates {
let filter_map: HashMap<String, AttributeValue> = if filter_pairs.is_empty() {
HashMap::new()
} else {
serde_json::from_str(&row.filter_json).map_err(|e| {
DynoxideError::InternalServerError(format!("Bad filter JSON in storage: {e}"))
})?
};
let passes = filter_pairs.iter().all(|pair| {
filter_map
.get(&pair.attr)
.is_some_and(|stored| expressions::condition::values_equal(stored, &pair.value))
});
if !passes {
continue;
}
let stored_vec: Vec<f32> = serde_json::from_str(&row.vector_json).map_err(|e| {
DynoxideError::InternalServerError(format!("Bad vector JSON in storage: {e}"))
})?;
scored.push((
score_candidate(&vix.distance_function, &query_vec, &stored_vec),
row,
));
}
if higher_is_better(&vix.distance_function) {
scored.sort_by(|a, b| b.0.total_cmp(&a.0));
} else {
scored.sort_by(|a, b| a.0.total_cmp(&b.0));
}
scored.truncate(request.top_k as usize);
let loop_tracker = expressions::TrackedExpressionAttributes::without_tracking(
&request.expression_attribute_names,
&request.expression_attribute_values,
);
let no_keys: &[String] = &[];
let survivor_keys: Vec<(String, String)> = scored
.iter()
.map(|(_, row)| (row.table_pk.clone(), row.table_sk.clone()))
.collect();
let items: HashMap<(String, String), String> = storage
.vector_items_for_keys(&request.table_name, &request.index_name, &survivor_keys)
.await?
.into_iter()
.map(|(pk, sk, item_json)| ((pk, sk), item_json))
.collect();
let mut search_results = Vec::with_capacity(scored.len());
for (score, row) in scored {
let key = (row.table_pk.clone(), row.table_sk.clone());
let Some(item_json) = items.get(&key) else {
continue;
};
let item: Item = serde_json::from_str(item_json).map_err(|e| {
DynoxideError::InternalServerError(format!("Bad item JSON in storage: {e}"))
})?;
let result_item = if let Some(ref proj) = projection {
expressions::projection::apply(&item, proj, &loop_tracker, no_keys)
.map_err(DynoxideError::ValidationException)?
} else {
let mut item = item;
item.remove(&vix.vector_attribute.attribute_name);
item
};
search_results.push(SearchResult {
item: result_item,
score,
});
}
tracker.check_unused()?;
let consumed_capacity =
crate::types::capacity_wanted(request.return_consumed_capacity.as_deref()).then(|| {
crate::types::VectorSearchCapacity {
vector_search_request_bytes: search_scanned_bytes(&candidates),
}
});
Ok(SearchVectorsResponse {
search_results,
consumed_capacity,
})
}