Skip to main content

clt_database/vector/operations/
slice.rs

1use crate::{
2    vector::vector_types::{Vector, VectorType},
3    LimboError, Result,
4};
5
6pub fn vector_slice(vector: &Vector, start: usize, end: usize) -> Result<Vector<'static>> {
7    if start > end {
8        return Err(LimboError::InvalidArgument(
9            "start index must not be greater than end index".into(),
10        ));
11    }
12    if end > vector.dims || end < start {
13        return Err(LimboError::ConversionError(
14            "vector_slice range out of bounds".into(),
15        ));
16    }
17    match vector.vector_type {
18        VectorType::Float32Dense => Ok(Vector {
19            vector_type: vector.vector_type,
20            dims: end - start,
21            owned: Some(vector.bin_data()[start * 4..end * 4].to_vec()),
22            refer: None,
23        }),
24        VectorType::Float64Dense => Ok(Vector {
25            vector_type: vector.vector_type,
26            dims: end - start,
27            owned: Some(vector.bin_data()[start * 8..end * 8].to_vec()),
28            refer: None,
29        }),
30        VectorType::Float32Sparse => {
31            let mut values = Vec::new();
32            let mut idx = Vec::new();
33            let sparse = vector.as_f32_sparse();
34            for (&i, &value) in sparse.idx.iter().zip(sparse.values.iter()) {
35                let i = i as usize;
36                if i < start || i >= end {
37                    continue;
38                }
39                values.extend_from_slice(&value.to_le_bytes());
40                idx.extend_from_slice(&((i - start) as u32).to_le_bytes());
41            }
42            values.extend_from_slice(&idx);
43            Ok(Vector {
44                vector_type: vector.vector_type,
45                dims: end - start,
46                owned: Some(values),
47                refer: None,
48            })
49        }
50        VectorType::Float1Bit | VectorType::Float8 => Err(LimboError::ConversionError(
51            "vector_slice is not supported for float1bit/float8 vectors".to_string(),
52        )),
53    }
54}
55
56#[cfg(clt_turso_tests)]
57mod tests {
58    use crate::vector::{
59        operations::slice::vector_slice,
60        vector_types::{Vector, VectorType},
61    };
62
63    fn float32_vec_from(slice: &[f32]) -> Vector {
64        let mut data = Vec::new();
65        for &v in slice {
66            data.extend_from_slice(&v.to_le_bytes());
67        }
68
69        Vector {
70            vector_type: VectorType::Float32Dense,
71            dims: slice.len(),
72            owned: Some(data),
73            refer: None,
74        }
75    }
76
77    fn f32_slice_from_vector(vector: &Vector) -> Vec<f32> {
78        vector.as_f32_slice().to_vec()
79    }
80
81    #[test]
82    fn test_vector_slice_normal_case() {
83        let input_vec = float32_vec_from(&[1.0, 2.0, 3.0, 4.0, 5.0]);
84        let result = vector_slice(&input_vec, 1, 4).unwrap();
85
86        assert_eq!(result.dims, 3);
87        assert_eq!(f32_slice_from_vector(&result), vec![2.0, 3.0, 4.0]);
88    }
89
90    #[test]
91    fn test_vector_slice_full_range() {
92        let input_vec = float32_vec_from(&[10.0, 20.0, 30.0]);
93        let result = vector_slice(&input_vec, 0, 3).unwrap();
94
95        assert_eq!(result.dims, 3);
96        assert_eq!(f32_slice_from_vector(&result), vec![10.0, 20.0, 30.0]);
97    }
98
99    #[test]
100    fn test_vector_slice_single_element() {
101        let input_vec = float32_vec_from(&[4.40, 2.71]);
102        let result = vector_slice(&input_vec, 1, 2).unwrap();
103
104        assert_eq!(result.dims, 1);
105        assert_eq!(f32_slice_from_vector(&result), vec![2.71]);
106    }
107
108    #[test]
109    fn test_vector_slice_empty_list() {
110        let input_vec = float32_vec_from(&[1.0, 2.0]);
111        let result = vector_slice(&input_vec, 2, 2).unwrap();
112
113        assert_eq!(result.dims, 0);
114    }
115
116    #[test]
117    fn test_vector_slice_zero_length() {
118        let input_vec = float32_vec_from(&[1.0, 2.0, 3.0]);
119        let err = vector_slice(&input_vec, 2, 1);
120        assert!(err.is_err(), "Expected error on zero-length range");
121    }
122
123    #[test]
124    fn test_vector_slice_out_of_bounds() {
125        let input_vec = float32_vec_from(&[1.0, 2.0]);
126        let err = vector_slice(&input_vec, 0, 5);
127        assert!(err.is_err());
128    }
129
130    #[test]
131    fn test_vector_slice_start_out_of_bounds() {
132        let input_vec = float32_vec_from(&[1.0, 2.0]);
133        let err = vector_slice(&input_vec, 5, 5);
134        assert!(err.is_err());
135    }
136
137    #[test]
138    fn test_vector_slice_end_out_of_bounds() {
139        let input_vec = float32_vec_from(&[1.0, 2.0]);
140        let err = vector_slice(&input_vec, 1, 3);
141        assert!(err.is_err());
142    }
143}