nodedb-vector 0.6.0

// SPDX-License-Identifier: Apache-2.0

//! Adaptive filtered search strategy selection.
//!
//! Selects the optimal search strategy based on filter selectivity:
//! - **Low selectivity (<50% filtered out):** pre-filter during HNSW traversal.
//! - **High selectivity (50-95% filtered out):** post-filter with over-fetch.
//! - **Extreme selectivity (>95% filtered out):** brute-force on matching IDs.

use roaring::RoaringBitmap;

use crate::hnsw::{HnswIndex, SearchResult};

/// Filter strategy thresholds.
pub struct FilterThresholds {
    /// Below this selectivity, use pre-filter (default strategy).
    pub high_selectivity: f64,
    /// Above this selectivity, use brute-force on matching IDs.
    pub extreme_selectivity: f64,
}

impl Default for FilterThresholds {
    fn default() -> Self {
        Self {
            high_selectivity: 0.50,
            extreme_selectivity: 0.95,
        }
    }
}

/// Selected search strategy.
#[derive(Debug, Clone, Copy, PartialEq)]
#[non_exhaustive]
pub enum FilterStrategy {
    /// HNSW traversal with bitmap pre-filter (standard).
    PreFilter,
    /// HNSW unfiltered with k×10 over-fetch, then post-filter.
    PostFilter { over_fetch_factor: usize },
    /// Skip HNSW. Brute-force distance on matching IDs only.
    BruteForceMatching,
}

/// Estimate filter selectivity from a Roaring bitmap.
///
/// Returns the fraction of vectors that are filtered OUT (not matching).
pub fn estimate_selectivity(bitmap: &RoaringBitmap, total_vectors: usize) -> f64 {
    if total_vectors == 0 {
        return 0.0;
    }
    let matching = bitmap.len() as usize;
    1.0 - (matching as f64 / total_vectors as f64)
}

/// Select the optimal search strategy based on filter selectivity.
pub fn select_strategy(selectivity: f64, thresholds: &FilterThresholds) -> FilterStrategy {
    if selectivity >= thresholds.extreme_selectivity {
        FilterStrategy::BruteForceMatching
    } else if selectivity >= thresholds.high_selectivity {
        FilterStrategy::PostFilter {
            over_fetch_factor: 10,
        }
    } else {
        FilterStrategy::PreFilter
    }
}

/// Execute adaptive filtered search on an HNSW index.
///
/// A query without the index dimension fails with
/// [`VectorError::DimensionMismatch`](crate::error::VectorError::DimensionMismatch)
/// under every strategy.
pub fn adaptive_search(
    index: &HnswIndex,
    query: &[f32],
    top_k: usize,
    ef: usize,
    bitmap: &RoaringBitmap,
    thresholds: &FilterThresholds,
) -> Result<Vec<SearchResult>, crate::error::VectorError> {
    crate::error::check_dim(index.dim(), query.len())?;
    let total = index.len();
    let selectivity = estimate_selectivity(bitmap, total);
    let strategy = select_strategy(selectivity, thresholds);

    match strategy {
        FilterStrategy::PreFilter => index.search_filtered(query, top_k, ef, bitmap),
        FilterStrategy::PostFilter { over_fetch_factor } => {
            let fetch_k = top_k * over_fetch_factor;
            let results = index.search(query, fetch_k, ef.max(fetch_k))?;
            let mut filtered: Vec<SearchResult> = results
                .into_iter()
                .filter(|r| bitmap.contains(r.id))
                .collect();
            filtered.truncate(top_k);
            Ok(filtered)
        }
        FilterStrategy::BruteForceMatching => {
            let metric = index.params().metric;
            let probe = index.probe(query);
            // A node without a vector has no distance and is not a hit.
            let mut results: Vec<SearchResult> = bitmap
                .iter()
                .filter(|&id| !index.is_deleted(id))
                .filter_map(|id| {
                    let distance = index.dist_to_node(&probe, id, metric)?;
                    Some(SearchResult { id, distance })
                })
                .collect();

            if results.len() > top_k {
                results.select_nth_unstable_by(top_k, |a, b| {
                    a.distance
                        .partial_cmp(&b.distance)
                        .unwrap_or(std::cmp::Ordering::Equal)
                });
                results.truncate(top_k);
            }
            results.sort_by(|a, b| {
                a.distance
                    .partial_cmp(&b.distance)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
            Ok(results)
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::distance::DistanceMetric;
    use crate::hnsw::{HnswIndex, HnswParams};

    fn build_test_index() -> HnswIndex {
        let mut idx = HnswIndex::with_seed(
            3,
            HnswParams {
                m: 8,
                m0: 16,
                ef_construction: 50,
                metric: DistanceMetric::L2,
                dtype: nodedb_types::vector_dtype::VectorStorageDtype::F32,
            },
            42,
        );
        for i in 0..1000 {
            idx.insert(vec![i as f32, 0.0, 0.0]).unwrap();
        }
        idx
    }

    #[test]
    fn low_selectivity_uses_prefilter() {
        let thresholds = FilterThresholds::default();
        let strategy = select_strategy(0.2, &thresholds);
        assert_eq!(strategy, FilterStrategy::PreFilter);
    }

    #[test]
    fn high_selectivity_uses_postfilter() {
        let thresholds = FilterThresholds::default();
        let strategy = select_strategy(0.8, &thresholds);
        assert!(matches!(strategy, FilterStrategy::PostFilter { .. }));
    }

    #[test]
    fn extreme_selectivity_uses_bruteforce() {
        let thresholds = FilterThresholds::default();
        let strategy = select_strategy(0.99, &thresholds);
        assert_eq!(strategy, FilterStrategy::BruteForceMatching);
    }

    #[test]
    fn adaptive_search_extreme_filter() {
        let idx = build_test_index();
        let thresholds = FilterThresholds::default();

        let mut bitmap = RoaringBitmap::new();
        for i in 500..510 {
            bitmap.insert(i);
        }

        let results =
            adaptive_search(&idx, &[505.0, 0.0, 0.0], 3, 64, &bitmap, &thresholds).unwrap();
        assert_eq!(results.len(), 3);
        for r in &results {
            assert!(bitmap.contains(r.id), "got filtered-out id {}", r.id);
        }
        assert_eq!(results[0].id, 505);
    }

    #[test]
    fn adaptive_search_low_filter() {
        let idx = build_test_index();
        let thresholds = FilterThresholds::default();

        let mut bitmap = RoaringBitmap::new();
        for i in 0..800 {
            bitmap.insert(i);
        }

        let results =
            adaptive_search(&idx, &[100.0, 0.0, 0.0], 5, 64, &bitmap, &thresholds).unwrap();
        assert_eq!(results.len(), 5);
        for r in &results {
            assert!(bitmap.contains(r.id));
        }
    }

    #[test]
    fn selectivity_estimation() {
        let mut bitmap = RoaringBitmap::new();
        for i in 0..100 {
            bitmap.insert(i);
        }
        let sel = estimate_selectivity(&bitmap, 1000);
        assert!((sel - 0.9).abs() < 0.01);
    }

    /// Every strategy refuses a query of the wrong dimension, including the
    /// brute-force path that never touches the graph.
    #[test]
    fn wrong_dimension_query_is_a_typed_error() {
        let idx = build_test_index();
        let thresholds = FilterThresholds::default();
        for n in [5u32, 500, 1000] {
            let bitmap: RoaringBitmap = (0..n).collect();
            let result = adaptive_search(&idx, &[1.0, 0.0], 3, 64, &bitmap, &thresholds);
            assert!(
                matches!(
                    result,
                    Err(crate::error::VectorError::DimensionMismatch {
                        expected: 3,
                        got: 2
                    })
                ),
                "{result:?}"
            );
        }
    }

    const SEED: u64 = 0x9E37_79B9_7F4A_7C15;

    /// Thresholds that force pre-filter, post-filter, and brute force in turn.
    const FORCED: [FilterThresholds; 3] = [
        FilterThresholds {
            high_selectivity: 2.0,
            extreme_selectivity: 2.0,
        },
        FilterThresholds {
            high_selectivity: -1.0,
            extreme_selectivity: 2.0,
        },
        BRUTE_FORCE,
    ];

    const BRUTE_FORCE: FilterThresholds = FilterThresholds {
        high_selectivity: -1.0,
        extreme_selectivity: -1.0,
    };

    fn every_tenth(n: u32) -> RoaringBitmap {
        (0..n).step_by(10).collect()
    }

    /// Brute force over a graph-only restore reads node vectors through the
    /// backing and ranks exactly like the source index.
    #[cfg(not(target_arch = "wasm32"))]
    #[test]
    fn brute_force_after_graph_only_restore_matches_the_source_index() {
        use crate::test_support::{graph_only_restore_with_backing, seeded_index, seeded_vectors};
        use nodedb_types::vector_dtype::VectorStorageDtype;

        let vectors = seeded_vectors(300, 8, SEED);
        let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32);
        let restored = graph_only_restore_with_backing(&src, &vectors);
        let bitmap = every_tenth(300);

        let expected = adaptive_search(&src, &vectors[7], 5, 64, &bitmap, &BRUTE_FORCE).unwrap();
        let got = adaptive_search(&restored, &vectors[7], 5, 64, &bitmap, &BRUTE_FORCE).unwrap();
        assert_eq!(expected.len(), 5);
        let pairs = |r: &[SearchResult]| -> Vec<(u32, f32)> {
            r.iter().map(|h| (h.id, h.distance)).collect()
        };
        assert_eq!(pairs(got.as_slice()), pairs(expected.as_slice()));
    }

    /// A graph-only restore searched before `with_backing` has no vectors. No
    /// strategy may return its nodes, at infinite distance or otherwise.
    #[test]
    fn no_strategy_returns_a_node_without_a_vector() {
        use crate::test_support::{graph_only_restore, seeded_index, seeded_vectors};
        use nodedb_types::vector_dtype::VectorStorageDtype;

        let vectors = seeded_vectors(300, 8, SEED);
        let src = seeded_index(&vectors, DistanceMetric::L2, VectorStorageDtype::F32);
        let idx = graph_only_restore(&src);
        let bitmap = every_tenth(300);
        for thresholds in &FORCED {
            let results = adaptive_search(&idx, &vectors[7], 5, 64, &bitmap, thresholds).unwrap();
            assert!(results.is_empty(), "got {results:?}");
        }
    }

    /// Brute force on an F16/BF16 index scores the encoded vectors under the
    /// build metric.
    #[test]
    fn brute_force_on_narrow_dtypes_matches_decoded_distances() {
        use crate::distance::distance;
        use crate::test_support::{seeded_index, seeded_vectors};
        use nodedb_types::vector_dtype::VectorStorageDtype;

        let vectors = seeded_vectors(300, 8, SEED);
        let bitmap = every_tenth(300);
        for dtype in [VectorStorageDtype::F16, VectorStorageDtype::BF16] {
            let idx = seeded_index(&vectors, DistanceMetric::Cosine, dtype);
            let encoded_query = crate::dtype::cast_from_f32(&vectors[7], dtype);
            let query = crate::dtype::cast_to_f32(&encoded_query, dtype, 8).unwrap();
            let mut reference: Vec<(u32, f32)> = bitmap
                .iter()
                .map(|id| {
                    let node = idx.materialize_vector(id).unwrap();
                    (id, distance(&query, &node, DistanceMetric::Cosine))
                })
                .collect();
            reference.sort_by(|a, b| a.1.total_cmp(&b.1));

            let got = adaptive_search(&idx, &vectors[7], 5, 64, &bitmap, &BRUTE_FORCE).unwrap();
            assert_eq!(got.len(), 5, "dtype={dtype:?}");
            for (hit, (id, d)) in got.iter().zip(&reference) {
                assert_eq!(hit.id, *id, "dtype={dtype:?}: got {got:?}");
                assert!(
                    (hit.distance - d).abs() < 1e-4,
                    "dtype={dtype:?}: {hit:?} vs {d}"
                );
            }
        }
    }
}