src/in_memory_index.rs

@@ -16,9 +16,9 @@ pub struct InMemoryIndex {
 #[pymethods]
 impl InMemoryIndex {
     #[new]
-    pub fn new(_py: Python, tokens: Vec<u16>, verbose: bool) -> Self {
+    pub fn new(_py: Python, tokens: Vec<u16>) -> Self {
         InMemoryIndex {
-            table: SuffixTable::new(tokens, verbose),
+            table: SuffixTable::new(tokens),
         }
     }
 
@@ -32,7 +32,6 @@ impl InMemoryIndex {
     #[staticmethod]
     pub fn from_token_file(
         path: String,
-        verbose: bool,
         token_limit: Option<usize>,
     ) -> PyResult<Self> {
         let mut buffer = Vec::new();
@@ -47,10 +46,38 @@ impl InMemoryIndex {
         };
 
         Ok(InMemoryIndex {
-            table: SuffixTable::new(transmute_slice(buffer.as_slice()), verbose),
+            table: SuffixTable::new(transmute_slice(buffer.as_slice())),
         })
     }
 
+    #[staticmethod]
+    pub fn from_token_file_sais(
+        path: String,
+        token_limit: Option<usize>,
+    ) -> PyResult<Self> {
+        let mut buffer = Vec::new();
+        let mut file = File::open(&path)?;
+
+        if let Some(max_tokens) = token_limit {
+            // Limit on the number of tokens to consider is provided
+            let max_bytes = max_tokens * std::mem::size_of::<u16>();
+            file.take(max_bytes as u64).read_to_end(&mut buffer)?;
+        } else {
+            file.read_to_end(&mut buffer)?;
+        };
+
+        Ok(InMemoryIndex {
+            table: SuffixTable::new_sais(transmute_slice(buffer.as_slice())),
+        })
+    }
+
+    #[staticmethod]
+    pub fn sais(tokens: Vec<u16>) -> Self {
+        InMemoryIndex {
+            table: SuffixTable::new_sais(tokens),
+        }
+    }
+
     pub fn contains(&self, query: Vec<u16>) -> bool {
         self.table.contains(&query)
     }

src/lib.rs

@@ -8,7 +8,7 @@ use pyo3::prelude::*;
 
 mod in_memory_index;
 mod memmap_index;
-mod par_quicksort;
+mod sais;
 mod table;
 mod util;
 

src/memmap_index.rs

@@ -1,10 +1,10 @@
 use pyo3::exceptions::PyValueError;
 use pyo3::prelude::*;
+use rayon::prelude::*;
 use std::fs::{File, OpenOptions};
 use std::time::Instant;
 
 use crate::mmap_slice::{MmapSlice, MmapSliceMut};
-use crate::par_quicksort::par_sort_unstable_by_key;
 use crate::table::SuffixTable;
 
 /// A memmap index exposes suffix table functionality over text corpora too large to fit in memory.
@@ -29,7 +29,7 @@ impl MemmapIndex {
     }
 
     #[staticmethod]
-    pub fn build(text_path: String, table_path: String, verbose: bool) -> PyResult<Self> {
+    pub fn build(text_path: String, table_path: String) -> PyResult<Self> {
         // Memory map the text as read-only
         let text_mmap = MmapSlice::new(&File::open(&text_path)?)?;
 
@@ -73,10 +73,8 @@ impl MemmapIndex {
                 // The unstable algorithm is critical for avoiding out-of-memory errors, since it does
                 // not allocate any more memory than the input and output slices.
                 println!("Sorting indices...");
-                par_sort_unstable_by_key(
-                    table_mmap.as_slice_mut(),
+                table_mmap.par_sort_unstable_by_key(
                     |&i| &text_mmap[i as usize..],
-                    verbose,
                 );
             });
         println!("Time elapsed: {:?}", start.elapsed());