Introduce a first basic working positions-based engine

Cargo.lock

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

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"] }

src/bin/indexer.rs

@@ -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,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() {
+        while let Some(word) = stream.next() {
-            key.truncate(1);
+            for i in 1..=word.len() {
-            key.extend_from_slice(prefix);
+                let prefix = &word[..i];
-            if let Some(attrs) = indexed.prefix_postings_ids.remove(prefix) {
+                key.truncate(1);
-                let attrs: BTreeMap<_, _> = attrs.into_iter().collect();
+                key.extend_from_slice(prefix);
-                // We iterate over all the attributes containing the documents ids
+                if let Some(attrs) = indexed.prefix_postings_ids.remove(prefix) {
-                for (attr, ids) in attrs {
+                    let attrs: BTreeMap<_, _> = attrs.into_iter().collect();
-                    // we postfix the word by the attribute id
+                    // We iterate over all the attributes containing the documents ids
-                    key.extend_from_slice(&attr.to_be_bytes());
+                    for (attr, ids) in attrs {
-                    // We serialize the document ids into a buffer
+                        // we postfix the word by the attribute id
-                    buffer.clear();
+                        key.extend_from_slice(&attr.to_be_bytes());
-                    ids.serialize_into(&mut buffer)?;
+                        // We serialize the document ids into a buffer
-                    // that we write under the generated key into MTBL
+                        buffer.clear();
-                    out.add(&key, &buffer).unwrap();
+                        ids.serialize_into(&mut buffer)?;
-                    // And cleanup the attribute id afterward (u32 = 4 * u8)
+                        // that we write under the generated key into MTBL
-                    key.truncate(key.len() - 4);
+                        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<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)?;

src/lib.rs

@@ -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;
+        let mut words_positions = Vec::new();
-        for (_word, _is_prefix, dfa) in &dfas {
+        for (word, is_prefix, dfa) in dfas {
-            let mut union_result = RoaringBitmap::default();
+            let mut count = 0;
-            let mut stream = fst.search(dfa).into_stream();
+            let mut union_positions = RoaringBitmap::default();
-            while let Some(word) = stream.next() {
+            if word.len() <= 4 && is_prefix {
-                let word = std::str::from_utf8(word)?;
+                if let Some(ids) = self.prefix_postings_attrs.get(rtxn, word.as_bytes())? {
-                if let Some(attrs) = self.postings_attrs.get(rtxn, word)? {
+                    let right = RoaringBitmap::deserialize_from(ids)?;
-                    let right = RoaringBitmap::deserialize_from(attrs)?;
+                    union_positions.union_with(&right);
-                    union_result.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 {
+            eprintln!("{} words for {:?} we have found positions {:?}", count, word, union_positions);
-                Some(left) => left.intersect_with(&union_result),
+            words_positions.push((word, is_prefix, dfa, union_positions));
-                None => intersect_attrs = Some(union_result),
-            }
         }
 
-        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
+        for (i, (word, is_prefix, dfa, _)) in words_positions.into_iter().take(2).enumerate() {
-        //      seen in the previous attrs loop
+            let mut count = 0;
-        for (word, is_prefix, dfa) in &dfas {
+            let mut union_docids = RoaringBitmap::default();
-            let mut union_result = RoaringBitmap::default();
-            for attr in intersect_attrs.as_ref().unwrap_or(&RoaringBitmap::default()) {
-                let before = Instant::now();
 
-                let mut count = 0;
+            if word.len() <= 4 && is_prefix {
-                if word.len() <= 4 && *is_prefix {
+                let mut key = word.as_bytes()[..word.len().min(5)].to_vec();
-                    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(&attr.to_be_bytes());
+                key.extend_from_slice(&pos.to_be_bytes());
-                    if let Some(ids) = self.prefix_postings_ids.get(rtxn, &key)? {
+                if let Some(ids) = self.prefix_postings_ids.get(rtxn, &key)? {
-                        let right = RoaringBitmap::deserialize_from(ids)?;
+                    let right = RoaringBitmap::deserialize_from(ids)?;
-                        union_result.union_with(&right);
+                    union_docids.union_with(&right);
-                        count = 1;
+                    count = 1;
-                    }
+                }
-                } else {
+            } else {
-                    let mut stream = fst.search(dfa).into_stream();
+                let mut stream = fst.search(dfa).into_stream();
-                    while let Some(word) = stream.next() {
+                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) => {
+                Some(left) => left.intersect_with(&union_docids),
-                    let before = Instant::now();
+                None => intersect_docids = Some(union_docids),
-                    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),
             }
         }