From 2483d08bca3f78a0b04456d052deb772064e6e5e Mon Sep 17 00:00:00 2001 From: Nikita Sivukhin Date: Tue, 21 Oct 2025 16:28:00 +0400 Subject: [PATCH] do not allocate if possible --- core/vector/mod.rs | 4 +- core/vector/operations/concat.rs | 26 ++++--- core/vector/operations/convert.rs | 20 +++-- core/vector/operations/serialize.rs | 17 +++-- core/vector/operations/slice.rs | 14 ++-- core/vector/operations/text.rs | 18 +++-- core/vector/vector_types.rs | 111 +++++++++++++++++++--------- 7 files changed, 136 insertions(+), 74 deletions(-) diff --git a/core/vector/mod.rs b/core/vector/mod.rs index 514780a78..14cb2a462 100644 --- a/core/vector/mod.rs +++ b/core/vector/mod.rs @@ -20,7 +20,7 @@ pub fn parse_vector(value: &Register, type_hint: Option) -> Result Err(LimboError::ConversionError( "Invalid vector type".to_string(), @@ -81,7 +81,7 @@ pub fn vector_extract(args: &[Register]) -> Result { return Ok(Value::build_text("[]")); } - let vector = Vector::from_blob(blob.to_vec())?; + let vector = Vector::from_vec(blob.to_vec())?; Ok(Value::build_text(operations::text::vector_to_text(&vector))) } diff --git a/core/vector/operations/concat.rs b/core/vector/operations/concat.rs index 3e7f6a4f1..178258838 100644 --- a/core/vector/operations/concat.rs +++ b/core/vector/operations/concat.rs @@ -3,7 +3,7 @@ use crate::{ LimboError, Result, }; -pub fn vector_concat(v1: &Vector, v2: &Vector) -> Result { +pub fn vector_concat(v1: &Vector, v2: &Vector) -> Result> { if v1.vector_type != v2.vector_type { return Err(LimboError::ConversionError( "Mismatched vector types".into(), @@ -12,17 +12,17 @@ pub fn vector_concat(v1: &Vector, v2: &Vector) -> Result { let data = match v1.vector_type { VectorType::Float32Dense | VectorType::Float64Dense => { - let mut data = Vec::with_capacity(v1.data.len() + v2.data.len()); - data.extend_from_slice(&v1.data); - data.extend_from_slice(&v2.data); + let mut data = Vec::with_capacity(v1.bin_len() + v2.bin_len()); + data.extend_from_slice(&v1.bin_data()); + data.extend_from_slice(&v2.bin_data()); data } VectorType::Float32Sparse => { - let mut data = Vec::with_capacity(v1.data.len() + v2.data.len()); - data.extend_from_slice(&v1.data[..v1.data.len() / 2]); - data.extend_from_slice(&v2.data[..v2.data.len() / 2]); - data.extend_from_slice(&v1.data[v1.data.len() / 2..]); - data.extend_from_slice(&v2.data[v2.data.len() / 2..]); + let mut data = Vec::with_capacity(v1.bin_len() + v2.bin_len()); + data.extend_from_slice(&v1.bin_data()[..v1.bin_len() / 2]); + data.extend_from_slice(&v2.bin_data()[..v2.bin_len() / 2]); + data.extend_from_slice(&v1.bin_data()[v1.bin_len() / 2..]); + data.extend_from_slice(&v2.bin_data()[v2.bin_len() / 2..]); data } }; @@ -30,7 +30,8 @@ pub fn vector_concat(v1: &Vector, v2: &Vector) -> Result { Ok(Vector { vector_type: v1.vector_type, dims: v1.dims + v2.dims, - data, + owned: Some(data), + refer: None, }) } @@ -41,7 +42,7 @@ mod tests { vector_types::{Vector, VectorType}, }; - fn float32_vec_from(slice: &[f32]) -> Vector { + fn float32_vec_from(slice: &[f32]) -> Vector<'static> { let mut data = Vec::new(); for &v in slice { data.extend_from_slice(&v.to_le_bytes()); @@ -50,7 +51,8 @@ mod tests { Vector { vector_type: VectorType::Float32Dense, dims: slice.len(), - data, + owned: Some(data), + refer: None, } } diff --git a/core/vector/operations/convert.rs b/core/vector/operations/convert.rs index 619c7be26..1db0ab99e 100644 --- a/core/vector/operations/convert.rs +++ b/core/vector/operations/convert.rs @@ -69,7 +69,7 @@ mod tests { fn assert_vectors(v1: &Vector, v2: &Vector) { assert_eq!(v1.vector_type, v2.vector_type); assert_eq!(v1.dims, v2.dims); - assert_eq!(v1.data, v2.data); + assert_eq!(v1.bin_data(), v2.bin_data()); } #[test] @@ -77,30 +77,33 @@ mod tests { let vf32 = Vector { vector_type: VectorType::Float32Dense, dims: 3, - data: concat(&[ + owned: Some(concat(&[ 1.0f32.to_le_bytes(), 0.0f32.to_le_bytes(), 2.0f32.to_le_bytes(), - ]), + ])), + refer: None, }; let vf64 = Vector { vector_type: VectorType::Float64Dense, dims: 3, - data: concat(&[ + owned: Some(concat(&[ 1.0f64.to_le_bytes(), 0.0f64.to_le_bytes(), 2.0f64.to_le_bytes(), - ]), + ])), + refer: None, }; let vf32_sparse = Vector { vector_type: VectorType::Float32Sparse, dims: 3, - data: concat(&[ + owned: Some(concat(&[ 1.0f32.to_le_bytes(), 2.0f32.to_le_bytes(), 0u32.to_le_bytes(), 2u32.to_le_bytes(), - ]), + ])), + refer: None, }; let vectors = [vf32, vf64, vf32_sparse]; @@ -110,7 +113,8 @@ mod tests { let v_copy = Vector { vector_type: v1.vector_type, dims: v1.dims, - data: v1.data.clone(), + owned: v1.owned.clone(), + refer: None, }; assert_vectors(&vector_convert(v_copy, v2.vector_type).unwrap(), v2); } diff --git a/core/vector/operations/serialize.rs b/core/vector/operations/serialize.rs index 622819a59..8f8c3af9d 100644 --- a/core/vector/operations/serialize.rs +++ b/core/vector/operations/serialize.rs @@ -3,17 +3,20 @@ use crate::{ Value, }; -pub fn vector_serialize(mut x: Vector) -> Value { +pub fn vector_serialize(x: Vector) -> Value { match x.vector_type { - VectorType::Float32Dense => Value::from_blob(x.data), + VectorType::Float32Dense => Value::from_blob(x.bin_eject()), VectorType::Float64Dense => { - x.data.push(2); - Value::from_blob(x.data) + let mut data = x.bin_eject(); + data.push(2); + Value::from_blob(data) } VectorType::Float32Sparse => { - x.data.extend_from_slice(&(x.dims as u32).to_le_bytes()); - x.data.push(9); - Value::from_blob(x.data) + let dims = x.dims; + let mut data = x.bin_eject(); + data.extend_from_slice(&(dims as u32).to_le_bytes()); + data.push(9); + Value::from_blob(data) } } } diff --git a/core/vector/operations/slice.rs b/core/vector/operations/slice.rs index 003195070..a1f6b99d8 100644 --- a/core/vector/operations/slice.rs +++ b/core/vector/operations/slice.rs @@ -3,7 +3,7 @@ use crate::{ LimboError, Result, }; -pub fn vector_slice(vector: &Vector, start: usize, end: usize) -> Result { +pub fn vector_slice(vector: &Vector, start: usize, end: usize) -> Result> { if start > end { return Err(LimboError::InvalidArgument( "start index must not be greater than end index".into(), @@ -18,12 +18,14 @@ pub fn vector_slice(vector: &Vector, start: usize, end: usize) -> Result VectorType::Float32Dense => Ok(Vector { vector_type: vector.vector_type, dims: end - start, - data: vector.data[start * 4..end * 4].to_vec(), + owned: Some(vector.bin_data()[start * 4..end * 4].to_vec()), + refer: None, }), VectorType::Float64Dense => Ok(Vector { vector_type: vector.vector_type, dims: end - start, - data: vector.data[start * 8..end * 8].to_vec(), + owned: Some(vector.bin_data()[start * 8..end * 8].to_vec()), + refer: None, }), VectorType::Float32Sparse => { let mut values = Vec::new(); @@ -41,7 +43,8 @@ pub fn vector_slice(vector: &Vector, start: usize, end: usize) -> Result Ok(Vector { vector_type: vector.vector_type, dims: end - start, - data: values, + owned: Some(values), + refer: None, }) } } @@ -63,7 +66,8 @@ mod tests { Vector { vector_type: VectorType::Float32Dense, dims: slice.len(), - data, + owned: Some(data), + refer: None, } } diff --git a/core/vector/operations/text.rs b/core/vector/operations/text.rs index c6403d812..810bd8bc0 100644 --- a/core/vector/operations/text.rs +++ b/core/vector/operations/text.rs @@ -56,7 +56,8 @@ pub fn vector_from_text(vector_type: VectorType, text: &str) -> Result { Vector { vector_type, dims: 0, - data: Vec::new(), + owned: Some(Vec::new()), + refer: None, } } }); @@ -69,7 +70,7 @@ pub fn vector_from_text(vector_type: VectorType, text: &str) -> Result { } } -fn vector32_from_text<'a>(tokens: impl Iterator) -> Result { +fn vector32_from_text<'a>(tokens: impl Iterator) -> Result> { let mut data = Vec::new(); for token in tokens { let value = token @@ -85,11 +86,12 @@ fn vector32_from_text<'a>(tokens: impl Iterator) -> Result(tokens: impl Iterator) -> Result { +fn vector64_from_text<'a>(tokens: impl Iterator) -> Result> { let mut data = Vec::new(); for token in tokens { let value = token @@ -105,11 +107,12 @@ fn vector64_from_text<'a>(tokens: impl Iterator) -> Result(tokens: impl Iterator) -> Result { +fn vector32_sparse_from_text<'a>(tokens: impl Iterator) -> Result> { let mut idx = Vec::new(); let mut values = Vec::new(); let mut dims = 0u32; @@ -135,6 +138,7 @@ fn vector32_sparse_from_text<'a>(tokens: impl Iterator) -> Resul Ok(Vector { vector_type: VectorType::Float32Sparse, dims: dims as usize, - data: values, + owned: Some(values), + refer: None, }) } diff --git a/core/vector/vector_types.rs b/core/vector/vector_types.rs index dafa2ca23..c6c599090 100644 --- a/core/vector/vector_types.rs +++ b/core/vector/vector_types.rs @@ -8,10 +8,11 @@ pub enum VectorType { } #[derive(Debug)] -pub struct Vector { +pub struct Vector<'a> { pub vector_type: VectorType, pub dims: usize, - pub data: Vec, + pub owned: Option>, + pub refer: Option<&'a [u8]>, } #[derive(Debug)] @@ -20,14 +21,14 @@ pub struct VectorSparse<'a, T: std::fmt::Debug> { pub values: &'a [T], } -impl Vector { - pub fn vector_type(mut blob: Vec) -> Result<(VectorType, Vec)> { +impl<'a> Vector<'a> { + pub fn vector_type(blob: &[u8]) -> Result<(VectorType, usize)> { // Even-sized blobs are always float32. if blob.len() % 2 == 0 { - return Ok((VectorType::Float32Dense, blob)); + return Ok((VectorType::Float32Dense, blob.len())); } // Odd-sized blobs have type byte at the end - let vector_type = blob.pop().unwrap(); + let vector_type = blob[blob.len() - 1]; /* vector types used by LibSQL: (see https://github.com/tursodatabase/libsql/blob/a55bf61192bdb89e97568de593c4af5b70d24bde/libsql-sqlite3/src/vectorInt.h#L52) @@ -39,12 +40,12 @@ impl Vector { #define VECTOR_TYPE_FLOATB16 6 */ match vector_type { - 1 => Ok((VectorType::Float32Dense, blob)), - 2 => Ok((VectorType::Float64Dense, blob)), + 1 => Ok((VectorType::Float32Dense, blob.len() - 1)), + 2 => Ok((VectorType::Float64Dense, blob.len() - 1)), 3..=6 => Err(LimboError::ConversionError( "unsupported vector type from LibSQL".to_string(), )), - 9 => Ok((VectorType::Float32Sparse, blob)), + 9 => Ok((VectorType::Float32Sparse, blob.len() - 1)), _ => Err(LimboError::ConversionError(format!( "unknown vector type: {vector_type}" ))), @@ -63,7 +64,8 @@ impl Vector { Self { vector_type: VectorType::Float32Dense, dims, - data: values, + owned: Some(values), + refer: None, } } pub fn from_f64(mut values_f64: Vec) -> Self { @@ -79,7 +81,8 @@ impl Vector { Self { vector_type: VectorType::Float64Dense, dims, - data: values, + owned: Some(values), + refer: None, } } pub fn from_f32_sparse(dims: usize, mut values_f32: Vec, mut idx_u32: Vec) -> Self { @@ -105,14 +108,27 @@ impl Vector { Self { vector_type: VectorType::Float32Sparse, dims, - data: values, + owned: Some(values), + refer: None, } } - pub fn from_blob(blob: Vec) -> Result { - let (vector_type, data) = Self::vector_type(blob)?; - Self::from_data(vector_type, data) + pub fn from_vec(mut blob: Vec) -> Result { + let (vector_type, len) = Self::vector_type(&blob)?; + blob.truncate(len); + Self::from_data(vector_type, Some(blob), None) } - pub fn from_data(vector_type: VectorType, mut data: Vec) -> Result { + pub fn from_slice(blob: &'a [u8]) -> Result { + let (vector_type, len) = Self::vector_type(&blob)?; + Self::from_data(vector_type, None, Some(&blob[..len])) + } + pub fn from_data( + vector_type: VectorType, + owned: Option>, + refer: Option<&'a [u8]>, + ) -> Result { + let owned_slice = owned.as_ref().map(|x| x.as_slice()); + let refer_slice = refer.as_ref().map(|&x| x); + let data = owned_slice.unwrap_or_else(|| refer_slice.unwrap()); match vector_type { VectorType::Float32Dense => { if data.len() % 4 != 0 { @@ -124,7 +140,8 @@ impl Vector { Ok(Vector { vector_type, dims: data.len() / 4, - data, + owned, + refer, }) } VectorType::Float64Dense => { @@ -137,7 +154,8 @@ impl Vector { Ok(Vector { vector_type, dims: data.len() / 8, - data, + owned, + refer, }) } VectorType::Float32Sparse => { @@ -147,17 +165,41 @@ impl Vector { data.len(), ))); } - let dims_bytes = data.split_off(data.len() - 4); + let original_len = data.len(); + let dims_bytes = &data[original_len - 4..]; let dims = u32::from_le_bytes(dims_bytes.try_into().unwrap()) as usize; + let owned = owned.map(|mut x| { + x.truncate(original_len - 4); + x + }); + let refer = refer.map(|x| &x[0..original_len - 4]); let vector = Vector { vector_type, dims, - data, + owned, + refer, }; Ok(vector) } } } + + pub fn bin_len(&self) -> usize { + let owned = self.owned.as_ref().map(|x| x.len()); + let refer = self.refer.as_ref().map(|x| x.len()); + owned.unwrap_or_else(|| refer.unwrap()) + } + + pub fn bin_data(&'a self) -> &'a [u8] { + let owned = self.owned.as_ref().map(|x| x.as_slice()); + let refer = self.refer.as_ref().map(|&x| x); + owned.unwrap_or_else(|| refer.unwrap()) + } + + pub fn bin_eject(self) -> Vec { + self.owned.unwrap_or_else(|| self.refer.unwrap().to_vec()) + } + /// # Safety /// /// This method is used to reinterpret the underlying `Vec` data @@ -171,12 +213,12 @@ impl Vector { } assert_eq!( - self.data.len(), + self.bin_len(), self.dims * std::mem::size_of::(), "data length must equal dims * size_of::()" ); - let ptr = self.data.as_ptr(); + let ptr = self.bin_data().as_ptr(); let align = std::mem::align_of::(); assert_eq!( ptr.align_offset(align), @@ -200,12 +242,12 @@ impl Vector { } assert_eq!( - self.data.len(), + self.bin_len(), self.dims * std::mem::size_of::(), "data length must equal dims * size_of::()" ); - let ptr = self.data.as_ptr(); + let ptr = self.bin_data().as_ptr(); let align = std::mem::align_of::(); assert_eq!( ptr.align_offset(align), @@ -218,14 +260,14 @@ impl Vector { pub fn as_f32_sparse(&self) -> VectorSparse<'_, f32> { debug_assert!(self.vector_type == VectorType::Float32Sparse); - let ptr = self.data.as_ptr(); + let ptr = self.bin_data().as_ptr(); let align = std::mem::align_of::(); assert_eq!( ptr.align_offset(align), 0, "data pointer must be aligned to {align} bytes for f32 access" ); - let length = self.data.len() / 4 / 2; + let length = self.bin_data().len() / 4 / 2; let values = unsafe { std::slice::from_raw_parts(ptr as *const f32, length) }; let idx = unsafe { std::slice::from_raw_parts((ptr as *const u32).add(length), length) }; debug_assert!(idx.is_sorted()); @@ -292,12 +334,13 @@ pub(crate) mod tests { } /// Convert an ArbitraryVector to a Vector. - impl From> for Vector { + impl From> for Vector<'static> { fn from(v: ArbitraryVector) -> Self { Vector { vector_type: v.vector_type, dims: DIMS, - data: v.data, + owned: Some(v.data), + refer: None, } } } @@ -357,7 +400,7 @@ pub(crate) mod tests { let vtype = v.vector_type; let value = operations::serialize::vector_serialize(v); let blob = value.to_blob().unwrap().to_vec(); - match Vector::vector_type(blob) { + match Vector::vector_type(&blob) { Ok((detected_type, _)) => detected_type == vtype, Err(_) => false, } @@ -396,12 +439,12 @@ pub(crate) mod tests { VectorType::Float32Dense => { let slice = v.as_f32_slice(); // Check if the slice length matches the dimensions and the data length is correct (4 bytes per float) - slice.len() == DIMS && (slice.len() * 4 == v.data.len()) + slice.len() == DIMS && (slice.len() * 4 == v.bin_len()) } VectorType::Float64Dense => { let slice = v.as_f64_slice(); // Check if the slice length matches the dimensions and the data length is correct (8 bytes per float) - slice.len() == DIMS && (slice.len() * 8 == v.data.len()) + slice.len() == DIMS && (slice.len() * 8 == v.bin_len()) } _ => unreachable!(), } @@ -454,12 +497,14 @@ pub(crate) mod tests { let a = Vector { vector_type: VectorType::Float32Dense, dims: 2, - data: vec![0, 0, 0, 0, 52, 208, 106, 63], + owned: Some(vec![0, 0, 0, 0, 52, 208, 106, 63]), + refer: None, }; let b = Vector { vector_type: VectorType::Float32Dense, dims: 2, - data: vec![0, 0, 0, 0, 58, 100, 45, 192], + owned: Some(vec![0, 0, 0, 0, 58, 100, 45, 192]), + refer: None, }; assert!( (operations::distance_cos::vector_distance_cos(&a, &b).unwrap() - 2.0).abs() <= 1e-6