diff --git a/Cargo.lock b/Cargo.lock index 69c627947..a7590c0de 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -790,6 +790,7 @@ dependencies = [ "fst", "fxhash", "heed", + "itertools", "jemallocator", "levenshtein_automata", "memmap", diff --git a/Cargo.toml b/Cargo.toml index ce09f5e06..7a2c0005d 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -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"] } diff --git a/src/bin/indexer.rs b/src/bin/indexer.rs index 5406d9d80..6786b257f 100644 --- a/src/bin/indexer.rs +++ b/src/bin/indexer.rs @@ -46,6 +46,7 @@ struct Opt { struct Indexed { fst: fst::Set>, postings_attrs: FastMap4, + prefix_postings_attrs: FastMap4, postings_ids: FastMap4>, prefix_postings_ids: FastMap4>, headers: Vec, @@ -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,30 +133,33 @@ 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() { - 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); + 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] = 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) -> anyhow::Result { 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) -> anyhow::Result { 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) -> anyhow::Result { 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)?; diff --git a/src/lib.rs b/src/lib.rs index 8017aa911..ff0d4a8ec 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -32,6 +32,7 @@ pub type AttributeId = u32; pub struct Index { pub main: PolyDatabase, pub postings_attrs: Database, + pub prefix_postings_attrs: Database, pub postings_ids: Database, pub prefix_postings_ids: Database, pub documents: Database, ByteSlice>, @@ -41,11 +42,12 @@ impl Index { pub fn new(env: &heed::Env) -> heed::Result { 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> { @@ -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 = None; - for (_word, _is_prefix, dfa) in &dfas { - 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(attrs) = self.postings_attrs.get(rtxn, word)? { - let right = RoaringBitmap::deserialize_from(attrs)?; - union_result.union_with(&right); + 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_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 = 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; + let mut union_docids = RoaringBitmap::default(); - let mut count = 0; - 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()); - if let Some(ids) = self.prefix_postings_ids.get(rtxn, &key)? { - let right = RoaringBitmap::deserialize_from(ids)?; - union_result.union_with(&right); - count = 1; - } - } else { - let mut stream = fst.search(dfa).into_stream(); - while let Some(word) = stream.next() { + if word.len() <= 4 && is_prefix { + let mut key = word.as_bytes()[..word.len().min(5)].to_vec(); + 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_docids.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)?; + let mut key = word.as_bytes().to_vec(); + 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; - 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); - } } } - - 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), } }