haystack_core/graph/

query_planner.rs

// Query planner — bitmap-accelerated filter evaluation.
//
// Two-phase approach:
// 1. **Bitmap phase**: For Has/Missing nodes, combine tag bitmaps with
//    set operations (AND/OR/NOT) to produce a candidate bitset.
// 2. **Scan phase**: Iterate set bits from candidates, evaluate the
//    full filter AST on each candidate entity.
//
// Returns `None` when bitmap optimization is not applicable (e.g.
// comparison nodes), causing the caller to fall back to a full scan.

use roaring::RoaringBitmap;

use crate::filter::{CmpOp, FilterNode};
use crate::kinds::Kind;

use super::adjacency::RefAdjacency;
use super::bitmap::TagBitmapIndex;
use super::value_index::ValueIndex;

/// Helper: convert a list of entity IDs into a roaring bitmap.
fn ids_to_bitmap(entity_ids: &[usize]) -> RoaringBitmap {
    entity_ids.iter().map(|&id| id as u32).collect()
}

/// Compute a bitmap of candidate entity ids that *may* match the given
/// filter node, using tag presence bitmaps only.
///
/// Returns `None` if the filter contains terms that cannot be resolved
/// via bitmap (e.g. comparison operators), meaning a full scan is needed.
pub fn bitmap_candidates(
    node: &FilterNode,
    tag_index: &TagBitmapIndex,
    max_id: usize,
) -> Option<RoaringBitmap> {
    match node {
        FilterNode::Has(path) => {
            if !path.is_single() {
                return None; // multi-segment paths need entity resolution
            }
            tag_index.has_tag(path.first()).cloned()
        }

        FilterNode::Missing(path) => {
            if !path.is_single() {
                return None;
            }
            match tag_index.has_tag(path.first()) {
                Some(bm) => Some(TagBitmapIndex::negate(bm, max_id)),
                None => {
                    // Tag never seen: all entities match "missing".
                    if max_id == 0 {
                        Some(RoaringBitmap::new())
                    } else {
                        Some(TagBitmapIndex::negate(&RoaringBitmap::new(), max_id))
                    }
                }
            }
        }

        FilterNode::And(left, right) => {
            let l = bitmap_candidates(left, tag_index, max_id);
            let r = bitmap_candidates(right, tag_index, max_id);
            match (l, r) {
                (Some(lb), Some(rb)) => {
                    if lb.is_empty() || rb.is_empty() {
                        return Some(RoaringBitmap::new());
                    }
                    Some(TagBitmapIndex::intersect(&[&lb, &rb]))
                }
                // If only one side is optimizable, use it as the candidate
                // set; the other side will be checked during the scan phase.
                (Some(bm), None) | (None, Some(bm)) => Some(bm),
                (None, None) => None,
            }
        }

        FilterNode::Or(left, right) => {
            let l = bitmap_candidates(left, tag_index, max_id);
            let r = bitmap_candidates(right, tag_index, max_id);
            match (l, r) {
                (Some(lb), Some(rb)) => Some(TagBitmapIndex::union(&[&lb, &rb])),
                // If either side is not optimizable, we cannot narrow.
                _ => None,
            }
        }

        FilterNode::Cmp { path, .. } => {
            // Comparison nodes cannot be resolved via bitmap, but we can
            // still use the tag existence bitmap to prune candidates.
            if path.is_single() {
                tag_index.has_tag(path.first()).cloned()
            } else {
                None
            }
        }

        FilterNode::SpecMatch(_) => None,
    }
}

/// Enhanced candidate computation that uses B-Tree value indexes and
/// ref adjacency for Cmp nodes in addition to tag bitmap indexes.
pub fn bitmap_candidates_with_values(
    node: &FilterNode,
    tag_index: &TagBitmapIndex,
    value_index: &ValueIndex,
    adjacency: &RefAdjacency,
    max_id: usize,
) -> Option<RoaringBitmap> {
    match node {
        // Ref equality: use adjacency reverse map for O(1) lookup.
        FilterNode::Cmp {
            path,
            op: CmpOp::Eq,
            val: Kind::Ref(r),
        } if path.is_single() => {
            let field = path.first();
            let entity_ids = adjacency.sources_for(field, &r.val);
            Some(ids_to_bitmap(&entity_ids))
        }

        FilterNode::Cmp { path, op, val }
            if path.is_single() && value_index.has_index(path.first()) =>
        {
            let field = path.first();
            let entity_ids = match op {
                CmpOp::Eq => value_index.eq_lookup(field, val),
                CmpOp::Ne => value_index.ne_lookup(field, val),
                CmpOp::Lt => value_index.lt_lookup(field, val),
                CmpOp::Le => value_index.le_lookup(field, val),
                CmpOp::Gt => value_index.gt_lookup(field, val),
                CmpOp::Ge => value_index.ge_lookup(field, val),
            };

            Some(ids_to_bitmap(&entity_ids))
        }

        FilterNode::And(left, right) => {
            let l = bitmap_candidates_with_values(left, tag_index, value_index, adjacency, max_id);
            let r = bitmap_candidates_with_values(right, tag_index, value_index, adjacency, max_id);
            match (l, r) {
                (Some(lb), Some(rb)) => {
                    if lb.is_empty() || rb.is_empty() {
                        return Some(RoaringBitmap::new());
                    }
                    Some(TagBitmapIndex::intersect(&[&lb, &rb]))
                }
                (Some(bm), None) | (None, Some(bm)) => Some(bm),
                (None, None) => None,
            }
        }

        FilterNode::Or(left, right) => {
            let l = bitmap_candidates_with_values(left, tag_index, value_index, adjacency, max_id);
            let r = bitmap_candidates_with_values(right, tag_index, value_index, adjacency, max_id);
            match (l, r) {
                (Some(lb), Some(rb)) => Some(TagBitmapIndex::union(&[&lb, &rb])),
                _ => None,
            }
        }

        // Delegate to base bitmap_candidates for all other node types.
        _ => bitmap_candidates(node, tag_index, max_id),
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use crate::filter::parse_filter;

    fn build_test_index() -> (TagBitmapIndex, usize) {
        let mut idx = TagBitmapIndex::new();
        // entity 0: site
        idx.add(0, &["site".into(), "dis".into()]);
        // entity 1: equip
        idx.add(1, &["equip".into(), "dis".into()]);
        // entity 2: point, sensor
        idx.add(2, &["point".into(), "sensor".into(), "dis".into()]);
        // entity 3: site, geoCity
        idx.add(3, &["site".into(), "geoCity".into(), "dis".into()]);
        (idx, 4) // max_id = 4
    }

    #[test]
    fn bitmap_candidates_for_has() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("site").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id).unwrap();
        let bits: Vec<usize> = TagBitmapIndex::iter_set_bits(&bm).collect();
        assert_eq!(bits, vec![0, 3]);
    }

    #[test]
    fn bitmap_candidates_for_missing() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("not site").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id).unwrap();
        let bits: Vec<usize> = TagBitmapIndex::iter_set_bits(&bm).collect();
        assert_eq!(bits, vec![1, 2]);
    }

    #[test]
    fn bitmap_candidates_for_and() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("site and geoCity").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id).unwrap();
        let bits: Vec<usize> = TagBitmapIndex::iter_set_bits(&bm).collect();
        assert_eq!(bits, vec![3]);
    }

    #[test]
    fn bitmap_candidates_for_or() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("site or equip").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id).unwrap();
        let bits: Vec<usize> = TagBitmapIndex::iter_set_bits(&bm).collect();
        assert_eq!(bits, vec![0, 1, 3]);
    }

    #[test]
    fn bitmap_candidates_for_comparison_uses_tag_existence() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("geoCity == \"Richmond\"").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id);
        assert!(bm.is_some());
        let bits: Vec<usize> = TagBitmapIndex::iter_set_bits(&bm.unwrap()).collect();
        assert_eq!(bits, vec![3]);
    }

    #[test]
    fn bitmap_candidates_returns_none_for_multi_segment_path() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("siteRef->area").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id);
        assert!(bm.is_none());
    }

    #[test]
    fn bitmap_candidates_returns_none_for_spec_match() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("ph::Site").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id);
        assert!(bm.is_none());
    }

    #[test]
    fn and_with_one_side_optimizable() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("site and dis == \"hello\"").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id);
        assert!(bm.is_some());
    }

    #[test]
    fn or_with_one_side_not_optimizable() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("site or siteRef->area").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id);
        assert!(bm.is_none());
    }

    #[test]
    fn empty_index() {
        let idx = TagBitmapIndex::new();
        let ast = parse_filter("site").unwrap();
        let bm = bitmap_candidates(&ast, &idx, 0);
        assert!(bm.is_none());
    }

    #[test]
    fn missing_unknown_tag_returns_all() {
        let (idx, max_id) = build_test_index();
        let ast = parse_filter("not unknownTag").unwrap();
        let bm = bitmap_candidates(&ast, &idx, max_id).unwrap();
        let bits: Vec<usize> = TagBitmapIndex::iter_set_bits(&bm).collect();
        assert_eq!(bits, vec![0, 1, 2, 3]);
    }
}