Fix the search subcommand document display loop

This commit is contained in:
Clément Renault 2021-01-16 15:51:23 +01:00
parent 2fa5808e3f
commit c35befbf38
No known key found for this signature in database
GPG Key ID: 92ADA4E935E71FA4

View File

@ -1,16 +1,15 @@
use std::collections::HashMap; use std::borrow::Cow;
use std::io::{self, BufRead}; use std::io::{self, BufRead, Write};
use std::iter::once; use std::iter::once;
use std::path::PathBuf; use std::path::PathBuf;
use std::time::Instant; use std::time::Instant;
use anyhow::Context;
use byte_unit::Byte; use byte_unit::Byte;
use heed::EnvOpenOptions; use heed::EnvOpenOptions;
use log::debug; use log::debug;
use structopt::StructOpt; use structopt::StructOpt;
use crate::Index; use crate::{Index, obkv_to_json};
#[derive(Debug, StructOpt)] #[derive(Debug, StructOpt)]
/// A simple search helper binary for the milli project. /// A simple search helper binary for the milli project.
@ -47,6 +46,11 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
// Open the LMDB database. // Open the LMDB database.
let index = Index::new(options, &opt.database)?; let index = Index::new(options, &opt.database)?;
let rtxn = index.read_txn()?; let rtxn = index.read_txn()?;
let fields_ids_map = index.fields_ids_map(&rtxn)?;
let displayed_fields = match index.displayed_fields(&rtxn)? {
Some(fields) => Cow::Borrowed(fields),
None => Cow::Owned(fields_ids_map.iter().map(|(id, _)| id).collect()),
};
let stdin = io::stdin(); let stdin = io::stdin();
let lines = match opt.query { let lines = match opt.query {
@ -54,27 +58,18 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
None => Box::new(stdin.lock().lines()) as Box<dyn Iterator<Item = _>>, None => Box::new(stdin.lock().lines()) as Box<dyn Iterator<Item = _>>,
}; };
let mut stdout = io::stdout();
for result in lines { for result in lines {
let before = Instant::now(); let before = Instant::now();
let query = result?; let query = result?;
let result = index.search(&rtxn).query(query).execute().unwrap(); let result = index.search(&rtxn).query(query).execute()?;
let mut stdout = io::stdout();
let fields_ids_map = index.fields_ids_map(&rtxn)?;
let documents = index.documents(&rtxn, result.documents_ids.iter().cloned())?; let documents = index.documents(&rtxn, result.documents_ids.iter().cloned())?;
for (_id, record) in documents { for (_id, record) in documents {
let document: anyhow::Result<HashMap<_, _>> = record.iter() let val = obkv_to_json(&displayed_fields, &fields_ids_map, record)?;
.map(|(k, v)| { serde_json::to_writer(&mut stdout, &val)?;
let key = fields_ids_map.name(k).context("field id not found")?; let _ = writeln!(&mut stdout);
let val = std::str::from_utf8(v)?;
Ok((key, val))
})
.collect();
let document = document?;
serde_json::to_writer(&mut stdout, &document)?;
} }
debug!("Took {:.02?} to find {} documents", before.elapsed(), result.documents_ids.len()); debug!("Took {:.02?} to find {} documents", before.elapsed(), result.documents_ids.len());