Feat: memcpy optimization

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.
@@ -329,6 +376,44 @@ 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", "<u4"), ("prefix", "<u4"), ("index", "<u4"), ("offset", "<u4")])
+
+
+@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 +570,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:

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<Engine> {
+        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)
+    }
+}
+
 #[pyclass]
 struct Engine {
     xs: Arc<[f64]>,
@@ -688,26 +713,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<u8>,
+        buffers: Vec<PyReadonlyArray1<u8>>,
+    ) -> PyResult<Self> {
+        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::<PyResult<_>>()?;
+        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