Add a separate option to run the API server using mmap

api/src/lib.rs

@@ -21,7 +21,7 @@ pub struct AppState {
     pub database: Arc<Database>,
     pub index: Arc<Index>
 }
-pub async fn start(index_location: &str, database_address: &str, port: u32) -> Result<(), errors::AppError> {
+pub async fn start(index_location: &str, database_address: &str, port: u32, use_mmap: bool) -> Result<(), errors::AppError> {
     let version = format!("{}/.version", index_location);
 
     let sampledata = format!("{}/datastore/sampledata.json", index_location);
@@ -33,7 +33,8 @@ pub async fn start(index_location: &str, database_address: &str, port: u32) -> R
     let taxons = format!("{}/datastore/taxons.tsv", index_location);
 
     let sa = format!("{}/sa.bin", index_location);
-    let proteins = format!("{}/proteins.tsv", index_location);
+    let proteins = format!("{}/proteins.bin", index_location);
+    let mappings = format!("{}/mapping.bin", index_location);
 
     let database = Database::try_from_url(database_address)?;
 
@@ -48,7 +49,7 @@ pub async fn start(index_location: &str, database_address: &str, port: u32) -> R
         &taxons
     )?;
 
-    let index = Index::try_from_files(&sa, &proteins)?;
+    let index = Index::try_from_files(&sa, &proteins, &mappings, use_mmap)?;
 
     let app_state = AppState {
         datastore: Arc::new(datastore),

api/src/main.rs

@@ -8,14 +8,16 @@ pub struct Arguments {
     #[arg(short, long)]
     database_address: String,
     #[arg(short, long)]
-    port: u32
+    port: u32,
+    #[arg(short, long, default_value_t = false)]
+    mmap: bool
 }
 
 #[tokio::main]
 async fn main() {
     let args = Arguments::parse();
 
-    if let Err(e) = start(&args.index_location, &args.database_address, args.port).await {
+    if let Err(e) = start(&args.index_location, &args.database_address, args.port, args.mmap).await {
         eprintln!("{}", e);
         std::process::exit(1);
     }

index/Cargo.toml

@@ -4,7 +4,7 @@ version = "2.3.2"
 edition = "2021"
 
 [dependencies]
-sa-compression = { git = "https://github.com/unipept/unipept-index.git" }
-sa-index = { git = "https://github.com/unipept/unipept-index.git" }
-sa-mappings = { git = "https://github.com/unipept/unipept-index.git" }
+sa-server = { git = "https://github.com/unipept/unipept-index.git", branch = "feature/speedup-loading-index" }
+sa-index = { git = "https://github.com/unipept/unipept-index.git", branch = "feature/speedup-loading-index" }
+sa-mappings = { git = "https://github.com/unipept/unipept-index.git", branch = "feature/speedup-loading-index" }
 thiserror = "1.0"

index/src/lib.rs

@@ -1,39 +1,35 @@
-use std::{
-    fs::File,
-    io::{BufReader, Read}
-};
-
 pub use errors::IndexError;
 use errors::LoadIndexError;
-use sa_compression::load_compressed_suffix_array;
+use sa_server::{load_mapping_file, load_proteins_file, load_suffix_array_file};
 pub use sa_index::peptide_search::ProteinInfo;
 pub use sa_index::peptide_search::SearchResult;
 use sa_index::{
-    binary::load_suffix_array,
     peptide_search::{search_all_peptides},
-    sa_searcher::BitVecSearcher,
-    SuffixArray
 };
-use sa_mappings::proteins::Proteins;
+use sa_index::sa_searcher::Searcher;
+use sa_index::suffix_to_protein_index::SuffixToProteinIndex;
 
 mod errors;
 
 pub struct Index {
-    searcher: BitVecSearcher
+    searcher: Searcher
 }
 
 impl Index {
-    pub fn try_from_files(index_file: &str, proteins_file: &str) -> Result<Self, IndexError> {
-
-        let proteins = Proteins::try_from_database_file(proteins_file)
-            .map_err(|_| LoadIndexError::LoadProteinsErrors(
-                LoadIndexError::FileNotFound(proteins_file.to_string()).to_string(),
-            ))?;
+    pub fn try_from_files(index_file: &str, proteins_file: &str, mapping_file: &str, use_mmap: bool) -> Result<Self, IndexError> {
+        eprintln!("Loading proteins from file: {}", proteins_file);
+        let proteins =
+            load_proteins_file(proteins_file, use_mmap).map_err(|err| LoadIndexError::LoadProteinsErrors(err.to_string()))?;
 
+        eprintln!("Loading suffix array from file: {}", index_file);
         let suffix_array =
-            load_index_file(index_file).map_err(|err| LoadIndexError::LoadSuffixArrayError(err.to_string()))?;
+            load_suffix_array_file(index_file, use_mmap).map_err(|err| LoadIndexError::LoadSuffixArrayError(err.to_string()))?;
 
-        let searcher = BitVecSearcher::new(suffix_array, proteins);
+        eprintln!("Creating searcher");
+        let suffix_to_protein_mapping =
+            load_mapping_file(mapping_file, use_mmap).map_err(|err| LoadIndexError::LoadProteinsErrors(err.to_string()))?;
+
+        let searcher = Searcher::new(suffix_array, proteins, suffix_to_protein_mapping.0);
 
         Ok(Self { searcher })
     }
@@ -42,27 +38,3 @@ impl Index {
         search_all_peptides(&self.searcher, peptides, cutoff.unwrap_or(10_000), equate_il, tryptic)
     }
 }
-
-fn load_index_file(index_file: &str) -> Result<SuffixArray, LoadIndexError> {
-    // Open the suffix array file
-    let mut sa_file = File::open(index_file).map_err(
-        |_| LoadIndexError::FileNotFound(index_file.to_string())
-    )?;
-
-    // Create a buffer reader for the file
-    let mut reader = BufReader::new(&mut sa_file);
-
-    // Read the bits per value from the binary file (1 byte)
-    let mut bits_per_value_buffer = [0_u8; 1];
-    reader.read_exact(&mut bits_per_value_buffer).map_err(|_| {
-        LoadIndexError::LoadSuffixArrayError("Could not read the flags from the binary file".to_string())
-    })?;
-    let bits_per_value = bits_per_value_buffer[0];
-
-    if bits_per_value == 64 {
-        Ok(load_suffix_array(&mut reader).map_err(|err| LoadIndexError::LoadSuffixArrayError(err.to_string()))?)
-    } else {
-        Ok(load_compressed_suffix_array(&mut reader, bits_per_value as usize)
-            .map_err(|err| LoadIndexError::LoadSuffixArrayError(err.to_string()))?)
-    }
-}