From 71602e0f1b4a125207ab823f10b15b26c284ad5e Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Tue, 17 Aug 2021 15:57:40 +0200
Subject: [PATCH] Add the sortable fields into the settings and in the index

---
 milli/src/index.rs           | 36 +++++++++++++++++++++++++++++++++++-
 milli/src/update/settings.rs | 24 ++++++++++++++++++++++++
 2 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/milli/src/index.rs b/milli/src/index.rs
index 120bcbadf..e2ab51a1c 100644
--- a/milli/src/index.rs
+++ b/milli/src/index.rs
@@ -28,6 +28,7 @@ pub mod main_key {
     pub const DISTINCT_FIELD_KEY: &str = "distinct-field-key";
     pub const DOCUMENTS_IDS_KEY: &str = "documents-ids";
     pub const FILTERABLE_FIELDS_KEY: &str = "filterable-fields";
+    pub const SORTABLE_FIELDS_KEY: &str = "sortable-fields";
     pub const FIELD_DISTRIBUTION_KEY: &str = "fields-distribution";
     pub const FIELDS_IDS_MAP_KEY: &str = "fields-ids-map";
     pub const HARD_EXTERNAL_DOCUMENTS_IDS_KEY: &str = "hard-external-documents-ids";
@@ -446,13 +447,45 @@ impl Index {
         Ok(fields_ids)
     }
 
+    /* sortable fields */
+
+    /// Writes the sortable fields names in the database.
+    pub(crate) fn put_sortable_fields(
+        &self,
+        wtxn: &mut RwTxn,
+        fields: &HashSet<String>,
+    ) -> heed::Result<()> {
+        self.main.put::<_, Str, SerdeJson<_>>(wtxn, main_key::SORTABLE_FIELDS_KEY, fields)
+    }
+
+    /// Deletes the sortable fields ids in the database.
+    pub(crate) fn delete_sortable_fields(&self, wtxn: &mut RwTxn) -> heed::Result<bool> {
+        self.main.delete::<_, Str>(wtxn, main_key::SORTABLE_FIELDS_KEY)
+    }
+
+    /// Returns the sortable fields names.
+    pub fn sortable_fields(&self, rtxn: &RoTxn) -> heed::Result<HashSet<String>> {
+        Ok(self
+            .main
+            .get::<_, Str, SerdeJson<_>>(rtxn, main_key::SORTABLE_FIELDS_KEY)?
+            .unwrap_or_default())
+    }
+
+    /// Identical to `sortable_fields`, but returns ids instead.
+    pub fn sortable_fields_ids(&self, rtxn: &RoTxn) -> Result<HashSet<FieldId>> {
+        let fields = self.sortable_fields(rtxn)?;
+        let fields_ids_map = self.fields_ids_map(rtxn)?;
+        Ok(fields.into_iter().filter_map(|name| fields_ids_map.id(&name)).collect())
+    }
+
     /* faceted documents ids */
 
     /// Returns the faceted fields names.
     ///
-    /// Faceted fields are the union of all the filterable, distinct, and Asc/Desc fields.
+    /// Faceted fields are the union of all the filterable, sortable, distinct, and Asc/Desc fields.
     pub fn faceted_fields(&self, rtxn: &RoTxn) -> Result<HashSet<String>> {
         let filterable_fields = self.filterable_fields(rtxn)?;
+        let sortable_fields = self.sortable_fields(rtxn)?;
         let distinct_field = self.distinct_field(rtxn)?;
         let asc_desc_fields =
             self.criteria(rtxn)?.into_iter().filter_map(|criterion| match criterion {
@@ -461,6 +494,7 @@ impl Index {
             });
 
         let mut faceted_fields = filterable_fields;
+        faceted_fields.extend(sortable_fields);
         faceted_fields.extend(asc_desc_fields);
         if let Some(field) = distinct_field {
             faceted_fields.insert(field.to_owned());
diff --git a/milli/src/update/settings.rs b/milli/src/update/settings.rs
index 156e566a5..c0b5e4549 100644
--- a/milli/src/update/settings.rs
+++ b/milli/src/update/settings.rs
@@ -75,6 +75,7 @@ pub struct Settings<'a, 't, 'u, 'i> {
     searchable_fields: Setting<Vec<String>>,
     displayed_fields: Setting<Vec<String>>,
     filterable_fields: Setting<HashSet<String>>,
+    sortable_fields: Setting<HashSet<String>>,
     criteria: Setting<Vec<String>>,
     stop_words: Setting<BTreeSet<String>>,
     distinct_field: Setting<String>,
@@ -102,6 +103,7 @@ impl<'a, 't, 'u, 'i> Settings<'a, 't, 'u, 'i> {
             searchable_fields: Setting::NotSet,
             displayed_fields: Setting::NotSet,
             filterable_fields: Setting::NotSet,
+            sortable_fields: Setting::NotSet,
             criteria: Setting::NotSet,
             stop_words: Setting::NotSet,
             distinct_field: Setting::NotSet,
@@ -135,6 +137,10 @@ impl<'a, 't, 'u, 'i> Settings<'a, 't, 'u, 'i> {
         self.filterable_fields = Setting::Set(names);
     }
 
+    pub fn set_sortable_fields(&mut self, names: HashSet<String>) {
+        self.sortable_fields = Setting::Set(names);
+    }
+
     pub fn reset_criteria(&mut self) {
         self.criteria = Setting::Reset;
     }
@@ -392,6 +398,23 @@ impl<'a, 't, 'u, 'i> Settings<'a, 't, 'u, 'i> {
         Ok(())
     }
 
+    fn update_sortable(&mut self) -> Result<()> {
+        match self.sortable_fields {
+            Setting::Set(ref fields) => {
+                let mut new_fields = HashSet::new();
+                for name in fields {
+                    new_fields.insert(name.clone());
+                }
+                self.index.put_sortable_fields(self.wtxn, &new_fields)?;
+            }
+            Setting::Reset => {
+                self.index.delete_sortable_fields(self.wtxn)?;
+            }
+            Setting::NotSet => (),
+        }
+        Ok(())
+    }
+
     fn update_criteria(&mut self) -> Result<()> {
         match self.criteria {
             Setting::Set(ref fields) => {
@@ -446,6 +469,7 @@ impl<'a, 't, 'u, 'i> Settings<'a, 't, 'u, 'i> {
 
         self.update_displayed()?;
         self.update_filterable()?;
+        self.update_sortable()?;
         self.update_distinct_field()?;
         self.update_criteria()?;
         self.update_primary_key()?;