Initial commit

src/bp_vec.rs

@@ -0,0 +1,197 @@
+use byteorder::{ByteOrder, NativeEndian};
+use bitpacking::{BitPacker, BitPacker4x};
+
+/// An append only bitpacked u32 vector that ignore order of insertion.
+#[derive(Default)]
+pub struct BpVec {
+    compressed: Vec<u8>,
+    uncompressed: Vec<u32>,
+}
+
+impl BpVec {
+    pub fn new() -> BpVec {
+        BpVec::default()
+    }
+
+    pub fn push(&mut self, elem: u32) {
+        self.uncompressed.push(elem);
+        if self.uncompressed.len() == BitPacker4x::BLOCK_LEN {
+            encode(&mut self.uncompressed[..], &mut self.compressed);
+            self.uncompressed.clear();
+        }
+    }
+
+    pub fn extend_from_slice(&mut self, elems: &[u32]) {
+        self.uncompressed.extend_from_slice(elems);
+        let remaining = self.uncompressed.len() % BitPacker4x::BLOCK_LEN;
+        for chunk in self.uncompressed[remaining..].chunks_exact_mut(BitPacker4x::BLOCK_LEN) {
+            encode(chunk, &mut self.compressed);
+        }
+        self.uncompressed.truncate(remaining);
+        self.uncompressed.shrink_to_fit();
+    }
+
+    pub fn to_vec(self) -> Vec<u32> {
+        let BpVec { compressed, mut uncompressed } = self;
+        decode(&compressed, &mut uncompressed);
+        uncompressed
+    }
+
+    pub fn capacity(&self) -> usize {
+        self.compressed.capacity() + self.uncompressed.capacity()
+    }
+}
+
+fn encode(items: &mut [u32], encoded: &mut Vec<u8>) {
+    assert_eq!(items.len(), BitPacker4x::BLOCK_LEN);
+
+    let bitpacker = BitPacker4x::new();
+
+    // We reserve enough space in the output buffer, filled with zeroes.
+    let len = encoded.len();
+    // initial_value + num_bits + encoded numbers
+    let max_possible_length = 4 + 1 + 4 * BitPacker4x::BLOCK_LEN;
+    encoded.resize(len + max_possible_length, 0);
+
+    // We sort the items to be able to efficiently bitpack them.
+    items.sort_unstable();
+    // We save the initial value to us for this block, the lowest one.
+    let initial_value = items[0];
+    // We compute the number of bits necessary to encode this block
+    let num_bits = bitpacker.num_bits_sorted(initial_value, items);
+
+    // We write the initial value for this block.
+    let buffer = &mut encoded[len..];
+    NativeEndian::write_u32(buffer, initial_value);
+    // We write the num_bits that will be read to decode this block
+    let buffer = &mut buffer[4..];
+    buffer[0] = num_bits;
+    // We encode the block numbers into the buffer using the num_bits
+    let buffer = &mut buffer[1..];
+    let compressed_len = bitpacker.compress_sorted(initial_value, items, buffer, num_bits);
+
+    // We truncate the buffer to the avoid leaking padding zeroes
+    encoded.truncate(len + 4 + 1 + compressed_len);
+}
+
+fn decode(mut encoded: &[u8], decoded: &mut Vec<u32>) {
+    let bitpacker = BitPacker4x::new();
+
+    // initial_value + num_bits
+    while let Some(header) = encoded.get(0..4 + 1) {
+        // We extract the header informations
+        let initial_value = NativeEndian::read_u32(header);
+        let num_bits = header[4];
+        let bytes = &encoded[4 + 1..];
+
+        // If the num_bits is equal to zero it means that all encoded numbers were zeroes
+        if num_bits == 0 {
+            decoded.resize(decoded.len() + BitPacker4x::BLOCK_LEN, initial_value);
+            encoded = bytes;
+            continue;
+        }
+
+        // We guess the block size based on the num_bits used for this block
+        let block_size = BitPacker4x::compressed_block_size(num_bits);
+
+        // We pad the decoded vector with zeroes
+        let new_len = decoded.len() + BitPacker4x::BLOCK_LEN;
+        decoded.resize(new_len, 0);
+
+        // Create a view into the decoded buffer and decode into it
+        let to_decompress = &mut decoded[new_len - BitPacker4x::BLOCK_LEN..new_len];
+        bitpacker.decompress_sorted(initial_value, &bytes[..block_size], to_decompress, num_bits);
+
+        // Advance the bytes offset to read the next block (+ num_bits)
+        encoded = &bytes[block_size..];
+    }
+}
+
+impl sdset::Collection<u32> for BpVec {
+    fn push(&mut self, elem: u32) {
+        BpVec::push(self, elem);
+    }
+
+    fn extend_from_slice(&mut self, elems: &[u32]) {
+        BpVec::extend_from_slice(self, elems);
+    }
+
+    fn extend<I>(&mut self, elems: I) where I: IntoIterator<Item=u32> {
+        elems.into_iter().for_each(|x| BpVec::push(self, x));
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    quickcheck! {
+        fn qc_push(xs: Vec<u32>) -> bool {
+            let mut xs: Vec<_> = xs.iter().cloned().cycle().take(1300).collect();
+
+            let mut bpvec = BpVec::new();
+            xs.iter().for_each(|x| bpvec.push(*x));
+            let mut result = bpvec.to_vec();
+
+            result.sort_unstable();
+            xs.sort_unstable();
+
+            xs == result
+        }
+    }
+
+    quickcheck! {
+        fn qc_extend_from_slice(xs: Vec<u32>) -> bool {
+            let mut xs: Vec<_> = xs.iter().cloned().cycle().take(1300).collect();
+
+            let mut bpvec = BpVec::new();
+            bpvec.extend_from_slice(&xs);
+            let mut result = bpvec.to_vec();
+
+            result.sort_unstable();
+            xs.sort_unstable();
+
+            xs == result
+        }
+    }
+
+    #[test]
+    fn empty() {
+        let mut bpvec = BpVec::new();
+        bpvec.extend_from_slice(&[]);
+        let result = bpvec.to_vec();
+
+        assert!(result.is_empty());
+    }
+
+    #[test]
+    fn one_zero() {
+        let mut bpvec = BpVec::new();
+        bpvec.extend_from_slice(&[0]);
+        let result = bpvec.to_vec();
+
+        assert_eq!(&[0], &*result);
+    }
+
+    #[test]
+    fn many_zeros() {
+        let xs: Vec<_> = std::iter::repeat(0).take(1300).collect();
+
+        let mut bpvec = BpVec::new();
+        bpvec.extend_from_slice(&xs);
+        let result = bpvec.to_vec();
+
+        assert_eq!(xs, result);
+    }
+
+    #[test]
+    fn many_ones() {
+        let xs: Vec<_> = std::iter::repeat(1).take(1300).collect();
+
+        let mut bpvec = BpVec::new();
+        bpvec.extend_from_slice(&xs);
+        let result = bpvec.to_vec();
+
+        assert_eq!(xs, result);
+    }
+}

src/codec/bitpacker_sorted.rs

@@ -0,0 +1,84 @@
+use bitpacking::{BitPacker, BitPacker4x};
+use byteorder::{ReadBytesExt, NativeEndian};
+use zerocopy::AsBytes;
+
+pub struct CodecBitPacker4xSorted;
+
+impl CodecBitPacker4xSorted {
+    pub fn bytes_encode(item: &[u32]) -> Option<Vec<u8>> {
+        // This is a hotfix to the SIGSEGV
+        // https://github.com/tantivy-search/bitpacking/issues/23
+        if item.is_empty() {
+            return Some(Vec::default())
+        }
+
+        let bitpacker = BitPacker4x::new();
+        let mut compressed = Vec::new();
+        let mut initial_value = 0;
+
+        // The number of remaining numbers that don't fit in the block size.
+        compressed.push((item.len() % BitPacker4x::BLOCK_LEN) as u8);
+
+        // we cannot use a mut slice here because of #68630, TooGeneric error.
+        // we can probably avoid this new allocation by directly using the compressed final Vec.
+        let mut buffer = vec![0u8; 4 * BitPacker4x::BLOCK_LEN];
+
+        for chunk in item.chunks(BitPacker4x::BLOCK_LEN) {
+            if chunk.len() == BitPacker4x::BLOCK_LEN {
+                // compute the number of bits necessary to encode this block
+                let num_bits = bitpacker.num_bits_sorted(initial_value, chunk);
+                // Encode the block numbers into the buffer using the num_bits
+                let compressed_len = bitpacker.compress_sorted(initial_value, chunk, &mut buffer, num_bits);
+                // Write the num_bits that will be read to decode this block
+                compressed.push(num_bits);
+                // Wrtie the bytes of the compressed block numbers
+                compressed.extend_from_slice(&buffer[..compressed_len]);
+                // Save the initial_value, which is the last value of the n-1 used for the n block
+                initial_value = *chunk.last().unwrap();
+            } else {
+                // Save the remaining numbers which don't fit inside of a BLOCK_LEN
+                compressed.extend_from_slice(chunk.as_bytes());
+            }
+        }
+
+        Some(compressed)
+    }
+
+    pub fn bytes_decode(bytes: &[u8]) -> Option<Vec<u32>> {
+        if bytes.is_empty() {
+            return Some(Vec::new())
+        }
+
+        let bitpacker = BitPacker4x::new();
+        let (remaining, bytes) = bytes.split_first().unwrap();
+        let remaining = *remaining as usize;
+
+        let (mut bytes, mut remaining_bytes) = bytes.split_at(bytes.len() - remaining * 4);
+        let mut decompressed = Vec::new();
+        let mut initial_value = 0;
+
+        while let Some(num_bits) = bytes.get(0) {
+            let block_size = BitPacker4x::compressed_block_size(*num_bits);
+
+            let new_len = decompressed.len() + BitPacker4x::BLOCK_LEN;
+            decompressed.resize(new_len, 0);
+
+            // Create a view into the decompressed buffer and decomress into it
+            let to_decompress = &mut decompressed[new_len - BitPacker4x::BLOCK_LEN..new_len];
+            bitpacker.decompress_sorted(initial_value, &bytes[1..block_size + 1], to_decompress, *num_bits);
+
+            // Set the new initial_value for the next block
+            initial_value = *decompressed.last().unwrap();
+            // Advance the bytes offset to read the next block (+ num_bits)
+            bytes = &bytes[block_size + 1..];
+        }
+
+        // We add the remaining uncompressed numbers.
+        let new_len = decompressed.len() + remaining;
+        decompressed.resize(new_len, 0);
+        let to_decompress = &mut decompressed[new_len - remaining..new_len];
+        remaining_bytes.read_u32_into::<NativeEndian>(to_decompress).ok()?;
+
+        Some(decompressed)
+    }
+}

src/codec/mod.rs

@@ -0,0 +1,3 @@
+mod bitpacker_sorted;
+
+pub use self::bitpacker_sorted::CodecBitPacker4xSorted;

src/main.rs

@@ -0,0 +1,186 @@
+#[cfg(test)]
+#[macro_use] extern crate quickcheck;
+
+mod codec;
+mod bp_vec;
+
+use std::collections::{HashMap, BTreeSet};
+use std::convert::TryFrom;
+use std::fs::File;
+use std::hash::BuildHasherDefault;
+use std::path::PathBuf;
+
+use anyhow::{ensure, Context};
+use fst::IntoStreamer;
+use fxhash::FxHasher32;
+use rayon::prelude::*;
+use sdset::{SetOperation, SetBuf};
+use slice_group_by::StrGroupBy;
+use structopt::StructOpt;
+
+use self::codec::CodecBitPacker4xSorted;
+use self::bp_vec::BpVec;
+
+pub type FastMap4<K, V> = HashMap<K, V, BuildHasherDefault<FxHasher32>>;
+pub type SmallString32 = smallstr::SmallString<[u8; 32]>;
+
+#[cfg(target_os = "linux")]
+#[global_allocator]
+static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
+
+#[derive(Debug, StructOpt)]
+#[structopt(name = "mm-indexer", about = "The server side of the daugt project.")]
+struct Opt {
+    /// The database path where the database is located.
+    /// It is created if it doesn't already exist.
+    #[structopt(long = "db", parse(from_os_str))]
+    database: PathBuf,
+
+    /// Files to index in parallel.
+    files_to_index: Vec<PathBuf>,
+}
+
+fn union_bitpacked_postings_ids(_key: &[u8], old_value: Option<&[u8]>, new_value: &[u8]) -> Option<Vec<u8>> {
+    if old_value.is_none() {
+        return Some(new_value.to_vec())
+    }
+
+    let old_value = old_value.unwrap_or_default();
+    let old_value = CodecBitPacker4xSorted::bytes_decode(&old_value).unwrap();
+    let new_value = CodecBitPacker4xSorted::bytes_decode(&new_value).unwrap();
+
+    let old_set = SetBuf::new(old_value).unwrap();
+    let new_set = SetBuf::new(new_value).unwrap();
+
+    let result = sdset::duo::Union::new(&old_set, &new_set).into_set_buf();
+    let compressed = CodecBitPacker4xSorted::bytes_encode(&result).unwrap();
+
+    Some(compressed)
+}
+
+fn union_words_fst(key: &[u8], old_value: Option<&[u8]>, new_value: &[u8]) -> Option<Vec<u8>> {
+    if key != b"words-fst" { unimplemented!() }
+
+    let old_value = match old_value {
+        Some(old_value) => old_value,
+        None => return Some(new_value.to_vec()),
+    };
+
+    eprintln!("old_words size: {}", old_value.len());
+    eprintln!("new_words size: {}", new_value.len());
+
+    let old_words = fst::Set::new(old_value).unwrap();
+    let new_words = fst::Set::new(new_value).unwrap();
+
+    // Do an union of the old and the new set of words.
+    let op = old_words.op().add(new_words.into_stream()).r#union();
+    let mut build = fst::SetBuilder::memory();
+    build.extend_stream(op.into_stream()).unwrap();
+
+    Some(build.into_inner().unwrap())
+}
+
+fn alphanumeric_tokens(string: &str) -> impl Iterator<Item = &str> {
+    let is_alphanumeric = |s: &&str| s.chars().next().map_or(false, char::is_alphanumeric);
+    string.linear_group_by_key(|c| c.is_alphanumeric()).filter(is_alphanumeric)
+}
+
+fn index_csv(tid: usize, db: sled::Db, mut rdr: csv::Reader<File>) -> anyhow::Result<usize> {
+    const MAX_POSITION: usize = 1000;
+    const MAX_ATTRIBUTES: usize = u32::max_value() as usize / MAX_POSITION;
+
+    let main = &*db;
+    let postings_ids = db.open_tree("postings-ids")?;
+    let documents = db.open_tree("documents")?;
+
+    let mut document = csv::StringRecord::new();
+    let mut new_postings_ids = FastMap4::default();
+    let mut new_words = BTreeSet::default();
+    let mut number_of_documents = 0;
+
+    // Write the headers into a Vec of bytes.
+    let headers = rdr.headers()?;
+    let mut writer = csv::WriterBuilder::new().has_headers(false).from_writer(Vec::new());
+    writer.write_byte_record(headers.as_byte_record())?;
+    let headers = writer.into_inner()?;
+
+    if let Some(old_headers) = main.insert("headers", headers.as_slice())? {
+        ensure!(old_headers == headers, "headers differs from the previous ones");
+    }
+
+    while rdr.read_record(&mut document)? {
+        let document_id = db.generate_id()?;
+        let document_id = u32::try_from(document_id).context("Generated id is too big")?;
+
+        for (_attr, content) in document.iter().enumerate().take(MAX_ATTRIBUTES) {
+            for (_pos, word) in alphanumeric_tokens(&content).enumerate().take(MAX_POSITION) {
+                new_postings_ids.entry(SmallString32::from(word)).or_insert_with(BpVec::new).push(document_id);
+            }
+        }
+
+        // We write the document in the database.
+        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.insert(document_id.to_be_bytes(), document)?;
+
+        number_of_documents += 1;
+        if number_of_documents % 100000 == 0 {
+            let postings_ids_size = new_postings_ids.iter().map(|(_, v)| v.capacity() * 4).sum::<usize>();
+            eprintln!("{}, documents seen {}, postings size {}",
+                tid, number_of_documents, postings_ids_size);
+        }
+    }
+
+    eprintln!("Start collecting the postings lists and words");
+
+    // We compute and store the postings list into the DB.
+    for (word, new_ids) in new_postings_ids {
+        let new_ids = SetBuf::from_dirty(new_ids.to_vec());
+        let compressed = CodecBitPacker4xSorted::bytes_encode(&new_ids)
+            .context("error while compressing using CodecBitPacker4xSorted")?;
+
+        postings_ids.merge(word.as_bytes(), compressed)?;
+
+        new_words.insert(word);
+    }
+
+    eprintln!("Finished collecting the postings lists and words");
+
+    eprintln!("Start merging the words-fst");
+
+    let new_words_fst = fst::Set::from_iter(new_words.iter().map(|s| s.as_str()))?;
+    drop(new_words);
+    main.merge("words-fst", new_words_fst.as_fst().as_bytes())?;
+
+    eprintln!("Finished merging the words-fst");
+
+    Ok(number_of_documents)
+}
+
+fn main() -> anyhow::Result<()> {
+    let opt = Opt::from_args();
+
+    let db = sled::open(opt.database)?;
+    let main = &*db;
+
+    // Setup the merge operators
+    main.set_merge_operator(union_words_fst);
+    let postings_ids = db.open_tree("postings-ids")?;
+    postings_ids.set_merge_operator(union_bitpacked_postings_ids);
+    // ...
+    let _documents = db.open_tree("documents")?;
+
+    let res = opt.files_to_index
+        .into_par_iter()
+        .enumerate()
+        .map(|(tid, path)| {
+            let rdr = csv::Reader::from_path(path)?;
+            index_csv(tid, db.clone(), rdr)
+        })
+        .try_reduce(|| 0, |a, b| Ok(a + b));
+
+    println!("{:?}", res);
+
+    Ok(())
+}