atagen

macltok

public
Star0
Watch0
cff1b0a4eb30290a6fe32fea825c816b696fe4f4
3
ssh://brewery@pls.lobotomise.me:33/atagen/macltok.git
src/search/index.rsRust826 lines27.2 KB
1
use std::{collections::HashMap, path::PathBuf, sync::Arc};
2
3
use sha2::{Digest, Sha256};
4
use tantivy::{
5
    DocAddress, DocId, DocSet, Index, IndexSettings, ReloadPolicy, Score, Searcher, SegmentReader,
6
    TantivyDocument, Term,
7
    collector::{Collector, SegmentCollector},
8
    directory::{Directory, OwnedBytes, RamDirectory},
9
    fastfield::{AliveBitSet, write_alive_bitset},
10
    indexer::{SingleSegmentIndexWriter, merge_filtered_segments},
11
    query::{BooleanQuery, Occur, Query, QueryParser, TermSetQuery},
12
    schema::{
13
        FAST, Field, IndexRecordOption, STORED, STRING, Schema, TextFieldIndexing, TextOptions,
14
        Value,
15
    },
16
    tokenizer::{
17
        AsciiFoldingFilter, LowerCaser, RemoveLongFilter, SimpleTokenizer, TextAnalyzer,
18
        TokenStream,
19
    },
20
};
21
use tantivy_common::BitSet;
22
23
use super::MessageIndex;
24
use crate::{Error, ErrorKind};
25
26
const TOKENIZER: &str = "macltok";
27
const WRITER_BUDGET: usize = 4_000_000;
28
/// segments of a size class merge once there are this many of them
29
const MERGE_FACTOR: usize = 6;
30
31
#[derive(Clone, Debug, Eq, PartialEq)]
32
pub(crate) struct Doc {
33
    pub(crate) key: String,
34
    pub(crate) room: String,
35
    pub(crate) sender: String,
36
    /// the original's, or the edit's while the original hasn't turned up
37
    pub(crate) ts: u64,
38
    /// the event whose text is indexed, the original or its latest edit
39
    pub(crate) content: String,
40
    pub(crate) content_ts: u64,
41
    /// only kept for provisional docs, read back empty otherwise
42
    pub(crate) text: String,
43
    /// an edit seen before its original
44
    pub(crate) provisional: bool,
45
}
46
47
#[derive(Clone, Debug, Eq, PartialEq)]
48
pub(crate) struct Hit {
49
    pub(crate) key: String,
50
    pub(crate) room: String,
51
    pub(crate) sender: String,
52
    pub(crate) content: String,
53
    pub(crate) ts: u64,
54
    pub(crate) tie: u64,
55
}
56
57
struct SegmentRecord {
58
    id: u64,
59
    del_gen: u64,
60
}
61
62
#[derive(Default)]
63
pub(crate) struct Changes {
64
    pub(crate) puts: Vec<(String, Vec<u8>)>,
65
    pub(crate) deletes: Vec<String>,
66
    files: Vec<u64>,
67
    dels: Vec<(u64, u64)>,
68
    removed: Vec<u64>,
69
}
70
71
#[derive(Clone, Copy)]
72
struct Fields {
73
    key: Field,
74
    room: Field,
75
    sender: Field,
76
    content: Field,
77
    text: Field,
78
    ts: Field,
79
    tie: Field,
80
    content_ts: Field,
81
    pending: Field,
82
}
83
84
struct Segment {
85
    id: u64,
86
    raw: RamDirectory,
87
    index: Index,
88
    searcher: Searcher,
89
    alive: Arc<BitSet>,
90
    max_doc: u32,
91
    del_gen: u64,
92
    persisted: bool,
93
    deletes_dirty: bool,
94
}
95
96
impl Segment {
97
    fn live(&self) -> usize {
98
        self.alive.len()
99
    }
100
}
101
102
pub(crate) struct KeywordIndex {
103
    schema: Schema,
104
    fields: Fields,
105
    settings: IndexSettings,
106
    parser: QueryParser,
107
    segments: Vec<Segment>,
108
    buffer: HashMap<String, Doc>,
109
    next_id: u64,
110
    persistent: bool,
111
    /// merged away since the last commit, with the delete generation they were persisted at
112
    removed: Vec<(u64, bool, u64)>,
113
}
114
115
pub(crate) fn analyzer() -> TextAnalyzer {
116
    TextAnalyzer::builder(SimpleTokenizer::default())
117
        .filter(RemoveLongFilter::limit(40))
118
        .filter(LowerCaser)
119
        .filter(AsciiFoldingFilter)
120
        .build()
121
}
122
123
/// the byte ranges in `text` whose analysed token is one of `terms`
124
pub(crate) fn match_ranges(text: &str, terms: &[String]) -> Vec<(usize, usize)> {
125
    if terms.is_empty() {
126
        return Vec::new();
127
    }
128
    let mut analyzer = analyzer();
129
    let mut stream = analyzer.token_stream(text);
130
    let mut ranges = Vec::new();
131
    while stream.advance() {
132
        let token = stream.token();
133
        if terms.iter().any(|term| *term == token.text) {
134
            ranges.push((token.offset_from, token.offset_to));
135
        }
136
    }
137
    ranges
138
}
139
140
fn schema() -> (Schema, Fields) {
141
    let mut builder = Schema::builder();
142
    let text_options = TextOptions::default().set_indexing_options(
143
        TextFieldIndexing::default()
144
            .set_tokenizer(TOKENIZER)
145
            .set_index_option(IndexRecordOption::WithFreqsAndPositions),
146
    );
147
    let fields = Fields {
148
        key: builder.add_text_field("key", STRING | STORED),
149
        room: builder.add_text_field("room", STRING | STORED),
150
        sender: builder.add_text_field("sender", STRING | STORED),
151
        content: builder.add_text_field("content", STRING | STORED),
152
        text: builder.add_text_field("text", text_options),
153
        ts: builder.add_u64_field("ts", FAST | STORED),
154
        tie: builder.add_u64_field("tie", FAST),
155
        content_ts: builder.add_u64_field("content_ts", STORED),
156
        pending: builder.add_text_field("pending", STORED),
157
    };
158
    (builder.build(), fields)
159
}
160
161
fn tie(key: &str) -> u64 {
162
    let digest = Sha256::digest(key.as_bytes());
163
    u64::from_le_bytes(digest[..8].try_into().expect("a sha256 digest is 32 bytes"))
164
}
165
166
fn settings() -> IndexSettings {
167
    // the dedicated compressor is a thread, and wasm32-unknown-unknown has none
168
    IndexSettings {
169
        docstore_compress_dedicated_thread: false,
170
        ..IndexSettings::default()
171
    }
172
}
173
174
fn tantivy_error(
175
    context: &'static str,
176
    error: impl std::error::Error + Send + Sync + 'static,
177
) -> Error {
178
    Error::with_source(ErrorKind::Store, context, error)
179
}
180
181
fn open_segment(id: u64, raw: RamDirectory, max_doc_hint: Option<u32>) -> Result<Segment, Error> {
182
    let index = Index::open(raw.clone())
183
        .map_err(|error| tantivy_error("opening a search segment", error))?;
184
    index.tokenizers().register(TOKENIZER, analyzer());
185
    let reader = index
186
        .reader_builder()
187
        .reload_policy(ReloadPolicy::Manual)
188
        .try_into()
189
        .map_err(|error| tantivy_error("reading a search segment", error))?;
190
    let searcher = reader.searcher();
191
    let max_doc = searcher
192
        .segment_readers()
193
        .iter()
194
        .map(SegmentReader::max_doc)
195
        .sum();
196
    if searcher.segment_readers().len() > 1 || max_doc_hint.is_some_and(|hint| hint != max_doc) {
197
        return Err(Error::new(
198
            ErrorKind::Store,
199
            "a search segment has an unexpected shape",
200
        ));
201
    }
202
    Ok(Segment {
203
        id,
204
        raw,
205
        index,
206
        searcher,
207
        alive: Arc::new(BitSet::with_max_value_and_full(max_doc)),
208
        max_doc,
209
        del_gen: 0,
210
        persisted: false,
211
        deletes_dirty: false,
212
    })
213
}
214
215
fn pack(files: &[(String, Vec<u8>)]) -> Vec<u8> {
216
    let mut packed = Vec::with_capacity(
217
        files
218
            .iter()
219
            .map(|(name, bytes)| name.len() + bytes.len() + 12)
220
            .sum(),
221
    );
222
    packed.extend_from_slice(&(files.len() as u32).to_le_bytes());
223
    for (name, bytes) in files {
224
        packed.extend_from_slice(&(name.len() as u32).to_le_bytes());
225
        packed.extend_from_slice(name.as_bytes());
226
        packed.extend_from_slice(&(bytes.len() as u64).to_le_bytes());
227
        packed.extend_from_slice(bytes);
228
    }
229
    packed
230
}
231
232
fn unpack(mut bytes: &[u8]) -> Option<Vec<(String, Vec<u8>)>> {
233
    fn take<'a>(bytes: &mut &'a [u8], len: usize) -> Option<&'a [u8]> {
234
        if bytes.len() < len {
235
            return None;
236
        }
237
        let (head, tail) = bytes.split_at(len);
238
        *bytes = tail;
239
        Some(head)
240
    }
241
    let count = u32::from_le_bytes(take(&mut bytes, 4)?.try_into().ok()?);
242
    let mut files = Vec::new();
243
    for _ in 0..count {
244
        let name_len = u32::from_le_bytes(take(&mut bytes, 4)?.try_into().ok()?) as usize;
245
        let name = std::str::from_utf8(take(&mut bytes, name_len)?)
246
            .ok()?
247
            .to_owned();
248
        let len =
249
            usize::try_from(u64::from_le_bytes(take(&mut bytes, 8)?.try_into().ok()?)).ok()?;
250
        files.push((name, take(&mut bytes, len)?.to_vec()));
251
    }
252
    bytes.is_empty().then_some(files)
253
}
254
255
pub(crate) fn segment_key(id: u64) -> String {
256
    format!("seg/{id}")
257
}
258
259
pub(crate) fn deletes_key(id: u64, generation: u64) -> String {
260
    format!("del/{id}/{generation}")
261
}
262
263
impl KeywordIndex {
264
    pub(crate) fn new(persistent: bool) -> Self {
265
        let (schema, fields) = schema();
266
        let parser_index = Index::create_in_ram(schema.clone());
267
        parser_index.tokenizers().register(TOKENIZER, analyzer());
268
        let mut parser = QueryParser::for_index(&parser_index, vec![fields.text]);
269
        parser.set_conjunction_by_default();
270
        Self {
271
            schema,
272
            fields,
273
            settings: settings(),
274
            parser,
275
            segments: Vec::new(),
276
            buffer: HashMap::new(),
277
            next_id: 1,
278
            persistent,
279
            removed: Vec::new(),
280
        }
281
    }
282
283
    /// `blob` hands back what an earlier commit put, any of it missing or unreadable fails the load
284
    pub(crate) fn load(
285
        manifest: &serde_json::Value,
286
        mut blob: impl FnMut(&str) -> Option<Vec<u8>>,
287
    ) -> Result<Self, Error> {
288
        let invalid = || Error::new(ErrorKind::Store, "the search manifest is unreadable");
289
        let records = manifest["segments"]
290
            .as_array()
291
            .ok_or_else(invalid)?
292
            .iter()
293
            .map(|record| {
294
                Some(SegmentRecord {
295
                    id: record.get(0)?.as_u64()?,
296
                    del_gen: record.get(1)?.as_u64()?,
297
                })
298
            })
299
            .collect::<Option<Vec<_>>>()
300
            .ok_or_else(invalid)?;
301
        let next_id = manifest["next"].as_u64().ok_or_else(invalid)?;
302
        let mut index = Self::new(true);
303
        let missing = || {
304
            Error::new(
305
                ErrorKind::Store,
306
                "a search segment is missing or unreadable",
307
            )
308
        };
309
        for record in &records {
310
            let files = blob(&segment_key(record.id))
311
                .and_then(|bytes| unpack(&bytes))
312
                .ok_or_else(missing)?;
313
            let raw = RamDirectory::create();
314
            for (name, bytes) in files {
315
                raw.atomic_write(&PathBuf::from(name), &bytes)
316
                    .map_err(|error| tantivy_error("loading a search segment", error))?;
317
            }
318
            let mut segment = open_segment(record.id, raw, None)?;
319
            if segment.index.schema() != index.schema {
320
                return Err(Error::new(
321
                    ErrorKind::Store,
322
                    "a search segment has another schema",
323
                ));
324
            }
325
            if record.del_gen > 0 {
326
                let bytes = blob(&deletes_key(record.id, record.del_gen)).ok_or_else(missing)?;
327
                let alive = AliveBitSet::open(OwnedBytes::new(bytes));
328
                let mut bits = BitSet::with_max_value(segment.max_doc);
329
                for doc in alive.iter_alive().filter(|doc| *doc < segment.max_doc) {
330
                    bits.insert(doc);
331
                }
332
                segment.alive = Arc::new(bits);
333
            }
334
            segment.del_gen = record.del_gen;
335
            segment.persisted = true;
336
            index.segments.push(segment);
337
        }
338
        index.next_id = next_id.max(
339
            records
340
                .iter()
341
                .map(|record| record.id + 1)
342
                .max()
343
                .unwrap_or(1),
344
        );
345
        Ok(index)
346
    }
347
348
    fn records(&self) -> Vec<SegmentRecord> {
349
        self.segments
350
            .iter()
351
            .filter(|segment| segment.persisted)
352
            .map(|segment| SegmentRecord {
353
                id: segment.id,
354
                del_gen: segment.del_gen,
355
            })
356
            .collect()
357
    }
358
359
    fn document(&self, doc: &Doc) -> TantivyDocument {
360
        let fields = self.fields;
361
        let mut document = TantivyDocument::default();
362
        document.add_text(fields.key, &doc.key);
363
        document.add_text(fields.room, &doc.room);
364
        document.add_text(fields.sender, &doc.sender);
365
        document.add_text(fields.content, &doc.content);
366
        document.add_text(fields.text, &doc.text);
367
        document.add_u64(fields.ts, doc.ts);
368
        document.add_u64(fields.tie, tie(&doc.key));
369
        document.add_u64(fields.content_ts, doc.content_ts);
370
        if doc.provisional {
371
            document.add_text(fields.pending, &doc.text);
372
        }
373
        document
374
    }
375
376
    fn stored(&self, document: &TantivyDocument) -> Option<Doc> {
377
        let fields = self.fields;
378
        let text = |field| {
379
            document
380
                .get_first(field)
381
                .and_then(|value| value.as_str())
382
                .map(str::to_owned)
383
        };
384
        let number = |field| document.get_first(field).and_then(|value| value.as_u64());
385
        let pending = text(fields.pending);
386
        Some(Doc {
387
            key: text(fields.key)?,
388
            room: text(fields.room)?,
389
            sender: text(fields.sender)?,
390
            ts: number(fields.ts)?,
391
            content: text(fields.content)?,
392
            content_ts: number(fields.content_ts)?,
393
            provisional: pending.is_some(),
394
            text: pending.unwrap_or_default(),
395
        })
396
    }
397
398
    /// the live docs in `segment` whose `field` holds `value`
399
    fn find(&self, segment: &Segment, field: Field, value: &str) -> Vec<DocId> {
400
        let term = Term::from_field_text(field, value);
401
        let mut found = Vec::new();
402
        for reader in segment.searcher.segment_readers() {
403
            let Ok(inverted) = reader.inverted_index(field) else {
404
                continue;
405
            };
406
            let Ok(Some(mut postings)) = inverted.read_postings(&term, IndexRecordOption::Basic)
407
            else {
408
                continue;
409
            };
410
            let mut doc = postings.doc();
411
            while doc != tantivy::TERMINATED {
412
                if segment.alive.contains(doc) {
413
                    found.push(doc);
414
                }
415
                doc = postings.advance();
416
            }
417
        }
418
        found
419
    }
420
421
    fn lookup(&self, field: Field, value: &str) -> Option<Doc> {
422
        self.segments.iter().rev().find_map(|segment| {
423
            let doc = *self.find(segment, field, value).first()?;
424
            let document: TantivyDocument = segment.searcher.doc(DocAddress::new(0, doc)).ok()?;
425
            self.stored(&document)
426
        })
427
    }
428
429
    fn flush_buffer(&mut self) -> Result<(), Error> {
430
        if self.buffer.is_empty() {
431
            return Ok(());
432
        }
433
        let raw = RamDirectory::create();
434
        let index = Index::create(raw.clone(), self.schema.clone(), self.settings.clone())
435
            .map_err(|error| tantivy_error("creating a search segment", error))?;
436
        index.tokenizers().register(TOKENIZER, analyzer());
437
        let mut writer = SingleSegmentIndexWriter::new(index, WRITER_BUDGET)
438
            .map_err(|error| tantivy_error("writing a search segment", error))?;
439
        let mut docs: Vec<_> = self.buffer.values().collect();
440
        docs.sort_by_key(|doc| doc.ts);
441
        for doc in docs {
442
            writer
443
                .add_document(self.document(doc))
444
                .map_err(|error| tantivy_error("writing a search segment", error))?;
445
        }
446
        writer
447
            .finalize()
448
            .map_err(|error| tantivy_error("finishing a search segment", error))?;
449
        let segment = open_segment(self.next_id, raw, Some(self.buffer.len() as u32))?;
450
        self.next_id += 1;
451
        self.segments.push(segment);
452
        self.buffer.clear();
453
        Ok(())
454
    }
455
456
    fn class(live: usize) -> u32 {
457
        live.max(1).ilog(MERGE_FACTOR)
458
    }
459
460
    /// merges a size class once it fills up, and rewrites a segment that's mostly deleted
461
    fn compact(&mut self) -> Result<(), Error> {
462
        self.segments.retain(|segment| {
463
            let empty = segment.live() == 0;
464
            if empty && self.persistent {
465
                self.removed
466
                    .push((segment.id, segment.persisted, segment.del_gen));
467
            }
468
            !empty
469
        });
470
        loop {
471
            let mut classes: HashMap<u32, Vec<usize>> = HashMap::new();
472
            for (position, segment) in self.segments.iter().enumerate() {
473
                classes
474
                    .entry(Self::class(segment.live()))
475
                    .or_default()
476
                    .push(position);
477
            }
478
            let full = classes
479
                .into_values()
480
                .find(|positions| positions.len() >= MERGE_FACTOR);
481
            let sparse = self
482
                .segments
483
                .iter()
484
                .position(|segment| {
485
                    segment.max_doc >= 1_000 && segment.live() * 2 < segment.max_doc as usize
486
                })
487
                .map(|position| vec![position]);
488
            let Some(positions) = full.or(sparse) else {
489
                return Ok(());
490
            };
491
            self.merge(&positions)?;
492
        }
493
    }
494
495
    fn merge(&mut self, positions: &[usize]) -> Result<(), Error> {
496
        let mut inputs = Vec::new();
497
        let mut alive = Vec::new();
498
        for position in positions {
499
            let segment = &self.segments[*position];
500
            let searchable = segment
501
                .index
502
                .searchable_segments()
503
                .map_err(|error| tantivy_error("merging search segments", error))?;
504
            let mut bytes = Vec::new();
505
            write_alive_bitset(&segment.alive, &mut bytes)
506
                .map_err(|error| tantivy_error("merging search segments", error))?;
507
            for tantivy_segment in searchable {
508
                inputs.push(tantivy_segment);
509
                alive.push(Some(AliveBitSet::open(OwnedBytes::new(bytes.clone()))));
510
            }
511
        }
512
        let raw = RamDirectory::create();
513
        merge_filtered_segments(&inputs, self.settings.clone(), alive, raw.clone())
514
            .map_err(|error| tantivy_error("merging search segments", error))?;
515
        let merged = open_segment(self.next_id, raw, None)?;
516
        self.next_id += 1;
517
        let mut position = 0;
518
        self.segments.retain(|segment| {
519
            let merging = positions.contains(&position);
520
            position += 1;
521
            if merging && self.persistent {
522
                self.removed
523
                    .push((segment.id, segment.persisted, segment.del_gen));
524
            }
525
            !merging
526
        });
527
        self.segments.push(merged);
528
        Ok(())
529
    }
530
531
    fn segment_files(segment: &Segment) -> Result<Vec<u8>, Error> {
532
        let mut names: Vec<PathBuf> = segment
533
            .index
534
            .searchable_segment_metas()
535
            .map_err(|error| tantivy_error("reading a search segment", error))?
536
            .iter()
537
            .flat_map(|meta| meta.list_files())
538
            .collect();
539
        names.push(PathBuf::from("meta.json"));
540
        let mut files = Vec::new();
541
        for name in names {
542
            if !segment.raw.exists(&name).unwrap_or(false) {
543
                continue;
544
            }
545
            let bytes = segment
546
                .raw
547
                .open_read(&name)
548
                .map_err(|error| tantivy_error("reading a search segment", error))?
549
                .read_bytes()
550
                .map_err(|error| tantivy_error("reading a search segment", error))?;
551
            files.push((
552
                name.to_string_lossy().into_owned(),
553
                bytes.as_slice().to_vec(),
554
            ));
555
        }
556
        Ok(pack(&files))
557
    }
558
}
559
560
impl MessageIndex for KeywordIndex {
561
    fn get(&self, key: &str) -> Option<Doc> {
562
        if let Some(doc) = self.buffer.get(key) {
563
            return Some(doc.clone());
564
        }
565
        self.lookup(self.fields.key, key)
566
    }
567
568
    fn by_content(&self, content: &str) -> Option<Doc> {
569
        if let Some(doc) = self.buffer.values().find(|doc| doc.content == content) {
570
            return Some(doc.clone());
571
        }
572
        self.lookup(self.fields.content, content)
573
    }
574
575
    fn put(&mut self, doc: Doc) {
576
        self.remove(&doc.key);
577
        self.buffer.insert(doc.key.clone(), doc);
578
    }
579
580
    fn remove(&mut self, key: &str) -> bool {
581
        let mut removed = self.buffer.remove(key).is_some();
582
        let key_field = self.fields.key;
583
        for position in 0..self.segments.len() {
584
            let docs = self.find(&self.segments[position], key_field, key);
585
            if docs.is_empty() {
586
                continue;
587
            }
588
            let segment = &mut self.segments[position];
589
            let alive = Arc::make_mut(&mut segment.alive);
590
            for doc in docs {
591
                alive.remove(doc);
592
            }
593
            segment.deletes_dirty = true;
594
            removed = true;
595
        }
596
        removed
597
    }
598
599
    fn buffered(&self) -> usize {
600
        self.buffer.len()
601
    }
602
603
    fn search(
604
        &mut self,
605
        query: &str,
606
        rooms: Option<&[String]>,
607
        before: Option<(u64, u64)>,
608
        limit: usize,
609
    ) -> Result<(Vec<Hit>, Vec<String>), Error> {
610
        self.flush_buffer()?;
611
        let (text_query, _) = self.parser.parse_query_lenient(query);
612
        let mut terms = Vec::new();
613
        text_query.query_terms(&mut |term, _| {
614
            if term.field() == self.fields.text
615
                && let Some(text) = term.value().as_str()
616
            {
617
                terms.push(text.to_owned());
618
            }
619
        });
620
        let query: Box<dyn Query> = match rooms {
621
            None => text_query,
622
            Some(rooms) => Box::new(BooleanQuery::new(vec![
623
                (Occur::Must, text_query),
624
                (
625
                    Occur::Must,
626
                    Box::new(TermSetQuery::new(
627
                        rooms
628
                            .iter()
629
                            .map(|room| Term::from_field_text(self.fields.room, room)),
630
                    )),
631
                ),
632
            ])),
633
        };
634
        let mut found = Vec::new();
635
        for segment in &self.segments {
636
            let collector = Newest {
637
                alive: segment.alive.clone(),
638
                before,
639
                limit,
640
            };
641
            let newest = segment
642
                .searcher
643
                .search(&*query, &collector)
644
                .map_err(|error| tantivy_error("searching messages", error))?;
645
            for (ts, tie, doc) in newest {
646
                found.push((ts, tie, segment.id, doc));
647
            }
648
        }
649
        found.sort_by(|a, b| (b.0, b.1).cmp(&(a.0, a.1)));
650
        found.truncate(limit);
651
        let mut hits = Vec::new();
652
        for (ts, tie, id, doc) in found {
653
            let segment = self
654
                .segments
655
                .iter()
656
                .find(|segment| segment.id == id)
657
                .expect("the segment was just searched");
658
            let document: TantivyDocument = segment
659
                .searcher
660
                .doc(doc)
661
                .map_err(|error| tantivy_error("reading a search hit", error))?;
662
            let Some(stored) = self.stored(&document) else {
663
                continue;
664
            };
665
            hits.push(Hit {
666
                key: stored.key,
667
                room: stored.room,
668
                sender: stored.sender,
669
                content: stored.content,
670
                ts,
671
                tie,
672
            });
673
        }
674
        terms.sort();
675
        terms.dedup();
676
        Ok((hits, terms))
677
    }
678
679
    fn maintain(&mut self) -> Result<(), Error> {
680
        self.flush_buffer()?;
681
        self.compact()
682
    }
683
684
    fn changes(&mut self) -> Result<Changes, Error> {
685
        self.maintain()?;
686
        let mut changes = Changes::default();
687
        for (id, persisted, del_gen) in &self.removed {
688
            if *persisted {
689
                changes.deletes.push(segment_key(*id));
690
            }
691
            if *del_gen > 0 {
692
                changes.deletes.push(deletes_key(*id, *del_gen));
693
            }
694
            changes.removed.push(*id);
695
        }
696
        for segment in &self.segments {
697
            if !segment.persisted {
698
                changes
699
                    .puts
700
                    .push((segment_key(segment.id), Self::segment_files(segment)?));
701
                changes.files.push(segment.id);
702
            }
703
            if segment.deletes_dirty && segment.live() < segment.max_doc as usize {
704
                let generation = segment.del_gen + 1;
705
                let mut bytes = Vec::new();
706
                write_alive_bitset(&segment.alive, &mut bytes)
707
                    .map_err(|error| tantivy_error("saving search deletes", error))?;
708
                changes
709
                    .puts
710
                    .push((deletes_key(segment.id, generation), bytes));
711
                if segment.del_gen > 0 {
712
                    changes
713
                        .deletes
714
                        .push(deletes_key(segment.id, segment.del_gen));
715
                }
716
                changes.dels.push((segment.id, generation));
717
            }
718
        }
719
        Ok(changes)
720
    }
721
722
    fn committed(&mut self, changes: &Changes) {
723
        self.removed
724
            .retain(|(id, ..)| !changes.removed.contains(id));
725
        for segment in &mut self.segments {
726
            if changes.files.contains(&segment.id) {
727
                segment.persisted = true;
728
            }
729
            if let Some((_, generation)) = changes.dels.iter().find(|(id, _)| *id == segment.id) {
730
                segment.del_gen = *generation;
731
                segment.deletes_dirty = false;
732
            }
733
        }
734
    }
735
736
    fn manifest(&self) -> serde_json::Value {
737
        let segments: Vec<_> = self
738
            .records()
739
            .into_iter()
740
            .map(|record| serde_json::json!([record.id, record.del_gen]))
741
            .collect();
742
        serde_json::json!({ "segments": segments, "next": self.next_id })
743
    }
744
}
745
746
/// the `limit` newest live docs older than `before`, by timestamp and then key hash
747
struct Newest {
748
    alive: Arc<BitSet>,
749
    before: Option<(u64, u64)>,
750
    limit: usize,
751
}
752
753
struct NewestSegment {
754
    alive: Arc<BitSet>,
755
    before: Option<(u64, u64)>,
756
    limit: usize,
757
    ts: tantivy::columnar::Column<u64>,
758
    tie: tantivy::columnar::Column<u64>,
759
    segment: u32,
760
    top: std::collections::BinaryHeap<std::cmp::Reverse<(u64, u64, DocAddress)>>,
761
}
762
763
impl Collector for Newest {
764
    type Fruit = Vec<(u64, u64, DocAddress)>;
765
    type Child = NewestSegment;
766
767
    fn for_segment(&self, segment: u32, reader: &SegmentReader) -> tantivy::Result<NewestSegment> {
768
        Ok(NewestSegment {
769
            alive: self.alive.clone(),
770
            before: self.before,
771
            limit: self.limit,
772
            ts: reader.fast_fields().u64("ts")?,
773
            tie: reader.fast_fields().u64("tie")?,
774
            segment,
775
            top: std::collections::BinaryHeap::new(),
776
        })
777
    }
778
779
    fn requires_scoring(&self) -> bool {
780
        false
781
    }
782
783
    fn merge_fruits(&self, fruits: Vec<Self::Fruit>) -> tantivy::Result<Self::Fruit> {
784
        let mut all: Vec<_> = fruits.into_iter().flatten().collect();
785
        all.sort_by(|a, b| (b.0, b.1).cmp(&(a.0, a.1)));
786
        all.truncate(self.limit);
787
        Ok(all)
788
    }
789
}
790
791
impl SegmentCollector for NewestSegment {
792
    type Fruit = Vec<(u64, u64, DocAddress)>;
793
794
    fn collect(&mut self, doc: DocId, _: Score) {
795
        if !self.alive.contains(doc) {
796
            return;
797
        }
798
        let (Some(ts), Some(tie)) = (self.ts.first(doc), self.tie.first(doc)) else {
799
            return;
800
        };
801
        if self.before.is_some_and(|before| (ts, tie) >= before) {
802
            return;
803
        }
804
        let entry = std::cmp::Reverse((ts, tie, DocAddress::new(self.segment, doc)));
805
        if self.top.len() < self.limit {
806
            self.top.push(entry);
807
        } else if self.top.peek().is_some_and(|smallest| entry < *smallest) {
808
            self.top.pop();
809
            self.top.push(entry);
810
        }
811
    }
812
813
    fn harvest(self) -> Self::Fruit {
814
        let mut top: Vec<_> = self
815
            .top
816
            .into_iter()
817
            .map(|std::cmp::Reverse(entry)| entry)
818
            .collect();
819
        top.sort_by(|a, b| (b.0, b.1).cmp(&(a.0, a.1)));
820
        top
821
    }
822
}
823
824
#[cfg(test)]
825
#[path = "index_tests.rs"]
826
mod tests;