292 lines
8.7 KiB
Rust
292 lines
8.7 KiB
Rust
use common::search::{
|
|
json_path_term, normalize_column_name, 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 targets: Vec<SearchConstraintTarget>,
|
|
pub query: String,
|
|
pub mode: ConstraintMode,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct SearchConstraintTarget {
|
|
pub table_name: Option<String>,
|
|
pub column: String,
|
|
}
|
|
|
|
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 = constraint_predicate(fields, constraint)?;
|
|
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::Must, 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 constraint_predicate(
|
|
fields: &SchemaFields,
|
|
constraint: &SearchConstraint,
|
|
) -> Result<Box<dyn Query>, Status> {
|
|
let mut alternatives = Vec::new();
|
|
|
|
for target in &constraint.targets {
|
|
let column_predicate = match constraint.mode {
|
|
ConstraintMode::Exact => exact_predicate(fields, &target.column, &constraint.query)?,
|
|
ConstraintMode::Fuzzy => {
|
|
fuzzy_predicate_scoped(fields, &target.column, &constraint.query)?
|
|
}
|
|
};
|
|
|
|
if let Some(table_name) = &target.table_name {
|
|
let table_term = Term::from_field_text(fields.table_name, table_name);
|
|
alternatives.push((
|
|
Occur::Should,
|
|
Box::new(BooleanQuery::new(vec![
|
|
(
|
|
Occur::Must,
|
|
Box::new(TermQuery::new(table_term, IndexRecordOption::Basic))
|
|
as Box<dyn Query>,
|
|
),
|
|
(Occur::Must, column_predicate),
|
|
])) as Box<dyn Query>,
|
|
));
|
|
} else {
|
|
alternatives.push((Occur::Should, column_predicate));
|
|
}
|
|
}
|
|
|
|
if alternatives.is_empty() {
|
|
return Err(Status::invalid_argument(
|
|
"constraint has no searchable column targets",
|
|
));
|
|
}
|
|
|
|
Ok(Box::new(BooleanQuery::new(alternatives)))
|
|
}
|
|
|
|
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 column = normalize_column_name(column);
|
|
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 column = normalize_column_name(column);
|
|
|
|
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 mut per_word_clauses: Vec<(Occur, Box<dyn Query>)> = Vec::new();
|
|
for word in words {
|
|
let term = Term::from_field_text(fields.all_text, 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, 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 parser = QueryParser::for_index(index, vec![fields.all_text]);
|
|
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.all_text]);
|
|
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),
|
|
}
|
|
}
|