From 6b14b20369fad16c220d3686f9686a888f965463 Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Mon, 13 Jul 2020 17:50:16 +0200
Subject: [PATCH] Introduce a method to retrieve the number of attributes of
 the documents

---
 src/lib.rs | 15 +++++++++++++--
 1 file changed, 13 insertions(+), 2 deletions(-)

diff --git a/src/lib.rs b/src/lib.rs
index fad74976f..61956891a 100644
--- a/src/lib.rs
+++ b/src/lib.rs
@@ -74,6 +74,17 @@ impl Index {
         self.main.get::<_, Str, ByteSlice>(rtxn, "headers")
     }
 
+    pub fn number_of_attributes<'t>(&self, rtxn: &'t heed::RoTxn) -> anyhow::Result<Option<usize>> {
+        match self.headers(rtxn)? {
+            Some(headers) => {
+                let mut rdr = csv::Reader::from_reader(headers);
+                let headers = rdr.headers()?;
+                Ok(Some(headers.len()))
+            }
+            None => Ok(None),
+        }
+    }
+
     pub fn put_fst<A: AsRef<[u8]>>(&self, wtxn: &mut heed::RwTxn, fst: &fst::Set<A>) -> anyhow::Result<()> {
         Ok(self.main.put::<_, Str, ByteSlice>(wtxn, "words-fst", fst.as_fst().as_bytes())?)
     }
@@ -146,9 +157,9 @@ impl Index {
         }
 
         let mut words_attributes_docids = Vec::new();
-        let number_attributes: u32 = 6;
+        let number_of_attributes = self.number_of_attributes(rtxn)?.map_or(0, |n| n as u32);
 
-        for i in 0..number_attributes {
+        for i in 0..number_of_attributes {
             let mut intersect_docids: Option<RoaringBitmap> = None;
             for derived_words in &words {
                 let mut union_docids = RoaringBitmap::new();