lance_index/scalar/inverted/
query.rs

1// SPDX-License-Identifier: Apache-2.0
2// SPDX-FileCopyrightText: Copyright The Lance Authors
3
4use crate::scalar::inverted::document_tokenizer::DocType;
5use crate::scalar::inverted::tokenizer::document_tokenizer::LanceTokenizer;
6use lance_core::{Error, Result};
7use serde::ser::SerializeMap;
8use serde::{Deserialize, Serialize};
9use std::collections::{HashMap, HashSet};
10
11#[derive(Debug, Clone, Serialize, Deserialize)]
12pub struct FtsSearchParams {
13    /// Controls result completeness for each recursively planned FTS node.
14    ///
15    /// `Some(k)` permits bounded top-k execution, while `None` requires complete
16    /// results so an outer compound query can combine every candidate exactly.
17    /// Future competitive-score pruning must use a separate planning contract
18    /// instead of inferring a bound from the ambient scanner limit.
19    pub limit: Option<usize>,
20    pub wand_factor: f32,
21    pub fuzziness: Option<u32>,
22    pub max_expansions: usize,
23    // None means not a phrase query
24    // Some(n) means a phrase query with slop n
25    pub phrase_slop: Option<u32>,
26    /// The number of beginning characters being unchanged for fuzzy matching.
27    pub prefix_length: u32,
28}
29
30impl FtsSearchParams {
31    pub fn new() -> Self {
32        Self {
33            limit: None,
34            wand_factor: 1.0,
35            fuzziness: Some(0),
36            max_expansions: 50,
37            phrase_slop: None,
38            prefix_length: 0,
39        }
40    }
41
42    pub fn with_limit(mut self, limit: Option<usize>) -> Self {
43        self.limit = limit;
44        self
45    }
46
47    pub fn with_wand_factor(mut self, factor: f32) -> Self {
48        self.wand_factor = factor;
49        self
50    }
51
52    pub fn with_fuzziness(mut self, fuzziness: Option<u32>) -> Self {
53        self.fuzziness = fuzziness;
54        self
55    }
56
57    pub fn with_max_expansions(mut self, max_expansions: usize) -> Self {
58        self.max_expansions = max_expansions;
59        self
60    }
61
62    pub fn with_phrase_slop(mut self, phrase_slop: Option<u32>) -> Self {
63        self.phrase_slop = phrase_slop;
64        self
65    }
66
67    pub fn with_prefix_length(mut self, prefix_length: u32) -> Self {
68        self.prefix_length = prefix_length;
69        self
70    }
71}
72
73impl Default for FtsSearchParams {
74    fn default() -> Self {
75        Self::new()
76    }
77}
78
79#[derive(Debug, Clone, Copy, PartialEq, Serialize, Deserialize, Default)]
80pub enum Operator {
81    And,
82    #[default]
83    Or,
84}
85
86impl TryFrom<&str> for Operator {
87    type Error = Error;
88    fn try_from(value: &str) -> Result<Self> {
89        match value.to_ascii_uppercase().as_str() {
90            "AND" => Ok(Self::And),
91            "OR" => Ok(Self::Or),
92            _ => Err(Error::invalid_input(format!("Invalid operator: {}", value))),
93        }
94    }
95}
96
97impl From<Operator> for &'static str {
98    fn from(operator: Operator) -> Self {
99        match operator {
100            Operator::And => "AND",
101            Operator::Or => "OR",
102        }
103    }
104}
105
106pub trait FtsQueryNode {
107    fn columns(&self) -> HashSet<String>;
108}
109
110#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
111#[serde(rename_all = "snake_case")]
112pub enum FtsQuery {
113    // leaf queries
114    Match(MatchQuery),
115    Phrase(PhraseQuery),
116
117    // compound queries
118    Boost(BoostQuery),
119    MultiMatch(MultiMatchQuery),
120    Boolean(BooleanQuery),
121}
122
123impl std::fmt::Display for FtsQuery {
124    fn fmt(&self, f: &mut std::fmt::Formatter) -> std::fmt::Result {
125        match self {
126            Self::Match(query) => write!(f, "Match({:?})", query),
127            Self::Phrase(query) => write!(f, "Phrase({:?})", query),
128            Self::Boost(query) => write!(
129                f,
130                "Boosting(positive={}, negative={}, negative_boost={})",
131                query.positive, query.negative, query.negative_boost
132            ),
133            Self::MultiMatch(query) => write!(f, "MultiMatch({:?})", query),
134            Self::Boolean(query) => {
135                write!(
136                    f,
137                    "Boolean(must={:?}, should={:?})",
138                    query.must, query.should
139                )
140            }
141        }
142    }
143}
144
145impl FtsQueryNode for FtsQuery {
146    fn columns(&self) -> HashSet<String> {
147        match self {
148            Self::Match(query) => query.columns(),
149            Self::Phrase(query) => query.columns(),
150            Self::Boost(query) => {
151                let mut columns = query.positive.columns();
152                columns.extend(query.negative.columns());
153                columns
154            }
155            Self::MultiMatch(query) => {
156                let mut columns = HashSet::new();
157                for match_query in &query.match_queries {
158                    columns.extend(match_query.columns());
159                }
160                columns
161            }
162            Self::Boolean(query) => {
163                let mut columns = HashSet::new();
164                for query in &query.must {
165                    columns.extend(query.columns());
166                }
167                for query in &query.should {
168                    columns.extend(query.columns());
169                }
170                columns
171            }
172        }
173    }
174}
175
176impl FtsQuery {
177    pub fn query(&self) -> String {
178        match self {
179            Self::Match(query) => query.terms.clone(),
180            Self::Phrase(query) => format!("\"{}\"", query.terms), // Phrase queries are quoted
181            Self::Boost(query) => query.positive.query(),
182            Self::MultiMatch(query) => query.match_queries[0].terms.clone(),
183            Self::Boolean(_) => {
184                // Bool queries don't have a single query string, they are composed of multiple queries
185                String::new()
186            }
187        }
188    }
189
190    pub fn is_missing_column(&self) -> bool {
191        match self {
192            Self::Match(query) => query.column.is_none(),
193            Self::Phrase(query) => query.column.is_none(),
194            Self::Boost(query) => {
195                query.positive.is_missing_column() || query.negative.is_missing_column()
196            }
197            Self::MultiMatch(query) => query.match_queries.iter().any(|q| q.column.is_none()),
198            Self::Boolean(query) => {
199                query.must.iter().any(|q| q.is_missing_column())
200                    || query.should.iter().any(|q| q.is_missing_column())
201            }
202        }
203    }
204
205    pub fn with_column(self, column: String) -> Self {
206        match self {
207            Self::Match(query) => Self::Match(query.with_column(Some(column))),
208            Self::Phrase(query) => Self::Phrase(query.with_column(Some(column))),
209            Self::Boost(query) => {
210                let positive = query.positive.with_column(column.clone());
211                let negative = query.negative.with_column(column);
212                Self::Boost(BoostQuery {
213                    positive: Box::new(positive),
214                    negative: Box::new(negative),
215                    negative_boost: query.negative_boost,
216                })
217            }
218            Self::MultiMatch(query) => {
219                let match_queries = query
220                    .match_queries
221                    .into_iter()
222                    .map(|q| q.with_column(Some(column.clone())))
223                    .collect();
224                Self::MultiMatch(MultiMatchQuery { match_queries })
225            }
226            Self::Boolean(query) => {
227                let must = query
228                    .must
229                    .into_iter()
230                    .map(|q| q.with_column(column.clone()))
231                    .collect();
232                let should = query
233                    .should
234                    .into_iter()
235                    .map(|q| q.with_column(column.clone()))
236                    .collect();
237                let must_not = query
238                    .must_not
239                    .into_iter()
240                    .map(|q| q.with_column(column.clone()))
241                    .collect();
242                Self::Boolean(BooleanQuery {
243                    must,
244                    should,
245                    must_not,
246                })
247            }
248        }
249    }
250}
251
252impl From<MatchQuery> for FtsQuery {
253    fn from(query: MatchQuery) -> Self {
254        Self::Match(query)
255    }
256}
257
258impl From<PhraseQuery> for FtsQuery {
259    fn from(query: PhraseQuery) -> Self {
260        Self::Phrase(query)
261    }
262}
263
264impl From<BoostQuery> for FtsQuery {
265    fn from(query: BoostQuery) -> Self {
266        Self::Boost(query)
267    }
268}
269
270impl From<MultiMatchQuery> for FtsQuery {
271    fn from(query: MultiMatchQuery) -> Self {
272        Self::MultiMatch(query)
273    }
274}
275
276impl From<BooleanQuery> for FtsQuery {
277    fn from(query: BooleanQuery) -> Self {
278        Self::Boolean(query)
279    }
280}
281
282#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
283pub struct MatchQuery {
284    // The column to search in.
285    // If None, it will be determined at query time.
286    pub column: Option<String>,
287    pub terms: String,
288
289    // literal default is not supported so we set it by function
290    #[serde(default = "MatchQuery::default_boost")]
291    pub boost: f32,
292
293    // The max edit distance for fuzzy matching.
294    // If Some(0), it will be exact match.
295    // If None, it will be determined automatically by the rules:
296    // - 0 for terms with length <= 2
297    // - 1 for terms with length <= 5
298    // - 2 for terms with length > 5
299    pub fuzziness: Option<u32>,
300
301    /// The maximum number of terms to expand for fuzzy matching.
302    /// Default to 50.
303    #[serde(default = "MatchQuery::default_max_expansions")]
304    pub max_expansions: usize,
305
306    /// The operator to use for combining terms.
307    /// This can be either `And` or `Or`, it's 'Or' by default.
308    /// - `And`: All terms must match.
309    /// - `Or`: At least one term must match.
310    #[serde(default)]
311    pub operator: Operator,
312
313    /// The number of beginning characters being unchanged for fuzzy matching.
314    /// Default to 0.
315    #[serde(default)]
316    pub prefix_length: u32,
317}
318
319impl MatchQuery {
320    pub fn new(terms: String) -> Self {
321        Self {
322            column: None,
323            terms,
324            boost: 1.0,
325            fuzziness: Some(0),
326            max_expansions: 50,
327            operator: Operator::Or,
328            prefix_length: 0,
329        }
330    }
331
332    pub(crate) fn default_boost() -> f32 {
333        1.0
334    }
335
336    pub(crate) fn default_max_expansions() -> usize {
337        50
338    }
339
340    pub fn with_column(mut self, column: Option<String>) -> Self {
341        self.column = column;
342        self
343    }
344
345    pub fn with_boost(mut self, boost: f32) -> Self {
346        self.boost = boost;
347        self
348    }
349
350    pub fn with_fuzziness(mut self, fuzziness: Option<u32>) -> Self {
351        self.fuzziness = fuzziness;
352        self
353    }
354
355    pub fn with_max_expansions(mut self, max_expansions: usize) -> Self {
356        self.max_expansions = max_expansions;
357        self
358    }
359
360    pub fn with_operator(mut self, operator: Operator) -> Self {
361        self.operator = operator;
362        self
363    }
364
365    pub fn with_prefix_length(mut self, prefix_length: u32) -> Self {
366        self.prefix_length = prefix_length;
367        self
368    }
369
370    pub fn auto_fuzziness(token: &str) -> u32 {
371        match token.len() {
372            0..=2 => 0,
373            3..=5 => 1,
374            _ => 2,
375        }
376    }
377}
378
379impl FtsQueryNode for MatchQuery {
380    fn columns(&self) -> HashSet<String> {
381        let mut columns = HashSet::new();
382        if let Some(column) = &self.column {
383            columns.insert(column.clone());
384        }
385        columns
386    }
387}
388
389#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
390pub struct PhraseQuery {
391    // The column to search in.
392    // If None, it will be determined at query time.
393    pub column: Option<String>,
394    pub terms: String,
395    #[serde(default = "u32::default")]
396    pub slop: u32,
397}
398
399impl PhraseQuery {
400    pub fn new(terms: String) -> Self {
401        Self {
402            column: None,
403            terms,
404            slop: 0,
405        }
406    }
407
408    pub fn with_column(mut self, column: Option<String>) -> Self {
409        self.column = column;
410        self
411    }
412
413    pub fn with_slop(mut self, slop: u32) -> Self {
414        self.slop = slop;
415        self
416    }
417}
418
419impl FtsQueryNode for PhraseQuery {
420    fn columns(&self) -> HashSet<String> {
421        let mut columns = HashSet::new();
422        if let Some(column) = &self.column {
423            columns.insert(column.clone());
424        }
425        columns
426    }
427}
428
429#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
430pub struct BoostQuery {
431    pub positive: Box<FtsQuery>,
432    pub negative: Box<FtsQuery>,
433    #[serde(default = "BoostQuery::default_negative_boost")]
434    pub negative_boost: f32,
435}
436
437impl BoostQuery {
438    pub fn new(positive: FtsQuery, negative: FtsQuery, negative_boost: Option<f32>) -> Self {
439        Self {
440            positive: Box::new(positive),
441            negative: Box::new(negative),
442            negative_boost: negative_boost.unwrap_or(0.5),
443        }
444    }
445
446    fn default_negative_boost() -> f32 {
447        0.5
448    }
449}
450
451impl FtsQueryNode for BoostQuery {
452    fn columns(&self) -> HashSet<String> {
453        let mut columns = self.positive.columns();
454        columns.extend(self.negative.columns());
455        columns
456    }
457}
458
459#[derive(Debug, Clone, PartialEq)]
460pub struct MultiMatchQuery {
461    // each query must be a match query with specified column
462    pub match_queries: Vec<MatchQuery>,
463}
464
465impl Serialize for MultiMatchQuery {
466    fn serialize<S>(&self, serializer: S) -> std::result::Result<S::Ok, S::Error>
467    where
468        S: serde::Serializer,
469    {
470        let mut map = serializer.serialize_map(Some(3))?;
471
472        let query = self.match_queries.first().ok_or(serde::ser::Error::custom(
473            "MultiMatchQuery must have at least one MatchQuery".to_string(),
474        ))?;
475        map.serialize_entry("query", &query.terms)?;
476        let columns = self
477            .match_queries
478            .iter()
479            .map(|q| q.column.as_ref().unwrap().clone())
480            .collect::<Vec<String>>();
481        map.serialize_entry("columns", &columns)?;
482        let boosts = self
483            .match_queries
484            .iter()
485            .map(|q| q.boost)
486            .collect::<Vec<f32>>();
487        map.serialize_entry("boost", &boosts)?;
488        map.end()
489    }
490}
491
492impl<'de> Deserialize<'de> for MultiMatchQuery {
493    fn deserialize<D>(deserializer: D) -> std::result::Result<Self, D::Error>
494    where
495        D: serde::Deserializer<'de>,
496    {
497        #[derive(Deserialize)]
498        struct MultiMatchQueryData {
499            query: String,
500            columns: Vec<String>,
501            boost: Option<Vec<f32>>,
502        }
503
504        let data = MultiMatchQueryData::deserialize(deserializer)?;
505        let boosts = data.boost.unwrap_or(vec![1.0; data.columns.len()]);
506
507        Self::try_new(data.query, data.columns)
508            .map_err(serde::de::Error::custom)?
509            .try_with_boosts(boosts)
510            .map_err(serde::de::Error::custom)
511    }
512}
513
514impl MultiMatchQuery {
515    pub fn try_new(query: String, columns: Vec<String>) -> Result<Self> {
516        if columns.is_empty() {
517            return Err(Error::invalid_input(
518                "Cannot create MultiMatchQuery with no columns".to_string(),
519            ));
520        }
521
522        let match_queries = columns
523            .into_iter()
524            .map(|column| MatchQuery::new(query.clone()).with_column(Some(column)))
525            .collect();
526        Ok(Self { match_queries })
527    }
528
529    pub fn try_with_boosts(mut self, boosts: Vec<f32>) -> Result<Self> {
530        if boosts.len() != self.match_queries.len() {
531            return Err(Error::invalid_input(
532                "The number of boosts must match the number of queries".to_string(),
533            ));
534        }
535
536        for (query, boost) in self.match_queries.iter_mut().zip(boosts) {
537            query.boost = boost;
538        }
539        Ok(self)
540    }
541
542    pub fn with_operator(mut self, operator: Operator) -> Self {
543        for query in &mut self.match_queries {
544            query.operator = operator;
545        }
546        self
547    }
548}
549
550impl FtsQueryNode for MultiMatchQuery {
551    fn columns(&self) -> HashSet<String> {
552        let mut columns = HashSet::with_capacity(self.match_queries.len());
553        for query in &self.match_queries {
554            columns.extend(query.columns());
555        }
556        columns
557    }
558}
559
560pub enum Occur {
561    Should,
562    Must,
563    MustNot,
564}
565
566impl TryFrom<&str> for Occur {
567    type Error = Error;
568    fn try_from(value: &str) -> Result<Self> {
569        match value.to_ascii_uppercase().as_str() {
570            "SHOULD" => Ok(Self::Should),
571            "MUST" => Ok(Self::Must),
572            "MUST_NOT" => Ok(Self::MustNot),
573            _ => Err(Error::invalid_input(format!(
574                "Invalid occur value: {}",
575                value
576            ))),
577        }
578    }
579}
580
581impl From<Occur> for &'static str {
582    fn from(occur: Occur) -> Self {
583        match occur {
584            Occur::Should => "SHOULD",
585            Occur::Must => "MUST",
586            Occur::MustNot => "MUST_NOT",
587        }
588    }
589}
590
591#[derive(Debug, Clone, PartialEq, Serialize, Deserialize)]
592pub struct BooleanQuery {
593    pub should: Vec<FtsQuery>,
594    pub must: Vec<FtsQuery>,
595    pub must_not: Vec<FtsQuery>,
596}
597
598impl BooleanQuery {
599    pub fn new(iter: impl IntoIterator<Item = (Occur, FtsQuery)>) -> Self {
600        let mut should = Vec::new();
601        let mut must = Vec::new();
602        let mut must_not = Vec::new();
603        for (occur, query) in iter {
604            match occur {
605                Occur::Should => should.push(query),
606                Occur::Must => must.push(query),
607                Occur::MustNot => must_not.push(query),
608            }
609        }
610        Self {
611            should,
612            must,
613            must_not,
614        }
615    }
616
617    pub fn with_should(mut self, query: FtsQuery) -> Self {
618        self.should.push(query);
619        self
620    }
621
622    pub fn with_must(mut self, query: FtsQuery) -> Self {
623        self.must.push(query);
624        self
625    }
626
627    pub fn with_must_not(mut self, query: FtsQuery) -> Self {
628        self.must_not.push(query);
629        self
630    }
631}
632
633#[derive(Debug, Clone, PartialEq)]
634#[cfg(test)]
635pub(crate) struct BooleanMatchPlan {
636    pub column: String,
637    pub should: Vec<MatchQuery>,
638    pub must: Vec<MatchQuery>,
639    pub must_not: Vec<MatchQuery>,
640}
641
642#[cfg(test)]
643impl BooleanMatchPlan {
644    pub(crate) fn try_build(query: &FtsQuery) -> Option<Self> {
645        match query {
646            FtsQuery::Match(match_query) => {
647                let mut column = None;
648                let mut should = Vec::new();
649                Self::push_match(&mut should, &mut column, match_query)?;
650                Some(Self {
651                    column: column?,
652                    should,
653                    must: Vec::new(),
654                    must_not: Vec::new(),
655                })
656            }
657            FtsQuery::Boolean(bool_query) => {
658                let mut column = None;
659                let should = Self::collect_matches(&bool_query.should, &mut column)?;
660                let must = Self::collect_matches(&bool_query.must, &mut column)?;
661                let must_not = Self::collect_matches(&bool_query.must_not, &mut column)?;
662
663                if should.is_empty() && must.is_empty() {
664                    return None;
665                }
666                Some(Self {
667                    column: column?,
668                    should,
669                    must,
670                    must_not,
671                })
672            }
673            _ => None,
674        }
675    }
676
677    fn push_match(
678        dest: &mut Vec<MatchQuery>,
679        column: &mut Option<String>,
680        query: &MatchQuery,
681    ) -> Option<()> {
682        let query_column = query.column.as_ref()?;
683        if let Some(existing) = column.as_ref() {
684            if existing != query_column {
685                return None;
686            }
687        } else {
688            *column = Some(query_column.clone());
689        }
690        dest.push(query.clone());
691        Some(())
692    }
693
694    fn collect_matches(
695        queries: &[FtsQuery],
696        column: &mut Option<String>,
697    ) -> Option<Vec<MatchQuery>> {
698        let mut matches = Vec::with_capacity(queries.len());
699        for query in queries {
700            let FtsQuery::Match(match_query) = query else {
701                return None;
702            };
703            Self::push_match(&mut matches, column, match_query)?;
704        }
705        Some(matches)
706    }
707}
708
709impl FtsQueryNode for BooleanQuery {
710    fn columns(&self) -> HashSet<String> {
711        let mut columns = HashSet::new();
712        for query in &self.should {
713            columns.extend(query.columns());
714        }
715        for query in &self.must {
716            columns.extend(query.columns());
717        }
718        for query in &self.must_not {
719            columns.extend(query.columns());
720        }
721        columns
722    }
723}
724
725#[derive(Clone)]
726pub struct Tokens {
727    tokens: Vec<String>,
728    positions: Vec<u32>,
729    tokens_map: HashMap<String, usize>,
730    token_type: DocType,
731}
732
733impl Tokens {
734    pub fn new(tokens: Vec<String>, token_type: DocType) -> Self {
735        let positions = (0..tokens.len() as u32).collect();
736        Self::with_positions(tokens, positions, token_type)
737    }
738
739    pub fn with_positions(tokens: Vec<String>, positions: Vec<u32>, token_type: DocType) -> Self {
740        debug_assert_eq!(tokens.len(), positions.len());
741        let mut tokens_vec = vec![];
742        let mut tokens_map = HashMap::new();
743        for (idx, token) in tokens.into_iter().enumerate() {
744            tokens_vec.push(token.clone());
745            tokens_map.insert(token, idx);
746        }
747
748        Self {
749            tokens: tokens_vec,
750            positions,
751            tokens_map,
752            token_type,
753        }
754    }
755
756    pub fn len(&self) -> usize {
757        self.tokens.len()
758    }
759
760    pub fn is_empty(&self) -> bool {
761        self.tokens.is_empty()
762    }
763
764    pub fn token_type(&self) -> &DocType {
765        &self.token_type
766    }
767
768    pub fn contains(&self, token: &str) -> bool {
769        self.tokens_map.contains_key(token)
770    }
771
772    pub fn token_index(&self, token: &str) -> Option<usize> {
773        self.tokens_map.get(token).copied()
774    }
775
776    pub fn get_token(&self, index: usize) -> &str {
777        &self.tokens[index]
778    }
779
780    pub fn position(&self, index: usize) -> u32 {
781        self.positions[index]
782    }
783}
784
785impl IntoIterator for Tokens {
786    type Item = String;
787    type IntoIter = std::vec::IntoIter<String>;
788
789    fn into_iter(self) -> Self::IntoIter {
790        self.tokens.into_iter()
791    }
792}
793
794impl<'a> IntoIterator for &'a Tokens {
795    type Item = &'a String;
796    type IntoIter = std::slice::Iter<'a, String>;
797
798    fn into_iter(self) -> Self::IntoIter {
799        self.tokens.iter()
800    }
801}
802
803pub fn collect_query_tokens(text: &str, tokenizer: &mut Box<dyn LanceTokenizer>) -> Tokens {
804    let token_type = tokenizer.doc_type();
805    let mut stream = tokenizer.token_stream_for_search(text);
806    let mut tokens = Vec::new();
807    let mut positions = Vec::new();
808    while let Some(token) = stream.next() {
809        tokens.push(token.text.clone());
810        positions.push(token.position as u32);
811    }
812    Tokens::with_positions(tokens, positions, token_type)
813}
814
815pub fn has_query_token(
816    text: &str,
817    tokenizer: &mut Box<dyn LanceTokenizer>,
818    query_tokens: &Tokens,
819) -> bool {
820    let mut stream = tokenizer.token_stream_for_doc(text);
821    while let Some(token) = stream.next() {
822        if query_tokens.contains(&token.text) {
823            return true;
824        }
825    }
826    false
827}
828
829pub fn fill_fts_query_column(
830    query: &FtsQuery,
831    columns: &[String],
832    replace: bool,
833) -> Result<FtsQuery> {
834    if !query.is_missing_column() && !replace {
835        return Ok(query.clone());
836    }
837    match query {
838        FtsQuery::Match(match_query) => {
839            match columns.len() {
840                0 => {
841                    Err(Error::invalid_input("Cannot perform full text search unless an INVERTED index has been created on at least one column".to_string()))
842                }
843                1 => {
844                    let column = columns[0].clone();
845                    let query = match_query.clone().with_column(Some(column));
846                    Ok(FtsQuery::Match(query))
847                }
848                _ => {
849                    // if there are multiple columns, we need to create a MultiMatch query
850                    let multi_match_query =
851                        MultiMatchQuery::try_new(match_query.terms.clone(), columns.to_vec())?;
852                    Ok(FtsQuery::MultiMatch(multi_match_query))
853                }
854            }
855        }
856        FtsQuery::Phrase(phrase_query) => {
857            match columns.len() {
858                0 => {
859                    Err(Error::invalid_input("Cannot perform full text search unless an INVERTED index has been created on at least one column".to_string()))
860                }
861                1 => {
862                    let column = columns[0].clone();
863                    let query = phrase_query.clone().with_column(Some(column));
864                    Ok(FtsQuery::Phrase(query))
865                }
866                _ => {
867                    Err(Error::invalid_input("the column must be specified in the query".to_string()))
868                }
869            }
870        }
871       FtsQuery::Boost(boost_query) => {
872            let positive = fill_fts_query_column(&boost_query.positive, columns, replace)?;
873            let negative = fill_fts_query_column(&boost_query.negative, columns, replace)?;
874            Ok(FtsQuery::Boost(BoostQuery {
875                positive: Box::new(positive),
876                negative: Box::new(negative),
877                negative_boost: boost_query.negative_boost,
878            }))
879        }
880        FtsQuery::MultiMatch(multi_match_query) => {
881            let match_queries = multi_match_query
882                .match_queries
883                .iter()
884                .map(|query| fill_fts_query_column(&FtsQuery::Match(query.clone()), columns, replace))
885                .map(|result| {
886                    result.map(|query| {
887                        if let FtsQuery::Match(match_query) = query {
888                            match_query
889                        } else {
890                            unreachable!("Expected MatchQuery")
891                        }
892                    })
893                })
894                .collect::<Result<Vec<_>>>()?;
895            Ok(FtsQuery::MultiMatch(MultiMatchQuery { match_queries }))
896       }
897        FtsQuery::Boolean(bool_query) => {
898            let must = bool_query
899                .must
900                .iter()
901                .map(|query| fill_fts_query_column(query, columns, replace))
902                .collect::<Result<Vec<_>>>()?;
903            let should = bool_query
904                .should
905                .iter()
906                .map(|query| fill_fts_query_column(query, columns, replace))
907                .collect::<Result<Vec<_>>>()?;
908            let must_not = bool_query
909                .must_not
910                .iter()
911                .map(|query| fill_fts_query_column(query, columns, replace))
912                .collect::<Result<Vec<_>>>()?;
913            Ok(FtsQuery::Boolean(BooleanQuery { must, should, must_not }))
914        }
915    }
916}
917
918#[cfg(test)]
919mod tests {
920    #[test]
921    fn test_match_query_serde() {
922        use super::*;
923        use serde_json::json;
924
925        let query = MatchQuery::new("hello world".to_string())
926            .with_column(Some("text".to_string()))
927            .with_boost(2.0)
928            .with_fuzziness(Some(1))
929            .with_max_expansions(10)
930            .with_operator(Operator::And);
931
932        let serialized = serde_json::to_value(&query).unwrap();
933        let expected = json!({
934            "column": "text",
935            "terms": "hello world",
936            "boost": 2.0,
937            "fuzziness": 1,
938            "max_expansions": 10,
939            "operator": "And",
940            "prefix_length": 0,
941        });
942        assert_eq!(serialized, expected);
943
944        let expected = json!({
945            "column": "text",
946            "terms": "hello world",
947            "fuzziness": 0,
948        });
949        let query = serde_json::from_str::<MatchQuery>(&expected.to_string()).unwrap();
950        assert_eq!(query.column, Some("text".to_owned()));
951        assert_eq!(query.terms, "hello world");
952        assert_eq!(query.boost, 1.0);
953        assert_eq!(query.fuzziness, Some(0));
954        assert_eq!(query.max_expansions, 50);
955        assert_eq!(query.operator, Operator::Or);
956        assert_eq!(query.prefix_length, 0);
957    }
958
959    #[test]
960    fn test_phrase_query_serde() {
961        use super::*;
962        use serde_json::json;
963
964        let query = json!({
965            "terms": "hello world",
966        });
967        let expected = PhraseQuery::new("hello world".to_string());
968        let query: PhraseQuery = serde_json::from_value(query).unwrap();
969        assert_eq!(query, expected);
970
971        let query = json!({
972            "terms": "hello world",
973            "column": "text",
974            "slop": 2,
975        });
976        let expected = PhraseQuery::new("hello world".to_string())
977            .with_column(Some("text".to_string()))
978            .with_slop(2);
979        let query: PhraseQuery = serde_json::from_value(query).unwrap();
980        assert_eq!(query, expected);
981    }
982
983    #[test]
984    fn test_boolean_match_plan_match_query() {
985        use super::*;
986
987        let query = MatchQuery::new("hello".to_string()).with_column(Some("text".to_string()));
988        let plan = BooleanMatchPlan::try_build(&FtsQuery::Match(query.clone())).unwrap();
989        assert_eq!(plan.column, "text");
990        assert_eq!(plan.should, vec![query]);
991        assert!(plan.must.is_empty());
992        assert!(plan.must_not.is_empty());
993    }
994
995    #[test]
996    fn test_boolean_match_plan_boolean_query() {
997        use super::*;
998
999        let should = MatchQuery::new("a".to_string()).with_column(Some("text".to_string()));
1000        let must = MatchQuery::new("b".to_string()).with_column(Some("text".to_string()));
1001        let must_not = MatchQuery::new("c".to_string()).with_column(Some("text".to_string()));
1002        let query = BooleanQuery::new(vec![
1003            (Occur::Should, should.clone().into()),
1004            (Occur::Must, must.clone().into()),
1005            (Occur::MustNot, must_not.clone().into()),
1006        ]);
1007        let plan = BooleanMatchPlan::try_build(&FtsQuery::Boolean(query)).unwrap();
1008        assert_eq!(plan.column, "text");
1009        assert_eq!(plan.should, vec![should]);
1010        assert_eq!(plan.must, vec![must]);
1011        assert_eq!(plan.must_not, vec![must_not]);
1012    }
1013
1014    #[test]
1015    fn test_boolean_match_plan_rejects_mixed_columns() {
1016        use super::*;
1017
1018        let should = MatchQuery::new("a".to_string()).with_column(Some("text".to_string()));
1019        let must = MatchQuery::new("b".to_string()).with_column(Some("title".to_string()));
1020        let query = BooleanQuery::new(vec![
1021            (Occur::Should, should.into()),
1022            (Occur::Must, must.into()),
1023        ]);
1024        assert!(BooleanMatchPlan::try_build(&FtsQuery::Boolean(query)).is_none());
1025    }
1026
1027    #[test]
1028    fn test_boolean_match_plan_rejects_non_match_queries() {
1029        use super::*;
1030
1031        let phrase =
1032            PhraseQuery::new("hello world".to_string()).with_column(Some("text".to_string()));
1033        let query = BooleanQuery::new(vec![(Occur::Should, phrase.into())]);
1034        assert!(BooleanMatchPlan::try_build(&FtsQuery::Boolean(query)).is_none());
1035    }
1036
1037    #[test]
1038    fn test_boolean_match_plan_rejects_only_must_not() {
1039        use super::*;
1040
1041        let must_not = MatchQuery::new("c".to_string()).with_column(Some("text".to_string()));
1042        let query = BooleanQuery::new(vec![(Occur::MustNot, must_not.into())]);
1043        assert!(BooleanMatchPlan::try_build(&FtsQuery::Boolean(query)).is_none());
1044    }
1045
1046    #[test]
1047    fn test_boolean_match_plan_rejects_missing_column() {
1048        use super::*;
1049
1050        let query = MatchQuery::new("hello".to_string());
1051        assert!(BooleanMatchPlan::try_build(&FtsQuery::Match(query)).is_none());
1052    }
1053}
lance_index/scalar/inverted/query.rs

lance_index/scalar/inverted/
query.rs
