meilisearch/src/bin/indexer.rs

652 lines
24 KiB
Rust
Raw Normal View History

use std::convert::{TryFrom, TryInto};
use std::fs::File;
2020-08-06 18:38:42 +08:00
use std::io::{self, Read, Write};
use std::iter::FromIterator;
2020-07-07 17:32:33 +08:00
use std::path::PathBuf;
2020-08-06 18:38:42 +08:00
use std::thread;
2020-07-02 04:45:43 +08:00
use std::time::Instant;
2020-07-07 17:32:33 +08:00
use anyhow::Context;
use arc_cache::ArcCache;
use bstr::ByteSlice as _;
2020-07-07 17:32:33 +08:00
use cow_utils::CowUtils;
use flate2::read::GzDecoder;
use fst::IntoStreamer;
2020-07-07 17:32:33 +08:00
use heed::EnvOpenOptions;
2020-05-30 21:35:33 +08:00
use heed::types::*;
use log::{debug, info};
use memmap::Mmap;
use oxidized_mtbl::{Reader, Writer, Merger, Sorter, CompressionType};
use rayon::prelude::*;
2020-06-30 01:48:02 +08:00
use roaring::RoaringBitmap;
2020-05-26 02:39:53 +08:00
use structopt::StructOpt;
2020-08-31 03:50:30 +08:00
use milli::{SmallVec32, Index, DocumentId, Position, Attribute, BEU32};
use milli::tokenizer::{simple_tokenizer, only_words};
2020-07-07 17:32:33 +08:00
const LMDB_MAX_KEY_LENGTH: usize = 511;
2020-07-07 17:32:33 +08:00
const ONE_MILLION: usize = 1_000_000;
const MAX_POSITION: usize = 1000;
const MAX_ATTRIBUTES: usize = u32::max_value() as usize / MAX_POSITION;
2020-05-26 02:39:53 +08:00
const HEADERS_KEY: &[u8] = b"\0headers";
const DOCUMENTS_IDS_KEY: &[u8] = b"\x04documents-ids";
const WORDS_FST_KEY: &[u8] = b"\x06words-fst";
const DOCUMENTS_IDS_BYTE: u8 = 4;
const WORD_ATTRIBUTE_DOCIDS_BYTE: u8 = 3;
const WORD_FOUR_POSITIONS_DOCIDS_BYTE: u8 = 5;
const WORD_POSITION_DOCIDS_BYTE: u8 = 2;
const WORD_POSITIONS_BYTE: u8 = 1;
2020-05-26 02:39:53 +08:00
#[cfg(target_os = "linux")]
#[global_allocator]
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
#[derive(Debug, StructOpt)]
#[structopt(name = "milli-indexer")]
/// The indexer binary of the milli project.
2020-05-26 02:39:53 +08:00
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,
/// The maximum size the database can take on disk. It is recommended to specify
/// the whole disk space (value must be a multiple of a page size).
#[structopt(long = "db-size", default_value = "107374182400")] // 100 GB
database_size: usize,
/// Number of parallel jobs, defaults to # of CPUs.
#[structopt(short, long)]
jobs: Option<usize>,
#[structopt(flatten)]
indexer: IndexerOpt,
/// Verbose mode (-v, -vv, -vvv, etc.)
#[structopt(short, long, parse(from_occurrences))]
verbose: usize,
/// CSV file to index, if unspecified the CSV is read from standard input.
///
/// You can also provide a ".gz" or ".gzip" CSV file, the indexer will figure out
/// how to decode and read it.
///
/// Note that it is much faster to index from a file as when the indexer reads from stdin
/// it will dedicate a thread for that and context switches could slow down the indexing jobs.
csv_file: Option<PathBuf>,
}
#[derive(Debug, StructOpt)]
struct IndexerOpt {
/// MTBL max number of chunks in bytes.
#[structopt(long)]
max_nb_chunks: Option<usize>,
/// MTBL max memory in bytes.
#[structopt(long)]
max_memory: Option<usize>,
/// Size of the ARC cache when indexing.
#[structopt(long, default_value = "43690")]
arc_cache_size: usize,
/// The name of the compression algorithm to use when compressing intermediate
/// chunks during indexing documents.
///
/// Choosing a fast algorithm will make the indexing faster but may consume more memory.
#[structopt(long, default_value = "snappy", possible_values = &["snappy", "zlib", "lz4", "lz4hc", "zstd"])]
chunk_compression_type: String,
/// The level of compression of the chosen algorithm.
#[structopt(long, requires = "chunk-compression-type")]
chunk_compression_level: Option<u32>,
}
fn compression_type_from_str(name: &str) -> CompressionType {
match name {
"snappy" => CompressionType::Snappy,
"zlib" => CompressionType::Zlib,
"lz4" => CompressionType::Lz4,
"lz4hc" => CompressionType::Lz4hc,
"zstd" => CompressionType::Zstd,
_ => panic!("invalid compression algorithm"),
}
2020-05-30 21:35:33 +08:00
}
fn lmdb_key_valid_size(key: &[u8]) -> bool {
!key.is_empty() && key.len() <= LMDB_MAX_KEY_LENGTH
}
type MergeFn = fn(&[u8], &[Vec<u8>]) -> Result<Vec<u8>, ()>;
struct Store {
word_positions: ArcCache<SmallVec32<u8>, RoaringBitmap>,
word_position_docids: ArcCache<(SmallVec32<u8>, Position), RoaringBitmap>,
word_four_positions_docids: ArcCache<(SmallVec32<u8>, Position), RoaringBitmap>,
word_attribute_docids: ArcCache<(SmallVec32<u8>, Attribute), RoaringBitmap>,
documents_ids: RoaringBitmap,
sorter: Sorter<MergeFn>,
documents_sorter: Sorter<MergeFn>,
2020-07-07 17:32:33 +08:00
}
impl Store {
fn new(
arc_cache_size: usize,
max_nb_chunks: Option<usize>,
max_memory: Option<usize>,
chunk_compression_type: CompressionType,
chunk_compression_level: Option<u32>,
) -> Store
{
let mut builder = Sorter::builder(merge as MergeFn);
builder.chunk_compression_type(chunk_compression_type);
if let Some(level) = chunk_compression_level {
builder.chunk_compression_level(level);
}
if let Some(nb_chunks) = max_nb_chunks {
builder.max_nb_chunks(nb_chunks);
}
if let Some(memory) = max_memory {
builder.max_memory(memory);
}
let mut documents_builder = Sorter::builder(docs_merge as MergeFn);
documents_builder.chunk_compression_type(chunk_compression_type);
if let Some(level) = chunk_compression_level {
builder.chunk_compression_level(level);
}
Store {
word_positions: ArcCache::new(arc_cache_size),
word_position_docids: ArcCache::new(arc_cache_size),
word_four_positions_docids: ArcCache::new(arc_cache_size),
word_attribute_docids: ArcCache::new(arc_cache_size),
documents_ids: RoaringBitmap::new(),
sorter: builder.build(),
documents_sorter: documents_builder.build(),
}
}
// Save the documents ids under the position and word we have seen it.
pub fn insert_word_position_docid(&mut self, word: &str, position: Position, id: DocumentId) -> anyhow::Result<()> {
let word_vec = SmallVec32::from(word.as_bytes());
let ids = RoaringBitmap::from_iter(Some(id));
let (_, lrus) = self.word_position_docids.insert((word_vec, position), ids, |old, new| old.union_with(&new));
Self::write_word_position_docids(&mut self.sorter, lrus)?;
self.insert_word_position(word, position)?;
self.insert_word_four_positions_docid(word, position, id)?;
self.insert_word_attribute_docid(word, position / MAX_POSITION as u32, id)?;
Ok(())
}
pub fn insert_word_four_positions_docid(&mut self, word: &str, position: Position, id: DocumentId) -> anyhow::Result<()> {
let position = position - position % 4;
let word_vec = SmallVec32::from(word.as_bytes());
let ids = RoaringBitmap::from_iter(Some(id));
let (_, lrus) = self.word_four_positions_docids.insert((word_vec, position), ids, |old, new| old.union_with(&new));
Self::write_word_four_positions_docids(&mut self.sorter, lrus)
}
// Save the positions where this word has been seen.
pub fn insert_word_position(&mut self, word: &str, position: Position) -> anyhow::Result<()> {
let word = SmallVec32::from(word.as_bytes());
let position = RoaringBitmap::from_iter(Some(position));
let (_, lrus) = self.word_positions.insert(word, position, |old, new| old.union_with(&new));
Self::write_word_positions(&mut self.sorter, lrus)
}
// Save the documents ids under the attribute and word we have seen it.
fn insert_word_attribute_docid(&mut self, word: &str, attribute: Attribute, id: DocumentId) -> anyhow::Result<()> {
let word = SmallVec32::from(word.as_bytes());
let ids = RoaringBitmap::from_iter(Some(id));
let (_, lrus) = self.word_attribute_docids.insert((word, attribute), ids, |old, new| old.union_with(&new));
Self::write_word_attribute_docids(&mut self.sorter, lrus)
}
pub fn write_headers(&mut self, headers: &[u8]) -> anyhow::Result<()> {
Ok(self.sorter.insert(HEADERS_KEY, headers)?)
}
2020-07-07 17:32:33 +08:00
pub fn write_document(&mut self, id: DocumentId, content: &[u8]) -> anyhow::Result<()> {
self.documents_ids.insert(id);
Ok(self.documents_sorter.insert(id.to_be_bytes(), content)?)
}
2020-07-07 17:32:33 +08:00
fn write_word_positions<I>(sorter: &mut Sorter<MergeFn>, iter: I) -> anyhow::Result<()>
where I: IntoIterator<Item=(SmallVec32<u8>, RoaringBitmap)>
{
// postings ids keys are all prefixed
let mut key = vec![WORD_POSITIONS_BYTE];
2020-07-07 17:32:33 +08:00
let mut buffer = Vec::new();
for (word, positions) in iter {
2020-07-07 17:32:33 +08:00
key.truncate(1);
key.extend_from_slice(&word);
// We serialize the positions into a buffer
buffer.clear();
buffer.reserve(positions.serialized_size());
positions.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
if lmdb_key_valid_size(&key) {
sorter.insert(&key, &buffer)?;
}
2020-07-07 17:32:33 +08:00
}
Ok(())
}
fn write_word_position_docids<I>(sorter: &mut Sorter<MergeFn>, iter: I) -> anyhow::Result<()>
where I: IntoIterator<Item=((SmallVec32<u8>, Position), RoaringBitmap)>
{
// postings positions ids keys are all prefixed
let mut key = vec![WORD_POSITION_DOCIDS_BYTE];
let mut buffer = Vec::new();
for ((word, pos), ids) in iter {
2020-07-07 17:32:33 +08:00
key.truncate(1);
key.extend_from_slice(&word);
// we postfix the word by the positions it appears in
2020-08-06 16:20:26 +08:00
key.extend_from_slice(&pos.to_be_bytes());
// We serialize the document ids into a buffer
buffer.clear();
buffer.reserve(ids.serialized_size());
ids.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
if lmdb_key_valid_size(&key) {
sorter.insert(&key, &buffer)?;
}
}
2020-07-07 17:32:33 +08:00
Ok(())
}
fn write_word_four_positions_docids<I>(sorter: &mut Sorter<MergeFn>, iter: I) -> anyhow::Result<()>
where I: IntoIterator<Item=((SmallVec32<u8>, Position), RoaringBitmap)>
{
// postings positions ids keys are all prefixed
let mut key = vec![WORD_FOUR_POSITIONS_DOCIDS_BYTE];
let mut buffer = Vec::new();
for ((word, pos), ids) in iter {
key.truncate(1);
key.extend_from_slice(&word);
// we postfix the word by the positions it appears in
key.extend_from_slice(&pos.to_be_bytes());
// We serialize the document ids into a buffer
buffer.clear();
buffer.reserve(ids.serialized_size());
ids.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
if lmdb_key_valid_size(&key) {
sorter.insert(&key, &buffer)?;
}
}
Ok(())
}
fn write_word_attribute_docids<I>(sorter: &mut Sorter<MergeFn>, iter: I) -> anyhow::Result<()>
where I: IntoIterator<Item=((SmallVec32<u8>, Attribute), RoaringBitmap)>
{
// postings attributes keys are all prefixed
let mut key = vec![WORD_ATTRIBUTE_DOCIDS_BYTE];
let mut buffer = Vec::new();
for ((word, attr), ids) in iter {
key.truncate(1);
key.extend_from_slice(&word);
// we postfix the word by the positions it appears in
key.extend_from_slice(&attr.to_be_bytes());
// We serialize the document ids into a buffer
buffer.clear();
buffer.reserve(ids.serialized_size());
ids.serialize_into(&mut buffer)?;
// that we write under the generated key into MTBL
if lmdb_key_valid_size(&key) {
sorter.insert(&key, &buffer)?;
}
}
Ok(())
}
fn write_documents_ids(sorter: &mut Sorter<MergeFn>, ids: RoaringBitmap) -> anyhow::Result<()> {
let mut buffer = Vec::with_capacity(ids.serialized_size());
ids.serialize_into(&mut buffer)?;
sorter.insert(DOCUMENTS_IDS_KEY, &buffer)?;
Ok(())
}
2020-08-07 19:11:31 +08:00
pub fn finish(mut self) -> anyhow::Result<(Reader<Mmap>, Reader<Mmap>)> {
Self::write_word_positions(&mut self.sorter, self.word_positions)?;
Self::write_word_position_docids(&mut self.sorter, self.word_position_docids)?;
Self::write_word_four_positions_docids(&mut self.sorter, self.word_four_positions_docids)?;
Self::write_word_attribute_docids(&mut self.sorter, self.word_attribute_docids)?;
Self::write_documents_ids(&mut self.sorter, self.documents_ids)?;
2020-07-07 17:32:33 +08:00
let mut wtr = tempfile::tempfile().map(Writer::new)?;
let mut builder = fst::SetBuilder::memory();
2020-07-07 17:32:33 +08:00
let mut iter = self.sorter.into_iter()?;
while let Some(result) = iter.next() {
let (key, val) = result?;
if let Some((&1, word)) = key.split_first() {
// This is a lexicographically ordered word position
// we use the key to construct the words fst.
builder.insert(word)?;
}
wtr.insert(key, val)?;
2020-07-07 17:32:33 +08:00
}
let fst = builder.into_set();
wtr.insert(WORDS_FST_KEY, fst.as_fst().as_bytes())?;
2020-07-07 17:32:33 +08:00
let mut docs_wtr = tempfile::tempfile().map(Writer::new)?;
self.documents_sorter.write_into(&mut docs_wtr)?;
let docs_file = docs_wtr.into_inner()?;
let docs_mmap = unsafe { Mmap::map(&docs_file)? };
2020-08-07 19:11:31 +08:00
let docs_reader = Reader::new(docs_mmap)?;
let file = wtr.into_inner()?;
let mmap = unsafe { Mmap::map(&file)? };
let reader = Reader::new(mmap)?;
2020-08-07 19:11:31 +08:00
Ok((reader, docs_reader))
}
}
2020-07-07 17:32:33 +08:00
fn docs_merge(key: &[u8], values: &[Vec<u8>]) -> Result<Vec<u8>, ()> {
let key = key.try_into().unwrap();
let id = u32::from_be_bytes(key);
panic!("documents must not conflict ({} with {} values)!", id, values.len())
}
fn merge(key: &[u8], values: &[Vec<u8>]) -> Result<Vec<u8>, ()> {
match key {
WORDS_FST_KEY => {
let fsts: Vec<_> = values.iter().map(|v| fst::Set::new(v).unwrap()).collect();
2020-08-31 03:50:30 +08:00
// Union of the FSTs
let mut op = fst::set::OpBuilder::new();
fsts.iter().for_each(|fst| op.push(fst.into_stream()));
let op = op.r#union();
let mut build = fst::SetBuilder::memory();
build.extend_stream(op.into_stream()).unwrap();
Ok(build.into_inner().unwrap())
},
HEADERS_KEY => {
assert!(values.windows(2).all(|vs| vs[0] == vs[1]));
Ok(values[0].to_vec())
},
key => match key[0] {
DOCUMENTS_IDS_BYTE
| WORD_POSITIONS_BYTE
| WORD_POSITION_DOCIDS_BYTE
| WORD_FOUR_POSITIONS_DOCIDS_BYTE
| WORD_ATTRIBUTE_DOCIDS_BYTE =>
{
2020-08-31 03:50:30 +08:00
let (head, tail) = values.split_first().unwrap();
2020-08-31 03:50:30 +08:00
let mut head = RoaringBitmap::deserialize_from(head.as_slice()).unwrap();
for value in tail {
let bitmap = RoaringBitmap::deserialize_from(value.as_slice()).unwrap();
2020-08-31 03:50:30 +08:00
head.union_with(&bitmap);
}
2020-07-07 17:32:33 +08:00
2020-08-31 03:50:30 +08:00
let mut vec = Vec::with_capacity(head.serialized_size());
head.serialize_into(&mut vec).unwrap();
Ok(vec)
},
otherwise => panic!("wut {:?}", otherwise),
2020-07-07 17:32:33 +08:00
}
}
}
// TODO merge with the previous values
// TODO store the documents in a compressed MTBL
fn lmdb_writer(wtxn: &mut heed::RwTxn, index: &Index, key: &[u8], val: &[u8]) -> anyhow::Result<()> {
if key == WORDS_FST_KEY {
// Write the words fst
index.main.put::<_, Str, ByteSlice>(wtxn, "words-fst", val)?;
}
else if key == HEADERS_KEY {
// Write the headers
index.main.put::<_, Str, ByteSlice>(wtxn, "headers", val)?;
}
else if key == DOCUMENTS_IDS_KEY {
// Write the documents ids list
index.main.put::<_, Str, ByteSlice>(wtxn, "documents-ids", val)?;
}
else if key.starts_with(&[WORD_POSITIONS_BYTE]) {
// Write the postings lists
index.word_positions.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[WORD_POSITION_DOCIDS_BYTE]) {
// Write the postings lists
index.word_position_docids.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[WORD_FOUR_POSITIONS_DOCIDS_BYTE]) {
// Write the postings lists
index.word_four_positions_docids.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
else if key.starts_with(&[WORD_ATTRIBUTE_DOCIDS_BYTE]) {
// Write the attribute postings lists
index.word_attribute_docids.as_polymorph()
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
}
Ok(())
}
fn merge_into_lmdb<F>(sources: Vec<Reader<Mmap>>, mut f: F) -> anyhow::Result<()>
where F: FnMut(&[u8], &[u8]) -> anyhow::Result<()>
{
debug!("Merging {} MTBL stores...", sources.len());
let before = Instant::now();
let mut builder = Merger::builder(merge);
builder.extend(sources);
let merger = builder.build();
let mut iter = merger.into_merge_iter()?;
while let Some(result) = iter.next() {
let (k, v) = result?;
(f)(&k, &v).with_context(|| format!("writing {:?} {:?} into LMDB", k.as_bstr(), k.as_bstr()))?;
}
debug!("MTBL stores merged in {:.02?}!", before.elapsed());
Ok(())
}
2020-07-07 17:32:33 +08:00
fn index_csv(
2020-08-06 18:38:42 +08:00
mut rdr: csv::Reader<Box<dyn Read + Send>>,
thread_index: usize,
2020-07-07 17:32:33 +08:00
num_threads: usize,
arc_cache_size: usize,
max_nb_chunks: Option<usize>,
max_memory: Option<usize>,
chunk_compression_type: CompressionType,
chunk_compression_level: Option<u32>,
2020-08-07 19:11:31 +08:00
) -> anyhow::Result<(Reader<Mmap>, Reader<Mmap>)>
{
debug!("{:?}: Indexing into a Store...", thread_index);
2020-05-26 02:39:53 +08:00
let mut store = Store::new(
arc_cache_size,
max_nb_chunks,
max_memory,
chunk_compression_type,
chunk_compression_level,
);
// Write the headers into a Vec of bytes and then into the store.
2020-05-26 02:39:53 +08:00
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()?;
store.write_headers(&headers)?;
2020-05-26 02:39:53 +08:00
let mut before = Instant::now();
2020-07-07 17:32:33 +08:00
let mut document_id: usize = 0;
let mut document = csv::StringRecord::new();
2020-05-26 02:39:53 +08:00
while rdr.read_record(&mut document)? {
2020-07-07 17:32:33 +08:00
// We skip documents that must not be indexed by this thread.
if document_id % num_threads == thread_index {
if document_id % ONE_MILLION == 0 {
debug!("We have seen {}m documents so far ({:.02?}).",
document_id / ONE_MILLION, before.elapsed());
before = Instant::now();
}
2020-07-02 04:45:43 +08:00
let document_id = DocumentId::try_from(document_id).context("generated id is too big")?;
for (attr, content) in document.iter().enumerate().take(MAX_ATTRIBUTES) {
2020-08-31 03:50:30 +08:00
for (pos, (_, token)) in simple_tokenizer(&content).filter(only_words).enumerate().take(MAX_POSITION) {
let word = token.cow_to_lowercase();
let position = (attr * MAX_POSITION + pos) as u32;
store.insert_word_position_docid(&word, position, document_id)?;
}
2020-05-26 02:39:53 +08:00
}
// 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()?;
store.write_document(document_id, &document)?;
2020-05-26 02:39:53 +08:00
}
// Compute the document id of the the next document.
document_id = document_id + 1;
2020-07-07 17:32:33 +08:00
}
2020-08-07 19:11:31 +08:00
let (reader, docs_reader) = store.finish()?;
debug!("{:?}: Store created!", thread_index);
2020-08-07 19:11:31 +08:00
Ok((reader, docs_reader))
2020-05-26 02:39:53 +08:00
}
2020-07-02 04:45:43 +08:00
fn main() -> anyhow::Result<()> {
let opt = Opt::from_args();
2020-06-29 19:54:47 +08:00
stderrlog::new()
.verbosity(opt.verbose)
.show_level(false)
.timestamp(stderrlog::Timestamp::Off)
.init()?;
if let Some(jobs) = opt.jobs {
rayon::ThreadPoolBuilder::new().num_threads(jobs).build_global()?;
}
2020-07-04 18:34:10 +08:00
std::fs::create_dir_all(&opt.database)?;
2020-07-07 17:32:33 +08:00
let env = EnvOpenOptions::new()
.map_size(opt.database_size)
2020-07-07 17:32:33 +08:00
.max_dbs(10)
2020-08-07 19:11:31 +08:00
.open(&opt.database)?;
2020-06-29 19:54:47 +08:00
let before_indexing = Instant::now();
let index = Index::new(&env)?;
2020-06-29 19:54:47 +08:00
2020-07-07 17:32:33 +08:00
let num_threads = rayon::current_num_threads();
let arc_cache_size = opt.indexer.arc_cache_size;
let max_nb_chunks = opt.indexer.max_nb_chunks;
let max_memory = opt.indexer.max_memory;
let chunk_compression_type = compression_type_from_str(&opt.indexer.chunk_compression_type);
let chunk_compression_level = opt.indexer.chunk_compression_level;
2020-08-06 18:38:42 +08:00
let csv_readers = match opt.csv_file {
Some(file_path) => {
// We open the file # jobs times.
(0..num_threads)
.map(|_| {
let file = File::open(&file_path)?;
// if the file extension is "gz" or "gzip" we can decode and read it.
let r = if file_path.extension().map_or(false, |ext| ext == "gz" || ext == "gzip") {
Box::new(GzDecoder::new(file)) as Box<dyn Read + Send>
} else {
Box::new(file) as Box<dyn Read + Send>
};
2020-08-06 18:38:42 +08:00
Ok(csv::Reader::from_reader(r)) as io::Result<_>
})
.collect::<Result<Vec<_>, _>>()?
},
None => {
let mut csv_readers = Vec::new();
let mut writers = Vec::new();
for (r, w) in (0..num_threads).map(|_| ringtail::io::pipe()) {
2020-08-06 18:38:42 +08:00
let r = Box::new(r) as Box<dyn Read + Send>;
csv_readers.push(csv::Reader::from_reader(r));
writers.push(w);
}
thread::spawn(move || {
let stdin = std::io::stdin();
let mut stdin = stdin.lock();
let mut buffer = [0u8; 4096];
loop {
match stdin.read(&mut buffer)? {
0 => return Ok(()) as io::Result<()>,
size => for w in &mut writers {
w.write_all(&buffer[..size])?;
}
}
}
});
csv_readers
},
};
2020-07-02 04:45:43 +08:00
2020-08-07 19:11:31 +08:00
let readers = csv_readers
.into_par_iter()
2020-07-07 17:32:33 +08:00
.enumerate()
.map(|(i, rdr)| {
index_csv(
rdr,
i,
num_threads,
arc_cache_size,
max_nb_chunks,
max_memory,
chunk_compression_type,
chunk_compression_level,
)
})
2020-08-07 19:11:31 +08:00
.collect::<Result<Vec<_>, _>>()?;
let mut stores = Vec::with_capacity(readers.len());
let mut docs_stores = Vec::with_capacity(readers.len());
readers.into_iter().for_each(|(s, d)| {
stores.push(s);
docs_stores.push(d);
});
2020-07-02 04:45:43 +08:00
let mut wtxn = env.write_txn()?;
// We merge the postings lists into LMDB.
debug!("We are writing the postings lists into LMDB on disk...");
merge_into_lmdb(stores, |k, v| lmdb_writer(&mut wtxn, &index, k, v))?;
// We merge the documents into LMDB.
debug!("We are writing the documents into LMDB on disk...");
merge_into_lmdb(docs_stores, |k, v| {
let id = k.try_into().map(u32::from_be_bytes)?;
Ok(index.documents.put(&mut wtxn, &BEU32::new(id), v)?)
})?;
// Retrieve the number of documents.
let count = index.number_of_documents(&wtxn)?;
wtxn.commit()?;
info!("Wrote {} documents in {:.02?}", count, before_indexing.elapsed());
2020-05-26 02:39:53 +08:00
Ok(())
}