From 7c0e544839d1cea94f45f07757251c8d25f6366e Mon Sep 17 00:00:00 2001
From: Gregory Conrad <gregorysconrad@gmail.com>
Date: Wed, 23 Nov 2022 21:18:58 -0500
Subject: [PATCH] feat: Add all_obkv_to_json function

---
 cli/src/main.rs              |  2 +-
 milli/src/lib.rs             | 30 +++++++++++++++++++++++++++++-
 milli/src/update/settings.rs |  2 +-
 3 files changed, 31 insertions(+), 3 deletions(-)

diff --git a/cli/src/main.rs b/cli/src/main.rs
index dd5489ebc..f2bbc0b23 100644
--- a/cli/src/main.rs
+++ b/cli/src/main.rs
@@ -450,7 +450,7 @@ impl Search {
         let documents = index.documents(&txn, result.documents_ids)?;
         let mut jsons = Vec::new();
         for (_, obkv) in documents {
-            let json = milli::obkv_to_json(&displayed_fields, &fields_ids_map, obkv)?;
+            let json = milli::obkv_to_json(&displayed_fields, &fields_ids_map, &obkv)?;
             jsons.push(json);
         }
 
diff --git a/milli/src/lib.rs b/milli/src/lib.rs
index c33aae9eb..7d74637f2 100644
--- a/milli/src/lib.rs
+++ b/milli/src/lib.rs
@@ -88,7 +88,7 @@ pub fn absolute_from_relative_position(field_id: FieldId, relative: RelativePosi
 pub fn obkv_to_json(
     displayed_fields: &[FieldId],
     fields_ids_map: &FieldsIdsMap,
-    obkv: obkv::KvReaderU16,
+    obkv: &obkv::KvReaderU16,
 ) -> Result<Object> {
     displayed_fields
         .iter()
@@ -105,6 +105,12 @@ pub fn obkv_to_json(
         .collect()
 }
 
+/// Transform every field of a raw obkv store into a JSON Object.
+pub fn all_obkv_to_json(obkv: &obkv::KvReaderU16, fields_ids_map: &FieldsIdsMap) -> Result<Object> {
+    let all_keys = obkv.iter().map(|(k, _v)| k).collect::<Vec<_>>();
+    obkv_to_json(all_keys.as_slice(), fields_ids_map, obkv)
+}
+
 /// Transform a JSON value into a string that can be indexed.
 pub fn json_to_string(value: &Value) -> Option<String> {
     fn inner(value: &Value, output: &mut String) -> bool {
@@ -285,4 +291,26 @@ mod tests {
         assert_eq!(0x12345678, absolute_from_relative_position(0x1234, 0x5678));
         assert_eq!(0xFFFFFFFF, absolute_from_relative_position(0xFFFF, 0xFFFF));
     }
+
+    #[test]
+    fn test_all_obkv_to_json() {
+        let mut fields_ids_map = FieldsIdsMap::new();
+        let id1 = fields_ids_map.insert("field1").unwrap();
+        let id2 = fields_ids_map.insert("field2").unwrap();
+
+        let mut writer = obkv::KvWriterU16::memory();
+        writer.insert(id1, b"1234").unwrap();
+        writer.insert(id2, b"4321").unwrap();
+        let contents = writer.into_inner().unwrap();
+        let obkv = obkv::KvReaderU16::new(&contents);
+
+        let expected = json!({
+            "field1": 1234,
+            "field2": 4321,
+        });
+        let expected = expected.as_object().unwrap();
+        let actual = all_obkv_to_json(&obkv, &fields_ids_map).unwrap();
+
+        assert_eq!(&actual, expected);
+    }
 }
diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs
index 6da32d73f..eb5f56c45 100644
--- a/milli/src/update/settings.rs
+++ b/milli/src/update/settings.rs
@@ -920,7 +920,7 @@ mod tests {
         let fidmap = index.fields_ids_map(&rtxn).unwrap();
         for document in index.all_documents(&rtxn).unwrap() {
             let document = document.unwrap();
-            let json = crate::obkv_to_json(&fidmap.ids().collect::<Vec<_>>(), &fidmap, document.1)
+            let json = crate::obkv_to_json(&fidmap.ids().collect::<Vec<_>>(), &fidmap, &document.1)
                 .unwrap();
             println!("json: {:?}", json);
         }