perf(index): batch bloom filter searches across row groups (#9361)

* perf(index): search bloom filters of all row groups in one read Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * perf(index): bound the filter bytes of each batched bloom search Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * perf(index): count batch filter bytes exactly as the search reads them Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * chore(index): allow single-range groups in the budget test Signed-off-by: Dennis Zhuang <killme2008@gmail.com> * test(index): cover filters shared across segments in the budget test Signed-off-by: Dennis Zhuang <killme2008@gmail.com> --------- Signed-off-by: Dennis Zhuang <killme2008@gmail.com>
2026-10-03 02:25:35 +00:00 · 2026-09-28 06:22:15 +00:00
parent 43eaea7a9a
commit 0a03e64d8c
3 changed files with 323 additions and 32 deletions
@@ -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<Range<usize>>],
+        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<Range<usize>>],
+        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<Range<usize>>],
+        metrics: Option<&mut BloomFilterReadMetrics>,
+    ) -> Result<()> {
+        let all = groups
+            .iter()
+            .flat_map(|g| g.iter().cloned())
+            .collect::<Vec<_>>();
+        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::<Vec<_>>();
+                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<Vec<u8>>,
+        reads: Arc<std::sync::Mutex<Vec<u64>>>,
+    }
+
+    #[async_trait::async_trait]
+    impl BloomFilterReader for RecordingReader {
+        async fn range_read(
+            &self,
+            offset: u64,
+            size: u32,
+            metrics: Option<&mut BloomFilterReadMetrics>,
+        ) -> Result<bytes::Bytes> {
+            self.inner.range_read(offset, size, metrics).await
+        }
+
+        async fn read_vec(
+            &self,
+            ranges: &[Range<u64>],
+            metrics: Option<&mut BloomFilterReadMetrics>,
+        ) -> Result<Vec<bytes::Bytes>> {
+            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<BloomFilterMeta> {
+            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::<Vec<_>>();
+        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::<Vec<_>>();
+            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::<Vec<_>>();
+        let mut refs = groups.iter_mut().collect::<Vec<_>>();
+        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() {
@@ -422,23 +422,14 @@ impl BloomFilterIndexApplier {
        reader: R,
        predicates: &[InListPredicate],
        output: &mut [(usize, Vec<Range<usize>>)],
-        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::<Vec<_>>();
+        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.
@@ -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::<Vec<_>>();
+        applier
+            .search_groups(
+                &predicates,
+                &mut row_groups,
+                metrics.map(|m| &mut m.bloom_filter_read_metrics),
+            )
+            .await
+            .context(ApplyBloomFilterIndexSnafu)?;

        Ok(true)
    }