From 5230ddb3ea7b2d86bcbca5df2822969351077fcc Mon Sep 17 00:00:00 2001
From: Kerollmops <clement@meilisearch.com>
Date: Wed, 12 Apr 2023 11:52:56 +0200
Subject: [PATCH] Resolve the attribute ranking rule conditions

---
 .../new/ranking_rule_graph/attribute/mod.rs   | 10 ++++--
 milli/src/search/new/resolve_query_graph.rs   | 36 +++++++++++++++++++
 2 files changed, 44 insertions(+), 2 deletions(-)

diff --git a/milli/src/search/new/ranking_rule_graph/attribute/mod.rs b/milli/src/search/new/ranking_rule_graph/attribute/mod.rs
index 2b25adc7e..4ca0b7dc4 100644
--- a/milli/src/search/new/ranking_rule_graph/attribute/mod.rs
+++ b/milli/src/search/new/ranking_rule_graph/attribute/mod.rs
@@ -4,7 +4,9 @@ use roaring::RoaringBitmap;
 use super::{ComputedCondition, RankingRuleGraphTrait};
 use crate::search::new::interner::{DedupInterner, Interned};
 use crate::search::new::query_term::LocatedQueryTermSubset;
-use crate::search::new::resolve_query_graph::compute_query_term_subset_docids;
+use crate::search::new::resolve_query_graph::{
+    compute_query_term_subset_docids, compute_query_term_subset_docids_within_field_id,
+};
 use crate::search::new::SearchContext;
 use crate::Result;
 
@@ -26,7 +28,11 @@ impl RankingRuleGraphTrait for AttributeGraph {
     ) -> Result<ComputedCondition> {
         let AttributeCondition { term, .. } = condition;
         // maybe compute_query_term_subset_docids should accept a universe as argument
-        let mut docids = compute_query_term_subset_docids(ctx, &term.term_subset)?;
+        let mut docids = compute_query_term_subset_docids_within_field_id(
+            ctx,
+            &term.term_subset,
+            condition.fid,
+        )?;
         docids &= universe;
 
         Ok(ComputedCondition {
diff --git a/milli/src/search/new/resolve_query_graph.rs b/milli/src/search/new/resolve_query_graph.rs
index f4938ca12..711497009 100644
--- a/milli/src/search/new/resolve_query_graph.rs
+++ b/milli/src/search/new/resolve_query_graph.rs
@@ -33,6 +33,8 @@ pub fn compute_query_term_subset_docids(
     ctx: &mut SearchContext,
     term: &QueryTermSubset,
 ) -> Result<RoaringBitmap> {
+    // TODO Use the roaring::MultiOps trait
+
     let mut docids = RoaringBitmap::new();
     for word in term.all_single_words_except_prefix_db(ctx)? {
         if let Some(word_docids) = ctx.word_docids(word)? {
@@ -52,6 +54,40 @@ pub fn compute_query_term_subset_docids(
     Ok(docids)
 }
 
+pub fn compute_query_term_subset_docids_within_field_id(
+    ctx: &mut SearchContext,
+    term: &QueryTermSubset,
+    fid: u16,
+) -> Result<RoaringBitmap> {
+    // TODO Use the roaring::MultiOps trait
+
+    let mut docids = RoaringBitmap::new();
+    for word in term.all_single_words_except_prefix_db(ctx)? {
+        if let Some(word_fid_docids) = ctx.get_db_word_fid_docids(word, fid)? {
+            docids |= CboRoaringBitmapCodec::bytes_decode(word_fid_docids)
+                .ok_or(heed::Error::Decoding)?;
+        }
+    }
+
+    for phrase in term.all_phrases(ctx)? {
+        for &word in phrase.words(ctx).iter().flatten() {
+            if let Some(word_fid_docids) = ctx.get_db_word_fid_docids(word, fid)? {
+                docids |= CboRoaringBitmapCodec::bytes_decode(word_fid_docids)
+                    .ok_or(heed::Error::Decoding)?;
+            }
+        }
+    }
+
+    if let Some(word_prefix) = term.use_prefix_db(ctx) {
+        if let Some(word_fid_docids) = ctx.get_db_word_prefix_fid_docids(word_prefix, fid)? {
+            docids |= CboRoaringBitmapCodec::bytes_decode(word_fid_docids)
+                .ok_or(heed::Error::Decoding)?;
+        }
+    }
+
+    Ok(docids)
+}
+
 pub fn compute_query_graph_docids(
     ctx: &mut SearchContext,
     q: &QueryGraph,