Skip to main content

ringdb/
query.rs

1use std::time::Duration;
2
3/// A single match returned by a query.
4#[derive(Debug, Clone, Copy, PartialEq)]
5pub struct Hit {
6    /// Insertion-order ID of the matching vector (first inserted = 0).
7    pub id: u32,
8    /// Squared Euclidean distance to the query vector. Use `.sqrt()` to get
9    /// the actual distance.
10    pub dist_sq: f32,
11}
12
13/// A ring query: find all vectors whose Euclidean distance to `query`
14/// lies within `[d - lambda, d + lambda]`.
15///
16/// Internally, ringdb uses squared L2 distances to avoid computing square
17/// roots. The ring bounds become:
18///
19/// ```text
20/// lower_sq = max(0, d - lambda)²
21/// upper_sq = (d + lambda)²
22/// ```
23pub struct RingQuery<'a> {
24    /// The query vector. Must have length equal to `RingDb::dims()`.
25    pub query: &'a [f32],
26    /// Target distance (centre of the ring).
27    pub d: f32,
28    /// Half-width of the ring.
29    pub lambda: f32,
30}
31
32/// A range query: find all vectors whose Euclidean distance to `query`
33/// lies within `[d_min, d_max]`.
34pub struct RangeQuery<'a> {
35    /// The query vector. Must have length equal to `RingDb::dims()`.
36    pub query: &'a [f32],
37    /// Lower bound of the distance interval (inclusive). Must be ≥ 0.
38    pub d_min: f32,
39    /// Upper bound of the distance interval (inclusive). Must be ≥ `d_min`.
40    pub d_max: f32,
41}
42
43/// A disk query: find all vectors whose Euclidean distance to `query`
44/// is at most `d_max` (i.e. the full disk/ball of radius `d_max`).
45///
46/// This is equivalent to a [`RangeQuery`] with `d_min = 0`.
47pub struct DiskQuery<'a> {
48    /// The query vector. Must have length equal to `RingDb::dims()`.
49    pub query: &'a [f32],
50    /// Radius of the disk (inclusive upper bound on distance). Must be ≥ 0.
51    pub d_max: f32,
52}
53
54/// A disk-intersection query: find all vectors inside every disk.
55///
56/// Each disk is a [`DiskQuery`] containing a center vector and an inclusive
57/// radius. The query must contain at least one disk. Returned [`Hit::dist_sq`]
58/// values are measured against the first disk's query vector.
59pub struct DiskIntersectionQuery<'a> {
60    /// Disks that all returned vectors must satisfy.
61    pub disks: &'a [DiskQuery<'a>],
62}
63
64/// Result of a ring/range/disk query.
65pub struct QueryResult {
66    /// All matching vectors together with their squared distances.
67    pub hits: Vec<Hit>,
68    /// Name of the backend that executed the query (e.g. `"cpu"`, `"wgpu"`, `"cuda"`).
69    pub backend_used: &'static str,
70    /// Wall-clock time for the query (excluding dataset upload).
71    pub elapsed: Duration,
72}
73
74impl QueryResult {
75    /// Convenience: collect just the IDs from `hits` into a new `Vec<u32>`.
76    ///
77    /// Useful when calling `fetch_payloads` or `fetch_pods`, which take `&[u32]`.
78    pub fn ids(&self) -> Vec<u32> {
79        self.hits.iter().map(|h| h.id).collect()
80    }
81}