Skip to main content

hawdb_executor/
external.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
15//! Host-provided physical read operators.
16
17#[doc(hidden)]
18pub mod seed;
19
20use crate::VectorExecutionReport;
21use hawdb_core::{HawDBError, Result, RuntimeTaskContext};
22use hawdb_plan_cypher::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_cypher::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}