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

duckdb/core/
data_chunk.rs
