vibespatial.spatial.spatial_index_knn_device¶

Device-side k-nearest-neighbor spatial query.

Provides spatial_index_knn_device, a GPU-accelerated k-NN spatial query that replaces the CPU STRtree nearest path in sjoin_nearest.

Pipeline:

Admit a capacity-bounded candidate and retained-top-k workspace
Expand unresolved query bounds with bounded query/target streaming
Compute exact/refined distances only for each newly entered radius shell
Merge deterministic per-query top-k rows with CCCL segmented sorts
Output device-resident (query_idx, target_idx, distance) triples

ADR-0002: METRIC precision plan with selective fp64 ranking refinement. ADR-0033 tier classification:

Tier 2 (CuPy): element-wise distance filtering, gather/scatter

Tier 3a (CCCL): segmented_sort (per-query ranking), exclusive_sum, compact_indices

Attributes¶

`cp`
`logger`

Classes¶

`DeviceKnnResult`	Device-resident k-NN query result.
`DeviceKnnTelemetry`	Host-known control telemetry for one bounded fixed-k execution.

Functions¶

spatial_index_knn_device(→ DeviceKnnResult | None)

Run fixed-k nearest with operation-local allocation telemetry.

Module Contents¶

vibespatial.spatial.spatial_index_knn_device.cp = None¶

vibespatial.spatial.spatial_index_knn_device.logger¶

class vibespatial.spatial.spatial_index_knn_device.DeviceKnnResult¶

Device-resident k-NN query result.

All arrays are CuPy device arrays to avoid D->H transfers when the result feeds directly into the next GPU pipeline stage (e.g., sjoin attribute assembly).

Attributes¶

d_query_idxdevice int32 array: Query geometry indices (one per result pair).
d_target_idxdevice int32 array: Target geometry indices (one per result pair).
d_distancesdevice float64 array: Exact distances for each (query, target) pair.
total_pairsint: Number of result pairs.
kint: Requested k value.

d_query_idx: Any¶

d_target_idx: Any¶

d_distances: Any¶

total_pairs: int¶

k: int¶

telemetry: DeviceKnnTelemetry | None = None¶

to_host() → tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]¶: Copy result to host as numpy arrays.

class vibespatial.spatial.spatial_index_knn_device.DeviceKnnTelemetry¶

Host-known control telemetry for one bounded fixed-k execution.

query_tile_rows: int¶

target_tile_rows: int¶

pair_capacity: int¶

admitted_workspace_bytes: int¶

radius_iterations: int = 0¶

candidate_pairs: int = 0¶

max_candidate_pairs: int = 0¶

query_tiles: int = 0¶

target_stream_tiles: int = 0¶

scalar_fences: int = 0¶

peak_device_bytes: int | None = None¶

allocation_count: int | None = None¶

d2h_count: int = 0¶

d2h_bytes: int = 0¶

materialization_count: int = 0¶

detail() → str¶

vibespatial.spatial.spatial_index_knn_device.spatial_index_knn_device(query_owned: vibespatial.geometry.owned.OwnedGeometryArray, tree_owned: vibespatial.geometry.owned.OwnedGeometryArray, query_bounds: Any, tree_bounds: Any, *, native_spatial_index: Any | None = None, k: int = 1, max_distance: float | None = None, exclusive: bool = False, return_all: bool = False, precision: vibespatial.runtime.precision.PrecisionMode | str = PrecisionMode.AUTO) → DeviceKnnResult | None¶: Run fixed-k nearest with operation-local allocation telemetry.