vortex_array/arrays/varbin/
array.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright the Vortex contributors
3
4use std::fmt::Display;
5use std::fmt::Formatter;
6
7use num_traits::AsPrimitive;
8use vortex_array::arrays::PrimitiveArray;
9use vortex_buffer::ByteBuffer;
10use vortex_error::VortexExpect;
11use vortex_error::VortexResult;
12use vortex_error::vortex_ensure;
13use vortex_error::vortex_err;
14
15use crate::ArrayRef;
16use crate::ArraySlots;
17use crate::VortexSessionExecute;
18use crate::array::Array;
19use crate::array::ArrayParts;
20use crate::array::TypedArrayRef;
21use crate::array::child_to_validity;
22use crate::array::validity_to_child;
23use crate::array_slots;
24use crate::arrays::VarBin;
25use crate::arrays::varbin::builder::VarBinBuilder;
26use crate::buffer::BufferHandle;
27use crate::dtype::DType;
28use crate::dtype::IntegerPType;
29use crate::dtype::Nullability;
30use crate::legacy_session;
31use crate::match_each_integer_ptype;
32use crate::validity::Validity;
33
34#[array_slots(VarBin)]
35pub struct VarBinSlots {
36    /// The offsets array defining the start/end of each variable-length binary element.
37    pub offsets: ArrayRef,
38    /// The validity bitmap indicating which elements are non-null.
39    pub validity: Option<ArrayRef>,
40}
41
42#[derive(Clone, Debug)]
43pub struct VarBinData {
44    pub(super) bytes: BufferHandle,
45}
46
47impl Display for VarBinData {
48    fn fmt(&self, _f: &mut Formatter<'_>) -> std::fmt::Result {
49        Ok(())
50    }
51}
52
53pub struct VarBinDataParts {
54    pub dtype: DType,
55    pub bytes: BufferHandle,
56    pub offsets: ArrayRef,
57    pub validity: Validity,
58}
59
60impl VarBinData {
61    /// Creates a new `VarBinArray`.
62    ///
63    /// # Panics
64    ///
65    /// Panics if the provided components do not satisfy the invariants documented
66    /// in `VarBinArray::new_unchecked`.
67    pub fn build(offsets: ArrayRef, bytes: ByteBuffer, dtype: DType, validity: Validity) -> Self {
68        Self::try_build(offsets, bytes, dtype, validity).vortex_expect("VarBinArray new")
69    }
70
71    /// Creates a new `VarBinArray`.
72    ///
73    /// # Panics
74    ///
75    /// Panics if the provided components do not satisfy the invariants documented
76    /// in `VarBinArray::new_unchecked`.
77    pub fn build_from_handle(
78        offset: ArrayRef,
79        bytes: BufferHandle,
80        dtype: DType,
81        validity: Validity,
82    ) -> Self {
83        Self::try_build_from_handle(offset, bytes, dtype, validity).vortex_expect("VarBinArray new")
84    }
85
86    pub(crate) fn make_slots(offsets: ArrayRef, validity: &Validity, len: usize) -> ArraySlots {
87        VarBinSlots {
88            offsets,
89            validity: validity_to_child(validity, len),
90        }
91        .into_slots()
92    }
93
94    /// Constructs a new `VarBinArray`.
95    ///
96    /// See `VarBinArray::new_unchecked` for more information.
97    ///
98    /// # Errors
99    ///
100    /// Returns an error if the provided components do not satisfy the invariants documented in
101    /// `VarBinArray::new_unchecked`.
102    pub fn try_build(
103        offsets: ArrayRef,
104        bytes: ByteBuffer,
105        dtype: DType,
106        validity: Validity,
107    ) -> VortexResult<Self> {
108        let bytes = BufferHandle::new_host(bytes);
109        Self::validate(&offsets, &bytes, &dtype, &validity)?;
110
111        // SAFETY: validate ensures all invariants are met.
112        Ok(unsafe { Self::new_unchecked_from_handle(bytes) })
113    }
114
115    /// Constructs a new `VarBinArray` from a `BufferHandle` of memory that may exist
116    /// on the CPU or GPU.
117    ///
118    /// See `VarBinArray::new_unchecked` for more information.
119    ///
120    /// # Errors
121    ///
122    /// Returns an error if the provided components do not satisfy the invariants documented in
123    /// `VarBinArray::new_unchecked`.
124    pub fn try_build_from_handle(
125        offsets: ArrayRef,
126        bytes: BufferHandle,
127        dtype: DType,
128        validity: Validity,
129    ) -> VortexResult<Self> {
130        Self::validate(&offsets, &bytes, &dtype, &validity)?;
131
132        // SAFETY: validate ensures all invariants are met.
133        Ok(unsafe { Self::new_unchecked_from_handle(bytes) })
134    }
135
136    /// Creates a new `VarBinArray` without validation from these components:
137    ///
138    /// * `offsets` is an array of byte offsets into the `bytes` buffer.
139    /// * `bytes` is a buffer containing all the variable-length data concatenated.
140    /// * `dtype` specifies whether this contains UTF-8 strings or binary data.
141    /// * `validity` holds the null values.
142    ///
143    /// # Safety
144    ///
145    /// The caller must ensure all of the following invariants are satisfied:
146    ///
147    /// ## Offsets Requirements
148    ///
149    /// - `offsets` must be a non-nullable integer array.
150    /// - `offsets` must contain at least 1 element (for empty array, it contains \[0\]).
151    /// - All values in `offsets` must be monotonically non-decreasing.
152    /// - The first value in `offsets` must be 0.
153    /// - No offset value may exceed `bytes.len()`.
154    ///
155    /// ## Type Requirements
156    ///
157    /// - `dtype` must be exactly [`DType::Binary`] or [`DType::Utf8`].
158    /// - If `dtype` is [`DType::Utf8`], every byte slice `bytes[offsets[i]..offsets[i+1]]` must be valid UTF-8.
159    /// - `dtype.is_nullable()` must match the nullability of `validity`.
160    ///
161    /// ## Validity Requirements
162    ///
163    /// - If `validity` is [`Validity::Array`], its length must exactly equal `offsets.len() - 1`.
164    pub unsafe fn new_unchecked(bytes: ByteBuffer) -> Self {
165        // SAFETY: `new_unchecked_from_handle` has same invariants which should be checked
166        //  by caller.
167        unsafe { Self::new_unchecked_from_handle(BufferHandle::new_host(bytes)) }
168    }
169
170    /// Creates a new `VarBinArray` without validation from its components, with string data
171    /// stored in a `BufferHandle` (CPU or GPU).
172    ///
173    /// # Safety
174    ///
175    /// The caller must ensure all the invariants documented in `new_unchecked` are satisfied.
176    pub unsafe fn new_unchecked_from_handle(bytes: BufferHandle) -> Self {
177        Self { bytes }
178    }
179
180    /// Validates the components that would be used to create a `VarBinArray`.
181    ///
182    /// This function checks all the invariants required by `VarBinArray::new_unchecked`.
183    pub fn validate(
184        offsets: &ArrayRef,
185        bytes: &BufferHandle,
186        dtype: &DType,
187        validity: &Validity,
188    ) -> VortexResult<()> {
189        // Check offsets are non-nullable integer
190        vortex_ensure!(
191            offsets.dtype().is_int() && !offsets.dtype().is_nullable(),
192            MismatchedTypes: "non nullable int", offsets.dtype()
193        );
194
195        // Check dtype is Binary or Utf8
196        vortex_ensure!(
197            matches!(dtype, DType::Binary(_) | DType::Utf8(_)),
198            MismatchedTypes: "utf8 or binary", dtype
199        );
200
201        // Check nullability matches
202        vortex_ensure!(
203            dtype.is_nullable() != matches!(validity, Validity::NonNullable),
204            InvalidArgument: "incorrect validity {:?} for dtype {}",
205            validity,
206            dtype
207        );
208
209        // Check offsets has at least one element
210        vortex_ensure!(
211            !offsets.is_empty(),
212            InvalidArgument: "Offsets must have at least one element"
213        );
214
215        // Check validity length
216        if let Some(validity_len) = validity.maybe_len() {
217            vortex_ensure!(
218                validity_len == offsets.len() - 1,
219                "Validity length {} doesn't match array length {}",
220                validity_len,
221                offsets.len() - 1
222            );
223        }
224
225        // Validate UTF-8 for Utf8 dtype. Skip when offsets/bytes are not host-resident.
226        if offsets.is_host()
227            && bytes.is_on_host()
228            && matches!(dtype, DType::Utf8(_))
229            && let Some(bytes) = bytes.as_host_opt()
230        {
231            Self::validate_utf8(offsets, bytes.as_ref(), validity)?;
232        }
233
234        Ok(())
235    }
236
237    /// Validates that every non-null value is valid UTF-8.
238    #[allow(clippy::disallowed_methods)]
239    fn validate_utf8(offsets: &ArrayRef, bytes: &[u8], validity: &Validity) -> VortexResult<()> {
240        let validate_at = |i: usize, start: usize, end: usize| -> VortexResult<()> {
241            let string_bytes = &bytes[start..end];
242            simdutf8::basic::from_utf8(string_bytes).map_err(|_| {
243                #[expect(clippy::unwrap_used)]
244                // run validation using `compat` package to get more detailed error message
245                let err = simdutf8::compat::from_utf8(string_bytes).unwrap_err();
246                vortex_err!("invalid utf-8: {err} at index {i}")
247            })?;
248            Ok(())
249        };
250
251        let mut ctx = legacy_session().create_execution_ctx();
252        // TODO(joe): update the created VarBin with this decompressed Array.
253        let primitive_offsets = offsets.clone().execute::<PrimitiveArray>(&mut ctx)?;
254
255        // Array-backed validity is the only variant that needs an execution context: execute it into
256        // a mask once. The constant variants resolve null-ness without one. Resolving this before
257        // the per-type dispatch keeps the dtype loop simple.
258        let mask = match validity {
259            Validity::Array(_) => {
260                Some(validity.execute_mask(primitive_offsets.len().saturating_sub(1), &mut ctx)?)
261            }
262            _ => None,
263        };
264        let all_invalid = validity.definitely_all_null();
265
266        match_each_integer_ptype!(primitive_offsets.dtype().as_ptype(), |O| {
267            let offsets_slice = primitive_offsets.as_slice::<O>();
268
269            let last_offset: usize = offsets_slice[offsets_slice.len() - 1].as_();
270            vortex_ensure!(
271                last_offset <= bytes.len(),
272                InvalidArgument: "Last offset {} exceeds bytes length {}",
273                last_offset,
274                bytes.len()
275            );
276
277            for (i, (start, end)) in offsets_slice
278                .windows(2)
279                .map(|o| (o[0].as_(), o[1].as_()))
280                .enumerate()
281            {
282                let valid = mask.as_ref().map_or(!all_invalid, |mask| mask.value(i));
283                if valid {
284                    validate_at(i, start, end)?;
285                }
286            }
287        });
288        Ok(())
289    }
290
291    /// Access the value bytes child buffer
292    ///
293    /// # Note
294    ///
295    /// Bytes child buffer is never sliced when the array is sliced so this can include values
296    /// that are not logically present in the array. Users should prefer `sliced_bytes`
297    /// unless they're resolving values via the offset child array.
298    #[inline]
299    pub fn bytes(&self) -> &ByteBuffer {
300        self.bytes.as_host()
301    }
302
303    /// Access the value bytes buffer handle.
304    #[inline]
305    pub fn bytes_handle(&self) -> &BufferHandle {
306        &self.bytes
307    }
308}
309
310pub trait VarBinArrayExt: VarBinArraySlotsExt {
311    fn dtype_parts(&self) -> (bool, Nullability) {
312        match self.as_ref().dtype() {
313            DType::Utf8(nullability) => (true, *nullability),
314            DType::Binary(nullability) => (false, *nullability),
315            _ => unreachable!("VarBinArrayExt requires a utf8 or binary dtype"),
316        }
317    }
318
319    fn is_utf8(&self) -> bool {
320        self.dtype_parts().0
321    }
322
323    fn nullability(&self) -> Nullability {
324        self.dtype_parts().1
325    }
326
327    fn varbin_validity(&self) -> Validity {
328        child_to_validity(
329            self.as_ref().slots()[VarBinSlots::VALIDITY].as_ref(),
330            self.nullability(),
331        )
332    }
333
334    #[allow(clippy::disallowed_methods)]
335    fn offset_at(&self, index: usize) -> usize {
336        assert!(
337            index <= self.as_ref().len(),
338            "Index {index} out of bounds 0..={}",
339            self.as_ref().len()
340        );
341
342        (&self
343            .offsets()
344            .execute_scalar(index, &mut legacy_session().create_execution_ctx())
345            .vortex_expect("offsets must support execute_scalar"))
346            .try_into()
347            .vortex_expect("Failed to convert offset to usize")
348    }
349
350    fn bytes_at(&self, index: usize) -> ByteBuffer {
351        let start = self.offset_at(index);
352        let end = self.offset_at(index + 1);
353        self.bytes().slice(start..end)
354    }
355
356    fn sliced_bytes(&self) -> ByteBuffer {
357        let first_offset: usize = self.offset_at(0);
358        let last_offset = self.offset_at(self.as_ref().len());
359        self.bytes().slice(first_offset..last_offset)
360    }
361}
362impl<T: TypedArrayRef<VarBin>> VarBinArrayExt for T {}
363
364/// Forwarding constructors for `VarBinArray` (= `Array<VarBin>`).
365impl Array<VarBin> {
366    pub fn from_vec<T: AsRef<[u8]>>(vec: Vec<T>, dtype: DType) -> Self {
367        let size: usize = vec.iter().map(|v| v.as_ref().len()).sum();
368        if size < u32::MAX as usize {
369            Self::from_vec_sized::<u32, T>(vec, dtype)
370        } else {
371            Self::from_vec_sized::<u64, T>(vec, dtype)
372        }
373    }
374
375    #[expect(
376        clippy::same_name_method,
377        reason = "intentionally named from_iter like Iterator::from_iter"
378    )]
379    pub fn from_iter<T: AsRef<[u8]>, I: IntoIterator<Item = Option<T>>>(
380        iter: I,
381        dtype: DType,
382    ) -> Self {
383        let iter = iter.into_iter();
384        let mut builder = VarBinBuilder::<u32>::with_capacity(iter.size_hint().0);
385        for v in iter {
386            builder.append(v.as_ref().map(|o| o.as_ref()));
387        }
388        builder.finish(dtype)
389    }
390
391    pub fn from_iter_nonnull<T: AsRef<[u8]>, I: IntoIterator<Item = T>>(
392        iter: I,
393        dtype: DType,
394    ) -> Self {
395        let iter = iter.into_iter();
396        let mut builder = VarBinBuilder::<u32>::with_capacity(iter.size_hint().0);
397        for v in iter {
398            builder.append_value(v);
399        }
400        builder.finish(dtype)
401    }
402
403    fn from_vec_sized<O, T>(vec: Vec<T>, dtype: DType) -> Self
404    where
405        O: IntegerPType,
406        T: AsRef<[u8]>,
407    {
408        let mut builder = VarBinBuilder::<O>::with_capacity(vec.len());
409        for v in vec {
410            builder.append_value(v.as_ref());
411        }
412        builder.finish(dtype)
413    }
414
415    /// Create from a vector of string slices.
416    pub fn from_strs(value: Vec<&str>) -> Self {
417        Self::from_vec(value, DType::Utf8(Nullability::NonNullable))
418    }
419
420    /// Create from a vector of optional string slices.
421    pub fn from_nullable_strs(value: Vec<Option<&str>>) -> Self {
422        Self::from_iter(value, DType::Utf8(Nullability::Nullable))
423    }
424
425    /// Create from a vector of byte slices.
426    pub fn from_bytes(value: Vec<&[u8]>) -> Self {
427        Self::from_vec(value, DType::Binary(Nullability::NonNullable))
428    }
429
430    /// Create from a vector of optional byte slices.
431    pub fn from_nullable_bytes(value: Vec<Option<&[u8]>>) -> Self {
432        Self::from_iter(value, DType::Binary(Nullability::Nullable))
433    }
434
435    pub fn into_data_parts(self) -> VarBinDataParts {
436        let dtype = self.dtype().clone();
437        let validity = self.varbin_validity();
438        let offsets = self.offsets().clone();
439        let data = self.into_data();
440        VarBinDataParts {
441            dtype,
442            bytes: data.bytes,
443            offsets,
444            validity,
445        }
446    }
447}
448
449impl Array<VarBin> {
450    /// Creates a new `VarBinArray`.
451    pub fn new(offsets: ArrayRef, bytes: ByteBuffer, dtype: DType, validity: Validity) -> Self {
452        let len = offsets.len().saturating_sub(1);
453        let slots = VarBinData::make_slots(offsets, &validity, len);
454        let data = VarBinData::build(
455            slots[VarBinSlots::OFFSETS]
456                .as_ref()
457                .vortex_expect("VarBinArray offsets slot")
458                .clone(),
459            bytes,
460            dtype.clone(),
461            validity,
462        );
463        unsafe {
464            Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data).with_slots(slots))
465        }
466    }
467
468    /// Creates a new `VarBinArray` without validation.
469    ///
470    /// # Safety
471    ///
472    /// See [`VarBinData::new_unchecked`].
473    pub unsafe fn new_unchecked(
474        offsets: ArrayRef,
475        bytes: ByteBuffer,
476        dtype: DType,
477        validity: Validity,
478    ) -> Self {
479        let len = offsets.len().saturating_sub(1);
480        let slots = VarBinData::make_slots(offsets, &validity, len);
481        let data = unsafe { VarBinData::new_unchecked(bytes) };
482        unsafe {
483            Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data).with_slots(slots))
484        }
485    }
486
487    /// Creates a new `VarBinArray` without validation from a [`BufferHandle`].
488    ///
489    /// # Safety
490    ///
491    /// See [`VarBinData::new_unchecked_from_handle`].
492    pub unsafe fn new_unchecked_from_handle(
493        offsets: ArrayRef,
494        bytes: BufferHandle,
495        dtype: DType,
496        validity: Validity,
497    ) -> Self {
498        let len = offsets.len().saturating_sub(1);
499        let slots = VarBinData::make_slots(offsets, &validity, len);
500        let data = unsafe { VarBinData::new_unchecked_from_handle(bytes) };
501        unsafe {
502            Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data).with_slots(slots))
503        }
504    }
505
506    /// Constructs a new `VarBinArray`.
507    pub fn try_new(
508        offsets: ArrayRef,
509        bytes: ByteBuffer,
510        dtype: DType,
511        validity: Validity,
512    ) -> VortexResult<Self> {
513        let len = offsets.len() - 1;
514        let bytes = BufferHandle::new_host(bytes);
515        VarBinData::validate(&offsets, &bytes, &dtype, &validity)?;
516        let slots = VarBinData::make_slots(offsets, &validity, len);
517        // SAFETY: validate ensures all invariants are met.
518        let data = unsafe { VarBinData::new_unchecked_from_handle(bytes) };
519        Ok(unsafe {
520            Array::from_parts_unchecked(ArrayParts::new(VarBin, dtype, len, data).with_slots(slots))
521        })
522    }
523}
524
525impl From<Vec<&[u8]>> for Array<VarBin> {
526    fn from(value: Vec<&[u8]>) -> Self {
527        Self::from_vec(value, DType::Binary(Nullability::NonNullable))
528    }
529}
530
531impl From<Vec<Vec<u8>>> for Array<VarBin> {
532    fn from(value: Vec<Vec<u8>>) -> Self {
533        Self::from_vec(value, DType::Binary(Nullability::NonNullable))
534    }
535}
536
537impl From<Vec<String>> for Array<VarBin> {
538    fn from(value: Vec<String>) -> Self {
539        Self::from_vec(value, DType::Utf8(Nullability::NonNullable))
540    }
541}
542
543impl From<Vec<&str>> for Array<VarBin> {
544    fn from(value: Vec<&str>) -> Self {
545        Self::from_vec(value, DType::Utf8(Nullability::NonNullable))
546    }
547}
548
549impl From<Vec<Option<&[u8]>>> for Array<VarBin> {
550    fn from(value: Vec<Option<&[u8]>>) -> Self {
551        Self::from_iter(value, DType::Binary(Nullability::Nullable))
552    }
553}
554
555impl From<Vec<Option<Vec<u8>>>> for Array<VarBin> {
556    fn from(value: Vec<Option<Vec<u8>>>) -> Self {
557        Self::from_iter(value, DType::Binary(Nullability::Nullable))
558    }
559}
560
561impl From<Vec<Option<String>>> for Array<VarBin> {
562    fn from(value: Vec<Option<String>>) -> Self {
563        Self::from_iter(value, DType::Utf8(Nullability::Nullable))
564    }
565}
566
567impl From<Vec<Option<&str>>> for Array<VarBin> {
568    fn from(value: Vec<Option<&str>>) -> Self {
569        Self::from_iter(value, DType::Utf8(Nullability::Nullable))
570    }
571}
572
573impl<'a> FromIterator<Option<&'a [u8]>> for Array<VarBin> {
574    fn from_iter<T: IntoIterator<Item = Option<&'a [u8]>>>(iter: T) -> Self {
575        Self::from_iter(iter, DType::Binary(Nullability::Nullable))
576    }
577}
578
579impl FromIterator<Option<Vec<u8>>> for Array<VarBin> {
580    fn from_iter<T: IntoIterator<Item = Option<Vec<u8>>>>(iter: T) -> Self {
581        Self::from_iter(iter, DType::Binary(Nullability::Nullable))
582    }
583}
584
585impl FromIterator<Option<String>> for Array<VarBin> {
586    fn from_iter<T: IntoIterator<Item = Option<String>>>(iter: T) -> Self {
587        Self::from_iter(iter, DType::Utf8(Nullability::Nullable))
588    }
589}
590
591impl<'a> FromIterator<Option<&'a str>> for Array<VarBin> {
592    fn from_iter<T: IntoIterator<Item = Option<&'a str>>>(iter: T) -> Self {
593        Self::from_iter(iter, DType::Utf8(Nullability::Nullable))
594    }
595}
vortex_array/arrays/varbin/array.rs

vortex_array/arrays/varbin/
array.rs
