1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
use std::iter::{repeat_with, FromIterator};
use std::str;
use fst::{SetBuilder, Streamer};
use heed::RwTxn;
use crate::{Index, Result, SmallString32};
pub struct WordsPrefixesFst<'t, 'i> {
wtxn: &'t mut RwTxn<'i>,
index: &'i Index,
threshold: usize,
max_prefix_length: usize,
}
impl<'t, 'i> WordsPrefixesFst<'t, 'i> {
pub fn new(wtxn: &'t mut RwTxn<'i>, index: &'i Index) -> WordsPrefixesFst<'t, 'i> {
WordsPrefixesFst { wtxn, index, threshold: 100, max_prefix_length: 4 }
}
/// Set the number of words required to make a prefix be part of the words prefixes
/// database. If a word prefix is supposed to match more than this number of words in the
/// dictionary, therefore this prefix is added to the words prefixes datastructures.
///
/// Default value is 100. This value must be higher than 50 and will be clamped
/// to this bound otherwise.
pub fn threshold(&mut self, value: usize) -> &mut Self {
self.threshold = value;
self
}
/// Set the maximum length of prefixes in bytes.
///
/// Default value is `4` bytes. This value must be between 1 and 25 will be clamped
/// to these bounds, otherwise.
pub fn max_prefix_length(&mut self, value: usize) -> &mut Self {
self.max_prefix_length = value;
self
}
#[tracing::instrument(
level = "trace",
skip_all,
target = "indexing::prefix",
name = "words_prefix_fst"
)]
pub fn execute(self) -> Result<()> {
let words_fst = self.index.words_fst(self.wtxn)?;
let mut current_prefix = vec![SmallString32::new(); self.max_prefix_length];
let mut current_prefix_count = vec![0; self.max_prefix_length];
let mut builders =
repeat_with(SetBuilder::memory).take(self.max_prefix_length).collect::<Vec<_>>();
let mut stream = words_fst.stream();
while let Some(bytes) = stream.next() {
for n in 0..self.max_prefix_length {
let current_prefix = &mut current_prefix[n];
let current_prefix_count = &mut current_prefix_count[n];
let builder = &mut builders[n];
// We try to get the first n bytes out of this string but we only want
// to split at valid characters bounds. If we try to split in the middle of
// a character we ignore this word and go to the next one.
let word = str::from_utf8(bytes)?;
let prefix = match word.get(..=n) {
Some(prefix) => prefix,
None => continue,
};
// This is the first iteration of the loop,
// or the current word doesn't starts with the current prefix.
if *current_prefix_count == 0 || prefix != current_prefix.as_str() {
*current_prefix = SmallString32::from(prefix);
*current_prefix_count = 0;
}
*current_prefix_count += 1;
// There is enough words corresponding to this prefix to add it to the cache.
if *current_prefix_count >= self.threshold {
builder.insert(prefix)?;
}
}
}
// We merge all of the previously computed prefixes into on final set.
let prefix_fsts: Vec<_> = builders.into_iter().map(|sb| sb.into_set()).collect();
let op = fst::set::OpBuilder::from_iter(prefix_fsts.iter());
let mut builder = fst::SetBuilder::memory();
builder.extend_stream(op.r#union())?;
let prefix_fst = builder.into_set();
// Set the words prefixes FST in the dtabase.
self.index.put_words_prefixes_fst(self.wtxn, &prefix_fst)?;
Ok(())
}
}