zeppelin_embed/lifecycle/
materialize.rs

1//! Scoped ownership of text returned from a ranked store query.
2
3use super::{PublishedSnapshot, QueryCancellation, QueryError, SnapshotLease, Store, StoreError};
4use crate::ingest::{
5    ActiveSegment, DocId, DocumentVersion, GlobalRowId, RowSource, SearchCandidate,
6};
7use std::cell::Cell;
8use std::num::NonZeroUsize;
9use std::sync::Arc;
10
11/// Per-hit hybrid snippets in rank order, or the failure to read a ranked
12/// row's text. `None` marks a hit without stored text or without a match.
13pub type HybridSnippets =
14    Result<Vec<Option<crate::fts::query::OwnedLexicalSnippet>>, MaterializationError>;
15
16/// Preserves a deferred vector producer's typed error separately from search.
17#[derive(Debug)]
18pub enum HybridPreparationError<E> {
19    /// The caller could not produce a validated query vector.
20    Preparation(E),
21    /// Admission, retrieval, fusion, cancellation or materialization failed.
22    Search(crate::fusion::FusionError),
23}
24
25impl<E> From<crate::fusion::FusionError> for HybridPreparationError<E> {
26    fn from(error: crate::fusion::FusionError) -> Self {
27        Self::Search(error)
28    }
29}
30
31impl<E> From<QueryError> for HybridPreparationError<E> {
32    fn from(error: QueryError) -> Self {
33        Self::Search(error.into())
34    }
35}
36
37impl<E: std::fmt::Display> std::fmt::Display for HybridPreparationError<E> {
38    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
39        match self {
40            Self::Preparation(error) => error.fmt(formatter),
41            Self::Search(error) => error.fmt(formatter),
42        }
43    }
44}
45
46impl<E: std::error::Error + 'static> std::error::Error for HybridPreparationError<E> {
47    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
48        match self {
49            Self::Preparation(error) => Some(error),
50            Self::Search(error) => Some(error),
51        }
52    }
53}
54
55#[cfg(any(test, feature = "test-seams"))]
56#[derive(Default)]
57pub(super) struct TestMaterializationWork {
58    pub(super) admissions: std::sync::atomic::AtomicU64,
59    pub(super) reverse_version_lookups: std::sync::atomic::AtomicU64,
60    pub(super) direct_row_lookups: std::sync::atomic::AtomicU64,
61    #[cfg(feature = "query-timing")]
62    pub(super) lexical_admission_locks: std::sync::atomic::AtomicU64,
63    #[cfg(feature = "query-timing")]
64    pub(super) lexical_admission_lock_nanos: std::sync::atomic::AtomicU64,
65}
66
67#[cfg(all(feature = "query-timing", any(test, feature = "test-seams")))]
68impl TestMaterializationWork {
69    pub(super) fn observe_lexical_lock(&self, started: std::time::Instant) {
70        use std::sync::atomic::Ordering::Relaxed;
71        self.lexical_admission_locks.fetch_add(1, Relaxed);
72        self.lexical_admission_lock_nanos.fetch_add(
73            started.elapsed().as_nanos().min(u128::from(u64::MAX)) as u64,
74            Relaxed,
75        );
76    }
77}
78
79/// Literal lifecycle and source-lookup observations for directed query tests.
80#[cfg(any(test, feature = "test-seams"))]
81#[derive(Clone, Copy, Debug, Eq, PartialEq)]
82pub struct MaterializationTestCounters {
83    /// Invocations of vector or lexical admission, including standalone stored_text.
84    pub admissions: u64,
85    /// Invocations of active or sealed reverse document-version lookup.
86    pub reverse_version_lookups: u64,
87    /// Ranked physical rows requested through the scoped materializer.
88    pub direct_row_lookups: u64,
89    /// Measured state, active, and snapshot lock acquisitions in lexical admission.
90    #[cfg(feature = "query-timing")]
91    pub lexical_admission_locks: u64,
92    /// Wall nanoseconds in those acquisitions, including uncontended call overhead.
93    #[cfg(feature = "query-timing")]
94    pub lexical_admission_lock_nanos: u64,
95}
96
97/// One owned source row corresponding to a ranked query hit.
98#[derive(Clone, Debug, Eq, PartialEq)]
99pub struct MaterializedRow {
100    /// The exact document version whose row was scored.
101    pub document: DocumentVersion,
102    /// Original physical row coordinates within the retrieval snapshot.
103    pub row_id: GlobalRowId,
104    /// One owned copy of the stored UTF-8 body.
105    pub text: String,
106}
107
108/// A ranked row could not be materialized faithfully.
109#[derive(Debug)]
110pub enum MaterializationError {
111    /// The caller requested a rank outside the returned list.
112    RankOutOfRange {
113        /// Zero-based rank requested by the caller.
114        rank: usize,
115        /// Number of ranked hits available.
116        returned: usize,
117    },
118    /// The ranked row has no application identity.
119    MissingIdentity {
120        /// Physical row lacking an application identity.
121        row_id: GlobalRowId,
122    },
123    /// The ranked document has no stored text.
124    MissingText {
125        /// Physical row lacking stored text.
126        row_id: GlobalRowId,
127        /// Ranked application identity.
128        document: DocumentVersion,
129    },
130    /// The ranked physical source is absent from the admitted snapshot.
131    MissingSource {
132        /// Physical row whose source is absent.
133        row_id: GlobalRowId,
134    },
135    /// A lexical source ordinal has no physical source in the admitted assembly.
136    InvalidLexicalSource {
137        /// Source ordinal from the ranked lexical hit.
138        source: u32,
139    },
140    /// A fused application key has no retained source identity.
141    MissingFusedIdentity {
142        /// Fused application key lacking source coordinates.
143        document: DocId,
144    },
145    /// The physical row does not contain the version that was ranked.
146    IdentityMismatch {
147        /// Physical row that failed identity validation.
148        row_id: GlobalRowId,
149        /// Application identity retained during ranking.
150        expected: DocumentVersion,
151        /// Application identity found at the physical row.
152        actual: Option<DocumentVersion>,
153    },
154    /// The result could not reserve its owned text buffer.
155    AllocationFailed {
156        /// UTF-8 text bytes requested from the allocator.
157        bytes: usize,
158    },
159    /// A materialization work count overflowed.
160    ArithmeticOverflow,
161    /// Cancellation or deadline prevented complete materialization.
162    Query(QueryError),
163    /// Storage validation or lifecycle failure.
164    Storage(StoreError),
165}
166
167#[derive(Clone, Copy, Debug, Eq, PartialEq)]
168pub(crate) struct VersionMismatch {
169    pub(crate) expected: DocumentVersion,
170    pub(crate) actual: Option<DocumentVersion>,
171}
172
173pub(crate) fn require_document_version(
174    expected: DocumentVersion,
175    actual: Option<DocumentVersion>,
176) -> Result<(), VersionMismatch> {
177    if actual == Some(expected) {
178        Ok(())
179    } else {
180        Err(VersionMismatch { expected, actual })
181    }
182}
183
184impl std::fmt::Display for MaterializationError {
185    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
186        match self {
187            Self::RankOutOfRange { rank, returned } => {
188                write!(f, "result rank {rank} is outside {returned} returned rows")
189            }
190            Self::MissingIdentity { row_id } => {
191                write!(f, "result row {row_id:?} has no document identity")
192            }
193            Self::MissingText { row_id, document } => write!(
194                f,
195                "result row {row_id:?} for {document:?} has no stored text"
196            ),
197            Self::MissingSource { row_id } => {
198                write!(f, "result row {row_id:?} is outside the admitted snapshot")
199            }
200            Self::InvalidLexicalSource { source } => write!(
201                f,
202                "lexical result source {source} is outside the admitted assembly"
203            ),
204            Self::MissingFusedIdentity { document } => write!(
205                f,
206                "fused document {document:?} has no ranked source identity"
207            ),
208            Self::IdentityMismatch {
209                row_id,
210                expected,
211                actual,
212            } => write!(
213                f,
214                "result row {row_id:?} expected {expected:?}, found {actual:?}"
215            ),
216            Self::AllocationFailed { bytes } => {
217                write!(f, "result text allocation of {bytes} bytes failed")
218            }
219            Self::ArithmeticOverflow => f.write_str("result materialization counter overflow"),
220            Self::Query(error) => error.fmt(f),
221            Self::Storage(error) => error.fmt(f),
222        }
223    }
224}
225impl std::error::Error for MaterializationError {
226    fn source(&self) -> Option<&(dyn std::error::Error + 'static)> {
227        match self {
228            Self::Storage(error) => Some(error),
229            Self::Query(error) => Some(error),
230            _ => None,
231        }
232    }
233}
234
235/// Rank-addressed text access available only inside a Store query callback.
236pub struct QueryMaterializer<'a> {
237    #[cfg(any(test, feature = "test-seams"))]
238    pub(super) store: &'a Store,
239    pub(super) snapshot: &'a PublishedSnapshot,
240    pub(super) active: &'a ActiveSegment,
241    pub(super) cancellation: &'a QueryCancellation<'a>,
242    pub(super) address: &'a dyn Fn(usize) -> Result<ResultAddress, MaterializationError>,
243    pub(super) counts: Cell<crate::diag::MaterializationCounters>,
244}
245
246#[derive(Clone, Copy, Eq, PartialEq)]
247pub(super) struct ResultAddress {
248    pub(super) row_id: GlobalRowId,
249    pub(super) document: Option<DocumentVersion>,
250}
251
252pub(super) type HybridAddresses = super::stats::Accounted<Vec<(DocId, ResultAddress)>>;
253
254impl From<&SearchCandidate> for ResultAddress {
255    fn from(candidate: &SearchCandidate) -> Self {
256        Self {
257            row_id: candidate.row_id(),
258            document: candidate.document(),
259        }
260    }
261}
262
263impl QueryMaterializer<'_> {
264    /// Copies the exact stored text and version associated with one result rank.
265    pub fn text(&self, rank: usize) -> Result<MaterializedRow, MaterializationError> {
266        self.check()?;
267        let address = (self.address)(rank)?;
268        self.materialize_address(address)
269    }
270
271    /// Plants a wrong or missing ranked identity at the actual physical-row validation seam.
272    #[cfg(any(test, feature = "test-seams"))]
273    pub fn text_with_test_document(
274        &self,
275        rank: usize,
276        document: Option<DocumentVersion>,
277    ) -> Result<MaterializedRow, MaterializationError> {
278        self.check()?;
279        let mut address = (self.address)(rank)?;
280        address.document = document;
281        self.materialize_address(address)
282    }
283
284    fn materialize_address(
285        &self,
286        address: ResultAddress,
287    ) -> Result<MaterializedRow, MaterializationError> {
288        let row_id = address.row_id;
289        let document = address
290            .document
291            .ok_or(MaterializationError::MissingIdentity { row_id })?;
292        #[cfg(any(test, feature = "test-seams"))]
293        self.store
294            .text_materialization_work
295            .direct_row_lookups
296            .fetch_add(1, std::sync::atomic::Ordering::Relaxed);
297        let mut counts = self.counts.get();
298        counts.row_lookups = counts
299            .row_lookups
300            .checked_add(1)
301            .ok_or(MaterializationError::ArithmeticOverflow)?;
302        self.counts.set(counts);
303        let row = row_id.local_row() as usize;
304        let (actual, source) = match row_id.source() {
305            RowSource::Active => (
306                self.active.document(row),
307                self.active
308                    .text(row)
309                    .map_err(MaterializationError::Storage)?,
310            ),
311            RowSource::Sealed(id) => {
312                let segment = self
313                    .snapshot
314                    .segments()
315                    .iter()
316                    .find(|segment| segment.meta().id == id)
317                    .ok_or(MaterializationError::MissingSource { row_id })?;
318                let actual = segment
319                    .document_version(row)
320                    .map_err(StoreError::Segment)
321                    .map_err(MaterializationError::Storage)?;
322                let source = segment
323                    .query_stored_text()
324                    .map_err(StoreError::Segment)
325                    .map_err(MaterializationError::Storage)?
326                    .and_then(|text| text.row(row).flatten());
327                (actual, source)
328            }
329        };
330        require_document_version(document, actual).map_err(|mismatch| {
331            MaterializationError::IdentityMismatch {
332                row_id,
333                expected: mismatch.expected,
334                actual: mismatch.actual,
335            }
336        })?;
337        let source = source.ok_or(MaterializationError::MissingText { row_id, document })?;
338        let mut text = String::new();
339        text.try_reserve_exact(source.len()).map_err(|_| {
340            MaterializationError::AllocationFailed {
341                bytes: source.len(),
342            }
343        })?;
344        text.push_str(source);
345        let mut counts = self.counts.get();
346        counts.text_copies = counts
347            .text_copies
348            .checked_add(1)
349            .ok_or(MaterializationError::ArithmeticOverflow)?;
350        counts.text_bytes = counts
351            .text_bytes
352            .checked_add(source.len() as u64)
353            .ok_or(MaterializationError::ArithmeticOverflow)?;
354        self.counts.set(counts);
355        self.check()?;
356        Ok(MaterializedRow {
357            document,
358            row_id,
359            text,
360        })
361    }
362
363    fn check(&self) -> Result<(), MaterializationError> {
364        self.cancellation
365            .check_graph()
366            .map_err(super::map_scan_error)
367            .map_err(MaterializationError::Query)
368    }
369}
370
371impl Store {
372    /// Starts lexical retrieval before obtaining a caller-owned query vector.
373    /// Both legs and returned text use one admission pinned before the producer
374    /// runs. All submitted lexical work is joined on producer failure or panic.
375    ///
376    /// The returned vector borrow must remain valid for the complete query.
377    /// Producer errors retain their original type; core failures are separate.
378    ///
379    /// # Errors
380    /// Returns the producer's error or a typed admission/search failure.
381    pub fn search_hybrid_with_text_deferred<'vector, E, R>(
382        &self,
383        prepare_vector: impl FnOnce() -> Result<crate::ingest::SearchRequest<'vector>, E>,
384        lexical: &crate::fts::search::TermQuery,
385        query: &crate::fusion::HybridQuery,
386        options: impl Into<super::SearchOptions>,
387        control: super::QueryControl,
388        finish: impl FnOnce(&crate::ingest::StoreHybridSearchOutcome, &QueryMaterializer<'_>) -> R,
389    ) -> Result<(crate::ingest::StoreHybridSearchOutcome, R), HybridPreparationError<E>> {
390        self.search_hybrid_pinned_with_text(
391            prepare_vector,
392            None,
393            super::PinnedLexicalQuery::Term(lexical),
394            query,
395            options.into(),
396            control,
397            finish,
398        )
399    }
400
401    #[allow(clippy::too_many_arguments)] // Pinned query inputs plus its scoped continuation.
402    fn search_hybrid_pinned_with_text<'vector, E, R>(
403        &self,
404        prepare_vector: impl FnOnce() -> Result<crate::ingest::SearchRequest<'vector>, E>,
405        filter: Option<&super::QueryFilter>,
406        lexical: super::PinnedLexicalQuery<'_>,
407        query: &crate::fusion::HybridQuery,
408        options: super::SearchOptions,
409        control: super::QueryControl,
410        finish: impl FnOnce(&crate::ingest::StoreHybridSearchOutcome, &QueryMaterializer<'_>) -> R,
411    ) -> Result<(crate::ingest::StoreHybridSearchOutcome, R), HybridPreparationError<E>> {
412        self.search_hybrid_prepared_then(
413            prepare_vector,
414            filter,
415            lexical,
416            query,
417            options,
418            control,
419            true,
420            |mut outcome, snapshot, active, cancellation, addresses| {
421                cancellation.check_graph().map_err(super::map_scan_error)?;
422                let address = |rank| {
423                    let hit =
424                        outcome
425                            .hits
426                            .get(rank)
427                            .ok_or(MaterializationError::RankOutOfRange {
428                                rank,
429                                returned: outcome.hits.len(),
430                            })?;
431                    let addresses = addresses
432                        .ok_or(MaterializationError::MissingFusedIdentity { document: hit.key })?;
433                    let index = addresses
434                        .binary_search_by_key(&hit.key, |(key, _)| *key)
435                        .map_err(|_| MaterializationError::MissingFusedIdentity {
436                            document: hit.key,
437                        })?;
438                    addresses
439                        .get(index)
440                        .map(|(_, address)| *address)
441                        .ok_or(MaterializationError::MissingFusedIdentity { document: hit.key })
442                };
443                let materializer = QueryMaterializer {
444                    #[cfg(any(test, feature = "test-seams"))]
445                    store: self,
446                    snapshot,
447                    active,
448                    cancellation,
449                    address: &address,
450                    counts: Cell::new(Default::default()),
451                };
452                let value = finish(&outcome, &materializer);
453                let counts = materializer.counts.get();
454                outcome.diagnostics.materialization = Some(counts);
455                cancellation.check_graph().map_err(super::map_scan_error)?;
456                Ok((outcome, value))
457            },
458        )
459    }
460
461    /// Reads literal counters without resetting concurrent observations.
462    #[cfg(any(test, feature = "test-seams"))]
463    #[must_use]
464    pub fn query_materialization_test_counters(&self) -> MaterializationTestCounters {
465        use std::sync::atomic::Ordering::Relaxed;
466        MaterializationTestCounters {
467            admissions: self.text_materialization_work.admissions.load(Relaxed),
468            reverse_version_lookups: self
469                .text_materialization_work
470                .reverse_version_lookups
471                .load(Relaxed),
472            direct_row_lookups: self
473                .text_materialization_work
474                .direct_row_lookups
475                .load(Relaxed),
476            #[cfg(feature = "query-timing")]
477            lexical_admission_locks: self
478                .text_materialization_work
479                .lexical_admission_locks
480                .load(Relaxed),
481            #[cfg(feature = "query-timing")]
482            lexical_admission_lock_nanos: self
483                .text_materialization_work
484                .lexical_admission_lock_nanos
485                .load(Relaxed),
486        }
487    }
488    /// Runs hybrid fusion and constructs caller-owned results before releasing its snapshot.
489    pub fn search_hybrid_with_text<R>(
490        &self,
491        vector: crate::ingest::SearchRequest<'_>,
492        lexical: &crate::fts::search::TermQuery,
493        query: &crate::fusion::HybridQuery,
494        options: impl Into<super::SearchOptions>,
495        control: super::QueryControl,
496        finish: impl FnOnce(&crate::ingest::StoreHybridSearchOutcome, &QueryMaterializer<'_>) -> R,
497    ) -> Result<(crate::ingest::StoreHybridSearchOutcome, R), crate::fusion::FusionError> {
498        self.search_hybrid_pinned_with_text(
499            || Ok::<_, std::convert::Infallible>(vector),
500            vector.filter,
501            super::PinnedLexicalQuery::Term(lexical),
502            query,
503            options.into(),
504            control,
505            finish,
506        )
507        .map_err(|error| match error {
508            HybridPreparationError::Preparation(never) => match never {},
509            HybridPreparationError::Search(error) => error,
510        })
511    }
512
513    /// Runs a flat-term lexical query and returns one snippet per candidate,
514    /// in rank order, cut from the text pinned for the query's generation.
515    ///
516    /// Highlights are the tokens the store's own analyzer produces from that
517    /// text whose analyzed term is one of the query's terms, which is exactly
518    /// what the lexical leg scored.
519    ///
520    /// # Errors
521    /// Returns the search failure, or [`crate::ingest::StoreLexicalError::MissingSnippetMatch`]
522    /// when a ranked row's stored text contains none of the query's terms.
523    pub fn search_lexical_with_snippets(
524        &self,
525        query: &crate::fts::search::TermQuery,
526        k: usize,
527        window: NonZeroUsize,
528        control: super::QueryControl,
529    ) -> Result<
530        (
531            crate::ingest::StoreLexicalSearchOutcome,
532            Vec<crate::fts::query::OwnedLexicalSnippet>,
533        ),
534        crate::ingest::StoreLexicalError,
535    > {
536        self.search_lexical_filtered_with_snippets(query, k, window, control, None)
537    }
538
539    /// Searches eligible lexical rows and materializes snippets at the same generation.
540    pub fn search_lexical_filtered_with_snippets(
541        &self,
542        query: &crate::fts::search::TermQuery,
543        k: usize,
544        window: NonZeroUsize,
545        control: super::QueryControl,
546        filter: Option<&super::QueryFilter>,
547    ) -> Result<
548        (
549            crate::ingest::StoreLexicalSearchOutcome,
550            Vec<crate::fts::query::OwnedLexicalSnippet>,
551        ),
552        crate::ingest::StoreLexicalError,
553    > {
554        self.search_lexical_then(
555            query,
556            k,
557            control,
558            filter,
559            |outcome, snapshot, active, assembly, hits, cancellation| {
560                let mut snippets = Vec::with_capacity(hits.len());
561                for hit in hits {
562                    cancellation.check_graph().map_err(super::map_scan_error)?;
563                    let (text, _) = super::structured_lexical_row(
564                        snapshot,
565                        active,
566                        &assembly.sources,
567                        hit.doc,
568                    )?;
569                    snippets.push(
570                        crate::fts::snippet::owned_window(
571                            &self.tokenizer,
572                            text,
573                            &query.terms,
574                            window,
575                        )
576                        .ok_or(
577                            crate::ingest::StoreLexicalError::MissingSnippetMatch {
578                                segment: hit.doc.segment,
579                                row: hit.doc.row,
580                            },
581                        )?,
582                    );
583                }
584                Ok((outcome, snippets))
585            },
586        )
587    }
588
589    /// Runs hybrid fusion over a flat-term lexical leg and returns one
590    /// optional snippet per fused hit, in rank order, cut from the text pinned
591    /// for the query's generation. A hit gets no snippet when its row has no
592    /// stored text or the text contains none of the query's terms.
593    ///
594    /// # Errors
595    /// Returns the search failure; the inner result carries a failure to read
596    /// a ranked row's text.
597    pub fn search_hybrid_with_snippets(
598        &self,
599        vector: crate::ingest::SearchRequest<'_>,
600        lexical: &crate::fts::search::TermQuery,
601        query: &crate::fusion::HybridQuery,
602        options: impl Into<super::SearchOptions>,
603        control: super::QueryControl,
604        window: NonZeroUsize,
605    ) -> Result<(crate::ingest::StoreHybridSearchOutcome, HybridSnippets), crate::fusion::FusionError>
606    {
607        self.hybrid_snippets(
608            vector,
609            super::PinnedLexicalQuery::Term(lexical),
610            query,
611            options.into(),
612            control,
613            window,
614        )
615    }
616
617    /// [`Self::search_hybrid_with_snippets`] for a structured lexical leg.
618    /// Highlights mark the terms of the leg's reported expansions, so a prefix
619    /// marks every indexed term it expanded to.
620    ///
621    /// # Errors
622    /// As [`Self::search_hybrid_with_snippets`].
623    pub fn search_hybrid_structured_with_snippets(
624        &self,
625        vector: crate::ingest::SearchRequest<'_>,
626        lexical: &crate::fts::query::LexicalQuery,
627        query: &crate::fusion::HybridQuery,
628        options: impl Into<super::SearchOptions>,
629        control: super::QueryControl,
630        window: NonZeroUsize,
631    ) -> Result<(crate::ingest::StoreHybridSearchOutcome, HybridSnippets), crate::fusion::FusionError>
632    {
633        self.hybrid_snippets(
634            vector,
635            super::PinnedLexicalQuery::Structured(lexical),
636            query,
637            options.into(),
638            control,
639            window,
640        )
641    }
642
643    fn hybrid_snippets(
644        &self,
645        vector: crate::ingest::SearchRequest<'_>,
646        lexical: super::PinnedLexicalQuery<'_>,
647        query: &crate::fusion::HybridQuery,
648        options: super::SearchOptions,
649        control: super::QueryControl,
650        window: NonZeroUsize,
651    ) -> Result<(crate::ingest::StoreHybridSearchOutcome, HybridSnippets), crate::fusion::FusionError>
652    {
653        let flat_terms = match lexical {
654            super::PinnedLexicalQuery::Term(query) => Some(&query.terms),
655            super::PinnedLexicalQuery::Structured(_) => None,
656        };
657        self.search_hybrid_pinned_with_text(
658            || Ok::<_, std::convert::Infallible>(vector),
659            vector.filter,
660            lexical,
661            query,
662            options,
663            control,
664            |outcome, materializer| {
665                let expanded;
666                let terms: &[Vec<u8>] = match flat_terms {
667                    Some(terms) => terms,
668                    None => {
669                        expanded = outcome
670                            .lexical_expansions
671                            .iter()
672                            .map(|expansion| expansion.term.clone())
673                            .collect::<Vec<_>>();
674                        &expanded
675                    }
676                };
677                let mut snippets = Vec::with_capacity(outcome.hits.len());
678                for rank in 0..outcome.hits.len() {
679                    let row = match materializer.text(rank) {
680                        Ok(row) => row,
681                        Err(MaterializationError::MissingText { .. }) => {
682                            snippets.push(None);
683                            continue;
684                        }
685                        Err(error) => return Err(error),
686                    };
687                    snippets.push(crate::fts::snippet::owned_window(
688                        &self.tokenizer,
689                        &row.text,
690                        terms,
691                        window,
692                    ));
693                }
694                Ok(snippets)
695            },
696        )
697        .map_err(|error| match error {
698            HybridPreparationError::Preparation(never) => match never {},
699            HybridPreparationError::Search(error) => error,
700        })
701    }
702
703    /// Runs a lexical query and constructs caller-owned results before releasing its snapshot.
704    pub fn search_lexical_with_text<R>(
705        &self,
706        query: &crate::fts::search::TermQuery,
707        k: usize,
708        control: super::QueryControl,
709        finish: impl FnOnce(&crate::ingest::StoreLexicalSearchOutcome, &QueryMaterializer<'_>) -> R,
710    ) -> Result<(crate::ingest::StoreLexicalSearchOutcome, R), crate::ingest::StoreLexicalError>
711    {
712        self.search_lexical_then(
713            query,
714            k,
715            control,
716            None,
717            |mut outcome, snapshot, active, assembly, hits, cancellation| {
718                cancellation.check_graph().map_err(super::map_scan_error)?;
719                let address = |rank| {
720                    let candidate = outcome.candidates.get(rank).ok_or(
721                        MaterializationError::RankOutOfRange {
722                            rank,
723                            returned: outcome.candidates.len(),
724                        },
725                    )?;
726                    let hit = hits.get(rank).ok_or(MaterializationError::RankOutOfRange {
727                        rank,
728                        returned: hits.len(),
729                    })?;
730                    Ok(ResultAddress {
731                        row_id: lexical_row_id(snapshot, &assembly.sources, hit.doc)?,
732                        document: Some(candidate.document),
733                    })
734                };
735                let materializer = QueryMaterializer {
736                    #[cfg(any(test, feature = "test-seams"))]
737                    store: self,
738                    snapshot,
739                    active,
740                    cancellation,
741                    address: &address,
742                    counts: Cell::new(Default::default()),
743                };
744                let value = finish(&outcome, &materializer);
745                let counts = materializer.counts.get();
746                outcome.diagnostics.materialization = Some(counts);
747                cancellation.check_graph().map_err(super::map_scan_error)?;
748                Ok((outcome, value))
749            },
750        )
751    }
752
753    /// Runs a dense query and constructs owned caller results through rank-addressed text access.
754    pub fn search_with_text<R>(
755        &self,
756        request: crate::ingest::SearchRequest<'_>,
757        k: usize,
758        options: impl Into<super::SearchOptions>,
759        control: super::QueryControl,
760        finish: impl FnOnce(&crate::ingest::SearchOutcome, &QueryMaterializer<'_>) -> R,
761    ) -> Result<(crate::ingest::SearchOutcome, R), super::QueryError> {
762        self.search_with_graph_bound_mode_then(
763            request,
764            k,
765            options.into(),
766            control,
767            super::GraphBoundMode::Shared,
768            |mut outcome, admitted, control| {
769                let lease =
770                    SnapshotLease::new_at(Arc::clone(&admitted.snapshot), admitted.generation);
771                let cancellation = QueryCancellation::new(control, &lease);
772                cancellation.check_graph().map_err(super::map_scan_error)?;
773                let address = |rank| {
774                    outcome.candidates.get(rank).map(ResultAddress::from).ok_or(
775                        MaterializationError::RankOutOfRange {
776                            rank,
777                            returned: outcome.candidates.len(),
778                        },
779                    )
780                };
781                let materializer = QueryMaterializer {
782                    #[cfg(any(test, feature = "test-seams"))]
783                    store: self,
784                    counts: Cell::new(Default::default()),
785                    snapshot: &admitted.snapshot,
786                    active: &admitted.active_segment,
787                    cancellation: &cancellation,
788                    address: &address,
789                };
790                let value = finish(&outcome, &materializer);
791                let counts = materializer.counts.get();
792                outcome.diagnostics.materialization = Some(counts);
793                cancellation.check_graph().map_err(super::map_scan_error)?;
794                Ok((outcome, value))
795            },
796        )
797    }
798}
799
800fn lexical_row_id(
801    snapshot: &PublishedSnapshot,
802    sources: &[super::StructuredLexicalSource],
803    doc: crate::fts::search::GlobalDocId,
804) -> Result<GlobalRowId, MaterializationError> {
805    let source =
806        sources
807            .get(doc.segment as usize)
808            .ok_or(MaterializationError::InvalidLexicalSource {
809                source: doc.segment,
810            })?;
811    let source = match source {
812        super::StructuredLexicalSource::Active => RowSource::Active,
813        super::StructuredLexicalSource::Sealed(ordinal) => RowSource::Sealed(
814            snapshot
815                .segments()
816                .get(*ordinal)
817                .ok_or(MaterializationError::InvalidLexicalSource {
818                    source: doc.segment,
819                })?
820                .meta()
821                .id,
822        ),
823    };
824    Ok(GlobalRowId::new(source, doc.row))
825}
zeppelin_embed/lifecycle/materialize.rs

zeppelin_embed/lifecycle/
materialize.rs
