1use std::{
6 collections::{BTreeMap, BTreeSet},
7 convert::Infallible,
8 time::{Duration, Instant},
9};
10
11use hyphae_core::VectorSpaceName;
12use hyphae_query::{FieldPath, Record, Value};
13use thiserror::Error;
14use unicode_casefold::UnicodeCaseFold;
15use unicode_normalization::UnicodeNormalization;
16
17pub const MAX_LEXICAL_TOKEN_BYTES: usize = 256;
19pub const MAX_LEXICAL_FIELD_WEIGHT_MICROS: u32 = 1_000_000_000;
21pub const MAX_LEXICAL_FIELDS: usize = 64;
23pub const MAX_LEXICAL_PATH_SEGMENTS: usize = 32;
25pub const MAX_LEXICAL_PATH_SEGMENT_BYTES: usize = 1_024;
27const WEIGHT_SCALE: f64 = 1_000_000.0;
28const K1: f64 = 1.2;
29const B: f64 = 0.75;
30
31#[derive(Clone, Debug, Eq, PartialEq)]
33pub struct LexicalField {
34 pub path: FieldPath,
36 pub weight_micros: u32,
38}
39
40#[derive(Clone, Debug, Eq, PartialEq)]
42pub struct LexicalIndexDefinition {
43 pub name: VectorSpaceName,
45 pub fields: Vec<LexicalField>,
47}
48
49impl LexicalIndexDefinition {
50 pub fn new(name: VectorSpaceName, mut fields: Vec<LexicalField>) -> Result<Self, LexicalError> {
57 if fields.is_empty() {
58 return Err(LexicalError::EmptyFields);
59 }
60 if fields.len() > MAX_LEXICAL_FIELDS {
61 return Err(LexicalError::TooManyFields);
62 }
63 if fields.iter().any(|field| field.path.segments().is_empty()) {
64 return Err(LexicalError::EmptyFieldPath);
65 }
66 if fields.iter().any(|field| {
67 field.path.segments().len() > MAX_LEXICAL_PATH_SEGMENTS
68 || field.path.segments().iter().any(|segment| {
69 segment.is_empty() || segment.len() > MAX_LEXICAL_PATH_SEGMENT_BYTES
70 })
71 }) {
72 return Err(LexicalError::InvalidFieldSegment);
73 }
74 if fields
75 .iter()
76 .any(|field| !(1..=MAX_LEXICAL_FIELD_WEIGHT_MICROS).contains(&field.weight_micros))
77 {
78 return Err(LexicalError::InvalidFieldWeight);
79 }
80 fields.sort_by(|left, right| left.path.cmp(&right.path));
81 if fields.windows(2).any(|pair| pair[0].path == pair[1].path) {
82 return Err(LexicalError::DuplicateFieldPath);
83 }
84 Ok(Self { name, fields })
85 }
86}
87
88#[derive(Clone, Debug, Eq, PartialEq)]
90pub struct LexicalRequest {
91 pub index: VectorSpaceName,
93 pub query: String,
95 pub limit: usize,
97}
98
99#[derive(Clone, Debug, Eq, PartialEq)]
101pub struct LexicalLimits {
102 pub max_documents: u64,
104 pub max_tokens: u64,
106 pub max_candidates: u64,
108 pub max_returned: usize,
110 pub timeout: Duration,
112}
113
114impl Default for LexicalLimits {
115 fn default() -> Self {
116 Self {
117 max_documents: 1_000_000,
118 max_tokens: 10_000_000,
119 max_candidates: 100_000,
120 max_returned: 1_000,
121 timeout: Duration::from_secs(30),
122 }
123 }
124}
125
126#[derive(Clone, Debug, Eq, PartialEq)]
128pub struct LexicalFieldContribution {
129 pub path: FieldPath,
131 pub term_frequency: u64,
133 pub field_length: u64,
135}
136
137#[derive(Clone, Debug, Eq, PartialEq)]
139pub struct LexicalTermContribution {
140 pub token: String,
142 pub document_frequency: u64,
144 pub score_nanos: i64,
146 pub fields: Vec<LexicalFieldContribution>,
148}
149
150#[derive(Clone, Debug, Eq, PartialEq)]
152pub struct LexicalMatch {
153 pub key: Vec<u8>,
155 pub score_nanos: i64,
157 pub terms: Vec<LexicalTermContribution>,
159}
160
161#[derive(Clone, Copy, Debug, Eq, PartialEq)]
163pub enum LexicalAbstentionReason {
164 NoCandidates,
166}
167
168#[derive(Clone, Debug, Eq, PartialEq)]
170pub struct LexicalAbstention {
171 pub reason: LexicalAbstentionReason,
173 pub scanned_documents: u64,
175 pub query_tokens: Vec<String>,
177}
178
179#[derive(Clone, Debug, Eq, PartialEq)]
181pub enum LexicalOutcome {
182 Matches {
184 matches: Vec<LexicalMatch>,
186 scanned_documents: u64,
188 matched_documents: u64,
190 query_tokens: Vec<String>,
192 },
193 Abstained(LexicalAbstention),
195}
196
197#[derive(Clone, Debug, Error, Eq, PartialEq)]
199pub enum LexicalError {
200 #[error("lexical definition requires at least one field")]
202 EmptyFields,
203 #[error("lexical definition exceeds 64 fields")]
205 TooManyFields,
206 #[error("lexical field path must be nonempty")]
208 EmptyFieldPath,
209 #[error("lexical field path contains an invalid segment")]
211 InvalidFieldSegment,
212 #[error("lexical field paths must be unique")]
214 DuplicateFieldPath,
215 #[error("lexical field weight is outside 1..=1000000000")]
217 InvalidFieldWeight,
218 #[error("lexical request index does not match the definition")]
220 IndexMismatch,
221 #[error("lexical query has no retained normalized tokens")]
223 EmptyQuery,
224 #[error("lexical result limit must be nonzero")]
226 ZeroLimit,
227 #[error("lexical result limit {requested} exceeds maximum {maximum}")]
229 ResultLimitExceeded {
230 requested: usize,
232 maximum: usize,
234 },
235 #[error("lexical document key must be nonempty")]
237 EmptyDocumentKey,
238 #[error("duplicate lexical document key")]
240 DuplicateDocumentKey,
241 #[error("lexical document budget exceeded: {maximum}")]
243 DocumentBudgetExceeded {
244 maximum: u64,
246 },
247 #[error("lexical token budget exceeded: {maximum}")]
249 TokenBudgetExceeded {
250 maximum: u64,
252 },
253 #[error("lexical candidate budget exceeded: {maximum}")]
255 CandidateBudgetExceeded {
256 maximum: u64,
258 },
259 #[error("materialized lexical projection is malformed")]
261 MalformedProjection,
262 #[error("lexical retrieval timed out")]
264 TimedOut,
265 #[error("lexical score arithmetic overflow or non-finite result")]
267 ArithmeticOverflow,
268}
269
270pub fn tokenize_v1(input: &str) -> Vec<String> {
272 match tokenize_v1_checked(
273 input,
274 || Ok::<(), Infallible>(()),
275 || Ok::<(), Infallible>(()),
276 ) {
277 Ok(tokens) => tokens,
278 Err(never) => match never {},
279 }
280}
281
282pub fn tokenize_v1_checked<E>(
292 input: &str,
293 mut checkpoint: impl FnMut() -> Result<(), E>,
294 mut accept_token: impl FnMut() -> Result<(), E>,
295) -> Result<Vec<String>, E> {
296 const CHECKPOINT_INTERVAL: usize = 256;
297
298 checkpoint()?;
299 let mut tokens = Vec::new();
300 let mut token = String::new();
301 let mut discarding_oversized_token = false;
302 for (index, character) in input.nfkc().case_fold().enumerate() {
303 if index % CHECKPOINT_INTERVAL == 0 {
304 checkpoint()?;
305 }
306 if character.is_alphanumeric() {
307 if !discarding_oversized_token {
308 let next_length = token.len().saturating_add(character.len_utf8());
309 if next_length <= MAX_LEXICAL_TOKEN_BYTES {
310 token.push(character);
311 } else {
312 token.clear();
313 discarding_oversized_token = true;
314 }
315 }
316 } else {
317 push_token_checked(
318 &mut tokens,
319 &mut token,
320 &mut discarding_oversized_token,
321 &mut accept_token,
322 )?;
323 }
324 }
325 checkpoint()?;
326 push_token_checked(
327 &mut tokens,
328 &mut token,
329 &mut discarding_oversized_token,
330 &mut accept_token,
331 )?;
332 Ok(tokens)
333}
334
335fn push_token_checked<E>(
336 tokens: &mut Vec<String>,
337 token: &mut String,
338 discarding_oversized_token: &mut bool,
339 accept_token: &mut impl FnMut() -> Result<(), E>,
340) -> Result<(), E> {
341 if !*discarding_oversized_token && !token.is_empty() {
342 accept_token()?;
343 tokens.push(std::mem::take(token));
344 } else {
345 token.clear();
346 }
347 *discarding_oversized_token = false;
348 Ok(())
349}
350
351struct AnalyzedDocument {
352 key: Vec<u8>,
353 fields: Vec<Vec<String>>,
354}
355
356#[derive(Clone, Copy)]
357struct LexicalDeadline {
358 started: Instant,
359 timeout: Duration,
360}
361
362impl LexicalDeadline {
363 fn check(self) -> Result<(), LexicalError> {
364 check_timeout(self.started, self.timeout)
365 }
366}
367
368struct ScoringContext<'a> {
369 document_count: u64,
370 averages: &'a [f64],
371 frequencies: &'a BTreeMap<String, u64>,
372 definition: &'a LexicalIndexDefinition,
373 query_tokens: &'a [String],
374 deadline: LexicalDeadline,
375}
376
377#[derive(Clone, Debug, Eq, PartialEq)]
380pub struct LexicalMaterializedDocument {
381 pub key: Vec<u8>,
383 pub field_lengths: Vec<u64>,
385 pub term_frequencies: BTreeMap<String, Vec<u64>>,
387}
388
389#[derive(Clone, Debug, Eq, PartialEq)]
391pub struct LexicalMaterializedCorpus {
392 pub document_count: u64,
394 pub token_count: u64,
396 pub total_field_lengths: Vec<u64>,
398 pub documents: Vec<LexicalMaterializedDocument>,
400}
401
402pub fn retrieve_lexical(
409 records: &[Record],
410 definition: &LexicalIndexDefinition,
411 request: &LexicalRequest,
412 limits: &LexicalLimits,
413) -> Result<LexicalOutcome, LexicalError> {
414 validate_request(definition, request, limits)?;
415 let started = Instant::now();
416 let query_tokens = tokenize_before_deadline(&request.query, started, limits.timeout)?
417 .into_iter()
418 .collect::<BTreeSet<_>>()
419 .into_iter()
420 .collect::<Vec<_>>();
421 if query_tokens.is_empty() {
422 return Err(LexicalError::EmptyQuery);
423 }
424 let mut token_count = u64::try_from(query_tokens.len()).unwrap_or(u64::MAX);
425 if token_count > limits.max_tokens {
426 return Err(LexicalError::TokenBudgetExceeded {
427 maximum: limits.max_tokens,
428 });
429 }
430 let mut keys = BTreeSet::new();
431 let mut documents = Vec::with_capacity(records.len());
432 let mut total_lengths = vec![0_u64; definition.fields.len()];
433 for record in records {
434 check_timeout(started, limits.timeout)?;
435 if u64::try_from(documents.len()).unwrap_or(u64::MAX) >= limits.max_documents {
436 return Err(LexicalError::DocumentBudgetExceeded {
437 maximum: limits.max_documents,
438 });
439 }
440 if record.key.is_empty() {
441 return Err(LexicalError::EmptyDocumentKey);
442 }
443 if !keys.insert(record.key.as_slice()) {
444 return Err(LexicalError::DuplicateDocumentKey);
445 }
446 let mut fields = Vec::with_capacity(definition.fields.len());
447 for (field_index, field) in definition.fields.iter().enumerate() {
448 let tokens = match field.path.resolve(&record.value) {
449 Some(Value::String(value)) => {
450 tokenize_with_limits(value, &mut token_count, started, limits)?
451 }
452 _ => Vec::new(),
453 };
454 let length = u64::try_from(tokens.len()).unwrap_or(u64::MAX);
455 total_lengths[field_index] = total_lengths[field_index]
456 .checked_add(length)
457 .ok_or(LexicalError::ArithmeticOverflow)?;
458 fields.push(tokens);
459 }
460 documents.push(AnalyzedDocument {
461 key: record.key.clone(),
462 fields,
463 });
464 }
465 score_documents(
466 &documents,
467 &total_lengths,
468 definition,
469 request,
470 limits,
471 &query_tokens,
472 started,
473 )
474}
475
476pub fn retrieve_lexical_materialized(
483 corpus: &LexicalMaterializedCorpus,
484 definition: &LexicalIndexDefinition,
485 request: &LexicalRequest,
486 limits: &LexicalLimits,
487) -> Result<LexicalOutcome, LexicalError> {
488 validate_request(definition, request, limits)?;
489 let started = Instant::now();
490 let query_tokens = tokenize_before_deadline(&request.query, started, limits.timeout)?
491 .into_iter()
492 .collect::<BTreeSet<_>>()
493 .into_iter()
494 .collect::<Vec<_>>();
495 if query_tokens.is_empty() {
496 return Err(LexicalError::EmptyQuery);
497 }
498 validate_materialized_corpus(corpus, definition, &query_tokens, limits, started)?;
499 let averages = corpus
500 .total_field_lengths
501 .iter()
502 .map(|length| {
503 if corpus.document_count == 0 {
504 0.0
505 } else {
506 bounded_count_as_f64(*length) / bounded_count_as_f64(corpus.document_count)
507 }
508 })
509 .collect::<Vec<_>>();
510 let mut frequencies = BTreeMap::new();
511 for token in &query_tokens {
512 check_timeout(started, limits.timeout)?;
513 let mut frequency = 0_u64;
514 for document in &corpus.documents {
515 check_timeout(started, limits.timeout)?;
516 if document
517 .term_frequencies
518 .get(token)
519 .is_some_and(|fields| fields.iter().any(|value| *value > 0))
520 {
521 frequency = frequency
522 .checked_add(1)
523 .ok_or(LexicalError::ArithmeticOverflow)?;
524 }
525 }
526 frequencies.insert(token.clone(), frequency);
527 }
528 let scoring = ScoringContext {
529 document_count: corpus.document_count,
530 averages: &averages,
531 frequencies: &frequencies,
532 definition,
533 query_tokens: &query_tokens,
534 deadline: LexicalDeadline {
535 started,
536 timeout: limits.timeout,
537 },
538 };
539 let mut matches = Vec::with_capacity(corpus.documents.len().min(request.limit));
540 for document in &corpus.documents {
541 scoring.deadline.check()?;
542 if let Some(matched) = score_materialized_document(document, &scoring)? {
543 matches.push(matched);
544 }
545 }
546 finish_ranking(
547 matches,
548 corpus.document_count,
549 &query_tokens,
550 request.limit,
551 started,
552 limits.timeout,
553 )
554}
555
556fn validate_materialized_corpus(
557 corpus: &LexicalMaterializedCorpus,
558 definition: &LexicalIndexDefinition,
559 query_tokens: &[String],
560 limits: &LexicalLimits,
561 started: Instant,
562) -> Result<(), LexicalError> {
563 check_timeout(started, limits.timeout)?;
564 if corpus.document_count > limits.max_documents {
565 return Err(LexicalError::DocumentBudgetExceeded {
566 maximum: limits.max_documents,
567 });
568 }
569 let total_tokens = corpus
570 .token_count
571 .checked_add(u64::try_from(query_tokens.len()).unwrap_or(u64::MAX))
572 .ok_or(LexicalError::TokenBudgetExceeded {
573 maximum: limits.max_tokens,
574 })?;
575 if total_tokens > limits.max_tokens {
576 return Err(LexicalError::TokenBudgetExceeded {
577 maximum: limits.max_tokens,
578 });
579 }
580 if u64::try_from(corpus.documents.len()).unwrap_or(u64::MAX) > limits.max_candidates {
581 return Err(LexicalError::CandidateBudgetExceeded {
582 maximum: limits.max_candidates,
583 });
584 }
585 if corpus.total_field_lengths.len() != definition.fields.len() {
586 return Err(LexicalError::MalformedProjection);
587 }
588 let mut keys = BTreeSet::new();
589 for document in &corpus.documents {
590 check_timeout(started, limits.timeout)?;
591 if document.key.is_empty() {
592 return Err(LexicalError::EmptyDocumentKey);
593 }
594 if !keys.insert(document.key.as_slice()) {
595 return Err(LexicalError::DuplicateDocumentKey);
596 }
597 if document.field_lengths.len() != definition.fields.len()
598 || document.term_frequencies.len() != query_tokens.len()
599 {
600 return Err(LexicalError::MalformedProjection);
601 }
602 for token in query_tokens {
603 check_timeout(started, limits.timeout)?;
604 if document
605 .term_frequencies
606 .get(token)
607 .is_none_or(|frequencies| frequencies.len() != definition.fields.len())
608 {
609 return Err(LexicalError::MalformedProjection);
610 }
611 }
612 }
613 Ok(())
614}
615
616fn tokenize_before_deadline(
617 input: &str,
618 started: Instant,
619 timeout: Duration,
620) -> Result<Vec<String>, LexicalError> {
621 tokenize_v1_checked(
622 input,
623 || check_timeout(started, timeout),
624 || check_timeout(started, timeout),
625 )
626}
627
628fn tokenize_with_limits(
629 input: &str,
630 token_count: &mut u64,
631 started: Instant,
632 limits: &LexicalLimits,
633) -> Result<Vec<String>, LexicalError> {
634 tokenize_v1_checked(
635 input,
636 || check_timeout(started, limits.timeout),
637 || {
638 *token_count = token_count
639 .checked_add(1)
640 .ok_or(LexicalError::TokenBudgetExceeded {
641 maximum: limits.max_tokens,
642 })?;
643 if *token_count > limits.max_tokens {
644 return Err(LexicalError::TokenBudgetExceeded {
645 maximum: limits.max_tokens,
646 });
647 }
648 Ok(())
649 },
650 )
651}
652
653fn validate_request(
654 definition: &LexicalIndexDefinition,
655 request: &LexicalRequest,
656 limits: &LexicalLimits,
657) -> Result<(), LexicalError> {
658 if request.index != definition.name {
659 return Err(LexicalError::IndexMismatch);
660 }
661 if request.limit == 0 {
662 return Err(LexicalError::ZeroLimit);
663 }
664 if request.limit > limits.max_returned {
665 return Err(LexicalError::ResultLimitExceeded {
666 requested: request.limit,
667 maximum: limits.max_returned,
668 });
669 }
670 Ok(())
671}
672
673fn score_documents(
674 documents: &[AnalyzedDocument],
675 total_lengths: &[u64],
676 definition: &LexicalIndexDefinition,
677 request: &LexicalRequest,
678 limits: &LexicalLimits,
679 query_tokens: &[String],
680 started: Instant,
681) -> Result<LexicalOutcome, LexicalError> {
682 check_timeout(started, limits.timeout)?;
683 let document_count = u64::try_from(documents.len()).unwrap_or(u64::MAX);
684 let averages = total_lengths
685 .iter()
686 .map(|length| {
687 if document_count == 0 {
688 0.0
689 } else {
690 bounded_count_as_f64(*length) / bounded_count_as_f64(document_count)
691 }
692 })
693 .collect::<Vec<_>>();
694 let mut frequencies = BTreeMap::new();
695 for token in query_tokens {
696 check_timeout(started, limits.timeout)?;
697 let mut count = 0_u64;
698 for document in documents {
699 check_timeout(started, limits.timeout)?;
700 let mut present = false;
701 'fields: for field in &document.fields {
702 for candidate in field {
703 check_timeout(started, limits.timeout)?;
704 if candidate == token {
705 present = true;
706 break 'fields;
707 }
708 }
709 }
710 if present {
711 count = count
712 .checked_add(1)
713 .ok_or(LexicalError::ArithmeticOverflow)?;
714 }
715 }
716 frequencies.insert(token.clone(), count);
717 }
718 let scoring = ScoringContext {
719 document_count,
720 averages: &averages,
721 frequencies: &frequencies,
722 definition,
723 query_tokens,
724 deadline: LexicalDeadline {
725 started,
726 timeout: limits.timeout,
727 },
728 };
729 let mut matches = Vec::new();
730 for document in documents {
731 scoring.deadline.check()?;
732 let document_match = score_document(document, &scoring)?;
733 if let Some(matched) = document_match {
734 if u64::try_from(matches.len()).unwrap_or(u64::MAX) >= limits.max_candidates {
735 return Err(LexicalError::CandidateBudgetExceeded {
736 maximum: limits.max_candidates,
737 });
738 }
739 matches.push(matched);
740 }
741 }
742 finish_ranking(
743 matches,
744 document_count,
745 query_tokens,
746 request.limit,
747 started,
748 limits.timeout,
749 )
750}
751
752fn score_document(
753 document: &AnalyzedDocument,
754 context: &ScoringContext<'_>,
755) -> Result<Option<LexicalMatch>, LexicalError> {
756 context.deadline.check()?;
757 let field_lengths = document
758 .fields
759 .iter()
760 .map(|field| u64::try_from(field.len()).unwrap_or(u64::MAX))
761 .collect::<Vec<_>>();
762 let mut term_frequencies = BTreeMap::new();
763 for token in context.query_tokens {
764 context.deadline.check()?;
765 let mut per_field = Vec::with_capacity(document.fields.len());
766 for field in &document.fields {
767 let mut frequency = 0_u64;
768 for candidate in field {
769 context.deadline.check()?;
770 if candidate == token {
771 frequency = frequency
772 .checked_add(1)
773 .ok_or(LexicalError::ArithmeticOverflow)?;
774 }
775 }
776 per_field.push(frequency);
777 }
778 term_frequencies.insert(token.clone(), per_field);
779 }
780 score_statistics(&document.key, &field_lengths, &term_frequencies, context)
781}
782
783fn score_materialized_document(
784 document: &LexicalMaterializedDocument,
785 context: &ScoringContext<'_>,
786) -> Result<Option<LexicalMatch>, LexicalError> {
787 score_statistics(
788 &document.key,
789 &document.field_lengths,
790 &document.term_frequencies,
791 context,
792 )
793}
794
795fn score_statistics(
796 key: &[u8],
797 field_lengths: &[u64],
798 term_frequencies: &BTreeMap<String, Vec<u64>>,
799 context: &ScoringContext<'_>,
800) -> Result<Option<LexicalMatch>, LexicalError> {
801 context.deadline.check()?;
802 let mut terms = Vec::new();
803 let mut score_nanos = 0_i64;
804 for token in context.query_tokens {
805 context.deadline.check()?;
806 let document_frequency = context.frequencies[token];
807 if document_frequency == 0 {
808 continue;
809 }
810 let mut combined_tf = 0.0_f64;
811 let mut fields = Vec::with_capacity(context.definition.fields.len());
812 for (index, definition_field) in context.definition.fields.iter().enumerate() {
813 context.deadline.check()?;
814 let term_frequency = term_frequencies[token][index];
815 let field_length = field_lengths[index];
816 fields.push(LexicalFieldContribution {
817 path: definition_field.path.clone(),
818 term_frequency,
819 field_length,
820 });
821 if term_frequency > 0 && context.averages[index] > 0.0 {
822 let normalization =
823 1.0 - B + B * bounded_count_as_f64(field_length) / context.averages[index];
824 combined_tf += (f64::from(definition_field.weight_micros) / WEIGHT_SCALE)
825 * bounded_count_as_f64(term_frequency)
826 / normalization;
827 }
828 }
829 if combined_tf == 0.0 {
830 continue;
831 }
832 let numerator =
833 bounded_count_as_f64(context.document_count.saturating_sub(document_frequency)) + 0.5;
834 let denominator = bounded_count_as_f64(document_frequency) + 0.5;
835 let idf = libm::log(1.0 + numerator / denominator);
836 let term_score = quantize_score(idf * combined_tf * (K1 + 1.0) / (combined_tf + K1))?;
837 score_nanos = score_nanos
838 .checked_add(term_score)
839 .ok_or(LexicalError::ArithmeticOverflow)?;
840 terms.push(LexicalTermContribution {
841 token: token.clone(),
842 document_frequency,
843 score_nanos: term_score,
844 fields,
845 });
846 }
847 Ok((score_nanos > 0).then(|| LexicalMatch {
848 key: key.to_vec(),
849 score_nanos,
850 terms,
851 }))
852}
853
854fn finish_ranking(
855 mut matches: Vec<LexicalMatch>,
856 document_count: u64,
857 query_tokens: &[String],
858 limit: usize,
859 started: Instant,
860 timeout: Duration,
861) -> Result<LexicalOutcome, LexicalError> {
862 check_timeout(started, timeout)?;
863 matches.sort_by(|left, right| {
864 right
865 .score_nanos
866 .cmp(&left.score_nanos)
867 .then_with(|| left.key.cmp(&right.key))
868 });
869 check_timeout(started, timeout)?;
870 let matched_documents = u64::try_from(matches.len()).unwrap_or(u64::MAX);
871 matches.truncate(limit);
872 Ok(if matches.is_empty() {
873 LexicalOutcome::Abstained(LexicalAbstention {
874 reason: LexicalAbstentionReason::NoCandidates,
875 scanned_documents: document_count,
876 query_tokens: query_tokens.to_vec(),
877 })
878 } else {
879 LexicalOutcome::Matches {
880 matches,
881 scanned_documents: document_count,
882 matched_documents,
883 query_tokens: query_tokens.to_vec(),
884 }
885 })
886}
887
888fn quantize_score(value: f64) -> Result<i64, LexicalError> {
889 if !value.is_finite() || value < 0.0 {
890 return Err(LexicalError::ArithmeticOverflow);
891 }
892 let scaled = value * 1_000_000_000.0;
893 if !scaled.is_finite() {
894 return Err(LexicalError::ArithmeticOverflow);
895 }
896 if scaled >= maximum_i64_as_f64() {
897 return Ok(i64::MAX);
898 }
899 Ok(rounded_nonnegative_f64_as_i64(scaled))
900}
901
902#[allow(clippy::cast_precision_loss)]
905fn bounded_count_as_f64(value: u64) -> f64 {
906 value as f64
907}
908
909#[allow(clippy::cast_precision_loss)]
910fn maximum_i64_as_f64() -> f64 {
911 i64::MAX as f64
912}
913
914#[allow(clippy::cast_possible_truncation)]
917fn rounded_nonnegative_f64_as_i64(value: f64) -> i64 {
918 libm::floor(value + 0.5) as i64
919}
920
921fn check_timeout(started: Instant, timeout: Duration) -> Result<(), LexicalError> {
922 if started.elapsed() >= timeout {
923 Err(LexicalError::TimedOut)
924 } else {
925 Ok(())
926 }
927}
928
929#[cfg(test)]
930mod tests {
931 use std::collections::{BTreeMap, BTreeSet};
932
933 use super::*;
934 use proptest::prelude::*;
935
936 fn record(key: &[u8], title: &str, body: &str) -> Record {
937 Record::new(
938 key,
939 Value::Object(BTreeMap::from([
940 ("title".into(), Value::String(title.into())),
941 ("body".into(), Value::String(body.into())),
942 ])),
943 )
944 }
945
946 fn definition() -> Result<LexicalIndexDefinition, LexicalError> {
947 LexicalIndexDefinition::new(
948 VectorSpaceName::new("docs").map_err(|_| LexicalError::EmptyFields)?,
949 vec![
950 LexicalField {
951 path: FieldPath::field("body"),
952 weight_micros: 1_000_000,
953 },
954 LexicalField {
955 path: FieldPath::field("title"),
956 weight_micros: 2_000_000,
957 },
958 ],
959 )
960 }
961
962 fn materialize_reference_corpus(
963 records: &[Record],
964 definition: &LexicalIndexDefinition,
965 query: &str,
966 ) -> LexicalMaterializedCorpus {
967 let query_tokens = tokenize_v1(query)
968 .into_iter()
969 .collect::<BTreeSet<_>>()
970 .into_iter()
971 .collect::<Vec<_>>();
972 let mut token_count = 0_u64;
973 let mut total_field_lengths = vec![0_u64; definition.fields.len()];
974 let mut documents = Vec::new();
975
976 for record in records {
977 let fields = definition
978 .fields
979 .iter()
980 .map(|field| match field.path.resolve(&record.value) {
981 Some(Value::String(value)) => tokenize_v1(value),
982 _ => Vec::new(),
983 })
984 .collect::<Vec<_>>();
985 let field_lengths = fields
986 .iter()
987 .map(|field| u64::try_from(field.len()).unwrap_or(u64::MAX))
988 .collect::<Vec<_>>();
989 for (total, length) in total_field_lengths.iter_mut().zip(&field_lengths) {
990 *total = total.saturating_add(*length);
991 token_count = token_count.saturating_add(*length);
992 }
993 let term_frequencies = query_tokens
994 .iter()
995 .map(|token| {
996 let frequencies = fields
997 .iter()
998 .map(|field| {
999 u64::try_from(
1000 field.iter().filter(|candidate| *candidate == token).count(),
1001 )
1002 .unwrap_or(u64::MAX)
1003 })
1004 .collect::<Vec<_>>();
1005 (token.clone(), frequencies)
1006 })
1007 .collect::<BTreeMap<_, _>>();
1008 if term_frequencies
1009 .values()
1010 .any(|frequencies| frequencies.iter().any(|frequency| *frequency > 0))
1011 {
1012 documents.push(LexicalMaterializedDocument {
1013 key: record.key.clone(),
1014 field_lengths,
1015 term_frequencies,
1016 });
1017 }
1018 }
1019
1020 LexicalMaterializedCorpus {
1021 document_count: u64::try_from(records.len()).unwrap_or(u64::MAX),
1022 token_count,
1023 total_field_lengths,
1024 documents,
1025 }
1026 }
1027
1028 #[test]
1029 fn tokenizer_pins_nfkc_casefold_and_alphanumeric_runs() {
1030 assert_eq!(
1031 tokenize_v1("Straße ABC—café"),
1032 vec!["strasse", "abc", "café"]
1033 );
1034 }
1035
1036 #[test]
1037 fn checked_tokenizer_matches_v1_and_checks_inside_one_long_token() {
1038 let input = "Straße ABC—café";
1039 let mut accepted = 0_usize;
1040 let checked = match tokenize_v1_checked(
1041 input,
1042 || Ok::<(), Infallible>(()),
1043 || {
1044 accepted += 1;
1045 Ok::<(), Infallible>(())
1046 },
1047 ) {
1048 Ok(tokens) => tokens,
1049 Err(never) => match never {},
1050 };
1051 assert_eq!(checked, tokenize_v1(input));
1052 assert_eq!(accepted, checked.len());
1053 assert_eq!(
1054 tokenize_v1(&format!("{} tail", "a".repeat(257))),
1055 vec!["tail"]
1056 );
1057
1058 let long_token = "a".repeat(2_048);
1059 let mut checkpoints = 0_usize;
1060 let interrupted = tokenize_v1_checked(
1061 &long_token,
1062 || {
1063 checkpoints += 1;
1064 if checkpoints == 3 {
1065 Err("stop")
1066 } else {
1067 Ok(())
1068 }
1069 },
1070 || Ok::<(), &'static str>(()),
1071 );
1072 assert_eq!(interrupted, Err("stop"));
1073 assert_eq!(checkpoints, 3);
1074 }
1075
1076 #[test]
1077 fn zero_timeout_includes_query_tokenization_for_both_scorers() -> Result<(), LexicalError> {
1078 let definition = definition()?;
1079 let request = LexicalRequest {
1080 index: definition.name.clone(),
1081 query: "rust".into(),
1082 limit: 1,
1083 };
1084 let limits = LexicalLimits {
1085 timeout: Duration::ZERO,
1086 ..LexicalLimits::default()
1087 };
1088 let corpus = materialize_reference_corpus(&[], &definition, &request.query);
1089
1090 assert_eq!(
1091 retrieve_lexical(&[], &definition, &request, &limits),
1092 Err(LexicalError::TimedOut)
1093 );
1094 assert_eq!(
1095 retrieve_lexical_materialized(&corpus, &definition, &request, &limits),
1096 Err(LexicalError::TimedOut)
1097 );
1098 Ok(())
1099 }
1100
1101 #[test]
1102 fn duplicate_query_tokens_count_once_against_token_budget() -> Result<(), LexicalError> {
1103 let definition = definition()?;
1104 let request = LexicalRequest {
1105 index: definition.name.clone(),
1106 query: "rust rust".into(),
1107 limit: 1,
1108 };
1109 let limits = LexicalLimits {
1110 max_tokens: 1,
1111 ..LexicalLimits::default()
1112 };
1113 let corpus = materialize_reference_corpus(&[], &definition, &request.query);
1114
1115 let reference = retrieve_lexical(&[], &definition, &request, &limits)?;
1116 let materialized = retrieve_lexical_materialized(&corpus, &definition, &request, &limits)?;
1117 assert_eq!(reference, materialized);
1118 assert!(matches!(
1119 reference,
1120 LexicalOutcome::Abstained(LexicalAbstention {
1121 query_tokens,
1122 ..
1123 }) if query_tokens == vec!["rust"]
1124 ));
1125 Ok(())
1126 }
1127
1128 #[test]
1129 fn bm25f_is_deterministic_and_binary_key_breaks_ties() -> Result<(), LexicalError> {
1130 let definition = definition()?;
1131 let outcome = retrieve_lexical(
1132 &[
1133 record(b"b", "Rust memory", "durable engine"),
1134 record(b"a", "Rust memory", "durable engine"),
1135 record(b"z", "other", "nothing"),
1136 ],
1137 &definition,
1138 &LexicalRequest {
1139 index: definition.name.clone(),
1140 query: "RUST rust".into(),
1141 limit: 10,
1142 },
1143 &LexicalLimits::default(),
1144 )?;
1145 let LexicalOutcome::Matches { matches, .. } = outcome else {
1146 return Err(LexicalError::ArithmeticOverflow);
1147 };
1148 assert_eq!(matches[0].key, b"a");
1149 assert_eq!(matches[1].key, b"b");
1150 assert_eq!(matches[0].score_nanos, matches[1].score_nanos);
1151 Ok(())
1152 }
1153
1154 proptest! {
1155 #![proptest_config(ProptestConfig::with_cases(64))]
1156
1157 #[test]
1158 fn materialized_scorer_matches_reference_for_random_corpora(
1159 generated in prop::collection::vec(
1160 ("[a-z ]{0,24}", "[a-z ]{0,48}"),
1161 0..32
1162 ),
1163 query in "(rust|durable|engine|memory)( (rust|durable|engine|memory)){0,2}",
1164 limit in 1_usize..16
1165 ) {
1166 let definition = definition().map_err(|error| TestCaseError::fail(error.to_string()))?;
1167 let records = generated
1168 .iter()
1169 .enumerate()
1170 .map(|(index, (title, body))| {
1171 record(&u64::try_from(index).unwrap_or(u64::MAX).to_be_bytes(), title, body)
1172 })
1173 .collect::<Vec<_>>();
1174 let request = LexicalRequest {
1175 index: definition.name.clone(),
1176 query: query.clone(),
1177 limit,
1178 };
1179 let limits = LexicalLimits::default();
1180 let reference = retrieve_lexical(&records, &definition, &request, &limits)
1181 .map_err(|error| TestCaseError::fail(error.to_string()))?;
1182 let corpus = materialize_reference_corpus(&records, &definition, &query);
1183 let materialized =
1184 retrieve_lexical_materialized(&corpus, &definition, &request, &limits)
1185 .map_err(|error| TestCaseError::fail(error.to_string()))?;
1186
1187 prop_assert_eq!(materialized, reference);
1188 }
1189 }
1190}