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        for growing in [false, true] {
38            let bytes = NonZeroUsize::new(1024).unwrap();
39            let ledger = QueryMemoryLedger::new(bytes);
40            let parent = ledger.account(QueryMemoryClass::ExternalRead, "text allowance", bytes);
41            let other = ledger.account(QueryMemoryClass::PipelineBatch, "downstream", bytes);
42            let admitted = parent.sub_account(bytes).unwrap();
43            let account = admitted.sibling(QueryMemoryClass::ExternalRead, "text result", bytes);
44            let budget = ExternalReadResultBudget {
45                max_rows: 2,
46                max_memory_bytes: bytes,
47            };
48            let mut output = if growing {
49                TextSeedExecutionOutput::with_initial_row_capacity(&account, budget, 0)
50            } else {
51                TextSeedExecutionOutput::new(&account, budget)
52            }
53            .unwrap();
54            output
55                .push("document:first", Some("canonical:first"), 1.0)
56                .unwrap();
57            output
58                .push("document:second", Some("canonical:second"), 2.0)
59                .unwrap();
60            output.validate(&account, budget).unwrap();
61            drop(account);
62            drop(admitted);
63            assert_eq!(ledger.snapshot().used_bytes, bytes.get());
64            assert!(
65                other.reserve(1).is_err(),
66                "external output detached its root allowance"
67            );
68            assert_eq!(
69                output.rows()[1].external_id.as_deref(),
70                Some("canonical:second")
71            );
72            drop(output);
73            assert_eq!(ledger.snapshot().used_bytes, 0);
74            assert!(other.reserve(bytes.get()).is_ok());
75        }
76    }
77
78    #[test]
79    fn text_result_builder_refuses_row_payload_and_nonfinite_failures_without_partial_success() {
80        for failure in ["row", "payload", "score"] {
81            let bytes = NonZeroUsize::new(256).unwrap();
82            let ledger = QueryMemoryLedger::new(bytes);
83            let account = ledger.account(QueryMemoryClass::ExternalRead, "text result", bytes);
84            let budget = ExternalReadResultBudget {
85                max_rows: if failure == "row" { 1 } else { 2 },
86                max_memory_bytes: bytes,
87            };
88            let mut output = TextSeedExecutionOutput::new(&account, budget).unwrap();
89            output.push("a", Some("canonical:a"), 1.0).unwrap();
90            let before = ledger.snapshot().used_bytes;
91            let oversized = "x".repeat(1024);
92            let (id, score) = match failure {
93                "row" => ("b", 2.0),
94                "payload" => (oversized.as_str(), 2.0),
95                _ => ("b", f64::INFINITY),
96            };
97            let error = output.push(id, Some("canonical:b"), score).unwrap_err();
98            if failure == "row" {
99                assert!(matches!(&error, HawDBError::ReadBudgetExceeded(cause)
100                    if cause.resource == hawdb_core::ReadBudgetResource::Rows && cause.limit == 1));
101            } else {
102                assert!(matches!(&error, HawDBError::Execution(_)));
103            }
104            assert_eq!(
105                ledger.snapshot().used_bytes,
106                before,
107                "{failure} charged a rejected copy"
108            );
109            assert_eq!(output.rows().len(), 1);
110            assert_eq!(output.validate(&account, budget).unwrap_err(), error);
111            assert_eq!(output.push("c", None, 3.0).unwrap_err(), error);
112            assert_eq!(ledger.snapshot().used_bytes, before);
113            drop(output);
114            assert_eq!(ledger.snapshot().used_bytes, 0);
115            assert!(account.reserve(bytes.get()).is_ok());
116        }
117    }
118
119    #[test]
120    fn text_result_owner_requires_the_exact_admitted_account_and_prepaid_row_storage() {
121        let bytes = NonZeroUsize::new(256).unwrap();
122        let ledger = QueryMemoryLedger::new(bytes);
123        let account = ledger.account(QueryMemoryClass::ExternalRead, "text result", bytes);
124        let foreign_ledger = QueryMemoryLedger::new(bytes);
125        let foreign = foreign_ledger.account(QueryMemoryClass::ExternalRead, "foreign", bytes);
126        let sibling = ledger.account(QueryMemoryClass::ExternalRead, "wrong sibling", bytes);
127        let budget = ExternalReadResultBudget {
128            max_rows: 1,
129            max_memory_bytes: bytes,
130        };
131        for wrong in [&foreign, &sibling] {
132            let output = TextSeedExecutionOutput::new(wrong, budget).unwrap();
133            assert!(output.validate(&account, budget).is_err());
134        }
135        let mut output = TextSeedExecutionOutput::new(&account, budget).unwrap();
136        output.push("a", None, 1.0).unwrap();
137        let before = ledger.snapshot().used_bytes;
138        let narrowed = ExternalReadResultBudget {
139            max_rows: 0,
140            ..budget
141        };
142        assert!(matches!(output.validate(&account, narrowed),
143            Err(HawDBError::ReadBudgetExceeded(cause))
144                if cause.resource == hawdb_core::ReadBudgetResource::Rows && cause.limit == 0));
145        assert_eq!(output.rows().len(), 1);
146        assert_eq!(ledger.snapshot().used_bytes, before);
147        drop(output);
148        assert_eq!(ledger.snapshot().used_bytes, 0);
149        let held = account.reserve(bytes.get()).unwrap();
150        assert!(TextSeedExecutionOutput::new(&account, budget).is_err());
151        assert_eq!(ledger.snapshot().used_bytes, bytes.get());
152        drop(held);
153        let oversized_rows = ExternalReadResultBudget {
154            max_rows: usize::MAX,
155            ..budget
156        };
157        assert!(TextSeedExecutionOutput::new(&account, oversized_rows).is_err());
158        assert_eq!(ledger.snapshot().used_bytes, 0);
159        assert!(TextSeedExecutionOutput::new(&account, budget).is_ok());
160        assert_eq!(ledger.snapshot().used_bytes, 0);
161    }
162
163    #[test]
164    fn growing_text_result_uses_actual_capacity_and_admits_replacement_overlap() {
165        let row_bytes = size_of::<TextSeedExecutionRow>();
166        let bytes = NonZeroUsize::new(6 * row_bytes + 3).unwrap();
167        let ledger = QueryMemoryLedger::new(bytes);
168        let account = ledger.account(QueryMemoryClass::ExternalRead, "growing text", bytes);
169        let budget = ExternalReadResultBudget {
170            max_rows: 1_000_000,
171            max_memory_bytes: bytes,
172        };
173        let mut output =
174            TextSeedExecutionOutput::with_initial_row_capacity(&account, budget, 0).unwrap();
175        assert_eq!(output.rows.capacity(), 0);
176        assert_eq!(ledger.snapshot().used_bytes, 0);
177        for (id, score) in [("a", 3.0), ("b", 2.0), ("c", 1.0)] {
178            output.push(id, None, score).unwrap();
179        }
180        output.validate(&account, budget).unwrap();
181        assert_eq!(output.rows.capacity(), 4);
182        assert_eq!(output.rows()[2].id, "c");
183        assert_eq!(ledger.snapshot().used_bytes, 4 * row_bytes + 3);
184        assert_eq!(ledger.snapshot().peak_bytes, 6 * row_bytes + 3);
185        drop(output);
186        assert_eq!(ledger.snapshot().used_bytes, 0);
187    }
188
189    #[test]
190    fn growing_text_result_refuses_before_replacement_and_keeps_failure_sticky() {
191        let row_bytes = size_of::<TextSeedExecutionRow>();
192        let bytes = NonZeroUsize::new(6 * row_bytes + 2).unwrap();
193        let ledger = QueryMemoryLedger::new(bytes);
194        let account = ledger.account(QueryMemoryClass::ExternalRead, "growing text", bytes);
195        let budget = ExternalReadResultBudget {
196            max_rows: 1_000_000,
197            max_memory_bytes: bytes,
198        };
199        let mut output =
200            TextSeedExecutionOutput::with_initial_row_capacity(&account, budget, 0).unwrap();
201        output.push("a", None, 3.0).unwrap();
202        output.push("b", None, 2.0).unwrap();
203        let before = ledger.snapshot();
204        let error = output.push("c", None, 1.0).unwrap_err();
205        assert!(matches!(&error, HawDBError::ResourceBudgetExceeded(cause)
206            if cause.resource == hawdb_core::ResourceBudgetKind::MemoryBytes
207                && cause.scope == hawdb_core::ResourceBudgetScope::Operator
208                && cause.requested == (4 * row_bytes + 1) as u64
209                && cause.used == (2 * row_bytes + 2) as u64
210                && cause.budget == bytes.get() as u64));
211        assert_eq!(output.rows.capacity(), 2);
212        assert_eq!(output.rows().len(), 2);
213        assert_eq!(ledger.snapshot().used_bytes, before.used_bytes);
214        assert_eq!(ledger.snapshot().peak_bytes, before.peak_bytes);
215        assert_eq!(output.validate(&account, budget).unwrap_err(), error);
216        assert_eq!(output.push("d", None, 0.0).unwrap_err(), error);
217        drop(output);
218        assert_eq!(ledger.snapshot().used_bytes, 0);
219    }
220
221    #[test]
222    fn growing_text_result_clamps_capacity_and_refuses_an_incomplete_prefix() {
223        let bytes = NonZeroUsize::new(1024).unwrap();
224        for capacity in 0..=3 {
225            let ledger = QueryMemoryLedger::new(bytes);
226            let account = ledger.account(QueryMemoryClass::ExternalRead, "growing text", bytes);
227            let budget = ExternalReadResultBudget {
228                max_rows: 3,
229                max_memory_bytes: bytes,
230            };
231            let mut output =
232                TextSeedExecutionOutput::with_initial_row_capacity(&account, budget, capacity)
233                    .unwrap();
234            for id in ["a", "b", "c"] {
235                output.push(id, None, 1.0).unwrap();
236            }
237            assert_eq!(output.rows.capacity(), 3);
238            assert_eq!(
239                output
240                    .rows()
241                    .iter()
242                    .map(|row| row.id.as_str())
243                    .collect::<Vec<_>>(),
244                ["a", "b", "c"]
245            );
246            let before = ledger.snapshot().used_bytes;
247            let error = output.push("d", None, 1.0).unwrap_err();
248            assert!(matches!(&error, HawDBError::ReadBudgetExceeded(cause)
249                if cause.resource == hawdb_core::ReadBudgetResource::Rows && cause.limit == 3));
250            assert_eq!(output.rows().len(), 3);
251            assert_eq!(output.rows.capacity(), 3);
252            assert_eq!(ledger.snapshot().used_bytes, before);
253            assert_eq!(output.validate(&account, budget).unwrap_err(), error);
254            drop(output);
255            assert_eq!(ledger.snapshot().used_bytes, 0);
256        }
257    }
258}
259
260#[derive(Debug)]
261pub struct TextSeedExecutionRow {
262    pub id: String,
263    pub external_id: Option<String>,
264    pub score: f64,
265}
266
267/// Rows admitted before copying and retained in the supplied result account.
268/// The caller can borrow rows, but cannot detach their owning reservation.
269#[derive(Debug)]
270pub struct TextSeedExecutionOutput {
271    rows: Vec<TextSeedExecutionRow>,
272    budget: ExternalReadResultBudget,
273    prepaid_rows: bool,
274    exhausted: bool,
275    source_report: Option<crate::TextSeedSourceReport>,
276    failure: Option<HawDBError>,
277    reservation: QueryMemoryLease,
278}
279
280impl TextSeedExecutionOutput {
281    pub fn new(account: &QueryMemoryAccount, budget: ExternalReadResultBudget) -> Result<Self> {
282        Self::with_initial_row_capacity(account, budget, budget.max_rows)
283    }
284
285    /// Construct a complete-domain builder without prepaying an unused ceiling.
286    /// The provider calls finish_complete only after exhausting its source.
287    pub fn new_complete(
288        account: &QueryMemoryAccount,
289        budget: ExternalReadResultBudget,
290    ) -> Result<Self> {
291        Self::with_initial_row_capacity(account, budget, 0)
292    }
293
294    pub fn finish_complete(mut self) -> Result<Self> {
295        if let Some(error) = &self.failure {
296            return Err(error.clone());
297        }
298        self.exhausted = true;
299        Ok(self)
300    }
301
302    pub fn is_complete(&self) -> bool {
303        self.exhausted
304    }
305
306    pub fn with_source_report(mut self, report: crate::TextSeedSourceReport) -> Self {
307        self.source_report = Some(report);
308        self
309    }
310
311    pub fn source_report(&self) -> Option<crate::TextSeedSourceReport> {
312        self.source_report
313    }
314
315    pub fn retained_bytes(&self) -> usize {
316        self.reservation.bytes()
317    }
318
319    fn with_initial_row_capacity(
320        account: &QueryMemoryAccount,
321        budget: ExternalReadResultBudget,
322        capacity: usize,
323    ) -> Result<Self> {
324        if capacity > budget.max_rows {
325            return Err(HawDBError::Execution(
326                "text seed initial capacity exceeds its row budget".into(),
327            ));
328        }
329        let bytes = capacity
330            .checked_mul(size_of::<TextSeedExecutionRow>())
331            .ok_or_else(|| {
332                HawDBError::Execution("text seed result storage size overflow".into())
333            })?;
334        if bytes > budget.max_memory_bytes.get() {
335            return Err(HawDBError::Execution(
336                "text seed result row storage exceeds its memory budget".into(),
337            ));
338        }
339        let reservation = account.reserve(bytes)?;
340        let rows = Vec::with_capacity(capacity);
341        Ok(Self {
342            rows,
343            budget,
344            prepaid_rows: capacity == budget.max_rows,
345            exhausted: false,
346            source_report: None,
347            failure: None,
348            reservation,
349        })
350    }
351
352    pub fn push(&mut self, id: &str, external_id: Option<&str>, score: f64) -> Result<()> {
353        if let Some(error) = &self.failure {
354            return Err(error.clone());
355        }
356        let result = (|| {
357            if self.exhausted {
358                return Err(HawDBError::Execution(
359                    "text seed source was already exhausted".into(),
360                ));
361            }
362            if !score.is_finite() {
363                return Err(HawDBError::Execution(
364                    "text seed score must be finite".into(),
365                ));
366            }
367            if self.rows.len() >= self.budget.max_rows {
368                return Err(HawDBError::read_budget_exceeded(
369                    hawdb_core::ReadBudgetResource::Rows,
370                    self.budget.max_rows,
371                    "text seed result exceeds its row budget",
372                ));
373            }
374            let bytes = id
375                .len()
376                .checked_add(external_id.map_or(0, str::len))
377                .ok_or_else(|| {
378                    HawDBError::Execution("text seed result payload size overflow".into())
379                })?;
380            let grow_to = (self.rows.len() == self.rows.capacity()).then(|| {
381                self.rows
382                    .capacity()
383                    .saturating_mul(2)
384                    .max(1)
385                    .min(self.budget.max_rows)
386            });
387            let replacement_bytes = grow_to
388                .map(|capacity| capacity.checked_mul(size_of::<TextSeedExecutionRow>()))
389                .unwrap_or(Some(0))
390                .ok_or_else(|| {
391                    HawDBError::Execution("text seed result storage size overflow".into())
392                })?;
393            let requested = bytes.checked_add(replacement_bytes).ok_or_else(|| {
394                HawDBError::Execution("text seed result payload size overflow".into())
395            })?;
396            let next = self
397                .reservation
398                .bytes()
399                .checked_add(requested)
400                .ok_or_else(|| {
401                    HawDBError::Execution("text seed result payload size overflow".into())
402                })?;
403            if next > self.budget.max_memory_bytes.get() {
404                return Err(if !self.prepaid_rows {
405                    self.reservation.operator_budget_error(
406                        requested,
407                        self.reservation.bytes(),
408                        self.budget.max_memory_bytes.get(),
409                        "ExternalReadResultBudget.max_memory_bytes",
410                        "text seed result growth exceeds its memory budget".into(),
411                    )
412                } else {
413                    HawDBError::Execution(
414                        "text seed result payload exceeds its memory budget".into(),
415                    )
416                });
417            }
418            self.reservation.grow(requested)?;
419            if let Some(capacity) = grow_to {
420                // Admit old and new storage together before reallocating; move
421                // payload owners without copying and release the old slots last.
422                let old_bytes = self.rows.capacity() * size_of::<TextSeedExecutionRow>();
423                let mut rows = Vec::with_capacity(capacity);
424                rows.append(&mut self.rows);
425                self.rows = rows;
426                self.reservation.shrink(old_bytes);
427            }
428            self.rows.push(TextSeedExecutionRow {
429                id: id.into(),
430                external_id: external_id.map(str::to_owned),
431                score,
432            });
433            Ok(())
434        })();
435        if let Err(error) = &result {
436            self.failure = Some(error.clone());
437        }
438        result
439    }
440
441    pub fn rows(&self) -> &[TextSeedExecutionRow] {
442        &self.rows
443    }
444
445    pub(crate) fn validate(
446        &self,
447        account: &QueryMemoryAccount,
448        budget: ExternalReadResultBudget,
449    ) -> Result<()> {
450        if !self.reservation.belongs_to(account) {
451            return Err(HawDBError::Execution(
452                "text seed output did not retain the admitted result account and budget".into(),
453            ));
454        }
455        if let Some(error) = &self.failure {
456            return Err(error.clone());
457        }
458        if self.rows.len() > budget.max_rows {
459            return Err(HawDBError::read_budget_exceeded(
460                hawdb_core::ReadBudgetResource::Rows,
461                budget.max_rows,
462                "text seed result exceeds its validated row budget",
463            ));
464        }
465        if self.reservation.bytes() > budget.max_memory_bytes.get() {
466            return Err(HawDBError::Execution(
467                "text seed output did not retain the admitted result account and budget".into(),
468            ));
469        }
470        Ok(())
471    }
472}
hawdb_executor/external/text.rs

hawdb_executor/external/
text.rs
