Search: spans with a finer granularity

2024-11-25 14:24:26 +01:00 · 2024-07-02 12:14:40 +02:00 · 2024-07-02 12:14:40 +02:00 · 128e6c7502
commit 128e6c7502
parent 3129f96603
15 changed files with 55 additions and 13 deletions
--- a/meilisearch/src/search.rs
+++ b/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));
--- a/milli/src/search/hybrid.rs
+++ b/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) => {
--- a/milli/src/search/new/bucket_sort.rs
+++ b/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);

--- a/milli/src/search/new/exact_attribute.rs
+++ b/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>,
--- a/milli/src/search/new/geo_sort.rs
+++ b/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;
--- a/milli/src/search/new/graph_based_ranking_rule.rs
+++ b/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>,
--- a/milli/src/search/new/mod.rs
+++ b/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>,
--- a/milli/src/search/new/ranking_rule_graph/exactness/mod.rs
+++ b/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))
    }
--- a/milli/src/search/new/ranking_rule_graph/fid/mod.rs
+++ b/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)
    }
--- a/milli/src/search/new/ranking_rule_graph/position/mod.rs
+++ b/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)
    }
--- a/milli/src/search/new/ranking_rule_graph/proximity/mod.rs
+++ b/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)
    }
--- a/milli/src/search/new/ranking_rule_graph/typo/mod.rs
+++ b/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))
    }
--- a/milli/src/search/new/ranking_rule_graph/words/mod.rs
+++ b/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))
    }
--- a/milli/src/search/new/sort.rs
+++ b/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>,
--- a/milli/src/search/new/vector_sort.rs
+++ b/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;
    }