mirror of
https://github.com/meilisearch/meilisearch.git
synced 2025-07-26 16:21:07 +00:00
Introduce a thread that write to heed
This commit is contained in:
Kerollmops
201
src/bp_vec.rs
201
src/bp_vec.rs
@ -1,201 +0,0 @@
|
||||
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 compressed_capacity(&self) -> usize {
|
||||
self.compressed.capacity()
|
||||
}
|
||||
|
||||
pub fn uncompressed_capacity(&self) -> usize {
|
||||
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);
|
||||
}
|
||||
}
|
||||
@ -1,90 +0,0 @@
|
||||
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) {
|
||||
if *num_bits == 0 {
|
||||
decompressed.resize(decompressed.len() + BitPacker4x::BLOCK_LEN, initial_value);
|
||||
bytes = &bytes[1..];
|
||||
continue;
|
||||
}
|
||||
|
||||
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)
|
||||
}
|
||||
}
|
||||
@ -1,3 +0,0 @@
|
||||
mod bitpacker_sorted;
|
||||
|
||||
pub use self::bitpacker_sorted::CodecBitPacker4xSorted;
|
||||
251
src/main.rs
251
src/main.rs
@ -1,37 +1,32 @@
|
||||
#[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 std::sync::Arc;
|
||||
use std::sync::atomic::{AtomicUsize, Ordering};
|
||||
use std::thread;
|
||||
|
||||
use anyhow::{ensure, Context};
|
||||
use roaring::RoaringBitmap;
|
||||
use crossbeam_channel::{select, Sender, Receiver};
|
||||
use fst::IntoStreamer;
|
||||
use fxhash::FxHasher32;
|
||||
use heed::{EnvOpenOptions, Database};
|
||||
use heed::types::*;
|
||||
use rayon::prelude::*;
|
||||
use sdset::{SetOperation, SetBuf};
|
||||
use slice_group_by::StrGroupBy;
|
||||
use structopt::StructOpt;
|
||||
use zerocopy::{LayoutVerified, AsBytes};
|
||||
|
||||
// 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]>;
|
||||
pub type BEU32 = heed::zerocopy::U32<heed::byteorder::BE>;
|
||||
pub type DocumentId = u32;
|
||||
|
||||
#[cfg(target_os = "linux")]
|
||||
#[global_allocator]
|
||||
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
|
||||
|
||||
static ID_GENERATOR: AtomicUsize = AtomicUsize::new(0);
|
||||
static ID_GENERATOR: AtomicUsize = AtomicUsize::new(0); // AtomicU32 ?
|
||||
|
||||
#[derive(Debug, StructOpt)]
|
||||
#[structopt(name = "mm-indexer", about = "The server side of the daugt project.")]
|
||||
@ -45,73 +40,24 @@ struct Opt {
|
||||
files_to_index: Vec<PathBuf>,
|
||||
}
|
||||
|
||||
fn bytes_to_u32s(bytes: &[u8]) -> Vec<u32> {
|
||||
fn aligned_to(bytes: &[u8], align: usize) -> bool {
|
||||
(bytes as *const _ as *const () as usize) % align == 0
|
||||
}
|
||||
|
||||
match LayoutVerified::new_slice(bytes) {
|
||||
Some(slice) => slice.to_vec(),
|
||||
None => {
|
||||
let len = bytes.len();
|
||||
|
||||
// ensure that it is the alignment that is wrong and the length is valid
|
||||
assert!(len % 4 == 0, "length is {} and is not modulo 4", len);
|
||||
assert!(!aligned_to(bytes, std::mem::align_of::<u32>()), "bytes are already aligned");
|
||||
|
||||
let elems = len / 4;
|
||||
let mut vec = Vec::<u32>::with_capacity(elems);
|
||||
|
||||
unsafe {
|
||||
let dst = vec.as_mut_ptr() as *mut u8;
|
||||
std::ptr::copy_nonoverlapping(bytes.as_ptr(), dst, len);
|
||||
vec.set_len(elems);
|
||||
}
|
||||
|
||||
vec
|
||||
fn union_postings_ids(_key: &[u8], old_value: Option<&[u8]>, new_value: RoaringBitmap) -> Option<Vec<u8>> {
|
||||
let result = match old_value {
|
||||
Some(bytes) => {
|
||||
let mut old_value = RoaringBitmap::deserialize_from(bytes).unwrap();
|
||||
old_value.union_with(&new_value);
|
||||
old_value
|
||||
},
|
||||
}
|
||||
None => new_value,
|
||||
};
|
||||
|
||||
let mut vec = Vec::new();
|
||||
result.serialize_into(&mut vec).unwrap();
|
||||
Some(vec)
|
||||
}
|
||||
|
||||
fn union_postings_ids(
|
||||
_key: &[u8],
|
||||
old_value: Option<&[u8]>,
|
||||
operands: &mut rocksdb::MergeOperands,
|
||||
) -> Option<Vec<u8>>
|
||||
{
|
||||
let mut sets_bufs = Vec::new();
|
||||
|
||||
if let Some(old_value) = old_value {
|
||||
let old_value = bytes_to_u32s(old_value);
|
||||
sets_bufs.push(SetBuf::new_unchecked(old_value.to_vec()));
|
||||
}
|
||||
|
||||
for operand in operands {
|
||||
let new_value = bytes_to_u32s(operand);
|
||||
sets_bufs.push(SetBuf::new_unchecked(new_value.to_vec()));
|
||||
}
|
||||
|
||||
let sets = sets_bufs.iter().map(|s| s.as_set()).collect();
|
||||
let result: SetBuf<u32> = sdset::multi::Union::new(sets).into_set_buf();
|
||||
|
||||
assert!(result.as_bytes().len() % 4 == 0);
|
||||
|
||||
Some(result.as_bytes().to_vec())
|
||||
}
|
||||
|
||||
fn union_words_fst(
|
||||
key: &[u8],
|
||||
old_value: Option<&[u8]>,
|
||||
operands: &mut rocksdb::MergeOperands,
|
||||
) -> Option<Vec<u8>>
|
||||
{
|
||||
fn union_words_fst(key: &[u8], old_value: Option<&[u8]>, new_value: &fst::Set<Vec<u8>>) -> Option<Vec<u8>> {
|
||||
if key != b"words-fst" { unimplemented!() }
|
||||
|
||||
let mut fst_operands = Vec::new();
|
||||
for operand in operands {
|
||||
fst_operands.push(fst::Set::new(operand).unwrap());
|
||||
}
|
||||
|
||||
// Do an union of the old and the new set of words.
|
||||
let mut builder = fst::set::OpBuilder::new();
|
||||
|
||||
@ -121,9 +67,7 @@ fn union_words_fst(
|
||||
builder.push(old_words);
|
||||
}
|
||||
|
||||
for new_words in &fst_operands {
|
||||
builder.push(new_words.into_stream());
|
||||
}
|
||||
builder.push(new_value);
|
||||
|
||||
let op = builder.r#union();
|
||||
let mut build = fst::SetBuilder::memory();
|
||||
@ -137,19 +81,94 @@ fn alphanumeric_tokens(string: &str) -> impl Iterator<Item = &str> {
|
||||
string.linear_group_by_key(|c| c.is_alphanumeric()).filter(is_alphanumeric)
|
||||
}
|
||||
|
||||
fn index_csv(
|
||||
tid: usize,
|
||||
db: Arc<rocksdb::DB>,
|
||||
mut rdr: csv::Reader<File>,
|
||||
) -> anyhow::Result<usize>
|
||||
{
|
||||
enum MainKey {
|
||||
WordsFst(fst::Set<Vec<u8>>),
|
||||
Headers(Vec<u8>),
|
||||
}
|
||||
|
||||
#[derive(Clone)]
|
||||
struct DbSender {
|
||||
main: Sender<MainKey>,
|
||||
postings_ids: Sender<(SmallString32, RoaringBitmap)>,
|
||||
documents: Sender<(DocumentId, Vec<u8>)>,
|
||||
}
|
||||
|
||||
struct DbReceiver {
|
||||
main: Receiver<MainKey>,
|
||||
postings_ids: Receiver<(SmallString32, RoaringBitmap)>,
|
||||
documents: Receiver<(DocumentId, Vec<u8>)>,
|
||||
}
|
||||
|
||||
fn thread_channel() -> (DbSender, DbReceiver) {
|
||||
let (sd_main, rc_main) = crossbeam_channel::bounded(4);
|
||||
let (sd_postings, rc_postings) = crossbeam_channel::bounded(10);
|
||||
let (sd_documents, rc_documents) = crossbeam_channel::bounded(10);
|
||||
|
||||
let sender = DbSender { main: sd_main, postings_ids: sd_postings, documents: sd_documents };
|
||||
let receiver = DbReceiver { main: rc_main, postings_ids: rc_postings, documents: rc_documents };
|
||||
|
||||
(sender, receiver)
|
||||
}
|
||||
|
||||
fn writer_thread(env: heed::Env, receiver: DbReceiver) -> anyhow::Result<()> {
|
||||
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 mut wtxn = env.write_txn()?;
|
||||
|
||||
loop {
|
||||
select! {
|
||||
recv(receiver.main) -> msg => {
|
||||
let msg = match msg {
|
||||
Err(_) => break,
|
||||
Ok(msg) => msg,
|
||||
};
|
||||
|
||||
match msg {
|
||||
MainKey::WordsFst(new_fst) => {
|
||||
let old_value = main.get::<_, Str, ByteSlice>(&wtxn, "words-fst")?;
|
||||
let new_value = union_words_fst(b"words-fst", old_value, &new_fst)
|
||||
.context("error while do a words-fst union")?;
|
||||
main.put::<_, Str, ByteSlice>(&mut wtxn, "words-fst", &new_value)?;
|
||||
},
|
||||
MainKey::Headers(headers) => {
|
||||
if let Some(old_headers) = main.get::<_, Str, ByteSlice>(&wtxn, "headers")? {
|
||||
ensure!(old_headers == &*headers, "headers differs from the previous ones");
|
||||
}
|
||||
main.put::<_, Str, ByteSlice>(&mut wtxn, "headers", &headers)?;
|
||||
},
|
||||
}
|
||||
},
|
||||
recv(receiver.postings_ids) -> msg => {
|
||||
let (word, postings) = match msg {
|
||||
Err(_) => break,
|
||||
Ok(msg) => msg,
|
||||
};
|
||||
|
||||
let old_value = postings_ids.get(&wtxn, &word)?;
|
||||
let new_value = union_postings_ids(word.as_bytes(), old_value, postings)
|
||||
.context("error while do a words-fst union")?;
|
||||
postings_ids.put(&mut wtxn, &word, &new_value)?;
|
||||
},
|
||||
recv(receiver.documents) -> msg => {
|
||||
let (id, content) = match msg {
|
||||
Err(_) => break,
|
||||
Ok(msg) => msg,
|
||||
};
|
||||
documents.put(&mut wtxn, &BEU32::new(id), &content)?;
|
||||
},
|
||||
}
|
||||
}
|
||||
|
||||
wtxn.commit()?;
|
||||
Ok(())
|
||||
}
|
||||
|
||||
fn index_csv(tid: usize, db_sender: DbSender, 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.cf_handle("main").context("cf \"main\" not found")?;
|
||||
let postings_ids = db.cf_handle("postings-ids").context("cf \"postings-ids\" not found")?;
|
||||
let documents = db.cf_handle("documents").context("cf \"documents\" not found")?;
|
||||
|
||||
let mut document = csv::StringRecord::new();
|
||||
let mut new_postings_ids = FastMap4::default();
|
||||
let mut new_words = BTreeSet::default();
|
||||
@ -160,19 +179,19 @@ fn index_csv(
|
||||
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) = db.get_cf(&main, "headers")? {
|
||||
ensure!(old_headers == headers, "headers differs from the previous ones");
|
||||
}
|
||||
db.put_cf(&main, "headers", headers.as_slice())?;
|
||||
db_sender.main.send(MainKey::Headers(headers))?;
|
||||
|
||||
while rdr.read_record(&mut document)? {
|
||||
let document_id = ID_GENERATOR.fetch_add(1, Ordering::SeqCst);
|
||||
let document_id = u32::try_from(document_id).context("Generated id is too big")?;
|
||||
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 alphanumeric_tokens(&content).enumerate().take(MAX_POSITION) {
|
||||
new_postings_ids.entry(SmallString32::from(word)).or_insert_with(BpVec::new).push(document_id);
|
||||
if !word.is_empty() && word.len() < 500 { // LMDB limits
|
||||
new_postings_ids.entry(SmallString32::from(word))
|
||||
.or_insert_with(RoaringBitmap::new)
|
||||
.insert(document_id);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
@ -180,15 +199,11 @@ fn index_csv(
|
||||
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()?;
|
||||
db.put_cf(&documents, document_id.to_be_bytes(), document)?;
|
||||
db_sender.documents.send((document_id, document))?;
|
||||
|
||||
number_of_documents += 1;
|
||||
if number_of_documents % 100000 == 0 {
|
||||
let postings_ids_size = new_postings_ids.iter().map(|(_, v)| {
|
||||
v.compressed_capacity() + v.uncompressed_capacity() * 4
|
||||
}).sum::<usize>();
|
||||
eprintln!("{}, documents seen {}, postings size {}",
|
||||
tid, number_of_documents, postings_ids_size);
|
||||
eprintln!("{}, documents seen {}", tid, number_of_documents);
|
||||
}
|
||||
}
|
||||
|
||||
@ -196,8 +211,7 @@ fn index_csv(
|
||||
|
||||
// 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());
|
||||
db.merge_cf(&postings_ids, word.as_bytes(), new_ids.as_bytes())?;
|
||||
db_sender.postings_ids.send((word.clone(), new_ids))?;
|
||||
new_words.insert(word);
|
||||
}
|
||||
|
||||
@ -207,7 +221,7 @@ fn index_csv(
|
||||
|
||||
let new_words_fst = fst::Set::from_iter(new_words.iter().map(|s| s.as_str()))?;
|
||||
drop(new_words);
|
||||
db.merge_cf(&main, "words-fst", new_words_fst.as_fst().as_bytes())?;
|
||||
db_sender.main.send(MainKey::WordsFst(new_words_fst))?;
|
||||
|
||||
eprintln!("Finished merging the words-fst");
|
||||
eprintln!("Total number of documents seen is {}", ID_GENERATOR.load(Ordering::Relaxed));
|
||||
@ -218,31 +232,30 @@ fn index_csv(
|
||||
fn main() -> anyhow::Result<()> {
|
||||
let opt = Opt::from_args();
|
||||
|
||||
let mut opts = rocksdb::Options::default();
|
||||
opts.create_if_missing(true);
|
||||
opts.create_missing_column_families(true);
|
||||
// Setup the merge operators
|
||||
opts.set_merge_operator("main", union_words_fst, None); // Some(union_words_fst));
|
||||
opts.set_merge_operator("postings-ids", union_postings_ids, None); // Some(union_postings_ids));
|
||||
std::fs::create_dir_all(&opt.database)?;
|
||||
let env = EnvOpenOptions::new()
|
||||
.map_size(100 * 1024 * 1024 * 1024) // 100 GB
|
||||
.max_readers(10)
|
||||
.max_dbs(5)
|
||||
.open(opt.database)?;
|
||||
|
||||
let mut db = rocksdb::DB::open(&opts, &opt.database)?;
|
||||
let (sender, receiver) = thread_channel();
|
||||
let writing_child = thread::spawn(move || writer_thread(env, receiver));
|
||||
|
||||
let cfs = &["main", "postings-ids", "documents"];
|
||||
for cf in cfs.into_iter() {
|
||||
db.create_cf(cf, &opts).unwrap();
|
||||
}
|
||||
|
||||
let db = Arc::new(db);
|
||||
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)
|
||||
index_csv(tid, sender.clone(), rdr)
|
||||
})
|
||||
.try_reduce(|| 0, |a, b| Ok(a + b));
|
||||
|
||||
println!("{:?}", res);
|
||||
|
||||
eprintln!("witing the writing thread...");
|
||||
writing_child.join().unwrap().unwrap();
|
||||
|
||||
println!("indexed {:?} documents", res);
|
||||
|
||||
Ok(())
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user