uqa_operators/vector/
diagnostics.rs1use uqa_core::{vector_execution::VectorSearchOperation, PostingList};
10use uqa_storage::{
11 read_control::StorageReadControl, vector_index::VectorQueryResult, StorageBackendResult,
12 VectorIndex,
13};
14
15pub trait VectorSearchObserver: Send + Sync {
17 fn control(&self) -> &StorageReadControl;
18 fn record(
19 &self,
20 field: &str,
21 operation: VectorSearchOperation,
22 result: &VectorQueryResult,
23 ) -> StorageBackendResult<()>;
24}
25
26pub fn search_knn(
27 index: &dyn VectorIndex,
28 field: &str,
29 query: &[f32],
30 k: usize,
31 observer: Option<&dyn VectorSearchObserver>,
32) -> StorageBackendResult<PostingList> {
33 let Some(observer) = observer else {
34 return index.search_knn(query, k);
35 };
36 let result = index.search_knn_with_statistics(query, k, Some(observer.control()))?;
37 observer.record(field, VectorSearchOperation::KNN { k }, &result)?;
38 Ok(result.postings)
39}
40
41pub fn search_threshold(
42 index: &dyn VectorIndex,
43 field: &str,
44 query: &[f32],
45 threshold: f32,
46 observer: Option<&dyn VectorSearchObserver>,
47) -> StorageBackendResult<PostingList> {
48 let Some(observer) = observer else {
49 return index.search_threshold(query, threshold);
50 };
51 let result =
52 index.search_threshold_with_statistics(query, threshold, Some(observer.control()))?;
53 observer.record(
54 field,
55 VectorSearchOperation::Threshold { threshold },
56 &result,
57 )?;
58 Ok(result.postings)
59}