Skip to main content

uqa_operators/vector/
diagnostics.rs

1//
2// Unified Query Algebra
3//
4// Copyright (c) 2023-2026 Cognica, Inc.
5//
6
7//! Optional observation of the same primitive that supplies an operator's postings.
8
9use uqa_core::{vector_execution::VectorSearchOperation, PostingList};
10use uqa_storage::{
11    read_control::StorageReadControl, vector_index::VectorQueryResult, StorageBackendResult,
12    VectorIndex,
13};
14
15/// The owner binds the relation and retains the original invocation allowance.
16pub 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}