hawdb_executor/
external.rs1#[doc(hidden)]
18pub mod seed;
19
20use crate::VectorExecutionReport;
21use hawdb_core::{HawDBError, Result, RuntimeTaskContext};
22use hawdb_plan::VectorPhysicalPlan;
23use std::collections::BTreeMap;
24use std::mem::size_of;
25use std::num::NonZeroUsize;
26
27#[derive(Debug, Clone, Copy, PartialEq, Eq)]
28pub struct ExternalReadResultBudget {
29 pub max_rows: usize,
30 pub max_memory_bytes: NonZeroUsize,
31}
32
33#[derive(Debug, Clone, Copy)]
34pub struct ExternalReadResourceContract<'a> {
35 pub priority: u8,
36 pub max_parallelism: NonZeroUsize,
37 pub max_working_memory_bytes: NonZeroUsize,
38 pub result: ExternalReadResultBudget,
39 pub task_context: Option<&'a RuntimeTaskContext>,
40}
41
42impl ExternalReadResourceContract<'_> {
43 pub fn checkpoint(&self) -> Result<()> {
44 match self.task_context {
45 Some(task_context) => task_context.checkpoint().map_err(|reason| {
46 HawDBError::Execution(format!("external read task stopped: {reason}"))
47 }),
48 None => Ok(()),
49 }
50 }
51
52 pub fn reserved_memory_bytes(&self) -> usize {
53 self.max_working_memory_bytes
54 .get()
55 .saturating_add(self.result.max_memory_bytes.get())
56 }
57}
58
59pub struct VectorSeedExecutionRequest<'a> {
60 pub embedding: &'a [f32],
61 pub metadata_filters: &'a BTreeMap<String, String>,
62 pub vector_plan: &'a VectorPhysicalPlan,
63 pub resources: ExternalReadResourceContract<'a>,
64}
65
66pub struct VectorSeedExecutionRow {
67 pub id: String,
68 pub external_id: Option<String>,
69 pub score: f64,
70}
71
72pub struct VectorSeedExecutionOutput {
73 pub rows: Vec<VectorSeedExecutionRow>,
74 pub report: VectorExecutionReport,
75}
76
77impl VectorSeedExecutionOutput {
78 pub fn estimated_memory_bytes(&self) -> usize {
79 self.rows
80 .capacity()
81 .saturating_mul(size_of::<VectorSeedExecutionRow>())
82 .saturating_add(self.rows.iter().fold(0usize, |total, row| {
83 total.saturating_add(row.id.capacity()).saturating_add(
84 row.external_id
85 .as_ref()
86 .map_or(0, |external_id| external_id.capacity()),
87 )
88 }))
89 }
90
91 pub fn validate_result_budget(&self, budget: ExternalReadResultBudget) -> Result<()> {
92 if self.rows.len() > budget.max_rows {
93 return Err(HawDBError::Execution(format!(
94 "external vector read returned {} rows, exceeding result row budget {}",
95 self.rows.len(),
96 budget.max_rows
97 )));
98 }
99 let memory_bytes = self.estimated_memory_bytes();
100 if memory_bytes > budget.max_memory_bytes.get() {
101 return Err(HawDBError::Execution(format!(
102 "external vector read returned {memory_bytes} estimated bytes, exceeding result memory budget {}",
103 budget.max_memory_bytes
104 )));
105 }
106 Ok(())
107 }
108}
109
110pub trait ExternalReadOperator {
111 fn execute_vector_seed(
112 &mut self,
113 request: VectorSeedExecutionRequest<'_>,
114 ) -> Result<VectorSeedExecutionOutput>;
115}
116
117#[doc(hidden)]
118pub struct NoExternalReadOperator;
119
120impl ExternalReadOperator for NoExternalReadOperator {
121 fn execute_vector_seed(
122 &mut self,
123 _request: VectorSeedExecutionRequest<'_>,
124 ) -> Result<VectorSeedExecutionOutput> {
125 Err(HawDBError::Execution(
126 "vector search capability is unavailable without a search projection".to_string(),
127 ))
128 }
129}
130
131#[cfg(test)]
132mod tests {
133 use super::*;
134 use crate::{
135 VectorCompressionMode, VectorExecutionBackend, VectorFallbackReasonCode, VectorScoreSource,
136 };
137 use hawdb_core::RuntimeCancellationToken;
138 use hawdb_plan::VectorCandidateSource;
139
140 pub(super) fn empty_report() -> VectorExecutionReport {
141 VectorExecutionReport {
142 backend: VectorExecutionBackend::ScalarFlat,
143 compression_mode: VectorCompressionMode::Disabled,
144 candidate_source: VectorCandidateSource::Scalar,
145 backend_selection_reason: None,
146 estimated_raw_vector_bytes: None,
147 filter_selectivity_per_million: None,
148 candidate_score_source: VectorScoreSource::RawVector,
149 final_score_source: VectorScoreSource::RawVector,
150 generated_candidate_count: 0,
151 descriptor_pruned_count: 0,
152 scalar_filtered_count: 0,
153 residual_filtered_count: 0,
154 candidate_scan_rounds: 0,
155 reranked_candidate_count: 0,
156 returned_count: 0,
157 raw_vector_bytes_read: 0,
158 candidate_scan_metrics: None,
159 index_covered_document_count: None,
160 index_candidate_document_count: None,
161 index_coverage_complete: None,
162 fallback_reason_codes: Vec::<VectorFallbackReasonCode>::new(),
163 }
164 }
165
166 #[test]
167 fn external_vector_output_fails_closed_on_row_and_memory_budgets() {
168 let output = VectorSeedExecutionOutput {
169 rows: vec![VectorSeedExecutionRow {
170 id: "memory-1".to_string(),
171 external_id: Some("external-1".to_string()),
172 score: 1.0,
173 }],
174 report: empty_report(),
175 };
176
177 assert!(output
178 .validate_result_budget(ExternalReadResultBudget {
179 max_rows: 0,
180 max_memory_bytes: NonZeroUsize::new(1024).unwrap(),
181 })
182 .is_err());
183 assert!(output
184 .validate_result_budget(ExternalReadResultBudget {
185 max_rows: 1,
186 max_memory_bytes: NonZeroUsize::MIN,
187 })
188 .is_err());
189 }
190
191 #[test]
192 fn external_read_contract_observes_cancellation() {
193 let cancellation = RuntimeCancellationToken::new();
194 let task_context = RuntimeTaskContext::without_deadline(cancellation.clone());
195 let contract = ExternalReadResourceContract {
196 priority: 128,
197 max_parallelism: NonZeroUsize::MIN,
198 max_working_memory_bytes: NonZeroUsize::MIN,
199 result: ExternalReadResultBudget {
200 max_rows: 1,
201 max_memory_bytes: NonZeroUsize::MIN,
202 },
203 task_context: Some(&task_context),
204 };
205
206 assert!(contract.checkpoint().is_ok());
207 assert!(cancellation.cancel());
208 assert!(contract.checkpoint().is_err());
209 }
210}