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

lance_index/scalar/inverted/
query.rs
