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