feat: q2 filter change + parallelized gridding for stats

bench/spatial_bench/__main__.py

@@ -343,7 +343,7 @@ def main(argv: list[str] | None = None) -> int:
         else:
             produced = []
 
-    # A failing released baseline is a result, so only the branch build is required
+    # A failing released baseline is a result and only the branch build is required
     required = ["branch"] if args.profile else list(statuses)
     if not all(statuses.get(node, False) for node in required) or not produced:
         logs = [p for p in paths if p.suffix == ".log"]

bench/spatial_bench/bootstrap.sh

@@ -76,8 +76,8 @@ log "installing uv"
 curl -LsSf https://astral.sh/uv/install.sh | sh
 source "$HOME/.local/bin/env"
 
-# Amazon Linux 2023 ships Python 3.9, below the project floor, so pin uv to a managed
-# 3.10 for every sync and run. It is the supported floor, so the box exercises it.
+# Amazon Linux 2023 ships Python 3.9 below the project floor
+# Pin uv to a managed 3.10 for every sync and run which is the supported floor
 export UV_PYTHON=3.10
 
 log "cloning ${REPO_URL} @ ${REPO_BRANCH}"
@@ -125,7 +125,7 @@ fi
 
 mkdir -p /data/scratch /opt/pycanopy/assets
 
-# /tmp is tmpfs (RAM) on Amazon Linux 2023, so spill out-of-core scratch and Polars sort to the EBS data volume.
+# /tmp is tmpfs on Amazon Linux 2023 so scratch and Polars sort spill to the EBS data volume
 export PYCANOPY_SCRATCH=/data/scratch
 export POLARS_TEMP_DIR=/data/scratch
 export TMPDIR=/data/scratch

bench/spatial_bench/config.py

@@ -31,8 +31,8 @@
 
 STORAGE_OPTIONS = {"skip_signature": "true"}
 
-# Mirror of the upstream Hugging Face dataset, which is what the committed answers were
-# generated from. Public read, so queries stay anonymous.
+# Mirror of the upstream Hugging Face dataset the committed answers were generated from
+# Public read keeps queries anonymous
 PUBLIC_DATA_ROOT = "s3://pycanopy-bench-data/spatialbench"
 PUBLIC_DATA_TEMPLATE = f"{PUBLIC_DATA_ROOT}/{DATASET_VERSION}/sf{{scale_factor}}"
 

bench/spatial_bench/driver_utils.py

@@ -239,7 +239,7 @@ def run_profile_suite(query_ids: list[str], data_dir: str, variant: str = "branc
         for query_id, result in results.items()
         if result["status"] != "ok" or result.get("verify") != "match"
     ]
-    # The released build is a baseline, so only the branch build has to verify clean
+    # The released build is a baseline and only the branch build has to verify clean
     if invalid and variant == "branch":
         raise RuntimeError(f"profile verification failed for: {', '.join(invalid)}")
     if invalid:

bench/spatial_bench/queries/pycanopy/q02.py

@@ -17,22 +17,22 @@
 
 
 def pycanopy(data_paths: dict[str, str]) -> pl.DataFrame:
+    # The zone filter belongs in the scan so Polars decodes one boundary instead of every one
     zone, trip = pl.collect_all(
         [
-            pl.scan_parquet(data_paths["zone"], storage_options=STORAGE_OPTIONS).select(
-                ["z_name", "z_boundary"]
-            ),
+            pl.scan_parquet(data_paths["zone"], storage_options=STORAGE_OPTIONS)
+            .filter(pl.col("z_name") == ZONE_NAME)
+            .select(["z_name", "z_boundary"])
+            .head(1),
             pl.scan_parquet(data_paths["trip"], storage_options=STORAGE_OPTIONS).select(
                 ["t_pickuploc"]
             ),
         ]
     )
-    target = zone.filter(pl.col("z_name") == ZONE_NAME).head(1)
     # from_wkb keeps a MultiPolygon whole rather than exploding it into parts
-    poly = shapely.from_wkb(target["z_boundary"][0])
+    poly = shapely.from_wkb(zone["z_boundary"][0])
 
     sf = SpatialFrame.from_wkb_points(trip, "t_pickuploc")
-    # Only the count is needed, so take the engine's matching indices directly and skip
-    # gathering the in-zone rows into a DataFrame.
+    # Only the count is needed so take the engine's matching indices without gathering rows
     idx = sf.engine.points_within_distance_of_polygon(poly, 0.0)
     return pl.DataFrame({"trip_count_in_coconino_county": [len(idx)]})

bench/spatial_bench/queries/pycanopy/q04.py

@@ -28,7 +28,7 @@ def pycanopy(data_paths: dict[str, str]) -> pl.DataFrame:
             ),
         ]
     )
-    # Only the top trips need their geometry, so this second scan runs after the first result
+    # Only the top trips need geometry and this second scan runs after the first result
     trip = (
         trip_scan.select(["t_tripkey", "t_pickuploc"])
         .filter(pl.col("t_tripkey").is_in(top["t_tripkey"].implode()))

bench/spatial_bench/queries/pycanopy/q11.py

@@ -46,7 +46,7 @@ def pycanopy(data_paths: dict[str, str]) -> pl.DataFrame:
         .collect_batched()
     )
 
-    # Aligned morsels carry the same trips on each side, so per-morsel counts sum to the global count
+    # Aligned morsels carry the same trips on each side and per-morsel counts sum to the global count
     count = 0
     for pickup, dropoff in zip(pickup_batches, dropoff_batches, strict=True):
         count += (

bench/spatial_bench/report_utils.py

@@ -417,7 +417,7 @@ def _profile_head() -> str:
 
 
 def _metadata_block(metadata: dict[str, str]) -> str:
-    # Run metadata, so a later reader can tell a code change from a machine change
+    # Run metadata lets a later reader tell a code change from a machine change
     lines = [_SEP, "Run metadata", _SUBSEP]
     lines.extend(f"{key:<22}{value}" for key, value in metadata.items())
     return "\n".join(lines)
@@ -467,7 +467,7 @@ def _reports_metrics(results: dict) -> bool:
 
 
 def _stage_order(stages: list[dict[str, float]]) -> list[str]:
-    # Union both builds' stages, keeping pipeline order and appending whatever only one has
+    # Union both builds' stages in pipeline order and append whatever only one has
     ordered: list[str] = []
     for stage in stages:
         for name in stage:

python/pycanopy/coordinates.py

@@ -41,7 +41,7 @@ def resolve_coordinate_system(
         raise ValueError(
             f"coordinate_system must be 'planar' or 'geographic', got {coordinate_system!r}"
         )
-    # Only geographic constrains its coordinates, so a planar frame skips the scan entirely
+    # Only geographic constrains its coordinates and a planar frame skips the scan entirely
     if coordinate_system == "geographic" and len(xs) > 0 and not _looks_geographic(xs, ys):
         warnings.warn(
             "coordinate_system='geographic' reads x/y as WGS84 lon/lat degrees, but these "
@@ -55,5 +55,5 @@ def resolve_coordinate_system(
 
 
 def _looks_geographic(xs: np.ndarray, ys: np.ndarray) -> bool:
-    # Lon/lat degrees are bounded, so anything outside those bounds cannot be WGS84
+    # Lon/lat degrees are bounded and anything outside them cannot be WGS84
     return bool(np.all(np.abs(xs) <= _MAX_ABS_LON) and np.all(np.abs(ys) <= _MAX_ABS_LAT))

python/pycanopy/engine.py

@@ -307,8 +307,8 @@ def _wkb_points_fast(arr: pa.Array) -> tuple[np.ndarray, np.ndarray] | None:
     if offsets_buf is None or data_buf is None:
         return None
 
-    # Offsets are int32 for binary, int64 for large_binary (what polars emits). A
-    # sliced array shares its parent's buffers, so index past the slice's offset.
+    # Offsets are int32 for binary and int64 for large_binary (what polars emits)
+    # A sliced array shares its parent's buffers so index past the slice's offset
     offset_dtype = "<i8" if pa.types.is_large_binary(arr.type) else "<i4"
     offsets = np.frombuffer(offsets_buf, dtype=offset_dtype)[arr.offset : arr.offset + n + 1]
 

python/pycanopy/executor.py

@@ -52,8 +52,8 @@ def _range_plugin_expr(
     max_y: float,
     engine,
 ) -> pl.Expr:
-    # Bounding-box filter via map_batches. is_elementwise=False is a barrier, so the closure
-    # sees only post-scalar-filter rows, masked against the global index in Rust.
+    # Bounding-box filter via map_batches with is_elementwise=False as a barrier
+    # The closure sees only post-scalar-filter rows masked against the global index in Rust
     def _apply(s: pl.Series) -> pl.Series:
         orig_idx = s.to_numpy()
         if len(orig_idx) == 0:
@@ -200,8 +200,8 @@ def _execute_body(
 
         has_joins = any(isinstance(n, _JOIN_TYPES) for n in plan)
 
-        # Large-probe joins stream the probe in morsels and concatenate, so the join
-        # intermediate is bounded by one morsel rather than the full result.
+        # Large-probe joins stream the probe in morsels and concatenate
+        # The join intermediate is bounded by one morsel rather than the full result
         if has_joins:
             morsel = batch_size if batch_size is not None else MORSEL_ROWS
             join_node = next(n for n in plan if isinstance(n, _JOIN_TYPES))
@@ -415,8 +415,8 @@ def _emit_fused(
     def _emit_points_within_distance_of_polygon(
         self, node: PointsWithinDistanceOfPolygonNode, sf, lf: pl.LazyFrame
     ) -> pl.LazyFrame:
-        # Keep points within node.distance of the query polygon. The polygon is queried
-        # against all points, so indices resolve once and the lf filters by original row.
+        # Keep points within node.distance of the query polygon
+        # It is queried against all points so indices resolve once and lf filters by original row
         indices = sf.engine.points_within_distance_of_polygon(node.polygon, node.distance)
         return self._filter_by_indices(lf, indices)
 

python/pycanopy/optimizer.py

@@ -133,7 +133,7 @@ def optimize(self, plan: Plan, engine) -> Plan:
         """
         if not plan:
             return plan
-        # A trailing SelectNode is a terminal projection, not a predicate, so the cost passes skip it
+        # A trailing SelectNode is a terminal projection rather than a predicate and cost passes skip it
         select_tail = None
         if isinstance(plan[-1], SelectNode):
             select_tail = plan[-1]
@@ -200,7 +200,7 @@ def _disk_selectivity(
         return min(1.0, self._disk_area(node, coordinate_system) / total_area)
 
     def _disk_area(self, node: WithinDistanceOfPointNode, coordinate_system: str) -> float:
-        # A geographic radius is meters against a degree extent, so it becomes a degree ellipse
+        # A geographic radius is meters against a degree extent and becomes a degree ellipse
         if coordinate_system != "geographic":
             return math.pi * node.distance * node.distance
         cos_lat = math.cos(math.radians(min(abs(node.cy), _POLE_GUARD_LAT)))

src/index/brute.rs

@@ -31,7 +31,7 @@ impl SpatialIndex for BruteForce {
     }
 
     fn nearest(&self, qx: f64, qy: f64, k: usize) -> Vec<usize> {
-        // Rank by point-to-MBR distance, which is exact for the degenerate boxes of a point
+        // Rank by point-to-MBR distance which is exact for the degenerate boxes of a point
         // dataset and the lower bound the polygon refinement pass needs.
         let n = self.bbox_min_x.len();
         let k = k.min(n);
@@ -181,9 +181,9 @@ mod tests {
 
     #[test]
     fn polygon_nearest_ranks_by_mbr_not_centroid() {
-        // A wide sliver spanning x 0..10 at y 0, and a small square at x 4..5, y 3..4. From
-        // (5, 0.5) the sliver's edge is 0.5 away but its centroid is 5 away, which centroid
-        // ranking put second. Ranking by MBR is the lower bound polygon refinement relies on.
+        // A wide sliver spans x 0..10 at y 0 and a small square sits at x 4..5 y 3..4
+        // From (5, 0.5) the sliver's edge is 0.5 away while its centroid is 5 away
+        // Centroid ranking put it second and MBR is the lower bound refinement relies on
         let xs = vec![
             0.0, 10.0, 10.0, 0.0, 0.0, // sliver ring
             4.0, 5.0, 5.0, 4.0, 4.0, // square ring

src/index/rtree.rs

@@ -38,8 +38,9 @@ impl Ord for NeighborNode {
 }
 
 thread_local! {
-    // Reused across every query on this worker. geo-index's own `neighbors` and `search` allocate
-    // a fresh queue and result vector per call, which dominates a kernel invoked once per query.
+    // Reused across every query on this worker
+    // geo-index's own `neighbors` and `search` allocate a fresh queue and vector per call
+    // That allocation dominates a kernel invoked once per query
     static NEIGHBOR_QUEUE: RefCell<BinaryHeap<Reverse<NeighborNode>>> =
         const { RefCell::new(BinaryHeap::new()) };
     static SEARCH_QUEUE: RefCell<VecDeque<usize>> = const { RefCell::new(VecDeque::new()) };
@@ -258,7 +259,7 @@ mod tests {
         assert_eq!(sorted(build(xs, ys).range(0.5, 0.5, 1.5, 1.5)), vec![4]);
     }
 
-    // Deterministic pseudo-random points, so a failure reproduces without a rand dependency
+    // Deterministic pseudo-random points reproduce a failure without a rand dependency
     fn scattered(n: usize) -> (Vec<f64>, Vec<f64>) {
         let mut state = 0x2545_F491_4F6C_DD1Du64;
         let mut next = || {
@@ -315,7 +316,7 @@ mod tests {
 
     #[test]
     fn nearest_into_handles_exact_distance_ties() {
-        // Four coincident points: every candidate ties, so only the set is well defined
+        // Four coincident points tie on every candidate and only the set is well defined
         let tree = build(vec![1.0, 1.0, 1.0, 1.0, 9.0], vec![1.0, 1.0, 1.0, 1.0, 9.0]);
         let mut got = Vec::new();
         tree.nearest_into(1.0, 1.0, 3, &mut got);

src/lib.rs

@@ -954,7 +954,7 @@ impl Engine {
             ));
         }
         let started = Instant::now();
-        // Geographic distances are meters over degrees, so the box has to be widened, not dilated
+        // Geographic distances are meters over degrees and need the box widened not dilated
         let (min_x, min_y, max_x, max_y) = match self.metric {
             DistanceMetric::Planar => (cx - distance, cy - distance, cx + distance, cy + distance),
             DistanceMetric::Haversine => conservative_degree_box(cx, cy, distance),
@@ -2241,8 +2241,8 @@ impl Engine {
         }
         let ring_off = self.ring_offsets.as_deref().unwrap();
         let poly_off = self.poly_offsets.as_deref().unwrap();
-        // For MultiPolygons the pair area is the sum over their part pairs: parts within
-        // one polygon are disjoint, so the part intersections partition the overlap.
+        // For MultiPolygons the pair area sums over their part pairs
+        // Parts within one polygon are disjoint and their intersections partition the overlap
         let areas: Vec<f64> = match &self.part_poly {
             Some(pp) => {
                 let (offsets, parts) = polygon_parts_csr(pp, self.n_polygons);

src/planner/selector.rs

@@ -291,7 +291,7 @@ mod tests {
 
     #[test]
     fn auto_builds_index_for_many_probes() {
-        // Many probes amortise the build, so the index wins
+        // Many probes amortise the build and the index wins
         let s = stats(1_000_000, GeometryKind::Point, Distribution::Clustered);
         let f = CostFactors::default();
         assert_eq!(
@@ -312,7 +312,7 @@ mod tests {
 
     #[test]
     fn best_available_reuses_built_index_for_few_probes() {
-        // Build cost already paid, so even 1 probe uses the built index
+        // Build cost already paid so even 1 probe uses the built index
         let s = stats(1_000_000, GeometryKind::Point, Distribution::Clustered);
         let f = CostFactors::default();
         assert_eq!(