From c35befbf381c93662c987132ffb4db4e33d9893d Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Cl=C3=A9ment=20Renault?= <clement@meilisearch.com>
Date: Sat, 16 Jan 2021 15:51:23 +0100
Subject: [PATCH] Fix the search subcommand document display loop

---
 src/subcommand/search.rs | 31 +++++++++++++------------------
 1 file changed, 13 insertions(+), 18 deletions(-)

diff --git a/src/subcommand/search.rs b/src/subcommand/search.rs
index a4bdae870..b073ceae9 100644
--- a/src/subcommand/search.rs
+++ b/src/subcommand/search.rs
@@ -1,16 +1,15 @@
-use std::collections::HashMap;
-use std::io::{self, BufRead};
+use std::borrow::Cow;
+use std::io::{self, BufRead, Write};
 use std::iter::once;
 use std::path::PathBuf;
 use std::time::Instant;
 
-use anyhow::Context;
 use byte_unit::Byte;
 use heed::EnvOpenOptions;
 use log::debug;
 use structopt::StructOpt;
 
-use crate::Index;
+use crate::{Index, obkv_to_json};
 
 #[derive(Debug, StructOpt)]
 /// A simple search helper binary for the milli project.
@@ -47,6 +46,11 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
     // Open the LMDB database.
     let index = Index::new(options, &opt.database)?;
     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 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 = _>>,
     };
 
+    let mut stdout = io::stdout();
     for result in lines {
         let before = Instant::now();
 
         let query = result?;
-        let result = index.search(&rtxn).query(query).execute().unwrap();
-
-        let mut stdout = io::stdout();
-        let fields_ids_map = index.fields_ids_map(&rtxn)?;
+        let result = index.search(&rtxn).query(query).execute()?;
         let documents = index.documents(&rtxn, result.documents_ids.iter().cloned())?;
 
         for (_id, record) in documents {
-            let document: anyhow::Result<HashMap<_, _>> = record.iter()
-                .map(|(k, v)| {
-                    let key = fields_ids_map.name(k).context("field id not found")?;
-                    let val = std::str::from_utf8(v)?;
-                    Ok((key, val))
-                })
-                .collect();
-
-            let document = document?;
-            serde_json::to_writer(&mut stdout, &document)?;
+            let val = obkv_to_json(&displayed_fields, &fields_ids_map, record)?;
+            serde_json::to_writer(&mut stdout, &val)?;
+            let _ = writeln!(&mut stdout);
         }
 
         debug!("Took {:.02?} to find {} documents", before.elapsed(), result.documents_ids.len());