574 lines
18 KiB
Rust
574 lines
18 KiB
Rust
use common::search::{
|
|
SchemaFields, json_path_term, normalize_column_name, normalize_exact, tokenize_ngram,
|
|
tokenize_word,
|
|
};
|
|
use common::proto::komp_ac::search::SearchVersionScope;
|
|
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 mode: ConstraintMode,
|
|
}
|
|
|
|
#[derive(Clone, Debug)]
|
|
pub struct SearchConstraintTarget {
|
|
pub table_name: Option<String>,
|
|
pub column: String,
|
|
pub query: String,
|
|
}
|
|
|
|
pub fn build_master_query(
|
|
index: &Index,
|
|
fields: &SchemaFields,
|
|
free_query: &str,
|
|
must: &[SearchConstraint],
|
|
table_filter: Option<&str>,
|
|
version_scope: SearchVersionScope,
|
|
) -> 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)),
|
|
));
|
|
}
|
|
|
|
let archived_rows = |archived| {
|
|
Box::new(TermQuery::new(
|
|
Term::from_field_u64(fields.is_archived, u64::from(archived)),
|
|
IndexRecordOption::Basic,
|
|
)) as Box<dyn Query>
|
|
};
|
|
match version_scope {
|
|
SearchVersionScope::Current => clauses.push((Occur::Must, archived_rows(false))),
|
|
SearchVersionScope::Archived => clauses.push((Occur::Must, archived_rows(true))),
|
|
SearchVersionScope::All => {}
|
|
}
|
|
|
|
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, &target.query)?,
|
|
ConstraintMode::Fuzzy => {
|
|
fuzzy_predicate_scoped(fields, &target.column, &target.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),
|
|
}
|
|
}
|
|
|
|
#[cfg(test)]
|
|
mod tests {
|
|
use super::*;
|
|
use common::search::{archived_search_row_key, create_search_schema, register_tokenizers};
|
|
use tantivy::collector::Count;
|
|
use tantivy::schema::OwnedValue;
|
|
use tantivy::TantivyDocument;
|
|
|
|
struct TestDocument<'a> {
|
|
row_key: String,
|
|
row_id: u64,
|
|
table_name: &'a str,
|
|
values: &'a [(&'a str, &'a str)],
|
|
}
|
|
|
|
fn index_documents(documents: &[TestDocument<'_>]) -> (Index, SchemaFields) {
|
|
let schema = create_search_schema();
|
|
let index = Index::create_in_ram(schema.clone());
|
|
register_tokenizers(&index).expect("tokenizers should register");
|
|
let fields = SchemaFields::from(&schema).expect("schema should match");
|
|
let mut writer = index.writer(50_000_000).expect("writer should open");
|
|
|
|
for source in documents {
|
|
let mut document = TantivyDocument::default();
|
|
document.add_u64(fields.pg_id, source.row_id);
|
|
document.add_u64(fields.row_revision, 1);
|
|
let archive = common::search::parse_archived_search_row_key(&source.row_key);
|
|
document.add_u64(fields.is_archived, u64::from(archive.is_some()));
|
|
document.add_u64(
|
|
fields.table_definition_id,
|
|
archive.map(|item| item.0 as u64).unwrap_or_default(),
|
|
);
|
|
document.add_u64(
|
|
fields.version_number,
|
|
archive.map(|item| item.2 as u64).unwrap_or_default(),
|
|
);
|
|
document.add_text(fields.table_name, source.table_name);
|
|
document.add_text(fields.row_key, &source.row_key);
|
|
let mut object = std::collections::BTreeMap::new();
|
|
for (column, value) in source.values {
|
|
document.add_text(fields.all_text, value);
|
|
object.insert((*column).to_string(), OwnedValue::from(*value));
|
|
}
|
|
document.add_object(fields.data_word, object.clone());
|
|
document.add_object(fields.data_ngram, object.clone());
|
|
document.add_object(fields.data_exact, object);
|
|
writer.add_document(document).expect("document should index");
|
|
}
|
|
writer.commit().expect("documents should commit");
|
|
(index, fields)
|
|
}
|
|
|
|
fn index_versions() -> (Index, SchemaFields) {
|
|
index_documents(&[
|
|
TestDocument {
|
|
row_key: "customers:7".to_string(),
|
|
row_id: 7,
|
|
table_name: "customers",
|
|
values: &[("1", "new")],
|
|
},
|
|
TestDocument {
|
|
row_key: archived_search_row_key(40, 7, 1),
|
|
row_id: 7,
|
|
table_name: "customers",
|
|
values: &[("1", "old")],
|
|
},
|
|
])
|
|
}
|
|
|
|
fn count(index: &Index, query: &dyn Query) -> usize {
|
|
index
|
|
.reader()
|
|
.expect("reader should open")
|
|
.searcher()
|
|
.search(query, &Count)
|
|
.expect("query should run")
|
|
}
|
|
|
|
#[test]
|
|
fn version_scope_separates_current_and_archived_documents() {
|
|
let (index, fields) = index_versions();
|
|
let constraint = SearchConstraint {
|
|
targets: vec![SearchConstraintTarget {
|
|
table_name: None,
|
|
column: "1".to_string(),
|
|
query: "old".to_string(),
|
|
}],
|
|
mode: ConstraintMode::Exact,
|
|
};
|
|
|
|
let current = build_master_query(
|
|
&index,
|
|
&fields,
|
|
"",
|
|
std::slice::from_ref(&constraint),
|
|
Some("customers"),
|
|
SearchVersionScope::Current,
|
|
)
|
|
.expect("current query should build");
|
|
let archived = build_master_query(
|
|
&index,
|
|
&fields,
|
|
"",
|
|
&[constraint],
|
|
Some("customers"),
|
|
SearchVersionScope::Archived,
|
|
)
|
|
.expect("archive query should build");
|
|
|
|
assert_eq!(count(&index, &*current), 0);
|
|
assert_eq!(count(&index, &*archived), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn physical_constraint_survives_any_display_alias_rename() {
|
|
let (index, fields) = index_versions();
|
|
for public_alias in ["adresar", "customer", "renamed_again"] {
|
|
let resolved_physical_column = if !public_alias.is_empty() { "1" } else { unreachable!() };
|
|
let query = build_master_query(
|
|
&index,
|
|
&fields,
|
|
"",
|
|
&[SearchConstraint {
|
|
targets: vec![SearchConstraintTarget {
|
|
table_name: None,
|
|
column: resolved_physical_column.to_string(),
|
|
query: "new".to_string(),
|
|
}],
|
|
mode: ConstraintMode::Exact,
|
|
}],
|
|
Some("customers"),
|
|
SearchVersionScope::Current,
|
|
)
|
|
.expect("renamed alias should resolve to the same physical query");
|
|
assert_eq!(count(&index, &*query), 1, "alias {public_alias}");
|
|
}
|
|
}
|
|
|
|
#[test]
|
|
fn exact_constraints_are_anded_and_table_scoped() {
|
|
let (index, fields) = index_documents(&[
|
|
TestDocument {
|
|
row_key: "customers:1".to_string(),
|
|
row_id: 1,
|
|
table_name: "customers",
|
|
values: &[("1", "Alice Example"), ("2", "Bratislava")],
|
|
},
|
|
TestDocument {
|
|
row_key: "customers:2".to_string(),
|
|
row_id: 2,
|
|
table_name: "customers",
|
|
values: &[("1", "Alice Example"), ("2", "Kosice")],
|
|
},
|
|
TestDocument {
|
|
row_key: "suppliers:3".to_string(),
|
|
row_id: 3,
|
|
table_name: "suppliers",
|
|
values: &[("1", "Alice Example"), ("2", "Bratislava")],
|
|
},
|
|
]);
|
|
let constraints = [
|
|
SearchConstraint {
|
|
targets: vec![SearchConstraintTarget {
|
|
table_name: None,
|
|
column: "1".to_string(),
|
|
query: "Alice Example".to_string(),
|
|
}],
|
|
mode: ConstraintMode::Exact,
|
|
},
|
|
SearchConstraint {
|
|
targets: vec![SearchConstraintTarget {
|
|
table_name: None,
|
|
column: "2".to_string(),
|
|
query: "Bratislava".to_string(),
|
|
}],
|
|
mode: ConstraintMode::Exact,
|
|
},
|
|
];
|
|
let query = build_master_query(
|
|
&index,
|
|
&fields,
|
|
"",
|
|
&constraints,
|
|
Some("customers"),
|
|
SearchVersionScope::Current,
|
|
)
|
|
.expect("exact query should build");
|
|
|
|
assert_eq!(count(&index, &*query), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn fuzzy_constraints_stay_within_the_resolved_physical_column() {
|
|
let (index, fields) = index_documents(&[
|
|
TestDocument {
|
|
row_key: "customers:1".to_string(),
|
|
row_id: 1,
|
|
table_name: "customers",
|
|
values: &[("1", "Alice"), ("2", "Bratislava")],
|
|
},
|
|
TestDocument {
|
|
row_key: "customers:2".to_string(),
|
|
row_id: 2,
|
|
table_name: "customers",
|
|
values: &[("1", "Bratislava"), ("2", "Kosice")],
|
|
},
|
|
]);
|
|
let constraint = |column: &str| SearchConstraint {
|
|
targets: vec![SearchConstraintTarget {
|
|
table_name: None,
|
|
column: column.to_string(),
|
|
query: "Bratislva".to_string(),
|
|
}],
|
|
mode: ConstraintMode::Fuzzy,
|
|
};
|
|
let address_query = build_master_query(
|
|
&index,
|
|
&fields,
|
|
"",
|
|
&[constraint("2")],
|
|
Some("customers"),
|
|
SearchVersionScope::Current,
|
|
)
|
|
.expect("fuzzy address query should build");
|
|
let name_query = build_master_query(
|
|
&index,
|
|
&fields,
|
|
"",
|
|
&[constraint("1")],
|
|
Some("customers"),
|
|
SearchVersionScope::Current,
|
|
)
|
|
.expect("fuzzy name query should build");
|
|
|
|
assert_eq!(count(&index, &*address_query), 1);
|
|
assert_eq!(count(&index, &*name_query), 1);
|
|
}
|
|
|
|
#[test]
|
|
fn free_text_searches_all_public_values_but_respects_table_filter() {
|
|
let (index, fields) = index_documents(&[
|
|
TestDocument {
|
|
row_key: "customers:1".to_string(),
|
|
row_id: 1,
|
|
table_name: "customers",
|
|
values: &[("1", "Alice Example")],
|
|
},
|
|
TestDocument {
|
|
row_key: "suppliers:2".to_string(),
|
|
row_id: 2,
|
|
table_name: "suppliers",
|
|
values: &[("8", "Alice Example")],
|
|
},
|
|
]);
|
|
let query = build_master_query(
|
|
&index,
|
|
&fields,
|
|
"Alice Example",
|
|
&[],
|
|
Some("customers"),
|
|
SearchVersionScope::Current,
|
|
)
|
|
.expect("free-text query should build");
|
|
|
|
assert_eq!(count(&index, &*query), 1);
|
|
}
|
|
}
|