refactor spawn_extraction_task

This commit is contained in:
ad hoc
2022-03-23 14:48:15 +01:00
parent f82d4b36eb
commit 5f9f82757d
3 changed files with 69 additions and 30 deletions

View File

@ -26,7 +26,7 @@ use self::extract_word_pair_proximity_docids::extract_word_pair_proximity_docids
use self::extract_word_position_docids::extract_word_position_docids;
use super::helpers::{
as_cloneable_grenad, keep_first_prefix_value_merge_roaring_bitmaps, merge_cbo_roaring_bitmaps,
merge_readers, merge_roaring_bitmaps, CursorClonableMmap, GrenadParameters, MergeFn,
merge_roaring_bitmaps, CursorClonableMmap, GrenadParameters, MergeFn, MergeableReader,
};
use super::{helpers, TypedChunk};
use crate::{FieldId, Result};
@ -66,7 +66,7 @@ pub(crate) fn data_from_obkv_documents(
(docid_fid_facet_numbers_chunks, docid_fid_facet_strings_chunks),
) = result?;
spawn_extraction_task(
spawn_extraction_task::<_, _, Vec<grenad::Reader<File>>>(
docid_word_positions_chunks.clone(),
indexer.clone(),
lmdb_writer_sx.clone(),
@ -76,7 +76,7 @@ pub(crate) fn data_from_obkv_documents(
"word-pair-proximity-docids",
);
spawn_extraction_task(
spawn_extraction_task::<_, _, Vec<grenad::Reader<File>>>(
docid_word_positions_chunks.clone(),
indexer.clone(),
lmdb_writer_sx.clone(),
@ -86,7 +86,7 @@ pub(crate) fn data_from_obkv_documents(
"field-id-wordcount-docids",
);
spawn_extraction_task(
spawn_extraction_task::<_, _, Vec<grenad::Reader<File>>>(
docid_word_positions_chunks.clone(),
indexer.clone(),
lmdb_writer_sx.clone(),
@ -96,7 +96,7 @@ pub(crate) fn data_from_obkv_documents(
"word-docids",
);
spawn_extraction_task(
spawn_extraction_task::<_, _, Vec<grenad::Reader<File>>>(
docid_word_positions_chunks.clone(),
indexer.clone(),
lmdb_writer_sx.clone(),
@ -106,7 +106,7 @@ pub(crate) fn data_from_obkv_documents(
"word-position-docids",
);
spawn_extraction_task(
spawn_extraction_task::<_, _, Vec<grenad::Reader<File>>>(
docid_fid_facet_strings_chunks.clone(),
indexer.clone(),
lmdb_writer_sx.clone(),
@ -116,7 +116,7 @@ pub(crate) fn data_from_obkv_documents(
"field-id-facet-string-docids",
);
spawn_extraction_task(
spawn_extraction_task::<_, _, Vec<grenad::Reader<File>>>(
docid_fid_facet_numbers_chunks.clone(),
indexer.clone(),
lmdb_writer_sx.clone(),
@ -133,7 +133,7 @@ pub(crate) fn data_from_obkv_documents(
/// Generated grenad chunks are merged using the merge_fn.
/// The result of merged chunks is serialized as TypedChunk using the serialize_fn
/// and sent into lmdb_writer_sx.
fn spawn_extraction_task<FE, FS>(
fn spawn_extraction_task<FE, FS, M>(
chunks: Vec<grenad::Reader<CursorClonableMmap>>,
indexer: GrenadParameters,
lmdb_writer_sx: Sender<Result<TypedChunk>>,
@ -142,19 +142,21 @@ fn spawn_extraction_task<FE, FS>(
serialize_fn: FS,
name: &'static str,
) where
FE: Fn(grenad::Reader<CursorClonableMmap>, GrenadParameters) -> Result<grenad::Reader<File>>
FE: Fn(grenad::Reader<CursorClonableMmap>, GrenadParameters) -> Result<M::Output>
+ Sync
+ Send
+ 'static,
FS: Fn(grenad::Reader<File>) -> TypedChunk + Sync + Send + 'static,
FS: Fn(M::Output) -> TypedChunk + Sync + Send + 'static,
M: MergeableReader + FromParallelIterator<M::Output> + Send + 'static,
M::Output: Send,
{
rayon::spawn(move || {
let chunks: Result<Vec<_>> =
let chunks: Result<M> =
chunks.into_par_iter().map(|chunk| extract_fn(chunk, indexer.clone())).collect();
rayon::spawn(move || match chunks {
Ok(chunks) => {
debug!("merge {} database", name);
let reader = merge_readers(chunks, merge_fn, indexer);
let reader = chunks.merge(merge_fn, &indexer);
let _ = lmdb_writer_sx.send(reader.map(|r| serialize_fn(r)));
}
Err(e) => {