use super::{PublishedSnapshot, QueryCancellation, QueryError, SnapshotLease, Store, StoreError};
use crate::ingest::{
ActiveSegment, DocId, DocumentVersion, GlobalRowId, RowSource, SearchCandidate,
};
use std::cell::Cell;
use std::num::NonZeroUsize;
use std::sync::Arc;
pub type HybridSnippets =
Result<Vec<Option<crate::fts::query::OwnedLexicalSnippet>>, MaterializationError>;
#[derive(Debug)]
pub enum HybridPreparationError<E> {
Preparation(E),
Search(crate::fusion::FusionError),
}
impl<E> From<crate::fusion::FusionError> for HybridPreparationError<E> {
fn from(error: crate::fusion::FusionError) -> Self {
Self::Search(error)
}
}
impl<E> From<QueryError> for HybridPreparationError<E> {
fn from(error: QueryError) -> Self {
Self::Search(error.into())
}
}
impl<E: std::fmt::Display> std::fmt::Display for HybridPreparationError<E> {
fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::Preparation(error) => error.fmt(formatter),
Self::Search(error) => error.fmt(formatter),
}
}
}
impl<E: std::error::Error + 'static> std::error::Error for HybridPreparationError<E> {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Preparation(error) => Some(error),
Self::Search(error) => Some(error),
}
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Default)]
pub(super) struct TestMaterializationWork {
pub(super) admissions: std::sync::atomic::AtomicU64,
pub(super) reverse_version_lookups: std::sync::atomic::AtomicU64,
pub(super) direct_row_lookups: std::sync::atomic::AtomicU64,
#[cfg(feature = "query-timing")]
pub(super) lexical_admission_locks: std::sync::atomic::AtomicU64,
#[cfg(feature = "query-timing")]
pub(super) lexical_admission_lock_nanos: std::sync::atomic::AtomicU64,
}
#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
impl TestMaterializationWork {
pub(super) fn observe_lexical_lock(&self, started: std::time::Instant) {
use std::sync::atomic::Ordering::Relaxed;
self.lexical_admission_locks.fetch_add(1, Relaxed);
self.lexical_admission_lock_nanos.fetch_add(
started.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64,
Relaxed,
);
}
}
#[cfg(any(test, feature = "test-seams"))]
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct MaterializationTestCounters {
pub admissions: u64,
pub reverse_version_lookups: u64,
pub direct_row_lookups: u64,
#[cfg(feature = "query-timing")]
pub lexical_admission_locks: u64,
#[cfg(feature = "query-timing")]
pub lexical_admission_lock_nanos: u64,
}
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct MaterializedRow {
pub document: DocumentVersion,
pub row_id: GlobalRowId,
pub text: String,
}
#[derive(Debug)]
pub enum MaterializationError {
RankOutOfRange {
rank: usize,
returned: usize,
},
MissingIdentity {
row_id: GlobalRowId,
},
MissingText {
row_id: GlobalRowId,
document: DocumentVersion,
},
MissingSource {
row_id: GlobalRowId,
},
InvalidLexicalSource {
source: u32,
},
MissingFusedIdentity {
document: DocId,
},
IdentityMismatch {
row_id: GlobalRowId,
expected: DocumentVersion,
actual: Option<DocumentVersion>,
},
AllocationFailed {
bytes: usize,
},
ArithmeticOverflow,
Query(QueryError),
Storage(StoreError),
}
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub(crate) struct VersionMismatch {
pub(crate) expected: DocumentVersion,
pub(crate) actual: Option<DocumentVersion>,
}
pub(crate) fn require_document_version(
expected: DocumentVersion,
actual: Option<DocumentVersion>,
) -> Result<(), VersionMismatch> {
if actual == Some(expected) {
Ok(())
} else {
Err(VersionMismatch { expected, actual })
}
}
impl std::fmt::Display for MaterializationError {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
match self {
Self::RankOutOfRange { rank, returned } => {
write!(f, "result rank {rank} is outside {returned} returned rows")
}
Self::MissingIdentity { row_id } => {
write!(f, "result row {row_id:?} has no document identity")
}
Self::MissingText { row_id, document } => write!(
f,
"result row {row_id:?} for {document:?} has no stored text"
),
Self::MissingSource { row_id } => {
write!(f, "result row {row_id:?} is outside the admitted snapshot")
}
Self::InvalidLexicalSource { source } => write!(
f,
"lexical result source {source} is outside the admitted assembly"
),
Self::MissingFusedIdentity { document } => write!(
f,
"fused document {document:?} has no ranked source identity"
),
Self::IdentityMismatch {
row_id,
expected,
actual,
} => write!(
f,
"result row {row_id:?} expected {expected:?}, found {actual:?}"
),
Self::AllocationFailed { bytes } => {
write!(f, "result text allocation of {bytes} bytes failed")
}
Self::ArithmeticOverflow => f.write_str("result materialization counter overflow"),
Self::Query(error) => error.fmt(f),
Self::Storage(error) => error.fmt(f),
}
}
}
impl std::error::Error for MaterializationError {
fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
match self {
Self::Storage(error) => Some(error),
Self::Query(error) => Some(error),
_ => None,
}
}
}
pub struct QueryMaterializer<'a> {
#[cfg(any(test, feature = "test-seams"))]
pub(super) store: &'a Store,
pub(super) snapshot: &'a PublishedSnapshot,
pub(super) active: &'a ActiveSegment,
pub(super) cancellation: &'a QueryCancellation<'a>,
pub(super) address: &'a dyn Fn(usize) -> Result<ResultAddress, MaterializationError>,
pub(super) counts: Cell<crate::diag::MaterializationCounters>,
}
#[derive(Clone, Copy, Eq, PartialEq)]
pub(super) struct ResultAddress {
pub(super) row_id: GlobalRowId,
pub(super) document: Option<DocumentVersion>,
}
pub(super) type HybridAddresses = super::stats::Accounted<Vec<(DocId, ResultAddress)>>;
impl From<&SearchCandidate> for ResultAddress {
fn from(candidate: &SearchCandidate) -> Self {
Self {
row_id: candidate.row_id(),
document: candidate.document(),
}
}
}
impl QueryMaterializer<'_> {
pub fn text(&self, rank: usize) -> Result<MaterializedRow, MaterializationError> {
self.check()?;
let address = (self.address)(rank)?;
self.materialize_address(address)
}
#[cfg(any(test, feature = "test-seams"))]
pub fn text_with_test_document(
&self,
rank: usize,
document: Option<DocumentVersion>,
) -> Result<MaterializedRow, MaterializationError> {
self.check()?;
let mut address = (self.address)(rank)?;
address.document = document;
self.materialize_address(address)
}
fn materialize_address(
&self,
address: ResultAddress,
) -> Result<MaterializedRow, MaterializationError> {
let row_id = address.row_id;
let document = address
.document
.ok_or(MaterializationError::MissingIdentity { row_id })?;
#[cfg(any(test, feature = "test-seams"))]
self.store
.text_materialization_work
.direct_row_lookups
.fetch_add(1, std::sync::atomic::Ordering::Relaxed);
let mut counts = self.counts.get();
counts.row_lookups = counts
.row_lookups
.checked_add(1)
.ok_or(MaterializationError::ArithmeticOverflow)?;
self.counts.set(counts);
let row = row_id.local_row() as usize;
let (actual, source) = match row_id.source() {
RowSource::Active => (
self.active.document(row),
self.active
.text(row)
.map_err(MaterializationError::Storage)?,
),
RowSource::Sealed(id) => {
let segment = self
.snapshot
.segments()
.iter()
.find(|segment| segment.meta().id == id)
.ok_or(MaterializationError::MissingSource { row_id })?;
let actual = segment
.document_version(row)
.map_err(StoreError::Segment)
.map_err(MaterializationError::Storage)?;
let source = segment
.query_stored_text()
.map_err(StoreError::Segment)
.map_err(MaterializationError::Storage)?
.and_then(|text| text.row(row).flatten());
(actual, source)
}
};
require_document_version(document, actual).map_err(|mismatch| {
MaterializationError::IdentityMismatch {
row_id,
expected: mismatch.expected,
actual: mismatch.actual,
}
})?;
let source = source.ok_or(MaterializationError::MissingText { row_id, document })?;
let mut text = String::new();
text.try_reserve_exact(source.len()).map_err(|_| {
MaterializationError::AllocationFailed {
bytes: source.len(),
}
})?;
text.push_str(source);
let mut counts = self.counts.get();
counts.text_copies = counts
.text_copies
.checked_add(1)
.ok_or(MaterializationError::ArithmeticOverflow)?;
counts.text_bytes = counts
.text_bytes
.checked_add(source.len() as u64)
.ok_or(MaterializationError::ArithmeticOverflow)?;
self.counts.set(counts);
self.check()?;
Ok(MaterializedRow {
document,
row_id,
text,
})
}
fn check(&self) -> Result<(), MaterializationError> {
self.cancellation
.check_graph()
.map_err(super::map_scan_error)
.map_err(MaterializationError::Query)
}
}
impl Store {
pub fn search_hybrid_with_text_deferred<'vector, E, R>(
&self,
prepare_vector: impl FnOnce() -> Result<crate::ingest::SearchRequest<'vector>, E>,
lexical: &crate::fts::search::TermQuery,
query: &crate::fusion::HybridQuery,
options: impl Into<super::SearchOptions>,
control: super::QueryControl,
finish: impl FnOnce(&crate::ingest::StoreHybridSearchOutcome, &QueryMaterializer<'_>) -> R,
) -> Result<(crate::ingest::StoreHybridSearchOutcome, R), HybridPreparationError<E>> {
self.search_hybrid_pinned_with_text(
prepare_vector,
None,
None,
super::PinnedLexicalQuery::Term(lexical),
query,
options.into(),
control,
finish,
)
}
#[allow(clippy::too_many_arguments)] fn search_hybrid_pinned_with_text<'vector, E, R>(
&self,
prepare_vector: impl FnOnce() -> Result<crate::ingest::SearchRequest<'vector>, E>,
filter: Option<&super::QueryFilter>,
eligible: Option<&[crate::ingest::DocId]>,
lexical: super::PinnedLexicalQuery<'_>,
query: &crate::fusion::HybridQuery,
options: super::SearchOptions,
control: super::QueryControl,
finish: impl FnOnce(&crate::ingest::StoreHybridSearchOutcome, &QueryMaterializer<'_>) -> R,
) -> Result<(crate::ingest::StoreHybridSearchOutcome, R), HybridPreparationError<E>> {
self.search_hybrid_prepared_then(
prepare_vector,
filter,
eligible,
lexical,
query,
options,
control,
true,
None,
|mut outcome, snapshot, active, cancellation, addresses| {
cancellation.check_graph().map_err(super::map_scan_error)?;
let address = |rank| {
let hit =
outcome
.hits
.get(rank)
.ok_or(MaterializationError::RankOutOfRange {
rank,
returned: outcome.hits.len(),
})?;
let addresses = addresses
.ok_or(MaterializationError::MissingFusedIdentity { document: hit.key })?;
let index = addresses
.binary_search_by_key(&hit.key, |(key, _)| *key)
.map_err(|_| MaterializationError::MissingFusedIdentity {
document: hit.key,
})?;
addresses
.get(index)
.map(|(_, address)| *address)
.ok_or(MaterializationError::MissingFusedIdentity { document: hit.key })
};
let materializer = QueryMaterializer {
#[cfg(any(test, feature = "test-seams"))]
store: self,
snapshot,
active,
cancellation,
address: &address,
counts: Cell::new(Default::default()),
};
let value = finish(&outcome, &materializer);
let counts = materializer.counts.get();
outcome.diagnostics.materialization = Some(counts);
cancellation.check_graph().map_err(super::map_scan_error)?;
Ok((outcome, value))
},
)
}
#[cfg(any(test, feature = "test-seams"))]
#[must_use]
pub fn query_materialization_test_counters(&self) -> MaterializationTestCounters {
use std::sync::atomic::Ordering::Relaxed;
MaterializationTestCounters {
admissions: self.text_materialization_work.admissions.load(Relaxed),
reverse_version_lookups: self
.text_materialization_work
.reverse_version_lookups
.load(Relaxed),
direct_row_lookups: self
.text_materialization_work
.direct_row_lookups
.load(Relaxed),
#[cfg(feature = "query-timing")]
lexical_admission_locks: self
.text_materialization_work
.lexical_admission_locks
.load(Relaxed),
#[cfg(feature = "query-timing")]
lexical_admission_lock_nanos: self
.text_materialization_work
.lexical_admission_lock_nanos
.load(Relaxed),
}
}
pub fn search_hybrid_with_text<R>(
&self,
vector: crate::ingest::SearchRequest<'_>,
lexical: &crate::fts::search::TermQuery,
query: &crate::fusion::HybridQuery,
options: impl Into<super::SearchOptions>,
control: super::QueryControl,
finish: impl FnOnce(&crate::ingest::StoreHybridSearchOutcome, &QueryMaterializer<'_>) -> R,
) -> Result<(crate::ingest::StoreHybridSearchOutcome, R), crate::fusion::FusionError> {
self.search_hybrid_pinned_with_text(
|| Ok::<_, std::convert::Infallible>(vector),
vector.filter,
vector.eligible,
super::PinnedLexicalQuery::Term(lexical),
query,
options.into(),
control,
finish,
)
.map_err(|error| match error {
HybridPreparationError::Preparation(never) => match never {},
HybridPreparationError::Search(error) => error,
})
}
pub fn search_lexical_with_snippets(
&self,
query: &crate::fts::search::TermQuery,
k: usize,
window: NonZeroUsize,
control: super::QueryControl,
) -> Result<
(
crate::ingest::StoreLexicalSearchOutcome,
Vec<crate::fts::query::OwnedLexicalSnippet>,
),
crate::ingest::StoreLexicalError,
> {
self.search_lexical_filtered_with_snippets(query, k, window, control, None)
}
pub fn search_lexical_filtered_with_snippets(
&self,
query: &crate::fts::search::TermQuery,
k: usize,
window: NonZeroUsize,
control: super::QueryControl,
filter: Option<&super::QueryFilter>,
) -> Result<
(
crate::ingest::StoreLexicalSearchOutcome,
Vec<crate::fts::query::OwnedLexicalSnippet>,
),
crate::ingest::StoreLexicalError,
> {
self.search_lexical_then(
query,
k,
control,
filter,
|outcome, snapshot, active, assembly, hits, cancellation| {
let mut snippets = Vec::with_capacity(hits.len());
for hit in hits {
cancellation.check_graph().map_err(super::map_scan_error)?;
let (text, _) = super::structured_lexical_row(
snapshot,
active,
&assembly.sources,
hit.doc,
)?;
snippets.push(
crate::fts::snippet::owned_window(
&self.tokenizer,
text,
&query.terms,
window,
)
.ok_or(
crate::ingest::StoreLexicalError::MissingSnippetMatch {
segment: hit.doc.segment,
row: hit.doc.row,
},
)?,
);
}
Ok((outcome, snippets))
},
)
}
pub fn search_hybrid_with_snippets(
&self,
vector: crate::ingest::SearchRequest<'_>,
lexical: &crate::fts::search::TermQuery,
query: &crate::fusion::HybridQuery,
options: impl Into<super::SearchOptions>,
control: super::QueryControl,
window: NonZeroUsize,
) -> Result<(crate::ingest::StoreHybridSearchOutcome, HybridSnippets), crate::fusion::FusionError>
{
self.hybrid_snippets(
vector,
super::PinnedLexicalQuery::Term(lexical),
query,
options.into(),
control,
window,
)
}
pub fn search_hybrid_structured_with_snippets(
&self,
vector: crate::ingest::SearchRequest<'_>,
lexical: &crate::fts::query::LexicalQuery,
query: &crate::fusion::HybridQuery,
options: impl Into<super::SearchOptions>,
control: super::QueryControl,
window: NonZeroUsize,
) -> Result<(crate::ingest::StoreHybridSearchOutcome, HybridSnippets), crate::fusion::FusionError>
{
self.hybrid_snippets(
vector,
super::PinnedLexicalQuery::Structured(lexical),
query,
options.into(),
control,
window,
)
}
fn hybrid_snippets(
&self,
vector: crate::ingest::SearchRequest<'_>,
lexical: super::PinnedLexicalQuery<'_>,
query: &crate::fusion::HybridQuery,
options: super::SearchOptions,
control: super::QueryControl,
window: NonZeroUsize,
) -> Result<(crate::ingest::StoreHybridSearchOutcome, HybridSnippets), crate::fusion::FusionError>
{
let flat_terms = match lexical {
super::PinnedLexicalQuery::Term(query) => Some(&query.terms),
super::PinnedLexicalQuery::Structured(_) => None,
};
self.search_hybrid_pinned_with_text(
|| Ok::<_, std::convert::Infallible>(vector),
vector.filter,
vector.eligible,
lexical,
query,
options,
control,
|outcome, materializer| {
let expanded;
let terms: &[Vec<u8>] = match flat_terms {
Some(terms) => terms,
None => {
expanded = outcome
.lexical_expansions
.iter()
.map(|expansion| expansion.term.clone())
.collect::<Vec<_>>();
&expanded
}
};
let mut snippets = Vec::with_capacity(outcome.hits.len());
for rank in 0..outcome.hits.len() {
let row = match materializer.text(rank) {
Ok(row) => row,
Err(MaterializationError::MissingText { .. }) => {
snippets.push(None);
continue;
}
Err(error) => return Err(error),
};
snippets.push(crate::fts::snippet::owned_window(
&self.tokenizer,
&row.text,
terms,
window,
));
}
Ok(snippets)
},
)
.map_err(|error| match error {
HybridPreparationError::Preparation(never) => match never {},
HybridPreparationError::Search(error) => error,
})
}
pub fn search_lexical_with_text<R>(
&self,
query: &crate::fts::search::TermQuery,
k: usize,
control: super::QueryControl,
finish: impl FnOnce(&crate::ingest::StoreLexicalSearchOutcome, &QueryMaterializer<'_>) -> R,
) -> Result<(crate::ingest::StoreLexicalSearchOutcome, R), crate::ingest::StoreLexicalError>
{
self.search_lexical_then(
query,
k,
control,
None,
|mut outcome, snapshot, active, assembly, hits, cancellation| {
cancellation.check_graph().map_err(super::map_scan_error)?;
let address = |rank| {
let candidate = outcome.candidates.get(rank).ok_or(
MaterializationError::RankOutOfRange {
rank,
returned: outcome.candidates.len(),
},
)?;
let hit = hits.get(rank).ok_or(MaterializationError::RankOutOfRange {
rank,
returned: hits.len(),
})?;
Ok(ResultAddress {
row_id: lexical_row_id(snapshot, &assembly.sources, hit.doc)?,
document: Some(candidate.document),
})
};
let materializer = QueryMaterializer {
#[cfg(any(test, feature = "test-seams"))]
store: self,
snapshot,
active,
cancellation,
address: &address,
counts: Cell::new(Default::default()),
};
let value = finish(&outcome, &materializer);
let counts = materializer.counts.get();
outcome.diagnostics.materialization = Some(counts);
cancellation.check_graph().map_err(super::map_scan_error)?;
Ok((outcome, value))
},
)
}
pub fn search_with_text<R>(
&self,
request: crate::ingest::SearchRequest<'_>,
k: usize,
options: impl Into<super::SearchOptions>,
control: super::QueryControl,
finish: impl FnOnce(&crate::ingest::SearchOutcome, &QueryMaterializer<'_>) -> R,
) -> Result<(crate::ingest::SearchOutcome, R), super::QueryError> {
self.search_with_graph_bound_mode_then(
request,
k,
options.into(),
control,
super::GraphBoundMode::Shared,
|mut outcome, admitted, control| {
let lease =
SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
let cancellation = QueryCancellation::new(control, &lease);
cancellation.check_graph().map_err(super::map_scan_error)?;
let address = |rank| {
outcome.candidates.get(rank).map(ResultAddress::from).ok_or(
MaterializationError::RankOutOfRange {
rank,
returned: outcome.candidates.len(),
},
)
};
let materializer = QueryMaterializer {
#[cfg(any(test, feature = "test-seams"))]
store: self,
counts: Cell::new(Default::default()),
snapshot: &admitted.snapshot,
active: &admitted.active_segment,
cancellation: &cancellation,
address: &address,
};
let value = finish(&outcome, &materializer);
let counts = materializer.counts.get();
outcome.diagnostics.materialization = Some(counts);
cancellation.check_graph().map_err(super::map_scan_error)?;
Ok((outcome, value))
},
)
}
}
fn lexical_row_id(
snapshot: &PublishedSnapshot,
sources: &[super::StructuredLexicalSource],
doc: crate::fts::search::GlobalDocId,
) -> Result<GlobalRowId, MaterializationError> {
let source =
sources
.get(doc.segment as usize)
.ok_or(MaterializationError::InvalidLexicalSource {
source: doc.segment,
})?;
let source = match source {
super::StructuredLexicalSource::Active => RowSource::Active,
super::StructuredLexicalSource::Sealed(ordinal) => RowSource::Sealed(
snapshot
.segments()
.get(*ordinal)
.ok_or(MaterializationError::InvalidLexicalSource {
source: doc.segment,
})?
.meta()
.id,
),
};
Ok(GlobalRowId::new(source, doc.row))
}