2020-05-31 22:09:34 +08:00
|
|
|
use std::io::{self, Write};
|
|
|
|
use std::path::PathBuf;
|
|
|
|
use std::time::Instant;
|
|
|
|
|
2020-05-31 23:48:13 +08:00
|
|
|
use heed::EnvOpenOptions;
|
2020-05-31 22:09:34 +08:00
|
|
|
use structopt::StructOpt;
|
2020-05-31 23:48:13 +08:00
|
|
|
use mega_mini_indexer::{Index, BEU32};
|
2020-05-31 22:09:34 +08:00
|
|
|
|
|
|
|
#[derive(Debug, StructOpt)]
|
2020-06-05 00:19:52 +08:00
|
|
|
#[structopt(name = "mm-search", about = "The server side of the MMI project.")]
|
2020-05-31 22:09:34 +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 query string to search for (doesn't support prefix search yet).
|
|
|
|
query: String,
|
|
|
|
}
|
|
|
|
|
|
|
|
fn main() -> anyhow::Result<()> {
|
|
|
|
let opt = Opt::from_args();
|
|
|
|
|
|
|
|
std::fs::create_dir_all(&opt.database)?;
|
|
|
|
let env = EnvOpenOptions::new()
|
|
|
|
.map_size(100 * 1024 * 1024 * 1024) // 100 GB
|
|
|
|
.max_readers(10)
|
|
|
|
.max_dbs(5)
|
|
|
|
.open(opt.database)?;
|
|
|
|
|
2020-05-31 23:48:13 +08:00
|
|
|
let index = Index::new(&env)?;
|
2020-05-31 22:09:34 +08:00
|
|
|
|
2020-05-31 23:48:13 +08:00
|
|
|
let before = Instant::now();
|
2020-05-31 22:09:34 +08:00
|
|
|
let rtxn = env.read_txn()?;
|
|
|
|
|
2020-05-31 23:48:13 +08:00
|
|
|
let documents_ids = index.search(&rtxn, &opt.query)?;
|
|
|
|
let headers = match index.headers(&rtxn)? {
|
|
|
|
Some(headers) => headers,
|
2020-05-31 23:01:11 +08:00
|
|
|
None => return Ok(()),
|
|
|
|
};
|
|
|
|
|
2020-05-31 22:09:34 +08:00
|
|
|
let mut stdout = io::stdout();
|
|
|
|
stdout.write_all(&headers)?;
|
|
|
|
|
2020-05-31 23:48:13 +08:00
|
|
|
for id in &documents_ids {
|
|
|
|
if let Some(content) = index.documents.get(&rtxn, &BEU32::new(*id))? {
|
2020-05-31 22:09:34 +08:00
|
|
|
stdout.write_all(&content)?;
|
|
|
|
}
|
|
|
|
}
|
|
|
|
|
2020-05-31 23:48:13 +08:00
|
|
|
eprintln!("Took {:.02?} to find {} documents", before.elapsed(), documents_ids.len());
|
2020-05-31 22:09:34 +08:00
|
|
|
|
|
|
|
Ok(())
|
|
|
|
}
|