mirror of
https://github.com/meilisearch/meilisearch.git
synced 2025-07-26 16:21:07 +00:00
Introduce MTBL parallel merging before LMDB writing
This commit is contained in:
Kerollmops
255
src/main.rs
255
src/main.rs
@ -1,23 +1,25 @@
|
||||
use std::collections::hash_map::Entry;
|
||||
use std::collections::{HashMap, BTreeSet};
|
||||
use std::convert::TryFrom;
|
||||
use std::convert::TryInto;
|
||||
use std::fs::File;
|
||||
use std::hash::BuildHasherDefault;
|
||||
use std::path::PathBuf;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
|
||||
use anyhow::{ensure, Context};
|
||||
use roaring::RoaringBitmap;
|
||||
use fst::IntoStreamer;
|
||||
use anyhow::Context;
|
||||
use fst::{Streamer, IntoStreamer};
|
||||
use fxhash::FxHasher32;
|
||||
use heed::{EnvOpenOptions, PolyDatabase, Database};
|
||||
use heed::types::*;
|
||||
use heed::{EnvOpenOptions, PolyDatabase, Database};
|
||||
use oxidized_mtbl::{Reader, ReaderOptions, Writer, Merger, MergerOptions};
|
||||
use rayon::prelude::*;
|
||||
use roaring::RoaringBitmap;
|
||||
use slice_group_by::StrGroupBy;
|
||||
use structopt::StructOpt;
|
||||
|
||||
pub type FastMap4<K, V> = HashMap<K, V, BuildHasherDefault<FxHasher32>>;
|
||||
pub type SmallString32 = smallstr::SmallString<[u8; 32]>;
|
||||
pub type SmallVec32 = smallvec::SmallVec<[u8; 32]>;
|
||||
pub type BEU32 = heed::zerocopy::U32<heed::byteorder::BE>;
|
||||
pub type DocumentId = u32;
|
||||
|
||||
@ -39,100 +41,126 @@ struct Opt {
|
||||
files_to_index: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
fn union_postings_ids(_key: &[u8], old_value: Option<&[u8]>, new_value: RoaringBitmap) -> Option<Vec<u8>> {
|
||||
let result = match old_value {
|
||||
Some(bytes) => {
|
||||
let mut old_value = RoaringBitmap::deserialize_from(bytes).unwrap();
|
||||
old_value.union_with(&new_value);
|
||||
old_value
|
||||
},
|
||||
None => new_value,
|
||||
};
|
||||
|
||||
let mut vec = Vec::new();
|
||||
result.serialize_into(&mut vec).unwrap();
|
||||
Some(vec)
|
||||
}
|
||||
|
||||
fn union_words_fst(key: &[u8], old_value: Option<&[u8]>, new_value: &fst::Set<Vec<u8>>) -> Option<Vec<u8>> {
|
||||
if key != b"words-fst" { unimplemented!() }
|
||||
|
||||
// Do an union of the old and the new set of words.
|
||||
let mut builder = fst::set::OpBuilder::new();
|
||||
|
||||
let old_words = old_value.map(|v| fst::Set::new(v).unwrap());
|
||||
let old_words = old_words.as_ref().map(|v| v.into_stream());
|
||||
if let Some(old_words) = old_words {
|
||||
builder.push(old_words);
|
||||
}
|
||||
|
||||
builder.push(new_value);
|
||||
|
||||
let op = builder.r#union();
|
||||
let mut build = fst::SetBuilder::memory();
|
||||
build.extend_stream(op.into_stream()).unwrap();
|
||||
|
||||
Some(build.into_inner().unwrap())
|
||||
}
|
||||
|
||||
fn alphanumeric_tokens(string: &str) -> impl Iterator<Item = &str> {
|
||||
let is_alphanumeric = |s: &&str| s.chars().next().map_or(false, char::is_alphanumeric);
|
||||
string.linear_group_by_key(|c| c.is_alphanumeric()).filter(is_alphanumeric)
|
||||
}
|
||||
|
||||
#[derive(Default)]
|
||||
struct Indexed {
|
||||
fst: fst::Set<Vec<u8>>,
|
||||
postings_ids: FastMap4<SmallString32, RoaringBitmap>,
|
||||
postings_ids: FastMap4<SmallVec32, RoaringBitmap>,
|
||||
headers: Vec<u8>,
|
||||
documents: Vec<(DocumentId, Vec<u8>)>,
|
||||
}
|
||||
|
||||
impl Indexed {
|
||||
fn merge_with(mut self, mut other: Indexed) -> Indexed {
|
||||
#[derive(Default)]
|
||||
struct MtblKvStore(Option<File>);
|
||||
|
||||
// Union of the two FSTs
|
||||
let op = fst::set::OpBuilder::new()
|
||||
.add(self.fst.into_stream())
|
||||
.add(other.fst.into_stream())
|
||||
.r#union();
|
||||
impl MtblKvStore {
|
||||
fn from_indexed(mut indexed: Indexed) -> anyhow::Result<MtblKvStore> {
|
||||
let outfile = tempfile::tempfile()?;
|
||||
let mut out = Writer::new(outfile, None)?;
|
||||
|
||||
let mut build = fst::SetBuilder::memory();
|
||||
build.extend_stream(op.into_stream()).unwrap();
|
||||
let fst = build.into_set();
|
||||
out.add(b"\0headers", indexed.headers)?;
|
||||
out.add(b"\0words-fst", indexed.fst.as_fst().as_bytes())?;
|
||||
|
||||
// Merge the postings by unions
|
||||
for (word, mut postings) in other.postings_ids {
|
||||
match self.postings_ids.entry(word) {
|
||||
Entry::Occupied(mut entry) => {
|
||||
let old = entry.get();
|
||||
postings.union_with(&old);
|
||||
entry.insert(postings);
|
||||
},
|
||||
Entry::Vacant(entry) => {
|
||||
entry.insert(postings);
|
||||
},
|
||||
// postings ids keys are all prefixed by a '1'
|
||||
let mut key = vec![1];
|
||||
let mut buffer = Vec::new();
|
||||
// We must write the postings ids in order for mtbl therefore
|
||||
// we iterate over the fst to read the words in order
|
||||
let mut stream = indexed.fst.stream();
|
||||
while let Some(word) = stream.next() {
|
||||
key.truncate(1);
|
||||
key.extend_from_slice(word);
|
||||
if let Some(ids) = indexed.postings_ids.remove(word) {
|
||||
buffer.clear();
|
||||
ids.serialize_into(&mut buffer)?;
|
||||
out.add(&key, &buffer).unwrap();
|
||||
}
|
||||
}
|
||||
|
||||
// assert headers are valid
|
||||
if !self.headers.is_empty() {
|
||||
assert_eq!(self.headers, other.headers);
|
||||
// postings ids keys are all prefixed by a '2'
|
||||
key[0] = 2;
|
||||
indexed.documents.sort_unstable();
|
||||
for (id, content) in indexed.documents {
|
||||
key.truncate(1);
|
||||
key.extend_from_slice(&id.to_be_bytes());
|
||||
out.add(&key, content).unwrap();
|
||||
}
|
||||
|
||||
// extend the documents
|
||||
self.documents.append(&mut other.documents);
|
||||
let out = out.into_inner()?;
|
||||
Ok(MtblKvStore(Some(out)))
|
||||
}
|
||||
|
||||
Indexed {
|
||||
fst,
|
||||
postings_ids: self.postings_ids,
|
||||
headers: self.headers,
|
||||
documents: self.documents,
|
||||
fn merge_with(self, other: MtblKvStore) -> anyhow::Result<MtblKvStore> {
|
||||
let (left, right) = match (self.0, other.0) {
|
||||
(Some(left), Some(right)) => (left, right),
|
||||
(Some(left), None) => return Ok(MtblKvStore(Some(left))),
|
||||
(None, Some(right)) => return Ok(MtblKvStore(Some(right))),
|
||||
(None, None) => return Ok(MtblKvStore(None)),
|
||||
};
|
||||
|
||||
let left = unsafe { memmap::Mmap::map(&left)? };
|
||||
let right = unsafe { memmap::Mmap::map(&right)? };
|
||||
|
||||
let left = Reader::new(&left, ReaderOptions::default()).unwrap();
|
||||
let right = Reader::new(&right, ReaderOptions::default()).unwrap();
|
||||
|
||||
fn merge(key: &[u8], left: &[u8], right: &[u8]) -> Option<Vec<u8>> {
|
||||
if key == b"\0words-fst" {
|
||||
let left_fst = fst::Set::new(left).unwrap();
|
||||
let right_fst = fst::Set::new(right).unwrap();
|
||||
|
||||
// Union of the two FSTs
|
||||
let op = fst::set::OpBuilder::new()
|
||||
.add(left_fst.into_stream())
|
||||
.add(right_fst.into_stream())
|
||||
.r#union();
|
||||
|
||||
let mut build = fst::SetBuilder::memory();
|
||||
build.extend_stream(op.into_stream()).unwrap();
|
||||
Some(build.into_inner().unwrap())
|
||||
}
|
||||
else if key == b"\0headers" {
|
||||
assert_eq!(left, right);
|
||||
Some(left.to_vec())
|
||||
}
|
||||
else if key.starts_with(&[1]) {
|
||||
let mut left = RoaringBitmap::deserialize_from(left).unwrap();
|
||||
let right = RoaringBitmap::deserialize_from(right).unwrap();
|
||||
left.union_with(&right);
|
||||
let mut vec = Vec::new();
|
||||
left.serialize_into(&mut vec).unwrap();
|
||||
Some(vec)
|
||||
}
|
||||
else if key.starts_with(&[2]) {
|
||||
assert_eq!(left, right);
|
||||
Some(left.to_vec())
|
||||
}
|
||||
else {
|
||||
panic!("wut? {:?}", key)
|
||||
}
|
||||
}
|
||||
|
||||
let outfile = tempfile::tempfile()?;
|
||||
let mut out = Writer::new(outfile, None)?;
|
||||
|
||||
let sources = vec![left, right];
|
||||
let opt = MergerOptions { merge };
|
||||
let mut merger = Merger::new(sources, opt);
|
||||
|
||||
let mut iter = merger.iter();
|
||||
while let Some((k, v)) = iter.next() {
|
||||
out.add(k, v).unwrap();
|
||||
}
|
||||
|
||||
let out = out.into_inner()?;
|
||||
Ok(MtblKvStore(Some(out)))
|
||||
}
|
||||
}
|
||||
|
||||
fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<Indexed> {
|
||||
fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
|
||||
const MAX_POSITION: usize = 1000;
|
||||
const MAX_ATTRIBUTES: usize = u32::max_value() as usize / MAX_POSITION;
|
||||
|
||||
@ -153,7 +181,7 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<Indexed> {
|
||||
for (_attr, content) in document.iter().enumerate().take(MAX_ATTRIBUTES) {
|
||||
for (_pos, word) in alphanumeric_tokens(&content).enumerate().take(MAX_POSITION) {
|
||||
if !word.is_empty() && word.len() < 500 { // LMDB limits
|
||||
postings_ids.entry(SmallString32::from(word))
|
||||
postings_ids.entry(SmallVec32::from(word.as_bytes()))
|
||||
.or_insert_with(RoaringBitmap::new)
|
||||
.insert(document_id);
|
||||
}
|
||||
@ -173,44 +201,51 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<Indexed> {
|
||||
new_words.insert(word.clone());
|
||||
}
|
||||
|
||||
let new_words_fst = fst::Set::from_iter(new_words.iter().map(SmallString32::as_str))?;
|
||||
let new_words_fst = fst::Set::from_iter(new_words.iter().map(SmallVec32::as_ref))?;
|
||||
|
||||
Ok(Indexed { fst: new_words_fst, headers, postings_ids, documents })
|
||||
let indexed = Indexed { fst: new_words_fst, headers, postings_ids, documents };
|
||||
|
||||
MtblKvStore::from_indexed(indexed)
|
||||
}
|
||||
|
||||
// TODO merge with the previous values
|
||||
fn writer(
|
||||
wtxn: &mut heed::RwTxn,
|
||||
main: PolyDatabase,
|
||||
postings_ids: Database<Str, ByteSlice>,
|
||||
documents: Database<OwnedType<BEU32>, ByteSlice>,
|
||||
indexed: Indexed,
|
||||
mtbl_store: MtblKvStore,
|
||||
) -> anyhow::Result<usize>
|
||||
{
|
||||
// Write and merge the words fst
|
||||
let old_value = main.get::<_, Str, ByteSlice>(wtxn, "words-fst")?;
|
||||
let new_value = union_words_fst(b"words-fst", old_value, &indexed.fst)
|
||||
.context("error while do a words-fst union")?;
|
||||
main.put::<_, Str, ByteSlice>(wtxn, "words-fst", &new_value)?;
|
||||
let mtbl_store = match mtbl_store.0 {
|
||||
Some(store) => unsafe { memmap::Mmap::map(&store)? },
|
||||
None => return Ok(0),
|
||||
};
|
||||
let mtbl_store = Reader::new(&mtbl_store, ReaderOptions::default()).unwrap();
|
||||
|
||||
// Write the words fst
|
||||
let fst = mtbl_store.get(b"\0words-fst").unwrap();
|
||||
let fst = fst::Set::new(fst)?;
|
||||
main.put::<_, Str, ByteSlice>(wtxn, "words-fst", &fst.as_fst().as_bytes())?;
|
||||
|
||||
// Write and merge the headers
|
||||
if let Some(old_headers) = main.get::<_, Str, ByteSlice>(wtxn, "headers")? {
|
||||
ensure!(old_headers == &*indexed.headers, "headers differs from the previous ones");
|
||||
}
|
||||
main.put::<_, Str, ByteSlice>(wtxn, "headers", &indexed.headers)?;
|
||||
let headers = mtbl_store.get(b"\0headers").unwrap();
|
||||
main.put::<_, Str, ByteSlice>(wtxn, "headers", headers.as_ref())?;
|
||||
|
||||
// Write and merge the postings lists
|
||||
for (word, postings) in indexed.postings_ids {
|
||||
let old_value = postings_ids.get(wtxn, word.as_str())?;
|
||||
let new_value = union_postings_ids(word.as_bytes(), old_value, postings)
|
||||
.context("error while do a words-fst union")?;
|
||||
postings_ids.put(wtxn, &word, &new_value)?;
|
||||
let mut iter = mtbl_store.iter_prefix(&[1]).unwrap();
|
||||
while let Some((word, postings)) = iter.next() {
|
||||
let word = std::str::from_utf8(&word[1..]).unwrap();
|
||||
postings_ids.put(wtxn, &word, &postings)?;
|
||||
}
|
||||
|
||||
let count = indexed.documents.len();
|
||||
|
||||
// Write the documents
|
||||
for (id, content) in indexed.documents {
|
||||
let mut count = 0;
|
||||
let mut iter = mtbl_store.iter_prefix(&[2]).unwrap();
|
||||
while let Some((id_bytes, content)) = iter.next() {
|
||||
let id = id_bytes[1..].try_into().map(u32::from_be_bytes).unwrap();
|
||||
documents.put(wtxn, &BEU32::new(id), &content)?;
|
||||
count += 1;
|
||||
}
|
||||
|
||||
Ok(count)
|
||||
@ -232,29 +267,23 @@ fn main() -> anyhow::Result<()> {
|
||||
|
||||
let res = opt.files_to_index
|
||||
.into_par_iter()
|
||||
.try_fold(|| Indexed::default(), |acc, path| {
|
||||
.try_fold(MtblKvStore::default, |acc, path| {
|
||||
let rdr = csv::Reader::from_path(path)?;
|
||||
let indexed = index_csv(rdr)?;
|
||||
Ok(acc.merge_with(indexed)) as anyhow::Result<Indexed>
|
||||
})
|
||||
.map(|indexed| match indexed {
|
||||
Ok(indexed) => {
|
||||
let tid = rayon::current_thread_index();
|
||||
eprintln!("{:?}: A new step to write into LMDB", tid);
|
||||
let mut wtxn = env.write_txn()?;
|
||||
let count = writer(&mut wtxn, main, postings_ids, documents, indexed)?;
|
||||
wtxn.commit()?;
|
||||
eprintln!("{:?}: Wrote {} documents into LMDB", tid, count);
|
||||
Ok(count)
|
||||
},
|
||||
Err(e) => Err(e),
|
||||
let mtbl_store = index_csv(rdr)?;
|
||||
acc.merge_with(mtbl_store)
|
||||
})
|
||||
.inspect(|_| {
|
||||
eprintln!("Total number of documents seen so far is {}", ID_GENERATOR.load(Ordering::Relaxed))
|
||||
})
|
||||
.try_reduce(|| 0, |a, b| Ok(a + b));
|
||||
.try_reduce(MtblKvStore::default, MtblKvStore::merge_with);
|
||||
|
||||
println!("indexed {:?} documents", res);
|
||||
let mtbl_store = res?;
|
||||
|
||||
eprintln!("We are writing into LMDB...");
|
||||
let mut wtxn = env.write_txn()?;
|
||||
let count = writer(&mut wtxn, main, postings_ids, documents, mtbl_store)?;
|
||||
wtxn.commit()?;
|
||||
eprintln!("Wrote {} documents into LMDB", count);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user