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
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
// SPDX-License-Identifier: MIT
// Copyright 2026 Tom F. <https://github.com/tomtom215/>
// My way of giving something small back to the open source community
// and encouraging more Rust development!
//! Ergonomic wrapper around `DuckDB` data chunks.
//!
//! [`DataChunk`] provides safe, convenient access to the vectors and metadata
//! of a `duckdb_data_chunk`, eliminating the raw FFI calls that extension
//! authors currently need to write in every scan callback.
//!
//! # Example
//!
//! ```rust,no_run
//! use quack_rs::data_chunk::DataChunk;
//! use quack_rs::vector::{VectorWriter, VectorReader};
//! use libduckdb_sys::{duckdb_function_info, duckdb_data_chunk};
//!
//! unsafe extern "C" fn my_scan(info: duckdb_function_info, output: duckdb_data_chunk) {
//! let chunk = unsafe { DataChunk::from_raw(output) };
//! let mut writer = unsafe { chunk.writer(0) };
//! unsafe { writer.write_i64(0, 42) };
//! unsafe { chunk.set_size(1) };
//! }
//! ```
use libduckdb_sys::{
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_vector, idx_t,
};
use crate::chunk_writer::ChunkWriter;
use crate::vector::complex::StructVector;
use crate::vector::{StructReader, StructWriter, VectorReader, VectorWriter};
/// A non-owning wrapper around a `duckdb_data_chunk`.
///
/// This wrapper does **not** destroy the chunk on drop — `DuckDB` owns the
/// chunk and manages its lifetime. `DataChunk` simply provides ergonomic
/// methods for accessing vectors and metadata within callback functions.
pub struct DataChunk {
raw: duckdb_data_chunk,
}
impl DataChunk {
/// Wraps a raw `duckdb_data_chunk` handle.
///
/// # Safety
///
/// `raw` must be a valid `duckdb_data_chunk` obtained from a `DuckDB`
/// callback (e.g., a scan callback's `output` parameter or an aggregate
/// `update` callback's `input` chunk). The chunk must remain valid for
/// the lifetime of this wrapper.
#[inline]
#[must_use]
pub const unsafe fn from_raw(raw: duckdb_data_chunk) -> Self {
Self { raw }
}
/// Returns the number of rows in this data chunk.
#[inline]
#[must_use]
pub fn size(&self) -> usize {
// SAFETY: self.raw is valid per constructor contract.
usize::try_from(unsafe { duckdb_data_chunk_get_size(self.raw) }).unwrap_or(0)
}
/// Sets the number of rows in this data chunk.
///
/// Call this in scan callbacks after writing output rows. Set to `0` to
/// signal end of stream.
///
/// # Safety
///
/// `size` must not exceed the chunk's capacity (typically 2048).
#[inline]
pub unsafe fn set_size(&self, size: usize) {
// SAFETY: self.raw is valid per constructor contract.
unsafe { duckdb_data_chunk_set_size(self.raw, size as idx_t) };
}
/// Returns the number of columns in this data chunk.
#[inline]
#[must_use]
pub fn column_count(&self) -> usize {
// SAFETY: self.raw is valid per constructor contract.
usize::try_from(unsafe { duckdb_data_chunk_get_column_count(self.raw) }).unwrap_or(0)
}
/// Returns the raw `duckdb_vector` handle for the given column index.
///
/// # Safety
///
/// `col_idx` must be less than [`column_count`][DataChunk::column_count].
#[inline]
#[must_use]
pub unsafe fn vector(&self, col_idx: usize) -> duckdb_vector {
// SAFETY: self.raw is valid and col_idx is in bounds per caller's contract.
unsafe { duckdb_data_chunk_get_vector(self.raw, col_idx as idx_t) }
}
/// Creates a [`VectorWriter`] for the given column index.
///
/// # Safety
///
/// - `col_idx` must be less than [`column_count`][DataChunk::column_count].
/// - The chunk must be a writable output chunk (not a read-only input chunk).
pub unsafe fn writer(&self, col_idx: usize) -> VectorWriter {
let vec = unsafe { self.vector(col_idx) };
// SAFETY: vec is a valid writable vector from the output chunk.
unsafe { VectorWriter::from_vector(vec) }
}
/// Creates a [`VectorReader`] for the given column index.
///
/// The reader's row count is set to this chunk's current [`size`][DataChunk::size].
///
/// # Safety
///
/// `col_idx` must be less than [`column_count`][DataChunk::column_count].
pub unsafe fn reader(&self, col_idx: usize) -> VectorReader {
// SAFETY: self.raw is valid; col_idx is in bounds per caller's contract.
unsafe { VectorReader::new(self.raw, col_idx) }
}
/// Creates a [`StructReader`] for a STRUCT column at the given index.
///
/// This is a convenience method that combines [`vector`][Self::vector] with
/// [`StructReader::new`].
///
/// # Safety
///
/// - `col_idx` must be less than [`column_count`][Self::column_count].
/// - The column at `col_idx` must have a STRUCT type with `field_count` fields.
pub unsafe fn struct_reader(&self, col_idx: usize, field_count: usize) -> StructReader {
let vec = unsafe { self.vector(col_idx) };
// SAFETY: vec is a valid STRUCT vector per caller's contract.
unsafe { StructReader::new(vec, field_count, self.size()) }
}
/// Creates a [`VectorReader`] for a field of a STRUCT column.
///
/// Convenience for accessing a specific field in a STRUCT input column.
///
/// # Safety
///
/// - `col_idx` must be less than [`column_count`][Self::column_count].
/// - The column at `col_idx` must have a STRUCT type.
/// - `field_idx` must be a valid field index within the STRUCT.
pub unsafe fn struct_field_reader(&self, col_idx: usize, field_idx: usize) -> VectorReader {
let vec = unsafe { self.vector(col_idx) };
// SAFETY: vec is a valid STRUCT vector per caller's contract.
unsafe { StructVector::field_reader(vec, field_idx, self.size()) }
}
/// Creates a [`StructWriter`] for a STRUCT column at the given index.
///
/// This is a convenience method that combines [`vector`][Self::vector] with
/// [`StructWriter::new`].
///
/// # Safety
///
/// - `col_idx` must be less than [`column_count`][Self::column_count].
/// - The column at `col_idx` must have a STRUCT type with `field_count` fields.
/// - The chunk must be a writable output chunk.
pub unsafe fn struct_writer(&self, col_idx: usize, field_count: usize) -> StructWriter {
let vec = unsafe { self.vector(col_idx) };
// SAFETY: vec is a valid STRUCT vector per caller's contract.
unsafe { StructWriter::new(vec, field_count) }
}
/// Creates a [`ChunkWriter`] for this output data chunk.
///
/// The [`ChunkWriter`] tracks rows via [`next_row()`][ChunkWriter::next_row]
/// and automatically calls `set_size` on drop.
///
/// # Safety
///
/// This chunk must be a valid, writable output chunk from a table function
/// scan callback.
pub const unsafe fn into_chunk_writer(self) -> ChunkWriter {
// SAFETY: self.raw is valid per constructor's contract.
unsafe { ChunkWriter::new(self.raw) }
}
/// Returns the raw `duckdb_data_chunk` handle.
#[inline]
#[must_use]
pub const fn as_raw(&self) -> duckdb_data_chunk {
self.raw
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn size_of_data_chunk() {
assert_eq!(
std::mem::size_of::<DataChunk>(),
std::mem::size_of::<usize>()
);
}
}