Search: spans with a finer granularity

This commit is contained in:
Louis Dureuil 2024-07-02 12:14:40 +02:00
parent 3129f96603
commit 128e6c7502
No known key found for this signature in database
15 changed files with 55 additions and 13 deletions

View File

@ -752,10 +752,15 @@ fn prepare_search<'t>(
SearchKind::SemanticOnly { embedder_name, embedder } => {
let vector = match query.vector.clone() {
Some(vector) => vector,
None => embedder
.embed_one(query.q.clone().unwrap())
.map_err(milli::vector::Error::from)
.map_err(milli::Error::from)?,
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)?
}
};
search.semantic(embedder_name.clone(), embedder.clone(), Some(vector));

View File

@ -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) => {

View File

@ -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);

View File

@ -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>,

View File

@ -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;

View File

@ -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>,

View File

@ -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>,

View File

@ -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))
}

View File

@ -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)
}

View File

@ -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)
}

View File

@ -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)
}

View File

@ -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))
}

View File

@ -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))
}

View File

@ -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>,

View File

@ -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;
}