Skip to main content

next_plaid/
lib.rs

1//! Next-Plaid: CPU-based PLAID implementation for multi-vector search
2//!
3//! This crate provides a pure Rust, CPU-only implementation of the PLAID algorithm
4//! for efficient multi-vector search (late interaction retrieval).
5
6use std::sync::OnceLock;
7
8// Link BLAS implementation when feature is enabled
9#[cfg(feature = "accelerate")]
10extern crate blas_src;
11
12#[cfg(feature = "openblas")]
13extern crate openblas_src;
14
15pub mod codec;
16#[cfg(feature = "_cuda")]
17pub mod cuda;
18pub mod delete;
19pub mod embeddings;
20pub mod error;
21pub mod filtering;
22pub mod index;
23pub mod kmeans;
24pub mod maxsim;
25pub mod mmap;
26pub mod search;
27pub mod text_search;
28pub mod update;
29pub mod utils;
30
31pub use codec::ResidualCodec;
32pub use delete::delete_from_index;
33pub use error::{Error, Result};
34pub use index::MmapIndex;
35pub use index::{
36    encode_index_chunk, prepare_codec_artifacts, write_index_from_encoded_chunks,
37    EncodedIndexChunk, IndexConfig, Metadata, PreparedCodecArtifacts,
38};
39pub use kmeans::{
40    compute_centroids, compute_centroids_from_documents, compute_kmeans, estimate_num_partitions,
41    ComputeKmeansConfig, FastKMeans, KMeansConfig,
42};
43pub use search::{QueryResult, SearchParameters};
44pub use text_search::FtsTokenizer;
45pub use update::UpdateConfig;
46
47const DEFAULT_START_FROM_SCRATCH: usize = 999;
48
49fn parse_usize(raw: &str) -> Option<usize> {
50    raw.trim().parse::<usize>().ok()
51}
52
53pub fn default_start_from_scratch() -> usize {
54    static VALUE: OnceLock<usize> = OnceLock::new();
55    *VALUE.get_or_init(|| {
56        std::env::var("INDEX_DEFAULT_START_FROM_SCRATCH")
57            .ok()
58            .as_deref()
59            .and_then(parse_usize)
60            .unwrap_or(DEFAULT_START_FROM_SCRATCH)
61    })
62}
63
64#[cfg(feature = "_cuda")]
65pub use cuda::{clear_cuda_broken, is_cuda_broken, mark_cuda_broken, CudaContext};
66
67/// Check if GPU-only mode is forced via environment variable.
68/// Only checks the canonical `NEXT_PLAID_FORCE_GPU` env var.
69/// The higher-level `colgrep` crate's `apply_acceleration_mode()` propagates
70/// CLI flags and `COLGREP_*`/`FORCE_*` vars into this canonical var.
71pub fn is_force_gpu() -> bool {
72    std::env::var("NEXT_PLAID_FORCE_GPU")
73        .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
74        .unwrap_or(false)
75}
76
77/// Check if CPU-only mode is forced via environment variable.
78/// Only checks the canonical `NEXT_PLAID_FORCE_CPU` env var.
79pub fn is_force_cpu() -> bool {
80    !is_force_gpu()
81        && std::env::var("NEXT_PLAID_FORCE_CPU")
82            .map(|v| v == "1" || v.eq_ignore_ascii_case("true"))
83            .unwrap_or(false)
84}