linkrs_query/executor/expression/functions/builtin/
vector.rs

1//! Vector Functions Module
2//!
3//! This module provides vector operations functions for SQL queries, including:
4//! - Similarity functions (cosine_similarity, dot_product, euclidean_distance, manhattan_distance)
5//! - Vector property functions (dimension, l2_norm, nnz, normalize)
6//! - Vector access functions (element access, slicing)
7
8use crate::executor::expression::functions::signature::{FunctionSignature, ValueType};
9use crate::executor::expression::{ExpressionError, ExpressionErrorType};
10use linkrs_core::Value;
11
12/// Vector function enumeration
13#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
14pub enum VectorFunction {
15    /// cosine_similarity(vec1, vec2) - Compute cosine similarity between two vectors
16    CosineSimilarity,
17    /// dot_product(vec1, vec2) - Compute dot product of two vectors
18    DotProduct,
19    /// euclidean_distance(vec1, vec2) - Compute Euclidean distance between two vectors
20    EuclideanDistance,
21    /// manhattan_distance(vec1, vec2) - Compute Manhattan distance between two vectors
22    ManhattanDistance,
23    /// dimension(vector) - Get the dimension of a vector
24    Dimension,
25    /// l2_norm(vector) - Compute L2 norm of a vector
26    L2Norm,
27    /// nnz(vector) - Get number of non-zero elements
28    Nnz,
29    /// normalize(vector) - Normalize a vector to unit length
30    Normalize,
31    /// array_value(...) - Create an array from values
32    ArrayValue,
33    /// array_cosine_similarity(arr1, arr2) - Compute cosine similarity between two arrays
34    ArrayCosineSimilarity,
35    /// array_distance(arr1, arr2) - Compute Euclidean distance between two arrays
36    ArrayDistance,
37    /// array_squared_distance(arr1, arr2) - Compute squared Euclidean distance between two arrays
38    ArraySquaredDistance,
39    /// array_inner_product(arr1, arr2) - Compute inner product of two arrays
40    ArrayInnerProduct,
41    /// array_dot_product(arr1, arr2) - Compute dot product of two arrays
42    ArrayDotProduct,
43}
44
45impl VectorFunction {
46    /// Get function name
47    pub fn name(&self) -> &'static str {
48        match self {
49            VectorFunction::CosineSimilarity => "cosine_similarity",
50            VectorFunction::DotProduct => "dot_product",
51            VectorFunction::EuclideanDistance => "euclidean_distance",
52            VectorFunction::ManhattanDistance => "manhattan_distance",
53            VectorFunction::Dimension => "dimension",
54            VectorFunction::L2Norm => "l2_norm",
55            VectorFunction::Nnz => "nnz",
56            VectorFunction::Normalize => "normalize",
57            VectorFunction::ArrayValue => "array_value",
58            VectorFunction::ArrayCosineSimilarity => "array_cosine_similarity",
59            VectorFunction::ArrayDistance => "array_distance",
60            VectorFunction::ArraySquaredDistance => "array_squared_distance",
61            VectorFunction::ArrayInnerProduct => "array_inner_product",
62            VectorFunction::ArrayDotProduct => "array_dot_product",
63        }
64    }
65
66    /// Get function signature
67    pub fn signature(&self) -> FunctionSignature {
68        match self {
69            VectorFunction::CosineSimilarity => FunctionSignature::new(
70                "cosine_similarity",
71                vec![ValueType::Any, ValueType::Any],
72                Some(ValueType::Float),
73                false,
74            ),
75            VectorFunction::DotProduct => FunctionSignature::new(
76                "dot_product",
77                vec![ValueType::Any, ValueType::Any],
78                Some(ValueType::Float),
79                false,
80            ),
81            VectorFunction::EuclideanDistance => FunctionSignature::new(
82                "euclidean_distance",
83                vec![ValueType::Any, ValueType::Any],
84                Some(ValueType::Float),
85                false,
86            ),
87            VectorFunction::ManhattanDistance => FunctionSignature::new(
88                "manhattan_distance",
89                vec![ValueType::Any, ValueType::Any],
90                Some(ValueType::Float),
91                false,
92            ),
93            VectorFunction::Dimension => FunctionSignature::new(
94                "dimension",
95                vec![ValueType::Any],
96                Some(ValueType::Int),
97                false,
98            ),
99            VectorFunction::L2Norm => FunctionSignature::new(
100                "l2_norm",
101                vec![ValueType::Any],
102                Some(ValueType::Float),
103                false,
104            ),
105            VectorFunction::Nnz => {
106                FunctionSignature::new("nnz", vec![ValueType::Any], Some(ValueType::Int), false)
107            }
108            VectorFunction::Normalize => FunctionSignature::new(
109                "normalize",
110                vec![ValueType::Any],
111                Some(ValueType::Any),
112                false,
113            ),
114            VectorFunction::ArrayValue => {
115                FunctionSignature::new("array_value", vec![], Some(ValueType::Any), true)
116            }
117            VectorFunction::ArrayCosineSimilarity => FunctionSignature::new(
118                "array_cosine_similarity",
119                vec![ValueType::Any, ValueType::Any],
120                Some(ValueType::Float),
121                false,
122            ),
123            VectorFunction::ArrayDistance => FunctionSignature::new(
124                "array_distance",
125                vec![ValueType::Any, ValueType::Any],
126                Some(ValueType::Float),
127                false,
128            ),
129            VectorFunction::ArraySquaredDistance => FunctionSignature::new(
130                "array_squared_distance",
131                vec![ValueType::Any, ValueType::Any],
132                Some(ValueType::Float),
133                false,
134            ),
135            VectorFunction::ArrayInnerProduct => FunctionSignature::new(
136                "array_inner_product",
137                vec![ValueType::Any, ValueType::Any],
138                Some(ValueType::Float),
139                false,
140            ),
141            VectorFunction::ArrayDotProduct => FunctionSignature::new(
142                "array_dot_product",
143                vec![ValueType::Any, ValueType::Any],
144                Some(ValueType::Float),
145                false,
146            ),
147        }
148    }
149
150    /// Get the number of parameters
151    pub fn arity(&self) -> usize {
152        match self {
153            VectorFunction::CosineSimilarity => 2,
154            VectorFunction::DotProduct => 2,
155            VectorFunction::EuclideanDistance => 2,
156            VectorFunction::ManhattanDistance => 2,
157            VectorFunction::Dimension => 1,
158            VectorFunction::L2Norm => 1,
159            VectorFunction::Nnz => 1,
160            VectorFunction::Normalize => 1,
161            VectorFunction::ArrayValue => 0,
162            VectorFunction::ArrayCosineSimilarity => 2,
163            VectorFunction::ArrayDistance => 2,
164            VectorFunction::ArraySquaredDistance => 2,
165            VectorFunction::ArrayInnerProduct => 2,
166            VectorFunction::ArrayDotProduct => 2,
167        }
168    }
169
170    /// Check if variable parameters are accepted
171    pub fn is_variadic(&self) -> bool {
172        matches!(self, VectorFunction::ArrayValue)
173    }
174
175    /// Get function description
176    pub fn description(&self) -> &'static str {
177        match self {
178            VectorFunction::CosineSimilarity => "Compute cosine similarity between two vectors",
179            VectorFunction::DotProduct => "Compute dot product of two vectors",
180            VectorFunction::EuclideanDistance => "Compute Euclidean distance between two vectors",
181            VectorFunction::ManhattanDistance => "Compute Manhattan distance between two vectors",
182            VectorFunction::Dimension => "Get the dimension of a vector",
183            VectorFunction::L2Norm => "Compute L2 norm of a vector",
184            VectorFunction::Nnz => "Get number of non-zero elements in a vector",
185            VectorFunction::Normalize => "Normalize a vector to unit length",
186            VectorFunction::ArrayValue => "Create an array from values",
187            VectorFunction::ArrayCosineSimilarity => "Compute cosine similarity between two arrays",
188            VectorFunction::ArrayDistance => "Compute Euclidean distance between two arrays",
189            VectorFunction::ArraySquaredDistance => {
190                "Compute squared Euclidean distance between two arrays"
191            }
192            VectorFunction::ArrayInnerProduct => "Compute inner product of two arrays",
193            VectorFunction::ArrayDotProduct => "Compute dot product of two arrays",
194        }
195    }
196
197    /// Execute the function
198    pub fn execute(&self, args: &[Value]) -> Result<Value, ExpressionError> {
199        if !self.is_variadic() && args.len() != self.arity() {
200            return Err(ExpressionError::invalid_arity(
201                self.name(),
202                self.arity(),
203                args.len(),
204            ));
205        }
206        match self {
207            VectorFunction::CosineSimilarity => self.execute_cosine_similarity(args),
208            VectorFunction::DotProduct => self.execute_dot_product(args),
209            VectorFunction::EuclideanDistance => self.execute_euclidean_distance(args),
210            VectorFunction::ManhattanDistance => self.execute_manhattan_distance(args),
211            VectorFunction::Dimension => self.execute_dimension(args),
212            VectorFunction::L2Norm => self.execute_l2_norm(args),
213            VectorFunction::Nnz => self.execute_nnz(args),
214            VectorFunction::Normalize => self.execute_normalize(args),
215            VectorFunction::ArrayValue => self.execute_array_value(args),
216            VectorFunction::ArrayCosineSimilarity => self.execute_cosine_similarity(args),
217            VectorFunction::ArrayDistance => self.execute_euclidean_distance(args),
218            VectorFunction::ArraySquaredDistance => self.execute_array_squared_distance(args),
219            VectorFunction::ArrayInnerProduct => self.execute_dot_product(args),
220            VectorFunction::ArrayDotProduct => self.execute_dot_product(args),
221        }
222    }
223
224    /// Execute cosine_similarity function
225    fn execute_cosine_similarity(&self, args: &[Value]) -> Result<Value, ExpressionError> {
226        let vec1 = extract_vector(&args[0])?;
227        let vec2 = extract_vector(&args[1])?;
228
229        let similarity = compute_cosine_similarity(&vec1, &vec2)?;
230        Ok(Value::Double(similarity as f64))
231    }
232
233    /// Execute dot_product function
234    fn execute_dot_product(&self, args: &[Value]) -> Result<Value, ExpressionError> {
235        let vec1 = extract_vector(&args[0])?;
236        let vec2 = extract_vector(&args[1])?;
237
238        let dot = compute_dot_product(&vec1, &vec2)?;
239        Ok(Value::Double(dot as f64))
240    }
241
242    /// Execute euclidean_distance function
243    fn execute_euclidean_distance(&self, args: &[Value]) -> Result<Value, ExpressionError> {
244        let vec1 = extract_vector(&args[0])?;
245        let vec2 = extract_vector(&args[1])?;
246
247        let distance = compute_euclidean_distance(&vec1, &vec2)?;
248        Ok(Value::Double(distance as f64))
249    }
250
251    /// Execute manhattan_distance function
252    fn execute_manhattan_distance(&self, args: &[Value]) -> Result<Value, ExpressionError> {
253        let vec1 = extract_vector(&args[0])?;
254        let vec2 = extract_vector(&args[1])?;
255
256        let distance = compute_manhattan_distance(&vec1, &vec2)?;
257        Ok(Value::Double(distance as f64))
258    }
259
260    /// Execute dimension function
261    fn execute_dimension(&self, args: &[Value]) -> Result<Value, ExpressionError> {
262        let vec = extract_vector(&args[0])?;
263        Ok(Value::BigInt(vec.len() as i64))
264    }
265
266    /// Execute l2_norm function
267    fn execute_l2_norm(&self, args: &[Value]) -> Result<Value, ExpressionError> {
268        let vec = extract_vector(&args[0])?;
269        let norm = compute_l2_norm(&vec);
270        Ok(Value::Double(norm as f64))
271    }
272
273    /// Execute nnz function
274    fn execute_nnz(&self, args: &[Value]) -> Result<Value, ExpressionError> {
275        let vec = extract_vector(&args[0])?;
276        let nnz = vec.iter().filter(|&&x| x != 0.0).count();
277        Ok(Value::BigInt(nnz as i64))
278    }
279
280    /// Execute normalize function
281    fn execute_normalize(&self, args: &[Value]) -> Result<Value, ExpressionError> {
282        let vec = extract_vector(&args[0])?;
283        let norm = compute_l2_norm(&vec);
284
285        if norm == 0.0 {
286            return Ok(Value::vector(vec)); // Return original vector if norm is zero
287        }
288
289        let normalized: Vec<f32> = vec.iter().map(|&x| x / norm).collect();
290        Ok(Value::vector(normalized))
291    }
292
293    /// Execute array_value function
294    fn execute_array_value(&self, args: &[Value]) -> Result<Value, ExpressionError> {
295        let values: Vec<Value> = args.to_vec();
296        Ok(Value::list(linkrs_core::value::list::List { values }))
297    }
298
299    /// Execute array_squared_distance function
300    fn execute_array_squared_distance(&self, args: &[Value]) -> Result<Value, ExpressionError> {
301        let vec1 = extract_vector(&args[0])?;
302        let vec2 = extract_vector(&args[1])?;
303
304        if vec1.len() != vec2.len() {
305            return Err(ExpressionError::new(
306                ExpressionErrorType::TypeError,
307                format!(
308                    "Array dimensions must match: {} vs {}",
309                    vec1.len(),
310                    vec2.len()
311                ),
312            ));
313        }
314
315        let sum: f32 = vec1
316            .iter()
317            .zip(vec2.iter())
318            .map(|(&a, &b)| (a - b).powi(2))
319            .sum();
320        Ok(Value::Double(sum as f64))
321    }
322}
323
324/// Extract vector from Value
325fn extract_vector(value: &Value) -> Result<Vec<f32>, ExpressionError> {
326    match value {
327        Value::Vector(vec) => Ok(vec.to_dense()),
328        Value::List(list) => {
329            // Try to convert List<Float> or List<Int> to vector
330            let mut vec = Vec::with_capacity(list.values.len());
331            for item in &list.values {
332                match item {
333                    Value::Float(f) => vec.push(*f),
334                    Value::Double(f) => vec.push(*f as f32),
335                    Value::SmallInt(i) => vec.push(*i as f32),
336                    Value::Int(i) => vec.push(*i as f32),
337                    Value::BigInt(i) => vec.push(*i as f32),
338                    _ => {
339                        return Err(ExpressionError::new(
340                            ExpressionErrorType::TypeError,
341                            format!("Vector elements must be numeric, got {:?}", item.get_type()),
342                        ))
343                    }
344                }
345            }
346            Ok(vec)
347        }
348        _ => Err(ExpressionError::new(
349            ExpressionErrorType::TypeError,
350            format!(
351                "Expected vector or list of numbers, got {:?}",
352                value.get_type()
353            ),
354        )),
355    }
356}
357
358/// Compute cosine similarity between two vectors
359fn compute_cosine_similarity(vec1: &[f32], vec2: &[f32]) -> Result<f32, ExpressionError> {
360    if vec1.len() != vec2.len() {
361        return Err(ExpressionError::new(
362            ExpressionErrorType::TypeError,
363            format!(
364                "Vector dimensions must match: {} vs {}",
365                vec1.len(),
366                vec2.len()
367            ),
368        ));
369    }
370
371    let dot_product: f32 = vec1.iter().zip(vec2.iter()).map(|(&a, &b)| a * b).sum();
372    let norm1: f32 = vec1.iter().map(|&x| x * x).sum::<f32>().sqrt();
373    let norm2: f32 = vec2.iter().map(|&x| x * x).sum::<f32>().sqrt();
374
375    if norm1 == 0.0 || norm2 == 0.0 {
376        return Ok(0.0);
377    }
378
379    Ok(dot_product / (norm1 * norm2))
380}
381
382/// Compute dot product of two vectors
383fn compute_dot_product(vec1: &[f32], vec2: &[f32]) -> Result<f32, ExpressionError> {
384    if vec1.len() != vec2.len() {
385        return Err(ExpressionError::new(
386            ExpressionErrorType::TypeError,
387            format!(
388                "Vector dimensions must match: {} vs {}",
389                vec1.len(),
390                vec2.len()
391            ),
392        ));
393    }
394
395    Ok(vec1.iter().zip(vec2.iter()).map(|(&a, &b)| a * b).sum())
396}
397
398/// Compute Euclidean distance between two vectors
399fn compute_euclidean_distance(vec1: &[f32], vec2: &[f32]) -> Result<f32, ExpressionError> {
400    if vec1.len() != vec2.len() {
401        return Err(ExpressionError::new(
402            ExpressionErrorType::TypeError,
403            format!(
404                "Vector dimensions must match: {} vs {}",
405                vec1.len(),
406                vec2.len()
407            ),
408        ));
409    }
410
411    let sum: f32 = vec1
412        .iter()
413        .zip(vec2.iter())
414        .map(|(&a, &b)| (a - b).powi(2))
415        .sum();
416    Ok(sum.sqrt())
417}
418
419/// Compute Manhattan distance between two vectors
420fn compute_manhattan_distance(vec1: &[f32], vec2: &[f32]) -> Result<f32, ExpressionError> {
421    if vec1.len() != vec2.len() {
422        return Err(ExpressionError::new(
423            ExpressionErrorType::TypeError,
424            format!(
425                "Vector dimensions must match: {} vs {}",
426                vec1.len(),
427                vec2.len()
428            ),
429        ));
430    }
431
432    Ok(vec1
433        .iter()
434        .zip(vec2.iter())
435        .map(|(&a, &b)| (a - b).abs())
436        .sum())
437}
438
439/// Compute L2 norm of a vector
440fn compute_l2_norm(vec: &[f32]) -> f32 {
441    vec.iter().map(|&x| x * x).sum::<f32>().sqrt()
442}
443
444/// Register vector functions
445pub fn register_vector_functions(
446    registry: &mut crate::executor::expression::functions::FunctionRegistry,
447) {
448    registry.register_builtin(
449        crate::executor::expression::functions::BuiltinFunction::Vector(
450            VectorFunction::CosineSimilarity,
451        ),
452    );
453
454    registry.register_builtin(
455        crate::executor::expression::functions::BuiltinFunction::Vector(VectorFunction::DotProduct),
456    );
457
458    registry.register_builtin(
459        crate::executor::expression::functions::BuiltinFunction::Vector(
460            VectorFunction::EuclideanDistance,
461        ),
462    );
463
464    registry.register_builtin(
465        crate::executor::expression::functions::BuiltinFunction::Vector(
466            VectorFunction::ManhattanDistance,
467        ),
468    );
469
470    registry.register_builtin(
471        crate::executor::expression::functions::BuiltinFunction::Vector(VectorFunction::Dimension),
472    );
473
474    registry.register_builtin(
475        crate::executor::expression::functions::BuiltinFunction::Vector(VectorFunction::L2Norm),
476    );
477
478    registry.register_builtin(
479        crate::executor::expression::functions::BuiltinFunction::Vector(VectorFunction::Nnz),
480    );
481
482    registry.register_builtin(
483        crate::executor::expression::functions::BuiltinFunction::Vector(VectorFunction::Normalize),
484    );
485
486    registry.register_builtin(
487        crate::executor::expression::functions::BuiltinFunction::Vector(VectorFunction::ArrayValue),
488    );
489
490    registry.register_builtin(
491        crate::executor::expression::functions::BuiltinFunction::Vector(
492            VectorFunction::ArrayCosineSimilarity,
493        ),
494    );
495
496    registry.register_builtin(
497        crate::executor::expression::functions::BuiltinFunction::Vector(
498            VectorFunction::ArrayDistance,
499        ),
500    );
501
502    registry.register_builtin(
503        crate::executor::expression::functions::BuiltinFunction::Vector(
504            VectorFunction::ArraySquaredDistance,
505        ),
506    );
507
508    registry.register_builtin(
509        crate::executor::expression::functions::BuiltinFunction::Vector(
510            VectorFunction::ArrayInnerProduct,
511        ),
512    );
513
514    registry.register_builtin(
515        crate::executor::expression::functions::BuiltinFunction::Vector(
516            VectorFunction::ArrayDotProduct,
517        ),
518    );
519}
520
521#[cfg(test)]
522mod tests {
523    use super::*;
524    use linkrs_core::value::list::List;
525
526    #[test]
527    fn test_cosine_similarity() {
528        let vec1 = Value::vector(vec![1.0, 0.0, 0.0]);
529        let vec2 = Value::vector(vec![0.0, 1.0, 0.0]);
530
531        let result = VectorFunction::CosineSimilarity
532            .execute(&[vec1, vec2])
533            .unwrap();
534
535        assert!(matches!(result, Value::Double(_)));
536        if let Value::Double(score) = result {
537            assert!((score - 0.0).abs() < 1e-6);
538        }
539    }
540
541    #[test]
542    fn test_dot_product() {
543        let vec1 = Value::vector(vec![1.0, 2.0, 3.0]);
544        let vec2 = Value::vector(vec![4.0, 5.0, 6.0]);
545
546        let result = VectorFunction::DotProduct.execute(&[vec1, vec2]).unwrap();
547
548        assert!(matches!(result, Value::Double(_)));
549        if let Value::Double(dot) = result {
550            assert!((dot - 32.0).abs() < 1e-6);
551        }
552    }
553
554    #[test]
555    fn test_euclidean_distance() {
556        let vec1 = Value::vector(vec![0.0, 0.0]);
557        let vec2 = Value::vector(vec![3.0, 4.0]);
558
559        let result = VectorFunction::EuclideanDistance
560            .execute(&[vec1, vec2])
561            .unwrap();
562
563        assert!(matches!(result, Value::Double(_)));
564        if let Value::Double(dist) = result {
565            assert!((dist - 5.0).abs() < 1e-6);
566        }
567    }
568
569    #[test]
570    fn test_manhattan_distance() {
571        let vec1 = Value::vector(vec![0.0, 0.0]);
572        let vec2 = Value::vector(vec![3.0, 4.0]);
573
574        let result = VectorFunction::ManhattanDistance
575            .execute(&[vec1, vec2])
576            .unwrap();
577
578        assert!(matches!(result, Value::Double(_)));
579        if let Value::Double(dist) = result {
580            assert!((dist - 7.0).abs() < 1e-6);
581        }
582    }
583
584    #[test]
585    fn test_dimension() {
586        let vec = Value::vector(vec![0.1; 1536]);
587
588        let result = VectorFunction::Dimension.execute(&[vec]).unwrap();
589
590        assert!(matches!(result, Value::BigInt(_)));
591        if let Value::BigInt(dim) = result {
592            assert_eq!(dim, 1536);
593        }
594    }
595
596    #[test]
597    fn test_l2_norm() {
598        let vec = Value::vector(vec![3.0, 4.0]);
599
600        let result = VectorFunction::L2Norm.execute(&[vec]).unwrap();
601
602        assert!(matches!(result, Value::Double(_)));
603        if let Value::Double(norm) = result {
604            assert!((norm - 5.0).abs() < 1e-6);
605        }
606    }
607
608    #[test]
609    fn test_nnz() {
610        let vec = Value::vector(vec![1.0, 0.0, 2.0, 0.0, 3.0]);
611
612        let result = VectorFunction::Nnz.execute(&[vec]).unwrap();
613
614        assert!(matches!(result, Value::BigInt(_)));
615        if let Value::BigInt(n) = result {
616            assert_eq!(n, 3);
617        }
618    }
619
620    #[test]
621    fn test_normalize() {
622        let vec = Value::vector(vec![3.0, 4.0]);
623
624        let result = VectorFunction::Normalize.execute(&[vec]).unwrap();
625
626        assert!(matches!(result, Value::Vector(_)));
627        if let Value::Vector(normalized) = result {
628            let data = normalized.to_dense();
629            assert!((data[0] - 0.6).abs() < 1e-6); // 3/5 = 0.6
630            assert!((data[1] - 0.8).abs() < 1e-6); // 4/5 = 0.8
631        }
632    }
633
634    #[test]
635    fn test_dimension_mismatch() {
636        let vec1 = Value::vector(vec![1.0, 2.0]);
637        let vec2 = Value::vector(vec![1.0, 2.0, 3.0]);
638
639        let result = VectorFunction::CosineSimilarity.execute(&[vec1, vec2]);
640        assert!(result.is_err());
641    }
642
643    #[test]
644    fn test_array_value() {
645        let result = VectorFunction::ArrayValue
646            .execute(&[Value::Int(1), Value::Int(2), Value::Int(3)])
647            .unwrap();
648        assert!(matches!(result, Value::List(_)));
649    }
650
651    #[test]
652    fn test_array_cosine_similarity() {
653        let arr1 = Value::list(List {
654            values: vec![Value::Float(1.0), Value::Float(0.0), Value::Float(0.0)],
655        });
656        let arr2 = Value::list(List {
657            values: vec![Value::Float(0.0), Value::Float(1.0), Value::Float(0.0)],
658        });
659
660        let result = VectorFunction::ArrayCosineSimilarity
661            .execute(&[arr1, arr2])
662            .unwrap();
663
664        assert!(matches!(result, Value::Double(_)));
665        if let Value::Double(score) = result {
666            assert!((score - 0.0).abs() < 1e-6);
667        }
668    }
669
670    #[test]
671    fn test_array_distance() {
672        let arr1 = Value::list(List {
673            values: vec![Value::Float(0.0), Value::Float(0.0)],
674        });
675        let arr2 = Value::list(List {
676            values: vec![Value::Float(3.0), Value::Float(4.0)],
677        });
678
679        let result = VectorFunction::ArrayDistance
680            .execute(&[arr1, arr2])
681            .unwrap();
682
683        assert!(matches!(result, Value::Double(_)));
684        if let Value::Double(dist) = result {
685            assert!((dist - 5.0).abs() < 1e-6);
686        }
687    }
688
689    #[test]
690    fn test_array_squared_distance() {
691        let arr1 = Value::list(List {
692            values: vec![Value::Float(0.0), Value::Float(0.0)],
693        });
694        let arr2 = Value::list(List {
695            values: vec![Value::Float(3.0), Value::Float(4.0)],
696        });
697
698        let result = VectorFunction::ArraySquaredDistance
699            .execute(&[arr1, arr2])
700            .unwrap();
701
702        assert!(matches!(result, Value::Double(_)));
703        if let Value::Double(dist) = result {
704            assert!((dist - 25.0).abs() < 1e-6);
705        }
706    }
707
708    #[test]
709    fn test_array_inner_product() {
710        let arr1 = Value::list(List {
711            values: vec![Value::Float(1.0), Value::Float(2.0), Value::Float(3.0)],
712        });
713        let arr2 = Value::list(List {
714            values: vec![Value::Float(4.0), Value::Float(5.0), Value::Float(6.0)],
715        });
716
717        let result = VectorFunction::ArrayInnerProduct
718            .execute(&[arr1, arr2])
719            .unwrap();
720
721        assert!(matches!(result, Value::Double(_)));
722        if let Value::Double(dot) = result {
723            assert!((dot - 32.0).abs() < 1e-6);
724        }
725    }
726
727    #[test]
728    fn test_array_dot_product() {
729        let arr1 = Value::list(List {
730            values: vec![Value::Float(1.0), Value::Float(2.0), Value::Float(3.0)],
731        });
732        let arr2 = Value::list(List {
733            values: vec![Value::Float(4.0), Value::Float(5.0), Value::Float(6.0)],
734        });
735
736        let result = VectorFunction::ArrayDotProduct
737            .execute(&[arr1, arr2])
738            .unwrap();
739
740        assert!(matches!(result, Value::Double(_)));
741        if let Value::Double(dot) = result {
742            assert!((dot - 32.0).abs() < 1e-6);
743        }
744    }
745}
linkrs_query/executor/expression/functions/builtin/vector.rs

linkrs_query/executor/expression/functions/builtin/
vector.rs
