| 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; |