better_duck_core/raw/
data_chunk.rs

1use std::{
2    ops::{Deref, DerefMut},
3    ptr,
4};
5
6use crate::ffi::{duckdb_data_chunk, duckdb_data_chunk_reset, duckdb_destroy_data_chunk};
7
8use super::result::DuckResult;
9use crate::{error::Result, ffi};
10
11/// An owned DuckDB `duckdb_data_chunk` — a batch of column vectors — that is
12/// destroyed on drop.
13///
14/// Obtain one from a query result with [`DataChunk::from_result`], then append it
15/// elsewhere with [`Appender::append_chunk`](crate::Appender::append_chunk).
16pub struct DataChunk(
17    pub(crate) duckdb_data_chunk,
18    pub(crate) u64, // current row index in chunk
19);
20
21impl DataChunk {
22    /// Takes ownership of a raw `duckdb_data_chunk` handle.
23    ///
24    /// # Errors
25    ///
26    /// Returns an error if `data_chunk` is null.
27    #[inline]
28    pub fn new(data_chunk: ffi::duckdb_data_chunk) -> Result<DataChunk> {
29        if data_chunk.is_null() {
30            return Err(crate::error::Error::DuckDBFailure(
31                ffi::Error::new(ffi::DuckDBError),
32                Some("data chunk is null".to_owned()),
33            ));
34        }
35        Ok(DataChunk(data_chunk, 0))
36    }
37    /// Fetches the next chunk of a query result, or `None` once the result is
38    /// exhausted.
39    #[inline]
40    pub fn from_result(result: &DuckResult) -> Option<Result<DataChunk>> {
41        // SAFETY: `result` is a valid duckdb_result; the returned chunk (if non-null)
42        // is owned by us and must be destroyed via `duckdb_destroy_data_chunk`.
43        let data_chunk = unsafe { ffi::duckdb_fetch_chunk(**result) };
44        if data_chunk.is_null() {
45            return None;
46        }
47        // SAFETY: `data_chunk` is non-null and freshly obtained from `duckdb_fetch_chunk`.
48        let res = DataChunk::new(data_chunk);
49        Some(res)
50    }
51
52    /// The current row cursor used by [`next_row`](DataChunk::next_row).
53    #[allow(unused)]
54    #[inline]
55    pub fn current_row(&self) -> u64 {
56        self.1
57    }
58    /// The number of rows in the chunk.
59    #[inline]
60    pub fn row_count(&self) -> u64 {
61        // SAFETY: `self.0` is a valid duckdb_data_chunk (enforced by the caller).
62        unsafe { ffi::duckdb_data_chunk_get_size(self.0) }
63    }
64
65    /// Resets the chunk to empty (row count 0), keeping its allocated capacity so it
66    /// can be refilled and re-appended.
67    #[inline]
68    pub fn reset(&mut self) {
69        // SAFETY: `self.0` is a valid, non-null duckdb_data_chunk owned by `self`.
70        unsafe { duckdb_data_chunk_reset(self.0) };
71        self.1 = 0;
72    }
73
74    /// Advances the row cursor, returning the next row index, or `None` at the end
75    /// (destroying the underlying chunk).
76    #[inline]
77    pub fn next_row(&mut self) -> Option<u64> {
78        // SAFETY: `self.0` is a valid duckdb_data_chunk (enforced by the caller).
79        if self.row_count() < 1 {
80            return None;
81        }
82        // SAFETY: same as above.
83        if self.1 >= self.row_count() {
84            // Reset the row index and fetch the next chunk
85            self.1 = 0;
86            // SAFETY: `self.0` is a valid, non-null duckdb_data_chunk; after destroy
87            // we null it so this path is never re-entered.
88            unsafe { duckdb_destroy_data_chunk(&mut (self.0)) };
89            self.0 = ptr::null_mut();
90            return None;
91        }
92        self.1 += 1;
93        Some(self.1 - 1)
94    }
95}
96
97impl Deref for DataChunk {
98    type Target = duckdb_data_chunk;
99
100    #[inline]
101    fn deref(&self) -> &Self::Target {
102        &self.0
103    }
104}
105
106impl DerefMut for DataChunk {
107    #[inline]
108    fn deref_mut(&mut self) -> &mut Self::Target {
109        &mut self.0
110    }
111}
112
113impl Drop for DataChunk {
114    fn drop(&mut self) {
115        if !self.0.is_null() {
116            // SAFETY: `self.0` is a valid non-null `duckdb_data_chunk`. The null guard
117            // ensures this path runs at most once.
118            unsafe { duckdb_destroy_data_chunk(&mut (self.0)) };
119        }
120    }
121}
122
123#[cfg(test)]
124mod tests {
125    use super::*;
126    use crate::types::LogicalType;
127
128    #[test]
129    fn reset_clears_the_row_count() {
130        // Build a one-column INTEGER chunk and give it three rows.
131        let int_ty = LogicalType::of::<i32>().unwrap();
132        let mut raw_types = [int_ty.as_raw()];
133        // SAFETY: `raw_types` holds one valid logical type handle for the call; DuckDB
134        // copies it. The returned chunk is wrapped in RAII (destroyed once on drop).
135        let raw = unsafe { ffi::duckdb_create_data_chunk(raw_types.as_mut_ptr(), 1) };
136        drop(int_ty);
137        let mut chunk = DataChunk::new(raw).unwrap();
138        // SAFETY: `chunk` is valid; sizing to 3 rows is within the default capacity.
139        unsafe { ffi::duckdb_data_chunk_set_size(*chunk, 3) };
140        assert_eq!(chunk.row_count(), 3);
141
142        chunk.reset();
143        assert_eq!(chunk.row_count(), 0, "reset must clear the chunk to empty");
144    }
145}
better_duck_core/raw/data_chunk.rs

better_duck_core/raw/
data_chunk.rs
