sphereql_graphql/
query.rs

1use std::sync::Arc;
2
3use async_graphql::{Context, Object, Result};
4use tokio::sync::RwLock;
5
6use sphereql_core::{
7    SphericalPoint, angular_distance, chord_distance, euclidean_distance, great_circle_distance,
8    spherical_to_cartesian,
9};
10use sphereql_index::{SpatialIndex, SpatialItem, SpatialQueryResult};
11
12use crate::types::{
13    BandInput, ConeInput, DistanceMetric, NearestResultOutput, RegionInput, ShellInput,
14    SpatialQueryResultOutput, SphericalPointInput, SphericalPointOutput,
15};
16
17#[derive(Debug, Clone)]
18pub struct PointItem {
19    pub id: String,
20    pub position: SphericalPoint,
21}
22
23impl SpatialItem for PointItem {
24    type Id = String;
25    fn id(&self) -> &String {
26        &self.id
27    }
28    fn position(&self) -> &SphericalPoint {
29        &self.position
30    }
31}
32
33pub type PointIndex = Arc<RwLock<SpatialIndex<PointItem>>>;
34
35/// Take a read lock on the index, run the supplied `query` closure, and
36/// package the result into a `SpatialQueryResultOutput` with optional
37/// per-call truncation.
38///
39/// The four `within_*` resolvers (`cone`, `shell`, `band`, `region`)
40/// differed only in which `SpatialIndex::query_*` method they called —
41/// this helper folds the shared prologue (context lookup, lock, map,
42/// truncate, wrap) so each resolver is two lines.
43async fn run_spatial_query<F>(
44    ctx: &Context<'_>,
45    limit: Option<i32>,
46    query: F,
47) -> Result<SpatialQueryResultOutput>
48where
49    F: FnOnce(&SpatialIndex<PointItem>) -> SpatialQueryResult<PointItem>,
50{
51    let index = ctx
52        .data::<PointIndex>()
53        .map_err(|_| async_graphql::Error::new("SpatialIndex not found in context"))?;
54    let idx = index.read().await;
55    let result = query(&idx);
56
57    let mut items: Vec<SphericalPointOutput> = result
58        .items
59        .iter()
60        .map(|item| SphericalPointOutput::from(item.position()))
61        .collect();
62    if let Some(n) = limit {
63        items.truncate(n.max(0) as usize);
64    }
65
66    Ok(SpatialQueryResultOutput {
67        items,
68        total_scanned: result.total_scanned as i32,
69    })
70}
71
72pub struct SphericalQueryRoot;
73
74#[Object]
75impl SphericalQueryRoot {
76    async fn within_cone(
77        &self,
78        ctx: &Context<'_>,
79        cone: ConeInput,
80        limit: Option<i32>,
81    ) -> Result<SpatialQueryResultOutput> {
82        let core_cone = cone.to_core()?;
83        run_spatial_query(ctx, limit, |idx| idx.query_cone(&core_cone)).await
84    }
85
86    async fn within_shell(
87        &self,
88        ctx: &Context<'_>,
89        shell: ShellInput,
90        limit: Option<i32>,
91    ) -> Result<SpatialQueryResultOutput> {
92        let core_shell = shell.to_core()?;
93        run_spatial_query(ctx, limit, |idx| idx.query_shell(&core_shell)).await
94    }
95
96    async fn within_band(
97        &self,
98        ctx: &Context<'_>,
99        band: BandInput,
100        limit: Option<i32>,
101    ) -> Result<SpatialQueryResultOutput> {
102        let core_band = band.to_core()?;
103        run_spatial_query(ctx, limit, |idx| idx.query_band(&core_band)).await
104    }
105
106    async fn within_region(
107        &self,
108        ctx: &Context<'_>,
109        region: RegionInput,
110        limit: Option<i32>,
111    ) -> Result<SpatialQueryResultOutput> {
112        let core_region = region.to_core()?;
113        run_spatial_query(ctx, limit, |idx| idx.query_region(&core_region)).await
114    }
115
116    async fn nearest_to(
117        &self,
118        ctx: &Context<'_>,
119        point: SphericalPointInput,
120        k: i32,
121        max_distance: Option<f64>,
122    ) -> Result<Vec<NearestResultOutput>> {
123        let core_point = point.to_core()?;
124        let index = ctx
125            .data::<PointIndex>()
126            .map_err(|_| async_graphql::Error::new("SpatialIndex not found in context"))?;
127        let idx = index.read().await;
128        let results = idx.nearest(&core_point, k.max(0) as usize);
129
130        let results: Vec<NearestResultOutput> = results
131            .into_iter()
132            .filter(|r| match max_distance {
133                Some(max) => r.distance <= max,
134                None => true,
135            })
136            .map(|r| NearestResultOutput {
137                point: SphericalPointOutput::from(r.item.position()),
138                distance: r.distance,
139            })
140            .collect();
141
142        Ok(results)
143    }
144
145    async fn distance_between(
146        &self,
147        _ctx: &Context<'_>,
148        a: SphericalPointInput,
149        b: SphericalPointInput,
150        metric: Option<DistanceMetric>,
151        radius: Option<f64>,
152    ) -> Result<f64> {
153        let core_a = a.to_core()?;
154        let core_b = b.to_core()?;
155        let metric = metric.unwrap_or(DistanceMetric::Angular);
156
157        let distance = match metric {
158            DistanceMetric::Angular => angular_distance(&core_a, &core_b),
159            DistanceMetric::GreatCircle => {
160                great_circle_distance(&core_a, &core_b, radius.unwrap_or(1.0))
161            }
162            DistanceMetric::Chord => chord_distance(&core_a, &core_b),
163            DistanceMetric::Euclidean => {
164                let ca = spherical_to_cartesian(&core_a);
165                let cb = spherical_to_cartesian(&core_b);
166                euclidean_distance(&ca, &cb)
167            }
168        };
169
170        Ok(distance)
171    }
172}
sphereql_graphql/query.rs

sphereql_graphql/
query.rs
