perf: parallelize ngram indexing (#3501)

total indexing time reduced from 23s to 5s
```
ngram_index(1000000)    time:   [5.1192 s 5.1756 s 5.2319 s]
                        change: [-78.163% -77.791% -77.410%] (p = 0.00 < 0.05)
                        Performance has improved.
```

---------

Signed-off-by: BubbleCal <bubble-cal@outlook.com>

Author: BubbleCal

rust/lance-index/benches/ngram.rs

@@ -58,18 +58,24 @@ fn bench_ngram(c: &mut Criterion) {
         vec![doc_col, row_id_col],
     )
     .unwrap();
-    let stream =
-        RecordBatchStreamAdapter::new(batch.schema(), stream::iter(vec![Ok(batch.clone())]));
-    let stream = Box::pin(stream);
 
-    rt.block_on(async {
-        let mut builder = NGramIndexBuilder::default();
-        builder.train(stream).await.unwrap();
-        builder.write(store.as_ref()).await.unwrap();
+    let batches = (0..1000).map(|i| batch.slice(i * 1000, 1000)).collect_vec();
+
+    c.bench_function(format!("ngram_index({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 index = rt.block_on(NGramIndex::load(store)).unwrap();
 
-    c.bench_function(format!("invert({TOTAL})").as_str(), |b| {
+    let index = rt.block_on(NGramIndex::load(store)).unwrap();
+    c.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/src/scalar/inverted.rs

@@ -1,7 +1,7 @@
 // SPDX-License-Identifier: Apache-2.0
 // SPDX-FileCopyrightText: Copyright The Lance Authors
 
-mod builder;
+pub mod builder;
 mod index;
 mod tokenizer;
 mod wand;

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

@@ -53,7 +53,7 @@ lazy_static! {
     // it doesn't mean higher value will result in better performance,
     // because the bottleneck can be the IO once the number of shards is large enough,
     // it's 8 by default
-    static ref LANCE_FTS_NUM_SHARDS: usize = std::env::var("LANCE_FTS_NUM_SHARDS")
+    pub static ref LANCE_FTS_NUM_SHARDS: usize = std::env::var("LANCE_FTS_NUM_SHARDS")
         .unwrap_or_else(|_| "8".to_string())
         .parse()
         .expect("failed to parse LANCE_FTS_NUM_SHARDS");

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

@@ -27,6 +27,7 @@ use crate::vector::VectorIndex;
 use crate::{Index, IndexType};
 
 use super::btree::TrainingSource;
+use super::inverted::builder::LANCE_FTS_NUM_SHARDS;
 use super::inverted::TokenSet;
 use super::{AnyQuery, IndexReader, IndexStore, ScalarIndex, SearchResult, TextQuery};
 
@@ -465,12 +466,52 @@ impl NGramIndexBuilder {
         let schema = data.schema();
         Self::validate_schema(schema.as_ref())?;
 
+        let num_shards = *LANCE_FTS_NUM_SHARDS;
+        let mut senders = Vec::with_capacity(num_shards);
+        let mut builders = Vec::with_capacity(num_shards);
+        for _ in 0..*LANCE_FTS_NUM_SHARDS {
+            let (send, mut recv) = tokio::sync::mpsc::channel(2);
+            senders.push(send);
+
+            let mut builder = Self::new();
+            let future = tokio::spawn(async move {
+                while let Some(batch) = recv.recv().await {
+                    builder.process_batch(&batch);
+                }
+                builder
+            });
+            builders.push(future);
+        }
+
+        let mut idx = 0;
         while let Some(batch) = data.try_next().await? {
-            self.process_batch(&batch);
+            senders[idx % num_shards].send(batch).await.unwrap();
+            idx += 1;
+        }
+
+        std::mem::drop(senders);
+        let builders = futures::future::try_join_all(builders).await?;
+        for builder in builders {
+            self.merge(builder);
         }
+
         Ok(())
     }
 
+    fn merge(&mut self, mut other: Self) {
+        for (token, new_token_id) in other.tokens_map {
+            if let Some(token_id) = self.tokens_map.get(&token) {
+                self.bitmaps[*token_id as usize] |=
+                    std::mem::take(&mut other.bitmaps[new_token_id as usize]);
+            } else {
+                // This is a new token
+                self.tokens_map.insert(token, self.bitmaps.len() as u32);
+                self.bitmaps
+                    .push(std::mem::take(&mut other.bitmaps[new_token_id as usize]));
+            }
+        }
+    }
+
     pub async fn write(self, store: &dyn IndexStore) -> Result<()> {
         let mut ordered_tokens = self.tokens_map.into_iter().collect::<Vec<_>>();
         ordered_tokens.sort_by_key(|(_, id)| *id);