diff --git a/src/index/io.rs b/src/index/io.rs index d8e55a2..1561652 100644 --- a/src/index/io.rs +++ b/src/index/io.rs @@ -7,6 +7,8 @@ use faiss_sys::*; use std::ffi::CString; use std::ptr; +use super::io_flags::IoFlags; + /// Write an index to a file. /// /// # Error @@ -42,7 +44,34 @@ where let f = file_name.as_ref(); let f = CString::new(f).map_err(|_| Error::BadFilePath)?; let mut inner = ptr::null_mut(); - faiss_try(faiss_read_index_fname(f.as_ptr(), 0, &mut inner))?; + faiss_try(faiss_read_index_fname( + f.as_ptr(), + IoFlags::MEM_RESIDENT.into(), + &mut inner, + ))?; + Ok(IndexImpl::from_inner_ptr(inner)) + } +} + +/// Read an index from a file with io flags. You can memory map some index types with this. +/// +/// # Error +/// +/// This function returns an error if the description contains any byte with the value `\0` (since +/// it cannot be converted to a C string), or if the internal index reading operation fails. +pub fn read_index_with_flags

(file_name: P, io_flags: IoFlags) -> Result +where + P: AsRef, +{ + unsafe { + let f = file_name.as_ref(); + let f = CString::new(f).map_err(|_| Error::BadFilePath)?; + let mut inner = ptr::null_mut(); + faiss_try(faiss_read_index_fname( + f.as_ptr(), + io_flags.into(), + &mut inner, + ))?; Ok(IndexImpl::from_inner_ptr(inner)) } } @@ -74,4 +103,11 @@ mod tests { assert_eq!(index.ntotal(), 5); ::std::fs::remove_file(&filepath).unwrap(); } + + #[test] + fn test_read_with_flags() { + let index = read_index_with_flags("file_name", IoFlags::MEM_MAP | IoFlags::READ_ONLY); + // we just want to ensure the method signature is right here + assert!(index.is_err()); + } } diff --git a/src/index/io_flags.rs b/src/index/io_flags.rs new file mode 100644 index 0000000..d65a2b9 --- /dev/null +++ b/src/index/io_flags.rs @@ -0,0 +1,53 @@ +//! Module containing the io flags. + +/// Io Flags used during index reading, not all flags applicable to all indices +// This is a set of constants rather than enum so that bitwise operations exist +#[derive(Debug, Copy, Clone, Eq, Hash, PartialEq)] +pub struct IoFlags(u16); + +impl IoFlags { + /// Load entire index into memory (default behavior) + pub const MEM_RESIDENT: Self = IoFlags(0x00); + /// Memory-map index + pub const MEM_MAP: Self = IoFlags(0x01); + /// Index is read-only + pub const READ_ONLY: Self = IoFlags(0x02); +} + +impl std::ops::BitOr for IoFlags { + type Output = Self; + + fn bitor(self, rhs: Self) -> Self::Output { + Self(self.0 | rhs.0) + } +} + +impl From for IoFlags { + fn from(n: i32) -> IoFlags { + IoFlags(n as u16) + } +} + +impl From for i32 { + fn from(io_flag: IoFlags) -> i32 { + io_flag.0 as i32 + } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn can_do_bitor() { + let mmap = IoFlags::MEM_MAP; + let ro = IoFlags::READ_ONLY; + assert_eq!(IoFlags(0x03), mmap | ro); + } + + #[test] + fn can_coerce_to_i32() { + let mmap = IoFlags::MEM_MAP; + assert_eq!(1, mmap.into()); + } +} diff --git a/src/index/mod.rs b/src/index/mod.rs index 39a97a5..a1ea8ce 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -25,6 +25,7 @@ pub mod autotune; pub mod flat; pub mod id_map; pub mod io; +pub mod io_flags; pub mod ivf_flat; pub mod lsh; pub mod pretransform;