Merge pull request #21 from meilisearch/displayed-attributes

Add the displayed attributes setting to an index

src/index.rs

@@ -3,7 +3,7 @@ use std::path::Path;
 
 use anyhow::Context;
 use heed::types::*;
-use heed::{PolyDatabase, Database};
+use heed::{PolyDatabase, Database, RwTxn, RoTxn};
 use roaring::RoaringBitmap;
 
 use crate::Search;
@@ -14,6 +14,7 @@ use crate::{
     BoRoaringBitmapCodec, CboRoaringBitmapCodec,
 };
 
+pub const DISPLAYED_FIELDS_KEY: &str = "displayed-fields";
 pub const DOCUMENTS_IDS_KEY: &str = "documents-ids";
 pub const FIELDS_IDS_MAP_KEY: &str = "fields-ids-map";
 pub const PRIMARY_KEY_KEY: &str = "primary-key";
@@ -51,12 +52,12 @@ impl Index {
     }
 
     /// Create a write transaction to be able to write into the index.
-    pub fn write_txn(&self) -> heed::Result<heed::RwTxn> {
+    pub fn write_txn(&self) -> heed::Result<RwTxn> {
         self.env.write_txn()
     }
 
     /// Create a read transaction to be able to read the index.
-    pub fn read_txn(&self) -> heed::Result<heed::RoTxn> {
+    pub fn read_txn(&self) -> heed::Result<RoTxn> {
         self.env.read_txn()
     }
 
@@ -74,65 +75,95 @@ impl Index {
         self.env.prepare_for_closing()
     }
 
+    /* documents ids */
+
     /// Writes the documents ids that corresponds to the user-ids-documents-ids FST.
-    pub fn put_documents_ids(&self, wtxn: &mut heed::RwTxn, docids: &RoaringBitmap) -> heed::Result<()> {
+    pub fn put_documents_ids(&self, wtxn: &mut RwTxn, docids: &RoaringBitmap) -> heed::Result<()> {
         self.main.put::<_, Str, RoaringBitmapCodec>(wtxn, DOCUMENTS_IDS_KEY, docids)
     }
 
     /// Returns the internal documents ids.
-    pub fn documents_ids(&self, rtxn: &heed::RoTxn) -> heed::Result<RoaringBitmap> {
+    pub fn documents_ids(&self, rtxn: &RoTxn) -> heed::Result<RoaringBitmap> {
         Ok(self.main.get::<_, Str, RoaringBitmapCodec>(rtxn, DOCUMENTS_IDS_KEY)?.unwrap_or_default())
     }
 
+    /* primary key */
+
     /// Writes the documents primary key, this is the field name that is used to store the id.
-    pub fn put_primary_key(&self, wtxn: &mut heed::RwTxn, primary_key: u8) -> heed::Result<()> {
+    pub fn put_primary_key(&self, wtxn: &mut RwTxn, primary_key: u8) -> heed::Result<()> {
         self.main.put::<_, Str, OwnedType<u8>>(wtxn, PRIMARY_KEY_KEY, &primary_key)
     }
 
     /// Delete the primary key of the documents, this can be done to reset indexes settings.
-    pub fn delete_primary_key(&self, wtxn: &mut heed::RwTxn) -> heed::Result<bool> {
+    pub fn delete_primary_key(&self, wtxn: &mut RwTxn) -> heed::Result<bool> {
         self.main.delete::<_, Str>(wtxn, PRIMARY_KEY_KEY)
     }
 
     /// Returns the documents primary key, `None` if it hasn't been defined.
-    pub fn primary_key(&self, rtxn: &heed::RoTxn) -> heed::Result<Option<u8>> {
+    pub fn primary_key(&self, rtxn: &RoTxn) -> heed::Result<Option<u8>> {
         self.main.get::<_, Str, OwnedType<u8>>(rtxn, PRIMARY_KEY_KEY)
     }
 
+    /* users ids documents ids */
+
     /// Writes the users ids documents ids, a user id is a byte slice (i.e. `[u8]`)
     /// and refers to an internal id (i.e. `u32`).
-    pub fn put_users_ids_documents_ids<A: AsRef<[u8]>>(&self, wtxn: &mut heed::RwTxn, fst: &fst::Map<A>) -> heed::Result<()> {
+    pub fn put_users_ids_documents_ids<A: AsRef<[u8]>>(&self, wtxn: &mut RwTxn, fst: &fst::Map<A>) -> heed::Result<()> {
         self.main.put::<_, Str, ByteSlice>(wtxn, USERS_IDS_DOCUMENTS_IDS_KEY, fst.as_fst().as_bytes())
     }
 
     /// Returns the user ids documents ids map which associate the user ids (i.e. `[u8]`)
     /// with the internal ids (i.e. `u32`).
-    pub fn users_ids_documents_ids<'t>(&self, rtxn: &'t heed::RoTxn) -> anyhow::Result<fst::Map<Cow<'t, [u8]>>> {
+    pub fn users_ids_documents_ids<'t>(&self, rtxn: &'t RoTxn) -> anyhow::Result<fst::Map<Cow<'t, [u8]>>> {
         match self.main.get::<_, Str, ByteSlice>(rtxn, USERS_IDS_DOCUMENTS_IDS_KEY)? {
             Some(bytes) => Ok(fst::Map::new(bytes)?.map_data(Cow::Borrowed)?),
             None => Ok(fst::Map::default().map_data(Cow::Owned)?),
         }
     }
 
+    /* fields ids map */
+
     /// Writes the fields ids map which associate the documents keys with an internal field id
     /// (i.e. `u8`), this field id is used to identify fields in the obkv documents.
-    pub fn put_fields_ids_map(&self, wtxn: &mut heed::RwTxn, map: &FieldsIdsMap) -> heed::Result<()> {
+    pub fn put_fields_ids_map(&self, wtxn: &mut RwTxn, map: &FieldsIdsMap) -> heed::Result<()> {
         self.main.put::<_, Str, SerdeJson<FieldsIdsMap>>(wtxn, FIELDS_IDS_MAP_KEY, map)
     }
 
     /// Returns the fields ids map which associate the documents keys with an internal field id
     /// (i.e. `u8`), this field id is used to identify fields in the obkv documents.
-    pub fn fields_ids_map(&self, rtxn: &heed::RoTxn) -> heed::Result<FieldsIdsMap> {
+    pub fn fields_ids_map(&self, rtxn: &RoTxn) -> heed::Result<FieldsIdsMap> {
         Ok(self.main.get::<_, Str, SerdeJson<FieldsIdsMap>>(rtxn, FIELDS_IDS_MAP_KEY)?.unwrap_or_default())
     }
 
+    /* displayed fields */
+
+    /// Writes the fields ids that must be displayed in the defined order.
+    /// There must be not be any duplicate field id.
+    pub fn put_displayed_fields(&self, wtxn: &mut RwTxn, fields: &[u8]) -> heed::Result<()> {
+        self.main.put::<_, Str, ByteSlice>(wtxn, DISPLAYED_FIELDS_KEY, fields)
+    }
+
+    /// Deletes the displayed fields ids, this will make the engine to display
+    /// all the documents attributes in the order of the `FieldsIdsMap`.
+    pub fn delete_displayed_fields(&self, wtxn: &mut RwTxn) -> heed::Result<bool> {
+        self.main.delete::<_, Str>(wtxn, DISPLAYED_FIELDS_KEY)
+    }
+
+    /// Returns the displayed fields ids in the order they must be returned. If it returns
+    /// `None` it means that all the attributes are displayed in the order of the `FieldsIdsMap`.
+    pub fn displayed_fields<'t>(&self, rtxn: &'t RoTxn) -> heed::Result<Option<&'t [u8]>> {
+        self.main.get::<_, Str, ByteSlice>(rtxn, DISPLAYED_FIELDS_KEY)
+    }
+
+    /* words fst */
+
     /// Writes the FST which is the words dictionnary of the engine.
-    pub fn put_words_fst<A: AsRef<[u8]>>(&self, wtxn: &mut heed::RwTxn, fst: &fst::Set<A>) -> heed::Result<()> {
+    pub fn put_words_fst<A: AsRef<[u8]>>(&self, wtxn: &mut RwTxn, fst: &fst::Set<A>) -> heed::Result<()> {
         self.main.put::<_, Str, ByteSlice>(wtxn, WORDS_FST_KEY, fst.as_fst().as_bytes())
     }
 
     /// Returns the FST which is the words dictionnary of the engine.
-    pub fn words_fst<'t>(&self, rtxn: &'t heed::RoTxn) -> anyhow::Result<fst::Set<Cow<'t, [u8]>>> {
+    pub fn words_fst<'t>(&self, rtxn: &'t RoTxn) -> anyhow::Result<fst::Set<Cow<'t, [u8]>>> {
         match self.main.get::<_, Str, ByteSlice>(rtxn, WORDS_FST_KEY)? {
             Some(bytes) => Ok(fst::Set::new(bytes)?.map_data(Cow::Borrowed)?),
             None => Ok(fst::Set::default().map_data(Cow::Owned)?),
@@ -142,7 +173,7 @@ impl Index {
     /// Returns a [`Vec`] of the requested documents. Returns an error if a document is missing.
     pub fn documents<'t>(
         &self,
-        rtxn: &'t heed::RoTxn,
+        rtxn: &'t RoTxn,
         ids: impl IntoIterator<Item=DocumentId>,
     ) -> anyhow::Result<Vec<(DocumentId, obkv::KvReader<'t>)>>
     {
@@ -158,11 +189,11 @@ impl Index {
     }
 
     /// Returns the number of documents indexed in the database.
-    pub fn number_of_documents(&self, rtxn: &heed::RoTxn) -> anyhow::Result<usize> {
+    pub fn number_of_documents(&self, rtxn: &RoTxn) -> anyhow::Result<usize> {
         Ok(self.documents_ids(rtxn).map(|docids| docids.len() as usize)?)
     }
 
-    pub fn search<'a>(&'a self, rtxn: &'a heed::RoTxn) -> Search<'a> {
+    pub fn search<'a>(&'a self, rtxn: &'a RoTxn) -> Search<'a> {
         Search::new(rtxn, self)
     }
 }

src/subcommand/serve.rs

@@ -1,11 +1,12 @@
+use std::borrow::Cow;
 use std::collections::HashSet;
 use std::fs::{File, create_dir_all};
-use std::{mem, io};
 use std::net::SocketAddr;
 use std::path::PathBuf;
 use std::str::FromStr;
 use std::sync::Arc;
 use std::time::Instant;
+use std::{mem, io};
 
 use askama_warp::Template;
 use flate2::read::GzDecoder;
@@ -14,7 +15,7 @@ use futures::{FutureExt, StreamExt};
 use grenad::CompressionType;
 use heed::EnvOpenOptions;
 use indexmap::IndexMap;
-use serde::{Serialize, Deserialize};
+use serde::{Serialize, Deserialize, Deserializer};
 use structopt::StructOpt;
 use tokio::fs::File as TFile;
 use tokio::io::AsyncWriteExt;
@@ -159,6 +160,7 @@ enum UpdateStatus<M, P, N> {
 enum UpdateMeta {
     DocumentsAddition { method: String, format: String },
     ClearDocuments,
+    Settings(Settings),
 }
 
 #[derive(Debug, Clone, Serialize, Deserialize)]
@@ -170,6 +172,24 @@ enum UpdateMetaProgress {
     },
 }
 
+#[derive(Debug, Clone, Serialize, Deserialize)]
+struct Settings {
+    #[serde(
+        default,
+        deserialize_with = "deserialize_some",
+        skip_serializing_if = "Option::is_none",
+    )]
+    displayed_attributes: Option<Option<Vec<String>>>,
+}
+
+// Any value that is present is considered Some value, including null.
+fn deserialize_some<'de, T, D>(deserializer: D) -> Result<Option<T>, D::Error>
+where T: Deserialize<'de>,
+      D: Deserializer<'de>
+{
+    Deserialize::deserialize(deserializer).map(Some)
+}
+
 pub fn run(opt: Opt) -> anyhow::Result<()> {
     stderrlog::new()
         .verbosity(opt.verbose)
@@ -263,6 +283,24 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
                     let mut wtxn = index_cloned.write_txn()?;
                     let builder = update_builder.clear_documents(&mut wtxn, &index_cloned);
 
+                    match builder.execute() {
+                        Ok(_count) => wtxn.commit().map_err(Into::into),
+                        Err(e) => Err(e.into())
+                    }
+                },
+                UpdateMeta::Settings(settings) => {
+                    // We must use the write transaction of the update here.
+                    let mut wtxn = index_cloned.write_txn()?;
+                    let mut builder = update_builder.settings(&mut wtxn, &index_cloned);
+
+                    // We transpose the settings JSON struct into a real setting update.
+                    if let Some(names) = settings.displayed_attributes {
+                        match names {
+                            Some(names) => builder.set_displayed_fields(names),
+                            None => builder.reset_displayed_fields(),
+                        }
+                    }
+
                     match builder.execute() {
                         Ok(_count) => wtxn.commit().map_err(Into::into),
                         Err(e) => Err(e.into())
@@ -440,9 +478,14 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
 
             let mut documents = Vec::new();
             let fields_ids_map = index.fields_ids_map(&rtxn).unwrap();
+            let displayed_fields = match index.displayed_fields(&rtxn).unwrap() {
+                Some(fields) => Cow::Borrowed(fields),
+                None => Cow::Owned(fields_ids_map.iter().map(|(id, _)| id).collect()),
+            };
 
             for (_id, record) in index.documents(&rtxn, documents_ids).unwrap() {
-                let mut record = record.iter()
+                let mut record = displayed_fields.iter()
+                    .flat_map(|&id| record.get(id).map(|val| (id, val)))
                     .map(|(key_id, value)| {
                         let key = fields_ids_map.name(key_id).unwrap().to_owned();
                         // TODO we must deserialize a Json Value and highlight it.
@@ -510,7 +553,7 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
 
     let update_store_cloned = update_store.clone();
     let update_status_sender_cloned = update_status_sender.clone();
-    let indexing_route_csv = warp::filters::method::post()
+    let indexing_csv_route = warp::filters::method::post()
         .and(warp::path!("documents"))
         .and(warp::header::exact_ignore_case("content-type", "text/csv"))
         .and(warp::filters::query::query())
@@ -527,7 +570,7 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
 
     let update_store_cloned = update_store.clone();
     let update_status_sender_cloned = update_status_sender.clone();
-    let indexing_route_json = warp::filters::method::post()
+    let indexing_json_route = warp::filters::method::post()
         .and(warp::path!("documents"))
         .and(warp::header::exact_ignore_case("content-type", "application/json"))
         .and(warp::filters::query::query())
@@ -544,7 +587,7 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
 
     let update_store_cloned = update_store.clone();
     let update_status_sender_cloned = update_status_sender.clone();
-    let indexing_route_json_stream = warp::filters::method::post()
+    let indexing_json_stream_route = warp::filters::method::post()
         .and(warp::path!("documents"))
         .and(warp::header::exact_ignore_case("content-type", "application/x-ndjson"))
         .and(warp::filters::query::query())
@@ -559,12 +602,26 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
             )
         });
 
+    let update_store_cloned = update_store.clone();
     let update_status_sender_cloned = update_status_sender.clone();
     let clearing_route = warp::filters::method::post()
         .and(warp::path!("clear-documents"))
         .map(move || {
             let meta = UpdateMeta::ClearDocuments;
-            let update_id = update_store.register_update(&meta, &[]).unwrap();
+            let update_id = update_store_cloned.register_update(&meta, &[]).unwrap();
+            let _ = update_status_sender_cloned.send(UpdateStatus::Pending { update_id, meta });
+            eprintln!("update {} registered", update_id);
+            Ok(warp::reply())
+        });
+
+    let update_store_cloned = update_store.clone();
+    let update_status_sender_cloned = update_status_sender.clone();
+    let change_settings_route = warp::filters::method::post()
+        .and(warp::path!("settings"))
+        .and(warp::body::json())
+        .map(move |settings: Settings| {
+            let meta = UpdateMeta::Settings(settings);
+            let update_id = update_store_cloned.register_update(&meta, &[]).unwrap();
             let _ = update_status_sender_cloned.send(UpdateStatus::Pending { update_id, meta });
             eprintln!("update {} registered", update_id);
             Ok(warp::reply())
@@ -612,10 +669,11 @@ pub fn run(opt: Opt) -> anyhow::Result<()> {
         .or(dash_logo_white_route)
         .or(dash_logo_black_route)
         .or(query_route)
-        .or(indexing_route_csv)
+        .or(indexing_csv_route)
-        .or(indexing_route_json)
+        .or(indexing_json_route)
-        .or(indexing_route_json_stream)
+        .or(indexing_json_stream_route)
         .or(clearing_route)
+        .or(change_settings_route)
         .or(update_ws_route);
 
     let addr = SocketAddr::from_str(&opt.http_listen_addr)?;

src/update/mod.rs

@@ -2,6 +2,7 @@ mod available_documents_ids;
 mod clear_documents;
 mod delete_documents;
 mod index_documents;
+mod settings;
 mod update_builder;
 mod update_store;
 
@@ -9,5 +10,6 @@ pub use self::available_documents_ids::AvailableDocumentsIds;
 pub use self::clear_documents::ClearDocuments;
 pub use self::delete_documents::DeleteDocuments;
 pub use self::index_documents::{IndexDocuments, IndexDocumentsMethod, UpdateFormat};
+pub use self::settings::Settings;
 pub use self::update_builder::UpdateBuilder;
 pub use self::update_store::UpdateStore;

src/update/settings.rs

@@ -0,0 +1,126 @@
+use anyhow::Context;
+use crate::Index;
+
+pub struct Settings<'t, 'u, 'i> {
+    wtxn: &'t mut heed::RwTxn<'i, 'u>,
+    index: &'i Index,
+    // If the field is set to `None` it means that it hasn't been set by the user,
+    // however if it is `Some(None)` it means that the user forced a reset of the setting.
+    displayed_fields: Option<Option<Vec<String>>>,
+}
+
+impl<'t, 'u, 'i> Settings<'t, 'u, 'i> {
+    pub fn new(wtxn: &'t mut heed::RwTxn<'i, 'u>, index: &'i Index) -> Settings<'t, 'u, 'i> {
+        Settings { wtxn, index, displayed_fields: None }
+    }
+
+    pub fn reset_displayed_fields(&mut self) {
+        self.displayed_fields = Some(None);
+    }
+
+    pub fn set_displayed_fields(&mut self, names: Vec<String>) {
+        self.displayed_fields = Some(Some(names));
+    }
+
+    pub fn execute(self) -> anyhow::Result<()> {
+        // Check that the displayed attributes parameters has been specified.
+        if let Some(value) = self.displayed_fields {
+            match value {
+                // If it has been set, and it was a list of fields names, we create
+                // or generate the fields ids corresponds to those names and store them
+                // in the database in the order they were specified.
+                Some(fields_names) => {
+                    let mut fields_ids_map = self.index.fields_ids_map(self.wtxn)?;
+
+                    // We create or generate the fields ids corresponding to those names.
+                    let mut fields_ids = Vec::new();
+                    for name in fields_names {
+                        let id = fields_ids_map.insert(&name).context("field id limit reached")?;
+                        fields_ids.push(id);
+                    }
+
+                    self.index.put_displayed_fields(self.wtxn, &fields_ids)?;
+                },
+                // If it was set to `null` it means that the user wants to get the default behavior
+                // which is displaying all the attributes in no specific order (FieldsIdsMap order),
+                // we just have to delete the displayed fields.
+                None => {
+                    self.index.delete_displayed_fields(self.wtxn)?;
+                },
+            }
+        }
+
+        Ok(())
+    }
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+    use crate::update::{IndexDocuments, UpdateFormat};
+    use heed::EnvOpenOptions;
+
+    #[test]
+    fn default_displayed_fields() {
+        let path = tempfile::tempdir().unwrap();
+        let mut options = EnvOpenOptions::new();
+        options.map_size(10 * 1024 * 1024); // 10 MB
+        let index = Index::new(options, &path).unwrap();
+
+        // First we send 3 documents with ids from 1 to 3.
+        let mut wtxn = index.write_txn().unwrap();
+        let content = &b"name,age\nkevin,23\nkevina,21\nbenoit,34\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set to `None` (default value).
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids, None);
+        drop(rtxn);
+    }
+
+    #[test]
+    fn set_and_reset_displayed_field() {
+        let path = tempfile::tempdir().unwrap();
+        let mut options = EnvOpenOptions::new();
+        options.map_size(10 * 1024 * 1024); // 10 MB
+        let index = Index::new(options, &path).unwrap();
+
+        // First we send 3 documents with ids from 1 to 3.
+        let mut wtxn = index.write_txn().unwrap();
+        let content = &b"name,age\nkevin,23\nkevina,21\nbenoit,34\n"[..];
+        let mut builder = IndexDocuments::new(&mut wtxn, &index);
+        builder.update_format(UpdateFormat::Csv);
+        builder.execute(content, |_, _| ()).unwrap();
+
+        // In the same transaction we change the displayed fields to be only the age.
+        let mut builder = Settings::new(&mut wtxn, &index);
+        builder.set_displayed_fields(vec!["age".into()]);
+        builder.execute().unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set to only the "age" field.
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids_map = index.fields_ids_map(&rtxn).unwrap();
+        let age_field_id = fields_ids_map.id("age").unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids.unwrap(), &[age_field_id][..]);
+        drop(rtxn);
+
+        // We reset the fields ids to become `None`, the default value.
+        let mut wtxn = index.write_txn().unwrap();
+        let mut builder = Settings::new(&mut wtxn, &index);
+        builder.reset_displayed_fields();
+        builder.execute().unwrap();
+        wtxn.commit().unwrap();
+
+        // Check that the displayed fields are correctly set to `None` (default value).
+        let rtxn = index.read_txn().unwrap();
+        let fields_ids = index.displayed_fields(&rtxn).unwrap();
+        assert_eq!(fields_ids, None);
+        drop(rtxn);
+    }
+}

src/update/update_builder.rs

@@ -1,9 +1,7 @@
 use grenad::CompressionType;
 
 use crate::Index;
-use super::clear_documents::ClearDocuments;
+use super::{ClearDocuments, DeleteDocuments, IndexDocuments, Settings};
-use super::delete_documents::DeleteDocuments;
-use super::index_documents::IndexDocuments;
 
 pub struct UpdateBuilder {
     pub(crate) log_every_n: Option<usize>,
@@ -99,4 +97,13 @@ impl UpdateBuilder {
 
         builder
     }
+
+    pub fn settings<'t, 'u, 'i>(
+        self,
+        wtxn: &'t mut heed::RwTxn<'i, 'u>,
+        index: &'i Index,
+    ) -> Settings<'t, 'u, 'i>
+    {
+        Settings::new(wtxn, index)
+    }
 }