use libduckdb_sys::{
duckdb_data_chunk, duckdb_data_chunk_get_size, duckdb_data_chunk_get_vector,
duckdb_validity_row_is_valid, duckdb_vector, duckdb_vector_get_data,
duckdb_vector_get_validity, idx_t,
};
pub struct VectorReader {
data: *const u8,
validity: *mut u64,
row_count: usize,
}
impl VectorReader {
pub unsafe fn new(chunk: duckdb_data_chunk, col_idx: usize) -> Self {
let row_count = usize::try_from(unsafe { duckdb_data_chunk_get_size(chunk) }).unwrap_or(0);
let vector = unsafe { duckdb_data_chunk_get_vector(chunk, col_idx as idx_t) };
let data = unsafe { duckdb_vector_get_data(vector) }.cast::<u8>();
let validity = unsafe { duckdb_vector_get_validity(vector) };
Self {
data,
validity,
row_count,
}
}
pub unsafe fn from_vector(vector: duckdb_vector, row_count: usize) -> Self {
let data = unsafe { duckdb_vector_get_data(vector) }.cast::<u8>();
let validity = unsafe { duckdb_vector_get_validity(vector) };
Self {
data,
validity,
row_count,
}
}
#[mutants::skip]
#[must_use]
#[inline]
pub const fn row_count(&self) -> usize {
self.row_count
}
#[inline]
pub unsafe fn is_valid(&self, idx: usize) -> bool {
if self.validity.is_null() {
return true;
}
unsafe { duckdb_validity_row_is_valid(self.validity, idx as idx_t) }
}
#[inline]
pub const unsafe fn read_i8(&self, idx: usize) -> i8 {
unsafe { core::ptr::read_unaligned(self.data.add(idx).cast::<i8>()) }
}
#[inline]
pub const unsafe fn read_i16(&self, idx: usize) -> i16 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 2).cast::<i16>()) }
}
#[inline]
pub const unsafe fn read_i32(&self, idx: usize) -> i32 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 4).cast::<i32>()) }
}
#[inline]
pub const unsafe fn read_i64(&self, idx: usize) -> i64 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 8).cast::<i64>()) }
}
#[inline]
pub const unsafe fn read_u8(&self, idx: usize) -> u8 {
unsafe { *self.data.add(idx) }
}
#[inline]
pub const unsafe fn read_u16(&self, idx: usize) -> u16 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 2).cast::<u16>()) }
}
#[inline]
pub const unsafe fn read_u32(&self, idx: usize) -> u32 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 4).cast::<u32>()) }
}
#[inline]
pub const unsafe fn read_u64(&self, idx: usize) -> u64 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 8).cast::<u64>()) }
}
#[inline]
pub const unsafe fn read_f32(&self, idx: usize) -> f32 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 4).cast::<f32>()) }
}
#[inline]
pub const unsafe fn read_f64(&self, idx: usize) -> f64 {
unsafe { core::ptr::read_unaligned(self.data.add(idx * 8).cast::<f64>()) }
}
#[inline]
pub const unsafe fn read_bool(&self, idx: usize) -> bool {
unsafe { *self.data.add(idx) != 0 }
}
#[inline]
pub const unsafe fn read_i128(&self, idx: usize) -> i128 {
let base = unsafe { self.data.add(idx * 16) };
let lower = unsafe { core::ptr::read_unaligned(base.cast::<u64>()) };
let upper = unsafe { core::ptr::read_unaligned(base.add(8).cast::<i64>()) };
#[allow(clippy::cast_lossless)]
let result = (upper as i128) << 64 | (lower as i128);
result
}
pub unsafe fn read_str(&self, idx: usize) -> &str {
unsafe { crate::vector::string::read_duck_string(self.data, idx) }
}
pub unsafe fn read_blob(&self, idx: usize) -> &[u8] {
unsafe { crate::vector::string::read_duck_string(self.data, idx).as_bytes() }
}
#[inline]
pub const unsafe fn read_uuid(&self, idx: usize) -> i128 {
unsafe { self.read_i128(idx) }
}
#[inline]
pub const unsafe fn read_date(&self, idx: usize) -> i32 {
unsafe { self.read_i32(idx) }
}
#[inline]
pub const unsafe fn read_timestamp(&self, idx: usize) -> i64 {
unsafe { self.read_i64(idx) }
}
#[inline]
pub const unsafe fn read_time(&self, idx: usize) -> i64 {
unsafe { self.read_i64(idx) }
}
#[inline]
pub const unsafe fn read_interval(&self, idx: usize) -> crate::interval::DuckInterval {
unsafe { crate::interval::read_interval_at(self.data, idx) }
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn bool_read_u8_pattern() {
let data: [u8; 4] = [0, 1, 2, 255];
let as_bools: Vec<bool> = data.iter().map(|&b| b != 0).collect();
assert_eq!(as_bools, [false, true, true, true]);
}
#[test]
fn row_count_is_zero_for_empty_state() {
let reader = VectorReader {
data: std::ptr::null(),
validity: std::ptr::null_mut(),
row_count: 0,
};
assert_eq!(reader.row_count(), 0);
}
#[test]
fn is_valid_when_validity_null() {
let reader = VectorReader {
data: std::ptr::null(),
validity: std::ptr::null_mut(),
row_count: 5,
};
assert!(unsafe { reader.is_valid(0) });
assert!(unsafe { reader.is_valid(4) });
}
}