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::{HawDBError, Result};
18use hawdb_plan_cypher::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        self.try_admit_payload(binding_payload_bytes(candidate), target_id, hop)
124    }
125
126    pub(crate) fn try_admit_payload(
127        &mut self,
128        candidate_bytes: usize,
129        target_id: Option<NodeId>,
130        hop: usize,
131    ) -> Result<bool> {
132        let Some(budget) = self.budget else {
133            self.returned_count = self.returned_count.saturating_add(1);
134            return Ok(true);
135        };
136        if self.returned_count >= budget.candidate_limit {
137            self.truncation_reason = Some(GraphExpansionTruncationReason::CandidateLimit);
138            return Err(HawDBError::GraphExpansionCandidateLimitExceeded {
139                requested: self.returned_count.saturating_add(1),
140                limit: budget.candidate_limit,
141            });
142        }
143        if self.payload_bytes_used.saturating_add(candidate_bytes) > budget.payload_byte_limit {
144            self.truncation_reason = Some(GraphExpansionTruncationReason::PayloadByteLimit);
145            return Err(HawDBError::GraphExpansionPayloadLimitExceeded {
146                requested: self.payload_bytes_used.saturating_add(candidate_bytes),
147                limit: budget.payload_byte_limit,
148            });
149        }
150        if let Some(target_id) = target_id
151            && !self.expanded_nodes.contains(&target_id)
152        {
153            const EXPANDED_NODE_MEMORY_BYTES: usize =
154                std::mem::size_of::<NodeId>() + 3 * std::mem::size_of::<usize>();
155            if let Some(lease) = self.expanded_nodes_lease.as_mut() {
156                lease.grow(EXPANDED_NODE_MEMORY_BYTES)?;
157            }
158            if !self.expanded_nodes.insert(target_id)
159                && let Some(lease) = self.expanded_nodes_lease.as_mut()
160            {
161                lease.shrink(EXPANDED_NODE_MEMORY_BYTES);
162            }
163        }
164        self.payload_bytes_used = self.payload_bytes_used.saturating_add(candidate_bytes);
165        self.returned_count = self.returned_count.saturating_add(1);
166        self.expanded_edge_count = self.expanded_edge_count.saturating_add(hop);
167        Ok(true)
168    }
169
170    pub fn record_seed(&mut self) {
171        self.seed_count = self.seed_count.saturating_add(1);
172    }
173
174    pub fn set_reranked_seed_count(&mut self, reranked_seed_count: usize) {
175        self.reranked_seed_count = reranked_seed_count;
176    }
177
178    pub fn returned_count(&self) -> usize {
179        self.returned_count
180    }
181
182    pub fn report(
183        &self,
184        rel_type: &str,
185        min_hops: usize,
186        max_hops: usize,
187        returned_count: usize,
188    ) -> Option<GraphExpansionExecutionReport> {
189        let budget = self.budget?;
190        Some(GraphExpansionExecutionReport {
191            seed_count: self.seed_count,
192            expanded_node_count: self.expanded_nodes.len(),
193            expanded_edge_count: self.expanded_edge_count,
194            relation_types: if rel_type.is_empty() {
195                Vec::new()
196            } else {
197                vec![rel_type.to_string()]
198            },
199            min_hops,
200            max_hops,
201            reranked_seed_count: self.reranked_seed_count,
202            candidate_limit: budget.candidate_limit,
203            payload_byte_limit: budget.payload_byte_limit,
204            payload_bytes_used: self.payload_bytes_used,
205            returned_count,
206            truncation_reason: self.truncation_reason,
207        })
208    }
209}
210
211#[cfg(test)]
212mod tests {
213    use super::*;
214
215    #[test]
216    fn truncation_reason_codes_are_stable() {
217        assert_eq!(
218            GraphExpansionTruncationReason::CandidateLimit.as_str(),
219            "candidate_limit"
220        );
221        assert_eq!(
222            GraphExpansionTruncationReason::PayloadByteLimit.as_str(),
223            "payload_byte_limit"
224        );
225    }
226
227    #[test]
228    fn reranked_seed_count_can_be_recorded_after_streaming_input() {
229        let mut state = GraphExpansionExecutionState::new(
230            Some(GraphExpansionBudget {
231                candidate_limit: 8,
232                payload_byte_limit: 1_024,
233            }),
234            0,
235            0,
236        );
237
238        state.set_reranked_seed_count(3);
239
240        let report = state.report("related", 1, 1, 0).expect("budgeted report");
241        assert_eq!(report.reranked_seed_count, 3);
242    }
243
244    #[test]
245    fn expansion_state_enforces_candidate_and_payload_budgets_before_push() {
246        let binding = Binding {
247            values: std::collections::BTreeMap::from([(
248                "value".to_string(),
249                hawdb_core::Value::String("payload".to_string()),
250            )]),
251            nodes: std::collections::BTreeMap::new(),
252            relationships: std::collections::BTreeMap::new(),
253        };
254        let mut state = GraphExpansionExecutionState::new(
255            Some(GraphExpansionBudget {
256                candidate_limit: 1,
257                payload_byte_limit: usize::MAX,
258            }),
259            1,
260            1,
261        );
262        let mut output = Vec::new();
263
264        assert!(state
265            .try_push(&mut output, binding.clone(), None, 1)
266            .unwrap());
267        let error = state.try_push(&mut output, binding, None, 1).unwrap_err();
268        assert_eq!(
269            error,
270            hawdb_core::HawDBError::GraphExpansionCandidateLimitExceeded {
271                requested: 2,
272                limit: 1,
273            }
274        );
275        assert_eq!(output.len(), 1);
276        assert!(error.to_string().contains("2 rows, limit 1 rows"));
277        assert!(error
278            .to_string()
279            .contains("graph_expansion_budget.candidate_limit"));
280        assert_eq!(
281            state.truncation_reason,
282            Some(GraphExpansionTruncationReason::CandidateLimit)
283        );
284
285        let binding = Binding {
286            values: std::collections::BTreeMap::from([(
287                "value".to_string(),
288                hawdb_core::Value::String("payload".to_string()),
289            )]),
290            nodes: std::collections::BTreeMap::new(),
291            relationships: std::collections::BTreeMap::new(),
292        };
293        let mut state = GraphExpansionExecutionState::new(
294            Some(GraphExpansionBudget {
295                candidate_limit: 2,
296                payload_byte_limit: binding_payload_bytes(&binding).saturating_sub(1),
297            }),
298            1,
299            1,
300        );
301        let mut output = Vec::new();
302
303        let requested = binding_payload_bytes(&binding);
304        let error = state.try_push(&mut output, binding, None, 1).unwrap_err();
305        assert_eq!(
306            error,
307            hawdb_core::HawDBError::GraphExpansionPayloadLimitExceeded {
308                requested,
309                limit: requested - 1,
310            }
311        );
312        assert!(error.to_string().contains("bytes"));
313        assert!(error
314            .to_string()
315            .contains("graph_expansion_budget.payload_byte_limit"));
316        assert!(output.is_empty());
317        assert_eq!(
318            state.truncation_reason,
319            Some(GraphExpansionTruncationReason::PayloadByteLimit)
320        );
321    }
322
323    #[test]
324    fn expansion_state_cumulative_payload_budget_refuses_before_append_and_recovers() {
325        let binding = Binding::values(std::collections::BTreeMap::from([(
326            "value".into(),
327            hawdb_core::Value::String("payload".into()),
328        )]));
329        let row_bytes = binding_payload_bytes(&binding);
330        assert!(row_bytes > 0);
331        let root_budget = std::num::NonZeroUsize::new(4096).unwrap();
332        let ledger = crate::QueryMemoryLedger::new(root_budget);
333        let account = ledger.account(
334            crate::QueryMemoryClass::BlockingState,
335            "cumulative expansion",
336            root_budget,
337        );
338        for allowed_rows in [2, 3] {
339            let mut state = GraphExpansionExecutionState::with_memory_account(
340                Some(GraphExpansionBudget {
341                    candidate_limit: 10,
342                    payload_byte_limit: row_bytes * allowed_rows,
343                }),
344                1,
345                0,
346                &account,
347            )
348            .unwrap();
349            let mut output = Vec::new();
350            for id in [NodeId(1), NodeId(2)] {
351                assert!(state
352                    .try_push(&mut output, binding.clone(), Some(id), 1)
353                    .unwrap());
354            }
355            assert_eq!(output.len(), 2);
356            let before = ledger.snapshot();
357            assert!(before.used_bytes > 0);
358            let result = state.try_push(&mut output, binding.clone(), Some(NodeId(3)), 1);
359            if allowed_rows == 2 {
360                assert_eq!(
361                    result.unwrap_err(),
362                    HawDBError::GraphExpansionPayloadLimitExceeded {
363                        requested: row_bytes * 3,
364                        limit: row_bytes * 2,
365                    }
366                );
367                assert_eq!(output.len(), 2);
368                assert_eq!(state.returned_count, 2);
369                assert_eq!(state.expanded_nodes.len(), 2);
370                assert_eq!(
371                    ledger.snapshot(),
372                    before,
373                    "refusal must not reserve or insert the third identity"
374                );
375            } else {
376                assert!(result.unwrap());
377                assert_eq!(output.len(), 3);
378                assert_eq!(state.payload_bytes_used, row_bytes * 3);
379            }
380            drop(output);
381            drop(state);
382            assert_eq!(ledger.snapshot().used_bytes, 0);
383        }
384    }
385
386    #[test]
387    fn expansion_node_set_uses_and_releases_the_query_root() {
388        let ledger = crate::QueryMemoryLedger::new(std::num::NonZeroUsize::new(64).unwrap());
389        let account = ledger.account(
390            crate::QueryMemoryClass::BlockingState,
391            "test expansion",
392            std::num::NonZeroUsize::new(64).unwrap(),
393        );
394        let mut state = GraphExpansionExecutionState::with_memory_account(
395            Some(GraphExpansionBudget {
396                candidate_limit: 2,
397                payload_byte_limit: usize::MAX,
398            }),
399            1,
400            0,
401            &account,
402        )
403        .unwrap();
404        let binding = Binding::values(std::collections::BTreeMap::new());
405
406        assert!(state.try_admit(&binding, Some(NodeId(7)), 1).unwrap());
407        assert!(ledger.snapshot().used_bytes > 0);
408        assert!(state.try_admit(&binding, Some(NodeId(7)), 1).unwrap());
409        assert_eq!(state.expanded_nodes.len(), 1);
410
411        drop(state);
412        assert_eq!(ledger.snapshot().used_bytes, 0);
413    }
414    #[test]
415    fn unlimited_expansion_retains_reports_and_fails_at_admitted_memory() {
416        let entry_bytes = std::mem::size_of::<NodeId>() + 3 * std::mem::size_of::<usize>();
417        for admitted_nodes in [2, 3] {
418            let limit = std::num::NonZeroUsize::new(entry_bytes * admitted_nodes).unwrap();
419            let ledger = crate::QueryMemoryLedger::new(limit);
420            let account = ledger.account(
421                crate::QueryMemoryClass::BlockingState,
422                "seeded expansion",
423                limit,
424            );
425            let mut state = GraphExpansionExecutionState::with_memory_account(
426                Some(GraphExpansionBudget {
427                    candidate_limit: usize::MAX,
428                    payload_byte_limit: usize::MAX,
429                }),
430                1,
431                0,
432                &account,
433            )
434            .unwrap();
435            let payload_bytes = 8 * 1024 * 1024 + 1;
436            for id in [NodeId(1), NodeId(2), NodeId(1)] {
437                assert!(state.try_admit_payload(payload_bytes, Some(id), 1).unwrap());
438            }
439            assert_eq!(ledger.snapshot().used_bytes, entry_bytes * 2);
440            let before = ledger.snapshot();
441            let result = state.try_admit_payload(payload_bytes, Some(NodeId(3)), 1);
442            if admitted_nodes == 2 {
443                let HawDBError::ResourceBudgetExceeded(cause) = result.unwrap_err() else {
444                    panic!("expected a typed memory refusal");
445                };
446                assert_eq!(cause.resource, hawdb_core::ResourceBudgetKind::MemoryBytes);
447                assert_eq!(cause.owner, "seeded expansion");
448                assert_eq!(cause.requested, entry_bytes as u64);
449                assert_eq!(cause.used, (entry_bytes * 2) as u64);
450                assert_eq!(cause.budget, limit.get() as u64);
451                assert_eq!(ledger.snapshot(), before);
452            } else {
453                assert!(result.unwrap());
454            }
455            let rows = admitted_nodes + 1;
456            let report = state.report("LINK", 1, 1, state.returned_count()).unwrap();
457            assert_eq!(report.seed_count, 1);
458            assert_eq!(report.expanded_node_count, admitted_nodes);
459            assert_eq!(report.returned_count, rows);
460            assert_eq!(report.expanded_edge_count, rows);
461            assert_eq!(report.payload_bytes_used, rows * payload_bytes);
462            assert!(!report.truncated());
463            drop(state);
464            assert_eq!(ledger.snapshot().used_bytes, 0);
465        }
466    }
467}
hawdb_executor/graph.rs

hawdb_executor/
graph.rs
