duckdb/core/
data_chunk.rs

1use super::{
2    logical_type::LogicalTypeHandle,
3    vector::{ArrayVector, FlatVector, ListVector, StructVector},
4};
5use crate::ffi::{
6    duckdb_create_data_chunk, duckdb_data_chunk, duckdb_data_chunk_get_column_count, duckdb_data_chunk_get_size,
7    duckdb_data_chunk_get_vector, duckdb_data_chunk_set_size, duckdb_destroy_data_chunk,
8};
9
10/// Handle to the DataChunk in DuckDB.
11pub struct DataChunkHandle {
12    /// Pointer to the DataChunk in duckdb C API.
13    ptr: duckdb_data_chunk,
14
15    /// Whether this `DataChunkHandle` owns `ptr`.
16    owned: bool,
17}
18
19impl Drop for DataChunkHandle {
20    fn drop(&mut self) {
21        if self.owned && !self.ptr.is_null() {
22            unsafe { duckdb_destroy_data_chunk(&mut self.ptr) }
23            self.ptr = std::ptr::null_mut();
24        }
25    }
26}
27
28impl DataChunkHandle {
29    /// Wrap a `duckdb_data_chunk` pointer owned elsewhere (e.g. by a DuckDB
30    /// callback frame) without taking ownership.
31    ///
32    /// # Safety
33    ///
34    /// `ptr` must be a valid `duckdb_data_chunk` that stays allocated and
35    /// unmutated by other code for the full lifetime of the returned handle
36    /// and any vectors derived from it.
37    //
38    // Known aliasing hole (#673 follow-up): `flat_vector` / `list_vector` /
39    // `array_vector` / `struct_vector` take `&self`, so safe code can get two
40    // writable wrappers over the same column and produce aliased `&mut [T]`
41    // slices.
42    #[cfg(feature = "vtab")]
43    pub(crate) unsafe fn new_unowned(ptr: duckdb_data_chunk) -> Self {
44        Self { ptr, owned: false }
45    }
46
47    /// Create a new [DataChunkHandle] with the given [LogicalTypeHandle]s.
48    pub fn new(logical_types: &[LogicalTypeHandle]) -> Self {
49        let num_columns = logical_types.len();
50        let mut c_types = Vec::with_capacity(num_columns);
51        c_types.extend(logical_types.iter().map(|t| t.ptr));
52        let ptr = unsafe { duckdb_create_data_chunk(c_types.as_mut_ptr(), num_columns as u64) };
53        Self { ptr, owned: true }
54    }
55
56    /// Get the vector at the specific column index: `idx`.
57    pub fn flat_vector(&self, idx: usize) -> FlatVector<'_> {
58        unsafe { FlatVector::from_raw(duckdb_data_chunk_get_vector(self.ptr, idx as u64)) }
59    }
60
61    /// Get a list vector from the column index.
62    pub fn list_vector(&self, idx: usize) -> ListVector<'_> {
63        unsafe { ListVector::from_raw(duckdb_data_chunk_get_vector(self.ptr, idx as u64)) }
64    }
65
66    /// Get a array vector from the column index.
67    pub fn array_vector(&self, idx: usize) -> ArrayVector<'_> {
68        unsafe { ArrayVector::from_raw(duckdb_data_chunk_get_vector(self.ptr, idx as u64)) }
69    }
70
71    /// Get struct vector at the column index: `idx`.
72    pub fn struct_vector(&self, idx: usize) -> StructVector<'_> {
73        unsafe { StructVector::from_raw(duckdb_data_chunk_get_vector(self.ptr, idx as u64)) }
74    }
75
76    /// Set the size of the data chunk
77    pub fn set_len(&self, new_len: usize) {
78        unsafe { duckdb_data_chunk_set_size(self.ptr, new_len as u64) };
79    }
80
81    /// Get the length / the number of rows in this [DataChunkHandle].
82    pub fn len(&self) -> usize {
83        unsafe { duckdb_data_chunk_get_size(self.ptr) as usize }
84    }
85
86    /// Check whether this [DataChunkHandle] is empty.
87    pub fn is_empty(&self) -> bool {
88        self.len() == 0
89    }
90
91    /// Get the number of columns in this [DataChunkHandle].
92    pub fn num_columns(&self) -> usize {
93        unsafe { duckdb_data_chunk_get_column_count(self.ptr) as usize }
94    }
95
96    /// Get the ptr of duckdb_data_chunk in this [DataChunkHandle].
97    pub fn get_ptr(&self) -> duckdb_data_chunk {
98        self.ptr
99    }
100}
101
102#[cfg(test)]
103mod test {
104    use super::{super::logical_type::LogicalTypeId, *};
105
106    #[test]
107    fn test_data_chunk_construction() {
108        let dc = DataChunkHandle::new(&[LogicalTypeHandle::from(LogicalTypeId::Integer)]);
109
110        assert_eq!(dc.num_columns(), 1);
111
112        drop(dc);
113    }
114
115    #[test]
116    fn test_vector() {
117        let datachunk = DataChunkHandle::new(&[LogicalTypeHandle::from(LogicalTypeId::Bigint)]);
118        let mut vector = datachunk.flat_vector(0);
119        let data = unsafe { vector.as_mut_slice::<i64>() };
120
121        data[0] = 42;
122    }
123
124    #[test]
125    fn test_logi() {
126        let key = LogicalTypeHandle::from(LogicalTypeId::Varchar);
127
128        let value = LogicalTypeHandle::from(LogicalTypeId::UTinyint);
129
130        let map = LogicalTypeHandle::map(&key, &value);
131
132        assert_eq!(map.id(), LogicalTypeId::Map);
133
134        // let union_ = LogicalType::new_union_type(HashMap::from([
135        //     ("number", LogicalType::new(LogicalTypeId::Bigint)),
136        //     ("string", LogicalType::new(LogicalTypeId::Varchar)),
137        // ]));
138        // assert_eq!(union_.type_id(), LogicalTypeId::Union);
139
140        // let struct_ = LogicalType::new_struct_type(HashMap::from([
141        //     ("number", LogicalType::new(LogicalTypeId::Bigint)),
142        //     ("string", LogicalType::new(LogicalTypeId::Varchar)),
143        // ]));
144        // assert_eq!(struct_.type_id(), LogicalTypeId::Struct);
145    }
146}
duckdb/core/data_chunk.rs

duckdb/core/
data_chunk.rs
