Skip to main content

mito2/sst/index/bloom_filter/
applier.rs

1// Copyright 2023 Greptime Team
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15mod 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/// Metrics for tracking bloom filter index apply operations.
55#[derive(Default, Clone)]
56pub struct BloomFilterIndexApplyMetrics {
57    /// Total time spent applying the index.
58    pub apply_elapsed: std::time::Duration,
59    /// Number of blob cache misses.
60    pub blob_cache_miss: usize,
61    /// Total size of blobs read (in bytes).
62    pub blob_read_bytes: u64,
63    /// Metrics for bloom filter read operations.
64    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    /// Returns true if the metrics are empty (contain no meaningful data).
97    pub fn is_empty(&self) -> bool {
98        self.apply_elapsed.is_zero()
99    }
100
101    /// Merges another metrics into this one.
102    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
112/// `BloomFilterIndexApplier` applies bloom filter predicates to the SST file.
113pub struct BloomFilterIndexApplier {
114    /// Directory of the table.
115    table_dir: String,
116
117    /// Path type for generating file paths.
118    path_type: PathType,
119
120    /// Object store to read the index file.
121    object_store: ObjectStore,
122
123    /// File cache to read the index file.
124    file_cache: Option<FileCacheRef>,
125
126    /// Factory to create puffin manager.
127    puffin_manager_factory: PuffinManagerFactory,
128
129    /// Cache for puffin metadata.
130    puffin_metadata_cache: Option<PuffinMetadataCacheRef>,
131
132    /// Cache for bloom filter index.
133    bloom_filter_index_cache: Option<BloomFilterIndexCacheRef>,
134
135    /// Bloom filter predicates.
136    /// For each column, the value will be retained only if it contains __all__ predicates.
137    default_predicates: Arc<BTreeMap<IndexTarget, Vec<InListPredicate>>>,
138
139    /// Expected predicate target types from the latest region metadata.
140    expected_predicate_target_types: BTreeMap<IndexTarget, ConcreteDataType>,
141}
142
143impl BloomFilterIndexApplier {
144    /// Creates a new `BloomFilterIndexApplier`.
145    ///
146    /// For each column, the value will be retained only if it contains __all__ predicates.
147    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    /// Applies bloom filter predicates to the provided SST file and returns a
191    /// list of row group ranges that match the predicates.
192    ///
193    /// The `row_groups` iterator provides the row group lengths and whether to search in the row group.
194    ///
195    /// Row group id existing in the returned result means that the row group is searched.
196    /// Empty ranges means that the row group is searched but no rows are found.
197    ///
198    ///
199    /// # Arguments
200    /// * `file_id` - The region file ID to apply predicates to
201    /// * `file_size_hint` - Optional hint for file size to avoid extra metadata reads
202    /// * `row_groups` - Iterator of row group lengths and whether to search in the row group
203    /// * `metrics` - Optional mutable reference to collect metrics on demand
204    #[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        // Calculates row groups' ranges based on start of the file.
219        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        // Initializes output with input ranges, but ranges are based on start of the file not the row group,
230        // so we need to adjust them later.
231        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            // Create appropriate reader based on whether we have caching enabled
246            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        // adjust ranges to be based on row group
272        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        // Record elapsed time to histogram and collect metrics if requested
281        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    /// Creates a blob reader from the cached or remote index file.
294    ///
295    /// Returns `None` if the target does not have an index.
296    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                    // Blob not found means no index for this target.
314                    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                    // Blob not found means no index for this target.
324                    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    /// Creates a blob reader from the cached index file
338    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    // TODO(ruihang): use the same util with the code in creator
378    fn target_blob_name(target: &IndexTarget) -> String {
379        let target_key = target.encode();
380        format!("{INDEX_BLOB_TYPE}-{target_key}")
381    }
382
383    /// Creates a blob reader from the remote index file
384    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 background download if file cache and file size are available
393        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    /// Returns compatible bloom filter predicates with the given SST metadata.
437    ///
438    /// Returns `None` when no compatible predicate remains for this SST.
439    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            // Other ordinary indexed columns are absent in this synthetic batch.
544            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                // Run twice to cover cold reads and cached index reads.
638                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        // Column 1 (tag_str): expected string — matches SST (compatible).
754        // Column 3 (field_u64): expected int64 — SST has uint64 (mismatched).
755        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()), // intentional mismatch
775        ]);
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        // The subset containing only the compatible column must be returned.
788        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        // tag_str:
844        //   - type: string
845        //   - index: bloom filter
846        //   - granularity: 2
847        //   - column_id: 1
848        //
849        // ts:
850        //   - type: timestamp
851        //   - index: time index
852        //   - column_id: 2
853        //
854        // field_u64:
855        //   - type: uint64
856        //   - index: bloom filter
857        //   - granularity: 4
858        //   - column_id: 3
859        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            &region_metadata,
872            intm_mgr,
873            memory_usage_threshold,
874            IndexBuildType::Flush,
875        )
876        .unwrap()
877        .unwrap();
878
879        // push 20 rows
880        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            &region_metadata,
898            factory.clone(),
899            file_id,
900        );
901
902        // rows        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
903        // row group: | o  row group |  o row group | o  row group     |  o row group     |
904        // tag_str:   |      o pred                 |   x pred                            |
905        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        // rows        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
913        // row group: | o  row group |  x row group | o  row group     |  o row group     |
914        // tag_str:   |      o pred                 |   x pred                            |
915        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        // rows        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
923        // row group: | o  row group |  o row group | o  row group     |  o row group     |
924        // tag_str:   |      o pred                 |   x pred                            |
925        // field_u64: | o pred   | x pred    |  x pred     |  x pred       | x pred       |
926        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        // rows        0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19
937        // row group: | o  row group |  o row group | x  row group     |  o row group     |
938        // field_u64: | o pred   | x pred    |  o pred     |  x pred       | x pred       |
939        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}