1mod builder;
16
17use std::collections::BTreeMap;
18use std::ops::Range;
19use std::sync::Arc;
20use std::time::Instant;
21
22use common_base::range_read::RangeReader;
23use common_telemetry::{tracing, warn};
24use datatypes::data_type::ConcreteDataType;
25use index::bloom_filter::applier::{BloomFilterApplier, InListPredicate};
26use index::bloom_filter::reader::{
27 BloomFilterReadMetrics, BloomFilterReader, BloomFilterReaderImpl,
28};
29use index::target::IndexTarget;
30use object_store::ObjectStore;
31use puffin::puffin_manager::cache::PuffinMetadataCacheRef;
32use puffin::puffin_manager::{PuffinManager, PuffinReader};
33use snafu::ResultExt;
34use store_api::metadata::RegionMetadataRef;
35use store_api::region_request::PathType;
36
37use crate::access_layer::{RegionFilePathFactory, WriteCachePathProvider};
38use crate::cache::file_cache::{FileCacheRef, FileType, IndexKey};
39use crate::cache::index::bloom_filter_index::{
40 BloomFilterIndexCacheRef, CachedBloomFilterIndexBlobReader, Tag,
41};
42use crate::error::{
43 ApplyBloomFilterIndexSnafu, Error, MetadataSnafu, PuffinBuildReaderSnafu, PuffinReadBlobSnafu,
44 Result,
45};
46use crate::metrics::INDEX_APPLY_ELAPSED;
47use crate::sst::file::RegionIndexId;
48use crate::sst::index::bloom_filter::INDEX_BLOB_TYPE;
49pub use crate::sst::index::bloom_filter::applier::builder::BloomFilterIndexApplierBuilder;
50use crate::sst::index::column::skipping_target_data_type;
51use crate::sst::index::puffin_manager::{BlobReader, PuffinManagerFactory};
52use crate::sst::index::{TYPE_BLOOM_FILTER_INDEX, trigger_index_background_download};
53
54#[derive(Default, Clone)]
56pub struct BloomFilterIndexApplyMetrics {
57 pub apply_elapsed: std::time::Duration,
59 pub blob_cache_miss: usize,
61 pub blob_read_bytes: u64,
63 pub read_metrics: BloomFilterReadMetrics,
65}
66
67impl std::fmt::Debug for BloomFilterIndexApplyMetrics {
68 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
69 let Self {
70 apply_elapsed,
71 blob_cache_miss,
72 blob_read_bytes,
73 read_metrics,
74 } = self;
75
76 if self.is_empty() {
77 return write!(f, "{{}}");
78 }
79 write!(f, "{{")?;
80
81 write!(f, "\"apply_elapsed\":\"{:?}\"", apply_elapsed)?;
82
83 if *blob_cache_miss > 0 {
84 write!(f, ", \"blob_cache_miss\":{}", blob_cache_miss)?;
85 }
86 if *blob_read_bytes > 0 {
87 write!(f, ", \"blob_read_bytes\":{}", blob_read_bytes)?;
88 }
89 write!(f, ", \"read_metrics\":{:?}", read_metrics)?;
90
91 write!(f, "}}")
92 }
93}
94
95impl BloomFilterIndexApplyMetrics {
96 pub fn is_empty(&self) -> bool {
98 self.apply_elapsed.is_zero()
99 }
100
101 pub fn merge_from(&mut self, other: &Self) {
103 self.apply_elapsed += other.apply_elapsed;
104 self.blob_cache_miss += other.blob_cache_miss;
105 self.blob_read_bytes += other.blob_read_bytes;
106 self.read_metrics.merge_from(&other.read_metrics);
107 }
108}
109
110pub(crate) type BloomFilterIndexApplierRef = Arc<BloomFilterIndexApplier>;
111
112pub struct BloomFilterIndexApplier {
114 table_dir: String,
116
117 path_type: PathType,
119
120 object_store: ObjectStore,
122
123 file_cache: Option<FileCacheRef>,
125
126 puffin_manager_factory: PuffinManagerFactory,
128
129 puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
131
132 bloom_filter_index_cache: Option<BloomFilterIndexCacheRef>,
134
135 default_predicates: Arc<BTreeMap<IndexTarget, Vec<InListPredicate>>>,
138
139 expected_predicate_target_types: BTreeMap<IndexTarget, ConcreteDataType>,
141}
142
143impl BloomFilterIndexApplier {
144 pub fn new(
148 table_dir: String,
149 path_type: PathType,
150 object_store: ObjectStore,
151 puffin_manager_factory: PuffinManagerFactory,
152 predicates: BTreeMap<IndexTarget, Vec<InListPredicate>>,
153 expected_predicate_target_types: BTreeMap<IndexTarget, ConcreteDataType>,
154 ) -> Self {
155 let default_predicates = Arc::new(predicates);
156 Self {
157 table_dir,
158 path_type,
159 object_store,
160 file_cache: None,
161 puffin_manager_factory,
162 puffin_metadata_cache: None,
163 bloom_filter_index_cache: None,
164 default_predicates,
165 expected_predicate_target_types,
166 }
167 }
168
169 pub fn with_file_cache(mut self, file_cache: Option<FileCacheRef>) -> Self {
170 self.file_cache = file_cache;
171 self
172 }
173
174 pub fn with_puffin_metadata_cache(
175 mut self,
176 puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
177 ) -> Self {
178 self.puffin_metadata_cache = puffin_metadata_cache;
179 self
180 }
181
182 pub fn with_bloom_filter_cache(
183 mut self,
184 bloom_filter_index_cache: Option<BloomFilterIndexCacheRef>,
185 ) -> Self {
186 self.bloom_filter_index_cache = bloom_filter_index_cache;
187 self
188 }
189
190 #[tracing::instrument(
205 skip_all,
206 fields(file_id = %file_id)
207 )]
208 pub async fn apply(
209 &self,
210 file_id: RegionIndexId,
211 file_size_hint: Option<u64>,
212 predicates: &BTreeMap<IndexTarget, Vec<InListPredicate>>,
213 row_groups: impl Iterator<Item = (usize, bool)>,
214 mut metrics: Option<&mut BloomFilterIndexApplyMetrics>,
215 ) -> Result<Vec<(usize, Vec<Range<usize>>)>> {
216 let apply_start = Instant::now();
217
218 let mut input = Vec::with_capacity(row_groups.size_hint().0);
220 let mut start = 0;
221 for (i, (len, to_search)) in row_groups.enumerate() {
222 let end = start + len;
223 if to_search {
224 input.push((i, start..end));
225 }
226 start = end;
227 }
228
229 let mut output = input
232 .iter()
233 .map(|(i, range)| (*i, vec![range.clone()]))
234 .collect::<Vec<_>>();
235
236 for (target, predicates) in predicates {
237 let blob = match self
238 .blob_reader(file_id, target, file_size_hint, metrics.as_deref_mut())
239 .await?
240 {
241 Some(blob) => blob,
242 None => continue,
243 };
244
245 if let Some(bloom_filter_cache) = &self.bloom_filter_index_cache {
247 let blob_size = blob.metadata().await.context(MetadataSnafu)?.content_length;
248 if let Some(m) = &mut metrics {
249 m.blob_read_bytes += blob_size;
250 }
251 let reader = CachedBloomFilterIndexBlobReader::new(
252 file_id.file_id(),
253 file_id.version,
254 target.clone(),
255 Tag::Skipping,
256 blob_size,
257 BloomFilterReaderImpl::new(blob),
258 bloom_filter_cache.clone(),
259 );
260 self.apply_predicates(reader, predicates, &mut output, metrics.as_deref_mut())
261 .await
262 .context(ApplyBloomFilterIndexSnafu)?;
263 } else {
264 let reader = BloomFilterReaderImpl::new(blob);
265 self.apply_predicates(reader, predicates, &mut output, metrics.as_deref_mut())
266 .await
267 .context(ApplyBloomFilterIndexSnafu)?;
268 }
269 }
270
271 for ((_, output), (_, input)) in output.iter_mut().zip(input) {
273 let start = input.start;
274 for range in output.iter_mut() {
275 range.start -= start;
276 range.end -= start;
277 }
278 }
279
280 let elapsed = apply_start.elapsed();
282 INDEX_APPLY_ELAPSED
283 .with_label_values(&[TYPE_BLOOM_FILTER_INDEX])
284 .observe(elapsed.as_secs_f64());
285
286 if let Some(m) = metrics {
287 m.apply_elapsed += elapsed;
288 }
289
290 Ok(output)
291 }
292
293 async fn blob_reader(
297 &self,
298 file_id: RegionIndexId,
299 target: &IndexTarget,
300 file_size_hint: Option<u64>,
301 metrics: Option<&mut BloomFilterIndexApplyMetrics>,
302 ) -> Result<Option<BlobReader>> {
303 let reader = match self
304 .cached_blob_reader(file_id, target, file_size_hint)
305 .await
306 {
307 Ok(Some(puffin_reader)) => puffin_reader,
308 other => {
309 if let Some(m) = metrics {
310 m.blob_cache_miss += 1;
311 }
312 if let Err(err) = other {
313 if is_blob_not_found(&err) {
315 return Ok(None);
316 }
317 warn!(err; "An unexpected error occurred while reading the cached index file. Fallback to remote index file.")
318 }
319 let res = self
320 .remote_blob_reader(file_id, target, file_size_hint)
321 .await;
322 if let Err(err) = res {
323 if is_blob_not_found(&err) {
325 return Ok(None);
326 }
327 return Err(err);
328 }
329
330 res?
331 }
332 };
333
334 Ok(Some(reader))
335 }
336
337 async fn cached_blob_reader(
339 &self,
340 file_id: RegionIndexId,
341 target: &IndexTarget,
342 file_size_hint: Option<u64>,
343 ) -> Result<Option<BlobReader>> {
344 let Some(file_cache) = &self.file_cache else {
345 return Ok(None);
346 };
347
348 let index_key = IndexKey::new(
349 file_id.region_id(),
350 file_id.file_id(),
351 FileType::Puffin(file_id.version),
352 );
353 if file_cache.get(index_key).await.is_none() {
354 return Ok(None);
355 };
356
357 let puffin_manager = self.puffin_manager_factory.build(
358 file_cache.local_store(),
359 WriteCachePathProvider::new(file_cache.clone()),
360 );
361 let blob_name = Self::target_blob_name(target);
362
363 let reader = puffin_manager
364 .reader(&file_id)
365 .await
366 .context(PuffinBuildReaderSnafu)?
367 .with_file_size_hint(file_size_hint)
368 .blob(&blob_name)
369 .await
370 .context(PuffinReadBlobSnafu)?
371 .reader()
372 .await
373 .context(PuffinBuildReaderSnafu)?;
374 Ok(Some(reader))
375 }
376
377 fn target_blob_name(target: &IndexTarget) -> String {
379 let target_key = target.encode();
380 format!("{INDEX_BLOB_TYPE}-{target_key}")
381 }
382
383 async fn remote_blob_reader(
385 &self,
386 file_id: RegionIndexId,
387 target: &IndexTarget,
388 file_size_hint: Option<u64>,
389 ) -> Result<BlobReader> {
390 let path_factory = RegionFilePathFactory::new(self.table_dir.clone(), self.path_type);
391
392 trigger_index_background_download(
394 self.file_cache.as_ref(),
395 &file_id,
396 file_size_hint,
397 &path_factory,
398 &self.object_store,
399 );
400
401 let puffin_manager = self
402 .puffin_manager_factory
403 .build(self.object_store.clone(), path_factory)
404 .with_puffin_metadata_cache(self.puffin_metadata_cache.clone());
405
406 let blob_name = Self::target_blob_name(target);
407
408 puffin_manager
409 .reader(&file_id)
410 .await
411 .context(PuffinBuildReaderSnafu)?
412 .with_file_size_hint(file_size_hint)
413 .blob(&blob_name)
414 .await
415 .context(PuffinReadBlobSnafu)?
416 .reader()
417 .await
418 .context(PuffinBuildReaderSnafu)
419 }
420
421 async fn apply_predicates<R: BloomFilterReader + Send + 'static>(
422 &self,
423 reader: R,
424 predicates: &[InListPredicate],
425 output: &mut [(usize, Vec<Range<usize>>)],
426 metrics: Option<&mut BloomFilterIndexApplyMetrics>,
427 ) -> std::result::Result<(), index::bloom_filter::error::Error> {
428 let mut applier = BloomFilterApplier::new(Box::new(reader)).await?;
429 let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::<Vec<_>>();
430 let read_metrics = metrics.map(|m| &mut m.read_metrics);
431 applier
432 .search_groups(predicates, &mut row_groups, read_metrics)
433 .await
434 }
435
436 pub fn compatible_predicate_for_sst(
440 &self,
441 sst_metadata: &RegionMetadataRef,
442 ) -> Option<Arc<BTreeMap<IndexTarget, Vec<InListPredicate>>>> {
443 let mut has_type_mismatch = false;
444 let mut compatible_targets = Vec::new();
445
446 for (target, expected) in &self.expected_predicate_target_types {
447 if skipping_target_data_type(sst_metadata, target).as_ref() != Some(expected) {
448 has_type_mismatch = true;
449 continue;
450 }
451 compatible_targets.push(target.clone());
452 }
453
454 if compatible_targets.is_empty() {
455 return None;
456 }
457
458 if !has_type_mismatch {
459 return Some(self.default_predicates.clone());
460 }
461
462 let mut compatible_predicates = BTreeMap::new();
463 for target in compatible_targets {
464 if let Some(predicates) = self.default_predicates.get(&target) {
465 compatible_predicates.insert(target, predicates.clone());
466 }
467 }
468
469 Some(Arc::new(compatible_predicates))
470 }
471}
472
473fn is_blob_not_found(err: &Error) -> bool {
474 matches!(
475 err,
476 Error::PuffinReadBlob {
477 source: puffin::error::Error::BlobNotFound { .. },
478 ..
479 }
480 )
481}
482
483#[cfg(test)]
484mod tests {
485 use std::collections::BTreeSet;
486
487 use datafusion_expr::{Expr, col, lit};
488 use futures::future::BoxFuture;
489 use index::Bytes;
490 use object_store::services::Memory;
491 use puffin::puffin_manager::PuffinWriter;
492 use store_api::metadata::RegionMetadata;
493 use store_api::storage::FileId;
494
495 use super::*;
496 use crate::sst::file::RegionFileId;
497 use crate::sst::index::IndexBuildType;
498 use crate::sst::index::bloom_filter::creator::BloomFilterIndexer;
499 use crate::sst::index::bloom_filter::creator::tests::{
500 mock_object_store, mock_region_metadata, new_batch, new_intm_mgr,
501 };
502
503 #[tokio::test]
504 async fn test_json_skipping_build_and_prune_both_formats() {
505 use crate::read::{Batch, BatchColumn};
506 use common_function::scalars::json::json_get::JsonGetWithType;
507 use common_function::scalars::udf::create_udf;
508 use datafusion_common::ScalarValue;
509 use datafusion_expr::expr::ScalarFunction;
510 use datatypes::arrow::array::{ArrayRef, Int64Array, StructArray};
511 use datatypes::arrow::buffer::NullBuffer;
512 use datatypes::arrow::datatypes::{Field, Schema};
513 use datatypes::arrow::record_batch::RecordBatch;
514 use datatypes::schema::SkippingIndexOptions;
515 use datatypes::vectors::{Helper, UInt8Vector, UInt64Vector};
516
517 for flat in [false, true] {
518 let (dir, factory) =
519 PuffinManagerFactory::new_for_test_async("json_skipping_index").await;
520 let mut metadata = (*mock_region_metadata()).clone();
521 let column = &mut metadata.column_metadatas[2].column_schema;
522 column.data_type =
523 ConcreteDataType::json2(datatypes::types::json_type::JsonNativeType::Object(
524 datatypes::types::json_type::JsonObjectType::new(),
525 ));
526 column.with_extension_type(&datatypes::extension::json::Json2ExtensionType::new(
527 Arc::new(datatypes::extension::json::JsonMetadata::new(
528 datatypes::json::JsonSettings::try_new(
529 vec![datatypes::json::JsonTypeHint {
530 path: vec!["value".into()],
531 data_type: ConcreteDataType::int64_datatype(),
532 inverted_index: false,
533 fulltext_index: None,
534 skipping_index: Some(
535 SkippingIndexOptions::new(2, 0.01, Default::default()).unwrap(),
536 ),
537 }],
538 Some(0),
539 )
540 .unwrap(),
541 )),
542 ));
543 for col in &mut metadata.column_metadatas {
545 col.column_schema.unset_skipping_options().unwrap();
546 }
547 let metadata = Arc::new(metadata);
548 let file_id = FileId::random();
549 let target = IndexTarget::new_json_path(3, vec!["value".into()]).unwrap();
550 let mut creator = BloomFilterIndexer::new(
551 file_id,
552 &metadata,
553 new_intm_mgr(dir.path().to_string_lossy()).await,
554 None,
555 IndexBuildType::Flush,
556 )
557 .unwrap()
558 .unwrap();
559 let values: ArrayRef = Arc::new(Int64Array::from(vec![Some(1), None, Some(1), None]));
560 let root: ArrayRef = Arc::new(StructArray::new(
561 vec![Field::new("value", values.data_type().clone(), true)].into(),
562 vec![values],
563 Some(NullBuffer::from(vec![true, false, true, true])),
564 ));
565 for start in [0, 2] {
566 let array = root.slice(start, 2);
567 if flat {
568 let batch = RecordBatch::try_new(
569 Arc::new(Schema::new(vec![Field::new(
570 "field_u64",
571 array.data_type().clone(),
572 true,
573 )])),
574 vec![array],
575 )
576 .unwrap();
577 creator.update_flat(&batch).await.unwrap();
578 } else {
579 let mut batch = Batch::new(
580 vec![],
581 Arc::new(UInt64Vector::from_vec(vec![0, 1])),
582 Arc::new(UInt64Vector::from_vec(vec![0, 0])),
583 Arc::new(UInt8Vector::from_vec(vec![1, 1])),
584 vec![BatchColumn {
585 column_id: 3,
586 data: Helper::try_into_vector(array).unwrap(),
587 }],
588 )
589 .unwrap();
590 creator.update(&mut batch).await.unwrap();
591 }
592 }
593
594 assert_eq!(creator.indexed_targets().collect::<Vec<_>>(), vec![&target]);
595 let object_store = mock_object_store();
596 let manager = factory.build(
597 object_store.clone(),
598 RegionFilePathFactory::new("json_test".into(), PathType::Bare),
599 );
600 let index_id = RegionIndexId::new(RegionFileId::new(metadata.region_id, file_id), 0);
601 let mut writer = manager.writer(&index_id).await.unwrap();
602 let (rows, bytes) = creator.finish(&mut writer).await.unwrap();
603 assert_eq!(rows, 4);
604 assert!(bytes > 0);
605 writer.finish().await.unwrap();
606 let json_expr = |path: &str, ty| {
607 Expr::ScalarFunction(ScalarFunction::new_udf(
608 Arc::new(create_udf(Arc::new(JsonGetWithType::default()))),
609 vec![col("field_u64"), lit(path), Expr::Literal(ty, None)],
610 ))
611 };
612 let cache = Arc::new(
613 crate::cache::index::bloom_filter_index::BloomFilterIndexCache::new(
614 1024 * 1024,
615 1024 * 1024,
616 1024,
617 ),
618 );
619 for expr in [
620 json_expr("$.value", ScalarValue::Int64(None)).eq(lit(9_i64)),
621 json_expr("$.value", ScalarValue::Int64(None))
622 .in_list(vec![lit(8_i64), lit(9_i64)], false),
623 ] {
624 let applier = BloomFilterIndexApplierBuilder::new(
625 "json_test".into(),
626 PathType::Bare,
627 object_store.clone(),
628 &metadata,
629 factory.clone(),
630 )
631 .with_bloom_filter_index_cache(Some(cache.clone()))
632 .build(&[expr])
633 .unwrap()
634 .unwrap();
635 let predicates = applier.compatible_predicate_for_sst(&metadata).unwrap();
636 assert!(predicates.contains_key(&target));
637 for _ in 0..2 {
639 let result = applier
640 .apply(index_id, None, &predicates, [(4, true)].into_iter(), None)
641 .await
642 .unwrap();
643 assert!(result[0].1.is_empty());
644 }
645 assert!(
646 applier
647 .compatible_predicate_for_sst(&mock_region_metadata())
648 .is_none()
649 );
650 }
651 let applier = BloomFilterIndexApplierBuilder::new(
652 "json_test".into(),
653 PathType::Bare,
654 object_store.clone(),
655 &metadata,
656 factory.clone(),
657 )
658 .build(&[json_expr("$.value", ScalarValue::Int64(None)).eq(lit(1_i64))])
659 .unwrap()
660 .unwrap();
661 let predicates = applier.compatible_predicate_for_sst(&metadata).unwrap();
662 let result = applier
663 .apply(index_id, None, &predicates, [(4, true)].into_iter(), None)
664 .await
665 .unwrap();
666 assert_eq!(result[0].1.len(), 1);
667 assert_eq!(result[0].1[0], 0..4);
668 for expr in [
669 json_expr("$.value", ScalarValue::Float64(None)),
670 json_expr("$.missing", ScalarValue::Int64(None)),
671 ] {
672 assert!(
673 BloomFilterIndexApplierBuilder::new(
674 "json_test".into(),
675 PathType::Bare,
676 object_store.clone(),
677 &metadata,
678 factory.clone()
679 )
680 .build(&[expr.eq(lit(1_i64))])
681 .unwrap()
682 .is_none()
683 );
684 }
685 }
686 }
687
688 #[tokio::test]
689 async fn test_compatible_predicate_for_sst() {
690 let (_d, puffin_manager_factory) =
691 PuffinManagerFactory::new_for_test_async("test_plan_for_sst_basic_").await;
692 let object_store = ObjectStore::new(Memory::default()).unwrap();
693 let table_dir = "table_dir".to_string();
694
695 let predicates = BTreeMap::from_iter([(
696 IndexTarget::ColumnId(1),
697 vec![InListPredicate {
698 list: BTreeSet::from_iter([Bytes::from("foo")]),
699 }],
700 )]);
701 let expected_predicate_target_types = BTreeMap::from_iter([(
702 IndexTarget::ColumnId(1),
703 ConcreteDataType::string_datatype(),
704 )]);
705
706 let applier = BloomFilterIndexApplier::new(
707 table_dir,
708 PathType::Bare,
709 object_store,
710 puffin_manager_factory,
711 predicates,
712 expected_predicate_target_types,
713 );
714 let predicates = applier.compatible_predicate_for_sst(&mock_region_metadata());
715 assert!(predicates.is_some());
716 }
717
718 #[tokio::test]
719 async fn test_compatible_predicate_for_sst_type_mismatch() {
720 let (_d, puffin_manager_factory) =
721 PuffinManagerFactory::new_for_test_async("test_plan_for_sst_type_mismatch_").await;
722 let object_store = ObjectStore::new(Memory::default()).unwrap();
723 let table_dir = "table_dir".to_string();
724
725 let predicates = BTreeMap::from_iter([(
726 IndexTarget::ColumnId(1),
727 vec![InListPredicate {
728 list: BTreeSet::from_iter([Bytes::from("foo")]),
729 }],
730 )]);
731 let expected_predicate_target_types =
732 BTreeMap::from_iter([(IndexTarget::ColumnId(1), ConcreteDataType::int64_datatype())]);
733
734 let applier = BloomFilterIndexApplier::new(
735 table_dir,
736 PathType::Bare,
737 object_store,
738 puffin_manager_factory,
739 predicates,
740 expected_predicate_target_types,
741 );
742 let predicates = applier.compatible_predicate_for_sst(&mock_region_metadata());
743 assert!(predicates.is_none());
744 }
745
746 #[tokio::test]
747 async fn test_compatible_predicate_for_sst_partial_type_mismatch() {
748 let (_d, puffin_manager_factory) =
749 PuffinManagerFactory::new_for_test_async("test_plan_for_sst_partial_mismatch_").await;
750 let object_store = ObjectStore::new(Memory::default()).unwrap();
751 let table_dir = "table_dir".to_string();
752
753 let predicates = BTreeMap::from_iter([
756 (
757 IndexTarget::ColumnId(1),
758 vec![InListPredicate {
759 list: BTreeSet::from_iter([Bytes::from("foo")]),
760 }],
761 ),
762 (
763 IndexTarget::ColumnId(3),
764 vec![InListPredicate {
765 list: BTreeSet::from_iter([Bytes::from("bar")]),
766 }],
767 ),
768 ]);
769 let expected_predicate_target_types = BTreeMap::from_iter([
770 (
771 IndexTarget::ColumnId(1),
772 ConcreteDataType::string_datatype(),
773 ),
774 (IndexTarget::ColumnId(3), ConcreteDataType::int64_datatype()), ]);
776
777 let applier = BloomFilterIndexApplier::new(
778 table_dir,
779 PathType::Bare,
780 object_store,
781 puffin_manager_factory,
782 predicates,
783 expected_predicate_target_types,
784 );
785 let result = applier.compatible_predicate_for_sst(&mock_region_metadata());
786
787 let result = result.expect("expected Some with compatible subset");
789 assert!(
790 result.contains_key(&IndexTarget::ColumnId(1)),
791 "compatible column 1 must be present"
792 );
793 assert!(
794 !result.contains_key(&IndexTarget::ColumnId(3)),
795 "mismatched column 3 must be absent"
796 );
797 assert_eq!(result.len(), 1, "only the compatible predicate must remain");
798 }
799
800 #[allow(clippy::type_complexity)]
801 fn tester(
802 table_dir: String,
803 object_store: ObjectStore,
804 metadata: &RegionMetadata,
805 puffin_manager_factory: PuffinManagerFactory,
806 file_id: RegionIndexId,
807 ) -> impl Fn(&[Expr], Vec<(usize, bool)>) -> BoxFuture<'static, Vec<(usize, Vec<Range<usize>>)>>
808 + use<'_> {
809 move |exprs, row_groups| {
810 let table_dir = table_dir.clone();
811 let object_store: ObjectStore = object_store.clone();
812 let metadata = metadata.clone();
813 let puffin_manager_factory = puffin_manager_factory.clone();
814 let exprs = exprs.to_vec();
815
816 Box::pin(async move {
817 let builder = BloomFilterIndexApplierBuilder::new(
818 table_dir,
819 PathType::Bare,
820 object_store,
821 &metadata,
822 puffin_manager_factory,
823 );
824
825 let applier = builder.build(&exprs).unwrap().unwrap();
826 let predicates = applier
827 .compatible_predicate_for_sst(&Arc::new(metadata.clone()))
828 .unwrap();
829 applier
830 .apply(file_id, None, &predicates, row_groups.into_iter(), None)
831 .await
832 .unwrap()
833 .into_iter()
834 .filter(|(_, ranges)| !ranges.is_empty())
835 .collect()
836 })
837 }
838 }
839
840 #[tokio::test]
841 #[allow(clippy::single_range_in_vec_init)]
842 async fn test_bloom_filter_applier() {
843 let region_metadata = mock_region_metadata();
860 let prefix = "test_bloom_filter_applier_";
861 let (d, factory) = PuffinManagerFactory::new_for_test_async(prefix).await;
862 let object_store = mock_object_store();
863 let intm_mgr = new_intm_mgr(d.path().to_string_lossy()).await;
864 let memory_usage_threshold = Some(1024);
865 let file_id = RegionFileId::new(region_metadata.region_id, FileId::random());
866 let file_id = RegionIndexId::new(file_id, 0);
867 let table_dir = "table_dir".to_string();
868
869 let mut indexer = BloomFilterIndexer::new(
870 file_id.file_id(),
871 ®ion_metadata,
872 intm_mgr,
873 memory_usage_threshold,
874 IndexBuildType::Flush,
875 )
876 .unwrap()
877 .unwrap();
878
879 let mut batch = new_batch("tag1", 0..10);
881 indexer.update(&mut batch).await.unwrap();
882 let mut batch = new_batch("tag2", 10..20);
883 indexer.update(&mut batch).await.unwrap();
884
885 let puffin_manager = factory.build(
886 object_store.clone(),
887 RegionFilePathFactory::new(table_dir.clone(), PathType::Bare),
888 );
889
890 let mut puffin_writer = puffin_manager.writer(&file_id).await.unwrap();
891 indexer.finish(&mut puffin_writer).await.unwrap();
892 puffin_writer.finish().await.unwrap();
893
894 let tester = tester(
895 table_dir.clone(),
896 object_store.clone(),
897 ®ion_metadata,
898 factory.clone(),
899 file_id,
900 );
901
902 let res = tester(
906 &[col("tag_str").eq(lit("tag1"))],
907 vec![(5, true), (5, true), (5, true), (5, true)],
908 )
909 .await;
910 assert_eq!(res, vec![(0, vec![0..5]), (1, vec![0..5])]);
911
912 let res = tester(
916 &[col("tag_str").eq(lit("tag1"))],
917 vec![(5, true), (5, false), (5, true), (5, true)],
918 )
919 .await;
920 assert_eq!(res, vec![(0, vec![0..5])]);
921
922 let res = tester(
927 &[
928 col("tag_str").eq(lit("tag1")),
929 col("field_u64").eq(lit(1u64)),
930 ],
931 vec![(5, true), (5, true), (5, true), (5, true)],
932 )
933 .await;
934 assert_eq!(res, vec![(0, vec![0..4])]);
935
936 let res = tester(
940 &[col("field_u64").in_list(vec![lit(1u64), lit(11u64)], false)],
941 vec![(5, true), (5, true), (5, false), (5, true)],
942 )
943 .await;
944 assert_eq!(res, vec![(0, vec![0..4]), (1, vec![3..5])]);
945 }
946}