#![allow(clippy::expect_used, clippy::indexing_slicing, clippy::panic)]
use proptest::prelude::*;
use proptest::test_runner::{Config as ProptestConfig, RngSeed};
use zeppelin_embed::fts::dict::{TermDictionary, TermInfo};
use zeppelin_embed::fts::index::{DEFAULT_FIELD, Document, SegmentIndex};
use zeppelin_embed::fts::phrase::{PhraseQuery, search_segment, streams_match};
use zeppelin_embed::fts::snippet::{best_window, match_offsets};
use zeppelin_embed::fts::tokenizer::{Analyzer, Profile, TokenizerConfig};
use zeppelin_embed::fts::{fuzzy, prefix};
fn code_analyzer() -> Analyzer {
Analyzer::new(Profile::Code.config()).expect("valid configuration")
}
fn text_analyzer() -> Analyzer {
Analyzer::new(TokenizerConfig::text_default()).expect("valid configuration")
}
#[test]
fn terms_with_prefix_finds_meeting_and_uses_divided_prefix_boosts() {
use zeppelin_embed::fts::query::LexicalQuery;
use zeppelin_embed::fts::search::{FieldWeights, TermQuery};
use zeppelin_embed::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
use zeppelin_embed::lifecycle::{CancelToken, OpenOptions, QueryControl, Store};
let directory = tempfile::tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("store");
store
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(1), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("meeting meek"),
IngestDocument::new(
DocumentVersion::new(DocId::new(2), Revision::new(1)),
vec![0.0, 1.0],
)
.with_text("meeting"),
]))
.expect("ingest");
store.seal().expect("seal");
let prefixed = LexicalQuery::TermsWithPrefix {
terms: Vec::new(),
prefix: b"mee".to_vec(),
fields: FieldWeights::flat(&[DEFAULT_FIELD]),
};
let result = store
.search_lexical_structured(&prefixed, 10, 64, QueryControl::Cancel(CancelToken::new()))
.expect("prefix search");
let meeting = result
.candidates
.iter()
.find(|candidate| candidate.document.doc_id() == DocId::new(1))
.expect("meeting document");
let exact_score = |term: &[u8]| {
store
.search_lexical(
&TermQuery::flat(vec![term.to_vec()], &[DEFAULT_FIELD]),
10,
QueryControl::Cancel(CancelToken::new()),
)
.expect("exact term")
.candidates
.into_iter()
.find(|candidate| candidate.document.doc_id() == DocId::new(1))
.expect("exact hit")
.score
};
let expected = (exact_score(b"meek") + exact_score(b"meet")) / 2.0;
assert_eq!(meeting.score.to_bits(), expected.to_bits());
assert_eq!(
result
.expansions
.iter()
.map(|expansion| (expansion.term.as_slice(), expansion.boost_thousandths))
.collect::<Vec<_>>(),
vec![(b"meek".as_slice(), 500), (b"meet".as_slice(), 500)]
);
store.close().expect("close");
}
#[test]
fn reverse_prefix_expansion_prevents_stemming_flicker() {
use zeppelin_embed::fts::query::LexicalQuery;
use zeppelin_embed::fts::search::FieldWeights;
use zeppelin_embed::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
use zeppelin_embed::lifecycle::{CancelToken, OpenOptions, QueryControl, Store};
let directory = tempfile::tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("store");
store
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(1), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("meeting"),
]))
.expect("ingest");
store.seal().expect("seal");
for prefix in [b"meeti".as_slice(), b"meetin"] {
let query = LexicalQuery::TermsWithPrefix {
terms: Vec::new(),
prefix: prefix.to_vec(),
fields: FieldWeights::flat(&[DEFAULT_FIELD]),
};
let result = store
.search_lexical_structured(&query, 10, 64, QueryControl::Cancel(CancelToken::new()))
.expect("reverse prefix search");
assert_eq!(
result
.candidates
.iter()
.map(|candidate| candidate.document.doc_id())
.collect::<Vec<_>>(),
vec![DocId::new(1)]
);
assert_eq!(
result
.expansions
.iter()
.map(|expansion| expansion.term.as_slice())
.collect::<Vec<_>>(),
vec![b"meet".as_slice()]
);
}
store.close().expect("close");
}
#[test]
fn astra_12_vocabulary_reused_until_exact_inputs_change() {
use zeppelin_embed::fts::{preparation_observer as observer, query::LexicalQuery};
use zeppelin_embed::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
use zeppelin_embed::lifecycle::{CancelToken, OpenOptions, QueryControl, Store};
let directory = tempfile::tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("store");
let ingest = |id, text: &str| {
store
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(id), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text(text),
]))
.expect("ingest")
};
ingest(1, "alpha alphabet beta");
store.seal().expect("seal");
let query = LexicalQuery::prefix(b"alph".to_vec(), DEFAULT_FIELD);
let search = || {
observer::begin();
let result = store
.search_lexical_structured(&query, 10, 64, QueryControl::Cancel(CancelToken::new()))
.expect("prefix");
let work = observer::vocabulary_work();
observer::take();
(result, work)
};
let (first, cold) = search();
let (second, warm) = search();
assert_eq!(first.candidates, second.candidates);
assert_eq!(cold.builds, 1);
assert_eq!(
warm.builds, 0,
"unchanged immutable inputs rebuilt vocabulary: {warm:?}"
);
assert_eq!(warm.copied_bytes, 0);
ingest(2, "alphonse gamma");
let (updated, changed) = search();
assert_eq!(changed.builds, 1);
assert_eq!(updated.candidates.len(), 2);
assert_eq!(search().1.builds, 0);
store.seal().expect("publish new snapshot");
assert_eq!(search().1.builds, 1);
assert_eq!(search().1.builds, 0);
store
.delete(zeppelin_embed::ingest::DeleteBatch::new(vec![DocId::new(
2,
)]))
.expect("tombstone");
let (deleted, work) = search();
assert_eq!(work.builds, 1);
assert_eq!(deleted.candidates.len(), 1);
assert_eq!(
deleted.expansions, updated.expansions,
"tombstones preserve all analyzed vocabulary expansions"
);
let purge = store.purge(&[DocId::new(2)]).expect("purge");
store.await_physical_purge(purge).expect("physical purge");
let (purged, work) = search();
assert_eq!(work.builds, 1);
assert_eq!(
purged.expansions, first.expansions,
"physical purge removes the deleted source vocabulary"
);
assert_eq!(store.stats().expect("stats").temporary_bytes, 0);
store.close().expect("close");
}
#[test]
fn astra_12_term_and_phrase_still_build_no_vocabulary() {
use zeppelin_embed::fts::{
preparation_observer as observer, query::LexicalQuery, search::TermQuery,
};
use zeppelin_embed::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
use zeppelin_embed::lifecycle::{CancelToken, OpenOptions, QueryControl, Store};
let directory = tempfile::tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("store");
store
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(1), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("alpha beta"),
]))
.expect("ingest");
store.seal().expect("seal");
for query in [
LexicalQuery::term(TermQuery::flat(vec![b"alpha".to_vec()], &[DEFAULT_FIELD])),
LexicalQuery::phrase(vec![b"alpha".to_vec(), b"beta".to_vec()], 0, DEFAULT_FIELD),
] {
observer::begin();
let result = store
.search_lexical_structured(&query, 1, 64, QueryControl::Cancel(CancelToken::new()))
.expect("query");
assert_eq!(result.candidates.len(), 1);
assert_eq!(
observer::vocabulary_work(),
observer::VocabularyWork::default()
);
observer::take();
}
store.close().expect("close");
}
#[test]
fn astra_14_phonetic_encodes_vocabulary_once() {
use zeppelin_embed::fts::{preparation_observer as observer, query::LexicalQuery};
use zeppelin_embed::ingest::{DocId, DocumentVersion, IngestBatch, IngestDocument, Revision};
use zeppelin_embed::lifecycle::{CancelToken, OpenOptions, QueryControl, Store};
let directory = tempfile::tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("store");
store
.ingest(IngestBatch::new(vec![
IngestDocument::new(
DocumentVersion::new(DocId::new(1), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("smith smyth night knight robert rupert alpha"),
]))
.expect("ingest");
store.seal().expect("seal");
let query = LexicalQuery::phonetic(b"night".to_vec(), DEFAULT_FIELD);
let search = || {
observer::begin();
let result = store
.search_lexical_structured(&query, 10, 64, QueryControl::Cancel(CancelToken::new()))
.expect("phonetic query");
let encodings = observer::phonetic_encoding_calls();
observer::take();
(result, encodings)
};
let (first, cold) = search();
let (second, warm) = search();
assert_eq!(first.candidates, second.candidates);
assert_eq!(
first
.expansions
.iter()
.map(|e| e.term.as_slice())
.collect::<Vec<_>>(),
vec![b"knight".as_slice(), b"night".as_slice()]
);
assert!(cold > 1, "first query builds the vocabulary map");
assert_eq!(
warm, 1,
"unchanged vocabulary was re-encoded: cold={cold}, warm={warm}"
);
store.close().expect("close");
}
const WORDS: [&str; 12] = [
"alpha",
"beta",
"gamma",
"delta",
"alpha",
"beta",
"receive",
"receipt",
"engine",
"engineer",
"put_if_match",
"Caf\u{00E9}",
];
const UNICODE_FRAGMENTS: [&str; 8] = [
"Caf\u{00E9}",
"cafe\u{0301}",
"na\u{00EF}ve",
"\u{1F680}",
"\u{4E2D}\u{6587}",
"stra\u{00DF}e",
"don\u{2019}t",
"\u{FB01}nd",
];
fn document_text() -> impl Strategy<Value = String> {
prop::collection::vec(
prop_oneof![
prop::sample::select(WORDS.as_slice()),
prop::sample::select(UNICODE_FRAGMENTS.as_slice()),
],
0..10,
)
.prop_map(|words| words.join(" "))
}
fn dictionary_from(terms: &[String]) -> TermDictionary {
let mut sorted: Vec<&str> = terms.iter().map(String::as_str).collect();
sorted.sort_unstable();
sorted.dedup();
let mut dictionary = TermDictionary::default();
for term in sorted {
if term.is_empty() {
continue;
}
let _ = dictionary.push(term.as_bytes(), TermInfo::default());
}
dictionary
}
fn levenshtein(left: &[u8], right: &[u8]) -> u32 {
let mut row: Vec<u32> = (0..=right.len()).map(|index| index as u32).collect();
for (i, a) in left.iter().enumerate() {
let mut next = vec![i as u32 + 1];
for (j, b) in right.iter().enumerate() {
let cost = u32::from(a != b);
next.push((row[j + 1] + 1).min(next[j] + 1).min(row[j] + cost));
}
row = next;
}
row.last().copied().unwrap_or(u32::MAX)
}
fn naive_phrase_match(streams: &[Vec<u32>], slop: u32) -> bool {
fn walk(
streams: &[Vec<u32>],
index: usize,
anchor: u32,
last: u32,
cost: u32,
slop: u32,
) -> bool {
let Some(stream) = streams.get(index) else {
return cost <= slop;
};
for position in stream {
if *position < last {
continue;
}
let expected = u32::try_from(index).unwrap_or(u32::MAX);
let relative = position.saturating_sub(anchor);
let added = relative.abs_diff(expected);
let total = cost.saturating_add(added);
if total > slop {
continue;
}
if walk(streams, index + 1, anchor, *position, total, slop) {
return true;
}
}
false
}
let Some(first) = streams.first() else {
return false;
};
if streams.iter().any(Vec::is_empty) {
return false;
}
first
.iter()
.any(|anchor| walk(streams, 1, *anchor, *anchor, 0, slop))
}
proptest! {
#![proptest_config(ProptestConfig {
rng_seed: RngSeed::Fixed(0x7ec4_0c15_0001),
cases: 512,
..ProptestConfig::default()
})]
#[test]
fn phrase_matches_iff_naive_token_stream_scan_matches(
streams in prop::collection::vec(
prop::collection::vec(0_u32..12, 0..4).prop_map(|mut positions| {
positions.sort_unstable();
positions.dedup();
positions
}),
1..4,
),
slop in 0_u32..5,
) {
prop_assert_eq!(
streams_match(&streams, slop),
naive_phrase_match(&streams, slop),
"streams {:?} at slop {}",
streams,
slop
);
}
#[test]
fn indexed_phrase_agrees_with_the_model_on_analyzed_documents(
texts in prop::collection::vec(document_text(), 1..6),
query in prop::collection::vec(prop::sample::select(WORDS.as_slice()), 1..3),
slop in 0_u32..4,
) {
let analyzer = code_analyzer();
let mut segment = SegmentIndex::new();
for text in &texts {
segment
.push_document(&analyzer, &Document::with_text(text))
.expect("indexable");
}
let mut terms: Vec<Vec<u8>> = Vec::new();
for word in &query {
if let Some(token) = analyzer.analyze(word).into_iter().next() {
terms.push(token.term.into_bytes());
}
}
prop_assume!(!terms.is_empty());
let phrase = PhraseQuery {
terms: terms.clone(),
slop,
field: DEFAULT_FIELD,
};
let actual = search_segment(&segment, &phrase).expect("valid phrase");
let mut expected: Vec<u32> = Vec::new();
for row in 0..segment.row_count() {
let mut streams: Vec<Vec<u32>> = Vec::new();
let mut complete = true;
for term in &terms {
let Some(list) = segment.posting_list(term, DEFAULT_FIELD) else {
complete = false;
break;
};
match list.postings().iter().find(|posting| posting.docid == row) {
Some(posting) => streams.push(posting.positions.clone()),
None => {
complete = false;
break;
}
}
}
if complete && naive_phrase_match(&streams, slop) {
expected.push(row);
}
}
prop_assert_eq!(actual, expected);
}
#[test]
fn prefix_results_equal_brute_force_over_the_dictionary(
terms in prop::collection::vec(
prop::sample::select(WORDS.as_slice()).prop_map(str::to_owned),
1..12,
),
length in 1_usize..5,
) {
let analyzed: Vec<String> = {
let analyzer = code_analyzer();
let mut out = Vec::new();
for term in &terms {
for token in analyzer.analyze(term) {
out.push(token.term);
}
}
out
};
prop_assume!(!analyzed.is_empty());
let dictionary = dictionary_from(&analyzed);
prop_assume!(!dictionary.is_empty());
let Some(sample) = dictionary.term_at(0) else {
return Ok(());
};
let cut = length.min(sample.len()).max(1);
let Some(prefix_bytes) = sample.get(..cut) else {
return Ok(());
};
let actual: Vec<Vec<u8>> = prefix::search(&dictionary, prefix_bytes)
.expect("non-empty prefix")
.into_iter()
.map(|entry| entry.term)
.collect();
let expected: Vec<Vec<u8>> = (0..dictionary.len())
.filter_map(|index| dictionary.term_at(index))
.filter(|term| term.starts_with(prefix_bytes))
.map(<[u8]>::to_vec)
.collect();
prop_assert_eq!(actual, expected);
}
#[test]
fn fuzzy_results_equal_brute_force_edit_distance_over_the_dictionary(
terms in prop::collection::vec(
prop::sample::select(WORDS.as_slice()).prop_map(str::to_owned),
1..12,
),
query in prop::sample::select(WORDS.as_slice()),
max_distance in 1_u32..3,
) {
let analyzer = code_analyzer();
let mut analyzed: Vec<String> = Vec::new();
for term in &terms {
for token in analyzer.analyze(term) {
analyzed.push(token.term);
}
}
prop_assume!(!analyzed.is_empty());
let dictionary = dictionary_from(&analyzed);
prop_assume!(!dictionary.is_empty());
let Some(query_token) = analyzer.analyze(query).into_iter().next() else {
return Ok(());
};
let query_bytes = query_token.term.into_bytes();
prop_assume!(!query_bytes.is_empty());
let (candidates, counters) = fuzzy::search(&dictionary, &query_bytes, max_distance);
let mut actual: Vec<Vec<u8>> =
candidates.into_iter().map(|candidate| candidate.term).collect();
actual.sort();
let mut expected: Vec<Vec<u8>> = (0..dictionary.len())
.filter_map(|index| dictionary.term_at(index))
.filter(|term| {
term.first() == query_bytes.first()
&& levenshtein(&query_bytes, term) <= max_distance
})
.map(<[u8]>::to_vec)
.collect();
expected.sort();
prop_assert_eq!(actual, expected);
prop_assert!(
counters.entries_visited <= dictionary.len() as u64,
"the walk visited more entries than exist"
);
}
#[test]
fn snippet_ranges_are_in_bounds_utf8_valid_and_contain_the_matched_surface_form(
text in document_text(),
word in prop::sample::select(WORDS.as_slice()),
) {
let analyzer = text_analyzer();
let Some(token) = analyzer.analyze(word).into_iter().next() else {
return Ok(());
};
let terms = vec![token.term.into_bytes()];
for highlight in match_offsets(&analyzer, &text, &terms) {
let start = usize::try_from(highlight.start).unwrap_or(usize::MAX);
let end = usize::try_from(highlight.end).unwrap_or(usize::MAX);
prop_assert!(end <= text.len(), "highlight escaped the field");
prop_assert!(text.is_char_boundary(start), "highlight split a character");
prop_assert!(text.is_char_boundary(end), "highlight split a character");
prop_assert!(text.get(start..end).is_some());
}
if let Some(snippet) = best_window(&analyzer, &text, &terms, 24, true).expect("stored") {
let window = snippet.text(&text);
prop_assert!(window.is_some(), "window is not a character boundary");
for highlight in &snippet.highlights {
prop_assert!(highlight.start >= snippet.window.start);
prop_assert!(highlight.end <= snippet.window.end);
}
}
}
#[test]
fn snippets_are_deterministic_across_runs(
text in document_text(),
word in prop::sample::select(WORDS.as_slice()),
) {
let analyzer = text_analyzer();
let Some(token) = analyzer.analyze(word).into_iter().next() else {
return Ok(());
};
let terms = vec![token.term.into_bytes()];
let first = best_window(&analyzer, &text, &terms, 24, true).expect("stored");
let second = best_window(&analyzer, &text, &terms, 24, true).expect("stored");
prop_assert_eq!(first, second);
}
}
#[test]
fn ze263_prefix_cap_selects_highest_live_df() {
use zeppelin_embed::fts::query::{LexicalMatchKind, LexicalQuery};
use zeppelin_embed::fts::search::FieldWeights;
use zeppelin_embed::ingest::{
DeleteBatch, DocId, DocumentVersion, IngestBatch, IngestDocument, Revision,
};
use zeppelin_embed::lifecycle::{CancelToken, OpenOptions, QueryControl, Store};
let directory = tempfile::tempdir().expect("directory");
let store = Store::open(directory.path(), OpenOptions::default()).expect("store");
let mut documents = Vec::new();
for row in 0..5 {
let terms = (0..40)
.filter(|term| row < if *term >= 20 { 4 } else { 1 })
.map(|term| format!("topic{term:02}"))
.collect::<Vec<_>>();
documents.push(
IngestDocument::new(
DocumentVersion::new(DocId::new(row), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text(terms.join(" ")),
);
}
documents.push(
IngestDocument::new(
DocumentVersion::new(DocId::new(99), Revision::new(1)),
vec![1.0, 0.0],
)
.with_text("topicdead"),
);
store.ingest(IngestBatch::new(documents)).expect("ingest");
store.seal().expect("seal");
store
.delete(DeleteBatch::new(vec![DocId::new(99)]))
.expect("delete");
let query = LexicalQuery::TermsWithPrefix {
terms: vec![b"topic39".to_vec(), b"topic39".to_vec()],
prefix: b"top".to_vec(),
fields: FieldWeights::flat(&[DEFAULT_FIELD]),
};
let result = store
.search_lexical_structured(&query, 10, 64, QueryControl::Cancel(CancelToken::new()))
.expect("search");
assert_eq!(result.expansions.len(), 34);
assert!(
result.expansions[..2]
.iter()
.all(|e| e.kind == LexicalMatchKind::Term
&& e.boost_thousandths == 1000
&& e.term == b"topic39")
);
let expected = std::iter::once(b"topic".to_vec())
.chain(
(0..11)
.chain(20..40)
.map(|term| format!("topic{term:02}").into_bytes()),
)
.collect::<Vec<_>>();
assert_eq!(
result.expansions[2..]
.iter()
.map(|e| e.term.clone())
.collect::<Vec<_>>(),
expected
);
assert!(
result.expansions[2..]
.iter()
.all(|e| e.boost_thousandths == 1000 / 32)
);
let exhaustive = store
.search_lexical_structured(
&LexicalQuery::prefix(b"top".to_vec(), DEFAULT_FIELD),
1,
64,
QueryControl::Cancel(CancelToken::new()),
)
.expect("explicit exhaustive prefix");
assert_eq!(exhaustive.expansions.len(), 42);
}