MeiliSearch/src/subcommand/search.rs

80 lines
2.3 KiB
Rust
Raw Normal View History

use std::io::{self, 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;
use structopt::StructOpt;
2020-05-31 16:09:34 +02:00
use crate::Index;
2020-05-31 16:09:34 +02:00
#[derive(Debug, StructOpt)]
/// 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,
/// 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,
/// 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).
query: Option<String>,
2020-05-31 16:09:34 +02:00
}
pub fn run(opt: Opt) -> anyhow::Result<()> {
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(opt.database_size)
.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.
let index = Index::new(&env)?;
2020-05-31 16:09:34 +02:00
let rtxn = env.read_txn()?;
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 = _>>,
};
for result in lines {
let before = Instant::now();
2020-05-31 16:09:34 +02:00
let query = result?;
2020-08-13 14:15:05 +02:00
let result = index.search(&rtxn).query(query).execute().unwrap();
let headers = match index.headers(&rtxn)? {
Some(headers) => headers,
None => return Ok(()),
};
let documents = index.documents(&rtxn, result.documents_ids.iter().cloned())?;
let mut wtr = csv::Writer::from_writer(io::stdout());
wtr.write_record(&headers)?;
for (_id, record) in documents {
wtr.write_record(&record)?;
2020-05-31 16:09:34 +02:00
}
wtr.flush()?;
2020-05-31 16:09:34 +02:00
2020-08-13 14:15:05 +02:00
debug!("Took {:.02?} to find {} documents", before.elapsed(), result.documents_ids.len());
}
2020-05-31 16:09:34 +02:00
Ok(())
}