Introduce a first basic working positions-based engine

This commit is contained in:
Kerollmops 2020-06-05 20:12:52 +02:00
parent f51a63e4ef
commit dfdaceb410
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4
4 changed files with 135 additions and 83 deletions

1
Cargo.lock generated
View File

@ -790,6 +790,7 @@ dependencies = [
"fst",
"fxhash",
"heed",
"itertools",
"jemallocator",
"levenshtein_automata",
"memmap",

View File

@ -27,6 +27,9 @@ smallvec = "1.4.0"
structopt = { version = "0.3.14", default-features = false }
tempfile = "3.1.0"
# to implement internally
itertools = "0.9.0"
# http server
serde = { version = "1.0", features = ["derive"] }
tokio = { version = "0.2.15", features = ["full"] }

View File

@ -46,6 +46,7 @@ struct Opt {
struct Indexed {
fst: fst::Set<Vec<u8>>,
postings_attrs: FastMap4<SmallVec32, RoaringBitmap>,
prefix_postings_attrs: FastMap4<SmallVec32, RoaringBitmap>,
postings_ids: FastMap4<SmallVec32, FastMap4<AttributeId, RoaringBitmap>>,
prefix_postings_ids: FastMap4<SmallVec32, FastMap4<AttributeId, RoaringBitmap>>,
headers: Vec<u8>,
@ -86,11 +87,31 @@ impl MtblKvStore {
}
}
// We must write the postings ids
// 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);
@ -112,9 +133,11 @@ impl MtblKvStore {
}
// We must write the prefix postings ids
key[0] = 3;
key[0] = 4;
let mut stream = indexed.fst.stream();
while let Some(prefix) = stream.next() {
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) {
@ -133,9 +156,10 @@ impl MtblKvStore {
}
}
}
}
// postings ids keys are all prefixed by a '4'
key[0] = 4;
key[0] = 5;
indexed.documents.sort_unstable();
for (id, content) in indexed.documents {
key.truncate(1);
@ -167,7 +191,7 @@ impl MtblKvStore {
Some(values[0].to_vec())
}
// We either merge postings attrs, prefix postings or postings ids.
else if key.starts_with(&[1]) || key.starts_with(&[2]) || key.starts_with(&[3]) {
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..] {
@ -179,7 +203,7 @@ impl MtblKvStore {
first.serialize_into(&mut vec).unwrap();
Some(vec)
}
else if key.starts_with(&[4]) {
else if key[0] == 5 {
assert!(values.windows(2).all(|vs| vs[0] == vs[1]));
Some(values[0].to_vec())
}
@ -219,6 +243,7 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
let mut document = csv::StringRecord::new();
let mut postings_attrs = FastMap4::default();
let mut prefix_postings_attrs = FastMap4::default();
let mut postings_ids = FastMap4::default();
let mut prefix_postings_ids = FastMap4::default();
let mut documents = Vec::new();
@ -234,21 +259,28 @@ fn index_csv(mut rdr: csv::Reader<File>) -> anyhow::Result<MtblKvStore> {
let document_id = DocumentId::try_from(document_id).context("Generated id is too big")?;
for (attr, content) in document.iter().enumerate().take(MAX_ATTRIBUTES) {
for (_pos, word) in simple_alphanumeric_tokens(&content).enumerate().take(MAX_POSITION) {
for (pos, word) in simple_alphanumeric_tokens(&content).enumerate().take(MAX_POSITION) {
if !word.is_empty() && word.len() < 500 { // LMDB limits
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(attr as u32); // attributes ids
.or_insert_with(RoaringBitmap::new).insert(position);
// 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(attr as u32) // attributes
.or_insert_with(FastMap4::default).entry(position) // positions
.or_insert_with(RoaringBitmap::new).insert(document_id); // document ids
// We save the documents ids under the position and prefix of the word we have seen it.
if let Some(prefix) = word.as_bytes().get(0..word.len().min(5)) {
for i in 0..=prefix.len() {
for i in 1..=prefix.len() {
prefix_postings_attrs.entry(SmallVec32::from(&prefix[..i]))
.or_insert_with(RoaringBitmap::new).insert(position);
prefix_postings_ids.entry(SmallVec32::from(&prefix[..i]))
.or_insert_with(FastMap4::default).entry(attr as u32) // attributes
.or_insert_with(FastMap4::default).entry(position) // positions
.or_insert_with(RoaringBitmap::new).insert(document_id); // document ids
}
}
@ -271,7 +303,15 @@ 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_attrs, postings_ids, prefix_postings_ids, documents };
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)
@ -293,16 +333,21 @@ fn writer(wtxn: &mut heed::RwTxn, index: &Index, key: &[u8], val: &[u8]) -> anyh
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[2]) {
// Write the prefix postings lists
index.prefix_postings_attrs.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[3]) {
// Write the postings lists
index.postings_ids.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[3]) {
else if key.starts_with(&[4]) {
// Write the prefix postings lists
index.prefix_postings_ids.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[4]) {
else if key.starts_with(&[5]) {
// Write the documents
index.documents.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;

View File

@ -32,6 +32,7 @@ pub type AttributeId = u32;
pub struct Index {
pub main: PolyDatabase,
pub postings_attrs: Database<Str, ByteSlice>,
pub prefix_postings_attrs: Database<ByteSlice, ByteSlice>,
pub postings_ids: Database<ByteSlice, ByteSlice>,
pub prefix_postings_ids: Database<ByteSlice, ByteSlice>,
pub documents: Database<OwnedType<BEU32>, ByteSlice>,
@ -41,11 +42,12 @@ impl Index {
pub fn new(env: &heed::Env) -> heed::Result<Index> {
let main = env.create_poly_database(None)?;
let postings_attrs = env.create_database(Some("postings-attrs"))?;
let prefix_postings_attrs = env.create_database(Some("prefix-postings-attrs"))?;
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_attrs, postings_ids, prefix_postings_ids, documents })
Ok(Index { main, postings_attrs, prefix_postings_attrs, postings_ids, prefix_postings_ids, documents })
}
pub fn headers<'t>(&self, rtxn: &'t heed::RoTxn) -> heed::Result<Option<&'t [u8]>> {
@ -63,7 +65,7 @@ impl Index {
let words: Vec<_> = QueryTokens::new(query).collect();
let ends_with_whitespace = query.chars().last().map_or(false, char::is_whitespace);
let number_of_words = words.len();
let dfas: Vec<_> = words.into_iter().enumerate().map(|(i, word)| {
let dfas = words.into_iter().enumerate().map(|(i, word)| {
let (word, quoted) = match word {
QueryToken::Free(word) => (word.cow_to_lowercase(), false),
QueryToken::Quoted(word) => (Cow::Borrowed(word), true),
@ -83,73 +85,74 @@ impl Index {
};
(word, is_prefix, dfa)
})
.collect();
});
let mut intersect_attrs: Option<RoaringBitmap> = None;
for (_word, _is_prefix, dfa) in &dfas {
let mut union_result = RoaringBitmap::default();
let mut stream = fst.search(dfa).into_stream();
let mut words_positions = Vec::new();
for (word, is_prefix, dfa) in dfas {
let mut count = 0;
let mut union_positions = RoaringBitmap::default();
if word.len() <= 4 && is_prefix {
if let Some(ids) = self.prefix_postings_attrs.get(rtxn, word.as_bytes())? {
let right = RoaringBitmap::deserialize_from(ids)?;
union_positions.union_with(&right);
count = 1;
}
} 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(attrs) = self.postings_attrs.get(rtxn, word)? {
let right = RoaringBitmap::deserialize_from(attrs)?;
union_result.union_with(&right);
union_positions.union_with(&right);
count += 1;
}
}
}
match &mut intersect_attrs {
Some(left) => left.intersect_with(&union_result),
None => intersect_attrs = Some(union_result),
}
eprintln!("{} words for {:?} we have found positions {:?}", count, word, union_positions);
words_positions.push((word, is_prefix, dfa, union_positions));
}
eprintln!("we should only look for documents with attrs {:?}", intersect_attrs);
use itertools::EitherOrBoth;
let (a, b) = (&words_positions[0].3, &words_positions[1].3);
let positions: Vec<_> = itertools::merge_join_by(a, b, |a, b| (a + 1).cmp(b)).flat_map(EitherOrBoth::both).collect();
if positions.is_empty() { return Ok(Vec::new()); }
let mut intersect_docids: Option<RoaringBitmap> = None;
// TODO would be faster to store and use the words
// seen in the previous attrs loop
for (word, is_prefix, dfa) in &dfas {
let mut union_result = RoaringBitmap::default();
for attr in intersect_attrs.as_ref().unwrap_or(&RoaringBitmap::default()) {
let before = Instant::now();
for (i, (word, is_prefix, dfa, _)) in words_positions.into_iter().take(2).enumerate() {
let mut count = 0;
if word.len() <= 4 && *is_prefix {
let mut union_docids = RoaringBitmap::default();
if word.len() <= 4 && is_prefix {
let mut key = word.as_bytes()[..word.len().min(5)].to_vec();
key.extend_from_slice(&attr.to_be_bytes());
let pos = if i == 0 { positions[0].0 } else { positions[0].1 };
key.extend_from_slice(&pos.to_be_bytes());
if let Some(ids) = self.prefix_postings_ids.get(rtxn, &key)? {
let right = RoaringBitmap::deserialize_from(ids)?;
union_result.union_with(&right);
union_docids.union_with(&right);
count = 1;
}
} else {
let mut stream = fst.search(dfa).into_stream();
while let Some(word) = stream.next() {
count += 1;
let word = std::str::from_utf8(word)?;
let mut key = word.as_bytes().to_vec();
key.extend_from_slice(&attr.to_be_bytes());
if let Some(ids) = self.postings_ids.get(rtxn, &key)? {
let right = RoaringBitmap::deserialize_from(ids)?;
union_result.union_with(&right);
let pos = if i == 0 { positions[0].0 } else { positions[0].1 };
key.extend_from_slice(&pos.to_be_bytes());
if let Some(attrs) = self.postings_ids.get(rtxn, &key)? {
let right = RoaringBitmap::deserialize_from(attrs)?;
union_docids.union_with(&right);
count += 1;
}
}
}
eprintln!("with {:?} similar words (for attr {}) union for {:?} gives {:?} took {:.02?}",
count, attr, word, union_result.len(), before.elapsed());
}
let _ = count;
match &mut intersect_docids {
Some(left) => {
let before = Instant::now();
let left_len = left.len();
left.intersect_with(&union_result);
eprintln!("intersect between {:?} and {:?} gives {:?} took {:.02?}",
left_len, union_result.len(), left.len(), before.elapsed());
},
None => intersect_docids = Some(union_result),
Some(left) => left.intersect_with(&union_docids),
None => intersect_docids = Some(union_docids),
}
}