2021-01-16 15:51:23 +01:00
|
|
|
use std::io::{self, BufRead, Write};
|
2020-06-20 12:01:15 +02:00
|
|
|
use std::iter::once;
|
2020-05-31 16:09:34 +02:00
|
|
|
use std::path::PathBuf;
|
|
|
|
use std::time::Instant;
|
|
|
|
|
2020-12-20 11:55:21 +01:00
|
|
|
use byte_unit::Byte;
|
2020-05-31 17:48:13 +02:00
|
|
|
use heed::EnvOpenOptions;
|
2020-07-12 10:55:09 +02:00
|
|
|
use log::debug;
|
|
|
|
use structopt::StructOpt;
|
2020-05-31 16:09:34 +02:00
|
|
|
|
2021-01-16 15:51:23 +01:00
|
|
|
use crate::{Index, obkv_to_json};
|
2020-06-10 22:05:01 +02:00
|
|
|
|
2020-05-31 16:09:34 +02:00
|
|
|
#[derive(Debug, StructOpt)]
|
2020-10-19 13:44:17 +02:00
|
|
|
/// A simple search helper binary for the milli project.
|
|
|
|
pub struct Opt {
|
2020-05-31 16:09:34 +02:00
|
|
|
/// 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 13:53:53 +02: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).
|
2020-12-20 11:55:21 +01:00
|
|
|
#[structopt(long = "db-size", default_value = "100 GiB")]
|
|
|
|
database_size: Byte,
|
2020-08-10 13:53:53 +02:00
|
|
|
|
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>,
|
2021-01-20 17:11:36 +01:00
|
|
|
|
|
|
|
/// Compute and print the facet distribution of all the faceted fields.
|
|
|
|
#[structopt(long)]
|
|
|
|
print_facet_distribution: bool,
|
2020-05-31 16:09:34 +02:00
|
|
|
}
|
|
|
|
|
2020-10-19 13:44:17 +02:00
|
|
|
pub fn run(opt: Opt) -> anyhow::Result<()> {
|
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)?;
|
2020-10-30 10:56:35 +01:00
|
|
|
let mut options = EnvOpenOptions::new();
|
2020-12-20 11:55:21 +01:00
|
|
|
options.map_size(opt.database_size.get_bytes() as usize);
|
2020-05-31 16:09:34 +02:00
|
|
|
|
2020-08-07 13:11:31 +02:00
|
|
|
// Open the LMDB database.
|
2020-10-30 10:56:35 +01:00
|
|
|
let index = Index::new(options, &opt.database)?;
|
|
|
|
let rtxn = index.read_txn()?;
|
2021-01-16 15:51:23 +01:00
|
|
|
let fields_ids_map = index.fields_ids_map(&rtxn)?;
|
2021-01-20 17:27:43 +01:00
|
|
|
let displayed_fields = match index.displayed_fields_ids(&rtxn)? {
|
|
|
|
Some(fields) => fields,
|
|
|
|
None => fields_ids_map.iter().map(|(id, _)| id).collect(),
|
2021-01-16 15:51:23 +01:00
|
|
|
};
|
2020-05-31 16:09:34 +02:00
|
|
|
|
2020-06-20 12:01:15 +02:00
|
|
|
let stdin = io::stdin();
|
|
|
|
let lines = match opt.query {
|
2020-12-01 14:25:17 +01:00
|
|
|
Some(query) => Box::new(once(Ok(query))),
|
2020-06-20 12:01:15 +02:00
|
|
|
None => Box::new(stdin.lock().lines()) as Box<dyn Iterator<Item = _>>,
|
2020-05-31 17:01:11 +02:00
|
|
|
};
|
|
|
|
|
2021-01-16 15:51:23 +01:00
|
|
|
let mut stdout = io::stdout();
|
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?;
|
2021-01-16 15:51:23 +01:00
|
|
|
let result = index.search(&rtxn).query(query).execute()?;
|
2020-08-28 15:38:05 +02:00
|
|
|
let documents = index.documents(&rtxn, result.documents_ids.iter().cloned())?;
|
2020-06-20 12:01:15 +02:00
|
|
|
|
2020-08-31 14:20:42 +02:00
|
|
|
for (_id, record) in documents {
|
2021-01-16 15:51:23 +01:00
|
|
|
let val = obkv_to_json(&displayed_fields, &fields_ids_map, record)?;
|
|
|
|
serde_json::to_writer(&mut stdout, &val)?;
|
|
|
|
let _ = writeln!(&mut stdout);
|
2020-05-31 16:09:34 +02:00
|
|
|
}
|
|
|
|
|
2021-01-20 17:11:36 +01:00
|
|
|
if opt.print_facet_distribution {
|
|
|
|
let facets = index.facets(&rtxn).candidates(result.candidates).execute()?;
|
|
|
|
serde_json::to_writer(&mut stdout, &facets)?;
|
|
|
|
let _ = writeln!(&mut stdout);
|
|
|
|
}
|
|
|
|
|
2020-08-13 14:15:05 +02:00
|
|
|
debug!("Took {:.02?} to find {} documents", before.elapsed(), result.documents_ids.len());
|
2020-06-20 12:01:15 +02:00
|
|
|
}
|
2020-05-31 16:09:34 +02:00
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|