From 59577a4247fa2a604a7395cff990ab9a3fa27037 Mon Sep 17 00:00:00 2001 From: Pranav Walimbe Date: Tue, 25 Aug 2026 12:10:28 -0700 Subject: [PATCH 1/4] perf: adopt decoded vectors into the engine without copying --- src/index/brute.rs | 11 ++++---- src/index/grid.rs | 9 +++--- src/index/kdtree.rs | 8 +++--- src/index/mod.rs | 7 ++--- src/index/rtree.rs | 4 +-- src/lib.rs | 59 ++++++++++++++++++++++----------------- tests/rust/index_tests.rs | 14 +++++----- 7 files changed, 60 insertions(+), 52 deletions(-) diff --git a/src/index/brute.rs b/src/index/brute.rs index b58134f..29c7276 100644 --- a/src/index/brute.rs +++ b/src/index/brute.rs @@ -3,6 +3,7 @@ use std::sync::Arc; use crate::index::{point_box_dist2, SpatialIndex}; +use crate::Shared; /// Linear scan index, used for small datasets or high-selectivity queries. /// @@ -14,14 +15,14 @@ pub struct BruteForce { /// Per-geometry bounding boxes. /// For point datasets all four are Arc::clone of the Engine's xs/ys (shared, zero cost). /// For polygon datasets these are new allocations derived from ring coords. - bbox_min_x: Arc<[f64]>, - bbox_min_y: Arc<[f64]>, - bbox_max_x: Arc<[f64]>, - bbox_max_y: Arc<[f64]>, + bbox_min_x: Shared, + bbox_min_y: Shared, + bbox_max_x: Shared, + bbox_max_y: Shared, } impl SpatialIndex for BruteForce { - fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { + fn build(xs: Shared, ys: Shared) -> Self { BruteForce { bbox_min_x: Arc::clone(&xs), bbox_min_y: Arc::clone(&ys), diff --git a/src/index/grid.rs b/src/index/grid.rs index ab0ba8b..9283096 100644 --- a/src/index/grid.rs +++ b/src/index/grid.rs @@ -1,10 +1,9 @@ //! Uniform grid index with CSR cell storage for large, uniformly distributed datasets. -use std::sync::Arc; - use rayon::prelude::*; use crate::index::SpatialIndex; +use crate::Shared; /// Uniform grid index with CSR (compressed sparse row) cell storage. /// Best for large datasets with uniform spatial distribution. @@ -15,8 +14,8 @@ pub struct UniformGrid { /// cell_offsets[i]..cell_offsets[i+1] is the slice of indices in cell i cell_offsets: Vec, indices: Vec, - xs: Arc<[f64]>, - ys: Arc<[f64]>, + xs: Shared, + ys: Shared, min_x: f64, min_y: f64, cell_w: f64, @@ -55,7 +54,7 @@ impl UniformGrid { } impl SpatialIndex for UniformGrid { - fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { + fn build(xs: Shared, ys: Shared) -> Self { let n = xs.len(); let (min_x, min_y, max_x, max_y) = xs.iter().zip(ys.iter()).fold( diff --git a/src/index/kdtree.rs b/src/index/kdtree.rs index afa42ed..ebe9b5a 100644 --- a/src/index/kdtree.rs +++ b/src/index/kdtree.rs @@ -1,12 +1,12 @@ //! Packed immutable KD-tree index optimised for point datasets. use std::f64::consts::PI; -use std::sync::Arc; use geo_index::kdtree::{KDTree, KDTreeBuilder, KDTreeIndex}; use crate::index::SpatialIndex; use crate::stats::types::SpatialHistogram; +use crate::Shared; /// Packed immutable KD-tree backed by geo-index, optimised for point datasets. /// @@ -15,14 +15,14 @@ use crate::stats::types::SpatialHistogram; /// traversal). The xs/ys Arcs are kept for the kNN distance refinement step. pub struct PackedKdTree { tree: KDTree, - xs: Arc<[f64]>, - ys: Arc<[f64]>, + xs: Shared, + ys: Shared, extent_area: f64, histogram: Option, } impl SpatialIndex for PackedKdTree { - fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { + fn build(xs: Shared, ys: Shared) -> Self { let n = xs.len(); let mut builder = KDTreeBuilder::::new(n as u32); diff --git a/src/index/mod.rs b/src/index/mod.rs index 7f5372b..bae45dc 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -5,14 +5,13 @@ pub mod grid; pub mod kdtree; pub mod rtree; -use std::sync::Arc; +use crate::Shared; /// Common interface for all spatial index backends. -/// Coordinates are passed as Arc<[f64]> so indexes can share Engine's allocation -/// without copying, since storing an Arc<[f64]> is an atomic refcount bump, not a memcpy. +/// Coordinates arrive as Shared so an index shares the Engine allocation by refcount. pub trait SpatialIndex: Send + Sync { /// Build an index over the given coordinate arrays - fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self + fn build(xs: Shared, ys: Shared) -> Self where Self: Sized; /// Indices of the k nearest entries to (qx, qy), sorted nearest-first. diff --git a/src/index/rtree.rs b/src/index/rtree.rs index ce5543b..4baab4e 100644 --- a/src/index/rtree.rs +++ b/src/index/rtree.rs @@ -3,7 +3,6 @@ use std::cell::RefCell; use std::cmp::{Ordering, Reverse}; use std::collections::{BinaryHeap, VecDeque}; -use std::sync::Arc; use rayon::prelude::*; @@ -11,6 +10,7 @@ use geo_index::rtree::sort::HilbertSort; use geo_index::rtree::{RTree, RTreeBuilder, RTreeIndex}; use crate::index::{point_box_dist2, SpatialIndex}; +use crate::Shared; /// One entry in the best-first traversal queue, ordered by squared distance to the node's box. /// @@ -96,7 +96,7 @@ pub struct PackedRTree { } impl SpatialIndex for PackedRTree { - fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { + fn build(xs: Shared, ys: Shared) -> Self { let n = xs.len() as u32; let mut builder = RTreeBuilder::::new(n); for (&x, &y) in xs.iter().zip(ys.iter()) { diff --git a/src/lib.rs b/src/lib.rs index 2208538..94eec75 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,12 +52,15 @@ use query::{ use rayon::prelude::*; use stats::{collector, types::GeometryKind}; +// Shared array storage that adopts a decoded Vec without copying +pub(crate) type Shared = Arc>; + #[pyclass] struct Engine { - xs: Arc<[f64]>, - ys: Arc<[f64]>, - ring_offsets: Option>, // coord range per ring, None for point datasets - poly_offsets: Option>, // ring range per polygon, first ring exterior, rest holes + xs: Shared, + ys: Shared, + ring_offsets: Option>, // coord range per ring, None for point datasets + poly_offsets: Option>, // ring range per polygon, first ring exterior, rest holes stats: stats::types::DatasetStats, brute: Option, rtree: Option, @@ -69,7 +72,7 @@ struct Engine { index_mode: IndexMode, // index selection policy: Eager (default) / None / Auto / Explicit cost_factors: CostFactors, prepared_polys: Option, // sub-linear PIP edge index, built lazily - part_poly: Option>, // logical polygon per index part, None if no MultiPolygons + part_poly: Option>, // logical polygon per index part, None if no MultiPolygons n_polygons: usize, // count of logical polygons, equals part count when part_poly is None metric: DistanceMetric, // how threshold distances are measured: Planar (default) / Haversine metrics: EngineMetrics, // fixed, always-on counters drained explicitly through Python @@ -254,7 +257,7 @@ impl PyPolygonBuilder { } impl Engine { - fn new_points(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Engine { + fn new_points(xs: Shared, ys: Shared) -> Engine { let stats_started = Instant::now(); let stats = collector::collect_points(&xs, &ys); let statistics_ns = elapsed_ns(stats_started); @@ -282,11 +285,11 @@ impl Engine { } fn new_polygons( - xs: Arc<[f64]>, - ys: Arc<[f64]>, - ring_offsets: Arc<[i64]>, - poly_offsets: Arc<[i64]>, - part_poly: Option>, + xs: Shared, + ys: Shared, + ring_offsets: Shared, + poly_offsets: Shared, + part_poly: Option>, n_polygons: usize, ) -> Engine { let stats_started = Instant::now(); @@ -568,7 +571,7 @@ impl Engine { #[pymethods] impl Engine { /// Construct from contiguous float64 numpy arrays of x and y coordinates. - /// Copies once into an Arc<[f64]> that all later index builds share via Arc::clone. + /// Copies the borrowed arrays once into a Shared that every index build then shares. #[staticmethod] fn from_points(xs: PyReadonlyArray1, ys: PyReadonlyArray1) -> PyResult { let xs_sl = xs @@ -581,7 +584,10 @@ impl Engine { return Err(PyValueError::new_err("xs and ys must have the same length")); } ensure_u32_indexable(xs_sl.len(), "point dataset")?; - Ok(Engine::new_points(xs_sl.into(), ys_sl.into())) + Ok(Engine::new_points( + Arc::new(xs_sl.to_vec()), + Arc::new(ys_sl.to_vec()), + )) } /// Construct from two-level polygon ring arrays (supports holes and MultiPolygons). @@ -645,7 +651,7 @@ impl Engine { } // part_poly maps parts to logical polygons. // None means each part is its own logical polygon (no MultiPolygons). - let (part_poly_arc, n_polygons): (Option>, usize) = match part_poly { + let (part_poly_arc, n_polygons): (Option>, usize) = match part_poly { Some(pp) => { let pp_sl = pp.as_slice().map_err(|_| { PyValueError::new_err("part_poly must be a contiguous int64 array") @@ -657,15 +663,18 @@ impl Engine { } let n_logical = pp_sl.iter().copied().max().map_or(0, |m| m as usize + 1); ensure_u32_indexable(n_logical, "logical polygon dataset")?; - (Some(pp_sl.iter().map(|&v| v as u32).collect()), n_logical) + ( + Some(Arc::new(pp_sl.iter().map(|&v| v as u32).collect())), + n_logical, + ) } None => (None, n_polys), }; Ok(Engine::new_polygons( - xs_sl.into(), - ys_sl.into(), - ring_sl.into(), - poly_sl.into(), + Arc::new(xs_sl.to_vec()), + Arc::new(ys_sl.to_vec()), + Arc::new(ring_sl.to_vec()), + Arc::new(poly_sl.to_vec()), part_poly_arc, n_polygons, )) @@ -1621,7 +1630,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let prepared = self.prepared_polys.as_ref(); - let part_poly = self.part_poly.as_deref(); + let part_poly = self.part_poly.as_deref().map(Vec::as_slice); let flat = match kind { IndexKind::BruteForce => par_contains( self.brute.as_ref().unwrap(), @@ -1727,7 +1736,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let prepared = self.prepared_polys.as_ref(); - let part_poly = self.part_poly.as_deref(); + let part_poly = self.part_poly.as_deref().map(Vec::as_slice); let state = match kind { IndexKind::BruteForce => par_contains_aggregate( self.brute.as_ref().unwrap(), @@ -1817,7 +1826,7 @@ impl Engine { let build_ns = self.build_index_if_needed(kind); let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); - let part_poly = self.part_poly.as_deref(); + let part_poly = self.part_poly.as_deref().map(Vec::as_slice); let flat = match kind { IndexKind::BruteForce => par_within_distance_to_polygons( self.brute.as_ref().unwrap(), @@ -1919,7 +1928,7 @@ impl Engine { let build_ns = self.build_index_if_needed(kind); let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); - let part_poly = self.part_poly.as_deref(); + let part_poly = self.part_poly.as_deref().map(Vec::as_slice); let state = match kind { IndexKind::BruteForce => par_within_distance_to_polygons_aggregate( self.brute.as_ref().unwrap(), @@ -2011,7 +2020,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let n_parts = poly_off.len().saturating_sub(1); - let part_poly = self.part_poly.as_deref(); + let part_poly = self.part_poly.as_deref().map(Vec::as_slice); let (idx, dist) = match kind { IndexKind::BruteForce => par_knn_to_polygons( self.brute.as_ref().unwrap(), @@ -2098,7 +2107,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let n_parts = poly_off.len().saturating_sub(1); - let part_poly = self.part_poly.as_deref(); + let part_poly = self.part_poly.as_deref().map(Vec::as_slice); let (q_idx, t_idx, dist) = match kind { IndexKind::BruteForce => par_knn_to_polygons_sorted( self.brute.as_ref().unwrap(), diff --git a/tests/rust/index_tests.rs b/tests/rust/index_tests.rs index 64ccd71..4e64cfe 100644 --- a/tests/rust/index_tests.rs +++ b/tests/rust/index_tests.rs @@ -9,10 +9,10 @@ use pycanopy::index::{ }; use pycanopy::query::range::{query_contains_polygons, query_range_polygons}; -fn five_point_grid() -> (Arc<[f64]>, Arc<[f64]>) { +fn five_point_grid() -> (Arc>, Arc>) { ( - Arc::from([0.0f64, 1.0, 2.0, 0.0, 1.0].as_slice()), - Arc::from([0.0f64, 0.0, 0.0, 1.0, 1.0].as_slice()), + Arc::new(vec![0.0f64, 1.0, 2.0, 0.0, 1.0]), + Arc::new(vec![0.0f64, 0.0, 0.0, 1.0, 1.0]), ) } @@ -192,8 +192,8 @@ fn range_empty_all_implementations_agree() { #[test] fn nearest_k5_on_larger_dataset_all_agree() { - let xs: Arc<[f64]> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); - let ys: Arc<[f64]> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); + let xs: Arc> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); + let ys: Arc> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); let oracle = as_set(BruteForce::build(Arc::clone(&xs), Arc::clone(&ys)).nearest(4.6, 3.2, 5)); assert_eq!( @@ -215,8 +215,8 @@ fn nearest_k5_on_larger_dataset_all_agree() { #[test] fn range_on_larger_dataset_all_agree() { - let xs: Arc<[f64]> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); - let ys: Arc<[f64]> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); + let xs: Arc> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); + let ys: Arc> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); let oracle = as_set(BruteForce::build(Arc::clone(&xs), Arc::clone(&ys)).range(2.0, 2.0, 5.0, 5.0)); From 96c2098e9d660ec645a718768ef3b149b569fcff Mon Sep 17 00:00:00 2001 From: Pranav Walimbe Date: Tue, 25 Aug 2026 12:19:25 -0700 Subject: [PATCH 2/4] perf: decode polygon WKB from arrow views without materialising bytes --- python/pycanopy/engine.py | 42 ++++++++ src/lib.rs | 68 +++++++++---- src/wkb.rs | 202 +++++++++++++++++++++++++++++++++----- 3 files changed, 266 insertions(+), 46 deletions(-) diff --git a/python/pycanopy/engine.py b/python/pycanopy/engine.py index aca557b..d47f59c 100644 --- a/python/pycanopy/engine.py +++ b/python/pycanopy/engine.py @@ -329,6 +329,41 @@ def _wkb_points_fast(arr: pa.Array) -> tuple[np.ndarray, np.ndarray] | None: ) +# Bytes per Arrow BinaryView descriptor +_VIEW_WIDTH = 16 + + +@lru_cache(maxsize=1) +def _newest_compat_level(): + # Polars compat level that keeps a Binary column in its view layout + import polars as pl # noqa: PLC0415 + + return pl.CompatLevel.newest() + + +def _wkb_view_buffers(column) -> tuple[np.ndarray, list[np.ndarray]] | None: + # Return a WKB column's Arrow view descriptors and data buffers without copying its bytes + to_arrow = getattr(column, "to_arrow", None) + is_view = getattr(pa.types, "is_binary_view", None) + if to_arrow is None or is_view is None: + return None + try: + array = to_arrow(compat_level=_newest_compat_level()) + except (TypeError, ValueError, AttributeError): + return None + if not isinstance(array, pa.Array) or not is_view(array.type): + return None + if array.null_count != 0: + return None + buffers = array.buffers() + if len(buffers) < 2 or buffers[1] is None: + return None + start = array.offset * _VIEW_WIDTH + views = np.frombuffer(buffers[1], dtype=np.uint8)[start : start + len(array) * _VIEW_WIDTH] + data = [np.frombuffer(buffer, dtype=np.uint8) for buffer in buffers[2:] if buffer is not None] + return views, data + + def _wkb_binary_buffers(column) -> tuple[np.ndarray, np.ndarray] | None: # Return zero-copy (data, offsets) numpy buffers of a WKB binary column, where data is # the concatenated value bytes and offsets the n+1 bounds, or None for null/non-binary. @@ -485,6 +520,13 @@ def from_wkb_polygons(cls, column) -> Engine: eng = cls.__new__(cls) eng._metrics_capture = None eng._metrics_capture_id = -1 + views = _wkb_view_buffers(column) + if views is not None: + try: + eng._core = _configure_core(_CoreEngine.from_wkb_polygon_views(*views)) + return _register_metrics_engine(eng) + except ValueError: + pass # unusual WKB variant -> contiguous buffers then shapely buffers = _wkb_binary_buffers(column) if buffers is not None: try: diff --git a/src/lib.rs b/src/lib.rs index 94eec75..7868908 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -52,6 +52,31 @@ use query::{ use rayon::prelude::*; use stats::{collector, types::GeometryKind}; +impl Engine { + // Build a polygon Engine from a decoded column and record what the decode cost + fn from_parsed_polygons(parsed: wkb::ParsedPolygons, wkb_decode_ns: u64) -> PyResult { + let n_polygons = match &parsed.part_poly { + Some(pp) => pp.iter().copied().max().map_or(0, |m| m as usize + 1), + None => parsed.poly_offsets.len().saturating_sub(1), + }; + ensure_u32_indexable( + parsed.poly_offsets.len().saturating_sub(1), + "polygon dataset", + )?; + ensure_u32_indexable(n_polygons, "logical polygon dataset")?; + let mut engine = Engine::new_polygons( + parsed.xs.into(), + parsed.ys.into(), + parsed.ring_offsets.into(), + parsed.poly_offsets.into(), + parsed.part_poly.map(Arc::new), + n_polygons, + ); + engine.metrics.wkb_decode_ns = wkb_decode_ns; + Ok(engine) + } +} + // Shared array storage that adopts a decoded Vec without copying pub(crate) type Shared = Arc>; @@ -697,26 +722,29 @@ impl Engine { .map_err(|_| PyValueError::new_err("offsets must be a contiguous int64 array"))?; let decode_started = Instant::now(); let parsed = wkb::parse_polygons(data_sl, off_sl).map_err(PyValueError::new_err)?; - let wkb_decode_ns = elapsed_ns(decode_started); - let n_polygons = match &parsed.part_poly { - Some(pp) => pp.iter().copied().max().map_or(0, |m| m as usize + 1), - None => parsed.poly_offsets.len().saturating_sub(1), - }; - ensure_u32_indexable( - parsed.poly_offsets.len().saturating_sub(1), - "polygon dataset", - )?; - ensure_u32_indexable(n_polygons, "logical polygon dataset")?; - let mut engine = Engine::new_polygons( - parsed.xs.into(), - parsed.ys.into(), - parsed.ring_offsets.into(), - parsed.poly_offsets.into(), - parsed.part_poly.map(Arc::from), - n_polygons, - ); - engine.metrics.wkb_decode_ns = wkb_decode_ns; - Ok(engine) + Engine::from_parsed_polygons(parsed, elapsed_ns(decode_started)) + } + + /// Construct from Arrow BinaryView descriptors and the buffers they name without copying + #[staticmethod] + fn from_wkb_polygon_views( + views: PyReadonlyArray1, + buffers: Vec>, + ) -> PyResult { + let views_sl = views + .as_slice() + .map_err(|_| PyValueError::new_err("views must be a contiguous uint8 array"))?; + let owned: Vec<&[u8]> = buffers + .iter() + .map(|b| { + b.as_slice().map_err(|_| { + PyValueError::new_err("each buffer must be a contiguous uint8 array") + }) + }) + .collect::>()?; + let decode_started = Instant::now(); + let parsed = wkb::parse_polygon_views(views_sl, &owned).map_err(PyValueError::new_err)?; + Engine::from_parsed_polygons(parsed, elapsed_ns(decode_started)) } /// Copy selected logical polygons into a compact Engine with no inherited indexes diff --git a/src/wkb.rs b/src/wkb.rs index 1abc5fe..91c34ef 100644 --- a/src/wkb.rs +++ b/src/wkb.rs @@ -187,17 +187,78 @@ fn count_geometry(r: &mut Reader) -> Result { } } -/// Total coordinate count of geometries `[g_start, g_end)`, used to size each chunk's slices -fn count_chunk( - data: &[u8], - offsets: &[i64], - g_start: usize, - g_end: usize, -) -> Result { +/// Addressing for the WKB bytes of each geometry in a column +pub trait Blobs: Sync { + /// Bytes of geometry `g` or an error when the source is malformed + fn blob(&self, g: usize) -> Result<&[u8], String>; + + /// Number of geometries the column holds + fn count(&self) -> usize; +} + +/// One concatenated value buffer where geometry g is `data[offsets[g]..offsets[g+1]]` +pub struct Contiguous<'a> { + pub data: &'a [u8], + pub offsets: &'a [i64], +} + +impl Blobs for Contiguous<'_> { + fn blob(&self, g: usize) -> Result<&[u8], String> { + let (start, end) = (self.offsets[g] as usize, self.offsets[g + 1] as usize); + self.data + .get(start..end) + .ok_or_else(|| "WKB offset out of bounds".to_string()) + } + + fn count(&self) -> usize { + self.offsets.len().saturating_sub(1) + } +} + +/// Arrow BinaryView descriptors naming a length and a buffer and an offset per geometry +pub struct Views<'a> { + pub views: &'a [u8], + pub buffers: &'a [&'a [u8]], +} + +/// Bytes per Arrow BinaryView descriptor +const VIEW_WIDTH: usize = 16; + +/// Length at or below which a view stores its bytes inline instead of naming a buffer +const VIEW_INLINE_MAX: u32 = 12; + +impl Blobs for Views<'_> { + fn blob(&self, g: usize) -> Result<&[u8], String> { + let base = g * VIEW_WIDTH; + let view = self + .views + .get(base..base + VIEW_WIDTH) + .ok_or_else(|| "WKB view out of bounds".to_string())?; + let length = u32::from_le_bytes([view[0], view[1], view[2], view[3]]) as usize; + if length <= VIEW_INLINE_MAX as usize { + return Ok(&view[4..4 + length]); + } + let index = u32::from_le_bytes([view[8], view[9], view[10], view[11]]) as usize; + let offset = u32::from_le_bytes([view[12], view[13], view[14], view[15]]) as usize; + let buffer = self + .buffers + .get(index) + .ok_or_else(|| "WKB view names a missing buffer".to_string())?; + buffer + .get(offset..offset + length) + .ok_or_else(|| "WKB view offset out of bounds".to_string()) + } + + fn count(&self) -> usize { + self.views.len() / VIEW_WIDTH + } +} + +/// Total coordinate count of geometries `[g_start, g_end)` that sizes each chunk's slices +fn count_chunk(blobs: &B, g_start: usize, g_end: usize) -> Result { let mut coords = 0usize; for g in g_start..g_end { - let (start, end) = (offsets[g] as usize, offsets[g + 1] as usize); - let slice = data.get(start..end).ok_or("WKB offset out of bounds")?; + let slice = blobs.blob(g)?; coords += count_geometry(&mut Reader::new(slice))?; } Ok(coords) @@ -283,9 +344,8 @@ struct ChunkSmall { /// Decode geometries `[g_start, g_end)` straight into the pre-sized `xs`/`ys` slices. The slices /// must hold exactly this chunk's coordinate count, which is asserted so the caller may set_len. -fn fill_chunk( - data: &[u8], - offsets: &[i64], +fn fill_chunk( + blobs: &B, g_start: usize, g_end: usize, xs: &mut [MaybeUninit], @@ -297,8 +357,7 @@ fn fill_chunk( let mut part_poly = Vec::new(); let mut multipart = false; for g in g_start..g_end { - let (start, end) = (offsets[g] as usize, offsets[g + 1] as usize); - let slice = data.get(start..end).ok_or("WKB offset out of bounds")?; + let slice = blobs.blob(g)?; let parts = fill_geometry( &mut Reader::new(slice), g, @@ -325,12 +384,11 @@ fn fill_chunk( /// Parse a WKB column in `chunk_size`-geometry chunks. A cheap count pass sizes the output, then /// the chunks decode in parallel into disjoint slices. The result is independent of `chunk_size`. -fn parse_polygons_chunked( - data: &[u8], - offsets: &[i64], +fn parse_polygons_chunked( + blobs: &B, chunk_size: usize, ) -> Result { - let n = offsets.len().saturating_sub(1); + let n = blobs.count(); if n == 0 { return Ok(ParsedPolygons { xs: Vec::new(), @@ -349,7 +407,7 @@ fn parse_polygons_chunked( // Pass 1: count coordinates per chunk so each chunk's output region is known before filling let coord_counts: Vec = bounds .par_iter() - .map(|&(s, e)| count_chunk(data, offsets, s, e)) + .map(|&(s, e)| count_chunk(blobs, s, e)) .collect::>()?; let total_coords: usize = coord_counts.iter().sum(); @@ -371,7 +429,7 @@ fn parse_polygons_chunked( } tasks .into_par_iter() - .map(|((s, e), xa, ya)| fill_chunk(data, offsets, s, e, xa, ya)) + .map(|((s, e), xa, ya)| fill_chunk(blobs, s, e, xa, ya)) .collect::>()? }; @@ -414,16 +472,87 @@ fn parse_polygons_chunked( /// Parse a whole WKB column. `data` is the concatenated value buffer, geometry g is /// `data[offsets[g]..offsets[g+1]]`. pub fn parse_polygons(data: &[u8], offsets: &[i64]) -> Result { - let n = offsets.len().saturating_sub(1); + parse_blobs(&Contiguous { data, offsets }) +} + +/// Parse a whole WKB column addressed by Arrow BinaryView descriptors over borrowed buffers +pub fn parse_polygon_views(views: &[u8], buffers: &[&[u8]]) -> Result { + parse_blobs(&Views { views, buffers }) +} + +/// Split a column into one chunk per thread and decode the chunks in parallel +fn parse_blobs(blobs: &B) -> Result { let n_threads = rayon::current_num_threads().max(1); - let chunk_size = n.div_ceil(n_threads).max(MIN_DECODE_CHUNK); - parse_polygons_chunked(data, offsets, chunk_size) + let chunk_size = blobs.count().div_ceil(n_threads).max(MIN_DECODE_CHUNK); + parse_polygons_chunked(blobs, chunk_size) } #[cfg(test)] mod tests { use super::*; + // Build a BinaryView descriptor naming a buffer or holding the bytes inline + fn view_descriptor(bytes: &[u8], buffer: u32, offset: u32) -> [u8; 16] { + let mut view = [0u8; 16]; + view[0..4].copy_from_slice(&(bytes.len() as u32).to_le_bytes()); + if bytes.len() <= 12 { + view[4..4 + bytes.len()].copy_from_slice(bytes); + return view; + } + view[4..8].copy_from_slice(&bytes[..4]); + view[8..12].copy_from_slice(&buffer.to_le_bytes()); + view[12..16].copy_from_slice(&offset.to_le_bytes()); + view + } + + #[test] + fn views_address_the_same_geometry_as_concatenated_offsets() { + let geoms: Vec> = (0..3).map(|g| square_with_hole(g as f64 * 10.0)).collect(); + let mut data = Vec::new(); + let mut offsets = vec![0i64]; + for geom in &geoms { + data.extend_from_slice(geom); + offsets.push(data.len() as i64); + } + let mut views = Vec::new(); + for (g, geom) in geoms.iter().enumerate() { + // One buffer per geometry proves the decoder follows the buffer index + views.extend_from_slice(&view_descriptor(geom, g as u32, 0)); + } + let borrowed: Vec<&[u8]> = geoms.iter().map(|g| g.as_slice()).collect(); + + let from_views = parse_polygon_views(&views, &borrowed).unwrap(); + let from_offsets = parse_polygons(&data, &offsets).unwrap(); + + assert_eq!(from_views.xs, from_offsets.xs); + assert_eq!(from_views.ys, from_offsets.ys); + assert_eq!(from_views.ring_offsets, from_offsets.ring_offsets); + assert_eq!(from_views.poly_offsets, from_offsets.poly_offsets); + } + + #[test] + fn an_inline_view_reads_its_bytes_out_of_the_descriptor() { + let inline = view_descriptor(&[1, 2, 3, 4], 0, 0); + let blobs = Views { + views: &inline, + buffers: &[], + }; + + assert_eq!(blobs.blob(0).unwrap(), &[1, 2, 3, 4]); + assert_eq!(blobs.count(), 1); + } + + #[test] + fn a_view_naming_a_missing_buffer_is_an_error() { + let view = view_descriptor(&[0u8; 40], 7, 0); + let blobs = Views { + views: &view, + buffers: &[], + }; + + assert!(blobs.blob(0).is_err()); + } + // Little-endian WKB for a unit square at (cx, cy): one closed 5-point ring fn le_square(cx: f64, cy: f64) -> Vec { let mut b = vec![1]; // little-endian @@ -522,9 +651,23 @@ mod tests { } let n = offsets.len() - 1; // chunk_size >= n is a single chunk, the serial reference the parallel paths must match - let serial = parse_polygons_chunked(&data, &offsets, n).unwrap(); + let serial = parse_polygons_chunked( + &Contiguous { + data: &data, + offsets: &offsets, + }, + n, + ) + .unwrap(); for chunk_size in [1usize, 2, 13, 64, 256] { - let p = parse_polygons_chunked(&data, &offsets, chunk_size).unwrap(); + let p = parse_polygons_chunked( + &Contiguous { + data: &data, + offsets: &offsets, + }, + chunk_size, + ) + .unwrap(); assert_eq!(p.xs, serial.xs, "xs differ at chunk_size {chunk_size}"); assert_eq!(p.ys, serial.ys, "ys differ at chunk_size {chunk_size}"); assert_eq!( @@ -631,6 +774,13 @@ mod tests { bad.truncate(bad.len() - 8); data.extend_from_slice(&bad); offsets.push(data.len() as i64); - assert!(parse_polygons_chunked(&data, &offsets, 1).is_err()); + assert!(parse_polygons_chunked( + &Contiguous { + data: &data, + offsets: &offsets, + }, + 1, + ) + .is_err()); } } From 7391cf64dd2a473629fd51381b2fdad30420447f Mon Sep 17 00:00:00 2001 From: Pranav Walimbe Date: Tue, 25 Aug 2026 12:26:10 -0700 Subject: [PATCH 3/4] perf: decode point WKB from arrow views without materialising bytes --- python/pycanopy/engine.py | 50 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 50 insertions(+) diff --git a/python/pycanopy/engine.py b/python/pycanopy/engine.py index d47f59c..0dcead3 100644 --- a/python/pycanopy/engine.py +++ b/python/pycanopy/engine.py @@ -207,6 +207,10 @@ def wkb_points_to_xy(points) -> tuple[np.ndarray, np.ndarray]: Returns: Pair (xs, ys) of contiguous float64 numpy arrays. """ + view_fast = _wkb_points_view_fast(points) + if view_fast is not None: + return view_fast + if hasattr(points, "to_arrow"): # e.g. a polars Series points = points.to_arrow() if isinstance(points, pa.ChunkedArray): @@ -292,6 +296,49 @@ def distance_to_point( ) +def _wkb_points_view_fast(column) -> tuple[np.ndarray, np.ndarray] | None: + # Read x/y from a WKB point column's Arrow view buffers without materialising its bytes + extracted = _wkb_view_buffers(column) + if extracted is None: + return None + descriptors, buffers = extracted + n = len(descriptors) // _VIEW_WIDTH + if n == 0: + return (np.empty(0, dtype=np.float64), np.empty(0, dtype=np.float64)) + + fields = descriptors.view(_VIEW_RECORD) + if not np.all(fields["length"] == _WKB_POINT_NBYTES): + return None + indices = fields["index"] + offsets = fields["offset"] + + # Each buffer holds one ascending run of views so a run decodes as one strided read + breaks = np.flatnonzero(np.diff(indices.astype(np.int64))) + 1 + starts = np.concatenate([[0], breaks]) + ends = np.concatenate([breaks, [n]]) + + xs = np.empty(n, dtype=np.float64) + ys = np.empty(n, dtype=np.float64) + for start, end in zip(starts, ends, strict=True): + span = end - start + base = offsets[start] + if not np.array_equal(offsets[start:end] - base, np.arange(span) * _WKB_POINT_NBYTES): + return None + buffer = buffers[indices[start]] + block = buffer[base : base + span * _WKB_POINT_NBYTES] + if block.size != span * _WKB_POINT_NBYTES: + return None + records = block.view(_WKB_POINT_RECORD) + if not ( + np.all(records["order"] == _WKB_LITTLE_ENDIAN) + and np.all(records["type"] == _WKB_POINT_TYPE) + ): + return None + xs[start:end] = records["x"] + ys[start:end] = records["y"] + return xs, ys + + def _wkb_points_fast(arr: pa.Array) -> tuple[np.ndarray, np.ndarray] | None: # Read x/y from a uniformly 21-byte WKB point column via one numpy view, or None for # nulls or any non-uniform or non-point layout so the caller can fall back to shapely. @@ -332,6 +379,9 @@ def _wkb_points_fast(arr: pa.Array) -> tuple[np.ndarray, np.ndarray] | None: # Bytes per Arrow BinaryView descriptor _VIEW_WIDTH = 16 +# Field layout of one Arrow BinaryView descriptor +_VIEW_RECORD = np.dtype([("length", " Date: Tue, 25 Aug 2026 14:47:53 -0700 Subject: [PATCH 4/4] revert: keep engine arrays in Arc<[T]> to protect index retrieval --- src/index/brute.rs | 11 ++++--- src/index/grid.rs | 9 +++--- src/index/kdtree.rs | 8 ++--- src/index/mod.rs | 7 +++-- src/index/rtree.rs | 4 +-- src/lib.rs | 61 +++++++++++++++++---------------------- tests/rust/index_tests.rs | 14 ++++----- 7 files changed, 53 insertions(+), 61 deletions(-) diff --git a/src/index/brute.rs b/src/index/brute.rs index 29c7276..b58134f 100644 --- a/src/index/brute.rs +++ b/src/index/brute.rs @@ -3,7 +3,6 @@ use std::sync::Arc; use crate::index::{point_box_dist2, SpatialIndex}; -use crate::Shared; /// Linear scan index, used for small datasets or high-selectivity queries. /// @@ -15,14 +14,14 @@ pub struct BruteForce { /// Per-geometry bounding boxes. /// For point datasets all four are Arc::clone of the Engine's xs/ys (shared, zero cost). /// For polygon datasets these are new allocations derived from ring coords. - bbox_min_x: Shared, - bbox_min_y: Shared, - bbox_max_x: Shared, - bbox_max_y: Shared, + bbox_min_x: Arc<[f64]>, + bbox_min_y: Arc<[f64]>, + bbox_max_x: Arc<[f64]>, + bbox_max_y: Arc<[f64]>, } impl SpatialIndex for BruteForce { - fn build(xs: Shared, ys: Shared) -> Self { + fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { BruteForce { bbox_min_x: Arc::clone(&xs), bbox_min_y: Arc::clone(&ys), diff --git a/src/index/grid.rs b/src/index/grid.rs index 9283096..ab0ba8b 100644 --- a/src/index/grid.rs +++ b/src/index/grid.rs @@ -1,9 +1,10 @@ //! Uniform grid index with CSR cell storage for large, uniformly distributed datasets. +use std::sync::Arc; + use rayon::prelude::*; use crate::index::SpatialIndex; -use crate::Shared; /// Uniform grid index with CSR (compressed sparse row) cell storage. /// Best for large datasets with uniform spatial distribution. @@ -14,8 +15,8 @@ pub struct UniformGrid { /// cell_offsets[i]..cell_offsets[i+1] is the slice of indices in cell i cell_offsets: Vec, indices: Vec, - xs: Shared, - ys: Shared, + xs: Arc<[f64]>, + ys: Arc<[f64]>, min_x: f64, min_y: f64, cell_w: f64, @@ -54,7 +55,7 @@ impl UniformGrid { } impl SpatialIndex for UniformGrid { - fn build(xs: Shared, ys: Shared) -> Self { + fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { let n = xs.len(); let (min_x, min_y, max_x, max_y) = xs.iter().zip(ys.iter()).fold( diff --git a/src/index/kdtree.rs b/src/index/kdtree.rs index ebe9b5a..afa42ed 100644 --- a/src/index/kdtree.rs +++ b/src/index/kdtree.rs @@ -1,12 +1,12 @@ //! Packed immutable KD-tree index optimised for point datasets. use std::f64::consts::PI; +use std::sync::Arc; use geo_index::kdtree::{KDTree, KDTreeBuilder, KDTreeIndex}; use crate::index::SpatialIndex; use crate::stats::types::SpatialHistogram; -use crate::Shared; /// Packed immutable KD-tree backed by geo-index, optimised for point datasets. /// @@ -15,14 +15,14 @@ use crate::Shared; /// traversal). The xs/ys Arcs are kept for the kNN distance refinement step. pub struct PackedKdTree { tree: KDTree, - xs: Shared, - ys: Shared, + xs: Arc<[f64]>, + ys: Arc<[f64]>, extent_area: f64, histogram: Option, } impl SpatialIndex for PackedKdTree { - fn build(xs: Shared, ys: Shared) -> Self { + fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { let n = xs.len(); let mut builder = KDTreeBuilder::::new(n as u32); diff --git a/src/index/mod.rs b/src/index/mod.rs index bae45dc..7f5372b 100644 --- a/src/index/mod.rs +++ b/src/index/mod.rs @@ -5,13 +5,14 @@ pub mod grid; pub mod kdtree; pub mod rtree; -use crate::Shared; +use std::sync::Arc; /// Common interface for all spatial index backends. -/// Coordinates arrive as Shared so an index shares the Engine allocation by refcount. +/// Coordinates are passed as Arc<[f64]> so indexes can share Engine's allocation +/// without copying, since storing an Arc<[f64]> is an atomic refcount bump, not a memcpy. pub trait SpatialIndex: Send + Sync { /// Build an index over the given coordinate arrays - fn build(xs: Shared, ys: Shared) -> Self + fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self where Self: Sized; /// Indices of the k nearest entries to (qx, qy), sorted nearest-first. diff --git a/src/index/rtree.rs b/src/index/rtree.rs index 4baab4e..ce5543b 100644 --- a/src/index/rtree.rs +++ b/src/index/rtree.rs @@ -3,6 +3,7 @@ use std::cell::RefCell; use std::cmp::{Ordering, Reverse}; use std::collections::{BinaryHeap, VecDeque}; +use std::sync::Arc; use rayon::prelude::*; @@ -10,7 +11,6 @@ use geo_index::rtree::sort::HilbertSort; use geo_index::rtree::{RTree, RTreeBuilder, RTreeIndex}; use crate::index::{point_box_dist2, SpatialIndex}; -use crate::Shared; /// One entry in the best-first traversal queue, ordered by squared distance to the node's box. /// @@ -96,7 +96,7 @@ pub struct PackedRTree { } impl SpatialIndex for PackedRTree { - fn build(xs: Shared, ys: Shared) -> Self { + fn build(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Self { let n = xs.len() as u32; let mut builder = RTreeBuilder::::new(n); for (&x, &y) in xs.iter().zip(ys.iter()) { diff --git a/src/lib.rs b/src/lib.rs index 7868908..dd188a8 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -69,7 +69,7 @@ impl Engine { parsed.ys.into(), parsed.ring_offsets.into(), parsed.poly_offsets.into(), - parsed.part_poly.map(Arc::new), + parsed.part_poly.map(Arc::from), n_polygons, ); engine.metrics.wkb_decode_ns = wkb_decode_ns; @@ -77,15 +77,12 @@ impl Engine { } } -// Shared array storage that adopts a decoded Vec without copying -pub(crate) type Shared = Arc>; - #[pyclass] struct Engine { - xs: Shared, - ys: Shared, - ring_offsets: Option>, // coord range per ring, None for point datasets - poly_offsets: Option>, // ring range per polygon, first ring exterior, rest holes + xs: Arc<[f64]>, + ys: Arc<[f64]>, + ring_offsets: Option>, // coord range per ring, None for point datasets + poly_offsets: Option>, // ring range per polygon, first ring exterior, rest holes stats: stats::types::DatasetStats, brute: Option, rtree: Option, @@ -97,7 +94,7 @@ struct Engine { index_mode: IndexMode, // index selection policy: Eager (default) / None / Auto / Explicit cost_factors: CostFactors, prepared_polys: Option, // sub-linear PIP edge index, built lazily - part_poly: Option>, // logical polygon per index part, None if no MultiPolygons + part_poly: Option>, // logical polygon per index part, None if no MultiPolygons n_polygons: usize, // count of logical polygons, equals part count when part_poly is None metric: DistanceMetric, // how threshold distances are measured: Planar (default) / Haversine metrics: EngineMetrics, // fixed, always-on counters drained explicitly through Python @@ -282,7 +279,7 @@ impl PyPolygonBuilder { } impl Engine { - fn new_points(xs: Shared, ys: Shared) -> Engine { + fn new_points(xs: Arc<[f64]>, ys: Arc<[f64]>) -> Engine { let stats_started = Instant::now(); let stats = collector::collect_points(&xs, &ys); let statistics_ns = elapsed_ns(stats_started); @@ -310,11 +307,11 @@ impl Engine { } fn new_polygons( - xs: Shared, - ys: Shared, - ring_offsets: Shared, - poly_offsets: Shared, - part_poly: Option>, + xs: Arc<[f64]>, + ys: Arc<[f64]>, + ring_offsets: Arc<[i64]>, + poly_offsets: Arc<[i64]>, + part_poly: Option>, n_polygons: usize, ) -> Engine { let stats_started = Instant::now(); @@ -596,7 +593,7 @@ impl Engine { #[pymethods] impl Engine { /// Construct from contiguous float64 numpy arrays of x and y coordinates. - /// Copies the borrowed arrays once into a Shared that every index build then shares. + /// Copies once into an Arc<[f64]> that all later index builds share via Arc::clone. #[staticmethod] fn from_points(xs: PyReadonlyArray1, ys: PyReadonlyArray1) -> PyResult { let xs_sl = xs @@ -609,10 +606,7 @@ impl Engine { return Err(PyValueError::new_err("xs and ys must have the same length")); } ensure_u32_indexable(xs_sl.len(), "point dataset")?; - Ok(Engine::new_points( - Arc::new(xs_sl.to_vec()), - Arc::new(ys_sl.to_vec()), - )) + Ok(Engine::new_points(xs_sl.into(), ys_sl.into())) } /// Construct from two-level polygon ring arrays (supports holes and MultiPolygons). @@ -676,7 +670,7 @@ impl Engine { } // part_poly maps parts to logical polygons. // None means each part is its own logical polygon (no MultiPolygons). - let (part_poly_arc, n_polygons): (Option>, usize) = match part_poly { + let (part_poly_arc, n_polygons): (Option>, usize) = match part_poly { Some(pp) => { let pp_sl = pp.as_slice().map_err(|_| { PyValueError::new_err("part_poly must be a contiguous int64 array") @@ -688,18 +682,15 @@ impl Engine { } let n_logical = pp_sl.iter().copied().max().map_or(0, |m| m as usize + 1); ensure_u32_indexable(n_logical, "logical polygon dataset")?; - ( - Some(Arc::new(pp_sl.iter().map(|&v| v as u32).collect())), - n_logical, - ) + (Some(pp_sl.iter().map(|&v| v as u32).collect()), n_logical) } None => (None, n_polys), }; Ok(Engine::new_polygons( - Arc::new(xs_sl.to_vec()), - Arc::new(ys_sl.to_vec()), - Arc::new(ring_sl.to_vec()), - Arc::new(poly_sl.to_vec()), + xs_sl.into(), + ys_sl.into(), + ring_sl.into(), + poly_sl.into(), part_poly_arc, n_polygons, )) @@ -1658,7 +1649,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let prepared = self.prepared_polys.as_ref(); - let part_poly = self.part_poly.as_deref().map(Vec::as_slice); + let part_poly = self.part_poly.as_deref(); let flat = match kind { IndexKind::BruteForce => par_contains( self.brute.as_ref().unwrap(), @@ -1764,7 +1755,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let prepared = self.prepared_polys.as_ref(); - let part_poly = self.part_poly.as_deref().map(Vec::as_slice); + let part_poly = self.part_poly.as_deref(); let state = match kind { IndexKind::BruteForce => par_contains_aggregate( self.brute.as_ref().unwrap(), @@ -1854,7 +1845,7 @@ impl Engine { let build_ns = self.build_index_if_needed(kind); let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); - let part_poly = self.part_poly.as_deref().map(Vec::as_slice); + let part_poly = self.part_poly.as_deref(); let flat = match kind { IndexKind::BruteForce => par_within_distance_to_polygons( self.brute.as_ref().unwrap(), @@ -1956,7 +1947,7 @@ impl Engine { let build_ns = self.build_index_if_needed(kind); let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); - let part_poly = self.part_poly.as_deref().map(Vec::as_slice); + let part_poly = self.part_poly.as_deref(); let state = match kind { IndexKind::BruteForce => par_within_distance_to_polygons_aggregate( self.brute.as_ref().unwrap(), @@ -2048,7 +2039,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let n_parts = poly_off.len().saturating_sub(1); - let part_poly = self.part_poly.as_deref().map(Vec::as_slice); + let part_poly = self.part_poly.as_deref(); let (idx, dist) = match kind { IndexKind::BruteForce => par_knn_to_polygons( self.brute.as_ref().unwrap(), @@ -2135,7 +2126,7 @@ impl Engine { let ring_off = self.ring_offsets.as_deref().unwrap(); let poly_off = self.poly_offsets.as_deref().unwrap(); let n_parts = poly_off.len().saturating_sub(1); - let part_poly = self.part_poly.as_deref().map(Vec::as_slice); + let part_poly = self.part_poly.as_deref(); let (q_idx, t_idx, dist) = match kind { IndexKind::BruteForce => par_knn_to_polygons_sorted( self.brute.as_ref().unwrap(), diff --git a/tests/rust/index_tests.rs b/tests/rust/index_tests.rs index 4e64cfe..64ccd71 100644 --- a/tests/rust/index_tests.rs +++ b/tests/rust/index_tests.rs @@ -9,10 +9,10 @@ use pycanopy::index::{ }; use pycanopy::query::range::{query_contains_polygons, query_range_polygons}; -fn five_point_grid() -> (Arc>, Arc>) { +fn five_point_grid() -> (Arc<[f64]>, Arc<[f64]>) { ( - Arc::new(vec![0.0f64, 1.0, 2.0, 0.0, 1.0]), - Arc::new(vec![0.0f64, 0.0, 0.0, 1.0, 1.0]), + Arc::from([0.0f64, 1.0, 2.0, 0.0, 1.0].as_slice()), + Arc::from([0.0f64, 0.0, 0.0, 1.0, 1.0].as_slice()), ) } @@ -192,8 +192,8 @@ fn range_empty_all_implementations_agree() { #[test] fn nearest_k5_on_larger_dataset_all_agree() { - let xs: Arc> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); - let ys: Arc> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); + let xs: Arc<[f64]> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); + let ys: Arc<[f64]> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); let oracle = as_set(BruteForce::build(Arc::clone(&xs), Arc::clone(&ys)).nearest(4.6, 3.2, 5)); assert_eq!( @@ -215,8 +215,8 @@ fn nearest_k5_on_larger_dataset_all_agree() { #[test] fn range_on_larger_dataset_all_agree() { - let xs: Arc> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); - let ys: Arc> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); + let xs: Arc<[f64]> = (0..100).map(|i| (i % 10) as f64).collect::>().into(); + let ys: Arc<[f64]> = (0..100).map(|i| (i / 10) as f64).collect::>().into(); let oracle = as_set(BruteForce::build(Arc::clone(&xs), Arc::clone(&ys)).range(2.0, 2.0, 5.0, 5.0));