2023-09-18 09:59:38 +02:00
|
|
|
use std::collections::{BTreeSet, HashSet};
|
2021-08-16 13:36:30 +02:00
|
|
|
use std::fs::File;
|
2023-09-28 16:26:01 +02:00
|
|
|
use std::io::{self, BufReader};
|
2021-08-16 13:36:30 +02:00
|
|
|
use std::iter::FromIterator;
|
|
|
|
|
2023-09-18 09:59:38 +02:00
|
|
|
use obkv::KvReaderU16;
|
2021-08-16 13:36:30 +02:00
|
|
|
use roaring::RoaringBitmap;
|
|
|
|
|
|
|
|
use super::helpers::{
|
2023-09-18 09:59:38 +02:00
|
|
|
create_sorter, merge_cbo_roaring_bitmaps, merge_roaring_bitmaps, serialize_roaring_bitmap,
|
|
|
|
sorter_into_reader, try_split_array_at, GrenadParameters,
|
2021-08-16 13:36:30 +02:00
|
|
|
};
|
2021-09-02 15:17:52 +02:00
|
|
|
use crate::error::SerializationError;
|
|
|
|
use crate::index::db_name::DOCID_WORD_POSITIONS;
|
2023-09-18 09:59:38 +02:00
|
|
|
use crate::update::MergeFn;
|
|
|
|
use crate::{DocumentId, FieldId, Result};
|
2021-08-16 13:36:30 +02:00
|
|
|
|
|
|
|
/// Extracts the word and the documents ids where this word appear.
|
|
|
|
///
|
|
|
|
/// Returns a grenad reader with the list of extracted words and
|
|
|
|
/// documents ids from the given chunk of docid word positions.
|
2022-03-24 15:22:57 +01:00
|
|
|
///
|
2022-04-04 20:59:20 +02:00
|
|
|
/// The first returned reader is the one for normal word_docids, and the second one is for
|
2022-03-24 15:22:57 +01:00
|
|
|
/// exact_word_docids
|
2021-08-24 13:55:53 +02:00
|
|
|
#[logging_timer::time]
|
2022-02-16 15:28:48 +01:00
|
|
|
pub fn extract_word_docids<R: io::Read + io::Seek>(
|
|
|
|
docid_word_positions: grenad::Reader<R>,
|
2021-08-16 13:36:30 +02:00
|
|
|
indexer: GrenadParameters,
|
2022-03-24 17:00:29 +01:00
|
|
|
exact_attributes: &HashSet<FieldId>,
|
2023-09-18 09:59:38 +02:00
|
|
|
) -> Result<(
|
|
|
|
grenad::Reader<BufReader<File>>,
|
|
|
|
grenad::Reader<BufReader<File>>,
|
|
|
|
grenad::Reader<BufReader<File>>,
|
|
|
|
)> {
|
2023-07-10 18:41:54 +02:00
|
|
|
puffin::profile_function!();
|
|
|
|
|
2021-08-16 13:36:30 +02:00
|
|
|
let max_memory = indexer.max_memory_by_thread();
|
|
|
|
|
|
|
|
let mut word_docids_sorter = create_sorter(
|
2022-09-13 10:40:37 +02:00
|
|
|
grenad::SortAlgorithm::Unstable,
|
2021-08-16 13:36:30 +02:00
|
|
|
merge_roaring_bitmaps,
|
|
|
|
indexer.chunk_compression_type,
|
|
|
|
indexer.chunk_compression_level,
|
|
|
|
indexer.max_nb_chunks,
|
2023-09-18 09:59:38 +02:00
|
|
|
max_memory.map(|x| x / 3),
|
2021-08-16 13:36:30 +02:00
|
|
|
);
|
|
|
|
|
2022-03-24 17:00:29 +01:00
|
|
|
let mut exact_word_docids_sorter = create_sorter(
|
2022-09-13 10:40:37 +02:00
|
|
|
grenad::SortAlgorithm::Unstable,
|
2022-03-24 17:00:29 +01:00
|
|
|
merge_roaring_bitmaps,
|
|
|
|
indexer.chunk_compression_type,
|
|
|
|
indexer.chunk_compression_level,
|
|
|
|
indexer.max_nb_chunks,
|
2023-09-18 09:59:38 +02:00
|
|
|
max_memory.map(|x| x / 3),
|
2022-03-24 17:00:29 +01:00
|
|
|
);
|
|
|
|
|
2023-09-18 09:59:38 +02:00
|
|
|
let mut word_fid_docids_sorter = create_sorter(
|
|
|
|
grenad::SortAlgorithm::Unstable,
|
|
|
|
merge_roaring_bitmaps,
|
|
|
|
indexer.chunk_compression_type,
|
|
|
|
indexer.chunk_compression_level,
|
|
|
|
indexer.max_nb_chunks,
|
|
|
|
max_memory.map(|x| x / 3),
|
|
|
|
);
|
|
|
|
|
|
|
|
let mut current_document_id = None;
|
|
|
|
let mut fid = 0;
|
|
|
|
let mut key_buffer = Vec::new();
|
2021-08-16 13:36:30 +02:00
|
|
|
let mut value_buffer = Vec::new();
|
2023-09-18 09:59:38 +02:00
|
|
|
let mut words = BTreeSet::new();
|
|
|
|
let mut exact_words = BTreeSet::new();
|
2022-02-16 15:28:48 +01:00
|
|
|
let mut cursor = docid_word_positions.into_cursor()?;
|
2023-09-18 09:59:38 +02:00
|
|
|
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 (fid_bytes, _) = try_split_array_at(key)
|
2022-10-13 22:02:54 +02:00
|
|
|
.ok_or(SerializationError::Decoding { db_name: Some(DOCID_WORD_POSITIONS) })?;
|
2021-08-16 13:36:30 +02:00
|
|
|
let document_id = u32::from_be_bytes(document_id_bytes);
|
2023-09-18 09:59:38 +02:00
|
|
|
fid = u16::from_be_bytes(fid_bytes);
|
|
|
|
|
|
|
|
// drain the btreemaps when we change document.
|
|
|
|
if current_document_id.map_or(false, |id| id != document_id) {
|
|
|
|
words_into_sorters(
|
|
|
|
document_id,
|
|
|
|
fid,
|
|
|
|
&mut key_buffer,
|
|
|
|
&mut value_buffer,
|
|
|
|
&mut exact_words,
|
|
|
|
&mut words,
|
|
|
|
&mut exact_word_docids_sorter,
|
|
|
|
&mut word_docids_sorter,
|
|
|
|
&mut word_fid_docids_sorter,
|
|
|
|
)?;
|
|
|
|
}
|
2021-08-16 13:36:30 +02:00
|
|
|
|
2023-09-18 09:59:38 +02:00
|
|
|
current_document_id = Some(document_id);
|
2022-03-24 17:00:29 +01:00
|
|
|
|
2023-09-18 09:59:38 +02:00
|
|
|
// every words contained in an attribute set to exact must be pushed in the exact_words list.
|
|
|
|
if exact_attributes.contains(&fid) {
|
|
|
|
for (_pos, word) in KvReaderU16::new(&value).iter() {
|
|
|
|
exact_words.insert(word.to_vec());
|
|
|
|
}
|
2022-03-24 17:00:29 +01:00
|
|
|
} else {
|
2023-09-18 09:59:38 +02:00
|
|
|
for (_pos, word) in KvReaderU16::new(&value).iter() {
|
|
|
|
words.insert(word.to_vec());
|
2022-03-24 17:00:29 +01:00
|
|
|
}
|
|
|
|
}
|
2021-08-16 13:36:30 +02:00
|
|
|
}
|
|
|
|
|
2023-09-18 09:59:38 +02:00
|
|
|
// We must make sure that don't lose the current document field id
|
|
|
|
if let Some(document_id) = current_document_id {
|
|
|
|
words_into_sorters(
|
|
|
|
document_id,
|
|
|
|
fid,
|
|
|
|
&mut key_buffer,
|
|
|
|
&mut value_buffer,
|
|
|
|
&mut exact_words,
|
|
|
|
&mut words,
|
|
|
|
&mut exact_word_docids_sorter,
|
|
|
|
&mut word_docids_sorter,
|
|
|
|
&mut word_fid_docids_sorter,
|
|
|
|
)?;
|
|
|
|
}
|
|
|
|
|
2022-03-24 15:22:57 +01:00
|
|
|
Ok((
|
|
|
|
sorter_into_reader(word_docids_sorter, indexer)?,
|
2022-03-24 17:00:29 +01:00
|
|
|
sorter_into_reader(exact_word_docids_sorter, indexer)?,
|
2023-09-18 09:59:38 +02:00
|
|
|
sorter_into_reader(word_fid_docids_sorter, indexer)?,
|
2022-03-24 15:22:57 +01:00
|
|
|
))
|
2021-08-16 13:36:30 +02:00
|
|
|
}
|
2023-09-18 09:59:38 +02:00
|
|
|
|
|
|
|
fn words_into_sorters(
|
|
|
|
document_id: DocumentId,
|
|
|
|
fid: FieldId,
|
|
|
|
key_buffer: &mut Vec<u8>,
|
|
|
|
value_buffer: &mut Vec<u8>,
|
|
|
|
exact_words: &mut BTreeSet<Vec<u8>>,
|
|
|
|
words: &mut BTreeSet<Vec<u8>>,
|
|
|
|
exact_word_docids_sorter: &mut grenad::Sorter<MergeFn>,
|
|
|
|
word_docids_sorter: &mut grenad::Sorter<MergeFn>,
|
|
|
|
word_fid_docids_sorter: &mut grenad::Sorter<MergeFn>,
|
|
|
|
) -> Result<()> {
|
|
|
|
puffin::profile_function!();
|
|
|
|
let bitmap = RoaringBitmap::from_iter(Some(document_id));
|
|
|
|
serialize_roaring_bitmap(&bitmap, value_buffer)?;
|
|
|
|
for word_bytes in exact_words.iter() {
|
|
|
|
exact_word_docids_sorter.insert(word_bytes, &mut *value_buffer)?;
|
|
|
|
}
|
|
|
|
|
|
|
|
for word_bytes in words.iter() {
|
|
|
|
word_docids_sorter.insert(word_bytes, &value_buffer)?;
|
|
|
|
}
|
|
|
|
|
|
|
|
for word_bytes in (&*words | &*exact_words).iter() {
|
|
|
|
key_buffer.clear();
|
|
|
|
key_buffer.extend_from_slice(&word_bytes);
|
|
|
|
key_buffer.push(0);
|
|
|
|
key_buffer.extend_from_slice(&fid.to_be_bytes());
|
|
|
|
word_fid_docids_sorter.insert(word_bytes, &value_buffer)?;
|
|
|
|
}
|
|
|
|
|
|
|
|
exact_words.clear();
|
|
|
|
words.clear();
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|