1use std::collections::{BTreeMap, BTreeSet, HashSet};
16use std::iter;
17use std::ops::Range;
18use std::sync::Arc;
19use std::time::Instant;
20
21use common_base::range_read::RangeReader;
22use common_telemetry::{tracing, warn};
23use index::bloom_filter::applier::{BloomFilterApplier, InListPredicate};
24use index::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReaderImpl};
25use index::fulltext_index::search::{FulltextIndexSearcher, RowId, TantivyFulltextIndexSearcher};
26use index::fulltext_index::tokenizer::{ChineseTokenizer, EnglishTokenizer, Tokenizer};
27use index::fulltext_index::{Analyzer, Config};
28use index::target::IndexTarget;
29use object_store::ObjectStore;
30use puffin::puffin_manager::cache::PuffinMetadataCacheRef;
31use puffin::puffin_manager::{GuardWithMetadata, PuffinManager, PuffinReader};
32use snafu::ResultExt;
33use store_api::region_request::PathType;
34use store_api::storage::ColumnId;
35
36use crate::access_layer::{RegionFilePathFactory, WriteCachePathProvider};
37use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey};
38use crate::cache::index::bloom_filter_index::{
39 BloomFilterIndexCacheRef, CachedBloomFilterIndexBlobReader, Tag,
40};
41use crate::cache::index::result_cache::PredicateKey;
42use crate::error::{
43 ApplyBloomFilterIndexSnafu, ApplyFulltextIndexSnafu, MetadataSnafu, PuffinBuildReaderSnafu,
44 PuffinReadBlobSnafu, Result,
45};
46use crate::metrics::INDEX_APPLY_ELAPSED;
47use crate::sst::file::RegionIndexId;
48use crate::sst::index::fulltext_index::applier::builder::{FulltextRequest, FulltextTerm};
49use crate::sst::index::fulltext_index::{INDEX_BLOB_TYPE_BLOOM, INDEX_BLOB_TYPE_TANTIVY};
50use crate::sst::index::puffin_manager::{
51 PuffinManagerFactory, SstPuffinBlob, SstPuffinDir, SstPuffinReader,
52};
53use crate::sst::index::{TYPE_FULLTEXT_INDEX, trigger_index_background_download};
54
55pub mod builder;
56
57#[derive(Default, Clone)]
59pub struct FulltextIndexApplyMetrics {
60 pub apply_elapsed: std::time::Duration,
62 pub blob_cache_miss: usize,
64 pub dir_cache_hit: usize,
66 pub dir_cache_miss: usize,
68 pub dir_init_elapsed: std::time::Duration,
70 pub bloom_filter_read_metrics: BloomFilterReadMetrics,
72}
73
74impl std::fmt::Debug for FulltextIndexApplyMetrics {
75 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
76 let Self {
77 apply_elapsed,
78 blob_cache_miss,
79 dir_cache_hit,
80 dir_cache_miss,
81 dir_init_elapsed,
82 bloom_filter_read_metrics,
83 } = self;
84
85 if self.is_empty() {
86 return write!(f, "{{}}");
87 }
88 write!(f, "{{")?;
89
90 write!(f, "\"apply_elapsed\":\"{:?}\"", apply_elapsed)?;
91
92 if *blob_cache_miss > 0 {
93 write!(f, ", \"blob_cache_miss\":{}", blob_cache_miss)?;
94 }
95 if *dir_cache_hit > 0 {
96 write!(f, ", \"dir_cache_hit\":{}", dir_cache_hit)?;
97 }
98 if *dir_cache_miss > 0 {
99 write!(f, ", \"dir_cache_miss\":{}", dir_cache_miss)?;
100 }
101 if !dir_init_elapsed.is_zero() {
102 write!(f, ", \"dir_init_elapsed\":\"{:?}\"", dir_init_elapsed)?;
103 }
104 write!(
105 f,
106 ", \"bloom_filter_read_metrics\":{:?}",
107 bloom_filter_read_metrics
108 )?;
109
110 write!(f, "}}")
111 }
112}
113
114impl FulltextIndexApplyMetrics {
115 pub fn is_empty(&self) -> bool {
117 self.apply_elapsed.is_zero()
118 }
119
120 pub fn collect_dir_metrics(
122 &mut self,
123 elapsed: std::time::Duration,
124 dir_metrics: puffin::puffin_manager::DirMetrics,
125 ) {
126 self.dir_init_elapsed += elapsed;
127 if dir_metrics.cache_hit {
128 self.dir_cache_hit += 1;
129 } else {
130 self.dir_cache_miss += 1;
131 }
132 }
133
134 pub fn merge_from(&mut self, other: &Self) {
136 self.apply_elapsed += other.apply_elapsed;
137 self.blob_cache_miss += other.blob_cache_miss;
138 self.dir_cache_hit += other.dir_cache_hit;
139 self.dir_cache_miss += other.dir_cache_miss;
140 self.dir_init_elapsed += other.dir_init_elapsed;
141 self.bloom_filter_read_metrics
142 .merge_from(&other.bloom_filter_read_metrics);
143 }
144}
145
146pub struct FulltextIndexApplier {
148 requests: Arc<BTreeMap<ColumnId, FulltextRequest>>,
150
151 index_source: IndexSource,
153
154 bloom_filter_index_cache: Option<BloomFilterIndexCacheRef>,
156
157 predicate_key: PredicateKey,
159}
160
161pub type FulltextIndexApplierRef = Arc<FulltextIndexApplier>;
162
163impl FulltextIndexApplier {
164 pub fn new(
166 table_dir: String,
167 path_type: PathType,
168 store: ObjectStore,
169 requests: BTreeMap<ColumnId, FulltextRequest>,
170 puffin_manager_factory: PuffinManagerFactory,
171 ) -> Self {
172 let requests = Arc::new(requests);
173 let index_source = IndexSource::new(table_dir, path_type, puffin_manager_factory, store);
174
175 Self {
176 predicate_key: PredicateKey::new_fulltext(requests.clone()),
177 requests,
178 index_source,
179 bloom_filter_index_cache: None,
180 }
181 }
182
183 pub fn with_file_cache(mut self, file_cache: Option<FileCacheRef>) -> Self {
185 self.index_source.set_file_cache(file_cache);
186 self
187 }
188
189 pub fn with_puffin_metadata_cache(
191 mut self,
192 puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
193 ) -> Self {
194 self.index_source
195 .set_puffin_metadata_cache(puffin_metadata_cache);
196 self
197 }
198
199 pub fn with_bloom_filter_cache(
201 mut self,
202 bloom_filter_index_cache: Option<BloomFilterIndexCacheRef>,
203 ) -> Self {
204 self.bloom_filter_index_cache = bloom_filter_index_cache;
205 self
206 }
207
208 pub fn predicate_key(&self) -> &PredicateKey {
210 &self.predicate_key
211 }
212}
213
214impl FulltextIndexApplier {
215 #[tracing::instrument(
223 skip_all,
224 fields(file_id = %file_id)
225 )]
226 pub async fn apply_fine(
227 &self,
228 file_id: RegionIndexId,
229 file_size_hint: Option<u64>,
230 mut metrics: Option<&mut FulltextIndexApplyMetrics>,
231 ) -> Result<Option<BTreeSet<RowId>>> {
232 let apply_start = Instant::now();
233
234 let mut row_ids: Option<BTreeSet<RowId>> = None;
235 for (column_id, request) in self.requests.iter() {
236 if request.queries.is_empty() && request.terms.is_empty() {
237 continue;
238 }
239
240 let Some(result) = self
241 .apply_fine_one_column(
242 file_size_hint,
243 file_id,
244 *column_id,
245 request,
246 metrics.as_deref_mut(),
247 )
248 .await?
249 else {
250 continue;
251 };
252
253 if let Some(ids) = row_ids.as_mut() {
254 ids.retain(|id| result.contains(id));
255 } else {
256 row_ids = Some(result);
257 }
258
259 if let Some(ids) = row_ids.as_ref()
260 && ids.is_empty()
261 {
262 break;
263 }
264 }
265
266 let elapsed = apply_start.elapsed();
268 INDEX_APPLY_ELAPSED
269 .with_label_values(&[TYPE_FULLTEXT_INDEX])
270 .observe(elapsed.as_secs_f64());
271
272 if let Some(m) = metrics {
273 m.apply_elapsed += elapsed;
274 }
275
276 Ok(row_ids)
277 }
278
279 async fn apply_fine_one_column(
280 &self,
281 file_size_hint: Option<u64>,
282 file_id: RegionIndexId,
283 column_id: ColumnId,
284 request: &FulltextRequest,
285 metrics: Option<&mut FulltextIndexApplyMetrics>,
286 ) -> Result<Option<BTreeSet<RowId>>> {
287 let blob_key = format!(
288 "{INDEX_BLOB_TYPE_TANTIVY}-{}",
289 IndexTarget::ColumnId(column_id)
290 );
291 let dir = self
292 .index_source
293 .dir(file_id, &blob_key, file_size_hint, metrics)
294 .await?;
295
296 let dir = match &dir {
297 Some(dir) => dir,
298 None => {
299 return Ok(None);
300 }
301 };
302
303 let config = Config::from_blob_metadata(dir.metadata()).context(ApplyFulltextIndexSnafu)?;
304 let path = dir.path();
305
306 let searcher =
307 TantivyFulltextIndexSearcher::new(path, config).context(ApplyFulltextIndexSnafu)?;
308 let mut row_ids: Option<BTreeSet<RowId>> = None;
309
310 for query in &request.queries {
312 let result = searcher
313 .search(&query.0)
314 .await
315 .context(ApplyFulltextIndexSnafu)?;
316
317 if let Some(ids) = row_ids.as_mut() {
318 ids.retain(|id| result.contains(id));
319 } else {
320 row_ids = Some(result);
321 }
322
323 if let Some(ids) = row_ids.as_ref()
324 && ids.is_empty()
325 {
326 break;
327 }
328 }
329
330 let query = request.terms_as_query(config.case_sensitive);
332 if !query.0.is_empty() {
333 let result = searcher
334 .search(&query.0)
335 .await
336 .context(ApplyFulltextIndexSnafu)?;
337
338 if let Some(ids) = row_ids.as_mut() {
339 ids.retain(|id| result.contains(id));
340 } else {
341 row_ids = Some(result);
342 }
343 }
344
345 Ok(row_ids)
346 }
347}
348
349impl FulltextIndexApplier {
350 #[allow(clippy::type_complexity)]
362 #[tracing::instrument(
363 skip_all,
364 fields(file_id = %file_id)
365 )]
366 pub async fn apply_coarse(
367 &self,
368 file_id: RegionIndexId,
369 file_size_hint: Option<u64>,
370 row_groups: impl Iterator<Item = (usize, bool)>,
371 mut metrics: Option<&mut FulltextIndexApplyMetrics>,
372 ) -> Result<Option<Vec<(usize, Vec<Range<usize>>)>>> {
373 let apply_start = Instant::now();
374
375 let (input, mut output) = Self::init_coarse_output(row_groups);
376 let mut applied = false;
377
378 for (column_id, request) in self.requests.iter() {
379 if request.terms.is_empty() {
380 continue;
382 }
383
384 applied |= self
385 .apply_coarse_one_column(
386 file_id,
387 file_size_hint,
388 *column_id,
389 &request.terms,
390 &mut output,
391 metrics.as_deref_mut(),
392 )
393 .await?;
394 }
395
396 if !applied {
397 return Ok(None);
398 }
399
400 Self::adjust_coarse_output(input, &mut output);
401
402 let elapsed = apply_start.elapsed();
404 INDEX_APPLY_ELAPSED
405 .with_label_values(&[TYPE_FULLTEXT_INDEX])
406 .observe(elapsed.as_secs_f64());
407
408 if let Some(m) = metrics {
409 m.apply_elapsed += elapsed;
410 }
411
412 Ok(Some(output))
413 }
414
415 async fn apply_coarse_one_column(
416 &self,
417 file_id: RegionIndexId,
418 file_size_hint: Option<u64>,
419 column_id: ColumnId,
420 terms: &[FulltextTerm],
421 output: &mut [(usize, Vec<Range<usize>>)],
422 mut metrics: Option<&mut FulltextIndexApplyMetrics>,
423 ) -> Result<bool> {
424 let blob_key = format!(
425 "{INDEX_BLOB_TYPE_BLOOM}-{}",
426 IndexTarget::ColumnId(column_id)
427 );
428 let Some(reader) = self
429 .index_source
430 .blob(file_id, &blob_key, file_size_hint, metrics.as_deref_mut())
431 .await?
432 else {
433 return Ok(false);
434 };
435 let config =
436 Config::from_blob_metadata(reader.metadata()).context(ApplyFulltextIndexSnafu)?;
437
438 let predicates = Self::terms_to_predicates(terms, &config);
439 if predicates.is_empty() {
440 return Ok(false);
441 }
442
443 let range_reader = reader.reader().await.context(PuffinBuildReaderSnafu)?;
444 let reader = if let Some(bloom_filter_cache) = &self.bloom_filter_index_cache {
445 let blob_size = range_reader
446 .metadata()
447 .await
448 .context(MetadataSnafu)?
449 .content_length;
450 let reader = CachedBloomFilterIndexBlobReader::new(
451 file_id.file_id(),
452 file_id.version,
453 column_id,
454 Tag::Fulltext,
455 blob_size,
456 BloomFilterReaderImpl::new(range_reader),
457 bloom_filter_cache.clone(),
458 );
459 Box::new(reader) as _
460 } else {
461 Box::new(BloomFilterReaderImpl::new(range_reader)) as _
462 };
463
464 let mut applier = BloomFilterApplier::new(reader)
465 .await
466 .context(ApplyBloomFilterIndexSnafu)?;
467 let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::<Vec<_>>();
468 applier
469 .search_groups(
470 &predicates,
471 &mut row_groups,
472 metrics.map(|m| &mut m.bloom_filter_read_metrics),
473 )
474 .await
475 .context(ApplyBloomFilterIndexSnafu)?;
476
477 Ok(true)
478 }
479
480 #[allow(clippy::type_complexity)]
488 fn init_coarse_output(
489 row_groups: impl Iterator<Item = (usize, bool)>,
490 ) -> (Vec<(usize, Range<usize>)>, Vec<(usize, Vec<Range<usize>>)>) {
491 let mut input = Vec::with_capacity(row_groups.size_hint().0);
493 let mut start = 0;
494 for (i, (len, to_search)) in row_groups.enumerate() {
495 let end = start + len;
496 if to_search {
497 input.push((i, start..end));
498 }
499 start = end;
500 }
501
502 let output = input
505 .iter()
506 .map(|(i, range)| (*i, vec![range.clone()]))
507 .collect::<Vec<_>>();
508
509 (input, output)
510 }
511
512 fn adjust_coarse_output(
514 input: Vec<(usize, Range<usize>)>,
515 output: &mut [(usize, Vec<Range<usize>>)],
516 ) {
517 for ((_, output), (_, input)) in output.iter_mut().zip(input) {
519 let start = input.start;
520 for range in output.iter_mut() {
521 range.start -= start;
522 range.end -= start;
523 }
524 }
525 }
526
527 fn terms_to_predicates(terms: &[FulltextTerm], config: &Config) -> Vec<InListPredicate> {
532 let mut probes = HashSet::new();
533 for term in terms {
534 if config.case_sensitive && term.col_lowered {
535 continue;
537 }
538 probes.extend(Self::term_to_probes(&term.term, config));
539 }
540
541 probes
542 .into_iter()
543 .map(|p| InListPredicate {
544 list: iter::once(p).collect(),
545 })
546 .collect::<Vec<_>>()
547 }
548
549 fn term_to_probes<'a>(term: &'a str, config: &'a Config) -> impl Iterator<Item = Vec<u8>> + 'a {
550 let tokens = match config.analyzer {
551 Analyzer::English => EnglishTokenizer {}.tokenize(term),
552 Analyzer::Chinese => ChineseTokenizer {}.tokenize(term),
553 };
554
555 tokens.into_iter().map(|t| {
556 if !config.case_sensitive {
557 t.to_lowercase()
558 } else {
559 t.to_string()
560 }
561 .into_bytes()
562 })
563 }
564}
565
566struct IndexSource {
568 table_dir: String,
569
570 path_type: PathType,
572
573 puffin_manager_factory: PuffinManagerFactory,
575
576 remote_store: ObjectStore,
578
579 file_cache: Option<FileCacheRef>,
581
582 puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
584}
585
586impl IndexSource {
587 fn new(
588 table_dir: String,
589 path_type: PathType,
590 puffin_manager_factory: PuffinManagerFactory,
591 remote_store: ObjectStore,
592 ) -> Self {
593 Self {
594 table_dir,
595 path_type,
596 puffin_manager_factory,
597 remote_store,
598 file_cache: None,
599 puffin_metadata_cache: None,
600 }
601 }
602
603 fn set_file_cache(&mut self, file_cache: Option<FileCacheRef>) {
604 self.file_cache = file_cache;
605 }
606
607 fn set_puffin_metadata_cache(&mut self, puffin_metadata_cache: Option<PuffinMetadataCacheRef>) {
608 self.puffin_metadata_cache = puffin_metadata_cache;
609 }
610
611 async fn blob(
615 &self,
616 file_id: RegionIndexId,
617 key: &str,
618 file_size_hint: Option<u64>,
619 metrics: Option<&mut FulltextIndexApplyMetrics>,
620 ) -> Result<Option<GuardWithMetadata<SstPuffinBlob>>> {
621 let (reader, fallbacked) = self.ensure_reader(file_id, file_size_hint).await?;
622
623 if fallbacked && let Some(m) = metrics {
625 m.blob_cache_miss += 1;
626 }
627
628 let res = reader.blob(key).await;
629 match res {
630 Ok(blob) => Ok(Some(blob)),
631 Err(err) if err.is_blob_not_found() => Ok(None),
632 Err(err) => {
633 if fallbacked {
634 Err(err).context(PuffinReadBlobSnafu)
635 } else {
636 warn!(err; "An unexpected error occurred while reading the cached index file. Fallback to remote index file.");
637 let reader = self.build_remote(file_id, file_size_hint).await?;
638 let res = reader.blob(key).await;
639 match res {
640 Ok(blob) => Ok(Some(blob)),
641 Err(err) if err.is_blob_not_found() => Ok(None),
642 Err(err) => Err(err).context(PuffinReadBlobSnafu),
643 }
644 }
645 }
646 }
647 }
648
649 async fn dir(
653 &self,
654 file_id: RegionIndexId,
655 key: &str,
656 file_size_hint: Option<u64>,
657 mut metrics: Option<&mut FulltextIndexApplyMetrics>,
658 ) -> Result<Option<GuardWithMetadata<SstPuffinDir>>> {
659 let (reader, fallbacked) = self.ensure_reader(file_id, file_size_hint).await?;
660
661 if fallbacked && let Some(m) = &mut metrics {
663 m.blob_cache_miss += 1;
664 }
665
666 let start = metrics.as_ref().map(|_| Instant::now());
667 let res = reader.dir(key).await;
668 match res {
669 Ok((dir, dir_metrics)) => {
670 if let Some(m) = metrics {
671 m.collect_dir_metrics(start.unwrap().elapsed(), dir_metrics);
673 }
674 Ok(Some(dir))
675 }
676 Err(err) if err.is_blob_not_found() => Ok(None),
677 Err(err) => {
678 if fallbacked {
679 Err(err).context(PuffinReadBlobSnafu)
680 } else {
681 warn!(err; "An unexpected error occurred while reading the cached index file. Fallback to remote index file.");
682 let reader = self.build_remote(file_id, file_size_hint).await?;
683 let start = metrics.as_ref().map(|_| Instant::now());
684 let res = reader.dir(key).await;
685 match res {
686 Ok((dir, dir_metrics)) => {
687 if let Some(m) = metrics {
688 m.collect_dir_metrics(start.unwrap().elapsed(), dir_metrics);
690 }
691 Ok(Some(dir))
692 }
693 Err(err) if err.is_blob_not_found() => Ok(None),
694 Err(err) => Err(err).context(PuffinReadBlobSnafu),
695 }
696 }
697 }
698 }
699 }
700
701 async fn ensure_reader(
703 &self,
704 file_id: RegionIndexId,
705 file_size_hint: Option<u64>,
706 ) -> Result<(SstPuffinReader, bool)> {
707 match self.build_local_cache(file_id, file_size_hint).await {
708 Ok(Some(r)) => Ok((r, false)),
709 Ok(None) => Ok((self.build_remote(file_id, file_size_hint).await?, true)),
710 Err(err) => Err(err),
711 }
712 }
713
714 async fn build_local_cache(
715 &self,
716 file_id: RegionIndexId,
717 file_size_hint: Option<u64>,
718 ) -> Result<Option<SstPuffinReader>> {
719 let Some(file_cache) = &self.file_cache else {
720 return Ok(None);
721 };
722
723 let index_key = IndexKey::new(
724 file_id.region_id(),
725 file_id.file_id(),
726 FileType::Puffin(file_id.version),
727 );
728 if file_cache.get(index_key).await.is_none() {
729 return Ok(None);
730 };
731
732 let puffin_manager = self
733 .puffin_manager_factory
734 .build(
735 file_cache.local_store(),
736 WriteCachePathProvider::new(file_cache.clone()),
737 )
738 .with_puffin_metadata_cache(self.puffin_metadata_cache.clone());
739 let reader = puffin_manager
740 .reader(&file_id)
741 .await
742 .context(PuffinBuildReaderSnafu)?
743 .with_file_size_hint(file_size_hint);
744 Ok(Some(reader))
745 }
746
747 async fn build_remote(
748 &self,
749 file_id: RegionIndexId,
750 file_size_hint: Option<u64>,
751 ) -> Result<SstPuffinReader> {
752 let path_factory = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type);
753
754 trigger_index_background_download(
756 self.file_cache.as_ref(),
757 &file_id,
758 file_size_hint,
759 &path_factory,
760 &self.remote_store,
761 );
762
763 let puffin_manager = self
764 .puffin_manager_factory
765 .build(self.remote_store.clone(), path_factory)
766 .with_puffin_metadata_cache(self.puffin_metadata_cache.clone());
767
768 let reader = puffin_manager
769 .reader(&file_id)
770 .await
771 .context(PuffinBuildReaderSnafu)?
772 .with_file_size_hint(file_size_hint);
773
774 Ok(reader)
775 }
776}