better_duck_core/raw/
selection_vector.rs1#![allow(clippy::not_unsafe_ptr_arg_deref)]
9
10use crate::ffi::{
11 duckdb_create_selection_vector, duckdb_destroy_selection_vector, duckdb_selection_vector,
12 duckdb_selection_vector_get_data_ptr, idx_t, sel_t,
13};
14
15pub struct SelectionVector {
17 ptr: duckdb_selection_vector,
18 size: u64,
19}
20
21impl SelectionVector {
22 #[must_use]
29 pub fn new(size: u64) -> Option<SelectionVector> {
30 let ptr = unsafe { duckdb_create_selection_vector(size as idx_t) };
32 if ptr.is_null() {
33 return None;
34 }
35 Some(SelectionVector { ptr, size })
36 }
37
38 #[must_use]
40 pub fn len(&self) -> u64 {
41 self.size
42 }
43
44 #[must_use]
46 pub fn is_empty(&self) -> bool {
47 self.size == 0
48 }
49
50 pub fn set(
56 &mut self,
57 at: u64,
58 row: u32,
59 ) {
60 assert!(at < self.size, "selection index {at} out of range (len {})", self.size);
61 unsafe { *duckdb_selection_vector_get_data_ptr(self.ptr).add(at as usize) = row as sel_t };
64 }
65
66 #[must_use]
68 pub fn get(
69 &self,
70 at: u64,
71 ) -> Option<u32> {
72 if at >= self.size {
73 return None;
74 }
75 Some(unsafe { *duckdb_selection_vector_get_data_ptr(self.ptr).add(at as usize) } as u32)
77 }
78
79 pub(crate) fn raw(&self) -> duckdb_selection_vector {
82 self.ptr
83 }
84}
85
86impl Drop for SelectionVector {
87 fn drop(&mut self) {
88 if !self.ptr.is_null() {
89 unsafe { duckdb_destroy_selection_vector(self.ptr) };
92 self.ptr = std::ptr::null_mut();
93 }
94 }
95}
96
97#[cfg(test)]
98mod tests {
99 use super::*;
100
101 #[test]
102 fn indices_round_trip_and_bounds_are_checked() {
103 let mut sel = SelectionVector::new(3).expect("selection vector");
104 assert_eq!(sel.len(), 3);
105 assert!(!sel.is_empty());
106 sel.set(0, 2);
107 sel.set(1, 0);
108 sel.set(2, 1);
109 assert_eq!(sel.get(0), Some(2));
110 assert_eq!(sel.get(2), Some(1));
111 assert_eq!(sel.get(3), None, "out-of-range read is None");
112 }
113
114 #[test]
115 #[should_panic(expected = "out of range")]
116 fn set_out_of_range_panics() {
117 let mut sel = SelectionVector::new(2).unwrap();
118 sel.set(5, 0);
119 }
120}