diff --git a/src/sql/scalar.rs b/src/sql/scalar.rs index 8a19324..3c3644d 100644 --- a/src/sql/scalar.rs +++ b/src/sql/scalar.rs @@ -38,12 +38,12 @@ pub(super) fn eval_scalar_function(name: &str, args: &[Value]) -> Value { if a.len() != b.len() { return Value::Null; } - let dot: f64 = a.iter().zip(&b).map(|(x, y)| x * y).sum(); + let dot: f64 = a.iter().zip(b.iter()).map(|(x, y)| x * y).sum(); match name.to_lowercase().as_str() { "dot_product" => Value::Float(dot), "l2_distance" => Value::Float( a.iter() - .zip(&b) + .zip(b.iter()) .map(|(x, y)| (x - y).powi(2)) .sum::() .sqrt(), @@ -568,15 +568,39 @@ fn eval_mq_scalar(program: &str, content: &str) -> Value { } } +type CachedVector = (String, std::rc::Rc<[f64]>); + +thread_local! { + /// The two most recently parsed vectors, most recent first. A similarity + /// query passes the same query vector for every row, so keeping it + /// alongside the current row's vector avoids re-parsing it each time. + static RECENT_VECTORS: std::cell::RefCell> = + const { std::cell::RefCell::new(Vec::new()) }; +} + /// Parses a non-empty JSON array of numbers; `None` for anything else. -fn parse_vector(value: &Value) -> Option> { - let serde_json::Value::Array(items) = serde_json::from_str(value.as_str()?).ok()? else { - return None; - }; - if items.is_empty() { +fn parse_vector(value: &Value) -> Option> { + let text = value.as_str()?; + let hit = RECENT_VECTORS.with(|c| { + let mut recent = c.borrow_mut(); + let pos = recent.iter().position(|(t, _)| t == text)?; + recent[..=pos].rotate_right(1); + Some(recent[0].1.clone()) + }); + if hit.is_some() { + return hit; + } + let parsed: Vec = serde_json::from_str(text).ok()?; + if parsed.is_empty() || !parsed.iter().all(|f| f.is_finite()) { return None; } - items.iter().map(serde_json::Value::as_f64).collect() + let vector: std::rc::Rc<[f64]> = parsed.into(); + RECENT_VECTORS.with(|c| { + let mut recent = c.borrow_mut(); + recent.insert(0, (text.to_string(), vector.clone())); + recent.truncate(2); + }); + Some(vector) } /// Looks up `path` (`key`, `$.a.b`, `a[0].b`) in `json`. diff --git a/src/sql/tests.rs b/src/sql/tests.rs index 933293e..c5c7589 100644 --- a/src/sql/tests.rs +++ b/src/sql/tests.rs @@ -3217,6 +3217,9 @@ fn test_sql_correlated_subquery_expression_forms(#[case] cond: &str, #[case] exp #[case("cosine_similarity('[]', '[]')", "NULL")] #[case("cosine_similarity('[1,\"a\"]', '[1,0]')", "NULL")] #[case("cosine_similarity('not json', '[1,0]')", "NULL")] +#[case("dot_product('[+1,02]', '[3,4]')", "NULL")] +#[case("dot_product('[.5,1.]', '[3,4]')", "NULL")] +#[case("dot_product('[1,2,]', '[3,4]')", "NULL")] #[case("cosine_similarity(NULL, '[1,0]')", "NULL")] fn test_sql_vector_functions(#[case] expr: &str, #[case] expected: &str) { let store = DocumentStore::new();