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] own the [DataChunk::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    #[allow(dead_code)]
30    pub(crate) unsafe fn new_unowned(ptr: duckdb_data_chunk) -> Self {
31        Self { ptr, owned: false }
32    }
33
34    /// Create a new [DataChunkHandle] with the given [LogicalTypeHandle]s.
35    pub fn new(logical_types: &[LogicalTypeHandle]) -> Self {
36        let num_columns = logical_types.len();
37        let mut c_types = Vec::with_capacity(num_columns);
38        c_types.extend(logical_types.iter().map(|t| t.ptr));
39        let ptr = unsafe { duckdb_create_data_chunk(c_types.as_mut_ptr(), num_columns as u64) };
40        Self { ptr, owned: true }
41    }
42
43    /// Get the vector at the specific column index: `idx`.
44    pub fn flat_vector(&self, idx: usize) -> FlatVector {
45        FlatVector::from(unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) })
46    }
47
48    /// Get a list vector from the column index.
49    pub fn list_vector(&self, idx: usize) -> ListVector {
50        ListVector::from(unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) })
51    }
52
53    /// Get a array vector from the column index.
54    pub fn array_vector(&self, idx: usize) -> ArrayVector {
55        ArrayVector::from(unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) })
56    }
57
58    /// Get struct vector at the column index: `idx`.
59    pub fn struct_vector(&self, idx: usize) -> StructVector {
60        StructVector::from(unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) })
61    }
62
63    /// Set the size of the data chunk
64    pub fn set_len(&self, new_len: usize) {
65        unsafe { duckdb_data_chunk_set_size(self.ptr, new_len as u64) };
66    }
67
68    /// Get the length / the number of rows in this [DataChunkHandle].
69    pub fn len(&self) -> usize {
70        unsafe { duckdb_data_chunk_get_size(self.ptr) as usize }
71    }
72
73    /// Check whether this [DataChunkHandle] is empty.
74    pub fn is_empty(&self) -> bool {
75        self.len() == 0
76    }
77
78    /// Get the number of columns in this [DataChunkHandle].
79    pub fn num_columns(&self) -> usize {
80        unsafe { duckdb_data_chunk_get_column_count(self.ptr) as usize }
81    }
82
83    /// Get the ptr of duckdb_data_chunk in this [DataChunkHandle].
84    pub fn get_ptr(&self) -> duckdb_data_chunk {
85        self.ptr
86    }
87}
88
89#[cfg(test)]
90mod test {
91    use super::{super::logical_type::LogicalTypeId, *};
92
93    #[test]
94    fn test_data_chunk_construction() {
95        let dc = DataChunkHandle::new(&[LogicalTypeHandle::from(LogicalTypeId::Integer)]);
96
97        assert_eq!(dc.num_columns(), 1);
98
99        drop(dc);
100    }
101
102    #[test]
103    fn test_vector() {
104        let datachunk = DataChunkHandle::new(&[LogicalTypeHandle::from(LogicalTypeId::Bigint)]);
105        let mut vector = datachunk.flat_vector(0);
106        let data = vector.as_mut_slice::<i64>();
107
108        data[0] = 42;
109    }
110
111    #[test]
112    fn test_logi() {
113        let key = LogicalTypeHandle::from(LogicalTypeId::Varchar);
114
115        let value = LogicalTypeHandle::from(LogicalTypeId::UTinyint);
116
117        let map = LogicalTypeHandle::map(&key, &value);
118
119        assert_eq!(map.id(), LogicalTypeId::Map);
120
121        // let union_ = LogicalType::new_union_type(HashMap::from([
122        //     ("number", LogicalType::new(LogicalTypeId::Bigint)),
123        //     ("string", LogicalType::new(LogicalTypeId::Varchar)),
124        // ]));
125        // assert_eq!(union_.type_id(), LogicalTypeId::Union);
126
127        // let struct_ = LogicalType::new_struct_type(HashMap::from([
128        //     ("number", LogicalType::new(LogicalTypeId::Bigint)),
129        //     ("string", LogicalType::new(LogicalTypeId::Varchar)),
130        // ]));
131        // assert_eq!(struct_.type_id(), LogicalTypeId::Struct);
132    }
133}
duckdb/core/data_chunk.rs

duckdb/core/
data_chunk.rs
