Introduce prefix postings ids for better perfs

This commit is contained in:
Kerollmops 2020-05-31 18:20:49 +02:00
parent a26553c90a
commit dde3e01a59
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
2 changed files with 59 additions and 32 deletions

View File

@ -7,15 +7,15 @@ use std::sync::atomic::{AtomicUsize, Ordering};
use anyhow::Context;
use cow_utils::CowUtils;
use fst::{Streamer, IntoStreamer};
use heed::EnvOpenOptions;
use heed::types::*;
use heed::{EnvOpenOptions, PolyDatabase, Database};
use oxidized_mtbl::{Reader, ReaderOptions, Writer, Merger, MergerOptions};
use rayon::prelude::*;
use roaring::RoaringBitmap;
use structopt::StructOpt;
use mega_mini_indexer::alphanumeric_tokens;
use mega_mini_indexer::{FastMap4, SmallVec32, BEU32, DocumentId};
use mega_mini_indexer::{FastMap4, SmallVec32, BEU32, Index, DocumentId};
#[cfg(target_os = "linux")]
#[global_allocator]
@ -38,6 +38,7 @@ struct Opt {
struct Indexed {
fst: fst::Set<Vec<u8>>,
postings_ids: FastMap4<SmallVec32, RoaringBitmap>,
prefix_postings_ids: FastMap4<SmallVec32, RoaringBitmap>,
headers: Vec<u8>,
documents: Vec<(DocumentId, Vec<u8>)>,
}
@ -69,8 +70,21 @@ impl MtblKvStore {
}
}
// postings ids keys are all prefixed by a '2'
// We must write the prefix postings ids
key[0] = 2;
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.prefix_postings_ids.remove(word) {
buffer.clear();
ids.serialize_into(&mut buffer)?;
out.add(&key, &buffer).unwrap();
}
}
// postings ids keys are all prefixed by a '2'
key[0] = 3;
indexed.documents.sort_unstable();
for (id, content) in indexed.documents {
key.truncate(1);
@ -115,7 +129,7 @@ impl MtblKvStore {
assert_eq!(left, right);
Some(left.to_vec())
}
else if key.starts_with(&[1]) {
else if key.starts_with(&[1]) || key.starts_with(&[2]) {
let mut left = RoaringBitmap::deserialize_from(left).unwrap();
let right = RoaringBitmap::deserialize_from(right).unwrap();
left.union_with(&right);
@ -123,7 +137,7 @@ impl MtblKvStore {
left.serialize_into(&mut vec).unwrap();
Some(vec)
}
else if key.starts_with(&[2]) {
else if key.starts_with(&[3]) {
assert_eq!(left, right);
Some(left.to_vec())
}
@ -155,6 +169,7 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
let mut document = csv::StringRecord::new();
let mut postings_ids = FastMap4::default();
let mut prefix_postings_ids = FastMap4::default();
let mut documents = Vec::new();
// Write the headers into a Vec of bytes.
@ -174,6 +189,11 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
postings_ids.entry(SmallVec32::from(word.as_bytes()))
.or_insert_with(RoaringBitmap::new)
.insert(document_id);
if let Some(prefix) = word.as_bytes().get(0..word.len().min(4)) {
prefix_postings_ids.entry(SmallVec32::from(prefix))
.or_insert_with(RoaringBitmap::new)
.insert(document_id);
}
}
}
}
@ -185,7 +205,7 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
documents.push((document_id, document));
}
// We compute and store the postings list into the DB.
// We store the words from the postings.
let mut new_words = BTreeSet::default();
for (word, _new_ids) in &postings_ids {
new_words.insert(word.clone());
@ -193,20 +213,13 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
let new_words_fst = fst::Set::from_iter(new_words.iter().map(SmallVec32::as_ref))?;
let indexed = Indexed { fst: new_words_fst, headers, postings_ids, documents };
let indexed = Indexed { fst: new_words_fst, headers, postings_ids, prefix_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>,
mtbl_store: MtblKvStore,
) -> anyhow::Result<usize>
{
fn writer(wtxn: &mut heed::RwTxn, index: Index, mtbl_store: MtblKvStore) -> anyhow::Result<usize> {
let mtbl_store = match mtbl_store.0 {
Some(store) => unsafe { memmap::Mmap::map(&store)? },
None => return Ok(0),
@ -216,25 +229,32 @@ fn writer(
// 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())?;
index.main.put::<_, Str, ByteSlice>(wtxn, "words-fst", &fst.as_fst().as_bytes())?;
// Write and merge the headers
let headers = mtbl_store.get(b"\0headers").unwrap();
main.put::<_, Str, ByteSlice>(wtxn, "headers", headers.as_ref())?;
index.main.put::<_, Str, ByteSlice>(wtxn, "headers", headers.as_ref())?;
// Write and merge the postings lists
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)?;
index.postings_ids.put(wtxn, &word, &postings)?;
}
// Write and merge the prefix postings lists
let mut iter = mtbl_store.iter_prefix(&[2]).unwrap();
while let Some((word, postings)) = iter.next() {
let word = std::str::from_utf8(&word[1..]).unwrap();
index.prefix_postings_ids.put(wtxn, &word, &postings)?;
}
// Write the documents
let mut count = 0;
let mut iter = mtbl_store.iter_prefix(&[2]).unwrap();
let mut iter = mtbl_store.iter_prefix(&[3]).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)?;
index.documents.put(wtxn, &BEU32::new(id), &content)?;
count += 1;
}
@ -251,10 +271,7 @@ fn main() -> anyhow::Result<()> {
.max_dbs(5)
.open(opt.database)?;
let main = env.create_poly_database(None)?;
let postings_ids: Database<Str, ByteSlice> = env.create_database(Some("postings-ids"))?;
let documents: Database<OwnedType<BEU32>, ByteSlice> = env.create_database(Some("documents"))?;
let index = Index::new(&env)?;
let res = opt.files_to_index
.into_par_iter()
.try_fold(MtblKvStore::default, |acc, path| {
@ -271,7 +288,7 @@ fn main() -> anyhow::Result<()> {
eprintln!("We are writing into LMDB...");
let mut wtxn = env.write_txn()?;
let count = writer(&mut wtxn, main, postings_ids, documents, mtbl_store)?;
let count = writer(&mut wtxn, index, mtbl_store)?;
wtxn.commit()?;
eprintln!("Wrote {} documents into LMDB", count);

View File

@ -26,6 +26,7 @@ pub fn alphanumeric_tokens(string: &str) -> impl Iterator<Item = &str> {
pub struct Index {
pub main: PolyDatabase,
pub postings_ids: Database<Str, ByteSlice>,
pub prefix_postings_ids: Database<Str, ByteSlice>,
pub documents: Database<OwnedType<BEU32>, ByteSlice>,
}
@ -33,11 +34,13 @@ impl Index {
pub fn new(env: &heed::Env) -> heed::Result<Index> {
let main = env.create_poly_database(None)?;
let postings_ids = env.create_database(Some("postings-ids"))?;
let prefix_postings_ids = env.create_database(Some("prefix-postings-ids"))?;
let documents = env.create_database(Some("documents"))?;
Ok(Index {
main,
postings_ids,
prefix_postings_ids,
documents,
})
}
@ -73,16 +76,23 @@ impl Index {
let mut intersect_result: Option<RoaringBitmap> = None;
for (word, dfa) in dfas {
let before = Instant::now();
let mut union_result = RoaringBitmap::default();
let mut stream = fst.search(dfa).into_stream();
while let Some(word) = stream.next() {
let word = std::str::from_utf8(word)?;
if let Some(ids) = self.postings_ids.get(rtxn, word)? {
let right = RoaringBitmap::deserialize_from(ids)?;
union_result.union_with(&right);
if word.len() <= 4 {
if let Some(ids) = self.prefix_postings_ids.get(rtxn, &word[..word.len().min(4)])? {
union_result = RoaringBitmap::deserialize_from(ids)?;
}
} else {
let mut stream = fst.search(dfa).into_stream();
while let Some(word) = stream.next() {
let word = std::str::from_utf8(word)?;
if let Some(ids) = self.postings_ids.get(rtxn, word)? {
let right = RoaringBitmap::deserialize_from(ids)?;
union_result.union_with(&right);
}
}
eprintln!("union for {:?} took {:.02?}", word, before.elapsed());
}
eprintln!("union for {:?} took {:.02?}", word, before.elapsed());
intersect_result = match intersect_result.take() {
Some(mut left) => {