use std::cell::Cell;
use crate::{
Result,
error::duckdb_failure_from_message,
ffi::{
DuckDBSuccess, duckdb_list_entry, duckdb_list_vector_get_child, duckdb_list_vector_get_size,
duckdb_list_vector_reserve, duckdb_list_vector_set_size, duckdb_state, duckdb_validity_set_row_invalid,
duckdb_vector, duckdb_vector_ensure_validity_writable, duckdb_vector_get_validity, duckdb_vector_size,
},
};
use super::{ArrayVector, FlatVector, StructVector};
pub(super) const MAX_VECTOR_SIZE: u64 = 1 << 37;
fn list_vector_state_result(state: duckdb_state, action: &str, size: usize) -> Result<()> {
if state == DuckDBSuccess {
Ok(())
} else {
Err(duckdb_failure_from_message(format!(
"failed to {action} {size} elements in DuckDB list vector"
)))
}
}
pub struct ListVector<'a> {
entries: FlatVector<'a>,
reserved_child_capacity: Cell<usize>,
}
impl<'a> ListVector<'a> {
pub(crate) unsafe fn from_raw(ptr: duckdb_vector) -> Self {
Self {
entries: unsafe { FlatVector::from_raw(ptr) },
reserved_child_capacity: Cell::new(0),
}
}
pub(crate) unsafe fn from_raw_with_capacity(ptr: duckdb_vector, capacity: usize) -> Self {
Self {
entries: unsafe { FlatVector::with_capacity(ptr, capacity) },
reserved_child_capacity: Cell::new(0),
}
}
pub fn len(&self) -> usize {
unsafe { duckdb_list_vector_get_size(self.entries.ptr) as usize }
}
pub fn is_empty(&self) -> bool {
self.len() == 0
}
pub fn child(&self, capacity: usize) -> FlatVector<'a> {
let ptr = self
.try_child_ptr_with_capacity(capacity)
.unwrap_or_else(|err| panic!("{err}"));
unsafe { FlatVector::with_capacity(ptr, capacity) }
}
pub fn struct_child(&self, capacity: usize) -> StructVector<'a> {
let ptr = self
.try_child_ptr_with_capacity(capacity)
.unwrap_or_else(|err| panic!("{err}"));
unsafe { StructVector::from_raw_with_capacity(ptr, capacity) }
}
fn child_capacity(&self) -> usize {
self.len()
.max(self.reserved_child_capacity.get())
.max(unsafe { duckdb_vector_size() as usize })
}
pub fn array_child(&self) -> ArrayVector<'a> {
unsafe {
ArrayVector::from_raw_with_capacity(duckdb_list_vector_get_child(self.entries.ptr), self.child_capacity())
}
}
pub fn list_child(&self) -> ListVector<'a> {
unsafe {
ListVector::from_raw_with_capacity(duckdb_list_vector_get_child(self.entries.ptr), self.child_capacity())
}
}
pub(crate) fn try_child_ptr_with_capacity(&self, capacity: usize) -> Result<duckdb_vector> {
self.try_reserve(capacity)?;
Ok(unsafe { duckdb_list_vector_get_child(self.entries.ptr) })
}
#[cfg(feature = "vtab-arrow")]
pub(crate) fn child_ptr(&self) -> duckdb_vector {
unsafe { duckdb_list_vector_get_child(self.entries.ptr) }
}
pub fn row_is_null(&self, row: u64) -> bool {
self.entries.row_is_null(row)
}
pub fn try_row_is_null(&self, row: u64) -> Result<bool> {
self.entries.try_row_is_null(row)
}
pub unsafe fn set_child<T: Copy>(&self, data: &[T]) {
unsafe { self.child(data.len()).copy(data) };
self.set_len(data.len());
}
pub fn set_entry(&mut self, idx: usize, offset: usize, length: usize) {
assert!(
idx < self.entries.capacity(),
"list entry row {idx} exceeds vector capacity {}",
self.entries.capacity()
);
unsafe {
self.entries
.as_mut_ptr::<duckdb_list_entry>()
.add(idx)
.write(duckdb_list_entry {
offset: offset as u64,
length: length as u64,
});
}
}
pub fn get_entry(&self, idx: usize) -> (usize, usize) {
self.try_get_entry(idx).unwrap_or_else(|err| panic!("{err}"))
}
pub fn try_get_entry(&self, idx: usize) -> Result<(usize, usize)> {
if idx >= self.entries.capacity() {
return Err(duckdb_failure_from_message(format!(
"list entry row {idx} exceeds vector capacity {}",
self.entries.capacity()
)));
}
let entry = unsafe { self.entries.as_mut_ptr::<duckdb_list_entry>().add(idx).read() };
let offset = usize::try_from(entry.offset).map_err(|_| {
duckdb_failure_from_message(format!(
"DuckDB list entry offset {} at row {idx} exceeds usize range",
entry.offset
))
})?;
let length = usize::try_from(entry.length).map_err(|_| {
duckdb_failure_from_message(format!(
"DuckDB list entry length {} at row {idx} exceeds usize range",
entry.length
))
})?;
Ok((offset, length))
}
pub fn set_null(&mut self, row: usize) {
unsafe {
duckdb_vector_ensure_validity_writable(self.entries.ptr);
let idx = duckdb_vector_get_validity(self.entries.ptr);
duckdb_validity_set_row_invalid(idx, row as u64);
}
}
pub fn try_reserve(&self, capacity: usize) -> Result<()> {
if capacity as u64 > MAX_VECTOR_SIZE {
return Err(duckdb_failure_from_message(format!(
"cannot reserve {capacity} elements in DuckDB list vector: exceeds maximum vector size {MAX_VECTOR_SIZE}"
)));
}
let state = unsafe { duckdb_list_vector_reserve(self.entries.ptr, capacity as u64) };
list_vector_state_result(state, "reserve child storage for", capacity)?;
self.reserved_child_capacity
.set(self.reserved_child_capacity.get().max(capacity));
Ok(())
}
pub fn set_len(&self, new_len: usize) {
self.try_set_len(new_len).unwrap_or_else(|err| panic!("{err}"));
}
pub fn try_set_len(&self, new_len: usize) -> Result<()> {
self.try_reserve(new_len)?;
let state = unsafe { duckdb_list_vector_set_size(self.entries.ptr, new_len as u64) };
list_vector_state_result(state, "set size to", new_len)
}
}