risingwave_common/array/
vector_array.rs

1// Copyright 2025 RisingWave Labs
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use std::fmt::Debug;
16use std::hash::Hash;
17use std::slice;
18use std::sync::LazyLock;
19
20use bytes::{Buf, BufMut};
21use itertools::{Itertools, repeat_n};
22use memcomparable::Error;
23use risingwave_common::types::F32;
24use risingwave_common_estimate_size::EstimateSize;
25use risingwave_pb::common::PbBuffer;
26use risingwave_pb::common::buffer::PbCompressionType;
27use risingwave_pb::data::{PbArray, PbArrayType, PbListArrayData};
28use serde::{Deserialize, Serialize};
29
30use super::{Array, ArrayBuilder};
31use crate::bitmap::{Bitmap, BitmapBuilder};
32use crate::types::{DataType, ListType, Scalar, ScalarRef, ToText};
33use crate::vector::{VectorInner, decode_vector_payload, encode_vector_payload};
34
35pub type VectorItemType = F32;
36pub type VectorDistanceType = f64;
37pub const VECTOR_ITEM_TYPE: DataType = DataType::Float32;
38pub const VECTOR_DISTANCE_TYPE: DataType = DataType::Float64;
39
40/// Sometimes we can interpret a vector as a list to reuse some code, pass this type around.
41pub static VECTOR_AS_LIST_TYPE: LazyLock<ListType> =
42    LazyLock::new(|| ListType::new(VECTOR_ITEM_TYPE));
43
44#[derive(Debug, Clone, EstimateSize)]
45pub struct VectorArrayBuilder {
46    bitmap: BitmapBuilder,
47    offsets: Vec<u32>,
48    inner: Vec<VectorItemType>,
49    elem_size: usize,
50}
51
52impl ArrayBuilder for VectorArrayBuilder {
53    type ArrayType = VectorArray;
54
55    #[cfg(not(test))]
56    fn new(_capacity: usize) -> Self {
57        panic!("please use `VectorArrayBuilder::with_type` instead");
58    }
59
60    #[cfg(test)]
61    fn new(capacity: usize) -> Self {
62        Self::with_type(capacity, VectorVal::test_type())
63    }
64
65    fn with_type(capacity: usize, ty: DataType) -> Self {
66        let DataType::Vector(elem_size) = ty else {
67            panic!("VectorArrayBuilder only supports Vector type");
68        };
69        let mut offsets = Vec::with_capacity(capacity + 1);
70        offsets.push(0);
71        Self {
72            bitmap: BitmapBuilder::with_capacity(capacity),
73            offsets,
74            inner: Vec::with_capacity(capacity * elem_size),
75            elem_size,
76        }
77    }
78
79    fn append_n(&mut self, n: usize, value: Option<VectorRef<'_>>) {
80        let last = self
81            .offsets
82            .last()
83            .cloned()
84            .expect("non-empty with an initial 0");
85        if let Some(value) = value {
86            assert_eq!(self.elem_size, value.inner.len());
87            self.inner.reserve(self.elem_size * n);
88            for _ in 0..n {
89                self.inner.extend_from_slice(value.inner);
90            }
91            self.offsets.reserve(n);
92            self.offsets.extend((1..=n).map(|i| {
93                last.checked_add((i * self.elem_size) as _)
94                    .expect("overflow")
95            }));
96            self.bitmap.append_n(n, true);
97        } else {
98            self.offsets.reserve(n);
99            self.offsets.extend(repeat_n(last, n));
100            self.bitmap.append_n(n, false);
101        }
102    }
103
104    fn append_array(&mut self, other: &VectorArray) {
105        assert_eq!(self.elem_size, other.elem_size);
106        self.bitmap.append_bitmap(&other.bitmap);
107        let last = self
108            .offsets
109            .last()
110            .cloned()
111            .expect("non-empty with an initial 0");
112        let other_offsets = &other.offsets[1..];
113        self.offsets.reserve(other_offsets.len());
114        self.offsets.extend(
115            other_offsets
116                .iter()
117                .map(|offset| last.checked_add(*offset).expect("overflow")),
118        );
119        self.inner.reserve(other.inner.len());
120        self.inner.extend_from_slice(&other.inner);
121    }
122
123    fn pop(&mut self) -> Option<()> {
124        if self.bitmap.pop().is_some() {
125            self.offsets
126                .pop()
127                .expect("non-empty when bitmap popped Some");
128            let last = self
129                .offsets
130                .last()
131                .cloned()
132                .expect("non-empty with initial 0");
133            self.inner.truncate(last as _);
134            Some(())
135        } else {
136            None
137        }
138    }
139
140    fn len(&self) -> usize {
141        self.bitmap.len()
142    }
143
144    fn finish(self) -> VectorArray {
145        VectorArray {
146            bitmap: self.bitmap.finish(),
147            offsets: self.offsets,
148            inner: self.inner,
149            elem_size: self.elem_size,
150        }
151    }
152}
153
154#[derive(Debug, Clone)]
155pub struct VectorArray {
156    bitmap: Bitmap,
157    /// Of size as `bitmap.len() + 1`. `(self.offsets[i]..self.offsets[i+1])` is the slice range of the i-th vector
158    /// if it's not null.
159    offsets: Vec<u32>,
160    inner: Vec<VectorItemType>,
161    elem_size: usize,
162}
163
164impl EstimateSize for VectorArray {
165    fn estimated_heap_size(&self) -> usize {
166        self.inner.estimated_heap_size()
167    }
168}
169
170impl Array for VectorArray {
171    type Builder = VectorArrayBuilder;
172    type OwnedItem = VectorVal;
173    type RefItem<'a> = VectorRef<'a>;
174
175    unsafe fn raw_value_at_unchecked(&self, idx: usize) -> Self::RefItem<'_> {
176        VectorRef {
177            inner: unsafe {
178                let start = self.inner.as_ptr().add(self.offsets[idx] as usize);
179                slice::from_raw_parts(start, self.elem_size)
180            },
181        }
182    }
183
184    fn len(&self) -> usize {
185        self.bitmap.len()
186    }
187
188    fn to_protobuf(&self) -> PbArray {
189        let mut payload = Vec::with_capacity(self.inner.len() * size_of::<VectorItemType>());
190        encode_vector_payload(self.inner.as_slice(), &mut payload);
191        PbArray {
192            array_type: PbArrayType::Vector as _,
193            null_bitmap: Some(self.bitmap.to_protobuf()),
194            values: vec![PbBuffer {
195                compression: PbCompressionType::None as _,
196                body: payload,
197            }],
198            struct_array_data: None,
199            list_array_data: Some(
200                PbListArrayData {
201                    offsets: self.offsets.clone(),
202                    value: None,
203                    value_type: Some(DataType::Float32.to_protobuf()),
204                    elem_size: Some(self.elem_size as _),
205                }
206                .into(),
207            ),
208        }
209    }
210
211    fn null_bitmap(&self) -> &Bitmap {
212        &self.bitmap
213    }
214
215    fn into_null_bitmap(self) -> Bitmap {
216        self.bitmap
217    }
218
219    fn set_bitmap(&mut self, bitmap: Bitmap) {
220        self.bitmap = bitmap;
221    }
222
223    fn data_type(&self) -> DataType {
224        DataType::Vector(self.elem_size)
225    }
226}
227
228impl VectorArray {
229    pub fn from_protobuf(
230        array: &risingwave_pb::data::PbArray,
231    ) -> super::ArrayResult<super::ArrayImpl> {
232        // reversing to_protobuf
233        assert_eq!(
234            array.array_type,
235            PbArrayType::Vector as i32,
236            "invalid array type for vector: {}",
237            array.array_type
238        );
239        let bitmap: Bitmap = array.get_null_bitmap()?.into();
240        let encoded_payload = &array.values[0].body;
241        let payload = decode_vector_payload(
242            {
243                let payload_len = encoded_payload.len();
244                let elem_size = size_of::<VectorItemType>();
245                // The payload length must be an exact multiple of the element size.
246                // Reject (panic) otherwise, instead of silently truncating the count.
247                // (Equivalent to the previously-used unstable `usize::checked_exact_div`.)
248                if payload_len.is_multiple_of(elem_size) {
249                    payload_len / elem_size
250                } else {
251                    panic!("invalid payload len {} for vector", payload_len)
252                }
253            },
254            array.values[0].body.as_slice(),
255        );
256        let array_data = array.get_list_array_data()?;
257        let elem_size = array_data.elem_size.expect("should exist for Vector") as usize;
258        let offsets = array_data.offsets.clone();
259        debug_assert_eq!(array_data.value_type, Some(DataType::Float32.to_protobuf()));
260        debug_assert_eq!(array_data.value, None);
261
262        Ok(VectorArray {
263            bitmap,
264            offsets,
265            inner: payload,
266            elem_size,
267        }
268        .into())
269    }
270
271    pub fn as_raw_slice(&self) -> &[f32] {
272        F32::inner_slice(&self.inner)
273    }
274
275    pub fn offsets(&self) -> &[u32] {
276        &self.offsets
277    }
278}
279
280pub type VectorVal = VectorInner<Box<[VectorItemType]>>;
281
282impl Debug for VectorVal {
283    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
284        self.to_ref().fmt(f)
285    }
286}
287
288impl Scalar for VectorVal {
289    type ScalarRefType<'a> = VectorRef<'a>;
290
291    fn as_scalar_ref(&self) -> VectorRef<'_> {
292        self.to_ref()
293    }
294}
295
296impl VectorVal {
297    #[cfg(test)]
298    pub const TEST_VECTOR_DIMENSION: usize = 3;
299
300    pub fn from_text(text: &str, size: usize) -> Result<Self, String> {
301        let text = text.trim();
302        let text = text
303            .strip_prefix('[')
304            .ok_or_else(|| "vector must start with [".to_owned())?
305            .strip_suffix(']')
306            .ok_or_else(|| "vector must end with ]".to_owned())?;
307        let inner = text
308            .split(',')
309            .map(|s| {
310                s.trim()
311                    .parse::<f32>()
312                    .map_err(|_| format!("invalid f32: {s}"))
313                    .and_then(|f| {
314                        if f.is_finite() {
315                            Ok(f.into())
316                        } else {
317                            Err(format!("{f} not allowed in vector"))
318                        }
319                    })
320            })
321            .collect::<Result<Vec<_>, _>>()?;
322        if inner.len() != size {
323            return Err(format!("expected {} dimensions, not {}", size, inner.len()));
324        }
325        Ok(Self {
326            inner: inner.into(),
327        })
328    }
329
330    #[cfg(test)]
331    pub fn test_type() -> DataType {
332        DataType::Vector(Self::TEST_VECTOR_DIMENSION)
333    }
334
335    pub fn to_ref(&self) -> VectorRef<'_> {
336        VectorRef { inner: &self.inner }
337    }
338}
339
340/// A `f32` without nan/inf/-inf. Added as intermediate type to `try_collect` `f32` values into a `VectorVal`.
341#[derive(Clone, Copy, Debug)]
342#[repr(transparent)]
343pub struct Finite32(f32);
344
345impl TryFrom<f32> for Finite32 {
346    type Error = String;
347
348    fn try_from(value: f32) -> Result<Self, Self::Error> {
349        if value.is_finite() {
350            Ok(Self(value))
351        } else {
352            Err(format!("{value} not allowed in vector"))
353        }
354    }
355}
356
357impl From<Vec<Finite32>> for VectorVal {
358    fn from(value: Vec<Finite32>) -> Self {
359        let (ptr, len, cap) = value.into_raw_parts();
360        // Safety: OrderedFloat is #[repr(transparent)] and has no invalid values.
361        Self {
362            inner: unsafe { Vec::from_raw_parts(ptr as *mut F32, len, cap).into_boxed_slice() },
363        }
364    }
365}
366
367impl FromIterator<Finite32> for VectorVal {
368    fn from_iter<I: IntoIterator<Item = Finite32>>(iter: I) -> Self {
369        let inner = iter.into_iter().collect_vec();
370        Self::from(inner)
371    }
372}
373
374pub type VectorRef<'a> = VectorInner<&'a [VectorItemType]>;
375
376impl Debug for VectorRef<'_> {
377    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
378        self.write_with_type(&DataType::Vector(self.dimension()), f)
379    }
380}
381
382impl ToText for VectorRef<'_> {
383    fn write<W: std::fmt::Write>(&self, f: &mut W) -> std::fmt::Result {
384        self.write_with_type(&DataType::Vector(self.dimension()), f)
385    }
386
387    fn write_with_type<W: std::fmt::Write>(&self, _ty: &DataType, f: &mut W) -> std::fmt::Result {
388        write!(f, "[")?;
389        for (i, item) in self.inner.iter().enumerate() {
390            if i > 0 {
391                write!(f, ",")?;
392            }
393            write!(f, "{}", item)?;
394        }
395        write!(f, "]")
396    }
397}
398
399impl<'a> ScalarRef<'a> for VectorRef<'a> {
400    type ScalarType = VectorVal;
401
402    fn to_owned_scalar(&self) -> VectorVal {
403        VectorVal {
404            inner: self.inner.to_vec().into_boxed_slice(),
405        }
406    }
407
408    fn hash_scalar<H: std::hash::Hasher>(&self, state: &mut H) {
409        self.inner.hash(state);
410    }
411}
412
413impl<'a> VectorRef<'a> {
414    /// Create a `VectorRef` from a slice of `VectorItemType` without checking the elements in the slice
415    /// is invalid, such as `inf` and `nan`.
416    pub fn from_slice_unchecked(inner: &'a [VectorItemType]) -> Self {
417        Self { inner }
418    }
419
420    pub fn memcmp_serialize(
421        self,
422        serializer: &mut memcomparable::Serializer<impl BufMut>,
423    ) -> memcomparable::Result<()> {
424        for item in self.inner {
425            item.serialize(&mut *serializer)?;
426        }
427        Ok(())
428    }
429
430    pub fn subvector(&self, start: usize, end: usize) -> VectorVal {
431        let slice = &self.inner[start..end];
432        VectorInner {
433            inner: slice.to_vec().into_boxed_slice(),
434        }
435    }
436}
437
438impl VectorVal {
439    pub fn memcmp_deserialize(
440        dimension: usize,
441        de: &mut memcomparable::Deserializer<impl Buf>,
442    ) -> memcomparable::Result<Self> {
443        let mut value = Vec::with_capacity(dimension);
444        for _ in 0..dimension {
445            value.push(Finite32::try_from(f32::deserialize(&mut *de)?).map_err(Error::Message)?)
446        }
447        Ok(VectorVal::from(value))
448    }
449}
risingwave_common/array/vector_array.rs

risingwave_common/array/
vector_array.rs
