feat: rework how we train ngram indices for better performance (#3518)

The biggest impactful change is that we can no longer apply ngram
indices when the string contains 0, 1, or 2 characters (these strings
are likely to match too many rows to make scalar indices useful
anyways).

Author: westonpace

Cargo.lock

@@ -151,6 +151,24 @@ impl Error {
     }
 }
 
+pub trait LanceOptionExt<T> {
+    /// Unwraps an option, returning an internal error if the option is None.
+    ///
+    /// Can be used when an option is expected to have a value.
+    fn expect_ok(self) -> Result<T>;
+}
+
+impl<T> LanceOptionExt<T> for Option<T> {
+    #[track_caller]
+    fn expect_ok(self) -> Result<T> {
+        let location = std::panic::Location::caller().to_snafu_location();
+        self.ok_or_else(|| Error::Internal {
+            message: "Expected option to have value".to_string(),
+            location,
+        })
+    }
+}
+
 trait ToSnafuLocation {
     fn to_snafu_location(&'static self) -> snafu::Location;
 }

rust/lance-core/src/error.rs

@@ -65,6 +65,7 @@ uuid.workspace = true
 approx.workspace = true
 clap = { workspace = true, features = ["derive"] }
 criterion.workspace = true
+env_logger = "0.11.6"
 lance-datagen.workspace = true
 lance-testing.workspace = true
 tempfile.workspace = true

rust/lance-index/Cargo.toml

@@ -12,7 +12,7 @@ use itertools::Itertools;
 use lance_core::cache::FileMetadataCache;
 use lance_core::ROW_ID;
 use lance_index::scalar::lance_format::LanceIndexStore;
-use lance_index::scalar::ngram::{NGramIndex, NGramIndexBuilder};
+use lance_index::scalar::ngram::{NGramIndex, NGramIndexBuilder, NGramIndexBuilderOptions};
 use lance_index::scalar::{ScalarIndex, TextQuery};
 use lance_io::object_store::ObjectStore;
 use object_store::path::Path;
@@ -22,6 +22,8 @@ use pprof::criterion::{Output, PProfProfiler};
 fn bench_ngram(c: &mut Criterion) {
     const TOTAL: usize = 1_000_000;
 
+    env_logger::init();
+
     let rt = tokio::runtime::Builder::new_multi_thread().build().unwrap();
 
     let tempdir = tempfile::tempdir().unwrap();
@@ -61,21 +63,35 @@ fn bench_ngram(c: &mut Criterion) {
 
     let batches = (0..1000).map(|i| batch.slice(i * 1000, 1000)).collect_vec();
 
-    c.bench_function(format!("ngram_index({TOTAL})").as_str(), |b| {
+    let mut group = c.benchmark_group("train");
+
+    group.sample_size(10);
+    group.bench_function(format!("ngram_train({TOTAL})").as_str(), |b| {
         b.to_async(&rt).iter(|| async {
             let stream = RecordBatchStreamAdapter::new(
                 batch.schema(),
                 stream::iter(batches.clone().into_iter().map(Ok)),
             );
             let stream = Box::pin(stream);
-            let mut builder = NGramIndexBuilder::default();
-            builder.train(stream).await.unwrap();
-            builder.write(store.as_ref()).await.unwrap();
+            let mut builder =
+                NGramIndexBuilder::try_new(NGramIndexBuilderOptions::default()).unwrap();
+            let num_spill_files = builder.train(stream).await.unwrap();
+            builder
+                .write_index(store.as_ref(), num_spill_files, None)
+                .await
+                .unwrap();
         })
     });
 
+    drop(group);
+
+    let mut group = c.benchmark_group("search");
+
+    group
+        .sample_size(10)
+        .measurement_time(Duration::from_secs(10));
     let index = rt.block_on(NGramIndex::load(store)).unwrap();
-    c.bench_function(format!("ngram_search({TOTAL})").as_str(), |b| {
+    group.bench_function(format!("ngram_search({TOTAL})").as_str(), |b| {
         b.to_async(&rt).iter(|| async {
             let sample_idx = rand::random::<usize>() % batch.num_rows();
             let sample = batch

rust/lance-index/benches/ngram.rs

@@ -211,6 +211,12 @@ pub trait IndexStore: std::fmt::Debug + Send + Sync + DeepSizeOf {
     ///
     /// This is often useful when remapping or updating
     async fn copy_index_file(&self, name: &str, dest_store: &dyn IndexStore) -> Result<()>;
+
+    /// Rename an index file
+    async fn rename_index_file(&self, name: &str, new_name: &str) -> Result<()>;
+
+    /// Delete an index file (used in the tmp spill store to keep tmp size down)
+    async fn delete_index_file(&self, name: &str) -> Result<()>;
 }
 
 /// Different scalar indices may support different kinds of queries
@@ -530,7 +536,7 @@ impl AnyQuery for TextQuery {
 }
 
 /// The result of a search operation against a scalar index
-#[derive(Debug)]
+#[derive(Debug, PartialEq)]
 pub enum SearchResult {
     /// The exact row ids that satisfy the query
     Exact(RowIdTreeMap),

rust/lance-index/src/scalar.rs

@@ -329,16 +329,6 @@ pub struct TokenSet {
 }
 
 impl TokenSet {
-    pub(crate) fn new(tokens: HashMap<String, u32>) -> Self {
-        let next_id = tokens.values().max().copied().unwrap_or(0) + 1;
-        let total_length = tokens.keys().map(|s| s.len()).sum();
-        Self {
-            tokens,
-            next_id,
-            total_length,
-        }
-    }
-
     pub fn num_tokens(&self) -> usize {
         self.tokens.len()
     }

rust/lance-index/src/scalar/inverted/index.rs

@@ -279,6 +279,18 @@ impl IndexStore for LanceIndexStore {
             Ok(())
         }
     }
+
+    async fn rename_index_file(&self, name: &str, new_name: &str) -> Result<()> {
+        let path = self.index_dir.child(name);
+        let new_path = self.index_dir.child(new_name);
+        self.object_store.copy(&path, &new_path).await?;
+        self.object_store.delete(&path).await
+    }
+
+    async fn delete_index_file(&self, name: &str) -> Result<()> {
+        let path = self.index_dir.child(name);
+        self.object_store.delete(&path).await
+    }
 }
 
 #[cfg(test)]