hawdb_executor/external/
text.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 super::{ExternalReadResourceContract, ExternalReadResultBudget};
16use crate::{QueryMemoryAccount, QueryMemoryLease};
17use hawdb_core::{HawDBError, Result};
18use std::collections::BTreeMap;
19use std::mem::size_of;
20
21pub struct TextSeedExecutionRequest<'a> {
22    pub query_text: &'a str,
23    pub metadata_filters: &'a BTreeMap<String, String>,
24    pub resources: ExternalReadResourceContract<'a>,
25    pub working_account: &'a QueryMemoryAccount,
26    pub result_account: &'a QueryMemoryAccount,
27}
28
29#[cfg(test)]
30mod tests {
31    use super::*;
32    use crate::{QueryMemoryClass, QueryMemoryLedger};
33    use std::num::NonZeroUsize;
34
35    #[test]
36    fn text_result_owner_retains_the_query_root_allowance_until_drop() {
37        let bytes = NonZeroUsize::new(1024).unwrap();
38        let ledger = QueryMemoryLedger::new(bytes);
39        let parent = ledger.account(QueryMemoryClass::ExternalRead, "text allowance", bytes);
40        let other = ledger.account(QueryMemoryClass::PipelineBatch, "downstream", bytes);
41        let admitted = parent.sub_account(bytes).unwrap();
42        let account = admitted.sibling(QueryMemoryClass::ExternalRead, "text result", bytes);
43        let budget = ExternalReadResultBudget {
44            max_rows: 2,
45            max_memory_bytes: bytes,
46        };
47        let mut output = TextSeedExecutionOutput::new(&account, budget).unwrap();
48        output
49            .push("document:first", Some("canonical:first"), 1.0)
50            .unwrap();
51        output
52            .push("document:second", Some("canonical:second"), 2.0)
53            .unwrap();
54        output.validate(&account, budget).unwrap();
55        drop(account);
56        drop(admitted);
57        assert_eq!(ledger.snapshot().used_bytes, bytes.get());
58        assert!(
59            other.reserve(1).is_err(),
60            "external output detached its root allowance"
61        );
62        assert_eq!(
63            output.rows()[1].external_id.as_deref(),
64            Some("canonical:second")
65        );
66        drop(output);
67        assert_eq!(ledger.snapshot().used_bytes, 0);
68        assert!(other.reserve(bytes.get()).is_ok());
69    }
70
71    #[test]
72    fn text_result_builder_refuses_row_payload_and_nonfinite_failures_without_partial_success() {
73        for failure in ["row", "payload", "score"] {
74            let bytes = NonZeroUsize::new(256).unwrap();
75            let ledger = QueryMemoryLedger::new(bytes);
76            let account = ledger.account(QueryMemoryClass::ExternalRead, "text result", bytes);
77            let budget = ExternalReadResultBudget {
78                max_rows: if failure == "row" { 1 } else { 2 },
79                max_memory_bytes: bytes,
80            };
81            let mut output = TextSeedExecutionOutput::new(&account, budget).unwrap();
82            output.push("a", Some("canonical:a"), 1.0).unwrap();
83            let before = ledger.snapshot().used_bytes;
84            let oversized = "x".repeat(1024);
85            let (id, score) = match failure {
86                "row" => ("b", 2.0),
87                "payload" => (oversized.as_str(), 2.0),
88                _ => ("b", f64::INFINITY),
89            };
90            let error = output.push(id, Some("canonical:b"), score).unwrap_err();
91            if failure == "row" {
92                assert!(matches!(&error, HawDBError::ReadBudgetExceeded(cause)
93                    if cause.resource == hawdb_core::ReadBudgetResource::Rows && cause.limit == 1));
94            } else {
95                assert!(matches!(&error, HawDBError::Execution(_)));
96            }
97            assert_eq!(
98                ledger.snapshot().used_bytes,
99                before,
100                "{failure} charged a rejected copy"
101            );
102            assert_eq!(output.rows().len(), 1);
103            assert_eq!(output.validate(&account, budget).unwrap_err(), error);
104            assert_eq!(output.push("c", None, 3.0).unwrap_err(), error);
105            assert_eq!(ledger.snapshot().used_bytes, before);
106            drop(output);
107            assert_eq!(ledger.snapshot().used_bytes, 0);
108            assert!(account.reserve(bytes.get()).is_ok());
109        }
110    }
111
112    #[test]
113    fn text_result_owner_requires_the_exact_admitted_account_and_prepaid_row_storage() {
114        let bytes = NonZeroUsize::new(256).unwrap();
115        let ledger = QueryMemoryLedger::new(bytes);
116        let account = ledger.account(QueryMemoryClass::ExternalRead, "text result", bytes);
117        let foreign_ledger = QueryMemoryLedger::new(bytes);
118        let foreign = foreign_ledger.account(QueryMemoryClass::ExternalRead, "foreign", bytes);
119        let sibling = ledger.account(QueryMemoryClass::ExternalRead, "wrong sibling", bytes);
120        let budget = ExternalReadResultBudget {
121            max_rows: 1,
122            max_memory_bytes: bytes,
123        };
124        for wrong in [&foreign, &sibling] {
125            let output = TextSeedExecutionOutput::new(wrong, budget).unwrap();
126            assert!(output.validate(&account, budget).is_err());
127        }
128        let mut output = TextSeedExecutionOutput::new(&account, budget).unwrap();
129        output.push("a", None, 1.0).unwrap();
130        let before = ledger.snapshot().used_bytes;
131        let narrowed = ExternalReadResultBudget {
132            max_rows: 0,
133            ..budget
134        };
135        assert!(matches!(output.validate(&account, narrowed),
136            Err(HawDBError::ReadBudgetExceeded(cause))
137                if cause.resource == hawdb_core::ReadBudgetResource::Rows && cause.limit == 0));
138        assert_eq!(output.rows().len(), 1);
139        assert_eq!(ledger.snapshot().used_bytes, before);
140        drop(output);
141        assert_eq!(ledger.snapshot().used_bytes, 0);
142        let held = account.reserve(bytes.get()).unwrap();
143        assert!(TextSeedExecutionOutput::new(&account, budget).is_err());
144        assert_eq!(ledger.snapshot().used_bytes, bytes.get());
145        drop(held);
146        let oversized_rows = ExternalReadResultBudget {
147            max_rows: usize::MAX,
148            ..budget
149        };
150        assert!(TextSeedExecutionOutput::new(&account, oversized_rows).is_err());
151        assert_eq!(ledger.snapshot().used_bytes, 0);
152        assert!(TextSeedExecutionOutput::new(&account, budget).is_ok());
153        assert_eq!(ledger.snapshot().used_bytes, 0);
154    }
155}
156
157#[derive(Debug)]
158pub struct TextSeedExecutionRow {
159    pub id: String,
160    pub external_id: Option<String>,
161    pub score: f64,
162}
163
164/// Rows admitted before copying and retained in the supplied result account.
165/// The caller can borrow rows, but cannot detach their owning reservation.
166#[derive(Debug)]
167pub struct TextSeedExecutionOutput {
168    rows: Vec<TextSeedExecutionRow>,
169    budget: ExternalReadResultBudget,
170    failure: Option<HawDBError>,
171    reservation: QueryMemoryLease,
172}
173
174impl TextSeedExecutionOutput {
175    pub fn new(account: &QueryMemoryAccount, budget: ExternalReadResultBudget) -> Result<Self> {
176        let bytes = budget
177            .max_rows
178            .checked_mul(size_of::<TextSeedExecutionRow>())
179            .ok_or_else(|| {
180                HawDBError::Execution("text seed result storage size overflow".into())
181            })?;
182        if bytes > budget.max_memory_bytes.get() {
183            return Err(HawDBError::Execution(
184                "text seed result row storage exceeds its memory budget".into(),
185            ));
186        }
187        let reservation = account.reserve(bytes)?;
188        let rows = Vec::with_capacity(budget.max_rows);
189        Ok(Self {
190            rows,
191            budget,
192            failure: None,
193            reservation,
194        })
195    }
196
197    pub fn push(&mut self, id: &str, external_id: Option<&str>, score: f64) -> Result<()> {
198        if let Some(error) = &self.failure {
199            return Err(error.clone());
200        }
201        let result = (|| {
202            if !score.is_finite() {
203                return Err(HawDBError::Execution(
204                    "text seed score must be finite".into(),
205                ));
206            }
207            if self.rows.len() >= self.budget.max_rows {
208                return Err(HawDBError::read_budget_exceeded(
209                    hawdb_core::ReadBudgetResource::Rows,
210                    self.budget.max_rows,
211                    "text seed result exceeds its row budget",
212                ));
213            }
214            let bytes = id
215                .len()
216                .checked_add(external_id.map_or(0, str::len))
217                .ok_or_else(|| {
218                    HawDBError::Execution("text seed result payload size overflow".into())
219                })?;
220            let next = self.reservation.bytes().checked_add(bytes).ok_or_else(|| {
221                HawDBError::Execution("text seed result payload size overflow".into())
222            })?;
223            if next > self.budget.max_memory_bytes.get() {
224                return Err(HawDBError::Execution(
225                    "text seed result payload exceeds its memory budget".into(),
226                ));
227            }
228            self.reservation.grow(bytes)?;
229            self.rows.push(TextSeedExecutionRow {
230                id: id.into(),
231                external_id: external_id.map(str::to_owned),
232                score,
233            });
234            Ok(())
235        })();
236        if let Err(error) = &result {
237            self.failure = Some(error.clone());
238        }
239        result
240    }
241
242    pub fn rows(&self) -> &[TextSeedExecutionRow] {
243        &self.rows
244    }
245
246    pub(crate) fn validate(
247        &self,
248        account: &QueryMemoryAccount,
249        budget: ExternalReadResultBudget,
250    ) -> Result<()> {
251        if !self.reservation.belongs_to(account) {
252            return Err(HawDBError::Execution(
253                "text seed output did not retain the admitted result account and budget".into(),
254            ));
255        }
256        if let Some(error) = &self.failure {
257            return Err(error.clone());
258        }
259        if self.rows.len() > budget.max_rows {
260            return Err(HawDBError::read_budget_exceeded(
261                hawdb_core::ReadBudgetResource::Rows,
262                budget.max_rows,
263                "text seed result exceeds its validated row budget",
264            ));
265        }
266        if self.reservation.bytes() > budget.max_memory_bytes.get() {
267            return Err(HawDBError::Execution(
268                "text seed output did not retain the admitted result account and budget".into(),
269            ));
270        }
271        Ok(())
272    }
273}
hawdb_executor/external/text.rs

hawdb_executor/external/
text.rs
