added rust radius utility

Author: pranav-walimbe

bench/ops/README.md

@@ -9,15 +9,16 @@ Measures the 9 `CostFactors` ns/op constants used by the query planner's cost mo
   `shapely.box` (polygons), then wraps it in a `SpatialFrame`/`Engine`.
 - Build cost times only `engine.build_index()` on a fresh engine. Probe cost times the repeated queries against a pre-built index.
 - The constant is `time / workload_term`, where the term is read directly off the cost
-  model formula in `cost.rs` (e.g. `Q * N` for `scan_ns_per_item`).
+  model formula in `cost.rs` (e.g. `Q * N` for `knn_scan_ns_per_item`).
 - This ratio is computed at multiple dataset sizes and we return the median as the suggested value. 
 
 
 ## What it measures
 
 | Constant | Op timed | Dataset | Term |
 |---|---|---|---|
-| `scan_ns_per_item` | brute-force kNN, no index | uniform points | `Q * N` |
+| `knn_scan_ns_per_item` | brute-force kNN, no index | uniform points | `Q * N` |
+| `bbox_scan_ns_per_item` | brute-force range, no index | uniform points | `Q * N` |
 | `grid_build_ns_per_item` | build | uniform points | `N` |
 | `kdtree_build_ns_per_item` | build | clustered points | `N * log2(N)` |
 | `rtree_build_ns_per_item` | build | polygons | `N * log2(N)` |
@@ -46,19 +47,20 @@ Flags:
 Suggested CostFactors (copy into src/planner/calibration.rs):
 
 Brute Force
-    scan_ns_per_item:            100.80,
+    knn_scan_ns_per_item:        3.07,
+    bbox_scan_ns_per_item:       0.59,
 
 Points
-    grid_build_ns_per_item:      84.74,
-    kdtree_build_ns_per_item:    4.36,
-    grid_range_ns:               211.38,
-    kdtree_range_ns:             81.19,
-    kdtree_knn_ns:               148.71,
+    grid_build_ns_per_item:      10.33,
+    kdtree_build_ns_per_item:    0.47,
+    grid_range_ns:               34.40,
+    kdtree_range_ns:             17.56,
+    kdtree_knn_ns:               24.92,
 
 Polygons
-    rtree_build_ns_per_item:     74.17,
-    rtree_range_ns:              176.16,
-    rtree_knn_ns:                1299.75,
+    rtree_build_ns_per_item:     60.08,
+    rtree_range_ns:              24.69,
+    rtree_knn_ns:                132.17,
 
-elapsed: 20.3 s   peak RSS: 268.9 MiB
+elapsed: 2.9 s   peak RSS: 281.2 MiB
 ```

bench/ops/__main__.py

@@ -39,7 +39,7 @@
 # Default destination for the written report
 _OUTPUT_PATH = Path(__file__).resolve().parents[2] / "assets" / "ops.txt"
 
-_SCAN_FIELDS = ["scan_ns_per_item"]
+_SCAN_FIELDS = ["knn_scan_ns_per_item", "bbox_scan_ns_per_item"]
 _POINT_FIELDS = [
     "grid_build_ns_per_item",
     "kdtree_build_ns_per_item",
@@ -135,6 +135,14 @@ def _scan_probe(
     return t, _Q * n
 
 
+def _bbox_scan_probe(n: int, runs: int, seed: int) -> tuple[float, float]:
+    # Time one no-index range scan over an empty box, isolating the per-item box test in a single call
+    sf = _point_frame(generate_points(n, seed), "none")
+    empty_box = (2.0, 2.0, 3.0, 3.0)
+    t = time_min(lambda: sf.engine.range_query(*empty_box), runs)
+    return t, n
+
+
 def _range_probe(
     n: int,
     runs: int,
@@ -198,9 +206,12 @@ def run(runs: int, seed: int) -> None:
     boxes = _query_boxes(_Q, seed, _BOX_SIDE)
 
     fits: dict[str, float] = {
-        "scan_ns_per_item": _measure_median(
+        "knn_scan_ns_per_item": _measure_median(
             lambda n: _scan_probe(n, runs, seed, qxs, qys), _SCAN_SIZES
         ),
+        "bbox_scan_ns_per_item": _measure_median(
+            lambda n: _bbox_scan_probe(n, runs, seed), _SCAN_SIZES
+        ),
         "grid_build_ns_per_item": _measure_median(
             lambda n: _build_probe(n, runs, seed, clustered=False), _POINT_SIZES
         ),

bench/spatial_bench/queries/q01.py

@@ -1,7 +1,7 @@
 """Q1: Trips starting within ~50km (0.45 degrees) of the Sedona city center.
 
-PyCanopy: a within-distance join of all trip pickup points against the single
-center point.
+PyCanopy: a single-point radius filter of all trip pickup points against the
+Sedona center.
 """
 
 from __future__ import annotations
@@ -21,12 +21,11 @@
 def pycanopy(tables) -> pl.DataFrame:
     trip = tables.table("trip", ["t_tripkey", "t_pickuploc", "t_pickuptime"])
     sf = tables.point_frame(trip, "t_pickuploc")
-    center_df = pl.DataFrame({"cx": [CENTER[0]], "cy": [CENTER[1]]})
-    joined = sf.lazy().within_distance_join(center_df, "cx", "cy", distance=RADIUS).collect()
+    near = sf.radius_query(CENTER[0], CENTER[1], RADIUS)
     return (
-        joined.with_columns(
+        near.with_columns(
             distance_to_center=(
-                (pl.col("_x") - pl.col("cx")) ** 2 + (pl.col("_y") - pl.col("cy")) ** 2
+                (pl.col("_x") - CENTER[0]) ** 2 + (pl.col("_y") - CENTER[1]) ** 2
             ).sqrt()
         )
         .select(

python/pycanopy/engine.py

@@ -774,6 +774,19 @@ def polygon_pairs_intersection_area(
             np.ascontiguousarray(j_idx, dtype=np.uint64),
         )
 
+    def radius_query(self, cx: float, cy: float, distance: float) -> np.ndarray:
+        """Return indices of engine points within `distance` of the center (cx, cy).
+
+        Args:
+            cx: Center x coordinate.
+            cy: Center y coordinate.
+            distance: Maximum Euclidean distance for a match.
+
+        Returns:
+            uint64 array of matching point indices.
+        """
+        return self._core.radius_query(cx, cy, distance)
+
     def points_within_distance_of_polygon(self, polygon, distance: float) -> np.ndarray:
         """Return indices of engine points within `distance` of a query polygon.
 

python/pycanopy/executor.py

@@ -25,6 +25,7 @@
     ScalarNode,
     SelectNode,
     WithinDistanceJoinNode,
+    WithinDistanceOfPointNode,
     WithinJoinNode,
 )
 
@@ -366,6 +367,8 @@ def _emit_node(self, node, sf, lf: pl.LazyFrame, plugin_path: PluginPath) -> pl.
             return self._emit_polygon_knn_join(node, sf, lf)
         if isinstance(node, PointsWithinDistanceOfPolygonNode):
             return self._emit_points_within_distance_of_polygon(node, sf, lf)
+        if isinstance(node, WithinDistanceOfPointNode):
+            return self._emit_within_distance_of_point(node, sf, lf)
         raise TypeError(f"Unknown plan node type: {type(node)}")
 
     # filter nodes
@@ -419,6 +422,13 @@ def _emit_points_within_distance_of_polygon(
         indices = sf.engine.points_within_distance_of_polygon(node.polygon, node.distance)
         return self._filter_by_indices(lf, indices)
 
+    def _emit_within_distance_of_point(
+        self, node: WithinDistanceOfPointNode, sf, lf: pl.LazyFrame
+    ) -> pl.LazyFrame:
+        # One radius query resolves the indices, then the lf filters by original row
+        indices = sf.engine.radius_query(node.cx, node.cy, node.distance)
+        return self._filter_by_indices(lf, indices)
+
     def _filter_by_indices(self, lf: pl.LazyFrame, indices) -> pl.LazyFrame:
         # Filter the lf to the given original row positions (list or numpy array), empty to nothing
         if len(indices) == 0:

python/pycanopy/frame.py

@@ -273,6 +273,20 @@ def intersects_pairs(self, key_col: str | None = None) -> pl.DataFrame:
             },
         )
 
+    def radius_query(self, cx: float, cy: float, distance: float) -> pl.DataFrame:
+        """Return the rows whose point lies within `distance` of the center (cx, cy).
+
+        Args:
+            cx: Center x coordinate.
+            cy: Center y coordinate.
+            distance: Maximum Euclidean distance for a match.
+
+        Returns:
+            The subset of this frame's DataFrame within the radius.
+        """
+        idx = self._engine.radius_query(cx, cy, distance)
+        return self._df[pl.Series(idx.astype(np.uint32))]
+
     def points_within_distance_of_polygon(self, polygon, distance: float) -> pl.DataFrame:
         """Return the rows whose point lies within `distance` of a polygon boundary (zero inside).
 

python/pycanopy/lazy.py

@@ -28,6 +28,7 @@
     ScalarNode,
     SelectNode,
     WithinDistanceJoinNode,
+    WithinDistanceOfPointNode,
     WithinJoinNode,
 )
 from pycanopy.optimizer import SpatialOptimizer
@@ -82,6 +83,11 @@ def _fmt_node(node) -> str:
         return f"POLY_KNN_JOIN [k={node.k}, query_rows={len(node.query_df):,}, barrier]"
     if isinstance(node, PointsWithinDistanceOfPolygonNode):
         return f"POINTS_WITHIN_DIST_OF_POLY [dist={node.distance:.4g}]"
+    if isinstance(node, WithinDistanceOfPointNode):
+        return (
+            f"WITHIN_DIST_OF_POINT [center=({node.cx:.4g}, {node.cy:.4g}), "
+            f"dist={node.distance:.4g}, sel={node.selectivity:.3g}]"
+        )
     if isinstance(node, IntersectsSelfJoinNode):
         return "INTERSECTS_SELF_JOIN [pairs, barrier]"
     return f"UNKNOWN [{type(node).__name__}]"
@@ -380,6 +386,25 @@ def points_within_distance_of_polygon(self, polygon, distance: float) -> Spatial
             [*self._plan, PointsWithinDistanceOfPolygonNode(polygon, distance)],
         )
 
+    def within_distance_of_point(self, cx: float, cy: float, distance: float) -> SpatialLazyFrame:
+        """Keep points within `distance` of a single center (point dataset).
+
+        Distance is Euclidean. The result is a subset of this frame's rows like a spatial
+        filter.
+
+        Args:
+            cx: Center x coordinate.
+            cy: Center y coordinate.
+            distance: Maximum Euclidean distance for a row to be kept.
+
+        Returns:
+            New SpatialLazyFrame with the within-distance-of-point node appended.
+        """
+        return SpatialLazyFrame(
+            self._sf,
+            [*self._plan, WithinDistanceOfPointNode(cx, cy, distance)],
+        )
+
     def intersects_pairs(self) -> SpatialLazyFrame:
         """Find all intersecting polygon pairs with overlap area and IoU (polygon dataset).
 

python/pycanopy/nodes.py

@@ -159,6 +159,20 @@ class PointsWithinDistanceOfPolygonNode:
     selectivity: float = 1.0
 
 
+@dataclass
+class WithinDistanceOfPointNode:
+    """Spatial filter: keep points within `distance` of a single center (point dataset).
+
+    Distance is Euclidean. Returns a subset of the frame's rows, so it behaves like
+    range / contains.
+    """
+
+    cx: float
+    cy: float
+    distance: float
+    selectivity: float = 1.0
+
+
 @dataclass
 class IntersectsSelfJoinNode:
     """Spatial self-join: all intersecting polygon pairs with overlap area and IoU.

python/pycanopy/optimizer.py

@@ -13,6 +13,7 @@
 
 import dataclasses
 import json
+import math
 
 from pycanopy.nodes import (
     ContainsNode,
@@ -28,6 +29,7 @@
     ScalarNode,
     SelectNode,
     WithinDistanceJoinNode,
+    WithinDistanceOfPointNode,
     WithinJoinNode,
 )
 
@@ -160,6 +162,8 @@ def _assign_selectivity(self, plan: Plan, engine) -> Plan:
                 node = dataclasses.replace(node, selectivity=1.0 / max(n, 1))
             elif isinstance(node, KnnNode):
                 node = dataclasses.replace(node, selectivity=min(1.0, node.k / max(n, 1)))
+            elif isinstance(node, WithinDistanceOfPointNode):
+                node = dataclasses.replace(node, selectivity=self._disk_selectivity(node, extent))
             elif isinstance(node, ScalarNode):
                 node = dataclasses.replace(node, cost=_scalar_cost(node.expr))
             # join nodes have no selectivity field
@@ -181,6 +185,20 @@ def _range_selectivity(
         query_area = max(0.0, node.max_x - node.min_x) * max(0.0, node.max_y - node.min_y)
         return min(1.0, query_area / total_area)
 
+    def _disk_selectivity(
+        self,
+        node: WithinDistanceOfPointNode,
+        extent: tuple[float, float, float, float] | None,
+    ) -> float:
+        # Selectivity of a radius query as its disk area divided by the dataset extent area
+        if extent is None:
+            return 1.0
+        min_x, min_y, max_x, max_y = extent
+        total_area = (max_x - min_x) * (max_y - min_y)
+        if total_area <= 0.0:
+            return 1.0
+        return min(1.0, math.pi * node.distance * node.distance / total_area)
+
     def _cost_sort(self, plan: Plan) -> Plan:
         # Reorder nodes so cheaper operations run first. Joins and KNN are barriers, and within
         # a run scalars go first (Polars handles them cheaply) then spatials by ascending selectivity.

src/lib.rs

@@ -30,7 +30,8 @@ use query::{
     batch::{
         par_bbox_filter, par_contains, par_knn, par_knn_to_polygons, par_knn_to_polygons_sorted,
         par_knn_with_delta, par_points_within_distance_of_polygon, par_polygon_intersects_join,
-        par_within_distance, par_within_distance_flipped, par_within_distance_to_polygons,
+        par_radius, par_within_distance, par_within_distance_flipped,
+        par_within_distance_to_polygons,
     },
     geometry::{convex_hull_area, polygon_area, polygon_intersection_area},
     multipoly::{dedup_indices, dedup_self_pairs, polygon_parts_csr, sum_part_areas},
@@ -643,6 +644,85 @@ impl Engine {
         Ok(result)
     }
 
+    /// Engine point indices within `distance` (Euclidean) of the center (cx, cy)
+    fn radius_query<'py>(
+        &mut self,
+        py: Python<'py>,
+        cx: f64,
+        cy: f64,
+        distance: f64,
+    ) -> PyResult<Bound<'py, PyArray1<u64>>> {
+        if self.ring_offsets.is_some() {
+            return Err(PyValueError::new_err(
+                "radius_query requires a point dataset",
+            ));
+        }
+        let bbox = Rect::new(
+            coord! { x: cx - distance, y: cy - distance },
+            coord! { x: cx + distance, y: cy + distance },
+        );
+        // Skip histogram early-exit when delta is non-empty: delta points are not in the histogram
+        if self.delta_xs.is_empty() {
+            if let Some(hist) = &self.stats.histogram {
+                if !hist.has_any_in_bbox(&bbox) {
+                    let empty: Vec<u64> = Vec::new();
+                    return Ok(PyArray1::from_vec(py, empty));
+                }
+            }
+        }
+        let kind = self.plan_index(&Query::Range { bbox }, 1);
+        self.build_index_if_needed(kind);
+        let mut result = match kind {
+            IndexKind::BruteForce => par_radius(
+                self.brute.as_ref().unwrap(),
+                &self.xs,
+                &self.ys,
+                cx,
+                cy,
+                distance,
+            ),
+            IndexKind::RTree => par_radius(
+                self.rtree.as_ref().unwrap(),
+                &self.xs,
+                &self.ys,
+                cx,
+                cy,
+                distance,
+            ),
+            IndexKind::KdTree => par_radius(
+                self.kdtree.as_ref().unwrap(),
+                &self.xs,
+                &self.ys,
+                cx,
+                cy,
+                distance,
+            ),
+            IndexKind::Grid => par_radius(
+                self.grid.as_ref().unwrap(),
+                &self.xs,
+                &self.ys,
+                cx,
+                cy,
+                distance,
+            ),
+        };
+        // Refine delta points by exact distance and address them in the combined index space
+        if !self.delta_xs.is_empty() {
+            let d2 = distance * distance;
+            let n_main = self.xs.len();
+            for (i, (&x, &y)) in self.delta_xs.iter().zip(self.delta_ys.iter()).enumerate() {
+                let dx = x - cx;
+                let dy = y - cy;
+                if dx * dx + dy * dy <= d2 {
+                    result.push((n_main + i) as u64);
+                }
+            }
+            self.delta_query_cost += self.delta_xs.len() as u64;
+            self.maybe_flush_on_cost(kind);
+        }
+        Ok(PyArray1::from_vec(py, result))
+    }
+
     /// Intersect multiple bbox queries via two-pointer sorted merge on the hit arrays.
     /// O(K * |H|) over K predicates and per-predicate hit count |H| rather than O(K * N)
     /// bitmap ANDs. Short-circuits once the running intersection is empty.