From 2885da801effab65ad456e1ccbb81256e57d0c06 Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Fri, 21 Oct 2022 16:36:33 +0200 Subject: [PATCH 01/11] Better static feature --- Cargo.toml | 1 - ngt-sys/build.rs | 3 +++ src/lib.rs | 3 --- 3 files changed, 3 insertions(+), 4 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 90e2f1d..98b08f7 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,6 @@ readme = "README.md" [dependencies] ngt-sys = { path = "ngt-sys", version = "1.14.8-static" } num_enum = "0.5" -openmp-sys = { version="1.2.3", features=["static"] } scopeguard = "1" [dev-dependencies] diff --git a/ngt-sys/build.rs b/ngt-sys/build.rs index 8e83626..63a1761 100644 --- a/ngt-sys/build.rs +++ b/ngt-sys/build.rs @@ -27,6 +27,9 @@ fn main() { #[cfg(not(feature = "static"))] println!("cargo:rustc-link-lib=dylib=ngt"); + #[cfg(feature = "static")] + println!("cargo:rustc-link-lib=gomp"); + let bindings = bindgen::Builder::default() .clang_arg(format!("-I{}/include", dst.display())) .header(format!("{}/include/NGT/NGTQ/Capi.h", dst.display())) diff --git a/src/lib.rs b/src/lib.rs index 5b5e46a..b0e0e2c 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -80,9 +80,6 @@ //! //! [ngt]: https://github.com/yahoojapan/NGT -// See: https://gitlab.com/kornelski/openmp-rs#1-adding-rust-dependency -extern crate openmp_sys; - mod error; mod index; pub mod optim; From 30a6aa119eb495f88ec84fccc0504ee497e852c1 Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Fri, 28 Oct 2022 22:43:38 +0200 Subject: [PATCH 02/11] Adapt to latest codebase --- .github/workflows/ci.yaml | 5 +- Cargo.toml | 7 +- README.md | 2 +- ngt-sys/Cargo.toml | 5 +- ngt-sys/NGT | 2 +- ngt-sys/build.rs | 42 ++- src/error.rs | 44 ++- src/lib.rs | 105 ++----- src/{ => ngt}/index.rs | 336 ++------------------ src/ngt/mod.rs | 72 +++++ src/{ => ngt}/optim.rs | 0 src/{ => ngt}/properties.rs | 62 ++-- src/qbg/index.rs | 604 ++++++++++++++++++++++++++++++++++++ src/qbg/mod.rs | 7 + src/qbg/properties.rs | 255 +++++++++++++++ src/qg/index.rs | 305 ++++++++++++++++++ src/qg/mod.rs | 5 + src/qg/properties.rs | 259 ++++++++++++++++ 18 files changed, 1672 insertions(+), 445 deletions(-) rename src/{ => ngt}/index.rs (59%) create mode 100644 src/ngt/mod.rs rename src/{ => ngt}/optim.rs (100%) rename src/{ => ngt}/properties.rs (87%) create mode 100644 src/qbg/index.rs create mode 100644 src/qbg/mod.rs create mode 100644 src/qbg/properties.rs create mode 100644 src/qg/index.rs create mode 100644 src/qg/mod.rs create mode 100644 src/qg/properties.rs diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index d48d305..910646f 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -20,8 +20,11 @@ jobs: - default - shared_mem - large_data - - shared_mem,large_data + - quantized + - large_data,shared_mem + - large_data,quantized - static + - static,quantized - static,shared_mem,large_data steps: - uses: actions/checkout@v3 diff --git a/Cargo.toml b/Cargo.toml index 98b08f7..7c7e66b 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ngt" -version = "0.4.5" +version = "0.5.0" authors = ["Romain Leroux "] edition = "2021" description = "Rust wrappers for NGT nearest neighbor search." @@ -11,7 +11,7 @@ license = "Apache-2.0" readme = "README.md" [dependencies] -ngt-sys = { path = "ngt-sys", version = "1.14.8-static" } +ngt-sys = { path = "ngt-sys", version = "2.0.10" } num_enum = "0.5" scopeguard = "1" @@ -20,7 +20,8 @@ rayon = "1" tempfile = "3" [features] -default = [] +default = ["quantized"] # TODO: should not be default static = ["ngt-sys/static"] shared_mem = ["ngt-sys/shared_mem"] large_data = ["ngt-sys/large_data"] +quantized = ["ngt-sys/quantized"] diff --git a/README.md b/README.md index 31e9cc7..a34b8c2 100644 --- a/README.md +++ b/README.md @@ -12,7 +12,7 @@ Building NGT requires `CMake`. By default `ngt-rs` will be built dynamically, wh means that you'll need to make the build artifact `libngt.so` available to your final binary. You'll also need to have `OpenMP` installed on the system where it will run. If you want to build `ngt-rs` statically, then use the `static` Cargo feature, note that in -this case `OpenMP` will be disabled when building NGT. +this case `OpenMP` will be linked statically too. Furthermore, NGT's shared memory and large dataset features are available through Cargo features `shared_mem` and `large_data` respectively. diff --git a/ngt-sys/Cargo.toml b/ngt-sys/Cargo.toml index 694c094..06b1a58 100644 --- a/ngt-sys/Cargo.toml +++ b/ngt-sys/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ngt-sys" -version = "1.14.8-static" +version = "2.0.10" authors = ["Romain Leroux "] edition = "2021" links = "ngt" @@ -18,4 +18,5 @@ cpp_build = { version = "0.5", optional = true } [features] static = ["dep:cpp_build"] shared_mem = [] -large_data = [] \ No newline at end of file +large_data = [] +quantized = [] diff --git a/ngt-sys/NGT b/ngt-sys/NGT index b843128..bc9b4dd 160000 --- a/ngt-sys/NGT +++ b/ngt-sys/NGT @@ -1 +1 @@ -Subproject commit b84312870d496d0f4e3ead449bc6424545d9f896 +Subproject commit bc9b4dd9822c1d39ffdd3cbfbe822306c51171cc diff --git a/ngt-sys/build.rs b/ngt-sys/build.rs index 63a1761..5fa1a6f 100644 --- a/ngt-sys/build.rs +++ b/ngt-sys/build.rs @@ -5,38 +5,48 @@ fn main() { let out_dir = env::var("OUT_DIR").unwrap(); let mut config = cmake::Config::new("NGT"); - if env::var("CARGO_FEATURE_SHARED_MEM").is_ok() { config.define("NGT_SHARED_MEMORY_ALLOCATOR", "ON"); } - if env::var("CARGO_FEATURE_LARGE_DATA").is_ok() { config.define("NGT_LARGE_DATASET", "ON"); } - + if env::var("CARGO_FEATURE_QUANTIZED").is_err() { + config.define("NGT_QBG_DISABLED", "ON"); + } else { + config.define("NGT_AVX2", "ON"); + config.define("CMAKE_BUILD_TYPE", "Release"); + } let dst = config.build(); - #[cfg(feature = "static")] - cpp_build::Config::new() - .include(format!("{}/lib", out_dir)) - .build("src/lib.rs"); - println!("cargo:rustc-link-search=native={}/lib", dst.display()); - #[cfg(feature = "static")] - println!("cargo:rustc-link-lib=static=ngt"); #[cfg(not(feature = "static"))] - println!("cargo:rustc-link-lib=dylib=ngt"); - + { + println!("cargo:rustc-link-lib=dylib=ngt"); + } #[cfg(feature = "static")] - println!("cargo:rustc-link-lib=gomp"); + { + cpp_build::Config::new() + .include(format!("{}/lib", out_dir)) + .build("src/lib.rs"); + println!("cargo:rustc-link-lib=static=ngt"); + println!("cargo:rustc-link-lib=gomp"); + println!("cargo:rustc-link-lib=blas"); + println!("cargo:rustc-link-lib=lapack"); + } + + let capi_header = if cfg!(feature = "quantized") { + format!("{}/include/NGT/NGTQ/Capi.h", dst.display()) + } else { + format!("{}/include/NGT/Capi.h", dst.display()) + }; + let out_path = PathBuf::from(out_dir); let bindings = bindgen::Builder::default() .clang_arg(format!("-I{}/include", dst.display())) - .header(format!("{}/include/NGT/NGTQ/Capi.h", dst.display())) + .header(capi_header) .generate() .expect("Unable to generate bindings"); - - let out_path = PathBuf::from(out_dir); bindings .write_to_file(out_path.join("bindings.rs")) .expect("Couldn't write bindings"); diff --git a/src/error.rs b/src/error.rs index 33f83ee..9af95c7 100644 --- a/src/error.rs +++ b/src/error.rs @@ -3,8 +3,6 @@ use std::fmt; use ngt_sys as sys; -use crate::properties::{DistanceType, ObjectType}; - pub type Result = std::result::Result; #[derive(Debug)] @@ -43,14 +41,48 @@ impl From for Error { } } -impl From> for Error { - fn from(source: num_enum::TryFromPrimitiveError) -> Self { +impl From for Error { + fn from(source: std::ffi::IntoStringError) -> Self { + Self(source.to_string()) + } +} + +impl From> for Error { + fn from(source: num_enum::TryFromPrimitiveError) -> Self { + Self(source.to_string()) + } +} + +impl From> for Error { + fn from(source: num_enum::TryFromPrimitiveError) -> Self { + Self(source.to_string()) + } +} + +#[cfg(feature = "quantized")] +impl From> for Error { + fn from(source: num_enum::TryFromPrimitiveError) -> Self { + Self(source.to_string()) + } +} + +#[cfg(feature = "quantized")] +impl From> for Error { + fn from(source: num_enum::TryFromPrimitiveError) -> Self { + Self(source.to_string()) + } +} + +#[cfg(feature = "quantized")] +impl From> for Error { + fn from(source: num_enum::TryFromPrimitiveError) -> Self { Self(source.to_string()) } } -impl From> for Error { - fn from(source: num_enum::TryFromPrimitiveError) -> Self { +#[cfg(feature = "quantized")] +impl From> for Error { + fn from(source: num_enum::TryFromPrimitiveError) -> Self { Self(source.to_string()) } } diff --git a/src/lib.rs b/src/lib.rs index b0e0e2c..335ea5b 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -4,90 +4,39 @@ //! Building NGT requires `CMake`. By default `ngt-rs` will be built dynamically, which //! means that you'll need to make the build artifact `libngt.so` available to your final //! binary. You'll also need to have `OpenMP` installed on the system where it will run. If -//! you want to build `ngt-rs` statically, then use the `static` Cargo feature, note that in -//! this case `OpenMP` will be disabled when building NGT. +//! you want to build `ngt-rs` statically, then use the `static` Cargo feature. //! //! Furthermore, NGT's shared memory and large dataset features are available through Cargo //! features `shared_mem` and `large_data` respectively. //! -//! ## Usage -//! -//! Defining the properties of a new index: -//! -//! ```rust -//! # fn main() -> Result<(), ngt::Error> { -//! use ngt::{Properties, DistanceType, ObjectType}; -//! -//! // Defaut properties with vectors of dimension 3 -//! let prop = Properties::dimension(3)?; -//! -//! // Or customize values (here are the defaults) -//! let prop = Properties::dimension(3)? -//! .creation_edge_size(10)? -//! .search_edge_size(40)? -//! .object_type(ObjectType::Float)? -//! .distance_type(DistanceType::L2)?; -//! -//! # Ok(()) -//! # } -//! ``` -//! -//! Creating/Opening an index and using it: -//! -//! ```rust -//! # fn main() -> Result<(), ngt::Error> { -//! use ngt::{Index, Properties, EPSILON}; -//! -//! // Create a new index -//! let prop = Properties::dimension(3)?; -//! let index = Index::create("target/path/to/index/dir", prop)?; -//! -//! // Open an existing index -//! let mut index = Index::open("target/path/to/index/dir")?; -//! -//! // Insert two vectors and get their id -//! let vec1 = vec![1.0, 2.0, 3.0]; -//! let vec2 = vec![4.0, 5.0, 6.0]; -//! let id1 = index.insert(vec1)?; -//! let id2 = index.insert(vec2)?; -//! -//! // Actually build the index (not yet persisted on disk) -//! // This is required in order to be able to search vectors -//! index.build(2)?; -//! -//! // Perform a vector search (with 1 result) -//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; -//! assert_eq!(res[0].id, id1); -//! assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); -//! -//! // Remove a vector and check that it is not present anymore -//! index.remove(id1)?; -//! let res = index.get_vec(id1); -//! assert!(matches!(res, Result::Err(_))); -//! -//! // Verify that now our search result is different -//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; -//! assert_eq!(res[0].id, id2); -//! assert_eq!(index.get_vec(id2)?, vec![4.0, 5.0, 6.0]); -//! -//! // Persist index on disk -//! index.persist()?; -//! -//! # std::fs::remove_dir_all("target/path/to/index/dir").unwrap(); -//! # Ok(()) -//! # } -//! ``` -//! //! [ngt]: https://github.com/yahoojapan/NGT +// TODO: consider include_str README + +#[cfg(all(feature = "quantized", feature = "shared_mem"))] +compile_error!("only one of ['quantized', 'shared_mem'] can be enabled"); + mod error; -mod index; -pub mod optim; -mod properties; +mod ngt; + +#[cfg(feature = "quantized")] +pub mod qg; + +#[cfg(feature = "quantized")] +pub mod qbg; + +pub type VecId = u32; + +#[derive(Debug, Clone, PartialEq)] +pub struct SearchResult { + pub id: VecId, + pub distance: f32, +} + +pub const EPSILON: f32 = 0.1; -pub use crate::error::Error; -pub use crate::index::{Index, SearchResult, VecId, EPSILON}; -pub use crate::properties::{DistanceType, ObjectType, Properties}; +pub use crate::error::{Error, Result}; +pub use crate::ngt::{NgtDistance, NgtIndex, NgtObject, NgtProperties}; -#[cfg(not(feature = "shared_mem"))] -pub use crate::index::{QGIndex, QGQuantizationParams, QGQuery}; +// TODO: search etc only f32, drop support for f64/Into +// TODO: what about float16 ? diff --git a/src/index.rs b/src/ngt/index.rs similarity index 59% rename from src/index.rs rename to src/ngt/index.rs index 92655d5..f51c71a 100644 --- a/src/index.rs +++ b/src/ngt/index.rs @@ -9,34 +9,25 @@ use std::ptr; use ngt_sys as sys; use scopeguard::defer; +use super::{NgtObject, NgtProperties}; use crate::error::{make_err, Error, Result}; -use crate::properties::{ObjectType, Properties}; - -pub const EPSILON: f32 = 0.1; - -pub type VecId = u32; - -#[derive(Debug, Clone, PartialEq)] -pub struct SearchResult { - pub id: VecId, - pub distance: f32, -} +use crate::{SearchResult, VecId}; #[derive(Debug)] -pub struct Index { +pub struct NgtIndex { pub(crate) path: CString, - pub(crate) prop: Properties, + pub(crate) prop: NgtProperties, pub(crate) index: sys::NGTIndex, ospace: sys::NGTObjectSpace, ebuf: sys::NGTError, } -unsafe impl Send for Index {} -unsafe impl Sync for Index {} +unsafe impl Send for NgtIndex {} +unsafe impl Sync for NgtIndex {} -impl Index { - /// Creates an empty ANNG index with the given [`Properties`](). - pub fn create>(path: P, prop: Properties) -> Result { +impl NgtIndex { + /// Creates an empty ANNG index with the given [`NgtProperties`](). + pub fn create>(path: P, prop: NgtProperties) -> Result { if cfg!(feature = "shared_mem") && path.as_ref().exists() { Err(Error(format!("Path {:?} already exists", path.as_ref())))? } @@ -67,7 +58,7 @@ impl Index { Err(make_err(ebuf))? } - Ok(Index { + Ok(NgtIndex { path, prop, index, @@ -99,9 +90,9 @@ impl Index { Err(make_err(ebuf))? } - let prop = Properties::from(index)?; + let prop = NgtProperties::from(index)?; - Ok(Index { + Ok(NgtIndex { path, prop, index, @@ -113,7 +104,7 @@ impl Index { /// Search the nearest vectors to the specified query vector. /// - /// **The index must have been [`built`](Index::build) beforehand**. + /// **The index must have been [`built`](NgtIndex::build) beforehand**. pub fn search(&self, vec: &[f64], res_size: u64, epsilon: f32) -> Result> { unsafe { let results = sys::ngt_create_empty_results(self.ebuf); @@ -156,7 +147,7 @@ impl Index { /// Search linearly the nearest vectors to the specified query vector. /// - /// **The index must have been [`built`](Index::build) beforehand**. + /// **The index must have been [`built`](NgtIndex::build) beforehand**. pub fn linear_search(&self, vec: &[f64], res_size: u64) -> Result> { unsafe { let results = sys::ngt_create_empty_results(self.ebuf); @@ -198,12 +189,10 @@ impl Index { /// Insert the specified vector into the index. However note that it is not /// discoverable yet. /// - /// **The method [`build`](Index::build) must be called after inserting vectors**. - pub fn insert>(&mut self, vec: Vec) -> Result { + /// **The method [`build`](NgtIndex::build) must be called after inserting vectors**. + pub fn insert(&mut self, mut vec: Vec) -> Result { unsafe { - let mut vec = vec.into_iter().map(Into::into).collect::>(); - - let id = sys::ngt_insert_index( + let id = sys::ngt_insert_index_as_float( self.index, vec.as_mut_ptr(), self.prop.dimension as u32, @@ -220,7 +209,7 @@ impl Index { /// Insert the multiple vectors into the index. However note that they are not /// discoverable yet. /// - /// **The method [`build`](Index::build) must be called after inserting vectors**. + /// **The method [`build`](NgtIndex::build) must be called after inserting vectors**. pub fn insert_batch>(&mut self, batch: Vec>) -> Result<()> { let batch_size = u32::try_from(batch.len())?; @@ -285,7 +274,7 @@ impl Index { pub fn get_vec(&self, id: VecId) -> Result> { unsafe { let results = match self.prop.object_type { - ObjectType::Float => { + NgtObject::Float => { let results = sys::ngt_get_object_as_float(self.ospace, id, self.ebuf); if results.is_null() { Err(make_err(self.ebuf))? @@ -300,7 +289,7 @@ impl Index { results.iter().copied().collect::>() } - ObjectType::Uint8 => { + NgtObject::Uint8 => { let results = sys::ngt_get_object_as_integer(self.ospace, id, self.ebuf); if results.is_null() { Err(make_err(self.ebuf))? @@ -326,13 +315,13 @@ impl Index { unsafe { sys::ngt_get_number_of_objects(self.index, self.ebuf) } } - /// The number of indexed vectors, available after [`build`](Index::build). + /// The number of indexed vectors, available after [`build`](NgtIndex::build). pub fn nb_indexed(&self) -> u32 { unsafe { sys::ngt_get_number_of_indexed_objects(self.index, self.ebuf) } } } -impl Drop for Index { +impl Drop for NgtIndex { fn drop(&mut self) { if !self.index.is_null() { unsafe { sys::ngt_close_index(self.index) }; @@ -345,235 +334,6 @@ impl Drop for Index { } } -#[cfg(not(feature = "shared_mem"))] -#[derive(Debug)] -pub struct QGIndex { - pub(crate) prop: Properties, - pub(crate) index: sys::NGTQGIndex, - ebuf: sys::NGTError, -} - -#[cfg(not(feature = "shared_mem"))] -impl QGIndex { - pub fn quantize(index: Index, params: QGQuantizationParams) -> Result { - unsafe { - let ebuf = sys::ngt_create_error_object(); - defer! { sys::ngt_destroy_error_object(ebuf); } - - let path = index.path.clone(); - - drop(index); // Close the index - if !sys::ngtqg_quantize(path.as_ptr(), params.into_raw(), ebuf) { - Err(make_err(ebuf))? - } - - QGIndex::open(path.to_str().unwrap()) - } - } - - /// Open the already existing quantized index at the specified path. - pub fn open>(path: P) -> Result { - if !path.as_ref().exists() { - Err(Error(format!("Path {:?} does not exist", path.as_ref())))? - } - - unsafe { - let ebuf = sys::ngt_create_error_object(); - defer! { sys::ngt_destroy_error_object(ebuf); } - - let path = CString::new(path.as_ref().as_os_str().as_bytes())?; - - let index = sys::ngtqg_open_index(path.as_ptr(), ebuf); - if index.is_null() { - Err(make_err(ebuf))? - } - - let prop = Properties::from(index)?; - - Ok(QGIndex { - prop, - index, - ebuf: sys::ngt_create_error_object(), - }) - } - } - - pub fn search(&self, query: QGQuery) -> Result> { - unsafe { - let results = sys::ngt_create_empty_results(self.ebuf); - if results.is_null() { - Err(make_err(self.ebuf))? - } - defer! { sys::ngt_destroy_results(results); } - - if !sys::ngtqg_search_index(self.index, query.into_raw(), results, self.ebuf) { - Err(make_err(self.ebuf))? - } - - let rsize = sys::ngt_get_result_size(results, self.ebuf); - let mut ret = Vec::with_capacity(rsize as usize); - - for i in 0..rsize as u32 { - let d = sys::ngt_get_result(results, i, self.ebuf); - if d.id == 0 && d.distance == 0.0 { - Err(make_err(self.ebuf))? - } else { - ret.push(SearchResult { - id: d.id, - distance: d.distance, - }); - } - } - - Ok(ret) - } - } - - /// Get the specified vector. - pub fn get_vec(&self, id: VecId) -> Result> { - unsafe { - let results = match self.prop.object_type { - ObjectType::Float => { - let ospace = sys::ngt_get_object_space(self.index, self.ebuf); - if ospace.is_null() { - Err(make_err(self.ebuf))? - } - - let results = sys::ngt_get_object_as_float(ospace, id, self.ebuf); - if results.is_null() { - Err(make_err(self.ebuf))? - } - - let results = Vec::from_raw_parts( - results as *mut f32, - self.prop.dimension as usize, - self.prop.dimension as usize, - ); - let results = mem::ManuallyDrop::new(results); - - results.iter().copied().collect::>() - } - ObjectType::Uint8 => { - let ospace = sys::ngt_get_object_space(self.index, self.ebuf); - if ospace.is_null() { - Err(make_err(self.ebuf))? - } - - let results = sys::ngt_get_object_as_integer(ospace, id, self.ebuf); - if results.is_null() { - Err(make_err(self.ebuf))? - } - - let results = Vec::from_raw_parts( - results as *mut u8, - self.prop.dimension as usize, - self.prop.dimension as usize, - ); - let results = mem::ManuallyDrop::new(results); - - results.iter().map(|byte| *byte as f32).collect::>() - } - }; - - Ok(results) - } - } -} - -#[cfg(not(feature = "shared_mem"))] -impl Drop for QGIndex { - fn drop(&mut self) { - if !self.index.is_null() { - unsafe { sys::ngtqg_close_index(self.index) }; - self.index = ptr::null_mut(); - } - if !self.ebuf.is_null() { - unsafe { sys::ngt_destroy_error_object(self.ebuf) }; - self.ebuf = ptr::null_mut(); - } - } -} - -#[cfg(not(feature = "shared_mem"))] -#[derive(Debug, Clone, PartialEq)] -pub struct QGQuantizationParams { - pub dimension_of_subvector: f32, - pub max_number_of_edges: u64, -} - -#[cfg(not(feature = "shared_mem"))] -impl Default for QGQuantizationParams { - fn default() -> Self { - Self { - dimension_of_subvector: 0.0, - max_number_of_edges: 128, - } - } -} - -#[cfg(not(feature = "shared_mem"))] -impl QGQuantizationParams { - unsafe fn into_raw(self) -> sys::NGTQGQuantizationParameters { - sys::NGTQGQuantizationParameters { - dimension_of_subvector: self.dimension_of_subvector, - max_number_of_edges: self.max_number_of_edges, - } - } -} - -#[cfg(not(feature = "shared_mem"))] -#[derive(Debug, Clone, PartialEq)] -pub struct QGQuery<'a> { - query: &'a [f32], - pub size: u64, - pub epsilon: f32, - pub result_expansion: f32, - pub radius: f32, -} - -#[cfg(not(feature = "shared_mem"))] -impl<'a> QGQuery<'a> { - pub fn new(query: &'a [f32]) -> Self { - Self { - query, - size: 20, - epsilon: 0.03, - result_expansion: 3.0, - radius: f32::MAX, - } - } - - pub fn size(mut self, size: u64) -> Self { - self.size = size; - self - } - - pub fn epsilon(mut self, epsilon: f32) -> Self { - self.epsilon = epsilon; - self - } - - pub fn result_expansion(mut self, result_expansion: f32) -> Self { - self.result_expansion = result_expansion; - self - } - - pub fn radius(mut self, radius: f32) -> Self { - self.radius = radius; - self - } - - unsafe fn into_raw(self) -> sys::NGTQGQuery { - sys::NGTQGQuery { - query: self.query.as_ptr() as *mut f32, - size: self.size, - epsilon: self.epsilon, - result_expansion: self.result_expansion, - radius: self.radius, - } - } -} - #[cfg(test)] mod tests { use std::error::Error as StdError; @@ -584,6 +344,7 @@ mod tests { use tempfile::tempdir; use super::*; + use crate::EPSILON; #[test] fn test_basics() -> StdResult<(), Box> { @@ -594,8 +355,8 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = Properties::dimension(3)?; - let mut index = Index::create(dir.path(), prop)?; + let prop = NgtProperties::dimension(3)?; + let mut index = NgtIndex::create(dir.path(), prop)?; // Insert two vectors and get their id let vec1 = vec![1.0, 2.0, 3.0]; @@ -635,7 +396,7 @@ mod tests { // Persist index on disk, and open it again index.persist()?; - index = Index::open(dir.path())?; + index = NgtIndex::open(dir.path())?; assert!(index.nb_inserted() == 1); assert!(index.nb_indexed() == 1); @@ -661,8 +422,8 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = Properties::dimension(3)?; - let mut index = Index::create(dir.path(), prop)?; + let prop = NgtProperties::dimension(3)?; + let mut index = NgtIndex::create(dir.path(), prop)?; // Batch insert 2 vectors, build and persist the index index.insert_batch(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]])?; @@ -686,8 +447,8 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = Properties::dimension(3)?; - let mut index = Index::create(dir.path(), prop)?; + let prop = NgtProperties::dimension(3)?; + let mut index = NgtIndex::create(dir.path(), prop)?; let vecs = vec![ vec![1.0, 2.0, 3.0], @@ -719,41 +480,4 @@ mod tests { dir.close()?; Ok(()) } - - #[cfg(not(feature = "shared_mem"))] - #[test] - fn test_quantize() -> StdResult<(), Box> { - // Get a temporary directory to store the index - let dir = tempdir()?; - if cfg!(feature = "shared_mem") { - std::fs::remove_dir(dir.path())?; - } - - // Create an index for vectors of dimension 3 - let prop = Properties::dimension(3)?; - let mut index = Index::create(dir.path(), prop)?; - - // Insert two vectors and get their id - let vec1 = vec![1.0, 2.0, 3.0]; - let vec2 = vec![4.0, 5.0, 6.0]; - let id1 = index.insert(vec1.clone())?; - let _id2 = index.insert(vec2.clone())?; - - // Build and persist the index - index.build(1)?; - index.persist()?; - - let params = QGQuantizationParams::default(); - let index = QGIndex::quantize(index, params)?; - - // Perform a vector search (with 1 result) - let vec = vec![1.1, 2.1, 3.1]; - let query = QGQuery::new(&vec).size(2); - let res = index.search(query)?; - assert_eq!(id1, res[0].id); - assert_eq!(vec1, index.get_vec(id1)?); - - dir.close()?; - Ok(()) - } } diff --git a/src/ngt/mod.rs b/src/ngt/mod.rs new file mode 100644 index 0000000..67cc417 --- /dev/null +++ b/src/ngt/mod.rs @@ -0,0 +1,72 @@ +//! Defining the properties of a new index: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! use ngt::{NgtProperties, NgtDistance, NgtObject}; +//! +//! // Defaut properties with vectors of dimension 3 +//! let prop = NgtProperties::dimension(3)?; +//! +//! // Or customize values (here are the defaults) +//! let prop = NgtProperties::dimension(3)? +//! .creation_edge_size(10)? +//! .search_edge_size(40)? +//! .object_type(NgtObject::Float)? +//! .distance_type(NgtDistance::L2)?; +//! +//! # Ok(()) +//! # } +//! ``` +//! +//! Creating/Opening an index and using it: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! use ngt::{NgtIndex, NgtProperties, EPSILON}; +//! +//! // Create a new index +//! let prop = NgtProperties::dimension(3)?; +//! let index = NgtIndex::create("target/path/to/index/dir", prop)?; +//! +//! // Open an existing index +//! let mut index = NgtIndex::open("target/path/to/index/dir")?; +//! +//! // Insert two vectors and get their id +//! let vec1 = vec![1.0, 2.0, 3.0]; +//! let vec2 = vec![4.0, 5.0, 6.0]; +//! let id1 = index.insert(vec1)?; +//! let id2 = index.insert(vec2)?; +//! +//! // Actually build the index (not yet persisted on disk) +//! // This is required in order to be able to search vectors +//! index.build(2)?; +//! +//! // Perform a vector search (with 1 result) +//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; +//! assert_eq!(res[0].id, id1); +//! assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); +//! +//! // Remove a vector and check that it is not present anymore +//! index.remove(id1)?; +//! let res = index.get_vec(id1); +//! assert!(matches!(res, Result::Err(_))); +//! +//! // Verify that now our search result is different +//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; +//! assert_eq!(res[0].id, id2); +//! assert_eq!(index.get_vec(id2)?, vec![4.0, 5.0, 6.0]); +//! +//! // Persist index on disk +//! index.persist()?; +//! +//! # std::fs::remove_dir_all("target/path/to/index/dir").unwrap(); +//! # Ok(()) +//! # } +//! ``` + +mod index; +// mod optim; +mod properties; + +pub use self::index::NgtIndex; +pub use self::properties::{NgtDistance, NgtObject, NgtProperties}; diff --git a/src/optim.rs b/src/ngt/optim.rs similarity index 100% rename from src/optim.rs rename to src/ngt/optim.rs diff --git a/src/properties.rs b/src/ngt/properties.rs similarity index 87% rename from src/properties.rs rename to src/ngt/properties.rs index c97413d..688b8c4 100644 --- a/src/properties.rs +++ b/src/ngt/properties.rs @@ -9,14 +9,14 @@ use crate::error::{make_err, Result}; #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] -pub enum ObjectType { +pub enum NgtObject { Uint8 = 1, Float = 2, } #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] -pub enum DistanceType { +pub enum NgtDistance { L1 = 0, L2 = 1, Angle = 2, @@ -32,25 +32,25 @@ pub enum DistanceType { } #[derive(Debug)] -pub struct Properties { +pub struct NgtProperties { pub(crate) dimension: i32, pub(crate) creation_edge_size: i16, pub(crate) search_edge_size: i16, - pub(crate) object_type: ObjectType, - pub(crate) distance_type: DistanceType, + pub(crate) object_type: NgtObject, + pub(crate) distance_type: NgtDistance, pub(crate) raw_prop: sys::NGTProperty, } -unsafe impl Send for Properties {} -unsafe impl Sync for Properties {} +unsafe impl Send for NgtProperties {} +unsafe impl Sync for NgtProperties {} -impl Properties { +impl NgtProperties { pub fn dimension(dimension: usize) -> Result { let dimension = i32::try_from(dimension)?; let creation_edge_size = 10; let search_edge_size = 40; - let object_type = ObjectType::Float; - let distance_type = DistanceType::L2; + let object_type = NgtObject::Float; + let distance_type = NgtDistance::L2; unsafe { let ebuf = sys::ngt_create_error_object(); @@ -138,13 +138,13 @@ impl Properties { if object_type < 0 { Err(make_err(ebuf))? } - let object_type = ObjectType::try_from(object_type)?; + let object_type = NgtObject::try_from(object_type)?; let distance_type = sys::ngt_get_property_distance_type(raw_prop, ebuf); if distance_type < 0 { Err(make_err(ebuf))? } - let distance_type = DistanceType::try_from(distance_type)?; + let distance_type = NgtDistance::try_from(distance_type)?; Ok(Self { dimension, @@ -204,23 +204,23 @@ impl Properties { Ok(()) } - pub fn object_type(mut self, object_type: ObjectType) -> Result { + pub fn object_type(mut self, object_type: NgtObject) -> Result { self.object_type = object_type; unsafe { Self::set_object_type(self.raw_prop, object_type)? }; Ok(self) } - unsafe fn set_object_type(raw_prop: sys::NGTProperty, object_type: ObjectType) -> Result<()> { + unsafe fn set_object_type(raw_prop: sys::NGTProperty, object_type: NgtObject) -> Result<()> { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } match object_type { - ObjectType::Uint8 => { + NgtObject::Uint8 => { if !sys::ngt_set_property_object_type_integer(raw_prop, ebuf) { Err(make_err(ebuf))? } } - ObjectType::Float => { + NgtObject::Float => { if !sys::ngt_set_property_object_type_float(raw_prop, ebuf) { Err(make_err(ebuf))? } @@ -230,7 +230,7 @@ impl Properties { Ok(()) } - pub fn distance_type(mut self, distance_type: DistanceType) -> Result { + pub fn distance_type(mut self, distance_type: NgtDistance) -> Result { self.distance_type = distance_type; unsafe { Self::set_distance_type(self.raw_prop, distance_type)? }; Ok(self) @@ -238,68 +238,68 @@ impl Properties { unsafe fn set_distance_type( raw_prop: sys::NGTProperty, - distance_type: DistanceType, + distance_type: NgtDistance, ) -> Result<()> { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } match distance_type { - DistanceType::L1 => { + NgtDistance::L1 => { if !sys::ngt_set_property_distance_type_l1(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::L2 => { + NgtDistance::L2 => { if !sys::ngt_set_property_distance_type_l2(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::Angle => { + NgtDistance::Angle => { if !sys::ngt_set_property_distance_type_angle(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::Hamming => { + NgtDistance::Hamming => { if !sys::ngt_set_property_distance_type_hamming(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::Cosine => { + NgtDistance::Cosine => { if !sys::ngt_set_property_distance_type_cosine(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::NormalizedAngle => { + NgtDistance::NormalizedAngle => { if !sys::ngt_set_property_distance_type_normalized_angle(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::NormalizedCosine => { + NgtDistance::NormalizedCosine => { if !sys::ngt_set_property_distance_type_normalized_cosine(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::Jaccard => { + NgtDistance::Jaccard => { if !sys::ngt_set_property_distance_type_jaccard(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::SparseJaccard => { + NgtDistance::SparseJaccard => { if !sys::ngt_set_property_distance_type_sparse_jaccard(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::NormalizedL2 => { + NgtDistance::NormalizedL2 => { if !sys::ngt_set_property_distance_type_normalized_l2(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::Poincare => { + NgtDistance::Poincare => { if !sys::ngt_set_property_distance_type_poincare(raw_prop, ebuf) { Err(make_err(ebuf))? } } - DistanceType::Lorentz => { + NgtDistance::Lorentz => { if !sys::ngt_set_property_distance_type_lorentz(raw_prop, ebuf) { Err(make_err(ebuf))? } @@ -310,7 +310,7 @@ impl Properties { } } -impl Drop for Properties { +impl Drop for NgtProperties { fn drop(&mut self) { if !self.raw_prop.is_null() { unsafe { sys::ngt_destroy_property(self.raw_prop) }; diff --git a/src/qbg/index.rs b/src/qbg/index.rs new file mode 100644 index 0000000..f57b12e --- /dev/null +++ b/src/qbg/index.rs @@ -0,0 +1,604 @@ +use std::ffi::CString; +use std::os::unix::ffi::OsStrExt; +use std::path::Path; +use std::{mem, ptr}; + +use ngt_sys as sys; +use num_enum::TryFromPrimitive; +use scopeguard::defer; + +use super::{QbgDistance, QbgObject}; +use crate::error::{make_err, Error, Result}; +use crate::{SearchResult, VecId}; + +#[derive(Debug)] +pub struct QbgIndex { + pub(crate) index: sys::QBGIndex, + path: CString, + _mode: T, + dimension: u32, + ebuf: sys::NGTError, +} + +impl QbgIndex { + pub fn create

(path: P, create_params: QbgConstructParams) -> Result + where + P: AsRef, + { + if !is_x86_feature_detected!("avx2") { + return Err(Error( + "Cannot quantize an index without AVX2 support".into(), + )); + } + + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let path = CString::new(path.as_ref().as_os_str().as_bytes())?; + + if !sys::qbg_create(path.as_ptr(), &mut create_params.into_raw() as *mut _, ebuf) { + Err(make_err(ebuf))? + } + + let index = sys::qbg_open_index(path.as_ptr(), false, ebuf); + if index.is_null() { + Err(make_err(ebuf))? + } + + let dimension = sys::qbg_get_dimension(index, ebuf) as u32; + if dimension == 0 { + Err(make_err(ebuf))? + } + + Ok(QbgIndex { + index, + path, + _mode: ModeWrite, + dimension, + ebuf: sys::ngt_create_error_object(), + }) + } + } + + pub fn insert(&mut self, mut vec: Vec) -> Result { + unsafe { + let id = + sys::qbg_append_object(self.index, vec.as_mut_ptr(), self.dimension, self.ebuf); + if id == 0 { + Err(make_err(self.ebuf))? + } + + Ok(id) + } + } + + pub fn build(&mut self, build_params: QbgBuildParams) -> Result<()> { + unsafe { + if !sys::qbg_build_index( + self.path.as_ptr(), + &mut build_params.into_raw() as *mut _, + self.ebuf, + ) { + Err(make_err(self.ebuf))? + } + Ok(()) + } + } + + pub fn persist(&mut self) -> Result<()> { + unsafe { + if !sys::qbg_save_index(self.index, self.ebuf) { + Err(make_err(self.ebuf))? + } + Ok(()) + } + } + + pub fn into_readable(self) -> Result> { + let path = self.path.clone(); + drop(self); + QbgIndex::open(path.into_string()?) + } +} + +impl QbgIndex { + pub fn open>(path: P) -> Result { + if !is_x86_feature_detected!("avx2") { + return Err(Error( + "Cannot use a quantized index without AVX2 support".into(), + )); + } + + if !path.as_ref().exists() { + Err(Error(format!("Path {:?} does not exist", path.as_ref())))? + } + + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let path = CString::new(path.as_ref().as_os_str().as_bytes())?; + let index = sys::qbg_open_index(path.as_ptr(), true, ebuf); + if index.is_null() { + Err(make_err(ebuf))? + } + + let dimension = sys::qbg_get_dimension(index, ebuf) as u32; + if dimension == 0 { + Err(make_err(ebuf))? + } + + Ok(QbgIndex { + index, + path, + _mode: ModeRead, + dimension, + ebuf: sys::ngt_create_error_object(), + }) + } + } + + pub fn search(&self, query: QbgQuery) -> Result> { + unsafe { + let results = sys::ngt_create_empty_results(self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + defer! { sys::qbg_destroy_results(results); } + + if !sys::qbg_search_index(self.index, query.into_raw(), results, self.ebuf) { + Err(make_err(self.ebuf))? + } + + let rsize = sys::qbg_get_result_size(results, self.ebuf); + let mut ret = Vec::with_capacity(rsize as usize); + + for i in 0..rsize as u32 { + let d = sys::qbg_get_result(results, i, self.ebuf); + if d.id == 0 && d.distance == 0.0 { + Err(make_err(self.ebuf))? + } else { + ret.push(SearchResult { + id: d.id, + distance: d.distance, + }); + } + } + + Ok(ret) + } + } + + pub fn into_writable(self) -> Result> { + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let path = self.path.clone(); + drop(self); + + let index = sys::qbg_open_index(path.as_ptr(), false, ebuf); + if index.is_null() { + Err(make_err(ebuf))? + } + + let dimension = sys::qbg_get_dimension(index, ebuf) as u32; + if dimension == 0 { + Err(make_err(ebuf))? + } + + Ok(QbgIndex { + index, + path, + _mode: ModeWrite, + dimension, + ebuf: sys::ngt_create_error_object(), + }) + } + } +} + +impl QbgIndex +where + T: IndexMode, +{ + pub fn get_vec(&self, id: VecId) -> Result> { + unsafe { + let results = sys::qbg_get_object(self.index, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + + let results = Vec::from_raw_parts( + results as *mut f32, + self.dimension as usize, + self.dimension as usize, + ); + let results = mem::ManuallyDrop::new(results); + let results = results.iter().copied().collect::>(); + + Ok(results) + } + } +} + +impl Drop for QbgIndex { + fn drop(&mut self) { + if !self.index.is_null() { + unsafe { sys::qbg_close_index(self.index) }; + self.index = ptr::null_mut(); + } + if !self.ebuf.is_null() { + unsafe { sys::ngt_destroy_error_object(self.ebuf) }; + self.ebuf = ptr::null_mut(); + } + } +} + +mod private { + pub trait Sealed {} +} + +pub trait IndexMode: private::Sealed {} + +#[derive(Debug, Clone, Copy)] +pub struct ModeRead; + +impl private::Sealed for ModeRead {} +impl IndexMode for ModeRead {} + +#[derive(Debug, Clone, Copy)] +pub struct ModeWrite; + +impl private::Sealed for ModeWrite {} +impl IndexMode for ModeWrite {} + +#[derive(Debug, Clone, PartialEq)] +pub struct QbgConstructParams { + extended_dimension: u64, + dimension: u64, + number_of_subvectors: u64, + number_of_blobs: u64, + internal_data_type: QbgObject, + data_type: QbgObject, + distance_type: QbgDistance, +} + +impl Default for QbgConstructParams { + fn default() -> Self { + Self { + extended_dimension: 0, + dimension: 0, + number_of_subvectors: 1, + number_of_blobs: 0, + internal_data_type: QbgObject::Float, + data_type: QbgObject::Float, + distance_type: QbgDistance::L2, + } + } +} + +impl QbgConstructParams { + pub fn extended_dimension(mut self, extended_dimension: u64) -> Self { + self.extended_dimension = extended_dimension; + self + } + + pub fn dimension(mut self, dimension: u64) -> Self { + self.dimension = dimension; + self + } + + pub fn number_of_subvectors(mut self, number_of_subvectors: u64) -> Self { + self.number_of_subvectors = number_of_subvectors; + self + } + + pub fn number_of_blobs(mut self, number_of_blobs: u64) -> Self { + self.number_of_blobs = number_of_blobs; + self + } + + pub fn internal_data_type(mut self, internal_data_type: QbgObject) -> Self { + self.internal_data_type = internal_data_type; + self + } + + pub fn data_type(mut self, data_type: QbgObject) -> Self { + self.data_type = data_type; + self + } + + pub fn distance_type(mut self, distance_type: QbgDistance) -> Self { + self.distance_type = distance_type; + self + } + + unsafe fn into_raw(self) -> sys::QBGConstructionParameters { + sys::QBGConstructionParameters { + extended_dimension: self.extended_dimension, + dimension: self.dimension, + number_of_subvectors: self.number_of_subvectors, + number_of_blobs: self.number_of_blobs, + internal_data_type: self.internal_data_type as i32, + data_type: self.data_type as i32, + distance_type: self.distance_type as i32, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgClusteringInitMode { + Head = 0, + Random = 1, + KmeansPlusPlus = 2, + RandomFixedSeed = 3, + KmeansPlusPlusFixedSeed = 4, + Best = 5, +} + +#[derive(Debug, Clone)] +pub struct QbgBuildParams { + // hierarchical kmeans + hierarchical_clustering_init_mode: QbgClusteringInitMode, + number_of_first_objects: u64, + number_of_first_clusters: u64, + number_of_second_objects: u64, + number_of_second_clusters: u64, + number_of_third_clusters: u64, + // optimization + number_of_objects: u64, + number_of_subvectors: u64, + optimization_clustering_init_mode: QbgClusteringInitMode, + rotation_iteration: u64, + subvector_iteration: u64, + number_of_matrices: u64, + rotation: bool, + repositioning: bool, +} + +impl Default for QbgBuildParams { + fn default() -> Self { + Self { + hierarchical_clustering_init_mode: QbgClusteringInitMode::KmeansPlusPlus, + number_of_first_objects: 0, + number_of_first_clusters: 0, + number_of_second_objects: 0, + number_of_second_clusters: 0, + number_of_third_clusters: 0, + number_of_objects: 1000, + number_of_subvectors: 1, + optimization_clustering_init_mode: QbgClusteringInitMode::KmeansPlusPlus, + rotation_iteration: 2000, + subvector_iteration: 400, + number_of_matrices: 3, + rotation: true, + repositioning: false, + } + } +} + +impl QbgBuildParams { + pub fn hierarchical_clustering_init_mode( + mut self, + clustering_init_mode: QbgClusteringInitMode, + ) -> Self { + self.hierarchical_clustering_init_mode = clustering_init_mode; + self + } + + pub fn number_of_first_objects(mut self, number_of_first_objects: u64) -> Self { + self.number_of_first_objects = number_of_first_objects; + self + } + + pub fn number_of_first_clusters(mut self, number_of_first_clusters: u64) -> Self { + self.number_of_first_clusters = number_of_first_clusters; + self + } + + pub fn number_of_second_objects(mut self, number_of_second_objects: u64) -> Self { + self.number_of_second_objects = number_of_second_objects; + self + } + + pub fn number_of_second_clusters(mut self, number_of_second_clusters: u64) -> Self { + self.number_of_second_clusters = number_of_second_clusters; + self + } + + pub fn number_of_third_clusters(mut self, number_of_third_clusters: u64) -> Self { + self.number_of_third_clusters = number_of_third_clusters; + self + } + + pub fn number_of_objects(mut self, number_of_objects: u64) -> Self { + self.number_of_objects = number_of_objects; + self + } + pub fn number_of_subvectors(mut self, number_of_subvectors: u64) -> Self { + self.number_of_subvectors = number_of_subvectors; + self + } + pub fn optimization_clustering_init_mode( + mut self, + clustering_init_mode: QbgClusteringInitMode, + ) -> Self { + self.optimization_clustering_init_mode = clustering_init_mode; + self + } + + pub fn rotation_iteration(mut self, rotation_iteration: u64) -> Self { + self.rotation_iteration = rotation_iteration; + self + } + + pub fn subvector_iteration(mut self, subvector_iteration: u64) -> Self { + self.subvector_iteration = subvector_iteration; + self + } + + pub fn number_of_matrices(mut self, number_of_matrices: u64) -> Self { + self.number_of_matrices = number_of_matrices; + self + } + + pub fn rotation(mut self, rotation: bool) -> Self { + self.rotation = rotation; + self + } + + pub fn repositioning(mut self, repositioning: bool) -> Self { + self.repositioning = repositioning; + self + } + + unsafe fn into_raw(self) -> sys::QBGBuildParameters { + sys::QBGBuildParameters { + hierarchical_clustering_init_mode: self.hierarchical_clustering_init_mode as i32, + number_of_first_objects: self.number_of_first_objects, + number_of_first_clusters: self.number_of_first_clusters, + number_of_second_objects: self.number_of_second_objects, + number_of_second_clusters: self.number_of_second_clusters, + number_of_third_clusters: self.number_of_third_clusters, + number_of_objects: self.number_of_objects, + number_of_subvectors: self.number_of_subvectors, + optimization_clustering_init_mode: self.optimization_clustering_init_mode as i32, + rotation_iteration: self.rotation_iteration, + subvector_iteration: self.subvector_iteration, + number_of_matrices: self.number_of_matrices, + rotation: self.rotation, + repositioning: self.repositioning, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct QbgQuery<'a> { + query: &'a [f32], + pub size: u64, + pub epsilon: f32, + pub blob_epsilon: f32, + pub result_expansion: f32, + pub number_of_explored_blobs: u64, + pub number_of_edges: u64, + pub radius: f32, +} + +impl<'a> QbgQuery<'a> { + pub fn new(query: &'a [f32]) -> Self { + Self { + query, + size: 20, + epsilon: 0.1, + blob_epsilon: 0.0, + result_expansion: 3.0, + number_of_explored_blobs: 256, + number_of_edges: 0, + radius: 0.0, + } + } + + pub fn size(mut self, size: u64) -> Self { + self.size = size; + self + } + + pub fn epsilon(mut self, epsilon: f32) -> Self { + self.epsilon = epsilon; + self + } + + pub fn blob_epsilon(mut self, blob_epsilon: f32) -> Self { + self.blob_epsilon = blob_epsilon; + self + } + + pub fn result_expansion(mut self, result_expansion: f32) -> Self { + self.result_expansion = result_expansion; + self + } + + pub fn number_of_explored_blobs(mut self, number_of_explored_blobs: u64) -> Self { + self.number_of_explored_blobs = number_of_explored_blobs; + self + } + + pub fn number_of_edges(mut self, number_of_edges: u64) -> Self { + self.number_of_edges = number_of_edges; + self + } + + pub fn radius(mut self, radius: f32) -> Self { + self.radius = radius; + self + } + + unsafe fn into_raw(self) -> sys::QBGQuery { + sys::QBGQuery { + query: self.query.as_ptr() as *mut f32, + number_of_results: self.size, + epsilon: self.epsilon, + blob_epsilon: self.blob_epsilon, + result_expansion: self.result_expansion, + number_of_explored_blobs: self.number_of_explored_blobs, + number_of_edges: self.number_of_edges, + radius: self.radius, + } + } +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + use std::iter::repeat; + use std::result::Result as StdResult; + + use tempfile::tempdir; + + use super::*; + + #[test] + fn test_qbg() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + std::fs::remove_dir(dir.path())?; + + // Create a QGB index + let ndims = 3; + let mut index = + QbgIndex::create(dir.path(), QbgConstructParams::default().dimension(ndims))?; + + // Insert vectors and get their ids + let nvecs = 16; + let ids = (1..ndims * nvecs) + .step_by(ndims as usize) + .map(|i| i as f32) + .map(|i| { + repeat(i) + .zip((0..ndims).map(|j| j as f32)) + .map(|(i, j)| i + j) + .collect() + }) + .map(|vector| index.insert(vector)) + .collect::>>()?; + + // Build and persist the index + index.build(QbgBuildParams::default())?; + index.persist()?; + + let index = index.into_readable()?; + + // Perform a vector search (with 2 results) + let v: Vec = (1..=ndims).into_iter().map(|x| x as f32).collect(); + let query = QbgQuery::new(&v).size(2); + let res = index.search(query)?; + assert_eq!(ids[0], res[0].id); + assert_eq!(v, index.get_vec(ids[0])?); + + dir.close()?; + Ok(()) + } +} diff --git a/src/qbg/mod.rs b/src/qbg/mod.rs new file mode 100644 index 0000000..272025f --- /dev/null +++ b/src/qbg/mod.rs @@ -0,0 +1,7 @@ +mod index; +mod properties; + +pub use self::index::{ + IndexMode, ModeRead, ModeWrite, QbgBuildParams, QbgConstructParams, QbgIndex, QbgQuery, +}; +pub use self::properties::{QbgDistance, QbgObject}; diff --git a/src/qbg/properties.rs b/src/qbg/properties.rs new file mode 100644 index 0000000..2a489a4 --- /dev/null +++ b/src/qbg/properties.rs @@ -0,0 +1,255 @@ +// use std::ptr; + +// use ngt_sys as sys; +use num_enum::TryFromPrimitive; +// use scopeguard::defer; + +// use crate::error::{make_err, Result}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgObject { + Uint8 = 0, + Float = 1, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgDistance { + L2 = 1, +} + +// TODO: move these to mod.rs ? + +// #[derive(Debug)] +// pub struct QbgProperties { +// pub(crate) dimension: i32, +// pub(crate) creation_edge_size: i16, +// pub(crate) search_edge_size: i16, +// pub(crate) object_type: QbgObject, +// pub(crate) distance_type: QbgDistance, +// pub(crate) raw_prop: sys::NGTProperty, +// } + +// unsafe impl Send for QbgProperties {} +// unsafe impl Sync for QbgProperties {} + +// impl QbgProperties { +// pub fn dimension(dimension: usize) -> Result { +// let dimension = i32::try_from(dimension)?; +// let creation_edge_size = 10; +// let search_edge_size = 40; +// let object_type = QbgObject::Float; +// let distance_type = QbgDistance::L2; + +// unsafe { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// let raw_prop = sys::ngt_create_property(ebuf); +// if raw_prop.is_null() { +// Err(make_err(ebuf))? +// } + +// Self::set_dimension(raw_prop, dimension)?; +// Self::set_creation_edge_size(raw_prop, creation_edge_size)?; +// Self::set_search_edge_size(raw_prop, search_edge_size)?; +// Self::set_object_type(raw_prop, object_type)?; +// Self::set_distance_type(raw_prop, distance_type)?; + +// Ok(Self { +// dimension, +// creation_edge_size, +// search_edge_size, +// object_type, +// distance_type, +// raw_prop, +// }) +// } +// } + +// pub fn try_clone(&self) -> Result { +// unsafe { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// let raw_prop = sys::ngt_create_property(ebuf); +// if raw_prop.is_null() { +// Err(make_err(ebuf))? +// } + +// Self::set_dimension(raw_prop, self.dimension)?; +// Self::set_creation_edge_size(raw_prop, self.creation_edge_size)?; +// Self::set_search_edge_size(raw_prop, self.search_edge_size)?; +// Self::set_object_type(raw_prop, self.object_type)?; +// Self::set_distance_type(raw_prop, self.distance_type)?; + +// Ok(Self { +// dimension: self.dimension, +// creation_edge_size: self.creation_edge_size, +// search_edge_size: self.search_edge_size, +// object_type: self.object_type, +// distance_type: self.distance_type, +// raw_prop, +// }) +// } +// } + +// pub(crate) fn from(index: sys::NGTIndex) -> Result { +// unsafe { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// let raw_prop = sys::ngt_create_property(ebuf); +// if raw_prop.is_null() { +// Err(make_err(ebuf))? +// } + +// if !sys::ngt_get_property(index, raw_prop, ebuf) { +// Err(make_err(ebuf))? +// } + +// let dimension = sys::ngt_get_property_dimension(raw_prop, ebuf); +// if dimension < 0 { +// Err(make_err(ebuf))? +// } + +// let creation_edge_size = sys::ngt_get_property_edge_size_for_creation(raw_prop, ebuf); +// if creation_edge_size < 0 { +// Err(make_err(ebuf))? +// } + +// let search_edge_size = sys::ngt_get_property_edge_size_for_search(raw_prop, ebuf); +// if search_edge_size < 0 { +// Err(make_err(ebuf))? +// } + +// let object_type = sys::ngt_get_property_object_type(raw_prop, ebuf); +// if object_type < 0 { +// Err(make_err(ebuf))? +// } +// let object_type = QbgObject::try_from(object_type)?; + +// let distance_type = sys::ngt_get_property_distance_type(raw_prop, ebuf); +// if distance_type < 0 { +// Err(make_err(ebuf))? +// } +// let distance_type = QbgDistance::try_from(distance_type)?; + +// Ok(Self { +// dimension, +// creation_edge_size, +// search_edge_size, +// object_type, +// distance_type, +// raw_prop, +// }) +// } +// } + +// unsafe fn set_dimension(raw_prop: sys::NGTProperty, dimension: i32) -> Result<()> { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// if !sys::ngt_set_property_dimension(raw_prop, dimension, ebuf) { +// Err(make_err(ebuf))? +// } + +// Ok(()) +// } + +// pub fn creation_edge_size(mut self, size: usize) -> Result { +// let size = i16::try_from(size)?; +// self.creation_edge_size = size; +// unsafe { Self::set_creation_edge_size(self.raw_prop, size)? }; +// Ok(self) +// } + +// unsafe fn set_creation_edge_size(raw_prop: sys::NGTProperty, size: i16) -> Result<()> { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// if !sys::ngt_set_property_edge_size_for_creation(raw_prop, size, ebuf) { +// Err(make_err(ebuf))? +// } + +// Ok(()) +// } + +// pub fn search_edge_size(mut self, size: usize) -> Result { +// let size = i16::try_from(size)?; +// self.search_edge_size = size; +// unsafe { Self::set_search_edge_size(self.raw_prop, size)? }; +// Ok(self) +// } + +// unsafe fn set_search_edge_size(raw_prop: sys::NGTProperty, size: i16) -> Result<()> { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// if !sys::ngt_set_property_edge_size_for_search(raw_prop, size, ebuf) { +// Err(make_err(ebuf))? +// } + +// Ok(()) +// } + +// pub fn object_type(mut self, object_type: QbgObject) -> Result { +// self.object_type = object_type; +// unsafe { Self::set_object_type(self.raw_prop, object_type)? }; +// Ok(self) +// } + +// unsafe fn set_object_type(raw_prop: sys::NGTProperty, object_type: QbgObject) -> Result<()> { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// match object_type { +// QbgObject::Uint8 => { +// if !sys::ngt_set_property_object_type_integer(raw_prop, ebuf) { +// Err(make_err(ebuf))? +// } +// } +// QbgObject::Float => { +// if !sys::ngt_set_property_object_type_float(raw_prop, ebuf) { +// Err(make_err(ebuf))? +// } +// } +// } + +// Ok(()) +// } + +// pub fn distance_type(mut self, distance_type: QbgDistance) -> Result { +// self.distance_type = distance_type; +// unsafe { Self::set_distance_type(self.raw_prop, distance_type)? }; +// Ok(self) +// } + +// unsafe fn set_distance_type( +// raw_prop: sys::NGTProperty, +// distance_type: QbgDistance, +// ) -> Result<()> { +// let ebuf = sys::ngt_create_error_object(); +// defer! { sys::ngt_destroy_error_object(ebuf); } + +// match distance_type { +// QbgDistance::L2 => { +// if !sys::ngt_set_property_distance_type_l2(raw_prop, ebuf) { +// Err(make_err(ebuf))? +// } +// } +// } + +// Ok(()) +// } +// } + +// impl Drop for QbgProperties { +// fn drop(&mut self) { +// if !self.raw_prop.is_null() { +// unsafe { sys::ngt_destroy_property(self.raw_prop) }; +// self.raw_prop = ptr::null_mut(); +// } +// } +// } diff --git a/src/qg/index.rs b/src/qg/index.rs new file mode 100644 index 0000000..08eb9c4 --- /dev/null +++ b/src/qg/index.rs @@ -0,0 +1,305 @@ +use std::ffi::CString; +use std::mem; +use std::os::unix::ffi::OsStrExt; +use std::path::Path; +use std::ptr; + +use ngt_sys as sys; +use scopeguard::defer; + +use super::{QgObject, QgProperties}; +use crate::error::{make_err, Error, Result}; +use crate::ngt::NgtIndex; +use crate::{SearchResult, VecId}; + +#[derive(Debug)] +pub struct QgIndex { + pub(crate) prop: QgProperties, + pub(crate) index: sys::NGTQGIndex, + ebuf: sys::NGTError, +} + +impl QgIndex { + pub fn quantize(index: NgtIndex, params: QgParams) -> Result { + if !is_x86_feature_detected!("avx2") { + return Err(Error( + "Cannot quantize an index without AVX2 support".into(), + )); + } + + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let path = index.path.clone(); + drop(index); // Close the index + if !sys::ngtqg_quantize(path.as_ptr(), params.into_raw(), ebuf) { + Err(make_err(ebuf))? + } + + QgIndex::open(path.into_string()?) + } + } + + /// Open the already existing quantized index at the specified path. + pub fn open>(path: P) -> Result { + if !is_x86_feature_detected!("avx2") { + return Err(Error( + "Cannot use a quantized index without AVX2 support".into(), + )); + } + + if !path.as_ref().exists() { + Err(Error(format!("Path {:?} does not exist", path.as_ref())))? + } + + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let path = CString::new(path.as_ref().as_os_str().as_bytes())?; + + let index = sys::ngtqg_open_index(path.as_ptr(), ebuf); + if index.is_null() { + Err(make_err(ebuf))? + } + + let prop = QgProperties::from(index)?; + + Ok(QgIndex { + prop, + index, + ebuf: sys::ngt_create_error_object(), + }) + } + } + + pub fn search(&self, query: QgQuery) -> Result> { + unsafe { + let results = sys::ngt_create_empty_results(self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + defer! { sys::ngt_destroy_results(results); } + + if !sys::ngtqg_search_index(self.index, query.into_raw(), results, self.ebuf) { + Err(make_err(self.ebuf))? + } + + let rsize = sys::ngt_get_result_size(results, self.ebuf); + let mut ret = Vec::with_capacity(rsize as usize); + + for i in 0..rsize as u32 { + let d = sys::ngt_get_result(results, i, self.ebuf); + if d.id == 0 && d.distance == 0.0 { + Err(make_err(self.ebuf))? + } else { + ret.push(SearchResult { + id: d.id, + distance: d.distance, + }); + } + } + + Ok(ret) + } + } + + /// Get the specified vector. + pub fn get_vec(&self, id: VecId) -> Result> { + unsafe { + let results = match self.prop.object_type { + QgObject::Float => { + let ospace = sys::ngt_get_object_space(self.index, self.ebuf); + if ospace.is_null() { + Err(make_err(self.ebuf))? + } + + let results = sys::ngt_get_object_as_float(ospace, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + + let results = Vec::from_raw_parts( + results as *mut f32, + self.prop.dimension as usize, + self.prop.dimension as usize, + ); + let results = mem::ManuallyDrop::new(results); + + results.iter().copied().collect::>() + } + QgObject::Uint8 => { + let ospace = sys::ngt_get_object_space(self.index, self.ebuf); + if ospace.is_null() { + Err(make_err(self.ebuf))? + } + + let results = sys::ngt_get_object_as_integer(ospace, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + + let results = Vec::from_raw_parts( + results as *mut u8, + self.prop.dimension as usize, + self.prop.dimension as usize, + ); + let results = mem::ManuallyDrop::new(results); + + results.iter().map(|byte| *byte as f32).collect::>() + } + }; + + Ok(results) + } + } +} + +impl Drop for QgIndex { + fn drop(&mut self) { + if !self.index.is_null() { + unsafe { sys::ngtqg_close_index(self.index) }; + self.index = ptr::null_mut(); + } + if !self.ebuf.is_null() { + unsafe { sys::ngt_destroy_error_object(self.ebuf) }; + self.ebuf = ptr::null_mut(); + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct QgParams { + pub dimension_of_subvector: f32, + pub max_number_of_edges: u64, +} + +impl Default for QgParams { + fn default() -> Self { + Self { + dimension_of_subvector: 0.0, + max_number_of_edges: 128, + } + } +} + +impl QgParams { + unsafe fn into_raw(self) -> sys::NGTQGQuantizationParameters { + sys::NGTQGQuantizationParameters { + dimension_of_subvector: self.dimension_of_subvector, + max_number_of_edges: self.max_number_of_edges, + } + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct QgQuery<'a> { + query: &'a [f32], + pub size: u64, + pub epsilon: f32, + pub result_expansion: f32, + pub radius: f32, +} + +impl<'a> QgQuery<'a> { + pub fn new(query: &'a [f32]) -> Self { + Self { + query, + size: 20, + epsilon: 0.03, + result_expansion: 3.0, + radius: f32::MAX, + } + } + + pub fn size(mut self, size: u64) -> Self { + self.size = size; + self + } + + pub fn epsilon(mut self, epsilon: f32) -> Self { + self.epsilon = epsilon; + self + } + + pub fn result_expansion(mut self, result_expansion: f32) -> Self { + self.result_expansion = result_expansion; + self + } + + pub fn radius(mut self, radius: f32) -> Self { + self.radius = radius; + self + } + + unsafe fn into_raw(self) -> sys::NGTQGQuery { + sys::NGTQGQuery { + query: self.query.as_ptr() as *mut f32, + size: self.size, + epsilon: self.epsilon, + result_expansion: self.result_expansion, + radius: self.radius, + } + } +} + +#[cfg(test)] +mod tests { + use std::error::Error as StdError; + use std::iter::repeat; + use std::result::Result as StdResult; + + use tempfile::tempdir; + + use super::*; + use crate::{NgtDistance, NgtObject, NgtProperties}; + + #[test] + fn test_qg() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + + // Create an NGT index for vectors + let ndims = 3; + let props = NgtProperties::dimension(ndims)? + .object_type(NgtObject::Uint8)? + .distance_type(NgtDistance::L2)?; + let mut index = NgtIndex::create(dir.path(), props)?; + + // Insert vectors and get their ids + let nvecs = 16; + let ids = (1..ndims * nvecs) + .step_by(ndims) + .map(|i| i as f32) + .map(|i| { + repeat(i) + .zip((0..ndims).map(|j| j as f32)) + .map(|(i, j)| i + j) + .collect() + }) + .map(|vector| index.insert(vector)) + .collect::>>()?; + + // Build and persist the index + index.build(1)?; + index.persist()?; + + // Quantize the index + let params = QgParams { + dimension_of_subvector: 1., + max_number_of_edges: 50, + }; + let index = QgIndex::quantize(index, params)?; + + // Perform a vector search (with 2 results) + let v: Vec = (1..=ndims).into_iter().map(|x| x as f32).collect(); + let query = QgQuery::new(&v).size(2); + let res = index.search(query)?; + assert_eq!(ids[0], res[0].id); + assert_eq!(v, index.get_vec(ids[0])?); + + dir.close()?; + Ok(()) + } +} diff --git a/src/qg/mod.rs b/src/qg/mod.rs new file mode 100644 index 0000000..c5a0e4f --- /dev/null +++ b/src/qg/mod.rs @@ -0,0 +1,5 @@ +mod index; +mod properties; + +pub use self::index::{QgIndex, QgParams, QgQuery}; +pub use self::properties::{QgDistance, QgObject, QgProperties}; diff --git a/src/qg/properties.rs b/src/qg/properties.rs new file mode 100644 index 0000000..8112080 --- /dev/null +++ b/src/qg/properties.rs @@ -0,0 +1,259 @@ +use std::ptr; + +use ngt_sys as sys; +use num_enum::TryFromPrimitive; +use scopeguard::defer; + +use crate::error::{make_err, Result}; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QgObject { + Uint8 = 1, + Float = 2, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QgDistance { + L2 = 1, + Cosine = 4, +} + +#[derive(Debug)] +pub struct QgProperties { + pub(crate) dimension: i32, + pub(crate) creation_edge_size: i16, + pub(crate) search_edge_size: i16, + pub(crate) object_type: QgObject, + pub(crate) distance_type: QgDistance, + pub(crate) raw_prop: sys::NGTProperty, +} + +unsafe impl Send for QgProperties {} +unsafe impl Sync for QgProperties {} + +impl QgProperties { + pub fn dimension(dimension: usize) -> Result { + let dimension = i32::try_from(dimension)?; + let creation_edge_size = 10; + let search_edge_size = 40; + let object_type = QgObject::Float; + let distance_type = QgDistance::L2; + + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let raw_prop = sys::ngt_create_property(ebuf); + if raw_prop.is_null() { + Err(make_err(ebuf))? + } + + Self::set_dimension(raw_prop, dimension)?; + Self::set_creation_edge_size(raw_prop, creation_edge_size)?; + Self::set_search_edge_size(raw_prop, search_edge_size)?; + Self::set_object_type(raw_prop, object_type)?; + Self::set_distance_type(raw_prop, distance_type)?; + + Ok(Self { + dimension, + creation_edge_size, + search_edge_size, + object_type, + distance_type, + raw_prop, + }) + } + } + + pub fn try_clone(&self) -> Result { + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let raw_prop = sys::ngt_create_property(ebuf); + if raw_prop.is_null() { + Err(make_err(ebuf))? + } + + Self::set_dimension(raw_prop, self.dimension)?; + Self::set_creation_edge_size(raw_prop, self.creation_edge_size)?; + Self::set_search_edge_size(raw_prop, self.search_edge_size)?; + Self::set_object_type(raw_prop, self.object_type)?; + Self::set_distance_type(raw_prop, self.distance_type)?; + + Ok(Self { + dimension: self.dimension, + creation_edge_size: self.creation_edge_size, + search_edge_size: self.search_edge_size, + object_type: self.object_type, + distance_type: self.distance_type, + raw_prop, + }) + } + } + + pub(crate) fn from(index: sys::NGTIndex) -> Result { + unsafe { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + let raw_prop = sys::ngt_create_property(ebuf); + if raw_prop.is_null() { + Err(make_err(ebuf))? + } + + if !sys::ngt_get_property(index, raw_prop, ebuf) { + Err(make_err(ebuf))? + } + + let dimension = sys::ngt_get_property_dimension(raw_prop, ebuf); + if dimension < 0 { + Err(make_err(ebuf))? + } + + let creation_edge_size = sys::ngt_get_property_edge_size_for_creation(raw_prop, ebuf); + if creation_edge_size < 0 { + Err(make_err(ebuf))? + } + + let search_edge_size = sys::ngt_get_property_edge_size_for_search(raw_prop, ebuf); + if search_edge_size < 0 { + Err(make_err(ebuf))? + } + + let object_type = sys::ngt_get_property_object_type(raw_prop, ebuf); + if object_type < 0 { + Err(make_err(ebuf))? + } + let object_type = QgObject::try_from(object_type)?; + + let distance_type = sys::ngt_get_property_distance_type(raw_prop, ebuf); + if distance_type < 0 { + Err(make_err(ebuf))? + } + let distance_type = QgDistance::try_from(distance_type)?; + + Ok(Self { + dimension, + creation_edge_size, + search_edge_size, + object_type, + distance_type, + raw_prop, + }) + } + } + + unsafe fn set_dimension(raw_prop: sys::NGTProperty, dimension: i32) -> Result<()> { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + if !sys::ngt_set_property_dimension(raw_prop, dimension, ebuf) { + Err(make_err(ebuf))? + } + + Ok(()) + } + + pub fn creation_edge_size(mut self, size: usize) -> Result { + let size = i16::try_from(size)?; + self.creation_edge_size = size; + unsafe { Self::set_creation_edge_size(self.raw_prop, size)? }; + Ok(self) + } + + unsafe fn set_creation_edge_size(raw_prop: sys::NGTProperty, size: i16) -> Result<()> { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + if !sys::ngt_set_property_edge_size_for_creation(raw_prop, size, ebuf) { + Err(make_err(ebuf))? + } + + Ok(()) + } + + pub fn search_edge_size(mut self, size: usize) -> Result { + let size = i16::try_from(size)?; + self.search_edge_size = size; + unsafe { Self::set_search_edge_size(self.raw_prop, size)? }; + Ok(self) + } + + unsafe fn set_search_edge_size(raw_prop: sys::NGTProperty, size: i16) -> Result<()> { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + if !sys::ngt_set_property_edge_size_for_search(raw_prop, size, ebuf) { + Err(make_err(ebuf))? + } + + Ok(()) + } + + pub fn object_type(mut self, object_type: QgObject) -> Result { + self.object_type = object_type; + unsafe { Self::set_object_type(self.raw_prop, object_type)? }; + Ok(self) + } + + unsafe fn set_object_type(raw_prop: sys::NGTProperty, object_type: QgObject) -> Result<()> { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + match object_type { + QgObject::Uint8 => { + if !sys::ngt_set_property_object_type_integer(raw_prop, ebuf) { + Err(make_err(ebuf))? + } + } + QgObject::Float => { + if !sys::ngt_set_property_object_type_float(raw_prop, ebuf) { + Err(make_err(ebuf))? + } + } + } + + Ok(()) + } + + pub fn distance_type(mut self, distance_type: QgDistance) -> Result { + self.distance_type = distance_type; + unsafe { Self::set_distance_type(self.raw_prop, distance_type)? }; + Ok(self) + } + + unsafe fn set_distance_type( + raw_prop: sys::NGTProperty, + distance_type: QgDistance, + ) -> Result<()> { + let ebuf = sys::ngt_create_error_object(); + defer! { sys::ngt_destroy_error_object(ebuf); } + + match distance_type { + QgDistance::L2 => { + if !sys::ngt_set_property_distance_type_l2(raw_prop, ebuf) { + Err(make_err(ebuf))? + } + } + QgDistance::Cosine => { + if !sys::ngt_set_property_distance_type_cosine(raw_prop, ebuf) { + Err(make_err(ebuf))? + } + } + } + + Ok(()) + } +} + +impl Drop for QgProperties { + fn drop(&mut self) { + if !self.raw_prop.is_null() { + unsafe { sys::ngt_destroy_property(self.raw_prop) }; + self.raw_prop = ptr::null_mut(); + } + } +} From dfdca9612dc5bdde29e4335dfae9568067e5fc96 Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Sun, 7 May 2023 00:14:30 +0200 Subject: [PATCH 03/11] Only support f32 vector (drop support for f64) --- src/lib.rs | 2 +- src/ngt/index.rs | 22 ++++++++-------------- 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 335ea5b..53f7365 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -38,5 +38,5 @@ pub const EPSILON: f32 = 0.1; pub use crate::error::{Error, Result}; pub use crate::ngt::{NgtDistance, NgtIndex, NgtObject, NgtProperties}; -// TODO: search etc only f32, drop support for f64/Into // TODO: what about float16 ? +// TODO: add doc (link to official blog posts) diff --git a/src/ngt/index.rs b/src/ngt/index.rs index f51c71a..0c7ebf5 100644 --- a/src/ngt/index.rs +++ b/src/ngt/index.rs @@ -105,7 +105,7 @@ impl NgtIndex { /// Search the nearest vectors to the specified query vector. /// /// **The index must have been [`built`](NgtIndex::build) beforehand**. - pub fn search(&self, vec: &[f64], res_size: u64, epsilon: f32) -> Result> { + pub fn search(&self, vec: &[f32], res_size: u64, epsilon: f32) -> Result> { unsafe { let results = sys::ngt_create_empty_results(self.ebuf); if results.is_null() { @@ -113,9 +113,9 @@ impl NgtIndex { } defer! { sys::ngt_destroy_results(results); } - if !sys::ngt_search_index( + if !sys::ngt_search_index_as_float( self.index, - vec.as_ptr() as *mut f64, + vec.as_ptr() as *mut f32, self.prop.dimension, res_size, epsilon, @@ -148,7 +148,7 @@ impl NgtIndex { /// Search linearly the nearest vectors to the specified query vector. /// /// **The index must have been [`built`](NgtIndex::build) beforehand**. - pub fn linear_search(&self, vec: &[f64], res_size: u64) -> Result> { + pub fn linear_search(&self, vec: &[f32], res_size: u64) -> Result> { unsafe { let results = sys::ngt_create_empty_results(self.ebuf); if results.is_null() { @@ -156,9 +156,9 @@ impl NgtIndex { } defer! { sys::ngt_destroy_results(results); } - if !sys::ngt_linear_search_index( + if !sys::ngt_linear_search_index_as_float( self.index, - vec.as_ptr() as *mut f64, + vec.as_ptr() as *mut f32, self.prop.dimension, res_size, results, @@ -210,7 +210,7 @@ impl NgtIndex { /// discoverable yet. /// /// **The method [`build`](NgtIndex::build) must be called after inserting vectors**. - pub fn insert_batch>(&mut self, batch: Vec>) -> Result<()> { + pub fn insert_batch(&mut self, batch: Vec>) -> Result<()> { let batch_size = u32::try_from(batch.len())?; if batch_size > 0 { @@ -226,16 +226,10 @@ impl NgtIndex { } unsafe { - let mut batch = batch - .into_iter() - .flatten() - .map(|v| v.into() as f32) - .collect::>(); - + let mut batch = batch.into_iter().flatten().collect::>(); if !sys::ngt_batch_append_index(self.index, batch.as_mut_ptr(), batch_size, self.ebuf) { Err(make_err(self.ebuf))? } - Ok(()) } } From fbf19f684b4104d6f630850942e450b382c8695b Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Sun, 7 May 2023 00:27:20 +0200 Subject: [PATCH 04/11] Adapt NGT optim --- src/lib.rs | 1 + src/ngt/mod.rs | 2 +- src/ngt/optim.rs | 50 +++++++++++++++++++++++------------------------- 3 files changed, 26 insertions(+), 27 deletions(-) diff --git a/src/lib.rs b/src/lib.rs index 53f7365..b019078 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -36,6 +36,7 @@ pub struct SearchResult { pub const EPSILON: f32 = 0.1; pub use crate::error::{Error, Result}; +pub use crate::ngt::optim; pub use crate::ngt::{NgtDistance, NgtIndex, NgtObject, NgtProperties}; // TODO: what about float16 ? diff --git a/src/ngt/mod.rs b/src/ngt/mod.rs index 67cc417..6a6e7b9 100644 --- a/src/ngt/mod.rs +++ b/src/ngt/mod.rs @@ -65,7 +65,7 @@ //! ``` mod index; -// mod optim; +pub mod optim; mod properties; pub use self::index::NgtIndex; diff --git a/src/ngt/optim.rs b/src/ngt/optim.rs index ddde83b..358033e 100644 --- a/src/ngt/optim.rs +++ b/src/ngt/optim.rs @@ -7,18 +7,18 @@ use ngt_sys as sys; use scopeguard::defer; use crate::error::{make_err, Result}; -use crate::index::Index; +use crate::ngt::index::NgtIndex; /// Optimizes the number of initial edges of an ANNG index. /// /// The default number of initial edges for each node in a default graph (ANNG) is a /// fixed value 10. To optimize this number, follow these steps: -/// 1. [`insert`](Index::insert) vectors in the ANNG index at `index_path`, don't -/// [`build`](Index::build) the index yet. -/// 2. When all vectors are inserted, [`persist`](Index::persist) the index. +/// 1. [`insert`](NgtIndex::insert) vectors in the ANNG index at `index_path`, don't +/// [`build`](NgtIndex::build) the index yet. +/// 2. When all vectors are inserted, [`persist`](NgtIndex::persist) the index. /// 3. Call this function with the same `index_path`. -/// 4. [`open`](Index::open) the index at `index_path` again, and now -/// [`build`](Index::build) it. +/// 4. [`open`](NgtIndex::open) the index at `index_path` again, and now +/// [`build`](NgtIndex::build) it. #[cfg(not(feature = "shared_mem"))] pub fn optimize_anng_edges_number>( index_path: P, @@ -53,9 +53,9 @@ pub fn optimize_anng_search_parameters>(index_path: P) -> Result< /// /// Improves accuracy of neighboring nodes for each node by searching with each /// node. Note that refinement takes a long processing time. An ANNG index can be -/// refined only after it has been [`built`](Index::build). +/// refined only after it has been [`built`](NgtIndex::build). #[cfg(not(feature = "shared_mem"))] -pub fn refine_anng(index: &mut Index, params: AnngRefineParams) -> Result<()> { +pub fn refine_anng(index: &mut NgtIndex, params: AnngRefineParams) -> Result<()> { unsafe { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } @@ -85,7 +85,7 @@ pub fn refine_anng(index: &mut Index, params: AnngRefineParams) -> Result<()> { /// /// If more performance is needed, a larger `creation_edge_size` can be set through /// [`Properties`](crate::Properties::creation_edge_size) at ANNG index -/// [`create`](Index::create) time. +/// [`create`](NgtIndex::create) time. /// /// Important [`GraphOptimParams`](GraphOptimParams) parameters are `nb_outgoing` edges /// and `nb_incoming` edges. The latter can be set to an even higher number than the @@ -254,7 +254,7 @@ impl GraphOptimizer { /// Optimize for the search parameters of an ANNG. fn adjust_search_coefficients>(&mut self, index_path: P) -> Result<()> { - let _ = Index::open(&index_path)?; + let _ = NgtIndex::open(&index_path)?; unsafe { let ebuf = sys::ngt_create_error_object(); @@ -276,7 +276,7 @@ impl GraphOptimizer { index_anng_in: P, index_onng_out: P, ) -> Result<()> { - let _ = Index::open(&index_anng_in)?; + let _ = NgtIndex::open(&index_anng_in)?; unsafe { let ebuf = sys::ngt_create_error_object(); @@ -310,9 +310,7 @@ mod tests { use tempfile::tempdir; - use crate::{DistanceType, Index, Properties}; - - use super::*; + use crate::{ngt::optim::*, ngt::*}; #[ignore] #[test] @@ -322,12 +320,12 @@ mod tests { let dir = tempdir()?; // Create an index for vectors of dimension 3 with cosine distance - let prop = Properties::dimension(3)?.distance_type(DistanceType::Cosine)?; - let mut index = Index::create(dir.path(), prop)?; + let prop = NgtProperties::dimension(3)?.distance_type(NgtDistance::Cosine)?; + let mut index = NgtIndex::create(dir.path(), prop)?; // Populate the index, but don't build it yet for i in 0..1_000_000 { - let _ = index.insert(vec![i, i + 1, i + 2])?; + let _ = index.insert(vec![i as f32, i as f32 + 1.0, i as f32 + 2.0])?; } index.persist()?; @@ -335,7 +333,7 @@ mod tests { optimize_anng_edges_number(dir.path(), AnngEdgeOptimParams::default())?; // Now build and persist again the optimized index - let mut index = Index::open(dir.path())?; + let mut index = NgtIndex::open(dir.path())?; index.build(4)?; index.persist()?; @@ -353,12 +351,12 @@ mod tests { let dir = tempdir()?; // Create an index for vectors of dimension 3 with cosine distance - let prop = Properties::dimension(3)?.distance_type(DistanceType::Cosine)?; - let mut index = Index::create(dir.path(), prop)?; + let prop = NgtProperties::dimension(3)?.distance_type(NgtDistance::Cosine)?; + let mut index = NgtIndex::create(dir.path(), prop)?; // Populate and build the index for i in 0..1000 { - let _ = index.insert(vec![i, i + 1, i + 2])?; + let _ = index.insert(vec![i as f32, i as f32 + 1.0, i as f32 + 2.0])?; } index.build(4)?; @@ -377,15 +375,15 @@ mod tests { let dir_in = tempdir()?; // Create an index for vectors of dimension 3 with cosine distance - let prop = Properties::dimension(3)? - .distance_type(DistanceType::Cosine)? + let prop = NgtProperties::dimension(3)? + .distance_type(NgtDistance::Cosine)? .creation_edge_size(100)?; // More than default value, improves the final ONNG - let mut index = Index::create(dir_in.path(), prop)?; + let mut index = NgtIndex::create(dir_in.path(), prop)?; // Populate and persist (but don't build yet) the index for i in 0..1000 { - let _ = index.insert(vec![i, i + 1, i + 2])?; + let _ = index.insert(vec![i as f32, i as f32 + 1.0, i as f32 + 2.0])?; } index.persist()?; @@ -393,7 +391,7 @@ mod tests { optimize_anng_edges_number(dir_in.path(), AnngEdgeOptimParams::default())?; // Now build and persist again the optimized index - let mut index = Index::open(dir_in.path())?; + let mut index = NgtIndex::open(dir_in.path())?; index.build(4)?; index.persist()?; From a90b66ba59f6fa2bd7de2e944675077274059972 Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Mon, 8 May 2023 18:37:46 +0200 Subject: [PATCH 05/11] Update doc, some refactoring --- README.md | 61 ++++++---- ngt-sys/build.rs | 7 +- src/lib.rs | 25 +---- src/ngt/optim.rs | 9 +- src/qbg/index.rs | 3 +- src/qbg/mod.rs | 17 ++- src/qbg/properties.rs | 255 ------------------------------------------ src/qg/index.rs | 11 +- src/qg/mod.rs | 2 +- 9 files changed, 81 insertions(+), 309 deletions(-) delete mode 100644 src/qbg/properties.rs diff --git a/README.md b/README.md index a34b8c2..2bf6522 100644 --- a/README.md +++ b/README.md @@ -6,46 +6,61 @@ [docs.rs]: https://docs.rs/ngt Rust wrappers for [NGT][], which provides high-speed approximate nearest neighbor -searches against a large volume of data. +searches against a large volume of data in high dimensional vector data space (several +ten to several thousand dimensions). -Building NGT requires `CMake`. By default `ngt-rs` will be built dynamically, which -means that you'll need to make the build artifact `libngt.so` available to your final -binary. You'll also need to have `OpenMP` installed on the system where it will run. If -you want to build `ngt-rs` statically, then use the `static` Cargo feature, note that in -this case `OpenMP` will be linked statically too. +This crate provides the following indexes: +* `NgtIndex`: Graph and tree-based index[^1] +* `QqIndex`: Quantized graph-based index[^2] +* `QbgIndex`: Quantized blob graph-based index -Furthermore, NGT's shared memory and large dataset features are available through Cargo -features `shared_mem` and `large_data` respectively. +The quantized indexes are available through the `quantized` Cargo feature. Note that +they rely on `BLAS` and `LAPACK` which thus have to be installed locally. The CPU +running the code must also support `AVX2` instructions. + +The `NgtIndex` default implementation is an ANNG, it can be optimized[^3] or converted +to an ONNG through the [`optim`][ngt-optim] module. + +By default `ngt-rs` will be built dynamically, which requires `CMake` to build NGT. This +means that you'll have to make the build artifact `libngt.so` available to your final +binary (see an example in the [CI][ngt-ci]). + +However the `static` feature will build and link NGT statically. Note that `OpenMP` will +also be linked statically. If the `quantized` feature is used, then `BLAS` and `LAPACK` +libraries will also be linked statically. + +Finally, NGT's [shared memory][ngt-sharedmem] and [large dataset][ngt-largedata] +features are available through the features `shared_mem` and `large_data` respectively. ## Usage Defining the properties of a new index: -```rust -use ngt::{Properties, DistanceType, ObjectType}; +```rust,ignore +use ngt::{NgtProperties, NgtDistance, NgtObject}; // Defaut properties with vectors of dimension 3 -let prop = Properties::dimension(3)?; +let prop = NgtProperties::dimension(3)?; // Or customize values (here are the defaults) -let prop = Properties::dimension(3)? +let prop = NgtProperties::dimension(3)? .creation_edge_size(10)? .search_edge_size(40)? - .object_type(ObjectType::Float)? - .distance_type(DistanceType::L2)?; + .object_type(NgtObject::Float)? + .distance_type(NgtDistance::L2)?; ``` Creating/Opening an index and using it: -```rust -use ngt::{Index, Properties, EPSILON}; +```rust,ignore +use ngt::{NgtIndex, NgtProperties, EPSILON}; // Create a new index -let prop = Properties::dimension(3)?; -let index = Index::create("target/path/to/index/dir", prop)?; +let prop = NgtProperties::dimension(3)?; +let index = NgtIndex::create("target/path/to/index/dir", prop)?; // Open an existing index -let mut index = Index::open("target/path/to/index/dir")?; +let mut index = NgtIndex::open("target/path/to/index/dir")?; // Insert two vectors and get their id let vec1 = vec![1.0, 2.0, 3.0]; @@ -77,3 +92,11 @@ index.persist()?; ``` [ngt]: https://github.com/yahoojapan/NGT +[ngt-sharedmem]: https://github.com/yahoojapan/NGT#shared-memory-use +[ngt-largedata]: https://github.com/yahoojapan/NGT#large-scale-data-use +[ngt-ci]: https://github.com/lerouxrgd/ngt-rs/blob/master/.github/workflows/ci.yaml +[ngt-optim]: https://docs.rs/ngt/latest/ngt/optim/index.html + +[^1]: https://opensource.com/article/19/10/ngt-open-source-library +[^2]: https://medium.com/@masajiro.iwasaki/fusion-of-graph-based-indexing-and-product-quantization-for-ann-search-7d1f0336d0d0 +[^3]: https://github.com/yahoojapan/NGT/wiki/Optimization-Examples-Using-Python diff --git a/ngt-sys/build.rs b/ngt-sys/build.rs index 5fa1a6f..8e427fa 100644 --- a/ngt-sys/build.rs +++ b/ngt-sys/build.rs @@ -31,8 +31,11 @@ fn main() { .build("src/lib.rs"); println!("cargo:rustc-link-lib=static=ngt"); println!("cargo:rustc-link-lib=gomp"); - println!("cargo:rustc-link-lib=blas"); - println!("cargo:rustc-link-lib=lapack"); + + if env::var("CARGO_FEATURE_QUANTIZED").is_ok() { + println!("cargo:rustc-link-lib=blas"); + println!("cargo:rustc-link-lib=lapack"); + } } let capi_header = if cfg!(feature = "quantized") { diff --git a/src/lib.rs b/src/lib.rs index b019078..393c267 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,29 +1,14 @@ -//! Rust wrappers for [NGT][], which provides high-speed approximate nearest neighbor -//! searches against a large volume of data. -//! -//! Building NGT requires `CMake`. By default `ngt-rs` will be built dynamically, which -//! means that you'll need to make the build artifact `libngt.so` available to your final -//! binary. You'll also need to have `OpenMP` installed on the system where it will run. If -//! you want to build `ngt-rs` statically, then use the `static` Cargo feature. -//! -//! Furthermore, NGT's shared memory and large dataset features are available through Cargo -//! features `shared_mem` and `large_data` respectively. -//! -//! [ngt]: https://github.com/yahoojapan/NGT - -// TODO: consider include_str README +#![doc = include_str!("../README.md")] #[cfg(all(feature = "quantized", feature = "shared_mem"))] compile_error!("only one of ['quantized', 'shared_mem'] can be enabled"); mod error; mod ngt; - -#[cfg(feature = "quantized")] -pub mod qg; - #[cfg(feature = "quantized")] pub mod qbg; +#[cfg(feature = "quantized")] +pub mod qg; pub type VecId = u32; @@ -36,8 +21,6 @@ pub struct SearchResult { pub const EPSILON: f32 = 0.1; pub use crate::error::{Error, Result}; -pub use crate::ngt::optim; -pub use crate::ngt::{NgtDistance, NgtIndex, NgtObject, NgtProperties}; +pub use crate::ngt::{optim, NgtDistance, NgtIndex, NgtObject, NgtProperties}; // TODO: what about float16 ? -// TODO: add doc (link to official blog posts) diff --git a/src/ngt/optim.rs b/src/ngt/optim.rs index 358033e..cd5da2b 100644 --- a/src/ngt/optim.rs +++ b/src/ngt/optim.rs @@ -325,7 +325,8 @@ mod tests { // Populate the index, but don't build it yet for i in 0..1_000_000 { - let _ = index.insert(vec![i as f32, i as f32 + 1.0, i as f32 + 2.0])?; + let i = i as f32; + let _ = index.insert(vec![i, i + 1.0, i + 2.0])?; } index.persist()?; @@ -356,7 +357,8 @@ mod tests { // Populate and build the index for i in 0..1000 { - let _ = index.insert(vec![i as f32, i as f32 + 1.0, i as f32 + 2.0])?; + let i = i as f32; + let _ = index.insert(vec![i, i + 1.0, i + 2.0])?; } index.build(4)?; @@ -383,7 +385,8 @@ mod tests { // Populate and persist (but don't build yet) the index for i in 0..1000 { - let _ = index.insert(vec![i as f32, i as f32 + 1.0, i as f32 + 2.0])?; + let i = i as f32; + let _ = index.insert(vec![i, i + 1.0, i + 2.0])?; } index.persist()?; diff --git a/src/qbg/index.rs b/src/qbg/index.rs index f57b12e..2edce50 100644 --- a/src/qbg/index.rs +++ b/src/qbg/index.rs @@ -7,10 +7,11 @@ use ngt_sys as sys; use num_enum::TryFromPrimitive; use scopeguard::defer; -use super::{QbgDistance, QbgObject}; use crate::error::{make_err, Error, Result}; use crate::{SearchResult, VecId}; +use super::{QbgDistance, QbgObject}; + #[derive(Debug)] pub struct QbgIndex { pub(crate) index: sys::QBGIndex, diff --git a/src/qbg/mod.rs b/src/qbg/mod.rs index 272025f..5f1059c 100644 --- a/src/qbg/mod.rs +++ b/src/qbg/mod.rs @@ -1,7 +1,20 @@ mod index; -mod properties; + +use num_enum::TryFromPrimitive; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgObject { + Uint8 = 0, + Float = 1, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgDistance { + L2 = 1, +} pub use self::index::{ IndexMode, ModeRead, ModeWrite, QbgBuildParams, QbgConstructParams, QbgIndex, QbgQuery, }; -pub use self::properties::{QbgDistance, QbgObject}; diff --git a/src/qbg/properties.rs b/src/qbg/properties.rs deleted file mode 100644 index 2a489a4..0000000 --- a/src/qbg/properties.rs +++ /dev/null @@ -1,255 +0,0 @@ -// use std::ptr; - -// use ngt_sys as sys; -use num_enum::TryFromPrimitive; -// use scopeguard::defer; - -// use crate::error::{make_err, Result}; - -#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] -#[repr(i32)] -pub enum QbgObject { - Uint8 = 0, - Float = 1, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] -#[repr(i32)] -pub enum QbgDistance { - L2 = 1, -} - -// TODO: move these to mod.rs ? - -// #[derive(Debug)] -// pub struct QbgProperties { -// pub(crate) dimension: i32, -// pub(crate) creation_edge_size: i16, -// pub(crate) search_edge_size: i16, -// pub(crate) object_type: QbgObject, -// pub(crate) distance_type: QbgDistance, -// pub(crate) raw_prop: sys::NGTProperty, -// } - -// unsafe impl Send for QbgProperties {} -// unsafe impl Sync for QbgProperties {} - -// impl QbgProperties { -// pub fn dimension(dimension: usize) -> Result { -// let dimension = i32::try_from(dimension)?; -// let creation_edge_size = 10; -// let search_edge_size = 40; -// let object_type = QbgObject::Float; -// let distance_type = QbgDistance::L2; - -// unsafe { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// let raw_prop = sys::ngt_create_property(ebuf); -// if raw_prop.is_null() { -// Err(make_err(ebuf))? -// } - -// Self::set_dimension(raw_prop, dimension)?; -// Self::set_creation_edge_size(raw_prop, creation_edge_size)?; -// Self::set_search_edge_size(raw_prop, search_edge_size)?; -// Self::set_object_type(raw_prop, object_type)?; -// Self::set_distance_type(raw_prop, distance_type)?; - -// Ok(Self { -// dimension, -// creation_edge_size, -// search_edge_size, -// object_type, -// distance_type, -// raw_prop, -// }) -// } -// } - -// pub fn try_clone(&self) -> Result { -// unsafe { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// let raw_prop = sys::ngt_create_property(ebuf); -// if raw_prop.is_null() { -// Err(make_err(ebuf))? -// } - -// Self::set_dimension(raw_prop, self.dimension)?; -// Self::set_creation_edge_size(raw_prop, self.creation_edge_size)?; -// Self::set_search_edge_size(raw_prop, self.search_edge_size)?; -// Self::set_object_type(raw_prop, self.object_type)?; -// Self::set_distance_type(raw_prop, self.distance_type)?; - -// Ok(Self { -// dimension: self.dimension, -// creation_edge_size: self.creation_edge_size, -// search_edge_size: self.search_edge_size, -// object_type: self.object_type, -// distance_type: self.distance_type, -// raw_prop, -// }) -// } -// } - -// pub(crate) fn from(index: sys::NGTIndex) -> Result { -// unsafe { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// let raw_prop = sys::ngt_create_property(ebuf); -// if raw_prop.is_null() { -// Err(make_err(ebuf))? -// } - -// if !sys::ngt_get_property(index, raw_prop, ebuf) { -// Err(make_err(ebuf))? -// } - -// let dimension = sys::ngt_get_property_dimension(raw_prop, ebuf); -// if dimension < 0 { -// Err(make_err(ebuf))? -// } - -// let creation_edge_size = sys::ngt_get_property_edge_size_for_creation(raw_prop, ebuf); -// if creation_edge_size < 0 { -// Err(make_err(ebuf))? -// } - -// let search_edge_size = sys::ngt_get_property_edge_size_for_search(raw_prop, ebuf); -// if search_edge_size < 0 { -// Err(make_err(ebuf))? -// } - -// let object_type = sys::ngt_get_property_object_type(raw_prop, ebuf); -// if object_type < 0 { -// Err(make_err(ebuf))? -// } -// let object_type = QbgObject::try_from(object_type)?; - -// let distance_type = sys::ngt_get_property_distance_type(raw_prop, ebuf); -// if distance_type < 0 { -// Err(make_err(ebuf))? -// } -// let distance_type = QbgDistance::try_from(distance_type)?; - -// Ok(Self { -// dimension, -// creation_edge_size, -// search_edge_size, -// object_type, -// distance_type, -// raw_prop, -// }) -// } -// } - -// unsafe fn set_dimension(raw_prop: sys::NGTProperty, dimension: i32) -> Result<()> { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// if !sys::ngt_set_property_dimension(raw_prop, dimension, ebuf) { -// Err(make_err(ebuf))? -// } - -// Ok(()) -// } - -// pub fn creation_edge_size(mut self, size: usize) -> Result { -// let size = i16::try_from(size)?; -// self.creation_edge_size = size; -// unsafe { Self::set_creation_edge_size(self.raw_prop, size)? }; -// Ok(self) -// } - -// unsafe fn set_creation_edge_size(raw_prop: sys::NGTProperty, size: i16) -> Result<()> { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// if !sys::ngt_set_property_edge_size_for_creation(raw_prop, size, ebuf) { -// Err(make_err(ebuf))? -// } - -// Ok(()) -// } - -// pub fn search_edge_size(mut self, size: usize) -> Result { -// let size = i16::try_from(size)?; -// self.search_edge_size = size; -// unsafe { Self::set_search_edge_size(self.raw_prop, size)? }; -// Ok(self) -// } - -// unsafe fn set_search_edge_size(raw_prop: sys::NGTProperty, size: i16) -> Result<()> { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// if !sys::ngt_set_property_edge_size_for_search(raw_prop, size, ebuf) { -// Err(make_err(ebuf))? -// } - -// Ok(()) -// } - -// pub fn object_type(mut self, object_type: QbgObject) -> Result { -// self.object_type = object_type; -// unsafe { Self::set_object_type(self.raw_prop, object_type)? }; -// Ok(self) -// } - -// unsafe fn set_object_type(raw_prop: sys::NGTProperty, object_type: QbgObject) -> Result<()> { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// match object_type { -// QbgObject::Uint8 => { -// if !sys::ngt_set_property_object_type_integer(raw_prop, ebuf) { -// Err(make_err(ebuf))? -// } -// } -// QbgObject::Float => { -// if !sys::ngt_set_property_object_type_float(raw_prop, ebuf) { -// Err(make_err(ebuf))? -// } -// } -// } - -// Ok(()) -// } - -// pub fn distance_type(mut self, distance_type: QbgDistance) -> Result { -// self.distance_type = distance_type; -// unsafe { Self::set_distance_type(self.raw_prop, distance_type)? }; -// Ok(self) -// } - -// unsafe fn set_distance_type( -// raw_prop: sys::NGTProperty, -// distance_type: QbgDistance, -// ) -> Result<()> { -// let ebuf = sys::ngt_create_error_object(); -// defer! { sys::ngt_destroy_error_object(ebuf); } - -// match distance_type { -// QbgDistance::L2 => { -// if !sys::ngt_set_property_distance_type_l2(raw_prop, ebuf) { -// Err(make_err(ebuf))? -// } -// } -// } - -// Ok(()) -// } -// } - -// impl Drop for QbgProperties { -// fn drop(&mut self) { -// if !self.raw_prop.is_null() { -// unsafe { sys::ngt_destroy_property(self.raw_prop) }; -// self.raw_prop = ptr::null_mut(); -// } -// } -// } diff --git a/src/qg/index.rs b/src/qg/index.rs index 08eb9c4..64fb291 100644 --- a/src/qg/index.rs +++ b/src/qg/index.rs @@ -20,7 +20,8 @@ pub struct QgIndex { } impl QgIndex { - pub fn quantize(index: NgtIndex, params: QgParams) -> Result { + /// Quantize an NGT index + pub fn quantize(index: NgtIndex, params: QgQuantizationParams) -> Result { if !is_x86_feature_detected!("avx2") { return Err(Error( "Cannot quantize an index without AVX2 support".into(), @@ -170,12 +171,12 @@ impl Drop for QgIndex { } #[derive(Debug, Clone, PartialEq)] -pub struct QgParams { +pub struct QgQuantizationParams { pub dimension_of_subvector: f32, pub max_number_of_edges: u64, } -impl Default for QgParams { +impl Default for QgQuantizationParams { fn default() -> Self { Self { dimension_of_subvector: 0.0, @@ -184,7 +185,7 @@ impl Default for QgParams { } } -impl QgParams { +impl QgQuantizationParams { unsafe fn into_raw(self) -> sys::NGTQGQuantizationParameters { sys::NGTQGQuantizationParameters { dimension_of_subvector: self.dimension_of_subvector, @@ -286,7 +287,7 @@ mod tests { index.persist()?; // Quantize the index - let params = QgParams { + let params = QgQuantizationParams { dimension_of_subvector: 1., max_number_of_edges: 50, }; diff --git a/src/qg/mod.rs b/src/qg/mod.rs index c5a0e4f..21d170a 100644 --- a/src/qg/mod.rs +++ b/src/qg/mod.rs @@ -1,5 +1,5 @@ mod index; mod properties; -pub use self::index::{QgIndex, QgParams, QgQuery}; +pub use self::index::{QgIndex, QgQuantizationParams, QgQuery}; pub use self::properties::{QgDistance, QgObject, QgProperties}; From 4087002b4c73269a1de13d6f1b5f0f2fc5cb2b0e Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Thu, 11 May 2023 00:24:08 +0200 Subject: [PATCH 06/11] Update NGT to 2.0.11 --- Cargo.toml | 2 +- ngt-sys/Cargo.toml | 2 +- ngt-sys/NGT | 2 +- 3 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 7c7e66b..ce8bb10 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,7 +11,7 @@ license = "Apache-2.0" readme = "README.md" [dependencies] -ngt-sys = { path = "ngt-sys", version = "2.0.10" } +ngt-sys = { path = "ngt-sys", version = "2.0.11" } num_enum = "0.5" scopeguard = "1" diff --git a/ngt-sys/Cargo.toml b/ngt-sys/Cargo.toml index 06b1a58..08150ef 100644 --- a/ngt-sys/Cargo.toml +++ b/ngt-sys/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ngt-sys" -version = "2.0.10" +version = "2.0.11" authors = ["Romain Leroux "] edition = "2021" links = "ngt" diff --git a/ngt-sys/NGT b/ngt-sys/NGT index bc9b4dd..a91abde 160000 --- a/ngt-sys/NGT +++ b/ngt-sys/NGT @@ -1 +1 @@ -Subproject commit bc9b4dd9822c1d39ffdd3cbfbe822306c51171cc +Subproject commit a91abde328aa70faad467b89545e99965eee0b2a From 8f0ecb089ef2141e65a1abc017a191a346f1458d Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Thu, 11 May 2023 01:10:06 +0200 Subject: [PATCH 07/11] Use random vectors in optim test --- Cargo.toml | 1 + src/ngt/optim.rs | 19 ++++++++++--------- 2 files changed, 11 insertions(+), 9 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index ce8bb10..35f93b1 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -16,6 +16,7 @@ num_enum = "0.5" scopeguard = "1" [dev-dependencies] +rand = "0.8" rayon = "1" tempfile = "3" diff --git a/src/ngt/optim.rs b/src/ngt/optim.rs index cd5da2b..5233b3e 100644 --- a/src/ngt/optim.rs +++ b/src/ngt/optim.rs @@ -308,6 +308,7 @@ mod tests { use std::error::Error as StdError; use std::result::Result as StdResult; + use rand::Rng; use tempfile::tempdir; use crate::{ngt::optim::*, ngt::*}; @@ -324,9 +325,9 @@ mod tests { let mut index = NgtIndex::create(dir.path(), prop)?; // Populate the index, but don't build it yet - for i in 0..1_000_000 { - let i = i as f32; - let _ = index.insert(vec![i, i + 1.0, i + 2.0])?; + let mut rng = rand::thread_rng(); + for _ in 0..25_000 { + index.insert(vec![rng.gen(); 3])?; } index.persist()?; @@ -356,9 +357,9 @@ mod tests { let mut index = NgtIndex::create(dir.path(), prop)?; // Populate and build the index - for i in 0..1000 { - let i = i as f32; - let _ = index.insert(vec![i, i + 1.0, i + 2.0])?; + let mut rng = rand::thread_rng(); + for _ in 0..1000 { + index.insert(vec![rng.gen(); 3])?; } index.build(4)?; @@ -384,9 +385,9 @@ mod tests { let mut index = NgtIndex::create(dir_in.path(), prop)?; // Populate and persist (but don't build yet) the index - for i in 0..1000 { - let i = i as f32; - let _ = index.insert(vec![i, i + 1.0, i + 2.0])?; + let mut rng = rand::thread_rng(); + for _ in 0..25_000 { + index.insert(vec![rng.gen(); 3])?; } index.persist()?; From cddfc423b00dc0b462760da1bf45a54e49837586 Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Thu, 11 May 2023 01:49:05 +0200 Subject: [PATCH 08/11] Basic f16 support --- Cargo.toml | 1 + src/lib.rs | 2 -- src/ngt/index.rs | 67 ++++++++++++++++++++++++++++++++++++++++++- src/ngt/properties.rs | 6 ++++ 4 files changed, 73 insertions(+), 3 deletions(-) diff --git a/Cargo.toml b/Cargo.toml index 35f93b1..0c87694 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -11,6 +11,7 @@ license = "Apache-2.0" readme = "README.md" [dependencies] +half = "2" ngt-sys = { path = "ngt-sys", version = "2.0.11" } num_enum = "0.5" scopeguard = "1" diff --git a/src/lib.rs b/src/lib.rs index 393c267..db7c43a 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,5 +22,3 @@ pub const EPSILON: f32 = 0.1; pub use crate::error::{Error, Result}; pub use crate::ngt::{optim, NgtDistance, NgtIndex, NgtObject, NgtProperties}; - -// TODO: what about float16 ? diff --git a/src/ngt/index.rs b/src/ngt/index.rs index 0c7ebf5..6bd9cdf 100644 --- a/src/ngt/index.rs +++ b/src/ngt/index.rs @@ -283,6 +283,21 @@ impl NgtIndex { results.iter().copied().collect::>() } + NgtObject::Float16 => { + let results = sys::ngt_get_object(self.ospace, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + + let results = Vec::from_raw_parts( + results as *mut half::f16, + self.prop.dimension as usize, + self.prop.dimension as usize, + ); + let results = mem::ManuallyDrop::new(results); + + results.iter().map(|f16| f16.to_f32()).collect::>() + } NgtObject::Uint8 => { let results = sys::ngt_get_object_as_integer(self.ospace, id, self.ebuf); if results.is_null() { @@ -296,7 +311,7 @@ impl NgtIndex { ); let results = mem::ManuallyDrop::new(results); - results.iter().map(|byte| *byte as f32).collect::>() + results.iter().map(|&byte| byte as f32).collect::>() } }; @@ -432,6 +447,56 @@ mod tests { Ok(()) } + #[test] + fn test_u8() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + if cfg!(feature = "shared_mem") { + std::fs::remove_dir(dir.path())?; + } + + // Create an index for vectors of dimension 3 + let prop = NgtProperties::dimension(3)?.object_type(NgtObject::Uint8)?; + let mut index = NgtIndex::create(dir.path(), prop)?; + + // Batch insert 2 vectors, build and persist the index + index.insert_batch(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]])?; + index.build(2)?; + index.persist()?; + + // Verify that the index was built correctly with a vector search + let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; + assert_eq!(1, res[0].id); + + dir.close()?; + Ok(()) + } + + #[test] + fn test_f16() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + if cfg!(feature = "shared_mem") { + std::fs::remove_dir(dir.path())?; + } + + // Create an index for vectors of dimension 3 + let prop = NgtProperties::dimension(3)?.object_type(NgtObject::Float16)?; + let mut index = NgtIndex::create(dir.path(), prop)?; + + // Batch insert 2 vectors, build and persist the index + index.insert_batch(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]])?; + index.build(2)?; + index.persist()?; + + // Verify that the index was built correctly with a vector search + let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; + assert_eq!(1, res[0].id); + + dir.close()?; + Ok(()) + } + #[test] fn test_multithreaded() -> StdResult<(), Box> { // Get a temporary directory to store the index diff --git a/src/ngt/properties.rs b/src/ngt/properties.rs index 688b8c4..05b294a 100644 --- a/src/ngt/properties.rs +++ b/src/ngt/properties.rs @@ -12,6 +12,7 @@ use crate::error::{make_err, Result}; pub enum NgtObject { Uint8 = 1, Float = 2, + Float16 = 3, } #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] @@ -225,6 +226,11 @@ impl NgtProperties { Err(make_err(ebuf))? } } + NgtObject::Float16 => { + if !sys::ngt_set_property_object_type_float16(raw_prop, ebuf) { + Err(make_err(ebuf))? + } + } } Ok(()) From 8dd09ac0b2fcf26d11a9b347be443cdeb443f1ff Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Fri, 9 Jun 2023 12:32:16 +0200 Subject: [PATCH 09/11] Indexes are now templated by their content types --- Cargo.toml | 2 +- README.md | 15 +- ngt-sys/Cargo.toml | 2 +- ngt-sys/NGT | 2 +- src/lib.rs | 3 + src/ngt/index.rs | 78 +++++++---- src/ngt/mod.rs | 15 +- src/ngt/optim.rs | 58 +++++--- src/ngt/properties.rs | 57 ++++++-- src/qbg/index.rs | 314 +++++++++--------------------------------- src/qbg/mod.rs | 23 +--- src/qbg/properties.rs | 260 ++++++++++++++++++++++++++++++++++ src/qg/index.rs | 78 ++++------- src/qg/mod.rs | 8 +- src/qg/properties.rs | 66 ++++++++- 15 files changed, 574 insertions(+), 407 deletions(-) create mode 100644 src/qbg/properties.rs diff --git a/Cargo.toml b/Cargo.toml index 0c87694..9257b7c 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ readme = "README.md" [dependencies] half = "2" -ngt-sys = { path = "ngt-sys", version = "2.0.11" } +ngt-sys = { path = "ngt-sys", version = "2.0.12" } num_enum = "0.5" scopeguard = "1" diff --git a/README.md b/README.md index 2bf6522..d127f19 100644 --- a/README.md +++ b/README.md @@ -11,7 +11,7 @@ ten to several thousand dimensions). This crate provides the following indexes: * `NgtIndex`: Graph and tree-based index[^1] -* `QqIndex`: Quantized graph-based index[^2] +* `QgIndex`: Quantized graph-based index[^2] * `QbgIndex`: Quantized blob graph-based index The quantized indexes are available through the `quantized` Cargo feature. Note that @@ -37,16 +37,15 @@ features are available through the features `shared_mem` and `large_data` respec Defining the properties of a new index: ```rust,ignore -use ngt::{NgtProperties, NgtDistance, NgtObject}; +use ngt::{NgtProperties, NgtDistance}; // Defaut properties with vectors of dimension 3 -let prop = NgtProperties::dimension(3)?; +let prop = NgtProperties::::dimension(3)?; // Or customize values (here are the defaults) -let prop = NgtProperties::dimension(3)? +let prop = NgtProperties::::dimension(3)? .creation_edge_size(10)? .search_edge_size(40)? - .object_type(NgtObject::Float)? .distance_type(NgtDistance::L2)?; ``` @@ -57,7 +56,7 @@ use ngt::{NgtIndex, NgtProperties, EPSILON}; // Create a new index let prop = NgtProperties::dimension(3)?; -let index = NgtIndex::create("target/path/to/index/dir", prop)?; +let index: NgtIndex = NgtIndex::create("target/path/to/index/dir", prop)?; // Open an existing index let mut index = NgtIndex::open("target/path/to/index/dir")?; @@ -68,7 +67,7 @@ let vec2 = vec![4.0, 5.0, 6.0]; let id1 = index.insert(vec1)?; let id2 = index.insert(vec2)?; -// Actually build the index (not yet persisted on disk) +// Build the index in RAM (not yet persisted on disk) // This is required in order to be able to search vectors index.build(2)?; @@ -80,7 +79,7 @@ assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); // Remove a vector and check that it is not present anymore index.remove(id1)?; let res = index.get_vec(id1); -assert!(matches!(res, Result::Err(_))); +assert!(res.is_err()); // Verify that now our search result is different let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; diff --git a/ngt-sys/Cargo.toml b/ngt-sys/Cargo.toml index 08150ef..9d3562b 100644 --- a/ngt-sys/Cargo.toml +++ b/ngt-sys/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ngt-sys" -version = "2.0.11" +version = "2.0.12" authors = ["Romain Leroux "] edition = "2021" links = "ngt" diff --git a/ngt-sys/NGT b/ngt-sys/NGT index a91abde..5920ef0 160000 --- a/ngt-sys/NGT +++ b/ngt-sys/NGT @@ -1 +1 @@ -Subproject commit a91abde328aa70faad467b89545e99965eee0b2a +Subproject commit 5920ef01016998e96ad106284ed35e684981250d diff --git a/src/lib.rs b/src/lib.rs index db7c43a..8d43036 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -22,3 +22,6 @@ pub const EPSILON: f32 = 0.1; pub use crate::error::{Error, Result}; pub use crate::ngt::{optim, NgtDistance, NgtIndex, NgtObject, NgtProperties}; + +#[doc(inline)] +pub use half; diff --git a/src/ngt/index.rs b/src/ngt/index.rs index 6bd9cdf..388bbe8 100644 --- a/src/ngt/index.rs +++ b/src/ngt/index.rs @@ -9,25 +9,28 @@ use std::ptr; use ngt_sys as sys; use scopeguard::defer; -use super::{NgtObject, NgtProperties}; +use super::{NgtObject, NgtObjectType, NgtProperties}; use crate::error::{make_err, Error, Result}; use crate::{SearchResult, VecId}; #[derive(Debug)] -pub struct NgtIndex { +pub struct NgtIndex { pub(crate) path: CString, - pub(crate) prop: NgtProperties, + pub(crate) prop: NgtProperties, pub(crate) index: sys::NGTIndex, ospace: sys::NGTObjectSpace, ebuf: sys::NGTError, } -unsafe impl Send for NgtIndex {} -unsafe impl Sync for NgtIndex {} +unsafe impl Send for NgtIndex {} +unsafe impl Sync for NgtIndex {} -impl NgtIndex { +impl NgtIndex +where + T: NgtObjectType, +{ /// Creates an empty ANNG index with the given [`NgtProperties`](). - pub fn create>(path: P, prop: NgtProperties) -> Result { + pub fn create>(path: P, prop: NgtProperties) -> Result { if cfg!(feature = "shared_mem") && path.as_ref().exists() { Err(Error(format!("Path {:?} already exists", path.as_ref())))? } @@ -190,19 +193,33 @@ impl NgtIndex { /// discoverable yet. /// /// **The method [`build`](NgtIndex::build) must be called after inserting vectors**. - pub fn insert(&mut self, mut vec: Vec) -> Result { + pub fn insert(&mut self, mut vec: Vec) -> Result { unsafe { - let id = sys::ngt_insert_index_as_float( - self.index, - vec.as_mut_ptr(), - self.prop.dimension as u32, - self.ebuf, - ); + let id = match self.prop.object_type { + NgtObject::Float => sys::ngt_insert_index_as_float( + self.index, + vec.as_mut_ptr() as *mut _, + self.prop.dimension as u32, + self.ebuf, + ), + NgtObject::Uint8 => sys::ngt_insert_index_as_uint8( + self.index, + vec.as_mut_ptr() as *mut _, + self.prop.dimension as u32, + self.ebuf, + ), + NgtObject::Float16 => sys::ngt_insert_index_as_float16( + self.index, + vec.as_mut_ptr() as *mut _, + self.prop.dimension as u32, + self.ebuf, + ), + }; if id == 0 { Err(make_err(self.ebuf))? + } else { + Ok(id) } - - Ok(id) } } @@ -265,9 +282,9 @@ impl NgtIndex { } /// Get the specified vector. - pub fn get_vec(&self, id: VecId) -> Result> { + pub fn get_vec(&self, id: VecId) -> Result> { unsafe { - let results = match self.prop.object_type { + match self.prop.object_type { NgtObject::Float => { let results = sys::ngt_get_object_as_float(self.ospace, id, self.ebuf); if results.is_null() { @@ -281,7 +298,8 @@ impl NgtIndex { ); let results = mem::ManuallyDrop::new(results); - results.iter().copied().collect::>() + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) } NgtObject::Float16 => { let results = sys::ngt_get_object(self.ospace, id, self.ebuf); @@ -296,7 +314,8 @@ impl NgtIndex { ); let results = mem::ManuallyDrop::new(results); - results.iter().map(|f16| f16.to_f32()).collect::>() + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) } NgtObject::Uint8 => { let results = sys::ngt_get_object_as_integer(self.ospace, id, self.ebuf); @@ -311,11 +330,10 @@ impl NgtIndex { ); let results = mem::ManuallyDrop::new(results); - results.iter().map(|&byte| byte as f32).collect::>() + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) } - }; - - Ok(results) + } } } @@ -330,7 +348,7 @@ impl NgtIndex { } } -impl Drop for NgtIndex { +impl Drop for NgtIndex { fn drop(&mut self) { if !self.index.is_null() { unsafe { sys::ngt_close_index(self.index) }; @@ -364,7 +382,7 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = NgtProperties::dimension(3)?; + let prop = NgtProperties::::dimension(3)?; let mut index = NgtIndex::create(dir.path(), prop)?; // Insert two vectors and get their id @@ -431,7 +449,7 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = NgtProperties::dimension(3)?; + let prop = NgtProperties::::dimension(3)?; let mut index = NgtIndex::create(dir.path(), prop)?; // Batch insert 2 vectors, build and persist the index @@ -456,7 +474,7 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = NgtProperties::dimension(3)?.object_type(NgtObject::Uint8)?; + let prop = NgtProperties::::dimension(3)?; let mut index = NgtIndex::create(dir.path(), prop)?; // Batch insert 2 vectors, build and persist the index @@ -481,7 +499,7 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = NgtProperties::dimension(3)?.object_type(NgtObject::Float16)?; + let prop = NgtProperties::::dimension(3)?; let mut index = NgtIndex::create(dir.path(), prop)?; // Batch insert 2 vectors, build and persist the index @@ -506,7 +524,7 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = NgtProperties::dimension(3)?; + let prop = NgtProperties::::dimension(3)?; let mut index = NgtIndex::create(dir.path(), prop)?; let vecs = vec![ diff --git a/src/ngt/mod.rs b/src/ngt/mod.rs index 6a6e7b9..ff9d15e 100644 --- a/src/ngt/mod.rs +++ b/src/ngt/mod.rs @@ -2,16 +2,15 @@ //! //! ```rust //! # fn main() -> Result<(), ngt::Error> { -//! use ngt::{NgtProperties, NgtDistance, NgtObject}; +//! use ngt::{NgtProperties, NgtDistance}; //! //! // Defaut properties with vectors of dimension 3 -//! let prop = NgtProperties::dimension(3)?; +//! let prop = NgtProperties::::dimension(3)?; //! //! // Or customize values (here are the defaults) -//! let prop = NgtProperties::dimension(3)? +//! let prop = NgtProperties::::dimension(3)? //! .creation_edge_size(10)? //! .search_edge_size(40)? -//! .object_type(NgtObject::Float)? //! .distance_type(NgtDistance::L2)?; //! //! # Ok(()) @@ -26,7 +25,7 @@ //! //! // Create a new index //! let prop = NgtProperties::dimension(3)?; -//! let index = NgtIndex::create("target/path/to/index/dir", prop)?; +//! let index: NgtIndex = NgtIndex::create("target/path/to/index/dir", prop)?; //! //! // Open an existing index //! let mut index = NgtIndex::open("target/path/to/index/dir")?; @@ -37,7 +36,7 @@ //! let id1 = index.insert(vec1)?; //! let id2 = index.insert(vec2)?; //! -//! // Actually build the index (not yet persisted on disk) +//! // Build the index in RAM (not yet persisted on disk) //! // This is required in order to be able to search vectors //! index.build(2)?; //! @@ -49,7 +48,7 @@ //! // Remove a vector and check that it is not present anymore //! index.remove(id1)?; //! let res = index.get_vec(id1); -//! assert!(matches!(res, Result::Err(_))); +//! assert!(res.is_err()); //! //! // Verify that now our search result is different //! let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; @@ -69,4 +68,4 @@ pub mod optim; mod properties; pub use self::index::NgtIndex; -pub use self::properties::{NgtDistance, NgtObject, NgtProperties}; +pub use self::properties::{NgtDistance, NgtObject, NgtObjectType, NgtProperties}; diff --git a/src/ngt/optim.rs b/src/ngt/optim.rs index 5233b3e..0fe7ad2 100644 --- a/src/ngt/optim.rs +++ b/src/ngt/optim.rs @@ -6,6 +6,7 @@ use std::ptr; use ngt_sys as sys; use scopeguard::defer; +use super::NgtObjectType; use crate::error::{make_err, Result}; use crate::ngt::index::NgtIndex; @@ -42,10 +43,14 @@ pub fn optimize_anng_edges_number>( /// /// Optimizes the search parameters about the explored edges and memory prefetch for the /// existing indexes. Does not modify the index data structure. -pub fn optimize_anng_search_parameters>(index_path: P) -> Result<()> { +pub fn optimize_anng_search_parameters(index_path: P) -> Result<()> +where + T: NgtObjectType, + P: AsRef, +{ let mut optimizer = GraphOptimizer::new(GraphOptimParams::default())?; optimizer.set_processing_modes(true, true, true)?; - optimizer.adjust_search_coefficients(index_path)?; + optimizer.adjust_search_coefficients::(index_path)?; Ok(()) } @@ -55,7 +60,10 @@ pub fn optimize_anng_search_parameters>(index_path: P) -> Result< /// node. Note that refinement takes a long processing time. An ANNG index can be /// refined only after it has been [`built`](NgtIndex::build). #[cfg(not(feature = "shared_mem"))] -pub fn refine_anng(index: &mut NgtIndex, params: AnngRefineParams) -> Result<()> { +pub fn refine_anng( + index: &mut NgtIndex, + params: AnngRefineParams, +) -> Result<()> { unsafe { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } @@ -90,13 +98,17 @@ pub fn refine_anng(index: &mut NgtIndex, params: AnngRefineParams) -> Result<()> /// Important [`GraphOptimParams`](GraphOptimParams) parameters are `nb_outgoing` edges /// and `nb_incoming` edges. The latter can be set to an even higher number than the /// `creation_edge_size` of the original ANNG. -pub fn convert_anng_to_onng>( +pub fn convert_anng_to_onng( index_anng_in: P, index_onng_out: P, params: GraphOptimParams, -) -> Result<()> { +) -> Result<()> +where + T: NgtObjectType, + P: AsRef, +{ let mut optimizer = GraphOptimizer::new(params)?; - optimizer.convert_anng_to_onng(index_anng_in, index_onng_out)?; + optimizer.convert_anng_to_onng::(index_anng_in, index_onng_out)?; Ok(()) } @@ -253,8 +265,12 @@ impl GraphOptimizer { } /// Optimize for the search parameters of an ANNG. - fn adjust_search_coefficients>(&mut self, index_path: P) -> Result<()> { - let _ = NgtIndex::open(&index_path)?; + fn adjust_search_coefficients(&mut self, index_path: P) -> Result<()> + where + P: AsRef, + T: NgtObjectType, + { + let _ = NgtIndex::::open(&index_path)?; unsafe { let ebuf = sys::ngt_create_error_object(); @@ -271,12 +287,12 @@ impl GraphOptimizer { } /// Converts the `index_in` ANNG to an ONNG at `index_out`. - fn convert_anng_to_onng>( - &mut self, - index_anng_in: P, - index_onng_out: P, - ) -> Result<()> { - let _ = NgtIndex::open(&index_anng_in)?; + fn convert_anng_to_onng(&mut self, index_anng_in: P, index_onng_out: P) -> Result<()> + where + T: NgtObjectType, + P: AsRef, + { + let _ = NgtIndex::::open(&index_anng_in)?; unsafe { let ebuf = sys::ngt_create_error_object(); @@ -321,7 +337,7 @@ mod tests { let dir = tempdir()?; // Create an index for vectors of dimension 3 with cosine distance - let prop = NgtProperties::dimension(3)?.distance_type(NgtDistance::Cosine)?; + let prop = NgtProperties::::dimension(3)?.distance_type(NgtDistance::Cosine)?; let mut index = NgtIndex::create(dir.path(), prop)?; // Populate the index, but don't build it yet @@ -335,12 +351,12 @@ mod tests { optimize_anng_edges_number(dir.path(), AnngEdgeOptimParams::default())?; // Now build and persist again the optimized index - let mut index = NgtIndex::open(dir.path())?; + let mut index = NgtIndex::::open(dir.path())?; index.build(4)?; index.persist()?; // Further optimize the index - optimize_anng_search_parameters(dir.path())?; + optimize_anng_search_parameters::(dir.path())?; dir.close()?; Ok(()) @@ -353,7 +369,7 @@ mod tests { let dir = tempdir()?; // Create an index for vectors of dimension 3 with cosine distance - let prop = NgtProperties::dimension(3)?.distance_type(NgtDistance::Cosine)?; + let prop = NgtProperties::::dimension(3)?.distance_type(NgtDistance::Cosine)?; let mut index = NgtIndex::create(dir.path(), prop)?; // Populate and build the index @@ -378,7 +394,7 @@ mod tests { let dir_in = tempdir()?; // Create an index for vectors of dimension 3 with cosine distance - let prop = NgtProperties::dimension(3)? + let prop = NgtProperties::::dimension(3)? .distance_type(NgtDistance::Cosine)? .creation_edge_size(100)?; // More than default value, improves the final ONNG @@ -395,7 +411,7 @@ mod tests { optimize_anng_edges_number(dir_in.path(), AnngEdgeOptimParams::default())?; // Now build and persist again the optimized index - let mut index = NgtIndex::open(dir_in.path())?; + let mut index = NgtIndex::::open(dir_in.path())?; index.build(4)?; index.persist()?; @@ -407,7 +423,7 @@ mod tests { let mut params = GraphOptimParams::default(); params.nb_outgoing = 10; params.nb_incoming = 100; // An even larger number of incoming edges can be specified - convert_anng_to_onng(dir_in.path(), dir_out.path(), params)?; + convert_anng_to_onng::(dir_in.path(), dir_out.path(), params)?; dir_out.close()?; dir_in.close()?; diff --git a/src/ngt/properties.rs b/src/ngt/properties.rs index 05b294a..8198a93 100644 --- a/src/ngt/properties.rs +++ b/src/ngt/properties.rs @@ -1,6 +1,7 @@ -use std::convert::TryFrom; use std::ptr; +use std::{convert::TryFrom, marker::PhantomData}; +use half::f16; use ngt_sys as sys; use num_enum::TryFromPrimitive; use scopeguard::defer; @@ -15,6 +16,35 @@ pub enum NgtObject { Float16 = 3, } +mod private { + pub trait Sealed {} +} + +pub trait NgtObjectType: private::Sealed { + fn as_obj() -> NgtObject; +} + +impl private::Sealed for f32 {} +impl NgtObjectType for f32 { + fn as_obj() -> NgtObject { + NgtObject::Float + } +} + +impl private::Sealed for u8 {} +impl NgtObjectType for u8 { + fn as_obj() -> NgtObject { + NgtObject::Uint8 + } +} + +impl private::Sealed for f16 {} +impl NgtObjectType for f16 { + fn as_obj() -> NgtObject { + NgtObject::Float16 + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] pub enum NgtDistance { @@ -33,24 +63,28 @@ pub enum NgtDistance { } #[derive(Debug)] -pub struct NgtProperties { +pub struct NgtProperties { pub(crate) dimension: i32, pub(crate) creation_edge_size: i16, pub(crate) search_edge_size: i16, pub(crate) object_type: NgtObject, pub(crate) distance_type: NgtDistance, pub(crate) raw_prop: sys::NGTProperty, + _marker: PhantomData, } -unsafe impl Send for NgtProperties {} -unsafe impl Sync for NgtProperties {} +unsafe impl Send for NgtProperties {} +unsafe impl Sync for NgtProperties {} -impl NgtProperties { +impl NgtProperties +where + T: NgtObjectType, +{ pub fn dimension(dimension: usize) -> Result { let dimension = i32::try_from(dimension)?; let creation_edge_size = 10; let search_edge_size = 40; - let object_type = NgtObject::Float; + let object_type = T::as_obj(); let distance_type = NgtDistance::L2; unsafe { @@ -75,6 +109,7 @@ impl NgtProperties { object_type, distance_type, raw_prop, + _marker: PhantomData, }) } } @@ -102,6 +137,7 @@ impl NgtProperties { object_type: self.object_type, distance_type: self.distance_type, raw_prop, + _marker: PhantomData, }) } } @@ -154,6 +190,7 @@ impl NgtProperties { object_type, distance_type, raw_prop, + _marker: PhantomData, }) } } @@ -205,12 +242,6 @@ impl NgtProperties { Ok(()) } - pub fn object_type(mut self, object_type: NgtObject) -> Result { - self.object_type = object_type; - unsafe { Self::set_object_type(self.raw_prop, object_type)? }; - Ok(self) - } - unsafe fn set_object_type(raw_prop: sys::NGTProperty, object_type: NgtObject) -> Result<()> { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } @@ -316,7 +347,7 @@ impl NgtProperties { } } -impl Drop for NgtProperties { +impl Drop for NgtProperties { fn drop(&mut self) { if !self.raw_prop.is_null() { unsafe { sys::ngt_destroy_property(self.raw_prop) }; diff --git a/src/qbg/index.rs b/src/qbg/index.rs index 2edce50..44d9b29 100644 --- a/src/qbg/index.rs +++ b/src/qbg/index.rs @@ -1,28 +1,33 @@ use std::ffi::CString; +use std::marker::PhantomData; use std::os::unix::ffi::OsStrExt; use std::path::Path; use std::{mem, ptr}; use ngt_sys as sys; -use num_enum::TryFromPrimitive; use scopeguard::defer; use crate::error::{make_err, Error, Result}; use crate::{SearchResult, VecId}; -use super::{QbgDistance, QbgObject}; +use super::{QbgBuildParams, QbgConstructParams, QbgObject, QbgObjectType}; #[derive(Debug)] -pub struct QbgIndex { +pub struct QbgIndex { pub(crate) index: sys::QBGIndex, path: CString, - _mode: T, + _mode: M, + obj_type: QbgObject, dimension: u32, ebuf: sys::NGTError, + _marker: PhantomData, } -impl QbgIndex { - pub fn create

(path: P, create_params: QbgConstructParams) -> Result +impl QbgIndex +where + T: QbgObjectType, +{ + pub fn create

(path: P, create_params: QbgConstructParams) -> Result where P: AsRef, { @@ -57,11 +62,14 @@ impl QbgIndex { path, _mode: ModeWrite, dimension, + obj_type: T::as_obj(), ebuf: sys::ngt_create_error_object(), + _marker: PhantomData, }) } } + // TODO: should be mut vec: Vec pub fn insert(&mut self, mut vec: Vec) -> Result { unsafe { let id = @@ -96,14 +104,17 @@ impl QbgIndex { } } - pub fn into_readable(self) -> Result> { + pub fn into_readable(self) -> Result> { let path = self.path.clone(); drop(self); QbgIndex::open(path.into_string()?) } } -impl QbgIndex { +impl QbgIndex +where + T: QbgObjectType, +{ pub fn open>(path: P) -> Result { if !is_x86_feature_detected!("avx2") { return Err(Error( @@ -134,8 +145,10 @@ impl QbgIndex { index, path, _mode: ModeRead, + obj_type: T::as_obj(), dimension, ebuf: sys::ngt_create_error_object(), + _marker: PhantomData, }) } } @@ -171,7 +184,7 @@ impl QbgIndex { } } - pub fn into_writable(self) -> Result> { + pub fn into_writable(self) -> Result> { unsafe { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } @@ -193,38 +206,63 @@ impl QbgIndex { index, path, _mode: ModeWrite, + obj_type: T::as_obj(), dimension, ebuf: sys::ngt_create_error_object(), + _marker: PhantomData, }) } } } -impl QbgIndex +impl QbgIndex where - T: IndexMode, + T: QbgObjectType, + M: IndexMode, { - pub fn get_vec(&self, id: VecId) -> Result> { + /// Get the specified vector. + pub fn get_vec(&self, id: VecId) -> Result> { unsafe { - let results = sys::qbg_get_object(self.index, id, self.ebuf); - if results.is_null() { - Err(make_err(self.ebuf))? + match self.obj_type { + QbgObject::Float => { + let results = sys::qbg_get_object(self.index, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + + let results = Vec::from_raw_parts( + results as *mut f32, + self.dimension as usize, + self.dimension as usize, + ); + let results = mem::ManuallyDrop::new(results); + + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) + } + QbgObject::Uint8 => { + // TODO: Would need some kind of qbg_get_object_as_integer + let results = sys::qbg_get_object(self.index, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + + let results = Vec::from_raw_parts( + results as *mut f32, + self.dimension as usize, + self.dimension as usize, + ); + let results = mem::ManuallyDrop::new(results); + + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) + } } - - let results = Vec::from_raw_parts( - results as *mut f32, - self.dimension as usize, - self.dimension as usize, - ); - let results = mem::ManuallyDrop::new(results); - let results = results.iter().copied().collect::>(); - - Ok(results) } } } -impl Drop for QbgIndex { +impl Drop for QbgIndex { fn drop(&mut self) { if !self.index.is_null() { unsafe { sys::qbg_close_index(self.index) }; @@ -255,227 +293,6 @@ pub struct ModeWrite; impl private::Sealed for ModeWrite {} impl IndexMode for ModeWrite {} -#[derive(Debug, Clone, PartialEq)] -pub struct QbgConstructParams { - extended_dimension: u64, - dimension: u64, - number_of_subvectors: u64, - number_of_blobs: u64, - internal_data_type: QbgObject, - data_type: QbgObject, - distance_type: QbgDistance, -} - -impl Default for QbgConstructParams { - fn default() -> Self { - Self { - extended_dimension: 0, - dimension: 0, - number_of_subvectors: 1, - number_of_blobs: 0, - internal_data_type: QbgObject::Float, - data_type: QbgObject::Float, - distance_type: QbgDistance::L2, - } - } -} - -impl QbgConstructParams { - pub fn extended_dimension(mut self, extended_dimension: u64) -> Self { - self.extended_dimension = extended_dimension; - self - } - - pub fn dimension(mut self, dimension: u64) -> Self { - self.dimension = dimension; - self - } - - pub fn number_of_subvectors(mut self, number_of_subvectors: u64) -> Self { - self.number_of_subvectors = number_of_subvectors; - self - } - - pub fn number_of_blobs(mut self, number_of_blobs: u64) -> Self { - self.number_of_blobs = number_of_blobs; - self - } - - pub fn internal_data_type(mut self, internal_data_type: QbgObject) -> Self { - self.internal_data_type = internal_data_type; - self - } - - pub fn data_type(mut self, data_type: QbgObject) -> Self { - self.data_type = data_type; - self - } - - pub fn distance_type(mut self, distance_type: QbgDistance) -> Self { - self.distance_type = distance_type; - self - } - - unsafe fn into_raw(self) -> sys::QBGConstructionParameters { - sys::QBGConstructionParameters { - extended_dimension: self.extended_dimension, - dimension: self.dimension, - number_of_subvectors: self.number_of_subvectors, - number_of_blobs: self.number_of_blobs, - internal_data_type: self.internal_data_type as i32, - data_type: self.data_type as i32, - distance_type: self.distance_type as i32, - } - } -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] -#[repr(i32)] -pub enum QbgClusteringInitMode { - Head = 0, - Random = 1, - KmeansPlusPlus = 2, - RandomFixedSeed = 3, - KmeansPlusPlusFixedSeed = 4, - Best = 5, -} - -#[derive(Debug, Clone)] -pub struct QbgBuildParams { - // hierarchical kmeans - hierarchical_clustering_init_mode: QbgClusteringInitMode, - number_of_first_objects: u64, - number_of_first_clusters: u64, - number_of_second_objects: u64, - number_of_second_clusters: u64, - number_of_third_clusters: u64, - // optimization - number_of_objects: u64, - number_of_subvectors: u64, - optimization_clustering_init_mode: QbgClusteringInitMode, - rotation_iteration: u64, - subvector_iteration: u64, - number_of_matrices: u64, - rotation: bool, - repositioning: bool, -} - -impl Default for QbgBuildParams { - fn default() -> Self { - Self { - hierarchical_clustering_init_mode: QbgClusteringInitMode::KmeansPlusPlus, - number_of_first_objects: 0, - number_of_first_clusters: 0, - number_of_second_objects: 0, - number_of_second_clusters: 0, - number_of_third_clusters: 0, - number_of_objects: 1000, - number_of_subvectors: 1, - optimization_clustering_init_mode: QbgClusteringInitMode::KmeansPlusPlus, - rotation_iteration: 2000, - subvector_iteration: 400, - number_of_matrices: 3, - rotation: true, - repositioning: false, - } - } -} - -impl QbgBuildParams { - pub fn hierarchical_clustering_init_mode( - mut self, - clustering_init_mode: QbgClusteringInitMode, - ) -> Self { - self.hierarchical_clustering_init_mode = clustering_init_mode; - self - } - - pub fn number_of_first_objects(mut self, number_of_first_objects: u64) -> Self { - self.number_of_first_objects = number_of_first_objects; - self - } - - pub fn number_of_first_clusters(mut self, number_of_first_clusters: u64) -> Self { - self.number_of_first_clusters = number_of_first_clusters; - self - } - - pub fn number_of_second_objects(mut self, number_of_second_objects: u64) -> Self { - self.number_of_second_objects = number_of_second_objects; - self - } - - pub fn number_of_second_clusters(mut self, number_of_second_clusters: u64) -> Self { - self.number_of_second_clusters = number_of_second_clusters; - self - } - - pub fn number_of_third_clusters(mut self, number_of_third_clusters: u64) -> Self { - self.number_of_third_clusters = number_of_third_clusters; - self - } - - pub fn number_of_objects(mut self, number_of_objects: u64) -> Self { - self.number_of_objects = number_of_objects; - self - } - pub fn number_of_subvectors(mut self, number_of_subvectors: u64) -> Self { - self.number_of_subvectors = number_of_subvectors; - self - } - pub fn optimization_clustering_init_mode( - mut self, - clustering_init_mode: QbgClusteringInitMode, - ) -> Self { - self.optimization_clustering_init_mode = clustering_init_mode; - self - } - - pub fn rotation_iteration(mut self, rotation_iteration: u64) -> Self { - self.rotation_iteration = rotation_iteration; - self - } - - pub fn subvector_iteration(mut self, subvector_iteration: u64) -> Self { - self.subvector_iteration = subvector_iteration; - self - } - - pub fn number_of_matrices(mut self, number_of_matrices: u64) -> Self { - self.number_of_matrices = number_of_matrices; - self - } - - pub fn rotation(mut self, rotation: bool) -> Self { - self.rotation = rotation; - self - } - - pub fn repositioning(mut self, repositioning: bool) -> Self { - self.repositioning = repositioning; - self - } - - unsafe fn into_raw(self) -> sys::QBGBuildParameters { - sys::QBGBuildParameters { - hierarchical_clustering_init_mode: self.hierarchical_clustering_init_mode as i32, - number_of_first_objects: self.number_of_first_objects, - number_of_first_clusters: self.number_of_first_clusters, - number_of_second_objects: self.number_of_second_objects, - number_of_second_clusters: self.number_of_second_clusters, - number_of_third_clusters: self.number_of_third_clusters, - number_of_objects: self.number_of_objects, - number_of_subvectors: self.number_of_subvectors, - optimization_clustering_init_mode: self.optimization_clustering_init_mode as i32, - rotation_iteration: self.rotation_iteration, - subvector_iteration: self.subvector_iteration, - number_of_matrices: self.number_of_matrices, - rotation: self.rotation, - repositioning: self.repositioning, - } - } -} - #[derive(Debug, Clone, PartialEq)] pub struct QbgQuery<'a> { query: &'a [f32], @@ -569,8 +386,7 @@ mod tests { // Create a QGB index let ndims = 3; - let mut index = - QbgIndex::create(dir.path(), QbgConstructParams::default().dimension(ndims))?; + let mut index = QbgIndex::create(dir.path(), QbgConstructParams::dimension(ndims))?; // Insert vectors and get their ids let nvecs = 16; diff --git a/src/qbg/mod.rs b/src/qbg/mod.rs index 5f1059c..0cb3da2 100644 --- a/src/qbg/mod.rs +++ b/src/qbg/mod.rs @@ -1,20 +1,9 @@ -mod index; - -use num_enum::TryFromPrimitive; +// TODO: Add module doc (specify available types) -#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] -#[repr(i32)] -pub enum QbgObject { - Uint8 = 0, - Float = 1, -} - -#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] -#[repr(i32)] -pub enum QbgDistance { - L2 = 1, -} +mod index; +mod properties; -pub use self::index::{ - IndexMode, ModeRead, ModeWrite, QbgBuildParams, QbgConstructParams, QbgIndex, QbgQuery, +pub use self::index::{IndexMode, ModeRead, ModeWrite, QbgIndex, QbgQuery}; +pub use self::properties::{ + QbgBuildParams, QbgConstructParams, QbgDistance, QbgObject, QbgObjectType, }; diff --git a/src/qbg/properties.rs b/src/qbg/properties.rs new file mode 100644 index 0000000..08cbe79 --- /dev/null +++ b/src/qbg/properties.rs @@ -0,0 +1,260 @@ +use std::marker::PhantomData; + +use ngt_sys as sys; +use num_enum::TryFromPrimitive; + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgObject { + Uint8 = 0, + Float = 1, +} + +mod private { + pub trait Sealed {} +} + +pub trait QbgObjectType: private::Sealed { + fn as_obj() -> QbgObject; +} + +impl private::Sealed for f32 {} +impl QbgObjectType for f32 { + fn as_obj() -> QbgObject { + QbgObject::Float + } +} + +impl private::Sealed for u8 {} +impl QbgObjectType for u8 { + fn as_obj() -> QbgObject { + QbgObject::Uint8 + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgDistance { + L2 = 1, +} + +#[derive(Debug, Clone, PartialEq)] +pub struct QbgConstructParams { + extended_dimension: u64, + dimension: u64, + number_of_subvectors: u64, + number_of_blobs: u64, + internal_data_type: QbgObject, + data_type: QbgObject, + distance_type: QbgDistance, + _marker: PhantomData, +} + +impl QbgConstructParams +where + T: QbgObjectType, +{ + pub fn dimension(dimension: u64) -> Self { + let extended_dimension = 0; + let number_of_subvectors = 1; + let number_of_blobs = 0; + let internal_data_type = QbgObject::Float; // TODO: Should be T::as_obj() ? + let data_type = T::as_obj(); + let distance_type = QbgDistance::L2; + + Self { + extended_dimension, + dimension, + number_of_subvectors, + number_of_blobs, + internal_data_type, + data_type, + distance_type, + _marker: PhantomData, + } + } + + pub fn extended_dimension(mut self, extended_dimension: u64) -> Self { + self.extended_dimension = extended_dimension; + self + } + + pub fn number_of_subvectors(mut self, number_of_subvectors: u64) -> Self { + self.number_of_subvectors = number_of_subvectors; + self + } + + pub fn number_of_blobs(mut self, number_of_blobs: u64) -> Self { + self.number_of_blobs = number_of_blobs; + self + } + + pub fn internal_data_type(mut self, internal_data_type: QbgObject) -> Self { + self.internal_data_type = internal_data_type; + self + } + + pub fn distance_type(mut self, distance_type: QbgDistance) -> Self { + self.distance_type = distance_type; + self + } + + pub(crate) unsafe fn into_raw(self) -> sys::QBGConstructionParameters { + sys::QBGConstructionParameters { + extended_dimension: self.extended_dimension, + dimension: self.dimension, + number_of_subvectors: self.number_of_subvectors, + number_of_blobs: self.number_of_blobs, + internal_data_type: self.internal_data_type as i32, + data_type: self.data_type as i32, + distance_type: self.distance_type as i32, + } + } +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] +#[repr(i32)] +pub enum QbgClusteringInitMode { + Head = 0, + Random = 1, + KmeansPlusPlus = 2, + RandomFixedSeed = 3, + KmeansPlusPlusFixedSeed = 4, + Best = 5, +} + +#[derive(Debug, Clone)] +pub struct QbgBuildParams { + // hierarchical kmeans + hierarchical_clustering_init_mode: QbgClusteringInitMode, + number_of_first_objects: u64, + number_of_first_clusters: u64, + number_of_second_objects: u64, + number_of_second_clusters: u64, + number_of_third_clusters: u64, + // optimization + number_of_objects: u64, + number_of_subvectors: u64, + optimization_clustering_init_mode: QbgClusteringInitMode, + rotation_iteration: u64, + subvector_iteration: u64, + number_of_matrices: u64, + rotation: bool, + repositioning: bool, +} + +impl Default for QbgBuildParams { + fn default() -> Self { + Self { + hierarchical_clustering_init_mode: QbgClusteringInitMode::KmeansPlusPlus, + number_of_first_objects: 0, + number_of_first_clusters: 0, + number_of_second_objects: 0, + number_of_second_clusters: 0, + number_of_third_clusters: 0, + number_of_objects: 1000, + number_of_subvectors: 1, + optimization_clustering_init_mode: QbgClusteringInitMode::KmeansPlusPlus, + rotation_iteration: 2000, + subvector_iteration: 400, + number_of_matrices: 3, + rotation: true, + repositioning: false, + } + } +} + +impl QbgBuildParams { + pub fn hierarchical_clustering_init_mode( + mut self, + clustering_init_mode: QbgClusteringInitMode, + ) -> Self { + self.hierarchical_clustering_init_mode = clustering_init_mode; + self + } + + pub fn number_of_first_objects(mut self, number_of_first_objects: u64) -> Self { + self.number_of_first_objects = number_of_first_objects; + self + } + + pub fn number_of_first_clusters(mut self, number_of_first_clusters: u64) -> Self { + self.number_of_first_clusters = number_of_first_clusters; + self + } + + pub fn number_of_second_objects(mut self, number_of_second_objects: u64) -> Self { + self.number_of_second_objects = number_of_second_objects; + self + } + + pub fn number_of_second_clusters(mut self, number_of_second_clusters: u64) -> Self { + self.number_of_second_clusters = number_of_second_clusters; + self + } + + pub fn number_of_third_clusters(mut self, number_of_third_clusters: u64) -> Self { + self.number_of_third_clusters = number_of_third_clusters; + self + } + + pub fn number_of_objects(mut self, number_of_objects: u64) -> Self { + self.number_of_objects = number_of_objects; + self + } + pub fn number_of_subvectors(mut self, number_of_subvectors: u64) -> Self { + self.number_of_subvectors = number_of_subvectors; + self + } + pub fn optimization_clustering_init_mode( + mut self, + clustering_init_mode: QbgClusteringInitMode, + ) -> Self { + self.optimization_clustering_init_mode = clustering_init_mode; + self + } + + pub fn rotation_iteration(mut self, rotation_iteration: u64) -> Self { + self.rotation_iteration = rotation_iteration; + self + } + + pub fn subvector_iteration(mut self, subvector_iteration: u64) -> Self { + self.subvector_iteration = subvector_iteration; + self + } + + pub fn number_of_matrices(mut self, number_of_matrices: u64) -> Self { + self.number_of_matrices = number_of_matrices; + self + } + + pub fn rotation(mut self, rotation: bool) -> Self { + self.rotation = rotation; + self + } + + pub fn repositioning(mut self, repositioning: bool) -> Self { + self.repositioning = repositioning; + self + } + + pub(crate) unsafe fn into_raw(self) -> sys::QBGBuildParameters { + sys::QBGBuildParameters { + hierarchical_clustering_init_mode: self.hierarchical_clustering_init_mode as i32, + number_of_first_objects: self.number_of_first_objects, + number_of_first_clusters: self.number_of_first_clusters, + number_of_second_objects: self.number_of_second_objects, + number_of_second_clusters: self.number_of_second_clusters, + number_of_third_clusters: self.number_of_third_clusters, + number_of_objects: self.number_of_objects, + number_of_subvectors: self.number_of_subvectors, + optimization_clustering_init_mode: self.optimization_clustering_init_mode as i32, + rotation_iteration: self.rotation_iteration, + subvector_iteration: self.subvector_iteration, + number_of_matrices: self.number_of_matrices, + rotation: self.rotation, + repositioning: self.repositioning, + } + } +} diff --git a/src/qg/index.rs b/src/qg/index.rs index 64fb291..9f142aa 100644 --- a/src/qg/index.rs +++ b/src/qg/index.rs @@ -7,21 +7,25 @@ use std::ptr; use ngt_sys as sys; use scopeguard::defer; -use super::{QgObject, QgProperties}; +use super::{QgObject, QgObjectType, QgProperties, QgQuantizationParams}; use crate::error::{make_err, Error, Result}; use crate::ngt::NgtIndex; use crate::{SearchResult, VecId}; #[derive(Debug)] -pub struct QgIndex { - pub(crate) prop: QgProperties, +pub struct QgIndex { + pub(crate) prop: QgProperties, pub(crate) index: sys::NGTQGIndex, ebuf: sys::NGTError, } -impl QgIndex { +impl QgIndex +where + T: QgObjectType, +{ /// Quantize an NGT index - pub fn quantize(index: NgtIndex, params: QgQuantizationParams) -> Result { + pub fn quantize(index: NgtIndex, params: QgQuantizationParams) -> Result { + // if !is_x86_feature_detected!("avx2") { return Err(Error( "Cannot quantize an index without AVX2 support".into(), @@ -75,7 +79,7 @@ impl QgIndex { } } - pub fn search(&self, query: QgQuery) -> Result> { + pub fn search(&self, query: QgQuery) -> Result> { unsafe { let results = sys::ngt_create_empty_results(self.ebuf); if results.is_null() { @@ -107,9 +111,9 @@ impl QgIndex { } /// Get the specified vector. - pub fn get_vec(&self, id: VecId) -> Result> { + pub fn get_vec(&self, id: VecId) -> Result> { unsafe { - let results = match self.prop.object_type { + match self.prop.object_type { QgObject::Float => { let ospace = sys::ngt_get_object_space(self.index, self.ebuf); if ospace.is_null() { @@ -128,7 +132,8 @@ impl QgIndex { ); let results = mem::ManuallyDrop::new(results); - results.iter().copied().collect::>() + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) } QgObject::Uint8 => { let ospace = sys::ngt_get_object_space(self.index, self.ebuf); @@ -148,16 +153,15 @@ impl QgIndex { ); let results = mem::ManuallyDrop::new(results); - results.iter().map(|byte| *byte as f32).collect::>() + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) } - }; - - Ok(results) + } } } } -impl Drop for QgIndex { +impl Drop for QgIndex { fn drop(&mut self) { if !self.index.is_null() { unsafe { sys::ngtqg_close_index(self.index) }; @@ -171,40 +175,16 @@ impl Drop for QgIndex { } #[derive(Debug, Clone, PartialEq)] -pub struct QgQuantizationParams { - pub dimension_of_subvector: f32, - pub max_number_of_edges: u64, -} - -impl Default for QgQuantizationParams { - fn default() -> Self { - Self { - dimension_of_subvector: 0.0, - max_number_of_edges: 128, - } - } -} - -impl QgQuantizationParams { - unsafe fn into_raw(self) -> sys::NGTQGQuantizationParameters { - sys::NGTQGQuantizationParameters { - dimension_of_subvector: self.dimension_of_subvector, - max_number_of_edges: self.max_number_of_edges, - } - } -} - -#[derive(Debug, Clone, PartialEq)] -pub struct QgQuery<'a> { - query: &'a [f32], +pub struct QgQuery<'a, T> { + query: &'a [T], pub size: u64, pub epsilon: f32, pub result_expansion: f32, pub radius: f32, } -impl<'a> QgQuery<'a> { - pub fn new(query: &'a [f32]) -> Self { +impl<'a, T> QgQuery<'a, T> { + pub fn new(query: &'a [T]) -> Self { Self { query, size: 20, @@ -254,7 +234,7 @@ mod tests { use tempfile::tempdir; use super::*; - use crate::{NgtDistance, NgtObject, NgtProperties}; + use crate::{NgtDistance, NgtProperties}; #[test] fn test_qg() -> StdResult<(), Box> { @@ -263,19 +243,17 @@ mod tests { // Create an NGT index for vectors let ndims = 3; - let props = NgtProperties::dimension(ndims)? - .object_type(NgtObject::Uint8)? - .distance_type(NgtDistance::L2)?; + let props = NgtProperties::::dimension(ndims)?.distance_type(NgtDistance::L2)?; let mut index = NgtIndex::create(dir.path(), props)?; // Insert vectors and get their ids - let nvecs = 16; + let nvecs = 64; let ids = (1..ndims * nvecs) .step_by(ndims) - .map(|i| i as f32) + .map(|i| i as u8) .map(|i| { repeat(i) - .zip((0..ndims).map(|j| j as f32)) + .zip((0..ndims).map(|j| j as u8)) .map(|(i, j)| i + j) .collect() }) @@ -294,7 +272,7 @@ mod tests { let index = QgIndex::quantize(index, params)?; // Perform a vector search (with 2 results) - let v: Vec = (1..=ndims).into_iter().map(|x| x as f32).collect(); + let v: Vec = (1..=ndims).into_iter().map(|x| x as u8).collect(); let query = QgQuery::new(&v).size(2); let res = index.search(query)?; assert_eq!(ids[0], res[0].id); diff --git a/src/qg/mod.rs b/src/qg/mod.rs index 21d170a..eb774b0 100644 --- a/src/qg/mod.rs +++ b/src/qg/mod.rs @@ -1,5 +1,9 @@ +// TODO: Add module doc (specify available types) + mod index; mod properties; -pub use self::index::{QgIndex, QgQuantizationParams, QgQuery}; -pub use self::properties::{QgDistance, QgObject, QgProperties}; +pub use self::index::{QgIndex, QgQuery}; +pub use self::properties::{ + QgDistance, QgObject, QgObjectType, QgProperties, QgQuantizationParams, +}; diff --git a/src/qg/properties.rs b/src/qg/properties.rs index 8112080..4cb23ea 100644 --- a/src/qg/properties.rs +++ b/src/qg/properties.rs @@ -1,3 +1,4 @@ +use std::marker::PhantomData; use std::ptr; use ngt_sys as sys; @@ -13,6 +14,28 @@ pub enum QgObject { Float = 2, } +mod private { + pub trait Sealed {} +} + +pub trait QgObjectType: private::Sealed { + fn as_obj() -> QgObject; +} + +impl private::Sealed for f32 {} +impl QgObjectType for f32 { + fn as_obj() -> QgObject { + QgObject::Float + } +} + +impl private::Sealed for u8 {} +impl QgObjectType for u8 { + fn as_obj() -> QgObject { + QgObject::Uint8 + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] pub enum QgDistance { @@ -21,24 +44,28 @@ pub enum QgDistance { } #[derive(Debug)] -pub struct QgProperties { +pub struct QgProperties { pub(crate) dimension: i32, pub(crate) creation_edge_size: i16, pub(crate) search_edge_size: i16, pub(crate) object_type: QgObject, pub(crate) distance_type: QgDistance, pub(crate) raw_prop: sys::NGTProperty, + _marker: PhantomData, } -unsafe impl Send for QgProperties {} -unsafe impl Sync for QgProperties {} +unsafe impl Send for QgProperties {} +unsafe impl Sync for QgProperties {} -impl QgProperties { +impl QgProperties +where + T: QgObjectType, +{ pub fn dimension(dimension: usize) -> Result { let dimension = i32::try_from(dimension)?; let creation_edge_size = 10; let search_edge_size = 40; - let object_type = QgObject::Float; + let object_type = T::as_obj(); let distance_type = QgDistance::L2; unsafe { @@ -63,6 +90,7 @@ impl QgProperties { object_type, distance_type, raw_prop, + _marker: PhantomData, }) } } @@ -90,6 +118,7 @@ impl QgProperties { object_type: self.object_type, distance_type: self.distance_type, raw_prop, + _marker: PhantomData, }) } } @@ -142,6 +171,7 @@ impl QgProperties { object_type, distance_type, raw_prop, + _marker: PhantomData, }) } } @@ -249,7 +279,7 @@ impl QgProperties { } } -impl Drop for QgProperties { +impl Drop for QgProperties { fn drop(&mut self) { if !self.raw_prop.is_null() { unsafe { sys::ngt_destroy_property(self.raw_prop) }; @@ -257,3 +287,27 @@ impl Drop for QgProperties { } } } + +#[derive(Debug, Clone, PartialEq)] +pub struct QgQuantizationParams { + pub dimension_of_subvector: f32, + pub max_number_of_edges: u64, +} + +impl Default for QgQuantizationParams { + fn default() -> Self { + Self { + dimension_of_subvector: 0.0, + max_number_of_edges: 128, + } + } +} + +impl QgQuantizationParams { + pub(crate) fn into_raw(self) -> sys::NGTQGQuantizationParameters { + sys::NGTQGQuantizationParameters { + dimension_of_subvector: self.dimension_of_subvector, + max_number_of_edges: self.max_number_of_edges, + } + } +} From 182747c4189b4b98a598bff2297268ed088ee1bf Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Wed, 5 Jul 2023 01:39:30 +0200 Subject: [PATCH 10/11] Improve typed API --- .github/workflows/ci.yaml | 2 + Cargo.toml | 5 +- README.md | 6 +- ngt-sys/Cargo.toml | 3 +- ngt-sys/NGT | 2 +- ngt-sys/build.rs | 5 +- src/ngt/index.rs | 71 ++++++++++---- src/qbg/index.rs | 200 +++++++++++++++++++++++++++++++------- src/qbg/properties.rs | 49 +++++++++- src/qg/index.rs | 171 +++++++++++++++++++++++++++++--- src/qg/properties.rs | 20 ++-- 11 files changed, 450 insertions(+), 84 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 910646f..3ea9141 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -21,10 +21,12 @@ jobs: - shared_mem - large_data - quantized + - quantized,qg_optim - large_data,shared_mem - large_data,quantized - static - static,quantized + - static,quantized,qg_optim - static,shared_mem,large_data steps: - uses: actions/checkout@v3 diff --git a/Cargo.toml b/Cargo.toml index 9257b7c..50f9742 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -12,7 +12,7 @@ readme = "README.md" [dependencies] half = "2" -ngt-sys = { path = "ngt-sys", version = "2.0.12" } +ngt-sys = { path = "ngt-sys", version = "2.1.2" } num_enum = "0.5" scopeguard = "1" @@ -22,8 +22,9 @@ rayon = "1" tempfile = "3" [features] -default = ["quantized"] # TODO: should not be default +default = ["quantized", "qg_optim"] # TODO: should not be default static = ["ngt-sys/static"] shared_mem = ["ngt-sys/shared_mem"] large_data = ["ngt-sys/large_data"] quantized = ["ngt-sys/quantized"] +qg_optim = ["quantized", "ngt-sys/qg_optim"] diff --git a/README.md b/README.md index d127f19..376afd1 100644 --- a/README.md +++ b/README.md @@ -14,9 +14,10 @@ This crate provides the following indexes: * `QgIndex`: Quantized graph-based index[^2] * `QbgIndex`: Quantized blob graph-based index -The quantized indexes are available through the `quantized` Cargo feature. Note that +Both quantized indexes are available through the `quantized` Cargo feature. Note that they rely on `BLAS` and `LAPACK` which thus have to be installed locally. The CPU -running the code must also support `AVX2` instructions. +running the code must also support `AVX2` instructions. Furthermore, `QgIndex` +performances can be [improved][qg-optim] by using the `qg_optim` Cargo feature. The `NgtIndex` default implementation is an ANNG, it can be optimized[^3] or converted to an ONNG through the [`optim`][ngt-optim] module. @@ -95,6 +96,7 @@ index.persist()?; [ngt-largedata]: https://github.com/yahoojapan/NGT#large-scale-data-use [ngt-ci]: https://github.com/lerouxrgd/ngt-rs/blob/master/.github/workflows/ci.yaml [ngt-optim]: https://docs.rs/ngt/latest/ngt/optim/index.html +[qg-optim]: https://github.com/yahoojapan/NGT#build-parameters-1 [^1]: https://opensource.com/article/19/10/ngt-open-source-library [^2]: https://medium.com/@masajiro.iwasaki/fusion-of-graph-based-indexing-and-product-quantization-for-ann-search-7d1f0336d0d0 diff --git a/ngt-sys/Cargo.toml b/ngt-sys/Cargo.toml index 9d3562b..3a628a5 100644 --- a/ngt-sys/Cargo.toml +++ b/ngt-sys/Cargo.toml @@ -1,6 +1,6 @@ [package] name = "ngt-sys" -version = "2.0.12" +version = "2.1.2" authors = ["Romain Leroux "] edition = "2021" links = "ngt" @@ -20,3 +20,4 @@ static = ["dep:cpp_build"] shared_mem = [] large_data = [] quantized = [] +qg_optim = [] diff --git a/ngt-sys/NGT b/ngt-sys/NGT index 5920ef0..eaee806 160000 --- a/ngt-sys/NGT +++ b/ngt-sys/NGT @@ -1 +1 @@ -Subproject commit 5920ef01016998e96ad106284ed35e684981250d +Subproject commit eaee8063c8bc5145670985f89341f94f9f6d5349 diff --git a/ngt-sys/build.rs b/ngt-sys/build.rs index 8e427fa..0387d8f 100644 --- a/ngt-sys/build.rs +++ b/ngt-sys/build.rs @@ -14,8 +14,11 @@ fn main() { if env::var("CARGO_FEATURE_QUANTIZED").is_err() { config.define("NGT_QBG_DISABLED", "ON"); } else { - config.define("NGT_AVX2", "ON"); config.define("CMAKE_BUILD_TYPE", "Release"); + if env::var("CARGO_FEATURE_QG_OPTIM").is_ok() { + config.define("NGTQG_NO_ROTATION", "ON"); + config.define("NGTQG_ZERO_GLOBAL", "ON"); + } } let dst = config.build(); diff --git a/src/ngt/index.rs b/src/ngt/index.rs index 388bbe8..619f9b4 100644 --- a/src/ngt/index.rs +++ b/src/ngt/index.rs @@ -132,7 +132,7 @@ where let rsize = sys::ngt_get_result_size(results, self.ebuf); let mut ret = Vec::with_capacity(rsize as usize); - for i in 0..rsize as u32 { + for i in 0..rsize { let d = sys::ngt_get_result(results, i, self.ebuf); if d.id == 0 && d.distance == 0.0 { Err(make_err(self.ebuf))? @@ -173,7 +173,7 @@ where let rsize = sys::ngt_get_result_size(results, self.ebuf); let mut ret = Vec::with_capacity(rsize as usize); - for i in 0..rsize as u32 { + for i in 0..rsize { let d = sys::ngt_get_result(results, i, self.ebuf); if d.id == 0 && d.distance == 0.0 { Err(make_err(self.ebuf))? @@ -198,13 +198,13 @@ where let id = match self.prop.object_type { NgtObject::Float => sys::ngt_insert_index_as_float( self.index, - vec.as_mut_ptr() as *mut _, + vec.as_mut_ptr() as *mut f32, self.prop.dimension as u32, self.ebuf, ), NgtObject::Uint8 => sys::ngt_insert_index_as_uint8( self.index, - vec.as_mut_ptr() as *mut _, + vec.as_mut_ptr() as *mut u8, self.prop.dimension as u32, self.ebuf, ), @@ -227,7 +227,7 @@ where /// discoverable yet. /// /// **The method [`build`](NgtIndex::build) must be called after inserting vectors**. - pub fn insert_batch(&mut self, batch: Vec>) -> Result<()> { + pub fn insert_batch(&mut self, batch: Vec>) -> Result<()> { let batch_size = u32::try_from(batch.len())?; if batch_size > 0 { @@ -243,9 +243,38 @@ where } unsafe { - let mut batch = batch.into_iter().flatten().collect::>(); - if !sys::ngt_batch_append_index(self.index, batch.as_mut_ptr(), batch_size, self.ebuf) { - Err(make_err(self.ebuf))? + let mut batch = batch.into_iter().flatten().collect::>(); + match self.prop.object_type { + NgtObject::Float => { + if !sys::ngt_batch_append_index( + self.index, + batch.as_mut_ptr() as *mut f32, + batch_size, + self.ebuf, + ) { + Err(make_err(self.ebuf))? + } + } + NgtObject::Uint8 => { + if !sys::ngt_batch_append_index_as_uint8( + self.index, + batch.as_mut_ptr() as *mut u8, + batch_size, + self.ebuf, + ) { + Err(make_err(self.ebuf))? + } + } + NgtObject::Float16 => { + if !sys::ngt_batch_append_index_as_float16( + self.index, + batch.as_mut_ptr() as *mut _, + batch_size, + self.ebuf, + ) { + Err(make_err(self.ebuf))? + } + } } Ok(()) } @@ -367,6 +396,7 @@ mod tests { use std::iter; use std::result::Result as StdResult; + use half::f16; use rayon::prelude::*; use tempfile::tempdir; @@ -374,7 +404,7 @@ mod tests { use crate::EPSILON; #[test] - fn test_basics() -> StdResult<(), Box> { + fn test_ngt_f32_basics() -> StdResult<(), Box> { // Get a temporary directory to store the index let dir = tempdir()?; if cfg!(feature = "shared_mem") { @@ -441,7 +471,7 @@ mod tests { } #[test] - fn test_batch() -> StdResult<(), Box> { + fn test_ngt_batch() -> StdResult<(), Box> { // Get a temporary directory to store the index let dir = tempdir()?; if cfg!(feature = "shared_mem") { @@ -466,7 +496,7 @@ mod tests { } #[test] - fn test_u8() -> StdResult<(), Box> { + fn test_ngt_u8() -> StdResult<(), Box> { // Get a temporary directory to store the index let dir = tempdir()?; if cfg!(feature = "shared_mem") { @@ -477,8 +507,9 @@ mod tests { let prop = NgtProperties::::dimension(3)?; let mut index = NgtIndex::create(dir.path(), prop)?; - // Batch insert 2 vectors, build and persist the index - index.insert_batch(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]])?; + // Insert 3 vectors, build and persist the index + index.insert_batch(vec![vec![1, 2, 3], vec![4, 5, 6]])?; + index.insert(vec![7, 8, 9])?; index.build(2)?; index.persist()?; @@ -491,7 +522,7 @@ mod tests { } #[test] - fn test_f16() -> StdResult<(), Box> { + fn test_ngt_f16() -> StdResult<(), Box> { // Get a temporary directory to store the index let dir = tempdir()?; if cfg!(feature = "shared_mem") { @@ -499,11 +530,15 @@ mod tests { } // Create an index for vectors of dimension 3 - let prop = NgtProperties::::dimension(3)?; + let prop = NgtProperties::::dimension(3)?; let mut index = NgtIndex::create(dir.path(), prop)?; - // Batch insert 2 vectors, build and persist the index - index.insert_batch(vec![vec![1.0, 2.0, 3.0], vec![4.0, 5.0, 6.0]])?; + // Insert 3 vectors, build and persist the index + index.insert_batch(vec![ + vec![1.0, 2.0, 3.0].into_iter().map(f16::from_f32).collect(), + vec![4.0, 5.0, 6.0].into_iter().map(f16::from_f32).collect(), + ])?; + index.insert(vec![7.0, 8.0, 9.0].into_iter().map(f16::from_f32).collect())?; index.build(2)?; index.persist()?; @@ -516,7 +551,7 @@ mod tests { } #[test] - fn test_multithreaded() -> StdResult<(), Box> { + fn test_ngt_multithreaded() -> StdResult<(), Box> { // Get a temporary directory to store the index let dir = tempdir()?; if cfg!(feature = "shared_mem") { diff --git a/src/qbg/index.rs b/src/qbg/index.rs index 44d9b29..053683e 100644 --- a/src/qbg/index.rs +++ b/src/qbg/index.rs @@ -4,6 +4,7 @@ use std::os::unix::ffi::OsStrExt; use std::path::Path; use std::{mem, ptr}; +use half::f16; use ngt_sys as sys; use scopeguard::defer; @@ -17,7 +18,6 @@ pub struct QbgIndex { pub(crate) index: sys::QBGIndex, path: CString, _mode: M, - obj_type: QbgObject, dimension: u32, ebuf: sys::NGTError, _marker: PhantomData, @@ -62,23 +62,39 @@ where path, _mode: ModeWrite, dimension, - obj_type: T::as_obj(), ebuf: sys::ngt_create_error_object(), _marker: PhantomData, }) } } - // TODO: should be mut vec: Vec - pub fn insert(&mut self, mut vec: Vec) -> Result { + pub fn insert(&mut self, mut vec: Vec) -> Result { unsafe { - let id = - sys::qbg_append_object(self.index, vec.as_mut_ptr(), self.dimension, self.ebuf); + let id = match T::as_obj() { + QbgObject::Float => sys::qbg_append_object( + self.index, + vec.as_mut_ptr() as *mut _, + self.dimension, + self.ebuf, + ), + QbgObject::Uint8 => sys::qbg_append_object_as_uint8( + self.index, + vec.as_mut_ptr() as *mut _, + self.dimension, + self.ebuf, + ), + QbgObject::Float16 => sys::qbg_append_object_as_float16( + self.index, + vec.as_mut_ptr() as *mut _, + self.dimension, + self.ebuf, + ), + }; if id == 0 { Err(make_err(self.ebuf))? + } else { + Ok(id) } - - Ok(id) } } @@ -145,7 +161,6 @@ where index, path, _mode: ModeRead, - obj_type: T::as_obj(), dimension, ebuf: sys::ngt_create_error_object(), _marker: PhantomData, @@ -153,7 +168,7 @@ where } } - pub fn search(&self, query: QbgQuery) -> Result> { + pub fn search(&self, query: QbgQuery) -> Result> { unsafe { let results = sys::ngt_create_empty_results(self.ebuf); if results.is_null() { @@ -161,14 +176,40 @@ where } defer! { sys::qbg_destroy_results(results); } - if !sys::qbg_search_index(self.index, query.into_raw(), results, self.ebuf) { - Err(make_err(self.ebuf))? + match T::as_obj() { + QbgObject::Float => { + let q = sys::QBGQueryFloat { + query: query.query.as_ptr() as *mut f32, + params: query.params(), + }; + if !sys::qbg_search_index_float(self.index, q, results, self.ebuf) { + Err(make_err(self.ebuf))? + } + } + QbgObject::Uint8 => { + let q = sys::QBGQueryUint8 { + query: query.query.as_ptr() as *mut u8, + params: query.params(), + }; + if !sys::qbg_search_index_uint8(self.index, q, results, self.ebuf) { + Err(make_err(self.ebuf))? + } + } + QbgObject::Float16 => { + let q = sys::QBGQueryFloat16 { + query: query.query.as_ptr() as *mut _, + params: query.params(), + }; + if !sys::qbg_search_index_float16(self.index, q, results, self.ebuf) { + Err(make_err(self.ebuf))? + } + } } let rsize = sys::qbg_get_result_size(results, self.ebuf); let mut ret = Vec::with_capacity(rsize as usize); - for i in 0..rsize as u32 { + for i in 0..rsize { let d = sys::qbg_get_result(results, i, self.ebuf); if d.id == 0 && d.distance == 0.0 { Err(make_err(self.ebuf))? @@ -206,7 +247,6 @@ where index, path, _mode: ModeWrite, - obj_type: T::as_obj(), dimension, ebuf: sys::ngt_create_error_object(), _marker: PhantomData, @@ -223,38 +263,41 @@ where /// Get the specified vector. pub fn get_vec(&self, id: VecId) -> Result> { unsafe { - match self.obj_type { + match T::as_obj() { QbgObject::Float => { let results = sys::qbg_get_object(self.index, id, self.ebuf); if results.is_null() { Err(make_err(self.ebuf))? } - let results = Vec::from_raw_parts( results as *mut f32, self.dimension as usize, self.dimension as usize, ); - let results = mem::ManuallyDrop::new(results); - - let results = results.iter().copied().collect::>(); Ok(mem::transmute::<_, Vec>(results)) } QbgObject::Uint8 => { - // TODO: Would need some kind of qbg_get_object_as_integer - let results = sys::qbg_get_object(self.index, id, self.ebuf); + let results = sys::qbg_get_object_as_uint8(self.index, id, self.ebuf); if results.is_null() { Err(make_err(self.ebuf))? } - let results = Vec::from_raw_parts( - results as *mut f32, + results as *mut u8, + self.dimension as usize, + self.dimension as usize, + ); + Ok(mem::transmute::<_, Vec>(results)) + } + QbgObject::Float16 => { + let results = sys::qbg_get_object_as_float16(self.index, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + let results = Vec::from_raw_parts( + results as *mut f16, self.dimension as usize, self.dimension as usize, ); - let results = mem::ManuallyDrop::new(results); - - let results = results.iter().copied().collect::>(); Ok(mem::transmute::<_, Vec>(results)) } } @@ -294,8 +337,8 @@ impl private::Sealed for ModeWrite {} impl IndexMode for ModeWrite {} #[derive(Debug, Clone, PartialEq)] -pub struct QbgQuery<'a> { - query: &'a [f32], +pub struct QbgQuery<'a, T> { + query: &'a [T], pub size: u64, pub epsilon: f32, pub blob_epsilon: f32, @@ -305,8 +348,11 @@ pub struct QbgQuery<'a> { pub radius: f32, } -impl<'a> QbgQuery<'a> { - pub fn new(query: &'a [f32]) -> Self { +impl<'a, T> QbgQuery<'a, T> +where + T: QbgObjectType, +{ + pub fn new(query: &'a [T]) -> Self { Self { query, size: 20, @@ -354,9 +400,8 @@ impl<'a> QbgQuery<'a> { self } - unsafe fn into_raw(self) -> sys::QBGQuery { - sys::QBGQuery { - query: self.query.as_ptr() as *mut f32, + unsafe fn params(&self) -> sys::QBGQueryParameters { + sys::QBGQueryParameters { number_of_results: self.size, epsilon: self.epsilon, blob_epsilon: self.blob_epsilon, @@ -379,7 +424,7 @@ mod tests { use super::*; #[test] - fn test_qbg() -> StdResult<(), Box> { + fn test_qbg_f32() -> StdResult<(), Box> { // Get a temporary directory to store the index let dir = tempdir()?; std::fs::remove_dir(dir.path())?; @@ -389,7 +434,7 @@ mod tests { let mut index = QbgIndex::create(dir.path(), QbgConstructParams::dimension(ndims))?; // Insert vectors and get their ids - let nvecs = 16; + let nvecs = 64; let ids = (1..ndims * nvecs) .step_by(ndims as usize) .map(|i| i as f32) @@ -418,4 +463,89 @@ mod tests { dir.close()?; Ok(()) } + + #[test] + fn test_qbg_f16() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + std::fs::remove_dir(dir.path())?; + + // Create a QGB index + let ndims = 3; + let mut index = QbgIndex::create(dir.path(), QbgConstructParams::dimension(ndims))?; + + // Insert vectors and get their ids + let nvecs = 64; + let ids = (1..ndims * nvecs) + .step_by(ndims as usize) + .map(|i| f16::from_f32(i as f32)) + .map(|i| { + repeat(i) + .zip((0..ndims).map(|j| f16::from_f32(j as f32))) + .map(|(i, j)| i + j) + .collect() + }) + .map(|vector| index.insert(vector)) + .collect::>>()?; + + // Build and persist the index + index.build(QbgBuildParams::default())?; + index.persist()?; + + let index = index.into_readable()?; + + // Perform a vector search (with 2 results) + let v: Vec = (1..=ndims) + .into_iter() + .map(|x| f16::from_f32(x as f32)) + .collect(); + let query = QbgQuery::new(&v).size(2); + let res = index.search(query)?; + assert_eq!(ids[0], res[0].id); + assert_eq!(v, index.get_vec(ids[0])?); + + dir.close()?; + Ok(()) + } + + #[test] + fn test_qbg_u8() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + std::fs::remove_dir(dir.path())?; + + // Create a QGB index + let ndims = 3; + let mut index = QbgIndex::create(dir.path(), QbgConstructParams::dimension(ndims))?; + + // Insert vectors and get their ids + let nvecs = 64; + let ids = (1..ndims * nvecs) + .step_by(ndims as usize) + .map(|i| i as u8) + .map(|i| { + repeat(i) + .zip((0..ndims).map(|j| j as u8)) + .map(|(i, j)| i + j) + .collect() + }) + .map(|vector| index.insert(vector)) + .collect::>>()?; + + // Build and persist the index + index.build(QbgBuildParams::default())?; + index.persist()?; + + let index = index.into_readable()?; + + // Perform a vector search (with 3 results) + let v: Vec = (1..=ndims).into_iter().map(|x| x as u8).collect(); + let query = QbgQuery::new(&v).size(3); + let res = index.search(query)?; + assert!(Vec::from_iter(res[0..3].iter().map(|r| r.id)).contains(&ids[0])); + assert!(v == index.get_vec(ids[0])?); + + dir.close()?; + Ok(()) + } } diff --git a/src/qbg/properties.rs b/src/qbg/properties.rs index 08cbe79..9d358b5 100644 --- a/src/qbg/properties.rs +++ b/src/qbg/properties.rs @@ -1,13 +1,17 @@ use std::marker::PhantomData; +use half::f16; use ngt_sys as sys; use num_enum::TryFromPrimitive; +use crate::error::Error; + #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] pub enum QbgObject { Uint8 = 0, Float = 1, + Float16 = 2, } mod private { @@ -32,6 +36,13 @@ impl QbgObjectType for u8 { } } +impl private::Sealed for f16 {} +impl QbgObjectType for f16 { + fn as_obj() -> QbgObject { + QbgObject::Float16 + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] pub enum QbgDistance { @@ -55,10 +66,10 @@ where T: QbgObjectType, { pub fn dimension(dimension: u64) -> Self { - let extended_dimension = 0; + let extended_dimension = next_multiple_of_16(dimension); let number_of_subvectors = 1; let number_of_blobs = 0; - let internal_data_type = QbgObject::Float; // TODO: Should be T::as_obj() ? + let internal_data_type = T::as_obj(); let data_type = T::as_obj(); let distance_type = QbgDistance::L2; @@ -74,9 +85,16 @@ where } } - pub fn extended_dimension(mut self, extended_dimension: u64) -> Self { - self.extended_dimension = extended_dimension; - self + pub fn extended_dimension(mut self, extended_dimension: u64) -> Result { + if extended_dimension % 16 == 0 && extended_dimension >= self.dimension { + self.extended_dimension = extended_dimension; + Ok(self) + } else { + Err(Error(format!( + "Invalid extended_dimension: {}, must be a multiple of 16 greater or equal to dimension", + extended_dimension + ))) + } } pub fn number_of_subvectors(mut self, number_of_subvectors: u64) -> Self { @@ -112,6 +130,10 @@ where } } +fn next_multiple_of_16(x: u64) -> u64 { + ((x + 15) / 16) * 16 +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] pub enum QbgClusteringInitMode { @@ -258,3 +280,20 @@ impl QbgBuildParams { } } } + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn test_qbg_params() { + let params = QbgConstructParams::::dimension(3); + assert_eq!(params.extended_dimension, 16); + + let params = QbgConstructParams::::dimension(16); + assert_eq!(params.extended_dimension, 16); + + let params = QbgConstructParams::::dimension(513); + assert_eq!(params.extended_dimension, 528); + } +} diff --git a/src/qg/index.rs b/src/qg/index.rs index 9f142aa..7c0589a 100644 --- a/src/qg/index.rs +++ b/src/qg/index.rs @@ -4,6 +4,7 @@ use std::os::unix::ffi::OsStrExt; use std::path::Path; use std::ptr; +use half::f16; use ngt_sys as sys; use scopeguard::defer; @@ -87,14 +88,40 @@ where } defer! { sys::ngt_destroy_results(results); } - if !sys::ngtqg_search_index(self.index, query.into_raw(), results, self.ebuf) { - Err(make_err(self.ebuf))? + match T::as_obj() { + QgObject::Float => { + let q = sys::NGTQGQueryFloat { + query: query.query.as_ptr() as *mut f32, + params: query.params(), + }; + if !sys::ngtqg_search_index_float(self.index, q, results, self.ebuf) { + Err(make_err(self.ebuf))? + } + } + QgObject::Uint8 => { + let q = sys::NGTQGQueryUint8 { + query: query.query.as_ptr() as *mut u8, + params: query.params(), + }; + if !sys::ngtqg_search_index_uint8(self.index, q, results, self.ebuf) { + Err(make_err(self.ebuf))? + } + } + QgObject::Float16 => { + let q = sys::NGTQGQueryFloat16 { + query: query.query.as_ptr() as *mut _, + params: query.params(), + }; + if !sys::ngtqg_search_index_float16(self.index, q, results, self.ebuf) { + Err(make_err(self.ebuf))? + } + } } let rsize = sys::ngt_get_result_size(results, self.ebuf); let mut ret = Vec::with_capacity(rsize as usize); - for i in 0..rsize as u32 { + for i in 0..rsize { let d = sys::ngt_get_result(results, i, self.ebuf); if d.id == 0 && d.distance == 0.0 { Err(make_err(self.ebuf))? @@ -153,6 +180,27 @@ where ); let results = mem::ManuallyDrop::new(results); + let results = results.iter().copied().collect::>(); + Ok(mem::transmute::<_, Vec>(results)) + } + QgObject::Float16 => { + let ospace = sys::ngt_get_object_space(self.index, self.ebuf); + if ospace.is_null() { + Err(make_err(self.ebuf))? + } + + let results = sys::ngt_get_object_as_float16(ospace, id, self.ebuf); + if results.is_null() { + Err(make_err(self.ebuf))? + } + + let results = Vec::from_raw_parts( + results as *mut f16, + self.prop.dimension as usize, + self.prop.dimension as usize, + ); + let results = mem::ManuallyDrop::new(results); + let results = results.iter().copied().collect::>(); Ok(mem::transmute::<_, Vec>(results)) } @@ -183,7 +231,10 @@ pub struct QgQuery<'a, T> { pub radius: f32, } -impl<'a, T> QgQuery<'a, T> { +impl<'a, T> QgQuery<'a, T> +where + T: QgObjectType, +{ pub fn new(query: &'a [T]) -> Self { Self { query, @@ -214,9 +265,8 @@ impl<'a, T> QgQuery<'a, T> { self } - unsafe fn into_raw(self) -> sys::NGTQGQuery { - sys::NGTQGQuery { - query: self.query.as_ptr() as *mut f32, + unsafe fn params(&self) -> sys::NGTQGQueryParameters { + sys::NGTQGQueryParameters { size: self.size, epsilon: self.epsilon, result_expansion: self.result_expansion, @@ -237,7 +287,102 @@ mod tests { use crate::{NgtDistance, NgtProperties}; #[test] - fn test_qg() -> StdResult<(), Box> { + fn test_qg_f32() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + + // Create an NGT index for vectors + let ndims = 3; + let props = NgtProperties::::dimension(ndims)?.distance_type(NgtDistance::L2)?; + let mut index = NgtIndex::create(dir.path(), props)?; + + // Insert vectors and get their ids + let nvecs = 64; + let ids = (1..ndims * nvecs) + .step_by(ndims) + .map(|i| i as f32) + .map(|i| { + repeat(i) + .zip((0..ndims).map(|j| j as f32)) + .map(|(i, j)| i + j) + .collect() + }) + .map(|vector| index.insert(vector)) + .collect::>>()?; + + // Build and persist the index + index.build(1)?; + index.persist()?; + + // Quantize the index + let params = QgQuantizationParams { + dimension_of_subvector: 1., + max_number_of_edges: 50, + }; + let index = QgIndex::quantize(index, params)?; + + // Perform a vector search (with 3 results) + let v: Vec = (1..=ndims).into_iter().map(|x| x as f32).collect(); + let query = QgQuery::new(&v).size(3); + let res = index.search(query)?; + assert!(ids[0] == res[0].id); + assert!(v == index.get_vec(ids[0])?); + + dir.close()?; + Ok(()) + } + + #[test] + fn test_qg_f16() -> StdResult<(), Box> { + // Get a temporary directory to store the index + let dir = tempdir()?; + + // Create an NGT index for vectors + let ndims = 3; + let props = NgtProperties::::dimension(ndims)?.distance_type(NgtDistance::L2)?; + let mut index = NgtIndex::create(dir.path(), props)?; + + // Insert vectors and get their ids + let nvecs = 64; + let ids = (1..ndims * nvecs) + .step_by(ndims) + .map(|i| f16::from_f32(i as f32)) + .map(|i| { + repeat(i) + .zip((0..ndims).map(|j| f16::from_f32(j as f32))) + .map(|(i, j)| i + j) + .collect() + }) + .map(|vector| index.insert(vector)) + .collect::>>()?; + + // Build and persist the index + index.build(1)?; + index.persist()?; + + // Quantize the index + let params = QgQuantizationParams { + dimension_of_subvector: 1., + max_number_of_edges: 50, + }; + let index = QgIndex::quantize(index, params)?; + + // Perform a vector search (with 3 results) + let v: Vec = (1..=ndims) + .into_iter() + .map(|x| f16::from_f32(x as f32)) + .collect(); + let query = QgQuery::new(&v).size(3); + let res = index.search(query)?; + assert!(ids[0] == res[0].id); + assert!(v == index.get_vec(ids[0])?); + + dir.close()?; + Ok(()) + } + + #[test] + fn test_qg_u8() -> StdResult<(), Box> { // Get a temporary directory to store the index let dir = tempdir()?; @@ -271,12 +416,12 @@ mod tests { }; let index = QgIndex::quantize(index, params)?; - // Perform a vector search (with 2 results) + // Perform a vector search (with 3 results) let v: Vec = (1..=ndims).into_iter().map(|x| x as u8).collect(); - let query = QgQuery::new(&v).size(2); - let res = index.search(query)?; - assert_eq!(ids[0], res[0].id); - assert_eq!(v, index.get_vec(ids[0])?); + let query = QgQuery::new(&v).size(3); + let res = &index.search(query)?; + assert!(Vec::from_iter(res[0..3].iter().map(|r| r.id)).contains(&ids[0])); + assert!(v == index.get_vec(ids[0])?); dir.close()?; Ok(()) diff --git a/src/qg/properties.rs b/src/qg/properties.rs index 4cb23ea..77b74f5 100644 --- a/src/qg/properties.rs +++ b/src/qg/properties.rs @@ -1,6 +1,7 @@ use std::marker::PhantomData; use std::ptr; +use half::f16; use ngt_sys as sys; use num_enum::TryFromPrimitive; use scopeguard::defer; @@ -12,6 +13,7 @@ use crate::error::{make_err, Result}; pub enum QgObject { Uint8 = 1, Float = 2, + Float16 = 3, } mod private { @@ -36,6 +38,13 @@ impl QgObjectType for u8 { } } +impl private::Sealed for f16 {} +impl QgObjectType for f16 { + fn as_obj() -> QgObject { + QgObject::Float16 + } +} + #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] pub enum QgDistance { @@ -223,12 +232,6 @@ where Ok(()) } - pub fn object_type(mut self, object_type: QgObject) -> Result { - self.object_type = object_type; - unsafe { Self::set_object_type(self.raw_prop, object_type)? }; - Ok(self) - } - unsafe fn set_object_type(raw_prop: sys::NGTProperty, object_type: QgObject) -> Result<()> { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } @@ -244,6 +247,11 @@ where Err(make_err(ebuf))? } } + QgObject::Float16 => { + if !sys::ngt_set_property_object_type_float16(raw_prop, ebuf) { + Err(make_err(ebuf))? + } + } } Ok(()) From fcfd3de8dfc496da1c823fc8ee0bca727f8a7c37 Mon Sep 17 00:00:00 2001 From: Romain Leroux Date: Fri, 18 Aug 2023 23:29:21 +0200 Subject: [PATCH 11/11] Update doc and improve QG setup --- .github/workflows/ci.yaml | 1 - Cargo.toml | 7 ++- README.md | 110 +++++++++++--------------------------- src/error.rs | 6 +++ src/lib.rs | 73 ++++++++++++++++++++++++- src/ngt/index.rs | 4 +- src/ngt/mod.rs | 65 ---------------------- src/ngt/optim.rs | 6 ++- src/qbg/index.rs | 16 +----- src/qbg/mod.rs | 64 +++++++++++++++++++++- src/qg/index.rs | 33 +++++------- src/qg/mod.rs | 68 ++++++++++++++++++++++- src/qg/properties.rs | 38 +++++++++++++ 13 files changed, 302 insertions(+), 189 deletions(-) diff --git a/.github/workflows/ci.yaml b/.github/workflows/ci.yaml index 3ea9141..7d27b72 100644 --- a/.github/workflows/ci.yaml +++ b/.github/workflows/ci.yaml @@ -17,7 +17,6 @@ jobs: matrix: os: [ubuntu-latest] feature: - - default - shared_mem - large_data - quantized diff --git a/Cargo.toml b/Cargo.toml index 50f9742..52d02c9 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -13,7 +13,7 @@ readme = "README.md" [dependencies] half = "2" ngt-sys = { path = "ngt-sys", version = "2.1.2" } -num_enum = "0.5" +num_enum = "0.7" scopeguard = "1" [dev-dependencies] @@ -22,9 +22,12 @@ rayon = "1" tempfile = "3" [features] -default = ["quantized", "qg_optim"] # TODO: should not be default static = ["ngt-sys/static"] shared_mem = ["ngt-sys/shared_mem"] large_data = ["ngt-sys/large_data"] quantized = ["ngt-sys/quantized"] qg_optim = ["quantized", "ngt-sys/qg_optim"] + +[package.metadata.docs.rs] +features = ["quantized"] +rustdoc-args = ["--cfg", "docsrs"] diff --git a/README.md b/README.md index 376afd1..3503d44 100644 --- a/README.md +++ b/README.md @@ -1,103 +1,55 @@ -# ngt-rs   [![Latest Version]][crates.io] [![Latest Doc]][docs.rs] +# ngt-rs -[Latest Version]: https://img.shields.io/crates/v/ngt.svg -[crates.io]: https://crates.io/crates/ngt -[Latest Doc]: https://docs.rs/ngt/badge.svg -[docs.rs]: https://docs.rs/ngt +[![crate]][crate-ngt] [![doc]][doc-ngt] + +[crate]: https://img.shields.io/crates/v/ngt.svg +[crate-ngt]: https://crates.io/crates/ngt +[doc]: https://docs.rs/ngt/badge.svg +[doc-ngt]: https://docs.rs/ngt Rust wrappers for [NGT][], which provides high-speed approximate nearest neighbor searches against a large volume of data in high dimensional vector data space (several -ten to several thousand dimensions). +ten to several thousand dimensions). The vector data can be `f32`, `u8`, or [f16][]. This crate provides the following indexes: -* `NgtIndex`: Graph and tree-based index[^1] -* `QgIndex`: Quantized graph-based index[^2] -* `QbgIndex`: Quantized blob graph-based index +* [`NgtIndex`][index-ngt]: Graph and tree based index[^1] +* [`QgIndex`][index-qg]: Quantized graph based index[^2] +* [`QbgIndex`][index-qbg]: Quantized blob graph based index Both quantized indexes are available through the `quantized` Cargo feature. Note that -they rely on `BLAS` and `LAPACK` which thus have to be installed locally. The CPU -running the code must also support `AVX2` instructions. Furthermore, `QgIndex` -performances can be [improved][qg-optim] by using the `qg_optim` Cargo feature. +they rely on `BLAS` and `LAPACK` which thus have to be installed locally. Furthermore, +`QgIndex` performances can be [improved][qg-optim] by using the `qg_optim` Cargo +feature. -The `NgtIndex` default implementation is an ANNG, it can be optimized[^3] or converted +The `NgtIndex` default implementation is an ANNG. It can be optimized[^3] or converted to an ONNG through the [`optim`][ngt-optim] module. By default `ngt-rs` will be built dynamically, which requires `CMake` to build NGT. This means that you'll have to make the build artifact `libngt.so` available to your final -binary (see an example in the [CI][ngt-ci]). - -However the `static` feature will build and link NGT statically. Note that `OpenMP` will -also be linked statically. If the `quantized` feature is used, then `BLAS` and `LAPACK` -libraries will also be linked statically. - -Finally, NGT's [shared memory][ngt-sharedmem] and [large dataset][ngt-largedata] -features are available through the features `shared_mem` and `large_data` respectively. - -## Usage - -Defining the properties of a new index: - -```rust,ignore -use ngt::{NgtProperties, NgtDistance}; - -// Defaut properties with vectors of dimension 3 -let prop = NgtProperties::::dimension(3)?; - -// Or customize values (here are the defaults) -let prop = NgtProperties::::dimension(3)? - .creation_edge_size(10)? - .search_edge_size(40)? - .distance_type(NgtDistance::L2)?; -``` +binary (see an example in the [CI][ngt-ci]). However the `static` feature will build and +link NGT statically. Note that `OpenMP` will also be linked statically. If the +`quantized` feature is used, then `BLAS` and `LAPACK` libraries will also be linked +statically. -Creating/Opening an index and using it: +NGT's [shared memory][ngt-sharedmem] and [large dataset][ngt-largedata] features are +available through the Cargo features `shared_mem` and `large_data` respectively. -```rust,ignore -use ngt::{NgtIndex, NgtProperties, EPSILON}; +[^1]: [Graph and tree based method explanation][ngt-desc] -// Create a new index -let prop = NgtProperties::dimension(3)?; -let index: NgtIndex = NgtIndex::create("target/path/to/index/dir", prop)?; +[^2]: [Quantized graph based method explanation][qg-desc] -// Open an existing index -let mut index = NgtIndex::open("target/path/to/index/dir")?; - -// Insert two vectors and get their id -let vec1 = vec![1.0, 2.0, 3.0]; -let vec2 = vec![4.0, 5.0, 6.0]; -let id1 = index.insert(vec1)?; -let id2 = index.insert(vec2)?; - -// Build the index in RAM (not yet persisted on disk) -// This is required in order to be able to search vectors -index.build(2)?; - -// Perform a vector search (with 1 result) -let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; -assert_eq!(res[0].id, id1); -assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); - -// Remove a vector and check that it is not present anymore -index.remove(id1)?; -let res = index.get_vec(id1); -assert!(res.is_err()); - -// Verify that now our search result is different -let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; -assert_eq!(res[0].id, id2); -assert_eq!(index.get_vec(id2)?, vec![4.0, 5.0, 6.0]); - -// Persist index on disk -index.persist()?; -``` +[^3]: [NGT index optimizations in Python][ngt-optim-py] [ngt]: https://github.com/yahoojapan/NGT +[ngt-desc]: https://opensource.com/article/19/10/ngt-open-source-library [ngt-sharedmem]: https://github.com/yahoojapan/NGT#shared-memory-use [ngt-largedata]: https://github.com/yahoojapan/NGT#large-scale-data-use [ngt-ci]: https://github.com/lerouxrgd/ngt-rs/blob/master/.github/workflows/ci.yaml [ngt-optim]: https://docs.rs/ngt/latest/ngt/optim/index.html +[ngt-optim-py]: https://github.com/yahoojapan/NGT/wiki/Optimization-Examples-Using-Python +[qg-desc]: https://medium.com/@masajiro.iwasaki/fusion-of-graph-based-indexing-and-product-quantization-for-ann-search-7d1f0336d0d0 [qg-optim]: https://github.com/yahoojapan/NGT#build-parameters-1 - -[^1]: https://opensource.com/article/19/10/ngt-open-source-library -[^2]: https://medium.com/@masajiro.iwasaki/fusion-of-graph-based-indexing-and-product-quantization-for-ann-search-7d1f0336d0d0 -[^3]: https://github.com/yahoojapan/NGT/wiki/Optimization-Examples-Using-Python +[f16]: https://docs.rs/half/latest/half/struct.f16.html +[index-ngt]: https://docs.rs/ngt/latest/ngt/#usage +[index-qg]: https://docs.rs/ngt/latest/ngt/qg/ +[index-qbg]: https://docs.rs/ngt/latest/ngt/qgb/ diff --git a/src/error.rs b/src/error.rs index 9af95c7..b648d1f 100644 --- a/src/error.rs +++ b/src/error.rs @@ -23,6 +23,12 @@ pub(crate) fn make_err(err: sys::NGTError) -> Error { Error(err_msg) } +impl From for Error { + fn from(err: String) -> Self { + Self(err) + } +} + impl From for Error { fn from(source: std::io::Error) -> Self { Self(source.to_string()) diff --git a/src/lib.rs b/src/lib.rs index 8d43036..d3f9b37 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -1,7 +1,77 @@ +#![cfg_attr(docsrs, feature(doc_auto_cfg))] #![doc = include_str!("../README.md")] +//! +//! # Usage +//! +//! Graph and tree based index (NGT Index) +//! +//! ## Defining the properties of a new NGT index: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! use ngt::{NgtProperties, NgtDistance}; +//! +//! // Defaut properties with vectors of dimension 3 +//! let prop = NgtProperties::::dimension(3)?; +//! +//! // Or customize values (here are the defaults) +//! let prop = NgtProperties::::dimension(3)? +//! .creation_edge_size(10)? +//! .search_edge_size(40)? +//! .distance_type(NgtDistance::L2)?; +//! +//! # Ok(()) +//! # } +//! ``` +//! +//! ## Creating/Opening a NGT index and using it: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! use ngt::{NgtIndex, NgtProperties}; +//! +//! // Create a new index +//! let prop = NgtProperties::dimension(3)?; +//! let index: NgtIndex = NgtIndex::create("target/path/to/ngt_index/dir", prop)?; +//! +//! // Open an existing index +//! let mut index = NgtIndex::open("target/path/to/ngt_index/dir")?; +//! +//! // Insert two vectors and get their id +//! let vec1 = vec![1.0, 2.0, 3.0]; +//! let vec2 = vec![4.0, 5.0, 6.0]; +//! let id1 = index.insert(vec1)?; +//! let id2 = index.insert(vec2)?; +//! +//! // Build the index in RAM (not yet persisted on disk) +//! // This is required in order to be able to search vectors +//! index.build(2)?; +//! +//! // Perform a vector search (with 1 result) +//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, ngt::EPSILON)?; +//! assert_eq!(res[0].id, id1); +//! assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); +//! +//! // Remove a vector and check that it is not present anymore +//! index.remove(id1)?; +//! let res = index.get_vec(id1); +//! assert!(res.is_err()); +//! +//! // Verify that now our search result is different +//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, ngt::EPSILON)?; +//! assert_eq!(res[0].id, id2); +//! assert_eq!(index.get_vec(id2)?, vec![4.0, 5.0, 6.0]); +//! +//! // Persist index on disk +//! index.persist()?; +//! +//! # std::fs::remove_dir_all("target/path/to/ngt_index/dir").unwrap(); +//! # Ok(()) +//! # } +//! ``` #[cfg(all(feature = "quantized", feature = "shared_mem"))] -compile_error!("only one of ['quantized', 'shared_mem'] can be enabled"); +compile_error!(r#"only one of ["quantized", "shared_mem"] can be enabled"#); mod error; mod ngt; @@ -23,5 +93,4 @@ pub const EPSILON: f32 = 0.1; pub use crate::error::{Error, Result}; pub use crate::ngt::{optim, NgtDistance, NgtIndex, NgtObject, NgtProperties}; -#[doc(inline)] pub use half; diff --git a/src/ngt/index.rs b/src/ngt/index.rs index 619f9b4..1bbc830 100644 --- a/src/ngt/index.rs +++ b/src/ngt/index.rs @@ -321,7 +321,7 @@ where } let results = Vec::from_raw_parts( - results as *mut f32, + results, self.prop.dimension as usize, self.prop.dimension as usize, ); @@ -353,7 +353,7 @@ where } let results = Vec::from_raw_parts( - results as *mut u8, + results, self.prop.dimension as usize, self.prop.dimension as usize, ); diff --git a/src/ngt/mod.rs b/src/ngt/mod.rs index ff9d15e..6403bd4 100644 --- a/src/ngt/mod.rs +++ b/src/ngt/mod.rs @@ -1,68 +1,3 @@ -//! Defining the properties of a new index: -//! -//! ```rust -//! # fn main() -> Result<(), ngt::Error> { -//! use ngt::{NgtProperties, NgtDistance}; -//! -//! // Defaut properties with vectors of dimension 3 -//! let prop = NgtProperties::::dimension(3)?; -//! -//! // Or customize values (here are the defaults) -//! let prop = NgtProperties::::dimension(3)? -//! .creation_edge_size(10)? -//! .search_edge_size(40)? -//! .distance_type(NgtDistance::L2)?; -//! -//! # Ok(()) -//! # } -//! ``` -//! -//! Creating/Opening an index and using it: -//! -//! ```rust -//! # fn main() -> Result<(), ngt::Error> { -//! use ngt::{NgtIndex, NgtProperties, EPSILON}; -//! -//! // Create a new index -//! let prop = NgtProperties::dimension(3)?; -//! let index: NgtIndex = NgtIndex::create("target/path/to/index/dir", prop)?; -//! -//! // Open an existing index -//! let mut index = NgtIndex::open("target/path/to/index/dir")?; -//! -//! // Insert two vectors and get their id -//! let vec1 = vec![1.0, 2.0, 3.0]; -//! let vec2 = vec![4.0, 5.0, 6.0]; -//! let id1 = index.insert(vec1)?; -//! let id2 = index.insert(vec2)?; -//! -//! // Build the index in RAM (not yet persisted on disk) -//! // This is required in order to be able to search vectors -//! index.build(2)?; -//! -//! // Perform a vector search (with 1 result) -//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; -//! assert_eq!(res[0].id, id1); -//! assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); -//! -//! // Remove a vector and check that it is not present anymore -//! index.remove(id1)?; -//! let res = index.get_vec(id1); -//! assert!(res.is_err()); -//! -//! // Verify that now our search result is different -//! let res = index.search(&vec![1.1, 2.1, 3.1], 1, EPSILON)?; -//! assert_eq!(res[0].id, id2); -//! assert_eq!(index.get_vec(id2)?, vec![4.0, 5.0, 6.0]); -//! -//! // Persist index on disk -//! index.persist()?; -//! -//! # std::fs::remove_dir_all("target/path/to/index/dir").unwrap(); -//! # Ok(()) -//! # } -//! ``` - mod index; pub mod optim; mod properties; diff --git a/src/ngt/optim.rs b/src/ngt/optim.rs index 0fe7ad2..4aa1156 100644 --- a/src/ngt/optim.rs +++ b/src/ngt/optim.rs @@ -1,3 +1,7 @@ +#![cfg_attr(feature = "shared_mem", allow(unused_imports))] + +//! Functions aimed at optimizing [`NgtIndex`](NgtIndex) + use std::ffi::CString; use std::os::unix::ffi::OsStrExt; use std::path::Path; @@ -92,7 +96,7 @@ pub fn refine_anng( /// [`optimize_anng_edges_number`](optimize_anng_edges_number). /// /// If more performance is needed, a larger `creation_edge_size` can be set through -/// [`Properties`](crate::Properties::creation_edge_size) at ANNG index +/// [`Properties`](crate::NgtProperties::creation_edge_size) at ANNG index /// [`create`](NgtIndex::create) time. /// /// Important [`GraphOptimParams`](GraphOptimParams) parameters are `nb_outgoing` edges diff --git a/src/qbg/index.rs b/src/qbg/index.rs index 053683e..a895288 100644 --- a/src/qbg/index.rs +++ b/src/qbg/index.rs @@ -31,12 +31,6 @@ where where P: AsRef, { - if !is_x86_feature_detected!("avx2") { - return Err(Error( - "Cannot quantize an index without AVX2 support".into(), - )); - } - unsafe { let ebuf = sys::ngt_create_error_object(); defer! { sys::ngt_destroy_error_object(ebuf); } @@ -132,12 +126,6 @@ where T: QbgObjectType, { pub fn open>(path: P) -> Result { - if !is_x86_feature_detected!("avx2") { - return Err(Error( - "Cannot use a quantized index without AVX2 support".into(), - )); - } - if !path.as_ref().exists() { Err(Error(format!("Path {:?} does not exist", path.as_ref())))? } @@ -270,7 +258,7 @@ where Err(make_err(self.ebuf))? } let results = Vec::from_raw_parts( - results as *mut f32, + results, self.dimension as usize, self.dimension as usize, ); @@ -282,7 +270,7 @@ where Err(make_err(self.ebuf))? } let results = Vec::from_raw_parts( - results as *mut u8, + results, self.dimension as usize, self.dimension as usize, ); diff --git a/src/qbg/mod.rs b/src/qbg/mod.rs index 0cb3da2..56a7aad 100644 --- a/src/qbg/mod.rs +++ b/src/qbg/mod.rs @@ -1,4 +1,66 @@ -// TODO: Add module doc (specify available types) +//! Quantized blob graph index (QBG Index) +//! +//! ## Defining the properties of a new QBG index: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! use ngt::qbg::{QbgConstructParams, QbgDistance}; +//! +//! // Defaut parameters with vectors of dimension 3 +//! let params = QbgConstructParams::::dimension(3); +//! +//! // Or customize values (here are the defaults) +//! let params = QbgConstructParams::::dimension(3) +//! .extended_dimension(16)? // next multiple of 16 after 3 +//! .number_of_subvectors(1) +//! .number_of_subvectors(0) +//! .distance_type(QbgDistance::L2); +//! +//! # Ok(()) +//! # } +//! ``` +//! +//! ## Creating/Opening a QBG index and using it: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! # std::fs::create_dir_all("target/path/to/qbg_index").unwrap(); +//! use ngt::qbg::{ +//! ModeRead, ModeWrite, QbgBuildParams, QbgConstructParams, QbgDistance, QbgIndex, QbgQuery, +//! }; +//! +//! // Create a new index +//! let params = QbgConstructParams::dimension(3); +//! let mut index: QbgIndex = +//! QbgIndex::create("target/path/to/qbg_index/dir", params)?; +//! +//! // Insert vectors and get their id +//! let vec1 = vec![1.0, 2.0, 3.0]; +//! let vec2 = vec![4.0, 5.0, 6.0]; +//! let id1 = index.insert(vec1)?; +//! let id2 = index.insert(vec2)?; +//! +//! // Add enough dummy vectors to build an index +//! for i in 0..64 { +//! index.insert(vec![100. + i as f32; 3])?; +//! } +//! // Build the index in RAM and persist it on disk +//! index.build(QbgBuildParams::default())?; +//! index.persist()?; +//! +//! // Open an existing index +//! let index: QbgIndex = QbgIndex::open("target/path/to/qbg_index/dir")?; +//! +//! // Perform a vector search (with 1 result) +//! let query = vec![1.1, 2.1, 3.1]; +//! let res = index.search(QbgQuery::new(&query).size(1))?; +//! assert_eq!(res[0].id, id1); +//! assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); +//! +//! # std::fs::remove_dir_all("target/path/to/qbg_index").unwrap(); +//! # Ok(()) +//! # } +//! ``` mod index; mod properties; diff --git a/src/qg/index.rs b/src/qg/index.rs index 7c0589a..be68cb3 100644 --- a/src/qg/index.rs +++ b/src/qg/index.rs @@ -11,6 +11,7 @@ use scopeguard::defer; use super::{QgObject, QgObjectType, QgProperties, QgQuantizationParams}; use crate::error::{make_err, Error, Result}; use crate::ngt::NgtIndex; +use crate::qg::QgDistance; use crate::{SearchResult, VecId}; #[derive(Debug)] @@ -26,12 +27,7 @@ where { /// Quantize an NGT index pub fn quantize(index: NgtIndex, params: QgQuantizationParams) -> Result { - // - if !is_x86_feature_detected!("avx2") { - return Err(Error( - "Cannot quantize an index without AVX2 support".into(), - )); - } + QgDistance::try_from(index.prop.distance_type)?; unsafe { let ebuf = sys::ngt_create_error_object(); @@ -49,12 +45,6 @@ where /// Open the already existing quantized index at the specified path. pub fn open>(path: P) -> Result { - if !is_x86_feature_detected!("avx2") { - return Err(Error( - "Cannot use a quantized index without AVX2 support".into(), - )); - } - if !path.as_ref().exists() { Err(Error(format!("Path {:?} does not exist", path.as_ref())))? } @@ -153,7 +143,7 @@ where } let results = Vec::from_raw_parts( - results as *mut f32, + results, self.prop.dimension as usize, self.prop.dimension as usize, ); @@ -174,7 +164,7 @@ where } let results = Vec::from_raw_parts( - results as *mut u8, + results, self.prop.dimension as usize, self.prop.dimension as usize, ); @@ -283,8 +273,9 @@ mod tests { use tempfile::tempdir; + use crate::qg::QgDistance; + use super::*; - use crate::{NgtDistance, NgtProperties}; #[test] fn test_qg_f32() -> StdResult<(), Box> { @@ -293,8 +284,8 @@ mod tests { // Create an NGT index for vectors let ndims = 3; - let props = NgtProperties::::dimension(ndims)?.distance_type(NgtDistance::L2)?; - let mut index = NgtIndex::create(dir.path(), props)?; + let props = QgProperties::::dimension(ndims)?.distance_type(QgDistance::L2)?; + let mut index = NgtIndex::create(dir.path(), props.try_into()?)?; // Insert vectors and get their ids let nvecs = 64; @@ -339,8 +330,8 @@ mod tests { // Create an NGT index for vectors let ndims = 3; - let props = NgtProperties::::dimension(ndims)?.distance_type(NgtDistance::L2)?; - let mut index = NgtIndex::create(dir.path(), props)?; + let props = QgProperties::::dimension(ndims)?.distance_type(QgDistance::L2)?; + let mut index = NgtIndex::create(dir.path(), props.try_into()?)?; // Insert vectors and get their ids let nvecs = 64; @@ -388,8 +379,8 @@ mod tests { // Create an NGT index for vectors let ndims = 3; - let props = NgtProperties::::dimension(ndims)?.distance_type(NgtDistance::L2)?; - let mut index = NgtIndex::create(dir.path(), props)?; + let props = QgProperties::::dimension(ndims)?.distance_type(QgDistance::L2)?; + let mut index = NgtIndex::create(dir.path(), props.try_into()?)?; // Insert vectors and get their ids let nvecs = 64; diff --git a/src/qg/mod.rs b/src/qg/mod.rs index eb774b0..9c22c13 100644 --- a/src/qg/mod.rs +++ b/src/qg/mod.rs @@ -1,4 +1,70 @@ -// TODO: Add module doc (specify available types) +//! Quantized graph index (QG Index) +//! +//! ## Defining the properties of a new QG index: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! use ngt::qg::{QgProperties, QgDistance}; +//! +//! // Defaut properties with vectors of dimension 3 +//! let prop = QgProperties::::dimension(3)?; +//! +//! // Or customize values (here are the defaults) +//! let prop = QgProperties::::dimension(3)? +//! .creation_edge_size(10)? +//! .search_edge_size(40)? +//! .distance_type(QgDistance::L2)?; +//! +//! # Ok(()) +//! # } +//! ``` +//! +//! ## Creating/Opening a QG index and using it: +//! +//! ```rust +//! # fn main() -> Result<(), ngt::Error> { +//! use ngt::NgtIndex; +//! use ngt::qg::{QgDistance, QgIndex, QgProperties, QgQuantizationParams, QgQuery}; +//! +//! // Create a new quantizable NGT index +//! let prop = QgProperties::dimension(3)?; +//! let mut index: NgtIndex = +//! NgtIndex::create("target/path/to/qg_index/dir", prop.try_into()?)?; +//! +//! // Insert two vectors and get their id +//! let vec1 = vec![1.0, 2.0, 3.0]; +//! let vec2 = vec![4.0, 5.0, 6.0]; +//! let id1 = index.insert(vec1)?; +//! let id2 = index.insert(vec2)?; +//! +//! // Add enough dummy vectors to build an index +//! for i in 0..64 { +//! index.insert(vec![100. + i as f32; 3])?; +//! } +//! // Build the index in RAM and persist it on disk +//! index.build(1)?; +//! index.persist()?; +//! +//! // Quantize the NGT index +//! let params = QgQuantizationParams { +//! dimension_of_subvector: 1., +//! max_number_of_edges: 50, +//! }; +//! let index = QgIndex::quantize(index, params)?; +//! +//! // Open an existing QG index +//! let index = QgIndex::open("target/path/to/qg_index/dir")?; +//! +//! // Perform a vector search (with 1 result) +//! let query = vec![1.1, 2.1, 3.1]; +//! let res = index.search(QgQuery::new(&query).size(1))?; +//! assert_eq!(res[0].id, id1); +//! assert_eq!(index.get_vec(id1)?, vec![1.0, 2.0, 3.0]); +//! +//! # std::fs::remove_dir_all("target/path/to/qg_index/dir").unwrap(); +//! # Ok(()) +//! # } +//! ``` mod index; mod properties; diff --git a/src/qg/properties.rs b/src/qg/properties.rs index 77b74f5..53abda3 100644 --- a/src/qg/properties.rs +++ b/src/qg/properties.rs @@ -7,6 +7,8 @@ use num_enum::TryFromPrimitive; use scopeguard::defer; use crate::error::{make_err, Result}; +use crate::ngt::NgtObjectType; +use crate::{NgtDistance, NgtProperties}; #[derive(Debug, Clone, Copy, PartialEq, Eq, TryFromPrimitive)] #[repr(i32)] @@ -52,6 +54,27 @@ pub enum QgDistance { Cosine = 4, } +impl From for NgtDistance { + fn from(d: QgDistance) -> Self { + match d { + QgDistance::L2 => NgtDistance::L2, + QgDistance::Cosine => NgtDistance::Cosine, + } + } +} + +impl TryFrom for QgDistance { + type Error = crate::Error; + + fn try_from(d: NgtDistance) -> Result { + match d { + NgtDistance::L2 => Ok(QgDistance::L2), + NgtDistance::Cosine => Ok(QgDistance::Cosine), + _ => Err(format!("Invalid distance {d:?} isn't supported for QG").into()), + } + } +} + #[derive(Debug)] pub struct QgProperties { pub(crate) dimension: i32, @@ -296,6 +319,21 @@ impl Drop for QgProperties { } } +impl TryFrom> for NgtProperties +where + T: QgObjectType, + T: NgtObjectType, +{ + type Error = crate::Error; + + fn try_from(prop: QgProperties) -> Result { + NgtProperties::dimension(prop.dimension as usize)? + .creation_edge_size(prop.creation_edge_size as usize)? + .search_edge_size(prop.search_edge_size as usize)? + .distance_type(prop.distance_type.into()) + } +} + #[derive(Debug, Clone, PartialEq)] pub struct QgQuantizationParams { pub dimension_of_subvector: f32,