2020-08-22 00:10:24 +08:00
|
|
|
use std::convert::{TryFrom, TryInto};
|
2020-08-28 21:38:05 +08:00
|
|
|
use std::fs::File;
|
2020-08-06 18:38:42 +08:00
|
|
|
use std::io::{self, Read, Write};
|
2020-08-04 21:19:21 +08:00
|
|
|
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;
|
2020-08-04 21:19:21 +08:00
|
|
|
use arc_cache::ArcCache;
|
2020-08-16 02:37:13 +08:00
|
|
|
use bstr::ByteSlice as _;
|
2020-07-07 17:32:33 +08:00
|
|
|
use cow_utils::CowUtils;
|
2020-08-22 00:10:24 +08:00
|
|
|
use flate2::read::GzDecoder;
|
2020-08-06 17:08:24 +08:00
|
|
|
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::*;
|
2020-08-28 21:38:05 +08:00
|
|
|
use log::{debug, info};
|
2020-08-04 21:19:21 +08:00
|
|
|
use memmap::Mmap;
|
|
|
|
use oxidized_mtbl::{Reader, Writer, Merger, Sorter, CompressionType};
|
2020-07-07 18:21:22 +08:00
|
|
|
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
|
|
|
|
2020-08-16 02:37:13 +08:00
|
|
|
const LMDB_MAX_KEY_LENGTH: usize = 511;
|
2020-07-07 17:32:33 +08:00
|
|
|
const ONE_MILLION: usize = 1_000_000;
|
2020-06-05 22:32:14 +08:00
|
|
|
|
|
|
|
const MAX_POSITION: usize = 1000;
|
|
|
|
const MAX_ATTRIBUTES: usize = u32::max_value() as usize / MAX_POSITION;
|
2020-05-26 02:39:53 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
const HEADERS_KEY: &[u8] = b"\0headers";
|
2020-08-29 21:14:04 +08:00
|
|
|
const DOCUMENTS_IDS_KEY: &[u8] = b"\x04documents-ids";
|
2020-08-29 22:29:18 +08:00
|
|
|
const WORDS_FST_KEY: &[u8] = b"\x06words-fst";
|
2020-08-29 21:14:04 +08:00
|
|
|
const DOCUMENTS_IDS_BYTE: u8 = 4;
|
2020-08-29 22:29:18 +08:00
|
|
|
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-08-04 21:19:21 +08:00
|
|
|
|
2020-05-26 02:39:53 +08:00
|
|
|
#[cfg(target_os = "linux")]
|
|
|
|
#[global_allocator]
|
|
|
|
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
|
|
|
|
|
|
|
|
#[derive(Debug, StructOpt)]
|
2020-08-21 22:41:26 +08:00
|
|
|
#[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,
|
|
|
|
|
2020-08-10 19:53:53 +08:00
|
|
|
/// 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,
|
|
|
|
|
2020-06-28 18:13:10 +08:00
|
|
|
/// Number of parallel jobs, defaults to # of CPUs.
|
|
|
|
#[structopt(short, long)]
|
|
|
|
jobs: Option<usize>,
|
|
|
|
|
2020-08-21 22:41:26 +08:00
|
|
|
#[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.
|
2020-08-21 22:08:32 +08:00
|
|
|
///
|
2020-08-22 00:10:24 +08:00
|
|
|
/// You can also provide a ".gz" or ".gzip" CSV file, the indexer will figure out
|
|
|
|
/// how to decode and read it.
|
|
|
|
///
|
2020-08-21 22:08:32 +08:00
|
|
|
/// 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.
|
2020-08-21 22:41:26 +08:00
|
|
|
csv_file: Option<PathBuf>,
|
|
|
|
}
|
|
|
|
|
|
|
|
#[derive(Debug, StructOpt)]
|
|
|
|
struct IndexerOpt {
|
2020-08-05 18:10:41 +08:00
|
|
|
/// 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>,
|
|
|
|
|
2020-08-06 17:08:24 +08:00
|
|
|
/// Size of the ARC cache when indexing.
|
2020-08-30 18:02:06 +08:00
|
|
|
#[structopt(long, default_value = "43690")]
|
2020-08-29 22:29:18 +08:00
|
|
|
arc_cache_size: usize,
|
2020-08-06 17:08:24 +08:00
|
|
|
|
2020-08-21 22:41:26 +08:00
|
|
|
/// 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,
|
2020-07-12 17:04:35 +08:00
|
|
|
|
2020-08-21 22:41:26 +08:00
|
|
|
/// 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
|
|
|
}
|
|
|
|
|
2020-08-05 18:10:41 +08:00
|
|
|
fn lmdb_key_valid_size(key: &[u8]) -> bool {
|
|
|
|
!key.is_empty() && key.len() <= LMDB_MAX_KEY_LENGTH
|
|
|
|
}
|
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
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>,
|
2020-08-29 22:29:18 +08:00
|
|
|
word_four_positions_docids: ArcCache<(SmallVec32<u8>, Position), RoaringBitmap>,
|
2020-08-06 17:08:24 +08:00
|
|
|
word_attribute_docids: ArcCache<(SmallVec32<u8>, Attribute), RoaringBitmap>,
|
2020-08-29 21:14:04 +08:00
|
|
|
documents_ids: RoaringBitmap,
|
2020-08-04 21:19:21 +08:00
|
|
|
sorter: Sorter<MergeFn>,
|
2020-08-07 00:19:10 +08:00
|
|
|
documents_sorter: Sorter<MergeFn>,
|
2020-07-07 17:32:33 +08:00
|
|
|
}
|
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
impl Store {
|
2020-08-21 22:41:26 +08:00
|
|
|
fn new(
|
2020-08-29 22:29:18 +08:00
|
|
|
arc_cache_size: usize,
|
2020-08-21 22:41:26 +08:00
|
|
|
max_nb_chunks: Option<usize>,
|
|
|
|
max_memory: Option<usize>,
|
|
|
|
chunk_compression_type: CompressionType,
|
|
|
|
chunk_compression_level: Option<u32>,
|
|
|
|
) -> Store
|
|
|
|
{
|
2020-08-05 18:10:41 +08:00
|
|
|
let mut builder = Sorter::builder(merge as MergeFn);
|
2020-08-21 22:41:26 +08:00
|
|
|
builder.chunk_compression_type(chunk_compression_type);
|
|
|
|
if let Some(level) = chunk_compression_level {
|
|
|
|
builder.chunk_compression_level(level);
|
|
|
|
}
|
2020-08-05 18:10:41 +08:00
|
|
|
if let Some(nb_chunks) = max_nb_chunks {
|
|
|
|
builder.max_nb_chunks(nb_chunks);
|
|
|
|
}
|
|
|
|
if let Some(memory) = max_memory {
|
|
|
|
builder.max_memory(memory);
|
|
|
|
}
|
2020-08-04 21:19:21 +08:00
|
|
|
|
2020-08-07 00:19:10 +08:00
|
|
|
let mut documents_builder = Sorter::builder(docs_merge as MergeFn);
|
2020-08-21 22:41:26 +08:00
|
|
|
documents_builder.chunk_compression_type(chunk_compression_type);
|
|
|
|
if let Some(level) = chunk_compression_level {
|
|
|
|
builder.chunk_compression_level(level);
|
|
|
|
}
|
2020-08-07 00:19:10 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
Store {
|
2020-08-06 17:08:24 +08:00
|
|
|
word_positions: ArcCache::new(arc_cache_size),
|
|
|
|
word_position_docids: ArcCache::new(arc_cache_size),
|
2020-08-29 22:29:18 +08:00
|
|
|
word_four_positions_docids: ArcCache::new(arc_cache_size),
|
2020-08-06 17:08:24 +08:00
|
|
|
word_attribute_docids: ArcCache::new(arc_cache_size),
|
2020-08-29 21:14:04 +08:00
|
|
|
documents_ids: RoaringBitmap::new(),
|
2020-08-05 18:10:41 +08:00
|
|
|
sorter: builder.build(),
|
2020-08-07 00:19:10 +08:00
|
|
|
documents_sorter: documents_builder.build(),
|
2020-08-04 21:19:21 +08:00
|
|
|
}
|
|
|
|
}
|
|
|
|
|
|
|
|
// 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<()> {
|
2020-08-06 17:08:24 +08:00
|
|
|
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)?;
|
2020-08-29 22:29:18 +08:00
|
|
|
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));
|
2020-08-30 18:02:06 +08:00
|
|
|
let (_, lrus) = self.word_four_positions_docids.insert((word_vec, position), ids, |old, new| old.union_with(&new));
|
2020-08-29 22:29:18 +08:00
|
|
|
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)
|
2020-08-06 17:08:24 +08:00
|
|
|
}
|
|
|
|
|
|
|
|
// 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<()> {
|
2020-08-04 21:19:21 +08:00
|
|
|
let word = SmallVec32::from(word.as_bytes());
|
|
|
|
let ids = RoaringBitmap::from_iter(Some(id));
|
2020-08-06 17:08:24 +08:00
|
|
|
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)
|
2020-07-07 22:48:49 +08:00
|
|
|
}
|
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
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
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
pub fn write_document(&mut self, id: DocumentId, content: &[u8]) -> anyhow::Result<()> {
|
2020-08-29 22:29:18 +08:00
|
|
|
self.documents_ids.insert(id);
|
2020-08-07 00:19:10 +08:00
|
|
|
Ok(self.documents_sorter.insert(id.to_be_bytes(), content)?)
|
2020-08-04 21:19:21 +08:00
|
|
|
}
|
2020-07-07 17:32:33 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
fn write_word_positions<I>(sorter: &mut Sorter<MergeFn>, iter: I) -> anyhow::Result<()>
|
|
|
|
where I: IntoIterator<Item=(SmallVec32<u8>, RoaringBitmap)>
|
|
|
|
{
|
2020-08-06 17:08:24 +08:00
|
|
|
// 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();
|
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
for (word, positions) in iter {
|
2020-07-07 17:32:33 +08:00
|
|
|
key.truncate(1);
|
2020-08-04 21:19:21 +08:00
|
|
|
key.extend_from_slice(&word);
|
|
|
|
// We serialize the positions into a buffer
|
|
|
|
buffer.clear();
|
2020-08-29 21:14:04 +08:00
|
|
|
buffer.reserve(positions.serialized_size());
|
2020-08-04 21:19:21 +08:00
|
|
|
positions.serialize_into(&mut buffer)?;
|
|
|
|
// that we write under the generated key into MTBL
|
2020-08-05 18:10:41 +08:00
|
|
|
if lmdb_key_valid_size(&key) {
|
|
|
|
sorter.insert(&key, &buffer)?;
|
|
|
|
}
|
2020-07-07 17:32:33 +08:00
|
|
|
}
|
2020-07-01 23:24:55 +08:00
|
|
|
|
2020-08-04 21:19:21 +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)>
|
|
|
|
{
|
2020-08-06 17:08:24 +08:00
|
|
|
// postings positions ids keys are all prefixed
|
|
|
|
let mut key = vec![WORD_POSITION_DOCIDS_BYTE];
|
2020-08-04 21:19:21 +08:00
|
|
|
let mut buffer = Vec::new();
|
|
|
|
|
|
|
|
for ((word, pos), ids) in iter {
|
2020-07-07 17:32:33 +08:00
|
|
|
key.truncate(1);
|
2020-08-04 21:19:21 +08:00
|
|
|
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());
|
2020-08-04 21:19:21 +08:00
|
|
|
// We serialize the document ids into a buffer
|
|
|
|
buffer.clear();
|
2020-08-29 21:14:04 +08:00
|
|
|
buffer.reserve(ids.serialized_size());
|
2020-08-04 21:19:21 +08:00
|
|
|
ids.serialize_into(&mut buffer)?;
|
|
|
|
// that we write under the generated key into MTBL
|
2020-08-05 18:10:41 +08:00
|
|
|
if lmdb_key_valid_size(&key) {
|
|
|
|
sorter.insert(&key, &buffer)?;
|
|
|
|
}
|
2020-07-01 23:24:55 +08:00
|
|
|
}
|
2020-07-07 17:32:33 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
Ok(())
|
2020-07-01 23:24:55 +08:00
|
|
|
}
|
|
|
|
|
2020-08-29 22:29:18 +08:00
|
|
|
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(())
|
|
|
|
}
|
|
|
|
|
2020-08-06 17:08:24 +08:00
|
|
|
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();
|
2020-08-29 21:14:04 +08:00
|
|
|
buffer.reserve(ids.serialized_size());
|
2020-08-06 17:08:24 +08:00
|
|
|
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(())
|
|
|
|
}
|
|
|
|
|
2020-08-29 21:14:04 +08:00
|
|
|
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>)> {
|
2020-08-04 21:19:21 +08:00
|
|
|
Self::write_word_positions(&mut self.sorter, self.word_positions)?;
|
|
|
|
Self::write_word_position_docids(&mut self.sorter, self.word_position_docids)?;
|
2020-08-29 22:29:18 +08:00
|
|
|
Self::write_word_four_positions_docids(&mut self.sorter, self.word_four_positions_docids)?;
|
2020-08-06 17:08:24 +08:00
|
|
|
Self::write_word_attribute_docids(&mut self.sorter, self.word_attribute_docids)?;
|
2020-08-29 21:14:04 +08:00
|
|
|
Self::write_documents_ids(&mut self.sorter, self.documents_ids)?;
|
2020-07-07 17:32:33 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
let mut wtr = tempfile::tempfile().map(Writer::new)?;
|
|
|
|
let mut builder = fst::SetBuilder::memory();
|
2020-07-07 17:32:33 +08:00
|
|
|
|
2020-08-04 21:19:21 +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
|
|
|
}
|
|
|
|
|
2020-08-04 21:19:21 +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
|
|
|
|
2020-08-07 00:19:10 +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)?;
|
2020-08-07 00:19:10 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
let file = wtr.into_inner()?;
|
|
|
|
let mmap = unsafe { Mmap::map(&file)? };
|
|
|
|
let reader = Reader::new(mmap)?;
|
2020-07-01 23:24:55 +08:00
|
|
|
|
2020-08-07 19:11:31 +08:00
|
|
|
Ok((reader, docs_reader))
|
2020-08-04 21:19:21 +08:00
|
|
|
}
|
|
|
|
}
|
2020-07-07 17:32:33 +08:00
|
|
|
|
2020-08-07 00:19:10 +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())
|
|
|
|
}
|
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
fn merge(key: &[u8], values: &[Vec<u8>]) -> Result<Vec<u8>, ()> {
|
2020-08-06 17:08:24 +08:00
|
|
|
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
|
2020-08-06 17:08:24 +08:00
|
|
|
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] {
|
2020-08-29 22:29:18 +08:00
|
|
|
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-06 17:08:24 +08:00
|
|
|
|
2020-08-31 03:50:30 +08:00
|
|
|
let mut head = RoaringBitmap::deserialize_from(head.as_slice()).unwrap();
|
|
|
|
for value in tail {
|
2020-08-06 17:08:24 +08:00
|
|
|
let bitmap = RoaringBitmap::deserialize_from(value.as_slice()).unwrap();
|
2020-08-31 03:50:30 +08:00
|
|
|
head.union_with(&bitmap);
|
2020-08-06 17:08:24 +08:00
|
|
|
}
|
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();
|
2020-08-06 17:08:24 +08:00
|
|
|
Ok(vec)
|
|
|
|
},
|
|
|
|
otherwise => panic!("wut {:?}", otherwise),
|
2020-07-07 17:32:33 +08:00
|
|
|
}
|
|
|
|
}
|
2020-07-01 23:24:55 +08:00
|
|
|
}
|
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
// TODO merge with the previous values
|
2020-08-06 17:08:24 +08:00
|
|
|
// TODO store the documents in a compressed MTBL
|
2020-08-04 21:19:21 +08:00
|
|
|
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)?;
|
|
|
|
}
|
2020-08-29 21:14:04 +08:00
|
|
|
else if key == DOCUMENTS_IDS_KEY {
|
|
|
|
// Write the documents ids list
|
|
|
|
index.main.put::<_, Str, ByteSlice>(wtxn, "documents-ids", val)?;
|
|
|
|
}
|
2020-08-06 17:08:24 +08:00
|
|
|
else if key.starts_with(&[WORD_POSITIONS_BYTE]) {
|
2020-08-04 21:19:21 +08:00
|
|
|
// Write the postings lists
|
|
|
|
index.word_positions.as_polymorph()
|
|
|
|
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
|
|
|
|
}
|
2020-08-06 17:08:24 +08:00
|
|
|
else if key.starts_with(&[WORD_POSITION_DOCIDS_BYTE]) {
|
2020-08-04 21:19:21 +08:00
|
|
|
// Write the postings lists
|
|
|
|
index.word_position_docids.as_polymorph()
|
|
|
|
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
|
|
|
|
}
|
2020-08-29 22:29:18 +08:00
|
|
|
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)?;
|
|
|
|
}
|
2020-08-06 17:08:24 +08:00
|
|
|
else if key.starts_with(&[WORD_ATTRIBUTE_DOCIDS_BYTE]) {
|
|
|
|
// Write the attribute postings lists
|
|
|
|
index.word_attribute_docids.as_polymorph()
|
2020-08-04 21:19:21 +08:00
|
|
|
.put::<_, ByteSlice, ByteSlice>(wtxn, &key[1..], val)?;
|
|
|
|
}
|
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|
2020-07-07 22:48:49 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
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();
|
2020-07-07 22:48:49 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
let mut builder = Merger::builder(merge);
|
|
|
|
builder.extend(sources);
|
|
|
|
let merger = builder.build();
|
2020-07-07 22:48:49 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
let mut iter = merger.into_merge_iter()?;
|
|
|
|
while let Some(result) = iter.next() {
|
|
|
|
let (k, v) = result?;
|
2020-08-16 02:37:13 +08:00
|
|
|
(f)(&k, &v).with_context(|| format!("writing {:?} {:?} into LMDB", k.as_bstr(), k.as_bstr()))?;
|
2020-08-04 21:19:21 +08:00
|
|
|
}
|
2020-07-07 22:48:49 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
debug!("MTBL stores merged in {:.02?}!", before.elapsed());
|
|
|
|
Ok(())
|
2020-07-07 22:48:49 +08:00
|
|
|
}
|
|
|
|
|
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>>,
|
2020-07-01 23:49:46 +08:00
|
|
|
thread_index: usize,
|
2020-07-07 17:32:33 +08:00
|
|
|
num_threads: usize,
|
2020-08-29 22:29:18 +08:00
|
|
|
arc_cache_size: usize,
|
2020-08-05 18:10:41 +08:00
|
|
|
max_nb_chunks: Option<usize>,
|
|
|
|
max_memory: Option<usize>,
|
2020-08-21 22:41:26 +08:00
|
|
|
chunk_compression_type: CompressionType,
|
|
|
|
chunk_compression_level: Option<u32>,
|
2020-08-07 19:11:31 +08:00
|
|
|
) -> anyhow::Result<(Reader<Mmap>, Reader<Mmap>)>
|
2020-07-01 23:49:46 +08:00
|
|
|
{
|
2020-08-07 00:19:10 +08:00
|
|
|
debug!("{:?}: Indexing into a Store...", thread_index);
|
2020-05-26 02:39:53 +08:00
|
|
|
|
2020-08-21 22:41:26 +08:00
|
|
|
let mut store = Store::new(
|
|
|
|
arc_cache_size,
|
|
|
|
max_nb_chunks,
|
|
|
|
max_memory,
|
|
|
|
chunk_compression_type,
|
|
|
|
chunk_compression_level,
|
|
|
|
);
|
2020-07-07 22:48:49 +08:00
|
|
|
|
2020-08-04 21:19:21 +08:00
|
|
|
// 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()?;
|
2020-08-04 21:19:21 +08:00
|
|
|
store.write_headers(&headers)?;
|
2020-05-26 02:39:53 +08:00
|
|
|
|
2020-08-06 17:08:24 +08:00
|
|
|
let mut before = Instant::now();
|
2020-07-07 17:32:33 +08:00
|
|
|
let mut document_id: usize = 0;
|
2020-07-07 22:48:49 +08:00
|
|
|
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
|
|
|
|
2020-08-29 16:56:40 +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
|
|
|
|
2020-08-29 16:56:40 +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();
|
2020-08-29 16:56:40 +08:00
|
|
|
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
|
|
|
}
|
2020-08-29 16:56:40 +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
|
|
|
}
|
|
|
|
|
2020-08-29 16:56:40 +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-07-04 23:02:27 +08:00
|
|
|
|
2020-08-07 19:11:31 +08:00
|
|
|
let (reader, docs_reader) = store.finish()?;
|
2020-08-04 21:19:21 +08:00
|
|
|
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
|
|
|
|
2020-07-12 17:04:35 +08:00
|
|
|
stderrlog::new()
|
|
|
|
.verbosity(opt.verbose)
|
|
|
|
.show_level(false)
|
|
|
|
.timestamp(stderrlog::Timestamp::Off)
|
|
|
|
.init()?;
|
|
|
|
|
2020-06-28 18:13:10 +08:00
|
|
|
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()
|
2020-08-10 19:53:53 +08:00
|
|
|
.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
|
|
|
|
2020-08-28 21:38:05 +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();
|
2020-08-21 22:41:26 +08:00
|
|
|
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-05 18:10:41 +08:00
|
|
|
|
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)?;
|
2020-08-22 00:10:24 +08:00
|
|
|
// 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();
|
2020-08-21 22:08:32 +08:00
|
|
|
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
|
2020-07-07 18:21:22 +08:00
|
|
|
.into_par_iter()
|
2020-07-07 17:32:33 +08:00
|
|
|
.enumerate()
|
2020-08-21 22:41:26 +08:00
|
|
|
.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
|
|
|
|
2020-08-28 21:38:05 +08:00
|
|
|
let mut wtxn = env.write_txn()?;
|
2020-08-29 21:14:04 +08:00
|
|
|
|
|
|
|
// We merge the postings lists into LMDB.
|
|
|
|
debug!("We are writing the postings lists into LMDB on disk...");
|
2020-08-28 21:38:05 +08:00
|
|
|
merge_into_lmdb(stores, |k, v| lmdb_writer(&mut wtxn, &index, k, v))?;
|
2020-08-29 21:14:04 +08:00
|
|
|
|
|
|
|
// 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.
|
2020-08-28 21:38:05 +08:00
|
|
|
let count = index.number_of_documents(&wtxn)?;
|
2020-08-29 21:14:04 +08:00
|
|
|
|
2020-08-28 21:38:05 +08:00
|
|
|
wtxn.commit()?;
|
|
|
|
|
|
|
|
info!("Wrote {} documents in {:.02?}", count, before_indexing.elapsed());
|
2020-05-26 02:39:53 +08:00
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|