snowykami/meilisearch
1
0
Fork 0
mirror of https://github.com/meilisearch/meilisearch.git synced 2025-07-27 00:31:02 +00:00
Code Issues Packages Projects Releases Wiki Activity

Index by writing directly into LMDB

Browse Source
This commit is contained in:
Kerollmops
2020-06-29 13:54:47 +02:00
parent 8453828a65
commit 5f0088594b
4 changed files with 59 additions and 405 deletions
Download Patch File Download Diff File Expand all files Collapse all files

343
src/bin/indexer.rs
View File

@ -1,22 +1,21 @@
use std::collections::hash_map::Entry;
use std::collections::{HashMap, BTreeSet, BTreeMap};
use std::collections::{HashMap, BTreeSet};
use std::convert::{TryFrom, TryInto};
use std::fs::File;
use std::io;
use std::iter::FromIterator;
use std::path::PathBuf;
use std::sync::atomic::{AtomicUsize, Ordering};
use anyhow::Context;
use cow_utils::CowUtils;
use fst::{Streamer, IntoStreamer};
use fst::Streamer;
use heed::EnvOpenOptions;
use heed::types::*;
use oxidized_mtbl::{Reader, ReaderOptions, Writer, Merger, MergerOptions};
use rayon::prelude::*;
use roaring::RoaringBitmap;
use slice_group_by::StrGroupBy;
use structopt::StructOpt;
use mega_mini_indexer::{FastMap4, SmallVec32, Index, DocumentId, AttributeId};
use mega_mini_indexer::{BEU32, Index, DocumentId};
const MAX_POSITION: usize = 1000;
const MAX_ATTRIBUTES: usize = u32::max_value() as usize / MAX_POSITION;
@ -40,218 +39,12 @@ struct Opt {
#[structopt(long = "db", parse(from_os_str))]
database: PathBuf,
/// Number of parallel jobs, defaults to # of CPUs.
#[structopt(short, long)]
jobs: Option<usize>,
/// Files to index in parallel.
files_to_index: Vec<PathBuf>,
/// CSV file to index.
csv_file: Option<PathBuf>,
}
struct Indexed {
fst: fst::Set<Vec<u8>>,
postings_attrs: FastMap4<SmallVec32<u8>, RoaringBitmap>,
prefix_postings_attrs: FastMap4<SmallVec32<u8>, RoaringBitmap>,
postings_ids: FastMap4<SmallVec32<u8>, FastMap4<AttributeId, RoaringBitmap>>,
prefix_postings_ids: FastMap4<SmallVec32<u8>, FastMap4<AttributeId, RoaringBitmap>>,
headers: Vec<u8>,
documents: Vec<(DocumentId, Vec<u8>)>,
}
#[derive(Default)]
struct MtblKvStore(Option<File>);
impl MtblKvStore {
fn from_indexed(mut indexed: Indexed) -> anyhow::Result<MtblKvStore> {
eprintln!("{:?}: Creating an MTBL store from an Indexed...", rayon::current_thread_index());
let outfile = tempfile::tempfile()?;
let mut out = Writer::new(outfile, None)?;
out.add(b"\0headers", indexed.headers)?;
out.add(b"\0words-fst", indexed.fst.as_fst().as_bytes())?;
// postings ids keys are all prefixed by a '1'
let mut key = vec![0];
let mut buffer = Vec::new();
// We must write the postings attrs
key[0] = 1;
// 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() {
if let Some(attrs) = indexed.postings_attrs.remove(word) {
key.truncate(1);
key.extend_from_slice(word);
// We serialize the attrs ids into a buffer
buffer.clear();
attrs.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
out.add(&key, &buffer).unwrap();
}
}
// We must write the prefix postings attrs
key[0] = 2;
// 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() {
for i in 1..=word.len() {
let prefix = &word[..i];
if let Some(attrs) = indexed.prefix_postings_attrs.remove(prefix) {
key.truncate(1);
key.extend_from_slice(prefix);
// We serialize the attrs ids into a buffer
buffer.clear();
attrs.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
out.add(&key, &buffer).unwrap();
}
}
}
// We must write the postings ids
key[0] = 3;
// 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(attrs) = indexed.postings_ids.remove(word) {
let attrs: BTreeMap<_, _> = attrs.into_iter().collect();
// We iterate over all the attributes containing the documents ids
for (attr, ids) in attrs {
// we postfix the word by the attribute id
key.extend_from_slice(&attr.to_be_bytes());
// We serialize the document ids into a buffer
buffer.clear();
ids.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
out.add(&key, &buffer).unwrap();
// And cleanup the attribute id afterward (u32 = 4 * u8)
key.truncate(key.len() - 4);
}
}
}
// We must write the prefix postings ids
key[0] = 4;
let mut stream = indexed.fst.stream();
while let Some(word) = stream.next() {
for i in 1..=word.len() {
let prefix = &word[..i];
key.truncate(1);
key.extend_from_slice(prefix);
if let Some(attrs) = indexed.prefix_postings_ids.remove(prefix) {
let attrs: BTreeMap<_, _> = attrs.into_iter().collect();
// We iterate over all the attributes containing the documents ids
for (attr, ids) in attrs {
// we postfix the word by the attribute id
key.extend_from_slice(&attr.to_be_bytes());
// We serialize the document ids into a buffer
buffer.clear();
ids.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
out.add(&key, &buffer).unwrap();
// And cleanup the attribute id afterward (u32 = 4 * u8)
key.truncate(key.len() - 4);
}
}
}
}
// postings ids keys are all prefixed by a '4'
key[0] = 5;
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();
}
let out = out.into_inner()?;
eprintln!("{:?}: MTBL store created!", rayon::current_thread_index());
Ok(MtblKvStore(Some(out)))
}
fn merge(key: &[u8], values: &[Vec<u8>]) -> Option<Vec<u8>> {
if key == b"\0words-fst" {
let fsts: Vec<_> = values.iter().map(|v| fst::Set::new(v).unwrap()).collect();
// Union of the two FSTs
let mut op = fst::set::OpBuilder::new();
fsts.iter().for_each(|fst| op.push(fst.into_stream()));
let op = op.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!(values.windows(2).all(|vs| vs[0] == vs[1]));
Some(values[0].to_vec())
}
// We either merge postings attrs, prefix postings or postings ids.
else if key[0] == 1 || key[0] == 2 || key[0] == 3 || key[0] == 4 {
let mut first = RoaringBitmap::deserialize_from(values[0].as_slice()).unwrap();
for value in &values[1..] {
let bitmap = RoaringBitmap::deserialize_from(value.as_slice()).unwrap();
first.union_with(&bitmap);
}
let mut vec = Vec::new();
first.serialize_into(&mut vec).unwrap();
Some(vec)
}
else if key[0] == 5 {
assert!(values.windows(2).all(|vs| vs[0] == vs[1]));
Some(values[0].to_vec())
}
else {
panic!("wut? {:?}", key)
}
}
fn from_many<F>(stores: Vec<MtblKvStore>, mut f: F) -> anyhow::Result<()>
where F: FnMut(&[u8], &[u8]) -> anyhow::Result<()>
{
eprintln!("{:?}: Merging {} MTBL stores...", rayon::current_thread_index(), stores.len());
let mmaps: Vec<_> = stores.iter().flat_map(|m| {
m.0.as_ref().map(|f| unsafe { memmap::Mmap::map(f).unwrap() })
}).collect();
let sources = mmaps.iter().map(|mmap| {
Reader::new(&mmap, ReaderOptions::default()).unwrap()
}).collect();
let opt = MergerOptions { merge: MtblKvStore::merge };
let mut merger = Merger::new(sources, opt);
let mut iter = merger.iter();
while let Some((k, v)) = iter.next() {
(f)(k, v)?;
}
eprintln!("{:?}: MTBL stores merged!", rayon::current_thread_index());
Ok(())
}
}
fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
eprintln!("{:?}: Indexing into an Indexed...", rayon::current_thread_index());
let mut document = csv::StringRecord::new();
let mut postings_attrs = FastMap4::default();
let prefix_postings_attrs = FastMap4::default();
let mut postings_ids = FastMap4::default();
let prefix_postings_ids = FastMap4::default();
let mut documents = Vec::new();
fn index_csv<R: io::Read>(wtxn: &mut heed::RwTxn, mut rdr: csv::Reader<R>, index: &Index) -> anyhow::Result<()> {
eprintln!("Indexing into LMDB...");
// Write the headers into a Vec of bytes.
let headers = rdr.headers()?;
@ -259,6 +52,8 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
writer.write_byte_record(headers.as_byte_record())?;
let headers = writer.into_inner()?;
let mut document = csv::StringRecord::new();
while rdr.read_record(&mut document)? {
let document_id = ID_GENERATOR.fetch_add(1, Ordering::SeqCst);
let document_id = DocumentId::try_from(document_id).context("Generated id is too big")?;
@ -269,14 +64,26 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
let word = word.cow_to_lowercase();
let position = (attr * 1000 + pos) as u32;
// We save the positions where this word has been seen.
postings_attrs.entry(SmallVec32::from(word.as_bytes()))
.or_insert_with(RoaringBitmap::new).insert(position);
// ------ merge word positions --------
// We save the documents ids under the position and word we have seen it.
postings_ids.entry(SmallVec32::from(word.as_bytes()))
.or_insert_with(FastMap4::default).entry(position) // positions
.or_insert_with(RoaringBitmap::new).insert(document_id); // document ids
let ids = match index.word_positions.get(wtxn, &word)? {
Some(mut ids) => { ids.insert(position); ids },
None => RoaringBitmap::from_iter(Some(position)),
};
index.word_positions.put(wtxn, &word, &ids)?;
// ------ merge word position documents ids --------
let mut key = word.as_bytes().to_vec();
key.extend_from_slice(&position.to_be_bytes());
let ids = match index.word_position_docids.get(wtxn, &key)? {
Some(mut ids) => { ids.insert(document_id); ids },
None => RoaringBitmap::from_iter(Some(document_id)),
};
index.word_position_docids.put(wtxn, &key, &ids)?;
}
}
}
@ -285,66 +92,21 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
let mut writer = csv::WriterBuilder::new().has_headers(false).from_writer(Vec::new());
writer.write_byte_record(document.as_byte_record())?;
let document = writer.into_inner()?;
documents.push((document_id, document));
index.documents.put(wtxn, &BEU32::new(document_id), &document)?;
}
// We store the words from the postings.
let mut new_words = BTreeSet::default();
for (word, _new_ids) in &postings_ids {
let iter = index.word_positions.as_polymorph().iter::<_, Str, DecodeIgnore>(wtxn)?;
for result in iter {
let (word, ()) = result?;
new_words.insert(word.clone());
}
let new_words_fst = fst::Set::from_iter(new_words.iter().map(SmallVec32::as_ref))?;
let new_words_fst = fst::Set::from_iter(new_words)?;
let indexed = Indexed {
fst: new_words_fst,
headers,
postings_attrs,
prefix_postings_attrs,
postings_ids,
prefix_postings_ids,
documents,
};
eprintln!("{:?}: Indexed created!", rayon::current_thread_index());
MtblKvStore::from_indexed(indexed)
}
// TODO merge with the previous values
fn writer(wtxn: &mut heed::RwTxn, index: &Index, key: &[u8], val: &[u8]) -> anyhow::Result<()> {
if key == b"\0words-fst" {
// Write the words fst
index.main.put::<_, Str, ByteSlice>(wtxn, "words-fst", val)?;
}
else if key == b"\0headers" {
// Write the headers
index.main.put::<_, Str, ByteSlice>(wtxn, "headers", val)?;
}
else if key.starts_with(&[1]) {
// Write the postings lists
index.word_positions.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[2]) {
// Write the prefix postings lists
index.prefix_word_positions.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[3]) {
// Write the postings lists
index.word_position_docids.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[4]) {
// Write the prefix postings lists
index.prefix_word_position_docids.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[5]) {
// Write the documents
index.documents.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
index.put_fst(wtxn, &new_words_fst)?;
index.put_headers(wtxn, &headers)?;
Ok(())
}
@ -392,10 +154,6 @@ fn compute_words_attributes_docids(wtxn: &mut heed::RwTxn, index: &Index) -> any
fn main() -> anyhow::Result<()> {
let opt = Opt::from_args();
if let Some(jobs) = opt.jobs {
rayon::ThreadPoolBuilder::new().num_threads(jobs).build_global()?;
}
std::fs::create_dir_all(&opt.database)?;
let env = EnvOpenOptions::new()
.map_size(100 * 1024 * 1024 * 1024) // 100 GB
@ -405,23 +163,24 @@ fn main() -> anyhow::Result<()> {
let index = Index::new(&env)?;
let stores: Vec<_> = opt.files_to_index
.into_par_iter()
.map(|path| {
let rdr = csv::Reader::from_path(path)?;
index_csv(rdr)
})
.inspect(|_| {
eprintln!("Total number of documents seen so far is {}", ID_GENERATOR.load(Ordering::Relaxed))
})
.collect::<Result<_, _>>()?;
eprintln!("We are writing into LMDB...");
let mut wtxn = env.write_txn()?;
MtblKvStore::from_many(stores, |k, v| writer(&mut wtxn, &index, k, v))?;
match opt.csv_file {
Some(path) => {
let rdr = csv::Reader::from_path(path)?;
index_csv(&mut wtxn, rdr, &index)?;
},
None => {
let rdr = csv::Reader::from_reader(io::stdin());
index_csv(&mut wtxn, rdr, &index)?;
}
};
compute_words_attributes_docids(&mut wtxn, &index)?;
let count = index.documents.len(&wtxn)?;
wtxn.commit()?;
eprintln!("Wrote {} documents into LMDB", count);
Ok(())

8
src/lib.rs
View File

@ -63,10 +63,18 @@ impl Index {
})
}
pub fn put_headers(&self, wtxn: &mut heed::RwTxn, headers: &[u8]) -> anyhow::Result<()> {
Ok(self.main.put::<_, Str, ByteSlice>(wtxn, "headers", headers)?)
}
pub fn headers<'t>(&self, rtxn: &'t heed::RoTxn) -> heed::Result<Option<&'t [u8]>> {
self.main.get::<_, Str, ByteSlice>(rtxn, "headers")
}
pub fn put_fst<A: AsRef<[u8]>>(&self, wtxn: &mut heed::RwTxn, fst: &fst::Set<A>) -> anyhow::Result<()> {
Ok(self.main.put::<_, Str, ByteSlice>(wtxn, "words-fst", fst.as_fst().as_bytes())?)
}
pub fn fst<'t>(&self, rtxn: &'t heed::RoTxn) -> anyhow::Result<Option<fst::Set<&'t [u8]>>> {
match self.main.get::<_, Str, ByteSlice>(rtxn, "words-fst")? {
Some(bytes) => Ok(Some(fst::Set::new(bytes)?)),