Skip to main content

uqa_planner/
text_top_k.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Physical planning for score-ordered text limits.
8
9use uqa_operators::{OperatorTree, TextTopKPlan, TextTopKStrategy};
10
11/// Storage and query facts needed to choose an exact text top-k algorithm.
12#[derive(Debug, Clone, Copy, PartialEq, Eq)]
13pub struct TextTopKCapabilities {
14    /// Number of analyzed query-term occurrences. Occurrences, rather than
15    /// unique terms, matter because duplicate query terms contribute twice.
16    pub analyzed_term_count: usize,
17    pub indexed_document_count: u64,
18}
19
20/// Push a score limit into a simple, field-bound text leaf.
21///
22/// Boolean/fusion trees deliberately remain exhaustive: cutting a child before
23/// its parent changes the carrier. Single-term and effectively unbounded
24/// searches also remain exhaustive because WAND cannot prune them profitably.
25#[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        // Execution validates scorer-versioned block bounds against the same transaction snapshot and falls back to exact WAND when unavailable.
60        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}