rabitq_rs/
python_bindings.rs

1//! Python bindings for MSTG index using PyO3
2#![allow(non_local_definitions)]
3
4#[cfg(feature = "python")]
5use numpy::{PyArray1, PyReadonlyArray1, PyReadonlyArray2};
6#[cfg(feature = "python")]
7use pyo3::prelude::*;
8
9use crate::mstg::{MstgConfig, MstgIndex, ScalarPrecision, SearchParams};
10use crate::rotation::RotatorType;
11use crate::Metric;
12
13#[cfg(feature = "python")]
14#[pyclass(name = "MstgIndex")]
15pub struct PyMstgIndex {
16    index: Option<MstgIndex>,
17    config: MstgConfig,
18    dimension: usize,
19}
20
21#[cfg(feature = "python")]
22#[pymethods]
23impl PyMstgIndex {
24    /// Create a new MSTG index with configuration
25    #[new]
26    #[allow(clippy::too_many_arguments)]
27    #[pyo3(signature = (
28        dimension,
29        metric="euclidean",
30        max_posting_size=16,
31        branching_factor=10,
32        balance_weight=1.0,
33        closure_epsilon=0.15,
34        max_replicas=8,
35        rabitq_bits=7,
36        faster_config=true,
37        hnsw_m=32,
38        hnsw_ef_construction=400,
39        centroid_precision="bf16",
40        default_ef_search=150,
41        pruning_epsilon=0.6
42    ))]
43    fn new(
44        dimension: usize,
45        metric: &str,
46        max_posting_size: usize,
47        branching_factor: usize,
48        balance_weight: f32,
49        closure_epsilon: f32,
50        max_replicas: usize,
51        rabitq_bits: usize,
52        faster_config: bool,
53        hnsw_m: usize,
54        hnsw_ef_construction: usize,
55        centroid_precision: &str,
56        default_ef_search: usize,
57        pruning_epsilon: f32,
58    ) -> PyResult<Self> {
59        let metric = match metric {
60            "euclidean" | "l2" => Metric::L2,
61            "angular" | "ip" | "inner_product" => Metric::InnerProduct,
62            _ => {
63                return Err(pyo3::exceptions::PyValueError::new_err(format!(
64                    "Invalid metric: {}. Use 'euclidean' or 'angular'",
65                    metric
66                )))
67            }
68        };
69
70        let centroid_precision = match centroid_precision {
71            "fp32" => ScalarPrecision::FP32,
72            "bf16" => ScalarPrecision::BF16,
73            "fp16" => ScalarPrecision::FP16,
74            "int8" => ScalarPrecision::INT8,
75            _ => {
76                return Err(pyo3::exceptions::PyValueError::new_err(format!(
77                    "Invalid precision: {}. Use 'fp32', 'bf16', 'fp16', or 'int8'",
78                    centroid_precision
79                )))
80            }
81        };
82
83        let config = MstgConfig {
84            max_posting_size,
85            branching_factor,
86            balance_weight,
87            closure_epsilon,
88            max_replicas,
89            rabitq_bits,
90            faster_config,
91            metric,
92            hnsw_m,
93            hnsw_ef_construction,
94            centroid_precision,
95            default_ef_search,
96            pruning_epsilon,
97        };
98
99        Ok(Self {
100            index: None,
101            config,
102            dimension,
103        })
104    }
105
106    /// Build index from numpy array (N x D)
107    fn fit(&mut self, data: PyReadonlyArray2<f32>) -> PyResult<()> {
108        let data = data.as_array();
109        let shape = data.shape();
110
111        if shape.len() != 2 {
112            return Err(pyo3::exceptions::PyValueError::new_err(
113                "Data must be 2D array (N x D)",
114            ));
115        }
116
117        if shape[1] != self.dimension {
118            return Err(pyo3::exceptions::PyValueError::new_err(format!(
119                "Data dimension {} does not match expected {}",
120                shape[1], self.dimension
121            )));
122        }
123
124        // Convert to Vec<Vec<f32>>
125        let n = shape[0];
126        let mut vectors = Vec::with_capacity(n);
127
128        for i in 0..n {
129            let row = data.row(i);
130            let vec: Vec<f32> = row.iter().copied().collect();
131            vectors.push(vec);
132        }
133
134        // Build index
135        match MstgIndex::build(&vectors, self.config.clone()) {
136            Ok(index) => {
137                self.index = Some(index);
138                Ok(())
139            }
140            Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
141                "Failed to build index: {}",
142                e
143            ))),
144        }
145    }
146
147    /// Set query-time search parameters
148    fn set_query_arguments(&mut self, ef_search: Option<usize>, pruning_epsilon: Option<f32>) {
149        if let Some(ef) = ef_search {
150            self.config.default_ef_search = ef;
151        }
152        if let Some(eps) = pruning_epsilon {
153            self.config.pruning_epsilon = eps;
154        }
155    }
156
157    /// Query for k nearest neighbors
158    /// Returns numpy array of shape (k, 2) with [id, distance]
159    fn query(&self, py: Python, query: PyReadonlyArray1<f32>, k: usize) -> PyResult<PyObject> {
160        let index = self.index.as_ref().ok_or_else(|| {
161            pyo3::exceptions::PyRuntimeError::new_err("Index not built yet. Call fit() first.")
162        })?;
163
164        let query = query.as_slice()?;
165
166        if query.len() != self.dimension {
167            return Err(pyo3::exceptions::PyValueError::new_err(format!(
168                "Query dimension {} does not match expected {}",
169                query.len(),
170                self.dimension
171            )));
172        }
173
174        let params = SearchParams::new(
175            self.config.default_ef_search,
176            self.config.pruning_epsilon,
177            k,
178        );
179
180        let results = index.search(query, &params);
181
182        // Convert to numpy array
183        let n = results.len();
184        let mut data = Vec::with_capacity(n * 2);
185        for result in &results {
186            data.push(result.vector_id as f32);
187            data.push(result.distance);
188        }
189
190        // Create 1D array then reshape to 2D
191        let array_1d = PyArray1::<f32>::from_vec(py, data);
192        let result_array = array_1d.reshape([n, 2]).unwrap();
193
194        Ok(result_array.to_owned().into_py(py))
195    }
196
197    /// Batch query for multiple queries
198    /// data: N x D array of queries
199    /// k: number of neighbors per query
200    /// Returns: list of N numpy arrays, each of shape (k, 2)
201    fn batch_query(
202        &self,
203        py: Python,
204        queries: PyReadonlyArray2<f32>,
205        k: usize,
206    ) -> PyResult<Vec<PyObject>> {
207        let index = self.index.as_ref().ok_or_else(|| {
208            pyo3::exceptions::PyRuntimeError::new_err("Index not built yet. Call fit() first.")
209        })?;
210
211        let queries = queries.as_array();
212        let shape = queries.shape();
213
214        if shape.len() != 2 {
215            return Err(pyo3::exceptions::PyValueError::new_err(
216                "Queries must be 2D array (N x D)",
217            ));
218        }
219
220        if shape[1] != self.dimension {
221            return Err(pyo3::exceptions::PyValueError::new_err(format!(
222                "Query dimension {} does not match expected {}",
223                shape[1], self.dimension
224            )));
225        }
226
227        let n_queries = shape[0];
228        let params = SearchParams::new(
229            self.config.default_ef_search,
230            self.config.pruning_epsilon,
231            k,
232        );
233
234        // Convert to Vec<Vec<f32>> for parallel batch_search
235        let query_vecs: Vec<Vec<f32>> = (0..n_queries)
236            .map(|i| {
237                let row = queries.row(i);
238                row.iter().copied().collect()
239            })
240            .collect();
241
242        // Parallel batch search (4-8x speedup)
243        let all_results = index.batch_search(&query_vecs, &params);
244
245        // Convert results to Python numpy arrays
246        all_results
247            .into_iter()
248            .map(|results| {
249                let n = results.len();
250                let mut data = Vec::with_capacity(n * 2);
251                for result in &results {
252                    data.push(result.vector_id as f32);
253                    data.push(result.distance);
254                }
255
256                let array_1d = PyArray1::<f32>::from_vec(py, data);
257                let result_array = array_1d.reshape([n, 2]).unwrap();
258                Ok(result_array.to_owned().into_py(py))
259            })
260            .collect()
261    }
262
263    /// Get memory usage in bytes
264    fn get_memory_usage(&self) -> PyResult<usize> {
265        let index = self
266            .index
267            .as_ref()
268            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("Index not built yet."))?;
269
270        // Estimate memory usage
271        let mem_mb = MstgIndex::estimate_memory_mb(&index.centroid_index, &index.posting_lists);
272        Ok((mem_mb * 1024.0 * 1024.0) as usize)
273    }
274
275    /// Get number of vectors in index
276    fn __len__(&self) -> PyResult<usize> {
277        let index = self
278            .index
279            .as_ref()
280            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("Index not built yet."))?;
281
282        let total: usize = index.posting_lists.iter().map(|p| p.len()).sum();
283        Ok(total)
284    }
285
286    fn __repr__(&self) -> String {
287        format!(
288            "MstgIndex(dimension={}, metric={:?}, built={})",
289            self.dimension,
290            self.config.metric,
291            self.index.is_some()
292        )
293    }
294
295    /// Save index to file
296    fn save(&self, path: &str) -> PyResult<()> {
297        let index = self
298            .index
299            .as_ref()
300            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("Index not built yet."))?;
301
302        index.save_to_path(path).map_err(|e| {
303            pyo3::exceptions::PyRuntimeError::new_err(format!("Save failed: {}", e))
304        })?;
305
306        Ok(())
307    }
308
309    /// Load index from file (InMemory mode - all data in RAM)
310    #[staticmethod]
311    fn load(path: &str) -> PyResult<Self> {
312        let index = MstgIndex::load_from_path(path).map_err(|e| {
313            pyo3::exceptions::PyRuntimeError::new_err(format!("Load failed: {}", e))
314        })?;
315
316        // Get dimension from centroid index (works for both InMemory and Mmap modes)
317        let dimension = index.centroid_index.dimension().unwrap_or(0);
318
319        let config = index.config.clone();
320
321        Ok(Self {
322            index: Some(index),
323            config,
324            dimension,
325        })
326    }
327
328    /// Load index from file using mmap (Disk mode - posting lists on disk)
329    ///
330    /// Only centroids and metadata are loaded into memory.
331    /// Posting lists are deserialized on-the-fly during search.
332    /// Uses much less memory but search is slower.
333    #[staticmethod]
334    fn load_mmap(path: &str) -> PyResult<Self> {
335        let index = MstgIndex::load_from_path_mmap(path).map_err(|e| {
336            pyo3::exceptions::PyRuntimeError::new_err(format!("Load mmap failed: {}", e))
337        })?;
338
339        let dimension = index.centroid_index.dimension().unwrap_or(0);
340
341        let config = index.config.clone();
342
343        Ok(Self {
344            index: Some(index),
345            config,
346            dimension,
347        })
348    }
349}
350
351// ============================================================================
352// IVF RaBitQ Index Python Bindings
353// ============================================================================
354
355#[cfg(feature = "python")]
356#[pyclass(name = "IvfRabitqIndex")]
357pub struct PyIvfRabitqIndex {
358    index: Option<crate::ivf::IvfRabitqIndex>,
359    dimension: usize,
360    metric: Metric,
361}
362
363#[cfg(feature = "python")]
364#[pymethods]
365impl PyIvfRabitqIndex {
366    /// Create a new IVF RaBitQ index
367    #[new]
368    #[pyo3(signature = (dimension, metric="euclidean"))]
369    fn new(dimension: usize, metric: &str) -> PyResult<Self> {
370        let metric = match metric {
371            "euclidean" | "l2" => Metric::L2,
372            "angular" | "ip" | "inner_product" => Metric::InnerProduct,
373            _ => {
374                return Err(pyo3::exceptions::PyValueError::new_err(format!(
375                    "Invalid metric: {}. Use 'euclidean' or 'angular'",
376                    metric
377                )))
378            }
379        };
380
381        Ok(Self {
382            index: None,
383            dimension,
384            metric,
385        })
386    }
387
388    /// Build index from numpy array (N x D) with automatic k-means clustering
389    #[pyo3(signature = (data, nlist, total_bits=7, rotator_type="random", seed=42, faster_config=true))]
390    fn fit(
391        &mut self,
392        data: PyReadonlyArray2<f32>,
393        nlist: usize,
394        total_bits: usize,
395        rotator_type: &str,
396        seed: u64,
397        faster_config: bool,
398    ) -> PyResult<()> {
399        let data = data.as_array();
400        let shape = data.shape();
401
402        if shape.len() != 2 {
403            return Err(pyo3::exceptions::PyValueError::new_err(
404                "Data must be 2D array (N x D)",
405            ));
406        }
407
408        if shape[1] != self.dimension {
409            return Err(pyo3::exceptions::PyValueError::new_err(format!(
410                "Data dimension {} does not match expected {}",
411                shape[1], self.dimension
412            )));
413        }
414
415        let rotator_type = match rotator_type {
416            "fht" | "random" => RotatorType::FhtKacRotator,
417            "matrix" | "identity" => RotatorType::MatrixRotator,
418            _ => {
419                return Err(pyo3::exceptions::PyValueError::new_err(format!(
420                    "Invalid rotator_type: {}. Use 'fht', 'random', 'matrix', or 'identity'",
421                    rotator_type
422                )))
423            }
424        };
425
426        // Convert to Vec<Vec<f32>>
427        let n = shape[0];
428        let mut vectors = Vec::with_capacity(n);
429
430        for i in 0..n {
431            let row = data.row(i);
432            let vec: Vec<f32> = row.iter().copied().collect();
433            vectors.push(vec);
434        }
435
436        // Build index
437        match crate::ivf::IvfRabitqIndex::train(
438            &vectors,
439            nlist,
440            total_bits,
441            self.metric,
442            rotator_type,
443            seed,
444            faster_config,
445        ) {
446            Ok(index) => {
447                self.index = Some(index);
448                Ok(())
449            }
450            Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
451                "Failed to build index: {}",
452                e
453            ))),
454        }
455    }
456
457    /// Build index from numpy array with pre-computed centroids and assignments
458    #[allow(clippy::too_many_arguments)]
459    #[pyo3(signature = (data, centroids, assignments, total_bits=7, rotator_type="random", seed=42, faster_config=true))]
460    fn fit_with_clusters(
461        &mut self,
462        data: PyReadonlyArray2<f32>,
463        centroids: PyReadonlyArray2<f32>,
464        assignments: PyReadonlyArray1<i32>,
465        total_bits: usize,
466        rotator_type: &str,
467        seed: u64,
468        faster_config: bool,
469    ) -> PyResult<()> {
470        let data = data.as_array();
471        let centroids = centroids.as_array();
472        let assignments = assignments.as_slice()?;
473
474        let data_shape = data.shape();
475        let centroids_shape = centroids.shape();
476
477        if data_shape.len() != 2 || centroids_shape.len() != 2 {
478            return Err(pyo3::exceptions::PyValueError::new_err(
479                "Data and centroids must be 2D arrays",
480            ));
481        }
482
483        if data_shape[1] != self.dimension || centroids_shape[1] != self.dimension {
484            return Err(pyo3::exceptions::PyValueError::new_err(format!(
485                "Data/centroids dimension must match expected {}",
486                self.dimension
487            )));
488        }
489
490        if data_shape[0] != assignments.len() {
491            return Err(pyo3::exceptions::PyValueError::new_err(
492                "Data and assignments must have same length",
493            ));
494        }
495
496        let rotator_type = match rotator_type {
497            "fht" | "random" => RotatorType::FhtKacRotator,
498            "matrix" | "identity" => RotatorType::MatrixRotator,
499            _ => {
500                return Err(pyo3::exceptions::PyValueError::new_err(format!(
501                    "Invalid rotator_type: {}. Use 'fht', 'random', 'matrix', or 'identity'",
502                    rotator_type
503                )))
504            }
505        };
506
507        // Convert data to Vec<Vec<f32>>
508        let n_data = data_shape[0];
509        let mut data_vecs = Vec::with_capacity(n_data);
510        for i in 0..n_data {
511            let row = data.row(i);
512            let vec: Vec<f32> = row.iter().copied().collect();
513            data_vecs.push(vec);
514        }
515
516        // Convert centroids to Vec<Vec<f32>>
517        let n_centroids = centroids_shape[0];
518        let mut centroid_vecs = Vec::with_capacity(n_centroids);
519        for i in 0..n_centroids {
520            let row = centroids.row(i);
521            let vec: Vec<f32> = row.iter().copied().collect();
522            centroid_vecs.push(vec);
523        }
524
525        // Convert assignments to Vec<usize>
526        let assignments_usize: Vec<usize> = assignments.iter().map(|&x| x as usize).collect();
527
528        // Build index
529        match crate::ivf::IvfRabitqIndex::train_with_clusters(
530            &data_vecs,
531            &centroid_vecs,
532            &assignments_usize,
533            total_bits,
534            self.metric,
535            rotator_type,
536            seed,
537            faster_config,
538        ) {
539            Ok(index) => {
540                self.index = Some(index);
541                Ok(())
542            }
543            Err(e) => Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
544                "Failed to build index with clusters: {}",
545                e
546            ))),
547        }
548    }
549
550    /// Query for k nearest neighbors
551    /// Returns numpy array of shape (k, 2) with [id, distance]
552    #[pyo3(signature = (query, k, nprobe=1))]
553    fn query(
554        &self,
555        py: Python,
556        query: PyReadonlyArray1<f32>,
557        k: usize,
558        nprobe: usize,
559    ) -> PyResult<PyObject> {
560        let index = self.index.as_ref().ok_or_else(|| {
561            pyo3::exceptions::PyRuntimeError::new_err("Index not built yet. Call fit() first.")
562        })?;
563
564        let query_slice = query.as_slice()?;
565
566        if query_slice.len() != self.dimension {
567            return Err(pyo3::exceptions::PyValueError::new_err(format!(
568                "Query dimension {} does not match expected {}",
569                query_slice.len(),
570                self.dimension
571            )));
572        }
573
574        let params = crate::ivf::SearchParams::new(k, nprobe);
575        let results = match index.search(query_slice, params) {
576            Ok(r) => r,
577            Err(e) => {
578                return Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
579                    "Search failed: {}",
580                    e
581                )))
582            }
583        };
584
585        // Pre-allocate exact size needed
586        let n = results.len();
587        let mut data = Vec::with_capacity(n * 2);
588
589        // Avoid iterator overhead for small vectors
590        for result in &results {
591            data.push(result.id as f32);
592            data.push(result.score);
593        }
594
595        // Create 1D array then reshape to 2D
596        let array_1d = PyArray1::<f32>::from_vec(py, data);
597        let result_array = array_1d.reshape([n, 2]).unwrap();
598
599        Ok(result_array.to_owned().into_py(py))
600    }
601
602    /// Batch query for multiple queries using parallel processing
603    /// data: N x D array of queries
604    /// k: number of neighbors per query
605    /// nprobe: number of clusters to probe
606    /// Returns: list of N numpy arrays, each of shape (k, 2)
607    ///
608    /// Performance: 2-8x faster than sequential queries (depends on CPU cores)
609    #[pyo3(signature = (queries, k, nprobe=1))]
610    fn batch_query(
611        &self,
612        py: Python,
613        queries: PyReadonlyArray2<f32>,
614        k: usize,
615        nprobe: usize,
616    ) -> PyResult<Vec<PyObject>> {
617        let index = self.index.as_ref().ok_or_else(|| {
618            pyo3::exceptions::PyRuntimeError::new_err("Index not built yet. Call fit() first.")
619        })?;
620
621        let queries_arr = queries.as_array();
622        let shape = queries_arr.shape();
623
624        if shape.len() != 2 {
625            return Err(pyo3::exceptions::PyValueError::new_err(
626                "Queries must be 2D array (N x D)",
627            ));
628        }
629
630        if shape[1] != self.dimension {
631            return Err(pyo3::exceptions::PyValueError::new_err(format!(
632                "Query dimension {} does not match expected {}",
633                shape[1], self.dimension
634            )));
635        }
636
637        let n_queries = shape[0];
638        let params = crate::ivf::SearchParams::new(k, nprobe);
639
640        // Convert queries to Vec of slices for batch_search
641        // This minimizes allocations - we only convert the numpy array to vectors once
642        let query_vecs: Vec<Vec<f32>> = (0..n_queries)
643            .map(|i| queries_arr.row(i).iter().copied().collect())
644            .collect();
645
646        let query_refs: Vec<&[f32]> = query_vecs.iter().map(|v| v.as_slice()).collect();
647
648        // Use parallel batch_search (2-8x speedup)
649        let all_results = index.batch_search(&query_refs, params);
650
651        // Pre-allocate result vector
652        let mut py_results = Vec::with_capacity(n_queries);
653
654        // Convert results to Python numpy arrays
655        for (i, result) in all_results.into_iter().enumerate() {
656            let results = match result {
657                Ok(r) => r,
658                Err(e) => {
659                    return Err(pyo3::exceptions::PyRuntimeError::new_err(format!(
660                        "Batch search failed at query {}: {}",
661                        i, e
662                    )))
663                }
664            };
665
666            // Pre-allocate exact size needed
667            let n = results.len();
668            let mut data = Vec::with_capacity(n * 2);
669            for result in &results {
670                data.push(result.id as f32);
671                data.push(result.score);
672            }
673
674            // Create 1D array then reshape to 2D
675            let array_1d = PyArray1::<f32>::from_vec(py, data);
676            let result_array = array_1d.reshape([n, 2]).unwrap();
677
678            py_results.push(result_array.to_owned().into_py(py));
679        }
680
681        Ok(py_results)
682    }
683
684    /// Save index to file
685    fn save(&self, path: &str) -> PyResult<()> {
686        let index = self.index.as_ref().ok_or_else(|| {
687            pyo3::exceptions::PyRuntimeError::new_err("Index not built yet. Call fit() first.")
688        })?;
689
690        index.save_to_path(path).map_err(|e| {
691            pyo3::exceptions::PyIOError::new_err(format!("Failed to save index: {}", e))
692        })
693    }
694
695    /// Load index from file
696    fn load(&mut self, path: &str) -> PyResult<()> {
697        let index = crate::ivf::IvfRabitqIndex::load_from_path(path).map_err(|e| {
698            pyo3::exceptions::PyIOError::new_err(format!("Failed to load index: {}", e))
699        })?;
700
701        self.index = Some(index);
702        Ok(())
703    }
704
705    /// Get number of vectors in index
706    fn __len__(&self) -> PyResult<usize> {
707        let index = self
708            .index
709            .as_ref()
710            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("Index not built yet."))?;
711
712        Ok(index.len())
713    }
714
715    /// Get number of clusters
716    fn cluster_count(&self) -> PyResult<usize> {
717        let index = self
718            .index
719            .as_ref()
720            .ok_or_else(|| pyo3::exceptions::PyRuntimeError::new_err("Index not built yet."))?;
721
722        Ok(index.cluster_count())
723    }
724
725    fn __repr__(&self) -> String {
726        format!(
727            "IvfRabitqIndex(dimension={}, metric={:?}, built={}, clusters={})",
728            self.dimension,
729            self.metric,
730            self.index.is_some(),
731            self.index
732                .as_ref()
733                .map(|idx| idx.cluster_count())
734                .unwrap_or(0)
735        )
736    }
737}
738
739/// Python module initialization
740#[cfg(feature = "python")]
741#[pymodule]
742fn rabitq_rs(_py: Python, m: &PyModule) -> PyResult<()> {
743    m.add_class::<PyMstgIndex>()?;
744    m.add_class::<PyIvfRabitqIndex>()?;
745    Ok(())
746}
rabitq_rs/python_bindings.rs

rabitq_rs/
python_bindings.rs
