uqa_planner/
text_top_k.rs1use uqa_operators::{OperatorTree, TextTopKPlan, TextTopKStrategy};
10
11#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct TextTopKCapabilities {
14 pub analyzed_term_count: usize,
17 pub indexed_document_count: u64,
18}
19
20#[must_use]
26pub fn plan_text_top_k(
27 tree: OperatorTree,
28 k: usize,
29 capabilities: TextTopKCapabilities,
30) -> OperatorTree {
31 let (query, field, scoring, top_k) = match tree {
32 OperatorTree::Term {
33 query,
34 field,
35 scoring,
36 top_k,
37 } => (query, field, scoring, top_k),
38 other => return other,
39 };
40
41 let eligible = field.is_some()
42 && scoring.is_some()
43 && top_k.is_none()
44 && capabilities.analyzed_term_count >= 2
45 && (k == 0 || (k as u128) < u128::from(capabilities.indexed_document_count));
46 if !eligible {
47 return OperatorTree::Term {
48 query,
49 field,
50 scoring,
51 top_k,
52 };
53 }
54
55 OperatorTree::Term {
56 query,
57 field,
58 scoring,
59 top_k: Some(TextTopKPlan {
61 k,
62 strategy: TextTopKStrategy::BlockMaxWand,
63 }),
64 }
65}
66
67#[cfg(test)]
68mod tests {
69 use super::*;
70 use uqa_operators::TextScoringMode;
71
72 #[test]
73 fn phrase_support_is_not_cut_off_by_bag_of_terms_top_k() {
74 let planned = plan_text_top_k(
75 OperatorTree::Phrase {
76 query: "red fox".into(),
77 field: Some("body".into()),
78 scoring: Some(TextScoringMode::BM25),
79 },
80 1,
81 TextTopKCapabilities {
82 analyzed_term_count: 2,
83 indexed_document_count: 100,
84 },
85 );
86 assert!(matches!(planned, OperatorTree::Phrase { query, .. } if query == "red fox"));
87 }
88
89 fn term() -> OperatorTree {
90 OperatorTree::Term {
91 query: "rust search".into(),
92 field: Some("body".into()),
93 scoring: Some(TextScoringMode::BM25),
94 top_k: None,
95 }
96 }
97
98 #[test]
99 fn eligible_query_defers_block_validation_to_execution() {
100 let planned = plan_text_top_k(
101 term(),
102 10,
103 TextTopKCapabilities {
104 analyzed_term_count: 2,
105 indexed_document_count: 100,
106 },
107 );
108 assert!(matches!(
109 planned,
110 OperatorTree::Term {
111 top_k: Some(TextTopKPlan {
112 strategy: TextTopKStrategy::BlockMaxWand,
113 ..
114 }),
115 ..
116 }
117 ));
118 }
119
120 #[test]
121 fn single_term_and_unbounded_inputs_stay_exhaustive() {
122 for (term_count, k, documents) in [(1, 10, 100), (2, 100, 100)] {
123 let planned = plan_text_top_k(
124 term(),
125 k,
126 TextTopKCapabilities {
127 analyzed_term_count: term_count,
128 indexed_document_count: documents,
129 },
130 );
131 assert!(matches!(planned, OperatorTree::Term { top_k: None, .. }));
132 }
133 }
134}