hawdb_executor/
graph.rs

1// Copyright 2026 Nowledge
2//
3// Licensed under the Apache License, Version 2.0 (the "License");
4// you may not use this file except in compliance with the License.
5// You may obtain a copy of the License at
6//
7//     http://www.apache.org/licenses/LICENSE-2.0
8//
9// Unless required by applicable law or agreed to in writing, software
10// distributed under the License is distributed on an "AS IS" BASIS,
11// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12// See the License for the specific language governing permissions and
13// limitations under the License.
14
15use crate::binding::{binding_payload_bytes, Binding};
16use crate::{QueryMemoryAccount, QueryMemoryLease};
17use hawdb_core::Result;
18use hawdb_plan::GraphExpansionBudget;
19use hawdb_storage::NodeId;
20use std::collections::BTreeSet;
21
22#[derive(Debug, Clone, Copy, PartialEq, Eq)]
23pub enum GraphExpansionTruncationReason {
24    CandidateLimit,
25    PayloadByteLimit,
26}
27
28impl GraphExpansionTruncationReason {
29    pub const fn as_str(self) -> &'static str {
30        match self {
31            Self::CandidateLimit => "candidate_limit",
32            Self::PayloadByteLimit => "payload_byte_limit",
33        }
34    }
35}
36
37#[derive(Debug, Clone, PartialEq, Eq)]
38pub struct GraphExpansionExecutionReport {
39    pub seed_count: usize,
40    pub expanded_node_count: usize,
41    pub expanded_edge_count: usize,
42    pub relation_types: Vec<String>,
43    pub min_hops: usize,
44    pub max_hops: usize,
45    pub reranked_seed_count: usize,
46    pub candidate_limit: usize,
47    pub payload_byte_limit: usize,
48    pub payload_bytes_used: usize,
49    pub returned_count: usize,
50    pub truncation_reason: Option<GraphExpansionTruncationReason>,
51}
52
53impl GraphExpansionExecutionReport {
54    pub fn truncated(&self) -> bool {
55        self.truncation_reason.is_some()
56    }
57}
58
59#[doc(hidden)]
60pub struct GraphExpansionExecutionState {
61    budget: Option<GraphExpansionBudget>,
62    seed_count: usize,
63    expanded_nodes: BTreeSet<NodeId>,
64    expanded_nodes_lease: Option<QueryMemoryLease>,
65    expanded_edge_count: usize,
66    reranked_seed_count: usize,
67    payload_bytes_used: usize,
68    returned_count: usize,
69    pub truncation_reason: Option<GraphExpansionTruncationReason>,
70}
71
72impl GraphExpansionExecutionState {
73    pub fn new(
74        budget: Option<GraphExpansionBudget>,
75        seed_count: usize,
76        reranked_seed_count: usize,
77    ) -> Self {
78        Self {
79            budget,
80            seed_count,
81            expanded_nodes: BTreeSet::new(),
82            expanded_nodes_lease: None,
83            expanded_edge_count: 0,
84            reranked_seed_count,
85            payload_bytes_used: 0,
86            returned_count: 0,
87            truncation_reason: None,
88        }
89    }
90
91    pub fn with_memory_account(
92        budget: Option<GraphExpansionBudget>,
93        seed_count: usize,
94        reranked_seed_count: usize,
95        memory_account: &QueryMemoryAccount,
96    ) -> Result<Self> {
97        Ok(Self {
98            expanded_nodes_lease: Some(memory_account.reserve(0)?),
99            ..Self::new(budget, seed_count, reranked_seed_count)
100        })
101    }
102
103    pub fn try_push(
104        &mut self,
105        output: &mut Vec<Binding>,
106        candidate: Binding,
107        target_id: Option<NodeId>,
108        hop: usize,
109    ) -> Result<bool> {
110        if !self.try_admit(&candidate, target_id, hop)? {
111            return Ok(false);
112        }
113        output.push(candidate);
114        Ok(true)
115    }
116
117    pub fn try_admit(
118        &mut self,
119        candidate: &Binding,
120        target_id: Option<NodeId>,
121        hop: usize,
122    ) -> Result<bool> {
123        let Some(budget) = self.budget else {
124            self.returned_count = self.returned_count.saturating_add(1);
125            return Ok(true);
126        };
127        if self.returned_count >= budget.candidate_limit {
128            self.truncation_reason = Some(GraphExpansionTruncationReason::CandidateLimit);
129            return Ok(false);
130        }
131        let candidate_bytes = binding_payload_bytes(candidate);
132        if self.payload_bytes_used.saturating_add(candidate_bytes) > budget.payload_byte_limit {
133            self.truncation_reason = Some(GraphExpansionTruncationReason::PayloadByteLimit);
134            return Ok(false);
135        }
136        if let Some(target_id) = target_id
137            && !self.expanded_nodes.contains(&target_id)
138        {
139            const EXPANDED_NODE_MEMORY_BYTES: usize =
140                std::mem::size_of::<NodeId>() + 3 * std::mem::size_of::<usize>();
141            if let Some(lease) = self.expanded_nodes_lease.as_mut() {
142                lease.grow(EXPANDED_NODE_MEMORY_BYTES)?;
143            }
144            if !self.expanded_nodes.insert(target_id)
145                && let Some(lease) = self.expanded_nodes_lease.as_mut()
146            {
147                lease.shrink(EXPANDED_NODE_MEMORY_BYTES);
148            }
149        }
150        self.payload_bytes_used = self.payload_bytes_used.saturating_add(candidate_bytes);
151        self.returned_count = self.returned_count.saturating_add(1);
152        self.expanded_edge_count = self.expanded_edge_count.saturating_add(hop);
153        Ok(true)
154    }
155
156    pub fn record_seed(&mut self) {
157        self.seed_count = self.seed_count.saturating_add(1);
158    }
159
160    pub fn set_reranked_seed_count(&mut self, reranked_seed_count: usize) {
161        self.reranked_seed_count = reranked_seed_count;
162    }
163
164    pub fn returned_count(&self) -> usize {
165        self.returned_count
166    }
167
168    pub fn report(
169        &self,
170        rel_type: &str,
171        min_hops: usize,
172        max_hops: usize,
173        returned_count: usize,
174    ) -> Option<GraphExpansionExecutionReport> {
175        let budget = self.budget?;
176        Some(GraphExpansionExecutionReport {
177            seed_count: self.seed_count,
178            expanded_node_count: self.expanded_nodes.len(),
179            expanded_edge_count: self.expanded_edge_count,
180            relation_types: if rel_type.is_empty() {
181                Vec::new()
182            } else {
183                vec![rel_type.to_string()]
184            },
185            min_hops,
186            max_hops,
187            reranked_seed_count: self.reranked_seed_count,
188            candidate_limit: budget.candidate_limit,
189            payload_byte_limit: budget.payload_byte_limit,
190            payload_bytes_used: self.payload_bytes_used,
191            returned_count,
192            truncation_reason: self.truncation_reason,
193        })
194    }
195}
196
197#[cfg(test)]
198mod tests {
199    use super::*;
200
201    #[test]
202    fn truncation_reason_codes_are_stable() {
203        assert_eq!(
204            GraphExpansionTruncationReason::CandidateLimit.as_str(),
205            "candidate_limit"
206        );
207        assert_eq!(
208            GraphExpansionTruncationReason::PayloadByteLimit.as_str(),
209            "payload_byte_limit"
210        );
211    }
212
213    #[test]
214    fn reranked_seed_count_can_be_recorded_after_streaming_input() {
215        let mut state = GraphExpansionExecutionState::new(
216            Some(GraphExpansionBudget {
217                candidate_limit: 8,
218                payload_byte_limit: 1_024,
219            }),
220            0,
221            0,
222        );
223
224        state.set_reranked_seed_count(3);
225
226        let report = state.report("related", 1, 1, 0).expect("budgeted report");
227        assert_eq!(report.reranked_seed_count, 3);
228    }
229
230    #[test]
231    fn expansion_state_enforces_candidate_and_payload_budgets_before_push() {
232        let binding = Binding {
233            values: std::collections::BTreeMap::from([(
234                "value".to_string(),
235                hawdb_core::Value::String("payload".to_string()),
236            )]),
237            nodes: std::collections::BTreeMap::new(),
238            relationships: std::collections::BTreeMap::new(),
239        };
240        let mut state = GraphExpansionExecutionState::new(
241            Some(GraphExpansionBudget {
242                candidate_limit: 1,
243                payload_byte_limit: usize::MAX,
244            }),
245            1,
246            1,
247        );
248        let mut output = Vec::new();
249
250        assert!(state
251            .try_push(&mut output, binding.clone(), None, 1)
252            .unwrap());
253        assert!(!state.try_push(&mut output, binding, None, 1).unwrap());
254        assert_eq!(
255            state.truncation_reason,
256            Some(GraphExpansionTruncationReason::CandidateLimit)
257        );
258
259        let binding = Binding {
260            values: std::collections::BTreeMap::from([(
261                "value".to_string(),
262                hawdb_core::Value::String("payload".to_string()),
263            )]),
264            nodes: std::collections::BTreeMap::new(),
265            relationships: std::collections::BTreeMap::new(),
266        };
267        let mut state = GraphExpansionExecutionState::new(
268            Some(GraphExpansionBudget {
269                candidate_limit: 2,
270                payload_byte_limit: binding_payload_bytes(&binding).saturating_sub(1),
271            }),
272            1,
273            1,
274        );
275        let mut output = Vec::new();
276
277        assert!(!state.try_push(&mut output, binding, None, 1).unwrap());
278        assert!(output.is_empty());
279        assert_eq!(
280            state.truncation_reason,
281            Some(GraphExpansionTruncationReason::PayloadByteLimit)
282        );
283    }
284
285    #[test]
286    fn expansion_node_set_uses_and_releases_the_query_root() {
287        let ledger = crate::QueryMemoryLedger::new(std::num::NonZeroUsize::new(64).unwrap());
288        let account = ledger.account(
289            crate::QueryMemoryClass::BlockingState,
290            "test expansion",
291            std::num::NonZeroUsize::new(64).unwrap(),
292        );
293        let mut state = GraphExpansionExecutionState::with_memory_account(
294            Some(GraphExpansionBudget {
295                candidate_limit: 2,
296                payload_byte_limit: usize::MAX,
297            }),
298            1,
299            0,
300            &account,
301        )
302        .unwrap();
303        let binding = Binding::values(std::collections::BTreeMap::new());
304
305        assert!(state.try_admit(&binding, Some(NodeId(7)), 1).unwrap());
306        assert!(ledger.snapshot().used_bytes > 0);
307        assert!(state.try_admit(&binding, Some(NodeId(7)), 1).unwrap());
308        assert_eq!(state.expanded_nodes.len(), 1);
309
310        drop(state);
311        assert_eq!(ledger.snapshot().used_bytes, 0);
312    }
313}
hawdb_executor/graph.rs

hawdb_executor/
graph.rs
