better_duck_core/udf/data_chunk.rs
1//! An owned-or-borrowed DuckDB data chunk.
2
3use crate::{
4 error::{Error, Result},
5 ffi::{
6 duckdb_create_data_chunk, duckdb_data_chunk, duckdb_data_chunk_get_column_count,
7 duckdb_data_chunk_get_size, duckdb_data_chunk_get_vector, duckdb_data_chunk_set_size,
8 duckdb_destroy_data_chunk, duckdb_vector_size,
9 },
10};
11
12use crate::types::LogicalType;
13
14use super::vector::{VectorMut, VectorRef};
15
16/// A DuckDB data chunk: a batch of rows laid out column-wise, one [`VectorRef`]/
17/// [`VectorMut`] per column.
18///
19/// Either owned by this handle (allocated via [`DataChunkHandle::new`] and
20/// destroyed on drop) or borrowed from a DuckDB callback frame (via a crate-
21/// internal constructor), in which case DuckDB owns the chunk itself and it
22/// must not be destroyed out from under it.
23///
24/// This is a separate type from the crate's internal, read-only
25/// `raw::data_chunk::DataChunk` — that type's `Drop` unconditionally destroys
26/// the chunk and its `next_row` destroys-and-nulls on exhaustion, both of which
27/// would double-free a chunk DuckDB owns.
28pub struct DataChunkHandle {
29 ptr: duckdb_data_chunk,
30 owned: bool,
31}
32
33impl DataChunkHandle {
34 /// Allocates a new data chunk with the given column types.
35 ///
36 /// # Errors
37 ///
38 /// Returns an error if DuckDB fails to allocate the chunk.
39 pub fn new(types: &[LogicalType]) -> Result<Self> {
40 let mut raw_types: Vec<_> = types.iter().map(LogicalType::as_raw).collect();
41 // SAFETY: `raw_types` is a valid, non-dangling array of `raw_types.len()`
42 // logical type handles, each owned (and kept alive) by `types` for the
43 // duration of this call; `duckdb_create_data_chunk` copies them.
44 let ptr =
45 unsafe { duckdb_create_data_chunk(raw_types.as_mut_ptr(), raw_types.len() as u64) };
46 if ptr.is_null() {
47 return Err(Error::ConversionError(
48 crate::error::DuckDBConversionError::ConversionError(
49 "duckdb_create_data_chunk returned null".to_owned(),
50 ),
51 ));
52 }
53 Ok(Self { ptr, owned: true })
54 }
55
56 /// Wraps a `duckdb_data_chunk` owned by a DuckDB callback frame, without
57 /// taking ownership of it.
58 ///
59 /// # Safety
60 ///
61 /// `ptr` must be a valid, non-null `duckdb_data_chunk` that stays allocated
62 /// and is not mutated by other code for the full lifetime of the returned
63 /// handle and any vectors derived from it.
64 pub(crate) unsafe fn borrowed(ptr: duckdb_data_chunk) -> Self {
65 Self { ptr, owned: false }
66 }
67
68 /// Returns a read-only view of column `idx`.
69 ///
70 /// # Errors
71 ///
72 /// Returns an error if `idx` is out of range.
73 pub fn vector(
74 &self,
75 idx: usize,
76 ) -> Result<VectorRef<'_>> {
77 self.check_col(idx)?;
78 // SAFETY: `self.ptr` is valid; `idx` was checked above. The returned
79 // vector pointer is valid for as long as `self.ptr` is (per the C API
80 // docs) and does not need separate destruction.
81 let vec_ptr = unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) };
82 // SAFETY: `vec_ptr` is valid for the lifetime of `&self`.
83 Ok(unsafe { VectorRef::new(vec_ptr) })
84 }
85
86 /// Returns an exclusive, writable view of column `idx`.
87 ///
88 /// # Errors
89 ///
90 /// Returns an error if `idx` is out of range.
91 pub fn vector_mut(
92 &mut self,
93 idx: usize,
94 ) -> Result<VectorMut<'_>> {
95 self.check_col(idx)?;
96 // SAFETY: `self.ptr` is valid; `idx` was checked above.
97 let vec_ptr = unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) };
98 // SAFETY: `vec_ptr` is valid for the lifetime of `&mut self`, and the
99 // `&mut self` borrow ensures no other view of this chunk's vectors is
100 // live at the same time.
101 Ok(unsafe { VectorMut::new(vec_ptr) })
102 }
103
104 /// Returns a writable view of every column, in column order.
105 ///
106 /// The views are disjoint by construction (distinct column indices), so
107 /// they may be held simultaneously despite each being individually
108 /// exclusive.
109 ///
110 /// # Errors
111 ///
112 /// Returns an error if a column vector cannot be obtained.
113 pub fn vectors_mut(&mut self) -> Result<Vec<VectorMut<'_>>> {
114 let n = self.num_columns();
115 (0..n)
116 .map(|idx| {
117 // SAFETY: `self.ptr` is valid; `idx < n` is in range.
118 let vec_ptr = unsafe { duckdb_data_chunk_get_vector(self.ptr, idx as u64) };
119 // SAFETY: each `vec_ptr` refers to a distinct column, so these
120 // views never alias each other even though several are held at
121 // once; all are bounded by the `&mut self` borrow.
122 Ok(unsafe { VectorMut::new(vec_ptr) })
123 })
124 .collect()
125 }
126
127 /// The number of rows currently in this chunk.
128 pub fn len(&self) -> usize {
129 // SAFETY: `self.ptr` is a valid data chunk.
130 unsafe { duckdb_data_chunk_get_size(self.ptr) as usize }
131 }
132
133 /// Returns `true` if this chunk has no rows.
134 pub fn is_empty(&self) -> bool {
135 self.len() == 0
136 }
137
138 /// The number of columns in this chunk.
139 pub fn num_columns(&self) -> usize {
140 // SAFETY: `self.ptr` is a valid data chunk.
141 unsafe { duckdb_data_chunk_get_column_count(self.ptr) as usize }
142 }
143
144 /// Sets the number of valid rows in this chunk.
145 ///
146 /// # Errors
147 ///
148 /// Returns an error if `n` exceeds [`DataChunkHandle::capacity`].
149 pub fn set_len(
150 &mut self,
151 n: usize,
152 ) -> Result<()> {
153 if n > self.capacity() {
154 return Err(Error::ConversionError(
155 crate::error::DuckDBConversionError::ConversionError(format!(
156 "row count {n} exceeds chunk capacity {}",
157 self.capacity()
158 )),
159 ));
160 }
161 // SAFETY: `self.ptr` is a valid data chunk; `n` was just bounds-checked
162 // against its capacity.
163 unsafe { duckdb_data_chunk_set_size(self.ptr, n as u64) };
164 Ok(())
165 }
166
167 /// The maximum number of rows this chunk can hold — DuckDB's build-time
168 /// vector size, not a fixed constant. Never hardcode a row-count limit;
169 /// always read it from here.
170 pub fn capacity(&self) -> usize {
171 // SAFETY: always safe to call; takes no arguments and reads no memory
172 // beyond DuckDB's own build-time configuration.
173 unsafe { duckdb_vector_size() as usize }
174 }
175
176 fn check_col(
177 &self,
178 idx: usize,
179 ) -> Result<()> {
180 if idx >= self.num_columns() {
181 return Err(Error::InvalidColumnIndex(idx));
182 }
183 Ok(())
184 }
185}
186
187impl Drop for DataChunkHandle {
188 fn drop(&mut self) {
189 if self.owned && !self.ptr.is_null() {
190 // SAFETY: `self.owned` guarantees this handle allocated `self.ptr`
191 // via `duckdb_create_data_chunk` and no one else destroys it;
192 // `self.ptr` is non-null (checked above) and destroyed exactly once.
193 unsafe { duckdb_destroy_data_chunk(&mut self.ptr) };
194 }
195 }
196}
197
198#[cfg(test)]
199mod tests {
200 use super::*;
201
202 #[test]
203 fn owned_chunk_reports_column_count_and_capacity() {
204 let types = [LogicalType::of::<i32>().unwrap(), LogicalType::of::<String>().unwrap()];
205 let chunk = DataChunkHandle::new(&types).unwrap();
206 assert_eq!(chunk.num_columns(), 2);
207 assert!(chunk.capacity() > 0);
208 assert_eq!(chunk.len(), 0);
209 }
210
211 #[test]
212 fn set_len_within_capacity_succeeds() {
213 let types = [LogicalType::of::<i32>().unwrap()];
214 let mut chunk = DataChunkHandle::new(&types).unwrap();
215 chunk.set_len(3).unwrap();
216 assert_eq!(chunk.len(), 3);
217 }
218
219 #[test]
220 fn set_len_beyond_capacity_errors() {
221 let types = [LogicalType::of::<i32>().unwrap()];
222 let mut chunk = DataChunkHandle::new(&types).unwrap();
223 let too_many = chunk.capacity() + 1;
224 assert!(chunk.set_len(too_many).is_err());
225 }
226
227 #[test]
228 fn vector_out_of_range_errors() {
229 let types = [LogicalType::of::<i32>().unwrap()];
230 let chunk = DataChunkHandle::new(&types).unwrap();
231 assert!(chunk.vector(1).is_err());
232 }
233
234 #[test]
235 fn vectors_mut_covers_every_column() {
236 let types = [LogicalType::of::<i32>().unwrap(), LogicalType::of::<i64>().unwrap()];
237 let mut chunk = DataChunkHandle::new(&types).unwrap();
238 let views = chunk.vectors_mut().unwrap();
239 assert_eq!(views.len(), 2);
240 }
241}