Implement size_hint for ngrams and optimize unknown lookup.

For BucketIndexer, the number of ngrams (and thus subword indices)
can be precomputed. Allocating a correctly sized buffer for
the subwords yields large bossts in benches.

Author: sebpuetz

src/chunks/vocab/subword.rs

@@ -13,6 +13,7 @@ use crate::subword::{
     BucketIndexer, ExplicitIndexer, FinalfusionHashIndexer, Indexer,
     SubwordIndices as StrSubwordIndices,
 };
+use crate::util::CollectWithCapacity;
 
 /// fastText vocabulary with hashed n-grams.
 pub type FastTextSubwordVocab = SubwordVocab<FastTextIndexer>;
@@ -81,7 +82,7 @@ where
     }
 
     fn bracket(word: impl AsRef<str>) -> String {
-        let mut bracketed = String::new();
+        let mut bracketed = String::with_capacity(word.as_ref().len() + 2);
         bracketed.push(Self::BOW);
         bracketed.push_str(word.as_ref());
         bracketed.push(Self::EOW);
@@ -164,11 +165,16 @@ where
     I: Indexer,
 {
     fn subword_indices(&self, word: &str) -> Option<Vec<usize>> {
-        let indices = Self::bracket(word)
+        let word = Self::bracket(word);
+        let indices = word
             .as_str()
             .subword_indices(self.min_n as usize, self.max_n as usize, &self.indexer)
-            .map(|idx| idx as usize + self.words_len())
-            .collect::<Vec<_>>();
+            .map(|idx| idx as usize + self.words_len());
+        if I::infallible() {
+            let size = indices.size_hint().1.unwrap();
+            return Some(indices.collect_with_capacity(size));
+        }
+        let indices = indices.collect::<Vec<_>>();
         if indices.is_empty() {
             None
         } else {

src/compat/fasttext/indexer.rs

@@ -48,6 +48,10 @@ impl Indexer for FastTextIndexer {
     fn upper_bound(&self) -> u64 {
         u64::from(self.buckets)
     }
+
+    fn infallible() -> bool {
+        true
+    }
 }
 
 /// fastText FNV-1a implementation.

src/subword.rs

@@ -21,6 +21,9 @@ pub trait Indexer {
 
     /// Return the (exclusive) upper bound of this indexer.
     fn upper_bound(&self) -> u64;
+
+    /// Indicates whether this Indexer never fails to produce an index.
+    fn infallible() -> bool;
 }
 
 /// N-Gram indexer with bucketing.
@@ -116,6 +119,10 @@ where
         // max val is <= 64
         2u64.pow(self.buckets_exp as u32)
     }
+
+    fn infallible() -> bool {
+        true
+    }
 }
 
 impl<H> PartialEq for HashIndexer<H> {
@@ -209,6 +216,10 @@ impl Indexer for ExplicitIndexer {
     fn upper_bound(&self) -> u64 {
         self.bound as u64
     }
+
+    fn infallible() -> bool {
+        false
+    }
 }
 
 /// A string reference with its length in characters.
@@ -341,6 +352,12 @@ impl<'a> Iterator for NGrams<'a> {
 
         Some(ngram_with_len)
     }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        let cap_approx = (self.max_n - self.min_n + 1) * self.char_offsets.len();
+        (cap_approx, Some(cap_approx))
+    }
 }
 
 /// Trait returning iterators over subwords and indices.
@@ -435,6 +452,11 @@ where
             .next()
             .map(|ngram| (ngram.inner, self.indexer.index_ngram(&ngram)))
     }
+
+    #[inline]
+    fn size_hint(&self) -> (usize, Option<usize>) {
+        self.ngrams.size_hint()
+    }
 }
 
 #[cfg(test)]