cuvs 26.4.0

RAPIDS vector search library
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

/*
 * SPDX-FileCopyrightText: Copyright (c) 2025-2026, NVIDIA CORPORATION.
 * SPDX-License-Identifier: Apache-2.0
 */

use crate::distance_type::DistanceType;
use crate::error::{check_cuvs, Result};
use std::fmt;
use std::io::{stderr, Write};

pub struct IndexParams(pub ffi::cuvsVamanaIndexParams_t);

impl IndexParams {
    /// Returns a new IndexParams
    pub fn new() -> Result<IndexParams> {
        unsafe {
            let mut params = std::mem::MaybeUninit::<ffi::cuvsVamanaIndexParams_t>::uninit();
            check_cuvs(ffi::cuvsVamanaIndexParamsCreate(params.as_mut_ptr()))?;
            Ok(IndexParams(params.assume_init()))
        }
    }

    /// DistanceType to use for building the index
    pub fn set_metric(self, metric: DistanceType) -> IndexParams {
        unsafe {
            (*self.0).metric = metric;
        }
        self
    }

    /// Maximum degree of output graph corresponds to the R parameter in the original Vamana
    /// literature.
    pub fn set_graph_degree(self, graph_degree: u32) -> IndexParams {
        unsafe {
            (*self.0).graph_degree = graph_degree;
        }
        self
    }

    /// Maximum number of visited nodes per search corresponds to the L parameter in the Vamana
    /// literature
    pub fn set_visited_size(self, visited_size: u32) -> IndexParams {
        unsafe {
            (*self.0).visited_size = visited_size;
        }
        self
    }

    /// Number of Vamana vector insertion iterations (each iteration inserts all vectors).
    pub fn set_vamana_iters(self, vamana_iters: f32) -> IndexParams {
        unsafe {
            (*self.0).vamana_iters = vamana_iters;
        }
        self
    }

    /// Alpha for pruning parameter
    pub fn set_alpha(self, alpha: f32) -> IndexParams {
        unsafe {
            (*self.0).alpha = alpha;
        }
        self
    }

    /// Maximum fraction of dataset inserted per batch.
    /// Larger max batch decreases graph quality, but improves speed
    pub fn set_max_fraction(self, max_fraction: f32) -> IndexParams {
        unsafe {
            (*self.0).max_fraction = max_fraction;
        }
        self
    }

    /// Base of growth rate of batch sizes
    pub fn set_batch_base(self, batch_base: f32) -> IndexParams {
        unsafe {
            (*self.0).batch_base = batch_base;
        }
        self
    }

    /// Size of candidate queue structure - should be (2^x)-1
    pub fn set_queue_size(self, queue_size: u32) -> IndexParams {
        unsafe {
            (*self.0).queue_size = queue_size;
        }
        self
    }

    /// Max batchsize of reverse edge processing (reduces memory footprint)
    pub fn set_reverse_batchsize(self, reverse_batchsize: u32) -> IndexParams {
        unsafe {
            (*self.0).reverse_batchsize = reverse_batchsize;
        }
        self
    }
}

impl fmt::Debug for IndexParams {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        // custom debug trait here, default value will show the pointer address
        // for the inner params object which isn't that useful.
        write!(f, "IndexParams({:?})", unsafe { *self.0 })
    }
}

impl Drop for IndexParams {
    fn drop(&mut self) {
        if let Err(e) = check_cuvs(unsafe { ffi::cuvsVamanaIndexParamsDestroy(self.0) }) {
            write!(
                stderr(),
                "failed to call cuvsVamanaIndexParamsDestroy {:?}",
                e
            )
            .expect("failed to write to stderr");
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_index_params() {
        let params = IndexParams::new()
            .unwrap()
            .set_alpha(1.0)
            .set_visited_size(128);

        unsafe {
            assert_eq!((*params.0).alpha, 1.0);
            assert_eq!((*params.0).visited_size, 128);
        }
    }
}