2020-08-07 13:11:31 +02:00
|
|
|
use std::fs::File;
|
2020-06-20 12:01:15 +02:00
|
|
|
use std::io::{self, Write, BufRead};
|
|
|
|
use std::iter::once;
|
2020-05-31 16:09:34 +02:00
|
|
|
use std::path::PathBuf;
|
|
|
|
use std::time::Instant;
|
|
|
|
|
2020-05-31 17:48:13 +02:00
|
|
|
use heed::EnvOpenOptions;
|
2020-07-12 10:55:09 +02:00
|
|
|
use log::debug;
|
2020-08-06 18:19:10 +02:00
|
|
|
use milli::Index;
|
2020-08-07 13:11:31 +02:00
|
|
|
use oxidized_mtbl::Reader;
|
2020-07-12 10:55:09 +02:00
|
|
|
use structopt::StructOpt;
|
2020-05-31 16:09:34 +02:00
|
|
|
|
2020-06-10 22:05:01 +02:00
|
|
|
#[cfg(target_os = "linux")]
|
|
|
|
#[global_allocator]
|
|
|
|
static ALLOC: jemallocator::Jemalloc = jemallocator::Jemalloc;
|
|
|
|
|
2020-05-31 16:09:34 +02:00
|
|
|
#[derive(Debug, StructOpt)]
|
2020-07-12 11:06:45 +02:00
|
|
|
#[structopt(name = "milli-search", about = "A simple search binary for milli project.")]
|
2020-05-31 16:09:34 +02: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-07-12 11:04:35 +02:00
|
|
|
/// Verbose mode (-v, -vv, -vvv, etc.)
|
|
|
|
#[structopt(short, long, parse(from_occurrences))]
|
|
|
|
verbose: usize,
|
|
|
|
|
2020-05-31 16:09:34 +02:00
|
|
|
/// The query string to search for (doesn't support prefix search yet).
|
2020-06-20 12:01:15 +02:00
|
|
|
query: Option<String>,
|
2020-05-31 16:09:34 +02:00
|
|
|
}
|
|
|
|
|
|
|
|
fn main() -> anyhow::Result<()> {
|
|
|
|
let opt = Opt::from_args();
|
|
|
|
|
2020-07-12 11:04:35 +02:00
|
|
|
stderrlog::new()
|
|
|
|
.verbosity(opt.verbose)
|
|
|
|
.show_level(false)
|
|
|
|
.timestamp(stderrlog::Timestamp::Off)
|
|
|
|
.init()?;
|
|
|
|
|
2020-05-31 16:09:34 +02:00
|
|
|
std::fs::create_dir_all(&opt.database)?;
|
|
|
|
let env = EnvOpenOptions::new()
|
|
|
|
.map_size(100 * 1024 * 1024 * 1024) // 100 GB
|
|
|
|
.max_readers(10)
|
2020-06-22 19:04:10 +02:00
|
|
|
.max_dbs(10)
|
2020-08-07 13:11:31 +02:00
|
|
|
.open(&opt.database)?;
|
2020-05-31 16:09:34 +02:00
|
|
|
|
2020-08-07 13:11:31 +02:00
|
|
|
// Open the LMDB database.
|
2020-05-31 17:48:13 +02:00
|
|
|
let index = Index::new(&env)?;
|
2020-05-31 16:09:34 +02:00
|
|
|
|
2020-08-07 13:11:31 +02:00
|
|
|
// Open the documents MTBL database.
|
|
|
|
let path = opt.database.join("documents.mtbl");
|
|
|
|
let file = File::open(path)?;
|
|
|
|
let mmap = unsafe { memmap::Mmap::map(&file)? };
|
|
|
|
let documents = Reader::new(mmap.as_ref())?;
|
|
|
|
|
2020-05-31 16:09:34 +02:00
|
|
|
let rtxn = env.read_txn()?;
|
|
|
|
|
2020-06-20 12:01:15 +02:00
|
|
|
let stdin = io::stdin();
|
|
|
|
let lines = match opt.query {
|
|
|
|
Some(query) => Box::new(once(Ok(query.to_string()))),
|
|
|
|
None => Box::new(stdin.lock().lines()) as Box<dyn Iterator<Item = _>>,
|
2020-05-31 17:01:11 +02:00
|
|
|
};
|
|
|
|
|
2020-06-20 12:01:15 +02:00
|
|
|
for result in lines {
|
|
|
|
let before = Instant::now();
|
2020-05-31 16:09:34 +02:00
|
|
|
|
2020-06-20 12:01:15 +02:00
|
|
|
let query = result?;
|
2020-07-14 11:51:21 +02:00
|
|
|
let (_, documents_ids) = index.search(&rtxn, &query)?;
|
2020-06-20 12:01:15 +02:00
|
|
|
let headers = match index.headers(&rtxn)? {
|
|
|
|
Some(headers) => headers,
|
|
|
|
None => return Ok(()),
|
|
|
|
};
|
|
|
|
|
|
|
|
let mut stdout = io::stdout();
|
|
|
|
stdout.write_all(&headers)?;
|
|
|
|
|
|
|
|
for id in &documents_ids {
|
2020-08-06 18:19:10 +02:00
|
|
|
let id_bytes = id.to_be_bytes();
|
|
|
|
if let Some(content) = documents.clone().get(&id_bytes)? {
|
|
|
|
stdout.write_all(content.as_ref())?;
|
2020-06-20 12:01:15 +02:00
|
|
|
}
|
2020-05-31 16:09:34 +02:00
|
|
|
}
|
|
|
|
|
2020-07-12 10:55:09 +02:00
|
|
|
debug!("Took {:.02?} to find {} documents", before.elapsed(), documents_ids.len());
|
2020-06-20 12:01:15 +02:00
|
|
|
}
|
2020-05-31 16:09:34 +02:00
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|