search with multiquery redesigned

search/src/query_builder.rs

@@ -0,0 +1,234 @@
+use common::search::{
+    json_path_term, normalize_exact, tokenize_ngram, tokenize_word, SchemaFields,
+};
+use tantivy::query::{
+    BooleanQuery, BoostQuery, EmptyQuery, FuzzyTermQuery, Occur, PhraseQuery, Query, QueryParser,
+    TermQuery,
+};
+use tantivy::schema::{IndexRecordOption, Term};
+use tantivy::Index;
+use tonic::Status;
+
+#[derive(Clone, Copy, Debug, Eq, PartialEq)]
+pub enum ConstraintMode {
+    Fuzzy,
+    Exact,
+}
+
+#[derive(Clone, Debug)]
+pub struct SearchConstraint {
+    pub column: String,
+    pub query: String,
+    pub mode: ConstraintMode,
+}
+
+pub fn build_master_query(
+    index: &Index,
+    fields: &SchemaFields,
+    free_query: &str,
+    must: &[SearchConstraint],
+    table_filter: Option<&str>,
+) -> Result<Box<dyn Query>, Status> {
+    let mut clauses: Vec<(Occur, Box<dyn Query>)> = Vec::new();
+    let mut has_search_clause = false;
+
+    for constraint in must {
+        let predicate = match constraint.mode {
+            ConstraintMode::Exact => exact_predicate(fields, &constraint.column, &constraint.query)?,
+            ConstraintMode::Fuzzy => {
+                fuzzy_predicate_scoped(fields, &constraint.column, &constraint.query)?
+            }
+        };
+        clauses.push((Occur::Must, predicate));
+        has_search_clause = true;
+    }
+
+    let free_words = tokenize_word(free_query);
+    if !free_words.is_empty() {
+        let predicate = fuzzy_predicate_unscoped(index, fields, &free_words)?;
+        clauses.push((Occur::Should, predicate));
+        has_search_clause = true;
+    }
+
+    if let Some(table_name) = table_filter {
+        let term = Term::from_field_text(fields.table_name, table_name);
+        clauses.push((
+            Occur::Must,
+            Box::new(TermQuery::new(term, IndexRecordOption::Basic)),
+        ));
+    }
+
+    if !has_search_clause {
+        return Ok(Box::new(EmptyQuery));
+    }
+
+    Ok(Box::new(BooleanQuery::new(clauses)))
+}
+
+fn exact_predicate(
+    fields: &SchemaFields,
+    column: &str,
+    query: &str,
+) -> Result<Box<dyn Query>, Status> {
+    let normalized_value = normalize_exact(query);
+    if normalized_value.is_empty() {
+        return Err(Status::invalid_argument(
+            "exact query is empty after normalization",
+        ));
+    }
+
+    let term = json_path_term(fields.data_exact, column, &normalized_value);
+    Ok(Box::new(TermQuery::new(term, IndexRecordOption::Basic)))
+}
+
+fn fuzzy_predicate_scoped(
+    fields: &SchemaFields,
+    column: &str,
+    query: &str,
+) -> Result<Box<dyn Query>, Status> {
+    let words = tokenize_word(query);
+    if words.is_empty() {
+        return Err(Status::invalid_argument(
+            "fuzzy query has no searchable tokens",
+        ));
+    }
+
+    let mut layers: Vec<(Occur, Box<dyn Query>)> = Vec::new();
+
+    let mut per_word_clauses: Vec<(Occur, Box<dyn Query>)> = Vec::new();
+    for word in &words {
+        let term = json_path_term(fields.data_word, column, word);
+        let mut alternates: Vec<(Occur, Box<dyn Query>)> = Vec::new();
+
+        alternates.push((
+            Occur::Should,
+            Box::new(BoostQuery::new(
+                Box::new(TermQuery::new(term.clone(), IndexRecordOption::WithFreqs)),
+                4.0,
+            )),
+        ));
+
+        alternates.push((
+            Occur::Should,
+            Box::new(BoostQuery::new(
+                Box::new(FuzzyTermQuery::new_prefix(term.clone(), 0, false)),
+                3.0,
+            )),
+        ));
+
+        if let Some(distance) = fuzzy_distance(word.chars().count()) {
+            alternates.push((
+                Occur::Should,
+                Box::new(BoostQuery::new(
+                    Box::new(FuzzyTermQuery::new(term.clone(), distance, true)),
+                    2.0,
+                )),
+            ));
+        }
+
+        per_word_clauses.push((Occur::Must, Box::new(BooleanQuery::new(alternates))));
+    }
+    layers.push((Occur::Should, Box::new(BooleanQuery::new(per_word_clauses))));
+
+    if words.len() > 1 {
+        let phrase_terms: Vec<(usize, Term)> = words
+            .iter()
+            .enumerate()
+            .map(|(offset, word)| (offset, json_path_term(fields.data_word, column, word)))
+            .collect();
+        let phrase = PhraseQuery::new_with_offset_and_slop(phrase_terms, 3);
+        layers.push((
+            Occur::Should,
+            Box::new(BoostQuery::new(Box::new(phrase), 2.0)),
+        ));
+    }
+
+    let ngrams = tokenize_ngram(query);
+    if !ngrams.is_empty() {
+        let ngram_clauses: Vec<(Occur, Box<dyn Query>)> = ngrams
+            .into_iter()
+            .map(|gram| {
+                let term = json_path_term(fields.data_ngram, column, &gram);
+                (
+                    Occur::Must,
+                    Box::new(TermQuery::new(term, IndexRecordOption::Basic)) as Box<dyn Query>,
+                )
+            })
+            .collect();
+        layers.push((
+            Occur::Should,
+            Box::new(BoostQuery::new(Box::new(BooleanQuery::new(ngram_clauses)), 1.0)),
+        ));
+    }
+
+    Ok(Box::new(BooleanQuery::new(layers)))
+}
+
+fn fuzzy_predicate_unscoped(
+    index: &Index,
+    fields: &SchemaFields,
+    words: &[String],
+) -> Result<Box<dyn Query>, Status> {
+    let mut layers: Vec<(Occur, Box<dyn Query>)> = Vec::new();
+
+    {
+        let parser = QueryParser::for_index(index, vec![fields.data_word]);
+        let query_string = words
+            .iter()
+            .map(|word| format!("+{}*", word))
+            .collect::<Vec<_>>()
+            .join(" ");
+        if let Ok(query) = parser.parse_query(&query_string) {
+            layers.push((Occur::Should, Box::new(BoostQuery::new(query, 4.0))));
+        }
+    }
+
+    {
+        let parser = QueryParser::for_index(index, vec![fields.data_word]);
+        let query_string = words
+            .iter()
+            .map(|word| match fuzzy_distance(word.chars().count()) {
+                Some(distance) => format!("+{}~{}", word, distance),
+                None => format!("+{}", word),
+            })
+            .collect::<Vec<_>>()
+            .join(" ");
+        if let Ok(query) = parser.parse_query(&query_string) {
+            layers.push((Occur::Should, Box::new(BoostQuery::new(query, 2.0))));
+        }
+    }
+
+    if words.len() > 1 {
+        let parser = QueryParser::for_index(index, vec![fields.data_word]);
+        let query_string = format!("\"{}\"~3", words.join(" "));
+        if let Ok(query) = parser.parse_query(&query_string) {
+            layers.push((Occur::Should, Box::new(BoostQuery::new(query, 2.0))));
+        }
+    }
+
+    {
+        let parser = QueryParser::for_index(index, vec![fields.data_ngram]);
+        let query_string = words
+            .iter()
+            .map(|word| format!("+{}", word))
+            .collect::<Vec<_>>()
+            .join(" ");
+        if let Ok(query) = parser.parse_query(&query_string) {
+            layers.push((Occur::Should, Box::new(BoostQuery::new(query, 1.0))));
+        }
+    }
+
+    if layers.is_empty() {
+        return Ok(Box::new(EmptyQuery));
+    }
+
+    Ok(Box::new(BooleanQuery::new(layers)))
+}
+
+fn fuzzy_distance(word_len: usize) -> Option<u8> {
+    match word_len {
+        0..=3 => None,
+        4..=6 => Some(1),
+        _ => Some(2),
+    }
+}