Skip to main content

nodedb_graph/csr/index/
scoped.rs

1// SPDX-License-Identifier: Apache-2.0
2
3//! Collection-scoped read paths.
4//!
5//! A CSR partition holds every collection's edges under one shared node space
6//! (nodes, node-labels and surrogates are shared across collections; only edges
7//! carry a collection id). These methods filter traversal to a single
8//! collection so a `MATCH ... IN '<collection>'` / GraphRAG read never crosses
9//! a collection boundary.
10
11use super::types::{CsrIndex, Direction};
12
13impl CsrIndex {
14    /// Collection-scoped neighbor lookup by string name.
15    ///
16    /// Like [`Self::neighbors`], with the same `label_filter` set semantics,
17    /// but only traverses edges tagged with
18    /// `collection`. A collection with no edges in this partition yields an
19    /// empty result (never falls back to the merged, cross-collection view).
20    /// Used by collection-scoped GraphRAG BFS so expansion never crosses a
21    /// collection boundary.
22    pub fn neighbors_in_collection(
23        &self,
24        node: &str,
25        label_filter: &[&str],
26        direction: Direction,
27        collection: &str,
28    ) -> Vec<(String, String)> {
29        let Some(&node_id) = self.node_to_id.get(node) else {
30            return Vec::new();
31        };
32        let Some(collection_id) = self.collection_id(collection) else {
33            return Vec::new();
34        };
35        self.record_access(node_id);
36        let labels = self.label_filter(label_filter);
37        let keep = |lid: u32| labels.keeps(lid);
38
39        let mut result = Vec::new();
40        if matches!(direction, Direction::Out | Direction::Both) {
41            for (lid, dst) in self.iter_out_edges_raw_in(node_id, collection_id) {
42                if keep(lid) {
43                    result.push((
44                        self.id_to_label[lid as usize].clone(),
45                        self.id_to_node[dst as usize].clone(),
46                    ));
47                }
48            }
49        }
50        if matches!(direction, Direction::In | Direction::Both) {
51            for (lid, src) in self.iter_in_edges_raw_in(node_id, collection_id) {
52                if keep(lid) {
53                    result.push((
54                        self.id_to_label[lid as usize].clone(),
55                        self.id_to_node[src as usize].clone(),
56                    ));
57                }
58            }
59        }
60        result
61    }
62
63    /// Collection-scoped raw outbound iteration: yields `(label_id, dst)` for
64    /// edges inserted under `collection_id` only (dense + buffer, minus
65    /// deleted). Used by the MATCH / RAG collection-scoped read paths so a
66    /// `... IN '<collection>'` query never traverses another collection's
67    /// edges even though all collections share one CSR partition.
68    pub fn iter_out_edges_raw_in(
69        &self,
70        node: u32,
71        collection_id: u32,
72    ) -> impl Iterator<Item = (u32, u32)> + '_ {
73        self.dense_scoped(node, collection_id, true).into_iter()
74    }
75
76    /// Collection-scoped raw inbound iteration (see `iter_out_edges_raw_in`).
77    pub fn iter_in_edges_raw_in(
78        &self,
79        node: u32,
80        collection_id: u32,
81    ) -> impl Iterator<Item = (u32, u32)> + '_ {
82        self.dense_scoped(node, collection_id, false).into_iter()
83    }
84
85    /// Build the collection-filtered `(label, neighbor)` list for a node.
86    /// `outbound = true` selects out-edges, `false` selects in-edges.
87    fn dense_scoped(&self, node: u32, collection_id: u32, outbound: bool) -> Vec<(u32, u32)> {
88        let idx = node as usize;
89        let (offsets, targets, labels, collections) = if outbound {
90            (
91                &self.out_offsets,
92                &self.out_targets,
93                &self.out_labels,
94                &self.out_collections,
95            )
96        } else {
97            (
98                &self.in_offsets,
99                &self.in_targets,
100                &self.in_labels,
101                &self.in_collections,
102            )
103        };
104        let mut result = Vec::new();
105        if idx + 1 < offsets.len() {
106            let start = offsets[idx] as usize;
107            let end = offsets[idx + 1] as usize;
108            for i in start..end {
109                if collections.get(i).copied().unwrap_or(0) != collection_id {
110                    continue;
111                }
112                let lid = labels[i];
113                let other = targets[i];
114                // Surviving edges here all have `collection_id` (filtered
115                // above), so deletion is keyed on the full identity.
116                let deleted = if outbound {
117                    self.deleted_edges
118                        .contains(&(node, lid, other, collection_id))
119                } else {
120                    self.deleted_edges
121                        .contains(&(other, lid, node, collection_id))
122                };
123                if !deleted {
124                    result.push((lid, other));
125                }
126            }
127        }
128        // Buffer edges for this node, filtered by collection.
129        let (buf, buf_coll) = if outbound {
130            (
131                self.buffer_out.get(idx),
132                self.buffer_out_collections.get(idx),
133            )
134        } else {
135            (self.buffer_in.get(idx), self.buffer_in_collections.get(idx))
136        };
137        if let (Some(buf), Some(buf_coll)) = (buf, buf_coll) {
138            for (k, &(lid, other)) in buf.iter().enumerate() {
139                if buf_coll.get(k).copied().unwrap_or(0) == collection_id {
140                    result.push((lid, other));
141                }
142            }
143        }
144        result
145    }
146}