mirror of
https://github.com/meilisearch/MeiliSearch
synced 2024-09-22 16:41:45 +02:00
245 lines
9.5 KiB
Rust
245 lines
9.5 KiB
Rust
use std::collections::{BTreeMap, VecDeque};
|
|
use std::fs::File;
|
|
use std::io::BufReader;
|
|
use std::{cmp, io};
|
|
|
|
use obkv::KvReaderU16;
|
|
|
|
use super::helpers::{
|
|
create_sorter, create_writer, merge_deladd_cbo_roaring_bitmaps, try_split_array_at,
|
|
writer_into_reader, GrenadParameters, MergeFn,
|
|
};
|
|
use crate::error::SerializationError;
|
|
use crate::index::db_name::DOCID_WORD_POSITIONS;
|
|
use crate::proximity::{index_proximity, MAX_DISTANCE};
|
|
use crate::update::del_add::{DelAdd, KvReaderDelAdd, KvWriterDelAdd};
|
|
use crate::{DocumentId, Result};
|
|
|
|
/// Extracts the best proximity between pairs of words and the documents ids where this pair appear.
|
|
///
|
|
/// Returns a grenad reader with the list of extracted word pairs proximities and
|
|
/// documents ids from the given chunk of docid word positions.
|
|
#[tracing::instrument(level = "trace", skip_all, target = "indexing::extract")]
|
|
pub fn extract_word_pair_proximity_docids<R: io::Read + io::Seek>(
|
|
docid_word_positions: grenad::Reader<R>,
|
|
indexer: GrenadParameters,
|
|
) -> Result<grenad::Reader<BufReader<File>>> {
|
|
puffin::profile_function!();
|
|
|
|
let max_memory = indexer.max_memory_by_thread();
|
|
|
|
let mut word_pair_proximity_docids_sorters: Vec<_> = (1..MAX_DISTANCE)
|
|
.map(|_| {
|
|
create_sorter(
|
|
grenad::SortAlgorithm::Unstable,
|
|
merge_deladd_cbo_roaring_bitmaps,
|
|
indexer.chunk_compression_type,
|
|
indexer.chunk_compression_level,
|
|
indexer.max_nb_chunks,
|
|
max_memory.map(|m| m / MAX_DISTANCE as usize),
|
|
)
|
|
})
|
|
.collect();
|
|
|
|
let mut del_word_positions: VecDeque<(String, u16)> =
|
|
VecDeque::with_capacity(MAX_DISTANCE as usize);
|
|
let mut add_word_positions: VecDeque<(String, u16)> =
|
|
VecDeque::with_capacity(MAX_DISTANCE as usize);
|
|
let mut del_word_pair_proximity = BTreeMap::new();
|
|
let mut add_word_pair_proximity = BTreeMap::new();
|
|
let mut current_document_id = None;
|
|
|
|
let mut cursor = docid_word_positions.into_cursor()?;
|
|
while let Some((key, value)) = cursor.move_on_next()? {
|
|
let (document_id_bytes, _fid_bytes) = try_split_array_at(key)
|
|
.ok_or(SerializationError::Decoding { db_name: Some(DOCID_WORD_POSITIONS) })?;
|
|
let document_id = u32::from_be_bytes(document_id_bytes);
|
|
|
|
// if we change document, we fill the sorter
|
|
if current_document_id.map_or(false, |id| id != document_id) {
|
|
puffin::profile_scope!("Document into sorter");
|
|
|
|
// FIXME: span inside of a hot loop might degrade performance and create big reports
|
|
let span = tracing::trace_span!(target: "indexing::details", "document_into_sorter");
|
|
let _entered = span.enter();
|
|
|
|
document_word_positions_into_sorter(
|
|
current_document_id.unwrap(),
|
|
&del_word_pair_proximity,
|
|
&add_word_pair_proximity,
|
|
&mut word_pair_proximity_docids_sorters,
|
|
)?;
|
|
del_word_pair_proximity.clear();
|
|
add_word_pair_proximity.clear();
|
|
}
|
|
|
|
current_document_id = Some(document_id);
|
|
|
|
let (del, add): (Result<_>, Result<_>) = rayon::join(
|
|
|| {
|
|
// deletions
|
|
if let Some(deletion) = KvReaderDelAdd::new(value).get(DelAdd::Deletion) {
|
|
for (position, word) in KvReaderU16::new(deletion).iter() {
|
|
// drain the proximity window until the head word is considered close to the word we are inserting.
|
|
while del_word_positions.front().map_or(false, |(_w, p)| {
|
|
index_proximity(*p as u32, position as u32) >= MAX_DISTANCE
|
|
}) {
|
|
word_positions_into_word_pair_proximity(
|
|
&mut del_word_positions,
|
|
&mut del_word_pair_proximity,
|
|
)?;
|
|
}
|
|
|
|
// insert the new word.
|
|
let word = std::str::from_utf8(word)?;
|
|
del_word_positions.push_back((word.to_string(), position));
|
|
}
|
|
|
|
while !del_word_positions.is_empty() {
|
|
word_positions_into_word_pair_proximity(
|
|
&mut del_word_positions,
|
|
&mut del_word_pair_proximity,
|
|
)?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
},
|
|
|| {
|
|
// additions
|
|
if let Some(addition) = KvReaderDelAdd::new(value).get(DelAdd::Addition) {
|
|
for (position, word) in KvReaderU16::new(addition).iter() {
|
|
// drain the proximity window until the head word is considered close to the word we are inserting.
|
|
while add_word_positions.front().map_or(false, |(_w, p)| {
|
|
index_proximity(*p as u32, position as u32) >= MAX_DISTANCE
|
|
}) {
|
|
word_positions_into_word_pair_proximity(
|
|
&mut add_word_positions,
|
|
&mut add_word_pair_proximity,
|
|
)?;
|
|
}
|
|
|
|
// insert the new word.
|
|
let word = std::str::from_utf8(word)?;
|
|
add_word_positions.push_back((word.to_string(), position));
|
|
}
|
|
|
|
while !add_word_positions.is_empty() {
|
|
word_positions_into_word_pair_proximity(
|
|
&mut add_word_positions,
|
|
&mut add_word_pair_proximity,
|
|
)?;
|
|
}
|
|
}
|
|
|
|
Ok(())
|
|
},
|
|
);
|
|
|
|
del?;
|
|
add?;
|
|
}
|
|
|
|
if let Some(document_id) = current_document_id {
|
|
puffin::profile_scope!("Final document into sorter");
|
|
// FIXME: span inside of a hot loop might degrade performance and create big reports
|
|
let span = tracing::trace_span!(target: "indexing::details", "final_document_into_sorter");
|
|
let _entered = span.enter();
|
|
|
|
document_word_positions_into_sorter(
|
|
document_id,
|
|
&del_word_pair_proximity,
|
|
&add_word_pair_proximity,
|
|
&mut word_pair_proximity_docids_sorters,
|
|
)?;
|
|
}
|
|
{
|
|
puffin::profile_scope!("sorter_into_reader");
|
|
// FIXME: span inside of a hot loop might degrade performance and create big reports
|
|
let span = tracing::trace_span!(target: "indexing::details", "sorter_into_reader");
|
|
let _entered = span.enter();
|
|
|
|
let mut writer = create_writer(
|
|
indexer.chunk_compression_type,
|
|
indexer.chunk_compression_level,
|
|
tempfile::tempfile()?,
|
|
);
|
|
|
|
for sorter in word_pair_proximity_docids_sorters {
|
|
sorter.write_into_stream_writer(&mut writer)?;
|
|
}
|
|
|
|
writer_into_reader(writer)
|
|
}
|
|
}
|
|
|
|
/// Fills the list of all pairs of words with the shortest proximity between 1 and 7 inclusive.
|
|
///
|
|
/// This list is used by the engine to calculate the documents containing words that are
|
|
/// close to each other.
|
|
fn document_word_positions_into_sorter(
|
|
document_id: DocumentId,
|
|
del_word_pair_proximity: &BTreeMap<(String, String), u8>,
|
|
add_word_pair_proximity: &BTreeMap<(String, String), u8>,
|
|
word_pair_proximity_docids_sorters: &mut [grenad::Sorter<MergeFn>],
|
|
) -> Result<()> {
|
|
use itertools::merge_join_by;
|
|
use itertools::EitherOrBoth::{Both, Left, Right};
|
|
|
|
let mut buffer = Vec::new();
|
|
let mut key_buffer = Vec::new();
|
|
for eob in
|
|
merge_join_by(del_word_pair_proximity.iter(), add_word_pair_proximity.iter(), |d, a| {
|
|
d.cmp(a)
|
|
})
|
|
{
|
|
buffer.clear();
|
|
let mut value_writer = KvWriterDelAdd::new(&mut buffer);
|
|
let ((w1, w2), prox) = match eob {
|
|
Left(key_value) => {
|
|
value_writer.insert(DelAdd::Deletion, document_id.to_ne_bytes()).unwrap();
|
|
key_value
|
|
}
|
|
Right(key_value) => {
|
|
value_writer.insert(DelAdd::Addition, document_id.to_ne_bytes()).unwrap();
|
|
key_value
|
|
}
|
|
Both(key_value, _) => {
|
|
value_writer.insert(DelAdd::Deletion, document_id.to_ne_bytes()).unwrap();
|
|
value_writer.insert(DelAdd::Addition, document_id.to_ne_bytes()).unwrap();
|
|
key_value
|
|
}
|
|
};
|
|
|
|
key_buffer.clear();
|
|
key_buffer.push(*prox);
|
|
key_buffer.extend_from_slice(w1.as_bytes());
|
|
key_buffer.push(0);
|
|
key_buffer.extend_from_slice(w2.as_bytes());
|
|
|
|
word_pair_proximity_docids_sorters[*prox as usize - 1]
|
|
.insert(&key_buffer, value_writer.into_inner().unwrap())?;
|
|
}
|
|
|
|
Ok(())
|
|
}
|
|
|
|
fn word_positions_into_word_pair_proximity(
|
|
word_positions: &mut VecDeque<(String, u16)>,
|
|
word_pair_proximity: &mut BTreeMap<(String, String), u8>,
|
|
) -> Result<()> {
|
|
let (head_word, head_position) = word_positions.pop_front().unwrap();
|
|
for (word, position) in word_positions.iter() {
|
|
let prox = index_proximity(head_position as u32, *position as u32) as u8;
|
|
if prox > 0 && prox < MAX_DISTANCE as u8 {
|
|
word_pair_proximity
|
|
.entry((head_word.clone(), word.clone()))
|
|
.and_modify(|p| {
|
|
*p = cmp::min(*p, prox);
|
|
})
|
|
.or_insert(prox);
|
|
}
|
|
}
|
|
Ok(())
|
|
}
|