nodedb_fts/index/writer/
document.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Document writes reuse ordered analyzer output and preserve surrogate bounds.
4
5use super::FtsIndex;
6use crate::{
7    backend::FtsBackend,
8    block::CompactPosting,
9    codec::smallfloat,
10    index::error::{FtsIndexError, MAX_INDEXABLE_SURROGATE},
11    lsm::segment_deletes::SegmentDeletes,
12    scope::IndexScope,
13};
14use nodedb_types::Surrogate;
15use std::collections::HashMap;
16use tracing::debug;
17
18impl<B: FtsBackend> FtsIndex<B> {
19    /// Index a document's text content into one index. The analyzer is the
20    /// one bound to the index's collection.
21    ///
22    /// Returns `Err(FtsIndexError::SurrogateOutOfRange)` for `Surrogate::ZERO`
23    /// or values exceeding `MAX_INDEXABLE_SURROGATE`. `Surrogate::ZERO` is the
24    /// unassigned sentinel. Fieldnorm arrays use raw `u32` surrogates as
25    /// indexes. Values near `u32::MAX` cause multi-GiB allocations. Runtime
26    /// bounds checks run before analysis and remain active in release builds.
27    pub fn index_document<'a>(
28        &self,
29        database_id: u64,
30        tid: u64,
31        index: impl Into<IndexScope<'a>>,
32        doc_id: Surrogate,
33        text: &str,
34    ) -> Result<(), FtsIndexError<B::Error>> {
35        let index = index.into();
36        Self::check_surrogate(doc_id)?;
37
38        let tokens = self
39            .analyze_for_collection(database_id, tid, index.collection(), text)
40            .map_err(FtsIndexError::backend)?;
41        self.index_analyzed_document(database_id, tid, index, doc_id, &tokens)
42    }
43
44    /// Index ordered tokens from this collection's current analyzer. Callers
45    /// preserve token order and hold analyzer config stable through indexing.
46    ///
47    /// Returns `Err(FtsIndexError::SurrogateOutOfRange)` for `Surrogate::ZERO`
48    /// or values exceeding `MAX_INDEXABLE_SURROGATE`.
49    pub fn index_analyzed_document<'a>(
50        &self,
51        database_id: u64,
52        tid: u64,
53        index: impl Into<IndexScope<'a>>,
54        doc_id: Surrogate,
55        tokens: &[String],
56    ) -> Result<(), FtsIndexError<B::Error>> {
57        let index = index.into();
58        Self::check_surrogate(doc_id)?;
59        if tokens.is_empty() {
60            return Ok(());
61        }
62
63        let mut term_data: HashMap<&str, (u32, Vec<u32>)> = HashMap::new();
64        for (pos, token) in tokens.iter().enumerate() {
65            let entry = term_data.entry(token.as_str()).or_insert((0, Vec::new()));
66            entry.0 += 1;
67            entry.1.push(pos as u32);
68        }
69
70        let doc_len = tokens.len() as u32;
71        let fieldnorm = smallfloat::encode(doc_len);
72
73        let term_count = term_data.len();
74        self.memtable.insert_doc(
75            database_id,
76            tid,
77            index,
78            doc_id,
79            doc_len,
80            term_data.into_iter().map(|(term, (freq, positions))| {
81                (
82                    term,
83                    CompactPosting {
84                        doc_id,
85                        term_freq: freq,
86                        fieldnorm,
87                        positions,
88                    },
89                )
90            }),
91        );
92
93        // Write document length, fieldnorm, and update incremental stats.
94        self.backend
95            .write_doc_length(database_id, tid, index, doc_id, doc_len)
96            .map_err(FtsIndexError::backend)?;
97        self.write_fieldnorm(database_id, tid, index, doc_id, doc_len)
98            .map_err(FtsIndexError::backend)?;
99        self.backend
100            .increment_stats(database_id, tid, index, doc_len)
101            .map_err(FtsIndexError::backend)?;
102
103        if self.memtable.should_flush() {
104            self.flush_all_memtables()?;
105        }
106
107        debug!(
108            database_id,
109            tid,
110            collection = index.collection(),
111            field = index.field_key(),
112            doc_id = doc_id.0,
113            tokens = tokens.len(),
114            terms = term_count,
115            "indexed document"
116        );
117        Ok(())
118    }
119
120    /// Remove a document from one index.
121    ///
122    /// The memtable drops the document's postings. Postings already flushed
123    /// to a segment stay in that immutable segment, so the document enters
124    /// the delete set of every segment the index holds: reads skip those
125    /// postings and a merge drops them. The corpus stats lose the document's
126    /// length. A document the index does not hold changes nothing.
127    pub fn remove_document<'a>(
128        &self,
129        database_id: u64,
130        tid: u64,
131        index: impl Into<IndexScope<'a>>,
132        doc_id: Surrogate,
133    ) -> Result<(), FtsIndexError<B::Error>> {
134        let index = index.into();
135        let doc_len = self
136            .backend
137            .read_doc_length(database_id, tid, index, doc_id)
138            .map_err(FtsIndexError::backend)?;
139
140        self.memtable.remove_doc(database_id, tid, index, doc_id);
141
142        let Some(len) = doc_len else {
143            return Ok(());
144        };
145        let segments = self
146            .backend
147            .list_segments(database_id, tid, index)
148            .map_err(FtsIndexError::backend)?;
149        if !segments.is_empty() {
150            let mut deletes =
151                SegmentDeletes::load(&self.backend, database_id, tid, index, &segments)?;
152            deletes.mark_removed(&segments, doc_id);
153            deletes.store(&self.backend, database_id, tid, index)?;
154        }
155        self.backend
156            .remove_doc_length(database_id, tid, index, doc_id)
157            .map_err(FtsIndexError::backend)?;
158        self.backend
159            .decrement_stats(database_id, tid, index, len)
160            .map_err(FtsIndexError::backend)?;
161
162        Ok(())
163    }
164
165    fn check_surrogate(doc_id: Surrogate) -> Result<(), FtsIndexError<B::Error>> {
166        let raw = doc_id.as_u32();
167        if raw == 0 || raw > MAX_INDEXABLE_SURROGATE {
168            return Err(FtsIndexError::SurrogateOutOfRange { surrogate: doc_id });
169        }
170        Ok(())
171    }
172}
173
174#[cfg(test)]
175mod tests {
176    use nodedb_types::Surrogate;
177
178    use crate::backend::memory::MemoryBackend;
179    use crate::test_support::test_governor;
180
181    use super::*;
182
183    const DB: u64 = 0;
184    const T: u64 = 1;
185    const DOCS: IndexScope<'static> = IndexScope::document("docs");
186
187    fn make_index() -> FtsIndex<MemoryBackend> {
188        FtsIndex::new(MemoryBackend::new(), test_governor())
189    }
190
191    #[test]
192    fn index_writes_to_memtable() {
193        let idx = make_index();
194        idx.index_document(DB, T, "docs", Surrogate(1), "hello world greeting")
195            .unwrap();
196
197        assert!(!idx.memtable.is_empty());
198        assert!(idx.memtable.posting_count() > 0);
199    }
200
201    #[test]
202    fn index_surrogate_stored() {
203        let idx = make_index();
204        // Surrogates must be in 1..=MAX_INDEXABLE_SURROGATE. Surrogate::ZERO is the unassigned sentinel and is rejected at index time.
205        idx.index_document(DB, T, "docs", Surrogate(10), "hello world greeting")
206            .unwrap();
207        idx.index_document(DB, T, "docs", Surrogate(11), "hello rust language")
208            .unwrap();
209
210        let (count, _) = idx.backend.collection_stats(DB, T, DOCS).unwrap();
211        assert_eq!(count, 2);
212    }
213
214    #[test]
215    fn remove_decrements_stats() {
216        let idx = make_index();
217        idx.index_document(DB, T, "docs", Surrogate(10), "hello world")
218            .unwrap();
219        idx.index_document(DB, T, "docs", Surrogate(11), "hello rust")
220            .unwrap();
221
222        idx.remove_document(DB, T, "docs", Surrogate(10)).unwrap();
223
224        let (count, _) = idx.backend.collection_stats(DB, T, DOCS).unwrap();
225        assert_eq!(count, 1);
226    }
227
228    fn hits(idx: &FtsIndex<MemoryBackend>, query: &str) -> Vec<u32> {
229        let mut ids: Vec<u32> = idx
230            .search(
231                DB,
232                T,
233                DOCS,
234                crate::FtsSearchParams {
235                    query,
236                    top_k: usize::MAX,
237                    fuzzy_enabled: false,
238                    mode: crate::posting::QueryMode::Or,
239                    prefilter: None,
240                },
241            )
242            .unwrap()
243            .into_iter()
244            .map(|r| r.doc_id.0)
245            .collect();
246        ids.sort_unstable();
247        ids
248    }
249
250    /// A document removed after its postings reached a segment stops
251    /// matching, leaves the document frequency and the corpus stats, and
252    /// leaves the segment at the next merge. Indexed again, it matches its
253    /// new text only.
254    #[test]
255    fn remove_after_flush_hides_segment_postings_and_merge_drops_them() {
256        use crate::lsm::compaction::{CompactLevelParams, SegmentMeta, compact_level, parse_level};
257        use crate::lsm::segment::reader::SegmentReader;
258        use crate::{DocScore, TextQuery};
259
260        let idx = make_index();
261        idx.index_document(DB, T, "docs", Surrogate(1), "rust tokio")
262            .unwrap();
263        idx.index_document(DB, T, "docs", Surrogate(2), "rust axum")
264            .unwrap();
265        idx.flush_memtable(DB, T, "docs").unwrap();
266        idx.index_document(DB, T, "docs", Surrogate(3), "rust serde")
267            .unwrap();
268        idx.flush_memtable(DB, T, "docs").unwrap();
269        assert_eq!(hits(&idx, "tokio"), vec![1]);
270
271        idx.remove_document(DB, T, "docs", Surrogate(1)).unwrap();
272        assert!(hits(&idx, "tokio").is_empty(), "a removed document matches");
273        assert_eq!(hits(&idx, "rust"), vec![2, 3]);
274        let blocks = idx.term_blocks(DB, T, DOCS, &["rust".into()]).unwrap();
275        assert_eq!(blocks[0].df, 2, "document frequency counts live documents");
276        assert_eq!(idx.backend.collection_stats(DB, T, DOCS).unwrap(), (2, 4));
277        let scorer = idx
278            .doc_scorer(
279                DB,
280                T,
281                DOCS,
282                TextQuery {
283                    query: "tokio",
284                    fuzzy_enabled: false,
285                    mode: crate::posting::QueryMode::Or,
286                },
287                None,
288                None,
289            )
290            .unwrap();
291        assert_eq!(
292            scorer.score(&[Surrogate(1)]).unwrap(),
293            vec![DocScore::Absent]
294        );
295
296        idx.index_document(DB, T, "docs", Surrogate(1), "rust hyper")
297            .unwrap();
298        idx.flush_memtable(DB, T, "docs").unwrap();
299        assert_eq!(hits(&idx, "hyper"), vec![1]);
300        assert!(hits(&idx, "tokio").is_empty());
301        assert_eq!(hits(&idx, "rust"), vec![1, 2, 3]);
302
303        let segments: Vec<SegmentMeta> = idx
304            .backend
305            .list_segments(DB, T, DOCS)
306            .unwrap()
307            .into_iter()
308            .map(|segment_id| SegmentMeta {
309                level: parse_level(&segment_id),
310                segment_id,
311                size: 0,
312            })
313            .collect();
314        assert_eq!(segments.len(), 3);
315        let governor = test_governor();
316        let (merged, merged_ids) = compact_level(CompactLevelParams {
317            backend: &idx.backend,
318            database_id: DB,
319            tid: T,
320            index: DOCS,
321            segments: &segments,
322            level: 0,
323            governor: &governor,
324        })
325        .unwrap()
326        .expect("three level-0 segments merge");
327        let reader = SegmentReader::open(merged.clone()).unwrap();
328        assert!(
329            reader.find_term("tokio").is_none(),
330            "the merge drops dead postings"
331        );
332        assert_eq!(reader.df("rust"), 3);
333
334        idx.backend
335            .write_segment(DB, T, DOCS, "L1:00000000000000ff", &merged)
336            .unwrap();
337        for id in &merged_ids {
338            idx.backend.remove_segment(DB, T, DOCS, id).unwrap();
339        }
340        assert!(hits(&idx, "tokio").is_empty());
341        assert_eq!(hits(&idx, "rust"), vec![1, 2, 3]);
342
343        // The merged segment carries no delete set: a later removal hides it.
344        idx.remove_document(DB, T, "docs", Surrogate(2)).unwrap();
345        assert_eq!(hits(&idx, "rust"), vec![1, 3]);
346    }
347
348    /// A flush after a restart does not reuse a stored segment id.
349    #[test]
350    fn a_flush_never_reuses_a_stored_segment_id() {
351        let idx = make_index();
352        idx.index_document(DB, T, "docs", Surrogate(1), "alpha")
353            .unwrap();
354        idx.flush_memtable(DB, T, "docs").unwrap();
355
356        let restarted = FtsIndex::new(idx.backend, test_governor());
357        restarted
358            .index_document(DB, T, "docs", Surrogate(2), "bravo")
359            .unwrap();
360        restarted.flush_memtable(DB, T, "docs").unwrap();
361        assert_eq!(
362            restarted.backend.list_segments(DB, T, DOCS).unwrap().len(),
363            2
364        );
365        assert_eq!(hits(&restarted, "alpha"), vec![1]);
366        assert_eq!(hits(&restarted, "bravo"), vec![2]);
367    }
368
369    #[test]
370    fn index_updates_stats() {
371        let idx = make_index();
372        idx.index_document(DB, T, "docs", Surrogate(10), "hello world greeting")
373            .unwrap();
374        idx.index_document(DB, T, "docs", Surrogate(11), "hello rust language")
375            .unwrap();
376
377        let (count, total) = idx.backend.collection_stats(DB, T, DOCS).unwrap();
378        assert_eq!(count, 2);
379        assert!(total > 0);
380    }
381
382    #[test]
383    fn field_index_keeps_its_own_postings_and_stats() {
384        let idx = make_index();
385        let title = IndexScope::field("docs", "title").unwrap();
386        idx.index_document(DB, T, title, Surrogate(1), "rust")
387            .unwrap();
388        idx.index_document(DB, T, "docs", Surrogate(1), "rust handbook")
389            .unwrap();
390        idx.index_document(DB, T, "docs", Surrogate(2), "rust nomicon")
391            .unwrap();
392
393        assert_eq!(idx.backend.collection_stats(DB, T, title).unwrap(), (1, 1));
394        assert_eq!(idx.backend.collection_stats(DB, T, DOCS).unwrap(), (2, 4));
395        assert_eq!(idx.memtable.get_postings(DB, T, title, "rust").len(), 1);
396        assert_eq!(idx.memtable.get_postings(DB, T, DOCS, "rust").len(), 2);
397
398        idx.remove_document(DB, T, title, Surrogate(1)).unwrap();
399        assert_eq!(idx.backend.collection_stats(DB, T, title).unwrap(), (0, 0));
400        assert!(idx.memtable.get_postings(DB, T, title, "rust").is_empty());
401        assert_eq!(idx.memtable.get_postings(DB, T, DOCS, "rust").len(), 2);
402    }
403
404    #[test]
405    fn empty_text_is_noop() {
406        let idx = make_index();
407        idx.index_document(DB, T, "docs", Surrogate(1), "the a is")
408            .unwrap();
409        assert_eq!(idx.backend.collection_stats(DB, T, DOCS).unwrap(), (0, 0));
410        assert!(idx.memtable.is_empty());
411    }
412
413    // ── Surrogate boundary tests ──────────────────────────────────────────────
414
415    /// Spec: Surrogate::ZERO (the unassigned sentinel) must be rejected at index
416    /// time with FtsIndexError::SurrogateOutOfRange, not written into the index.
417    #[test]
418    fn index_document_rejects_zero_surrogate() {
419        let idx = make_index();
420        let err = idx
421            .index_document(DB, T, "docs", Surrogate(0), "hello world")
422            .unwrap_err();
423        assert!(
424            matches!(err, FtsIndexError::SurrogateOutOfRange { surrogate } if surrogate == Surrogate(0)),
425            "expected SurrogateOutOfRange(sur:0), got {err}"
426        );
427    }
428
429    /// Spec: Surrogate(u32::MAX) must be rejected — it is reserved as a sentinel
430    /// and would also cause a 4 GiB fieldnorm array resize.
431    #[test]
432    fn index_document_rejects_u32_max_surrogate() {
433        let idx = make_index();
434        let err = idx
435            .index_document(DB, T, "docs", Surrogate(u32::MAX), "hello world")
436            .unwrap_err();
437        assert!(
438            matches!(err, FtsIndexError::SurrogateOutOfRange { .. }),
439            "expected SurrogateOutOfRange, got {err}"
440        );
441    }
442
443    /// Check the last valid surrogate constant and a representative valid input.
444    #[test]
445    fn index_document_accepts_max_indexable_surrogate() {
446        // Indexing MAX_INDEXABLE_SURROGATE requires multi-GiB fieldnorm arrays.
447        // This fixture uses Surrogate(1) and checks the constant and sentinel boundary separately.
448        let idx = make_index();
449        // Check a valid surrogate without allocating the largest fieldnorm array.
450        idx.index_document(DB, T, "docs", Surrogate(1), "boundary check")
451            .unwrap();
452        // Confirm the constant is correct.
453        assert_eq!(
454            crate::index::error::MAX_INDEXABLE_SURROGATE,
455            u32::MAX - 1,
456            "MAX_INDEXABLE_SURROGATE must be u32::MAX - 1"
457        );
458    }
459
460    /// Spec: the SurrogateOutOfRange error message must be informative.
461    #[test]
462    fn surrogate_out_of_range_error_is_informative() {
463        let err: FtsIndexError<crate::backend::memory::MemoryError> =
464            FtsIndexError::SurrogateOutOfRange {
465                surrogate: Surrogate(0),
466            };
467        let msg = err.to_string();
468        assert!(
469            msg.contains("out of the indexable range"),
470            "error message must mention range: {msg}"
471        );
472        assert!(
473            msg.contains("unassigned sentinel"),
474            "error message must explain zero sentinel: {msg}"
475        );
476    }
477    #[test]
478    fn analyzed_tokens_match_text_indexing_and_surrogate_validation() {
479        let text_index = make_index();
480        let token_index = make_index();
481        let text = "Alpha alpha beta gamma";
482        let tokens = token_index
483            .analyze_for_collection(DB, T, "docs", text)
484            .unwrap();
485        text_index
486            .index_document(DB, T, "docs", Surrogate(1), text)
487            .unwrap();
488        token_index
489            .index_analyzed_document(DB, T, "docs", Surrogate(1), &tokens)
490            .unwrap();
491        assert_eq!(
492            text_index.memtable.stats(DB, T, DOCS),
493            token_index.memtable.stats(DB, T, DOCS)
494        );
495        let mut terms = text_index.memtable.terms(DB, T, DOCS);
496        terms.sort();
497        for term in terms {
498            let expected = text_index.memtable.get_postings(DB, T, DOCS, &term);
499            let actual = token_index.memtable.get_postings(DB, T, DOCS, &term);
500            assert_eq!(actual.len(), expected.len());
501            for (actual, expected) in actual.iter().zip(expected) {
502                assert_eq!(actual.doc_id, expected.doc_id);
503                assert_eq!(actual.term_freq, expected.term_freq);
504                assert_eq!(actual.fieldnorm, expected.fieldnorm);
505                assert_eq!(actual.positions, expected.positions);
506            }
507        }
508        assert_eq!(
509            text_index.backend.collection_stats(DB, T, DOCS).unwrap(),
510            token_index.backend.collection_stats(DB, T, DOCS).unwrap()
511        );
512        for id in [Surrogate::ZERO, Surrogate(u32::MAX)] {
513            assert!(matches!(
514                token_index.index_analyzed_document(DB, T, "docs", id, &tokens),
515                Err(FtsIndexError::SurrogateOutOfRange { .. })
516            ));
517            assert!(matches!(
518                token_index.index_analyzed_document(DB, T, "docs", id, &[]),
519                Err(FtsIndexError::SurrogateOutOfRange { .. })
520            ));
521        }
522        let empty = make_index();
523        empty
524            .index_analyzed_document(DB, T, "docs", Surrogate(1), &[])
525            .unwrap();
526        assert!(empty.memtable.is_empty());
527    }
528}
nodedb_fts/index/writer/document.rs

nodedb_fts/index/writer/
document.rs
