Skip to main content

better_duck_core/raw/
selection_vector.rs

1//! RAII wrapper for a `duckdb_selection_vector` — an array of row indices used to
2//! slice a vector or drive a selected copy.
3//!
4//! A selection vector holds `size` `u32` indices. [`SelectionVector`] owns one
5//! (destroying it exactly once on drop) and exposes bounds-checked
6//! [`set`](SelectionVector::set)/[`get`](SelectionVector::get) over its index array.
7// FFI pointer args are used safely inside `unsafe` blocks.
8#![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
15/// An owned `duckdb_selection_vector` of `size` `u32` indices.
16pub struct SelectionVector {
17    ptr: duckdb_selection_vector,
18    size: u64,
19}
20
21impl SelectionVector {
22    /// Creates a selection vector holding `size` indices (initially unspecified;
23    /// fill them with [`set`](SelectionVector::set)).
24    ///
25    /// # Errors
26    ///
27    /// Returns `None` if DuckDB fails to allocate.
28    #[must_use]
29    pub fn new(size: u64) -> Option<SelectionVector> {
30        // SAFETY: always valid to call; a null return (allocation failure) → None.
31        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    /// The number of indices this selection vector holds.
39    #[must_use]
40    pub fn len(&self) -> u64 {
41        self.size
42    }
43
44    /// Whether the selection vector holds no indices.
45    #[must_use]
46    pub fn is_empty(&self) -> bool {
47        self.size == 0
48    }
49
50    /// Sets the index at position `at` to `row`.
51    ///
52    /// # Panics
53    ///
54    /// Panics if `at >= len()`.
55    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        // SAFETY: `at < size`, so `.add(at)` is in-bounds for the `size`-element array
62        // `duckdb_selection_vector_get_data_ptr` returns; the write is aligned for `u32`.
63        unsafe { *duckdb_selection_vector_get_data_ptr(self.ptr).add(at as usize) = row as sel_t };
64    }
65
66    /// The index at position `at`, or `None` if `at >= len()`.
67    #[must_use]
68    pub fn get(
69        &self,
70        at: u64,
71    ) -> Option<u32> {
72        if at >= self.size {
73            return None;
74        }
75        // SAFETY: `at < size`, so `.add(at)` is in-bounds and aligned for `u32`.
76        Some(unsafe { *duckdb_selection_vector_get_data_ptr(self.ptr).add(at as usize) } as u32)
77    }
78
79    /// The raw handle, borrowed for the duration of `&self`. The caller must not
80    /// destroy it.
81    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            // SAFETY: `self.ptr` is a valid, non-null selection vector created by
90            // `duckdb_create_selection_vector` and not yet destroyed; destroyed once.
91            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}