diff --git a/infos/src/main.rs b/infos/src/main.rs index da15251b0..bb09d7234 100644 --- a/infos/src/main.rs +++ b/infos/src/main.rs @@ -207,6 +207,24 @@ enum Command { word2: String, }, + /// Outputs a CSV with the proximities for the two specified words and + /// the documents ids where these relations appears. + /// + /// `word1`, `prefix` defines the word pair specified *in this specific order*. + /// `proximity` defines the proximity between the two specified words. + /// `documents_ids` defines the documents ids where the relation appears. + WordPrefixPairProximitiesDocids { + /// Display the whole documents ids in details. + #[structopt(long)] + full_display: bool, + + /// First word of the word pair. + word1: String, + + /// Second word of the word pair. + prefix: String, + }, + /// Outputs the words FST to standard output. /// /// One can use the FST binary helper to dissect and analyze it, @@ -282,6 +300,9 @@ fn main() -> anyhow::Result<()> { WordPairProximitiesDocids { full_display, word1, word2 } => { word_pair_proximities_docids(&index, &rtxn, !full_display, word1, word2) } + WordPrefixPairProximitiesDocids { full_display, word1, prefix } => { + word_prefix_pair_proximities_docids(&index, &rtxn, !full_display, word1, prefix) + } ExportWordsFst => export_words_fst(&index, &rtxn), ExportWordsPrefixFst => export_words_prefix_fst(&index, &rtxn), ExportDocuments { internal_documents_ids } => { @@ -1131,3 +1152,46 @@ fn word_pair_proximities_docids( Ok(wtr.flush()?) } + +fn word_prefix_pair_proximities_docids( + index: &Index, + rtxn: &heed::RoTxn, + debug: bool, + word1: String, + word_prefix: String, +) -> anyhow::Result<()> { + use heed::types::ByteSlice; + use milli::RoaringBitmapCodec; + + let stdout = io::stdout(); + let mut wtr = csv::Writer::from_writer(stdout.lock()); + wtr.write_record(&["word1", "word_prefix", "proximity", "documents_ids"])?; + + // Create the prefix key with only the pair of words. + let mut prefix = Vec::with_capacity(word1.len() + word_prefix.len() + 1); + prefix.extend_from_slice(word1.as_bytes()); + prefix.push(0); + prefix.extend_from_slice(word_prefix.as_bytes()); + + let db = index.word_prefix_pair_proximity_docids.as_polymorph(); + let iter = db.prefix_iter::<_, ByteSlice, RoaringBitmapCodec>(rtxn, &prefix)?; + for result in iter { + let (key, docids) = result?; + + // Skip keys that are longer than the requested one, + // a longer key means that the second word is a prefix of the request word. + if key.len() != prefix.len() + 1 { + continue; + } + + let proximity = key.last().unwrap(); + let docids = if debug { + format!("{:?}", docids) + } else { + format!("{:?}", docids.iter().collect::>()) + }; + wtr.write_record(&[&word1, &word_prefix, &proximity.to_string(), &docids])?; + } + + Ok(wtr.flush()?) +} diff --git a/milli/src/search/criteria/exactness.rs b/milli/src/search/criteria/exactness.rs index 22dcb9782..1e4d4e7a2 100644 --- a/milli/src/search/criteria/exactness.rs +++ b/milli/src/search/criteria/exactness.rs @@ -180,10 +180,6 @@ fn resolve_state( if let Some(attribute_allowed_docids) = ctx.field_id_word_count_docids(id, query_len)? { - println!( - "found candidates that have the good count: {:?}", - attribute_allowed_docids - ); let mut attribute_candidates_array = attribute_start_with_docids(ctx, id as u32, query)?; attribute_candidates_array.push(attribute_allowed_docids); diff --git a/milli/src/search/criteria/mod.rs b/milli/src/search/criteria/mod.rs index 61b0fe049..2a883de67 100644 --- a/milli/src/search/criteria/mod.rs +++ b/milli/src/search/criteria/mod.rs @@ -461,13 +461,18 @@ fn query_pair_proximity_docids( let prefix = right.prefix; match (&left.kind, &right.kind) { (QueryKind::Exact { word: left, .. }, QueryKind::Exact { word: right, .. }) => { - if prefix && ctx.in_prefix_cache(&right) { - Ok(ctx - .word_prefix_pair_proximity_docids(left.as_str(), right.as_str(), proximity)? - .unwrap_or_default()) - } else if prefix { - let r_words = word_derivations(&right, true, 0, ctx.words_fst(), wdcache)?; - all_word_pair_proximity_docids(ctx, &[(left, 0)], &r_words, proximity) + if prefix { + match ctx.word_prefix_pair_proximity_docids( + left.as_str(), + right.as_str(), + proximity, + )? { + Some(docids) => Ok(docids), + None => { + let r_words = word_derivations(&right, true, 0, ctx.words_fst(), wdcache)?; + all_word_pair_proximity_docids(ctx, &[(left, 0)], &r_words, proximity) + } + } } else { Ok(ctx .word_pair_proximity_docids(left.as_str(), right.as_str(), proximity)? @@ -477,22 +482,24 @@ fn query_pair_proximity_docids( (QueryKind::Tolerant { typo, word: left }, QueryKind::Exact { word: right, .. }) => { let l_words = word_derivations(&left, false, *typo, ctx.words_fst(), wdcache)?.to_owned(); - if prefix && ctx.in_prefix_cache(&right) { + if prefix { let mut docids = RoaringBitmap::new(); for (left, _) in l_words { - let current_docids = ctx - .word_prefix_pair_proximity_docids( - left.as_ref(), - right.as_ref(), - proximity, - )? - .unwrap_or_default(); + let current_docids = match ctx.word_prefix_pair_proximity_docids( + left.as_str(), + right.as_str(), + proximity, + )? { + Some(docids) => Ok(docids), + None => { + let r_words = + word_derivations(&right, true, 0, ctx.words_fst(), wdcache)?; + all_word_pair_proximity_docids(ctx, &[(left, 0)], &r_words, proximity) + } + }?; docids |= current_docids; } Ok(docids) - } else if prefix { - let r_words = word_derivations(&right, true, 0, ctx.words_fst(), wdcache)?; - all_word_pair_proximity_docids(ctx, &l_words, &r_words, proximity) } else { all_word_pair_proximity_docids(ctx, &l_words, &[(right, 0)], proximity) } diff --git a/milli/src/update/index_documents/mod.rs b/milli/src/update/index_documents/mod.rs index c9f5da0c1..b7fa1492c 100644 --- a/milli/src/update/index_documents/mod.rs +++ b/milli/src/update/index_documents/mod.rs @@ -21,7 +21,7 @@ use typed_chunk::{write_typed_chunk_into_index, TypedChunk}; pub use self::helpers::{ create_sorter, create_writer, merge_cbo_roaring_bitmaps, merge_roaring_bitmaps, - sorter_into_lmdb_database, write_into_lmdb_database, writer_into_reader, + sorter_into_lmdb_database, write_into_lmdb_database, writer_into_reader, MergeFn, }; use self::helpers::{grenad_obkv_into_chunks, GrenadParameters}; pub use self::transform::{Transform, TransformOutput}; @@ -81,7 +81,7 @@ pub struct IndexDocuments<'t, 'u, 'i, 'a> { pub(crate) thread_pool: Option<&'a ThreadPool>, facet_level_group_size: Option, facet_min_level_size: Option, - words_prefix_threshold: Option, + words_prefix_threshold: Option, max_prefix_length: Option, words_positions_level_group_size: Option, words_positions_min_level_size: Option, diff --git a/milli/src/update/word_prefix_pair_proximity_docids.rs b/milli/src/update/word_prefix_pair_proximity_docids.rs index 8f04c23cf..cabe1053b 100644 --- a/milli/src/update/word_prefix_pair_proximity_docids.rs +++ b/milli/src/update/word_prefix_pair_proximity_docids.rs @@ -1,15 +1,13 @@ -use std::str; +use std::collections::HashMap; -use fst::automaton::{Automaton, Str}; -use fst::{IntoStreamer, Streamer}; +use fst::IntoStreamer; use grenad::CompressionType; use heed::types::ByteSlice; -use heed::BytesEncode; use log::debug; +use slice_group_by::GroupBy; -use crate::heed_codec::StrStrU8Codec; use crate::update::index_documents::{ - create_sorter, merge_cbo_roaring_bitmaps, sorter_into_lmdb_database, WriteMethod, + create_sorter, merge_cbo_roaring_bitmaps, sorter_into_lmdb_database, MergeFn, WriteMethod, }; use crate::{Index, Result}; @@ -20,6 +18,7 @@ pub struct WordPrefixPairProximityDocids<'t, 'u, 'i> { pub(crate) chunk_compression_level: Option, pub(crate) max_nb_chunks: Option, pub(crate) max_memory: Option, + threshold: u32, } impl<'t, 'u, 'i> WordPrefixPairProximityDocids<'t, 'u, 'i> { @@ -34,16 +33,26 @@ impl<'t, 'u, 'i> WordPrefixPairProximityDocids<'t, 'u, 'i> { chunk_compression_level: None, max_nb_chunks: None, max_memory: None, + threshold: 100, } } + /// 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 + /// dictionnary, 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 these bound otherwise. + pub fn threshold(&mut self, value: u32) -> &mut Self { + self.threshold = value.max(50); + self + } + pub fn execute(self) -> Result<()> { debug!("Computing and writing the word prefix pair proximity docids into LMDB on disk..."); self.index.word_prefix_pair_proximity_docids.clear(self.wtxn)?; - let prefix_fst = self.index.words_prefixes_fst(self.wtxn)?; - // Here we create a sorter akin to the previous one. let mut word_prefix_pair_proximity_docids_sorter = create_sorter( merge_cbo_roaring_bitmaps, @@ -53,22 +62,59 @@ impl<'t, 'u, 'i> WordPrefixPairProximityDocids<'t, 'u, 'i> { self.max_memory, ); - // We insert all the word pairs corresponding to the word-prefix pairs - // where the prefixes appears in the prefix FST previously constructed. - let db = self.index.word_pair_proximity_docids.remap_data_type::(); - for result in db.iter(self.wtxn)? { - let ((word1, word2, prox), data) = result?; - let automaton = Str::new(word2).starts_with(); - let mut matching_prefixes = prefix_fst.search(automaton).into_stream(); - while let Some(prefix) = matching_prefixes.next() { - let prefix = str::from_utf8(prefix)?; - let pair = (word1, prefix, prox); - let bytes = StrStrU8Codec::bytes_encode(&pair).unwrap(); - word_prefix_pair_proximity_docids_sorter.insert(bytes, data)?; + let prefix_fst = self.index.words_prefixes_fst(self.wtxn)?; + let prefix_fst_keys = prefix_fst.into_stream().into_bytes(); + let prefix_fst_keys: Vec<_> = prefix_fst_keys + .as_slice() + .linear_group_by_key(|x| std::str::from_utf8(&x).unwrap().chars().nth(0).unwrap()) + .collect(); + + let mut db = + self.index.word_pair_proximity_docids.remap_data_type::().iter(self.wtxn)?; + + let mut buffer = Vec::new(); + let mut current_prefixes: Option<&&[Vec]> = None; + let mut prefixes_cache = HashMap::new(); + while let Some(((w1, w2, prox), data)) = db.next().transpose()? { + current_prefixes = match current_prefixes.take() { + Some(prefixes) if w2.as_bytes().starts_with(&prefixes[0]) => Some(prefixes), + _otherwise => { + write_prefixes_in_sorter( + &mut prefixes_cache, + &mut word_prefix_pair_proximity_docids_sorter, + self.threshold, + )?; + prefix_fst_keys.iter().find(|prefixes| w2.as_bytes().starts_with(&prefixes[0])) + } + }; + + if let Some(prefixes) = current_prefixes { + buffer.clear(); + buffer.extend_from_slice(w1.as_bytes()); + buffer.push(0); + for prefix in prefixes.iter().filter(|prefix| w2.as_bytes().starts_with(prefix)) { + buffer.truncate(w1.len() + 1); + buffer.extend_from_slice(prefix); + buffer.push(prox); + + match prefixes_cache.get_mut(&buffer) { + Some(value) => value.push(data), + None => { + prefixes_cache.insert(buffer.clone(), vec![data]); + } + } + } } } + write_prefixes_in_sorter( + &mut prefixes_cache, + &mut word_prefix_pair_proximity_docids_sorter, + self.threshold, + )?; + drop(prefix_fst); + drop(db); // We finally write the word prefix pair proximity docids into the LMDB database. sorter_into_lmdb_database( @@ -82,3 +128,25 @@ impl<'t, 'u, 'i> WordPrefixPairProximityDocids<'t, 'u, 'i> { Ok(()) } } + +fn write_prefixes_in_sorter( + prefixes: &mut HashMap, Vec<&[u8]>>, + sorter: &mut grenad::Sorter, + min_word_per_prefix: u32, +) -> Result<()> { + for (i, (key, data_slices)) in prefixes.drain().enumerate() { + // if the number of words prefixed by the prefix is higher than the threshold, + // we insert it in the sorter. + if data_slices.len() > min_word_per_prefix as usize { + for data in data_slices { + sorter.insert(&key, data)?; + } + // if the first prefix isn't elligible for insertion, + // then the other prefixes can't be elligible. + } else if i == 0 { + break; + } + } + + Ok(()) +} diff --git a/milli/src/update/words_prefixes_fst.rs b/milli/src/update/words_prefixes_fst.rs index f35dea10d..be33c156b 100644 --- a/milli/src/update/words_prefixes_fst.rs +++ b/milli/src/update/words_prefixes_fst.rs @@ -8,7 +8,7 @@ use crate::{Index, Result, SmallString32}; pub struct WordsPrefixesFst<'t, 'u, 'i> { wtxn: &'t mut heed::RwTxn<'i, 'u>, index: &'i Index, - threshold: f64, + threshold: u32, max_prefix_length: usize, _update_id: u64, } @@ -22,20 +22,20 @@ impl<'t, 'u, 'i> WordsPrefixesFst<'t, 'u, 'i> { WordsPrefixesFst { wtxn, index, - threshold: 0.1 / 100.0, // .01% + threshold: 100, max_prefix_length: 4, _update_id: update_id, } } - /// Set the ratio of concerned words required to make a prefix be part of the words prefixes + /// 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 /// dictionnary, therefore this prefix is added to the words prefixes datastructures. /// - /// Default value is `0.01` or `1%`. This value must be between 0 and 1 and will be clamped - /// to these bounds otherwise. - pub fn threshold(&mut self, value: f64) -> &mut Self { - self.threshold = value.min(1.0).max(0.0); // clamp [0, 1] + /// 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: u32) -> &mut Self { + self.threshold = value.max(50); self } @@ -50,8 +50,6 @@ impl<'t, 'u, 'i> WordsPrefixesFst<'t, 'u, 'i> { pub fn execute(self) -> Result<()> { let words_fst = self.index.words_fst(&self.wtxn)?; - let number_of_words = words_fst.len(); - let min_number_of_words = (number_of_words as f64 * self.threshold) as usize; let mut prefix_fsts = Vec::with_capacity(self.max_prefix_length); for n in 1..=self.max_prefix_length { @@ -80,7 +78,7 @@ impl<'t, 'u, 'i> WordsPrefixesFst<'t, 'u, 'i> { current_prefix_count += 1; // There is enough words corresponding to this prefix to add it to the cache. - if current_prefix_count == min_number_of_words { + if current_prefix_count >= self.threshold { builder.insert(prefix)?; } }