diff --git a/src/index/src/bloom_filter/applier.rs b/src/index/src/bloom_filter/applier.rs index db219c9e611..ebd44cf84e6 100644 --- a/src/index/src/bloom_filter/applier.rs +++ b/src/index/src/bloom_filter/applier.rs @@ -23,6 +23,11 @@ use crate::Bytes; use crate::bloom_filter::error::Result; use crate::bloom_filter::reader::{BloomFilterReadMetrics, BloomFilterReader}; +/// Filter bytes one batch of [`BloomFilterApplier::search_groups`] reads. A single row +/// group larger than this is still read as one batch, so this is not a memory limit; a +/// batch holds both its raw bytes and the decoded filters. +const MAX_BATCH_FILTER_BYTES: u64 = 8 * 1024 * 1024; + /// `InListPredicate` contains a list of acceptable values. A value needs to match at least /// one of the elements (logical OR semantic) for the predicate to be satisfied. #[derive(Debug, Clone, PartialEq, Eq, Hash)] @@ -43,6 +48,93 @@ impl BloomFilterApplier { Ok(Self { reader, meta }) } + /// Runs [`Self::search`] over several groups of ranges (e.g. row groups) and keeps in + /// each group only its matching ranges. + /// + /// Consecutive groups are searched together, one read per batch, as long as the + /// filters of a batch stay within [`MAX_BATCH_FILTER_BYTES`]. This saves a round trip + /// per group on object storage without holding the filters of a whole file at once. + /// Groups must be ordered and their ranges sorted and disjoint, as for `search`. + pub async fn search_groups( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result<()> { + self.search_groups_in_batches(predicates, groups, metrics, MAX_BATCH_FILTER_BYTES) + .await + } + + async fn search_groups_in_batches( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + mut metrics: Option<&mut BloomFilterReadMetrics>, + max_batch_bytes: u64, + ) -> Result<()> { + let mut start = 0; + while start < groups.len() { + // Mirrors `load_bloom_filters`: segments map to filter locations in order and + // only consecutive equal locations share a read, so this is the exact number of + // bytes the batch requests. + let mut last_loc = None; + let mut batch_bytes = 0; + let mut end = start; + while end < groups.len() { + let mut group_last = last_loc; + let mut bytes = 0; + for seg in self.row_ranges_to_segments(groups[end]) { + let loc = self.meta.segment_loc_indices[seg]; + if group_last != Some(loc) { + bytes += self.meta.bloom_filter_locs[loc as usize].size; + group_last = Some(loc); + } + } + // A group whose filters alone exceed the budget still forms its own batch. + if end > start && batch_bytes + bytes > max_batch_bytes { + break; + } + last_loc = group_last; + batch_bytes += bytes; + end += 1; + } + self.search_batch(predicates, &mut groups[start..end], metrics.as_deref_mut()) + .await?; + start = end; + } + Ok(()) + } + + /// Searches `groups` with one `search`, i.e. one read for all their filters. + async fn search_batch( + &mut self, + predicates: &[InListPredicate], + groups: &mut [&mut Vec>], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result<()> { + let all = groups + .iter() + .flat_map(|g| g.iter().cloned()) + .collect::>(); + if all.is_empty() { + return Ok(()); + } + // Each matched range lies within one input range, so it belongs to one group. + let mut matched = self + .search(predicates, &all, metrics) + .await? + .into_iter() + .peekable(); + for group in groups.iter_mut() { + for range in std::mem::take(*group) { + while let Some(m) = matched.next_if(|m| m.start < range.end) { + group.push(m); + } + } + } + Ok(()) + } + /// Searches ranges of rows that match all the given predicates in the search ranges. /// Each predicate represents an OR condition of probes, and all predicates must match (AND semantics). /// The logic is: (probe1 OR probe2 OR ...) AND (probe3 OR probe4 OR ...) @@ -217,6 +309,222 @@ mod tests { use crate::bloom_filter::reader::BloomFilterReaderImpl; use crate::external_provider::MockExternalTempFileProvider; + #[tokio::test] + async fn test_search_groups_matches_per_group_search() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + // Row i holds "v{i / 3}", so values straddle segment boundaries. + for i in 0..40 { + creator + .push_row_elems([format!("v{}", i / 3).into_bytes()]) + .await + .unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let bytes = writer.into_inner(); + + // Row groups of 10 rows, some already narrowed by other predicates. + let groups = vec![ + vec![0..10], + vec![10..13, 15..20], + vec![], + vec![30..33, 37..40], + ]; + for values in [ + vec!["v1"], + vec!["v4", "v5"], + vec!["v3", "v10", "v12"], + vec!["x"], + ] { + let predicates = vec![InListPredicate { + list: values.iter().map(|v| v.as_bytes().to_vec()).collect(), + }]; + let mut applier = + BloomFilterApplier::new(Box::new(BloomFilterReaderImpl::new(bytes.clone()))) + .await + .unwrap(); + let filter_size = applier.meta.bloom_filter_locs[0].size; + let mut expected = Vec::new(); + for group in &groups { + expected.push(if group.is_empty() { + vec![] + } else { + applier.search(&predicates, group, None).await.unwrap() + }); + } + // Budgets of zero (a batch per group), one filter and unlimited (one batch). + for budget in [0, filter_size, u64::MAX] { + let mut actual = groups.clone(); + let mut refs = actual.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, budget) + .await + .unwrap(); + assert_eq!(actual, expected, "values: {values:?}, budget: {budget}"); + } + } + } + + /// Records the bytes of every `read_vec`, i.e. of every batch. + struct RecordingReader { + inner: BloomFilterReaderImpl>, + reads: Arc>>, + } + + #[async_trait::async_trait] + impl BloomFilterReader for RecordingReader { + async fn range_read( + &self, + offset: u64, + size: u32, + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result { + self.inner.range_read(offset, size, metrics).await + } + + async fn read_vec( + &self, + ranges: &[Range], + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result> { + let bytes = ranges.iter().map(|r| r.end - r.start).sum(); + self.reads.lock().unwrap().push(bytes); + self.inner.read_vec(ranges, metrics).await + } + + async fn metadata( + &self, + metrics: Option<&mut BloomFilterReadMetrics>, + ) -> Result { + self.inner.metadata(metrics).await + } + } + + #[tokio::test] + #[allow(clippy::single_range_in_vec_init)] + async fn test_search_groups_respects_budget() { + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + // Distinct values everywhere: 10 segments with filters of the same size. + for i in 0..40 { + creator + .push_row_elems([format!("v{i}").into_bytes()]) + .await + .unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let bytes = writer.into_inner(); + let predicates = vec![InListPredicate { + list: BTreeSet::from([b"v1".to_vec()]), + }]; + + let reads = Arc::new(std::sync::Mutex::new(Vec::new())); + let reader = RecordingReader { + inner: BloomFilterReaderImpl::new(bytes), + reads: reads.clone(), + }; + let mut applier = BloomFilterApplier::new(Box::new(reader)).await.unwrap(); + let filter_size = applier.meta.bloom_filter_locs[0].size; + + // Row groups of 6 rows: most share a boundary segment with their neighbor, which + // a batch reads only once. + let groups = (0..40) + .step_by(6) + .map(|s| vec![s..(s + 6).min(40)]) + .collect::>(); + for (budget, expected_reads) in [ + (u64::MAX, vec![10 * filter_size]), + // Batches close before exceeding the budget; shared boundary filters count once. + ( + 5 * filter_size, + vec![5 * filter_size, 5 * filter_size, filter_size], + ), + // Smaller than any row group: one row group per batch, still above the budget. + ( + filter_size, + vec![ + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + 2 * filter_size, + filter_size, + ], + ), + ] { + reads.lock().unwrap().clear(); + let mut actual = groups.clone(); + let mut refs = actual.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, budget) + .await + .unwrap(); + assert_eq!(*reads.lock().unwrap(), expected_reads, "budget: {budget}"); + } + + // Segments 0..3 hold the same value and share one filter, which a batch counts + // and reads once even across row groups. + let mut creator = BloomFilterCreator::new( + 4, + 0.01, + Arc::new(MockExternalTempFileProvider::new()), + Arc::new(AtomicUsize::new(0)), + None, + ); + for i in 0..24 { + let value = if i < 12 { + "a".to_string() + } else { + format!("v{i}") + }; + creator.push_row_elems([value.into_bytes()]).await.unwrap(); + } + let mut writer = Cursor::new(Vec::new()); + creator.finish(&mut writer).await.unwrap(); + let reader = RecordingReader { + inner: BloomFilterReaderImpl::new(writer.into_inner()), + reads: reads.clone(), + }; + let mut applier = BloomFilterApplier::new(Box::new(reader)).await.unwrap(); + assert_eq!(applier.meta.bloom_filter_locs.len(), 4); + let filter_size = applier.meta.bloom_filter_locs[0].size; + assert!( + applier + .meta + .bloom_filter_locs + .iter() + .all(|l| l.size == filter_size) + ); + reads.lock().unwrap().clear(); + let mut groups = (0..24) + .step_by(6) + .map(|s| vec![s..s + 6]) + .collect::>(); + let mut refs = groups.iter_mut().collect::>(); + applier + .search_groups_in_batches(&predicates, &mut refs, None, 2 * filter_size) + .await + .unwrap(); + // Row groups 0 and 1 need only the shared filter; 2 and 3 need two each. + assert_eq!( + *reads.lock().unwrap(), + vec![filter_size, 2 * filter_size, 2 * filter_size] + ); + } + #[tokio::test] #[allow(clippy::single_range_in_vec_init)] async fn test_appliter() { diff --git a/src/mito2/src/sst/index/bloom_filter/applier.rs b/src/mito2/src/sst/index/bloom_filter/applier.rs index e2d25e4bf82..1faae466c89 100644 --- a/src/mito2/src/sst/index/bloom_filter/applier.rs +++ b/src/mito2/src/sst/index/bloom_filter/applier.rs @@ -422,23 +422,14 @@ impl BloomFilterIndexApplier { reader: R, predicates: &[InListPredicate], output: &mut [(usize, Vec>)], - mut metrics: Option<&mut BloomFilterIndexApplyMetrics>, + metrics: Option<&mut BloomFilterIndexApplyMetrics>, ) -> std::result::Result<(), index::bloom_filter::error::Error> { let mut applier = BloomFilterApplier::new(Box::new(reader)).await?; - - for (_, row_group_output) in output.iter_mut() { - // All rows are filtered out, skip the search - if row_group_output.is_empty() { - continue; - } - - let read_metrics = metrics.as_deref_mut().map(|m| &mut m.read_metrics); - *row_group_output = applier - .search(predicates, row_group_output, read_metrics) - .await?; - } - - Ok(()) + let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::>(); + let read_metrics = metrics.map(|m| &mut m.read_metrics); + applier + .search_groups(predicates, &mut row_groups, read_metrics) + .await } /// Returns compatible bloom filter predicates with the given SST metadata. diff --git a/src/mito2/src/sst/index/fulltext_index/applier.rs b/src/mito2/src/sst/index/fulltext_index/applier.rs index 5ccf4e57a4b..4efb2a5ccb9 100644 --- a/src/mito2/src/sst/index/fulltext_index/applier.rs +++ b/src/mito2/src/sst/index/fulltext_index/applier.rs @@ -464,23 +464,15 @@ impl FulltextIndexApplier { let mut applier = BloomFilterApplier::new(reader) .await .context(ApplyBloomFilterIndexSnafu)?; - for (_, row_group_output) in output.iter_mut() { - // All rows are filtered out, skip the search - if row_group_output.is_empty() { - continue; - } - - *row_group_output = applier - .search( - &predicates, - row_group_output, - metrics - .as_deref_mut() - .map(|m| &mut m.bloom_filter_read_metrics), - ) - .await - .context(ApplyBloomFilterIndexSnafu)?; - } + let mut row_groups = output.iter_mut().map(|(_, r)| r).collect::>(); + applier + .search_groups( + &predicates, + &mut row_groups, + metrics.map(|m| &mut m.bloom_filter_read_metrics), + ) + .await + .context(ApplyBloomFilterIndexSnafu)?; Ok(true) }