Skip to main content

duckdb/core/vector/
array.rs

1use std::marker::PhantomData;
2
3use libduckdb_sys::{duckdb_array_type_array_size, duckdb_array_vector_get_child};
4
5use crate::{
6    Result,
7    core::LogicalTypeHandle,
8    ffi::{
9        duckdb_validity_set_row_invalid, duckdb_vector, duckdb_vector_ensure_validity_writable,
10        duckdb_vector_get_column_type, duckdb_vector_get_validity, duckdb_vector_size,
11    },
12};
13
14use super::{FlatVector, try_vector_row_is_null, vector_row_is_null};
15
16/// A fixed-size list vector borrowed from a
17/// [`DataChunkHandle`][crate::core::DataChunkHandle].
18///
19/// Exposes a fixed-width list whose child storage is contiguous across all rows.
20pub struct ArrayVector<'a> {
21    ptr: duckdb_vector,
22    capacity: usize,
23    _phantom: PhantomData<&'a ()>,
24}
25
26impl<'a> ArrayVector<'a> {
27    /// Wrap a raw `duckdb_vector` pointer.
28    ///
29    /// # Safety
30    /// `ptr` must be a valid `duckdb_vector` that remains valid for all of `'a`.
31    pub(crate) unsafe fn from_raw(ptr: duckdb_vector) -> Self {
32        unsafe { Self::from_raw_with_capacity(ptr, duckdb_vector_size() as usize) }
33    }
34
35    /// Wrap a raw array vector pointer with a caller-supplied row capacity.
36    ///
37    /// # Safety
38    /// `ptr` must be a valid array `duckdb_vector` that remains valid for all
39    /// of `'a`, and `capacity` must not exceed the backing row allocation.
40    pub(crate) unsafe fn from_raw_with_capacity(ptr: duckdb_vector, capacity: usize) -> Self {
41        Self {
42            ptr,
43            capacity,
44            _phantom: PhantomData,
45        }
46    }
47
48    /// Get the logical type of this ArrayVector.
49    pub fn logical_type(&self) -> LogicalTypeHandle {
50        unsafe { LogicalTypeHandle::new(duckdb_vector_get_column_type(self.ptr)) }
51    }
52
53    /// Returns the size of the array type.
54    pub fn get_array_size(&self) -> u64 {
55        let ty = self.logical_type();
56        unsafe { duckdb_array_type_array_size(ty.ptr) as u64 }
57    }
58
59    /// Returns true if the row at the given index is null.
60    ///
61    /// # Panics
62    /// Panics if `row` is outside the vector capacity.
63    pub fn row_is_null(&self, row: u64) -> bool {
64        vector_row_is_null(self.ptr, row, self.capacity)
65    }
66
67    /// Returns whether the row is null, or an error if it is out of range.
68    pub fn try_row_is_null(&self, row: u64) -> Result<bool> {
69        try_vector_row_is_null(self.ptr, row, self.capacity)
70    }
71
72    /// Returns the child vector.
73    /// capacity should be a multiple of the array size.
74    ///
75    /// # Panics
76    /// Panics if `capacity` is not a multiple of the fixed array size or exceeds
77    /// the child allocation derived from the parent capacity.
78    // TODO: not ideal interface. Where should we keep count.
79    pub fn child(&self, capacity: usize) -> FlatVector<'a> {
80        let array_size = self.get_array_size() as usize;
81        assert_eq!(
82            capacity % array_size,
83            0,
84            "array child capacity must be a multiple of the fixed array size"
85        );
86        let child_capacity = self
87            .capacity
88            .checked_mul(array_size)
89            .expect("array child capacity overflows usize");
90        assert!(
91            capacity <= child_capacity,
92            "array child capacity {capacity} exceeds backing capacity {child_capacity}"
93        );
94        // SAFETY: fixed-array children contain `array_size` values per parent
95        // row, and the checks above bound `capacity` by that allocation.
96        unsafe { FlatVector::with_capacity(self.child_ptr(), capacity) }
97    }
98
99    pub(crate) fn child_ptr(&self) -> duckdb_vector {
100        unsafe { duckdb_array_vector_get_child(self.ptr) }
101    }
102
103    /// Set primitive data to the child node.
104    ///
105    /// # Panics
106    /// Panics if `data.len()` is not a multiple of the fixed array size.
107    ///
108    /// # Safety
109    /// The caller must ensure `T` matches the child vector's physical storage
110    /// and no other references to the same storage exist during the copy. The
111    /// caller must also ensure the child backing allocation has space for all
112    /// `data` elements; nested callers can arrange this by reserving capacity
113    /// on the parent list vector before taking the array child.
114    pub unsafe fn set_child<T: Copy>(&self, data: &[T]) {
115        unsafe { self.child(data.len()).copy(data) };
116    }
117
118    /// Set row as null
119    pub fn set_null(&mut self, row: usize) {
120        unsafe {
121            duckdb_vector_ensure_validity_writable(self.ptr);
122            let idx = duckdb_vector_get_validity(self.ptr);
123            duckdb_validity_set_row_invalid(idx, row as u64);
124        }
125    }
126}