1use std::sync::Arc;
4
5use crate::dsl::Field;
6use crate::segment::SegmentReader;
7use crate::structures::BlockPostingList;
8use crate::structures::TERMINATED;
9use crate::{DocId, Score};
10
11use super::{CountFuture, EmptyScorer, GlobalStats, Query, Scorer, ScorerFuture, TermQueryInfo};
12
13#[derive(Clone)]
15pub struct TermQuery {
16 pub field: Field,
17 pub term: Vec<u8>,
18 global_stats: Option<Arc<GlobalStats>>,
20}
21
22impl std::fmt::Debug for TermQuery {
23 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
24 f.debug_struct("TermQuery")
25 .field("field", &self.field)
26 .field("term", &String::from_utf8_lossy(&self.term))
27 .field("has_global_stats", &self.global_stats.is_some())
28 .finish()
29 }
30}
31
32impl std::fmt::Display for TermQuery {
33 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
34 write!(
35 f,
36 "Term({}:\"{}\")",
37 self.field.0,
38 String::from_utf8_lossy(&self.term)
39 )
40 }
41}
42
43impl TermQuery {
44 pub fn new(field: Field, term: impl Into<Vec<u8>>) -> Self {
45 Self {
46 field,
47 term: term.into(),
48 global_stats: None,
49 }
50 }
51
52 pub fn text(field: Field, text: &str) -> Self {
53 Self {
54 field,
55 term: text.to_lowercase().into_bytes(),
56 global_stats: None,
57 }
58 }
59
60 pub fn with_global_stats(field: Field, text: &str, stats: Arc<GlobalStats>) -> Self {
62 Self {
63 field,
64 term: text.to_lowercase().into_bytes(),
65 global_stats: Some(stats),
66 }
67 }
68
69 pub fn set_global_stats(&mut self, stats: Arc<GlobalStats>) {
71 self.global_stats = Some(stats);
72 }
73}
74
75fn compute_term_idf(
77 posting_list: &BlockPostingList,
78 field: Field,
79 reader: &SegmentReader,
80 global_stats: Option<&Arc<GlobalStats>>,
81 term: &[u8],
82) -> (f32, f32) {
83 if let Some(stats) = global_stats {
84 let term_str = String::from_utf8_lossy(term);
85 let global_idf = stats.text_idf(field, &term_str);
86 if global_idf > 0.0 {
87 return (global_idf, stats.avg_field_len(field));
88 }
89 }
90 let num_docs = reader.num_docs() as f32;
91 let doc_freq = posting_list.doc_count() as f32;
92 (
93 super::bm25_idf(doc_freq, num_docs),
94 reader.avg_field_len(field),
95 )
96}
97
98macro_rules! term_plan {
105 ($field:expr, $term:expr, $global_stats:expr, $reader:expr, $limit:expr,
106 $load_positions:expr, $get_postings_fn:ident, $get_positions_fn:ident
107 $(, $aw:tt)*) => {{
108 let field: Field = $field;
109 let term: &[u8] = $term;
110 let global_stats: Option<&Arc<GlobalStats>> = $global_stats;
111 let reader: &SegmentReader = $reader;
112 let limit: usize = $limit;
113
114 let is_indexed = reader.schema().get_field_entry(field).is_none_or(|e| e.indexed);
116 if !is_indexed {
117 let term_str = String::from_utf8_lossy(term);
118 if let Some(scorer) = FastFieldTextScorer::try_new(reader, field, &term_str) {
119 return Ok(Box::new(scorer) as Box<dyn Scorer + '_>);
120 }
121 return Ok(Box::new(EmptyScorer) as Box<dyn Scorer + '_>);
122 }
123
124 let postings = reader.$get_postings_fn(field, term) $(. $aw)* ?;
125
126 match postings {
127 Some(posting_list) if reader.is_chunked_field(field) => {
130 let num_chunks = reader.text_corpus_size(field);
131 let idf = super::bm25_idf(posting_list.doc_count() as f32, num_chunks);
132 super::planner::finish_chunked_text_maxscore(
133 vec![(posting_list, idf)],
134 limit,
135 reader,
136 field,
137 )
138 }
139 Some(posting_list) => {
140 let (idf, avg_field_len) =
141 compute_term_idf(&posting_list, field, reader, global_stats, term);
142
143 let positions = if $load_positions {
144 reader.$get_positions_fn(field, term) $(. $aw)* ?
145 } else {
146 None
147 };
148
149 let mut scorer = TermScorer::new(posting_list, idf, avg_field_len, 1.0);
150 if let Some(pos) = positions {
151 scorer = scorer.with_positions(field.0, pos);
152 }
153 Ok(Box::new(scorer) as Box<dyn Scorer + '_>)
154 }
155 None => {
156 let term_str = String::from_utf8_lossy(term);
157 if let Some(scorer) = FastFieldTextScorer::try_new(reader, field, &term_str) {
158 Ok(Box::new(scorer) as Box<dyn Scorer + '_>)
159 } else {
160 Ok(Box::new(EmptyScorer) as Box<dyn Scorer + '_>)
161 }
162 }
163 }
164 }};
165}
166
167impl Query for TermQuery {
168 fn scorer<'a>(&self, reader: &'a SegmentReader, limit: usize) -> ScorerFuture<'a> {
169 self.scorer_with_options(reader, limit, super::ScorerOptions::with_positions())
170 }
171
172 fn scorer_with_options<'a>(
173 &self,
174 reader: &'a SegmentReader,
175 limit: usize,
176 options: super::ScorerOptions,
177 ) -> ScorerFuture<'a> {
178 let field = self.field;
179 let term = self.term.clone();
180 let global_stats = self.global_stats.clone();
181 let load_positions = options.collect_positions;
182 Box::pin(async move {
183 term_plan!(
184 field,
185 &term,
186 global_stats.as_ref(),
187 reader,
188 limit,
189 load_positions,
190 get_postings,
191 get_positions,
192 await
193 )
194 })
195 }
196
197 fn count_estimate<'a>(&self, reader: &'a SegmentReader) -> CountFuture<'a> {
198 let field = self.field;
199 let term = self.term.clone();
200 Box::pin(async move {
201 match reader.get_postings(field, &term).await? {
202 Some(list) => Ok(list.doc_count()),
203 None => Ok(0),
204 }
205 })
206 }
207
208 #[cfg(feature = "sync")]
209 fn scorer_sync<'a>(
210 &self,
211 reader: &'a SegmentReader,
212 limit: usize,
213 ) -> crate::Result<Box<dyn Scorer + 'a>> {
214 self.scorer_sync_with_options(reader, limit, super::ScorerOptions::with_positions())
215 }
216
217 #[cfg(feature = "sync")]
218 fn scorer_sync_with_options<'a>(
219 &self,
220 reader: &'a SegmentReader,
221 limit: usize,
222 options: super::ScorerOptions,
223 ) -> crate::Result<Box<dyn Scorer + 'a>> {
224 term_plan!(
225 self.field,
226 &self.term,
227 self.global_stats.as_ref(),
228 reader,
229 limit,
230 options.collect_positions,
231 get_postings_sync,
232 get_positions_sync
233 )
234 }
235
236 fn as_doc_predicate<'a>(&self, reader: &'a SegmentReader) -> Option<super::DocPredicate<'a>> {
237 let fast_field = reader.fast_field(self.field.0)?;
238 let term_str = String::from_utf8_lossy(&self.term);
239 match fast_field.text_ordinal(&term_str) {
240 Some(target_ordinal) => Some(Box::new(move |doc_id: DocId| -> bool {
241 fast_field.get_u64(doc_id) == target_ordinal
242 })),
243 None => Some(Box::new(|_| false)),
245 }
246 }
247
248 #[cfg(feature = "sync")]
249 fn bitset_cardinality_estimate(&self, reader: &SegmentReader) -> Option<u64> {
250 if reader.is_chunked_field(self.field) {
252 return None;
253 }
254 let pl = reader.get_postings_sync(self.field, &self.term).ok()??;
256 Some(pl.doc_count() as u64)
257 }
258
259 #[cfg(feature = "sync")]
260 fn as_doc_bitset(&self, reader: &SegmentReader) -> Option<super::DocBitset> {
261 if reader.is_chunked_field(self.field) {
264 return None;
265 }
266 let pl = reader.get_postings_sync(self.field, &self.term).ok()??;
269 let mut bitset = super::DocBitset::new(reader.num_docs());
270 let mut iter = pl.iterator();
271 loop {
272 let doc = iter.doc();
273 if doc == crate::structures::TERMINATED {
274 break;
275 }
276 bitset.set(doc);
277 iter.advance();
278 }
279 Some(bitset)
280 }
281
282 fn decompose(&self) -> super::QueryDecomposition {
283 super::QueryDecomposition::TextTerm(TermQueryInfo {
284 field: self.field,
285 term: self.term.clone(),
286 })
287 }
288}
289
290struct TermScorer {
291 iterator: crate::structures::BlockPostingIterator<'static>,
292 idf: f32,
293 avg_field_len: f32,
295 field_boost: f32,
297 field_id: u32,
299 positions: Option<crate::structures::PositionPostingList>,
301}
302
303impl TermScorer {
304 pub fn new(
305 posting_list: BlockPostingList,
306 idf: f32,
307 avg_field_len: f32,
308 field_boost: f32,
309 ) -> Self {
310 Self {
311 iterator: posting_list.into_iterator(),
312 idf,
313 avg_field_len,
314 field_boost,
315 field_id: 0,
316 positions: None,
317 }
318 }
319
320 pub fn with_positions(
321 mut self,
322 field_id: u32,
323 positions: crate::structures::PositionPostingList,
324 ) -> Self {
325 self.field_id = field_id;
326 self.positions = Some(positions);
327 self
328 }
329}
330
331impl super::docset::DocSet for TermScorer {
332 fn doc(&self) -> DocId {
333 self.iterator.doc()
334 }
335
336 fn advance(&mut self) -> DocId {
337 self.iterator.advance()
338 }
339
340 fn seek(&mut self, target: DocId) -> DocId {
341 self.iterator.seek(target)
342 }
343
344 fn size_hint(&self) -> u32 {
345 0
346 }
347}
348
349struct FastFieldTextScorer<'a> {
355 fast_field: &'a crate::structures::fast_field::FastFieldReader,
356 target_ordinal: u64,
357 current: u32,
358 num_docs: u32,
359}
360
361impl<'a> FastFieldTextScorer<'a> {
362 fn try_new(reader: &'a SegmentReader, field: Field, text: &str) -> Option<Self> {
363 let fast_field = reader.fast_field(field.0)?;
364 let target_ordinal = fast_field.text_ordinal(text)?;
365 let num_docs = reader.num_docs();
366 let mut scorer = Self {
367 fast_field,
368 target_ordinal,
369 current: 0,
370 num_docs,
371 };
372 if num_docs > 0 && fast_field.get_u64(0) != target_ordinal {
374 scorer.scan_forward();
375 }
376 Some(scorer)
377 }
378
379 fn scan_forward(&mut self) {
380 loop {
381 self.current += 1;
382 if self.current >= self.num_docs {
383 self.current = self.num_docs;
384 return;
385 }
386 if self.fast_field.get_u64(self.current) == self.target_ordinal {
387 return;
388 }
389 }
390 }
391}
392
393impl super::docset::DocSet for FastFieldTextScorer<'_> {
394 fn doc(&self) -> DocId {
395 if self.current >= self.num_docs {
396 TERMINATED
397 } else {
398 self.current
399 }
400 }
401
402 fn advance(&mut self) -> DocId {
403 self.scan_forward();
404 self.doc()
405 }
406
407 fn seek(&mut self, target: DocId) -> DocId {
408 if target > self.current {
409 self.current = target;
410 if self.current < self.num_docs
411 && self.fast_field.get_u64(self.current) != self.target_ordinal
412 {
413 self.scan_forward();
414 }
415 }
416 self.doc()
417 }
418
419 fn size_hint(&self) -> u32 {
420 0
421 }
422}
423
424impl Scorer for FastFieldTextScorer<'_> {
425 fn score(&self) -> Score {
426 1.0
427 }
428}
429
430impl Scorer for TermScorer {
431 fn score(&self) -> Score {
432 let tf = self.iterator.term_freq() as f32;
433 super::bm25f_score(tf, self.idf, tf, self.avg_field_len, self.field_boost)
436 }
437
438 fn matched_positions(&self) -> Option<super::MatchedPositions> {
439 let positions = self.positions.as_ref()?;
440 let doc_id = self.iterator.doc();
441 let pos = positions.get_positions(doc_id)?;
442 let score = self.score();
443 let per_position_score = if pos.is_empty() {
445 0.0
446 } else {
447 score / pos.len() as f32
448 };
449 let scored_positions: Vec<super::ScoredPosition> = pos
450 .iter()
451 .map(|&p| super::ScoredPosition::new(p, per_position_score))
452 .collect();
453 Some(vec![(self.field_id, scored_positions)])
454 }
455}