49 - allow read_index_with_flags for memory mapping some index types

Cargo.toml

@@ -1,7 +1,7 @@
 [package]
 name = "faiss"
 description = "High-level bindings for Faiss, the vector similarity search engine"
-version = "0.10.1-alpha.0"
+version = "0.11.1-alpha.0"
 authors = ["Eduardo Pinho <[email protected]>"]
 license = "MIT/Apache-2.0"
 repository = "https://github.com/Enet4/faiss-rs"

src/index/io.rs

@@ -7,6 +7,8 @@ use faiss_sys::*;
 use std::ffi::CString;
 use std::ptr;
 
+use super::io_flags::io_flag;
+
 /// Write an index to a file.
 ///
 /// # Error
@@ -42,7 +44,57 @@ 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(),
+            io_flag::MEM_RESIDENT as i32,
+            &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<P>(file_name: P, io_flags: u8) -> Result<IndexImpl>
+where
+    P: AsRef<str>,
+{
+    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 as i32,
+            &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<P>(file_name: P, io_flags: u8) -> Result<IndexImpl>
+where
+    P: AsRef<str>,
+{
+    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 as i32,
+            &mut inner,
+        ))?;
         Ok(IndexImpl::from_inner_ptr(inner))
     }
 }
@@ -51,6 +103,7 @@ where
 mod tests {
     use super::*;
     use crate::index::flat::FlatIndex;
+    use crate::index::io_flags::io_flag;
     use crate::index::Index;
     const D: u32 = 8;
 
@@ -74,4 +127,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", io_flag::MEM_MAP | io_flag::READ_ONLY);
+        // we just want to ensure the method signature is right here
+        assert!(index.is_err());
+    }
 }

src/index/io_flags.rs

@@ -0,0 +1,24 @@
+//! 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
+pub mod io_flag {
+    /// Load entire index into memory (default behavior)
+    pub const MEM_RESIDENT: u8 = 0x00;
+    /// Memory-map index
+    pub const MEM_MAP: u8 = 0x01;
+    /// Index is read-only
+    pub const READ_ONLY: u8 = 0x02;
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    #[test]
+    fn can_do_bitor() {
+        let mmap = io_flag::MEM_MAP;
+        let ro = io_flag::READ_ONLY;
+        assert_eq!(3, mmap | ro);
+    }
+}

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;