Search: spans with a finer granularity

meilisearch/src/search.rs

@@ -752,10 +752,15 @@ fn prepare_search<'t>(
         SearchKind::SemanticOnly { embedder_name, embedder } => {
             let vector = match query.vector.clone() {
                 Some(vector) => vector,
-                None => embedder
+                None => {
+                    let span = tracing::trace_span!(target: "search::vector", "embed_one");
+                    let _entered = span.enter();
+
+                    embedder
                         .embed_one(query.q.clone().unwrap())
                         .map_err(milli::vector::Error::from)
-                    .map_err(milli::Error::from)?,
+                        .map_err(milli::Error::from)?
+                }
             };
 
             search.semantic(embedder_name.clone(), embedder.clone(), Some(vector));

milli/src/search/hybrid.rs

@@ -17,6 +17,7 @@ struct ScoreWithRatioResult {
 
 type ScoreWithRatio = (Vec<ScoreDetails>, f32);
 
+#[tracing::instrument(level = "trace", skip_all, target = "search::hybrid")]
 fn compare_scores(
     &(ref left_scores, left_ratio): &ScoreWithRatio,
     &(ref right_scores, right_ratio): &ScoreWithRatio,
@@ -84,6 +85,7 @@ impl ScoreWithRatioResult {
         }
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::hybrid")]
     fn merge(
         vector_results: Self,
         keyword_results: Self,
@@ -150,6 +152,7 @@ impl ScoreWithRatioResult {
 }
 
 impl<'a> Search<'a> {
+    #[tracing::instrument(level = "trace", skip_all, target = "search::hybrid")]
     pub fn execute_hybrid(&self, semantic_ratio: f32) -> Result<(SearchResult, Option<u32>)> {
         // TODO: find classier way to achieve that than to reset vector and query params
         // create separate keyword and semantic searches
@@ -194,6 +197,9 @@ impl<'a> Search<'a> {
             Some(vector_query) => vector_query,
             None => {
                 // attempt to embed the vector
+                let span = tracing::trace_span!(target: "search::hybrid", "embed_one");
+                let _entered = span.enter();
+
                 match embedder.embed_one(query) {
                     Ok(embedding) => embedding,
                     Err(error) => {

milli/src/search/new/bucket_sort.rs

@@ -213,9 +213,6 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
             continue;
         }
 
-        let span = tracing::trace_span!(target: "search::bucket_sort", "next_bucket", id = ranking_rules[cur_ranking_rule_index].id());
-        let entered = span.enter();
-
         let Some(next_bucket) = ranking_rules[cur_ranking_rule_index].next_bucket(
             ctx,
             logger,
@@ -225,7 +222,6 @@ pub fn bucket_sort<'ctx, Q: RankingRuleQueryTrait>(
             back!();
             continue;
         };
-        drop(entered);
 
         ranking_rule_scores.push(next_bucket.score);
 

milli/src/search/new/exact_attribute.rs

@@ -27,6 +27,7 @@ impl<'ctx> RankingRule<'ctx, QueryGraph> for ExactAttribute {
         "exact_attribute".to_owned()
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::exact_attribute")]
     fn start_iteration(
         &mut self,
         ctx: &mut SearchContext<'ctx>,
@@ -38,6 +39,7 @@ impl<'ctx> RankingRule<'ctx, QueryGraph> for ExactAttribute {
         Ok(())
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::exact_attribute")]
     fn next_bucket(
         &mut self,
         _ctx: &mut SearchContext<'ctx>,
@@ -51,6 +53,7 @@ impl<'ctx> RankingRule<'ctx, QueryGraph> for ExactAttribute {
         Ok(output)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::exact_attribute")]
     fn end_iteration(
         &mut self,
         _ctx: &mut SearchContext<'ctx>,

milli/src/search/new/geo_sort.rs

@@ -209,6 +209,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for GeoSort<Q> {
         "geo_sort".to_owned()
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::geo_sort")]
     fn start_iteration(
         &mut self,
         ctx: &mut SearchContext<'ctx>,
@@ -234,6 +235,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for GeoSort<Q> {
         Ok(())
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::geo_sort")]
     #[allow(clippy::only_used_in_recursion)]
     fn next_bucket(
         &mut self,
@@ -285,6 +287,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for GeoSort<Q> {
         self.next_bucket(ctx, logger, universe)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::geo_sort")]
     fn end_iteration(&mut self, _ctx: &mut SearchContext<'ctx>, _logger: &mut dyn SearchLogger<Q>) {
         // we do not reset the rtree here, it could be used in a next iteration
         self.query = None;

milli/src/search/new/graph_based_ranking_rule.rs

@@ -127,6 +127,8 @@ impl<'ctx, G: RankingRuleGraphTrait> RankingRule<'ctx, QueryGraph> for GraphBase
     fn id(&self) -> String {
         self.id.clone()
     }
+
+    #[tracing::instrument(level = "trace", skip_all, target = "search::graph_based")]
     fn start_iteration(
         &mut self,
         ctx: &mut SearchContext<'ctx>,
@@ -209,6 +211,7 @@ impl<'ctx, G: RankingRuleGraphTrait> RankingRule<'ctx, QueryGraph> for GraphBase
         Ok(())
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::graph_based")]
     fn next_bucket(
         &mut self,
         ctx: &mut SearchContext<'ctx>,
@@ -358,6 +361,7 @@ impl<'ctx, G: RankingRuleGraphTrait> RankingRule<'ctx, QueryGraph> for GraphBase
         Ok(Some(RankingRuleOutput { query: next_query_graph, candidates: bucket, score }))
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::graph_based")]
     fn end_iteration(
         &mut self,
         _ctx: &mut SearchContext<'ctx>,

milli/src/search/new/mod.rs

@@ -212,7 +212,7 @@ fn resolve_maximally_reduced_query_graph(
     Ok(docids)
 }
 
-#[tracing::instrument(level = "trace", skip_all, target = "search")]
+#[tracing::instrument(level = "trace", skip_all, target = "search::universe")]
 fn resolve_universe(
     ctx: &mut SearchContext,
     initial_universe: &RoaringBitmap,
@@ -229,7 +229,7 @@ fn resolve_universe(
     )
 }
 
-#[tracing::instrument(level = "trace", skip_all, target = "search")]
+#[tracing::instrument(level = "trace", skip_all, target = "search::query")]
 fn resolve_negative_words(
     ctx: &mut SearchContext,
     negative_words: &[Word],
@@ -243,7 +243,7 @@ fn resolve_negative_words(
     Ok(negative_bitmap)
 }
 
-#[tracing::instrument(level = "trace", skip_all, target = "search")]
+#[tracing::instrument(level = "trace", skip_all, target = "search::query")]
 fn resolve_negative_phrases(
     ctx: &mut SearchContext,
     negative_phrases: &[LocatedQueryTerm],
@@ -548,7 +548,7 @@ fn resolve_sort_criteria<'ctx, Query: RankingRuleQueryTrait>(
     Ok(())
 }
 
-#[tracing::instrument(level = "trace", skip_all, target = "search")]
+#[tracing::instrument(level = "trace", skip_all, target = "search::universe")]
 pub fn filtered_universe(
     index: &Index,
     txn: &RoTxn<'_>,
@@ -620,7 +620,7 @@ pub fn execute_vector_search(
 }
 
 #[allow(clippy::too_many_arguments)]
-#[tracing::instrument(level = "trace", skip_all, target = "search")]
+#[tracing::instrument(level = "trace", skip_all, target = "search::main")]
 pub fn execute_search(
     ctx: &mut SearchContext,
     query: Option<&str>,

milli/src/search/new/ranking_rule_graph/exactness/mod.rs

@@ -44,6 +44,7 @@ fn compute_docids(
 impl RankingRuleGraphTrait for ExactnessGraph {
     type Condition = ExactnessCondition;
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::exactness")]
     fn resolve_condition(
         ctx: &mut SearchContext,
         condition: &Self::Condition,
@@ -71,6 +72,7 @@ impl RankingRuleGraphTrait for ExactnessGraph {
         })
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::exactness")]
     fn build_edges(
         _ctx: &mut SearchContext,
         conditions_interner: &mut DedupInterner<Self::Condition>,
@@ -86,6 +88,7 @@ impl RankingRuleGraphTrait for ExactnessGraph {
         Ok(vec![(0, exact_condition), (dest_node.term_ids.len() as u32, skip_condition)])
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::exactness")]
     fn rank_to_score(rank: Rank) -> ScoreDetails {
         ScoreDetails::ExactWords(score_details::ExactWords::from_rank(rank))
     }

milli/src/search/new/ranking_rule_graph/fid/mod.rs

@@ -20,6 +20,7 @@ pub enum FidGraph {}
 impl RankingRuleGraphTrait for FidGraph {
     type Condition = FidCondition;
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::fid")]
     fn resolve_condition(
         ctx: &mut SearchContext,
         condition: &Self::Condition,
@@ -44,6 +45,7 @@ impl RankingRuleGraphTrait for FidGraph {
         })
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::fid")]
     fn build_edges(
         ctx: &mut SearchContext,
         conditions_interner: &mut DedupInterner<Self::Condition>,
@@ -101,6 +103,7 @@ impl RankingRuleGraphTrait for FidGraph {
         Ok(edges)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::fid")]
     fn rank_to_score(rank: Rank) -> ScoreDetails {
         ScoreDetails::Fid(rank)
     }

milli/src/search/new/ranking_rule_graph/position/mod.rs

@@ -20,6 +20,7 @@ pub enum PositionGraph {}
 impl RankingRuleGraphTrait for PositionGraph {
     type Condition = PositionCondition;
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::position")]
     fn resolve_condition(
         ctx: &mut SearchContext,
         condition: &Self::Condition,
@@ -44,6 +45,7 @@ impl RankingRuleGraphTrait for PositionGraph {
         })
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::position")]
     fn build_edges(
         ctx: &mut SearchContext,
         conditions_interner: &mut DedupInterner<Self::Condition>,
@@ -117,6 +119,7 @@ impl RankingRuleGraphTrait for PositionGraph {
         Ok(edges)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::position")]
     fn rank_to_score(rank: Rank) -> ScoreDetails {
         ScoreDetails::Position(rank)
     }

milli/src/search/new/ranking_rule_graph/proximity/mod.rs

@@ -21,6 +21,7 @@ pub enum ProximityGraph {}
 impl RankingRuleGraphTrait for ProximityGraph {
     type Condition = ProximityCondition;
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::proximity")]
     fn resolve_condition(
         ctx: &mut SearchContext,
         condition: &Self::Condition,
@@ -29,6 +30,7 @@ impl RankingRuleGraphTrait for ProximityGraph {
         compute_docids::compute_docids(ctx, condition, universe)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::proximity")]
     fn build_edges(
         ctx: &mut SearchContext,
         conditions_interner: &mut DedupInterner<Self::Condition>,
@@ -38,6 +40,7 @@ impl RankingRuleGraphTrait for ProximityGraph {
         build::build_edges(ctx, conditions_interner, source_term, dest_term)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::proximity")]
     fn rank_to_score(rank: Rank) -> ScoreDetails {
         ScoreDetails::Proximity(rank)
     }

milli/src/search/new/ranking_rule_graph/typo/mod.rs

@@ -19,6 +19,7 @@ pub enum TypoGraph {}
 impl RankingRuleGraphTrait for TypoGraph {
     type Condition = TypoCondition;
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::typo")]
     fn resolve_condition(
         ctx: &mut SearchContext,
         condition: &Self::Condition,
@@ -37,6 +38,7 @@ impl RankingRuleGraphTrait for TypoGraph {
         })
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::typo")]
     fn build_edges(
         ctx: &mut SearchContext,
         conditions_interner: &mut DedupInterner<Self::Condition>,
@@ -77,6 +79,7 @@ impl RankingRuleGraphTrait for TypoGraph {
         Ok(edges)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::typo")]
     fn rank_to_score(rank: Rank) -> ScoreDetails {
         ScoreDetails::Typo(score_details::Typo::from_rank(rank))
     }

milli/src/search/new/ranking_rule_graph/words/mod.rs

@@ -18,6 +18,7 @@ pub enum WordsGraph {}
 impl RankingRuleGraphTrait for WordsGraph {
     type Condition = WordsCondition;
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::words")]
     fn resolve_condition(
         ctx: &mut SearchContext,
         condition: &Self::Condition,
@@ -36,6 +37,7 @@ impl RankingRuleGraphTrait for WordsGraph {
         })
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::words")]
     fn build_edges(
         _ctx: &mut SearchContext,
         conditions_interner: &mut DedupInterner<Self::Condition>,
@@ -45,6 +47,7 @@ impl RankingRuleGraphTrait for WordsGraph {
         Ok(vec![(0, conditions_interner.insert(WordsCondition { term: to_term.clone() }))])
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::words")]
     fn rank_to_score(rank: Rank) -> ScoreDetails {
         ScoreDetails::Words(score_details::Words::from_rank(rank))
     }

milli/src/search/new/sort.rs

@@ -88,6 +88,8 @@ impl<'ctx, Query: RankingRuleQueryTrait> RankingRule<'ctx, Query> for Sort<'ctx,
         let Self { field_name, is_ascending, .. } = self;
         format!("{field_name}:{}", if *is_ascending { "asc" } else { "desc" })
     }
+
+    #[tracing::instrument(level = "trace", skip_all, target = "search::sort")]
     fn start_iteration(
         &mut self,
         ctx: &mut SearchContext<'ctx>,
@@ -186,6 +188,7 @@ impl<'ctx, Query: RankingRuleQueryTrait> RankingRule<'ctx, Query> for Sort<'ctx,
         Ok(())
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::sort")]
     fn next_bucket(
         &mut self,
         _ctx: &mut SearchContext<'ctx>,
@@ -211,6 +214,7 @@ impl<'ctx, Query: RankingRuleQueryTrait> RankingRule<'ctx, Query> for Sort<'ctx,
         }
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::sort")]
     fn end_iteration(
         &mut self,
         _ctx: &mut SearchContext<'ctx>,

milli/src/search/new/vector_sort.rs

@@ -73,6 +73,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for VectorSort<Q> {
         "vector_sort".to_owned()
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::vector_sort")]
     fn start_iteration(
         &mut self,
         ctx: &mut SearchContext<'ctx>,
@@ -89,6 +90,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for VectorSort<Q> {
     }
 
     #[allow(clippy::only_used_in_recursion)]
+    #[tracing::instrument(level = "trace", skip_all, target = "search::vector_sort")]
     fn next_bucket(
         &mut self,
         ctx: &mut SearchContext<'ctx>,
@@ -139,6 +141,7 @@ impl<'ctx, Q: RankingRuleQueryTrait> RankingRule<'ctx, Q> for VectorSort<Q> {
         self.next_bucket(ctx, _logger, universe)
     }
 
+    #[tracing::instrument(level = "trace", skip_all, target = "search::vector_sort")]
     fn end_iteration(&mut self, _ctx: &mut SearchContext<'ctx>, _logger: &mut dyn SearchLogger<Q>) {
         self.query = None;
     }