Apply pin rules to all collection indexes with correct locale

Cargo.toml

@@ -54,7 +54,7 @@ strum_macros = "0.26.4"
 tracing = "0.1"
 tracing-subscriber = { version = "0.3", features = ["env-filter", "json"] }
 
-oramacore_lib = { version = "0.1.7" }
+oramacore_lib = { version = "0.1.8" }
 
 # JS
 oxc_parser = "0.47.1"

src/collection_manager/sides/read/index/mod.rs

@@ -339,6 +339,10 @@
         })
     }
 
+    pub fn get_text_parser(&self) -> &TextParser {
+        self.text_parser.as_ref()
+    }
+
     pub async fn get_search_store<'index>(&'index self) -> IndexSearchStore<'index> {
         let (committed_fields, uncommitted_fields) = tokio::join!(
             self.committed_fields.read("get_all_document_ids"),
@@ -1122,7 +1126,7 @@
     string_filter_field_ids: Vec<(FieldId, StringFilterFieldInfo)>,
     string_field_ids: Vec<(FieldId, StringFieldInfo)>,
     vector_field_ids: Vec<(FieldId, VectorFieldInfo)>,
     path_to_index_id_map: Vec<(Box<[String]>, (FieldId, FieldType))>,
     created_at: DateTime<Utc>,
     updated_at: DateTime<Utc>,
     #[serde(default)]

src/collection_manager/sides/read/search.rs

@@ -105,8 +105,8 @@ impl<'collection, 'analytics_storage> Search<'collection, 'analytics_storage> {
         let pin_rule_consequences = extract_pin_rules(
             &*collection.get_pin_rules_reader("search").await,
             &search_params,
-        )
-        .await;
+            &indexes,
+        );
 
         let has_pin_rules = !pin_rule_consequences.is_empty();
 
@@ -239,12 +239,30 @@ pub fn extract_term_from_search_mode(search_mode: &SearchMode) -> &str {
     }
 }
 
-async fn extract_pin_rules(
+fn extract_pin_rules<'collection>(
     rules: &PinRulesReader<DocumentId>,
     search_params: &SearchParams,
+    indexes: &ReadIndexesLockGuard<'collection>,
 ) -> Vec<Consequence<DocumentId>> {
     let term = extract_term_from_search_mode(&search_params.mode);
-    rules.apply(term)
+    let mut consequences: Vec<_> = indexes
+        .iter()
+        .flat_map(|index| {
+            let text_parser = index.get_text_parser();
+            rules.apply(term, text_parser)
+        })
+        .collect();
+
+    consequences.sort_by(|a, b| {
+        a.promote
+            .iter()
+            .map(|item| (item.position, item.doc_id))
+            .cmp(b.promote.iter().map(|item| (item.position, item.doc_id)))
+    });
+
+    consequences.dedup();
+
+    consequences
 }
 
 async fn search_on_indexes(

src/collection_manager/sides/read/sort.rs

@@ -279,20 +279,20 @@ const MAX_PROMOTED_ITEMS: usize = 10_000;
 
 /// Internal shared implementation for applying pin rules with optional document filtering
 fn apply_pin_rules_internal<F>(
-    pins: &[Consequence<DocumentId>],
+    pin_rules: &[Consequence<DocumentId>],
     token_scores: &HashMap<DocumentId, f32>,
     mut top: Vec<TokenScore>,
     document_filter: Option<F>,
 ) -> Vec<TokenScore>
 where
     F: Fn(&DocumentId) -> bool,
 {
-    if pins.is_empty() {
+    if pin_rules.is_empty() {
         return top;
     }
 
     // Estimate the total number of promote items to avoid excessive memory allocation
-    let estimated_promote_items: usize = pins.iter().map(|c| c.promote.len()).sum();
+    let estimated_promote_items: usize = pin_rules.iter().map(|c| c.promote.len()).sum();
     if estimated_promote_items > MAX_PROMOTED_ITEMS {
         // Log warning but continue with a truncated set
         warn!(
@@ -304,7 +304,7 @@ where
 
     let mut promote_items = Vec::with_capacity(estimated_promote_items.min(MAX_PROMOTED_ITEMS));
     if let Some(ref filter) = document_filter {
-        for consequence in pins {
+        for consequence in pin_rules {
             promote_items.extend(
                 consequence
                     .promote
@@ -314,7 +314,7 @@ where
             );
         }
     } else {
-        for consequence in pins {
+        for consequence in pin_rules {
             promote_items.extend(consequence.promote.iter().cloned());
         }
     }
@@ -374,12 +374,12 @@ fn apply_pin_rules(
 
 fn apply_pin_rules_to_group(
     all_document_ids: HashSet<DocumentId>,
-    pins: &[Consequence<DocumentId>],
+    pin_rules: &[Consequence<DocumentId>],
     token_scores: &HashMap<DocumentId, f32>,
     top: Vec<TokenScore>,
 ) -> Vec<TokenScore> {
     apply_pin_rules_internal(
-        pins,
+        pin_rules,
         token_scores,
         top,
         Some(|doc_id: &DocumentId| all_document_ids.contains(doc_id)),

src/collection_manager/sides/write/index/fields.rs

@@ -781,7 +781,7 @@ impl StringScoreField {
                 }
             };
 
-            for stemmed in stemmeds {
+            if let Some(stemmed) = stemmeds {
                 let stemmed = Term(stemmed);
                 match terms.entry(stemmed) {
                     Entry::Occupied(mut entry) => {

src/tests/pin_rules.rs

@@ -16,6 +16,7 @@ async fn test_pin_rules_after_insert_simple() {
             json!({
                 "id": format!("{}", i),
                 "c": format!("c-{}", i),
+                "run": format!("run-{}", i),
             })
         })
         .collect();
@@ -32,6 +33,11 @@ async fn test_pin_rules_after_insert_simple() {
                     {
                         "pattern": "c",
                         "anchoring": "is"
+                    },
+                    {
+                        "pattern": "running",
+                        "anchoring": "is",
+                        "normalization": "stem",
                     }
                 ],
                 "consequence": {
@@ -80,6 +86,153 @@ async fn test_pin_rules_after_insert_simple() {
     let ids = extrapolate_ids_from_result(&result);
     assert_eq!(&ids, &["0", "5", "7", "1", "2", "3", "4", "6", "8", "9"]);
 
+    let stemmed_result = collection_client
+        .search(
+            json!({
+                "term": "runs"
+            })
+            .try_into()
+            .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    let stemmed_ids = extrapolate_ids_from_result(&stemmed_result);
+    assert_eq!(&stemmed_ids, &ids);
+
+    drop(test_context);
+}
+
+#[tokio::test(flavor = "multi_thread")]
+async fn test_pin_rules_multiple_indexes() {
+    init_log();
+
+    let test_context = TestContext::new().await;
+    let collection_client = test_context.create_collection().await.unwrap();
+    let index_client1 = collection_client.create_index().await.unwrap();
+    let index_client2 = collection_client.create_index().await.unwrap();
+
+    let docs: Vec<_> = (0_u8..10_u8)
+        .map(|i| {
+            json!({
+                "id": format!("{}", i),
+                "c": format!("c-{}", i),
+            })
+        })
+        .collect();
+    index_client1
+        .insert_documents(docs.try_into().unwrap())
+        .await
+        .unwrap();
+
+    let docs: Vec<_> = (10_u8..20_u8)
+        .map(|i| {
+            json!({
+                "id": format!("{}", i),
+                "c": format!("c-{}", i),
+            })
+        })
+        .collect();
+    index_client2
+        .insert_documents(docs.try_into().unwrap())
+        .await
+        .unwrap();
+
+    index_client1
+        .insert_pin_rules(
+            json!({
+                "id": "rule-1",
+                "conditions": [
+                    {
+                        "pattern": "c",
+                        "anchoring": "is"
+                    },
+                ],
+                "consequence": {
+                    "promote": [
+                        {
+                            "doc_id": "5",
+                            "position": 1
+                        },
+                        {
+                            "doc_id": "7",
+                            "position": 4
+                        }
+                    ]
+                }
+            })
+            .try_into()
+            .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    index_client2
+        .insert_pin_rules(
+            json!({
+                "id": "rule-2",
+                "conditions": [
+                    {
+                        "pattern": "c",
+                        "anchoring": "startsWith",
+                    }
+                ],
+                "consequence": {
+                    "promote": [
+                        {
+                            "doc_id": "11",
+                            "position": 2
+                        },
+                        {
+                            "doc_id": "15",
+                            "position": 3
+                        }
+                    ]
+                }
+            })
+            .try_into()
+            .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    let result = collection_client
+        .search(
+            json!({
+                "term": "c"
+            })
+            .try_into()
+            .unwrap(),
+        )
+        .await
+        .unwrap();
+
+    // The pin rule should promote document with id "5" to position 1
+    // The actual ID will be in format "index_id:5"
+    assert!(
+        result.hits[1].id.ends_with(":5"),
+        "Expected document ID ending with ':5', got: {}",
+        result.hits[1].id
+    );
+    assert!(
+        result.hits[2].id.ends_with(":11"),
+        "Expected document ID ending with ':11', got: {}",
+        result.hits[2].id
+    );
+    assert!(
+        result.hits[3].id.ends_with(":15"),
+        "Expected document ID ending with ':15', got: {}",
+        result.hits[3].id
+    );
+    assert!(
+        result.hits[4].id.ends_with(":7"),
+        "Expected document ID ending with ':7', got: {}",
+        result.hits[4].id
+    );
+
+    let ids = extrapolate_ids_from_result(&result);
+    assert_eq!(&ids, &["0", "5", "11", "15", "7", "1", "2", "3", "4", "6"]);
+
     drop(test_context);
 }