feat: more detailed perf profiling

CONTRIBUTING.md

@@ -2,7 +2,7 @@
 
 ## Build setup
 
-You need Rust (stable) and Python 3.10–3.12.
+You need Rust (stable), `cargo-nextest`, and Python 3.10–3.12. Install the Rust test runner with `cargo install cargo-nextest` if it is not already available.
 
 ```bash
 # Clone and set up
@@ -13,7 +13,7 @@ cd PyCanopy
 uv sync --group dev
 
 # Build the Rust extension and install in editable mode
-maturin develop
+uv run maturin develop
 
 # Full check: format + build + lint + test
 make check
@@ -22,44 +22,67 @@ make check
 For a release build (needed for accurate benchmark numbers):
 
 ```bash
-maturin develop --release
+uv run maturin develop --release
 ```
 
 ## Make targets
 
 | Command | What it does |
 |:--------|:-------------|
-| `make check` | fmt + build + lint + test |
-| `make test` | Run the Python test suite |
+| `make check` | Format, build, lint and run every test. The default target |
 | `make build` | Debug build |
-| `make build-prod` | Release build |
+| `make build-prod` | Release build, needed for accurate benchmark numbers |
+| `make tune-engine` | Calibrate planner costs and update the bundled profile |
+| `make profile` | Two-build SF1 profile, writes `assets/profile.txt` |
+| `make sf1`, `make sf10` | SpatialBench at that scale factor, all four engines |
 | `make clean` | Remove build artifacts |
 
+`profile`, `sf1` and `sf10` launch EC2 instances and need AWS credentials. Narrow the engine
+list with `make sf1 engines=pycanopy`.
+
 ## Running tests
 
 ```bash
-make test
+make check
 # or directly
-pytest tests/python -x -q
+cargo nextest run
+uv run pytest tests/python -x -q
 ```
 
 ## Style
 
 After every code change, run:
 
 ```bash
-ruff format && ruff check
+uv run ruff format && uv run ruff check
 cargo fmt && cargo clippy
 ```
 
-To avoid a slopocolypse, I like these guidelines:
+`scripts/check_comments.py` enforces the rules below that ruff and clippy cannot express. Run
+`uv run python scripts/check_comments.py` to list violations, or with `--fix` to strip trailing
+periods from single-line comments.
+
+For coding style, I like these guidelines:
+
+**Comments**
 
+- Comments annotate code in one line. Use a multi-line block only when one line truly cannot carry it.
+- Comments should use near-zero commas. Say the one thing the reader needs and stop.
+- A comment states its fact and stops. Never trail a justification clause off a comma, such as "so this is exact" or "the way the old code did".
 - No em dashes, no semicolons in comments or docstrings.
-- All Python imports at module level.
-- Public Python functions use Google-style docstrings (`Args:`, `Returns:`).
-- Private Python functions use a `#` comment as the first line in the body.
-- Rust `pub` items require `///` doc comments; every module file requires `//!`.
 - Single-line comments have no trailing period, multi-line comment blocks end each sentence with a period.
-- Comments annotate code in one line. Reach for a multi-line block only when one line truly cannot carry it.
-- Comments use minimal commas. Say the one thing the reader needs and stop.
-- Docstrings carry no `Raises:` section.
+- Write a TODO as `// TODO(name):` or `# TODO(name):`.
+
+**Python**
+
+- All imports at module level.
+- Public functions use Google-style docstrings with `Args:`, `Returns:` and `Yields:` as applicable.
+- Docstrings carry no `Raises:` section and no line for a `None` return or input.
+- Private functions carry no docstring. They use a `#` comment as the first line in the body.
+
+**Rust**
+
+- `pub` items require a one-line `///` doc comment. Private `fn` usually carry none.
+- `///` docs are free prose. `Args:` and `Returns:` headings are Python-only.
+- Every module file requires a single-line `//!` module doc.
+- Every `unsafe` block requires a `// SAFETY:` comment.

Makefile

@@ -41,15 +41,15 @@ build-prod:
 tune-engine: build-prod
 	uv run python -m bench.ops
 
-.PHONY: profile ## Two-build SF1 profile on EC2, writes assets/profile.txt
+.PHONY: profile
 profile:
 	uv run --group bench python -m bench.spatial_bench --profile
 
-.PHONY: sf1 ## SpatialBench SF1 on EC2, all four engines unless engines= is set
+.PHONY: sf1
 sf1:
 	uv run --group bench python -m bench.spatial_bench --scale-factor 1 --engine $(engines)
 
-.PHONY: sf10 ## SpatialBench SF10 on EC2, all four engines unless engines= is set
+.PHONY: sf10
 sf10:
 	uv run --group bench python -m bench.spatial_bench --scale-factor 10 --engine $(engines)
 

bench/spatial_bench/profiler_utils.py

@@ -2,6 +2,7 @@
 
 from __future__ import annotations
 
+import inspect
 import os
 import resource
 import sys
@@ -73,6 +74,69 @@ def stop(self) -> None:
         self._observe()
 
 
+# Host calls no Engine metric can see paired with the stage each belongs to
+_HOST_STAGES = (
+    ("read", "pl", "read_parquet"),
+    ("read", "pl", "collect_all"),
+    ("wkb decode", None, "wkb_points_to_xy"),
+    ("wkb decode", None, "wkb_point_distance"),
+    ("frame build", "SpatialFrame", "from_wkb_points"),
+    ("frame build", "SpatialFrame", "from_wkb_polygons"),
+)
+
+
+def _timed_call(profiler: StageProfiler, stage: str, call):
+    # Wrap one host call and accumulate its wall time under a named stage
+    def wrapper(*args, **kwargs):
+        with profiler.stage(stage):
+            return call(*args, **kwargs)
+
+    return wrapper
+
+
+def _patch_target(profiler: StageProfiler, module, stage: str, holder_name, attr):
+    # Replace one name with a timed wrapper and return what restores it
+    if holder_name is None:
+        original = module.__dict__.get(attr)
+        if original is None:
+            return None
+        module.__dict__[attr] = _timed_call(profiler, stage, original)
+        return (module.__dict__, attr, original)
+    holder = module.__dict__.get(holder_name)
+    original = getattr(holder, attr, None) if holder is not None else None
+    if original is None:
+        return None
+    wrapper = _timed_call(profiler, stage, original)
+    restore = vars(holder).get(attr, original)
+    setattr(holder, attr, staticmethod(wrapper) if inspect.isclass(holder) else wrapper)
+    return (holder, attr, restore)
+
+
+@contextmanager
+def instrument_host_stages(profiler: StageProfiler, module):
+    """Time the read and decode work that runs outside the Engine.
+
+    Args:
+        profiler: Stage profiler that accumulates the timings.
+        module: Query module whose globals hold the names to patch.
+
+    Yields:
+        None, for the duration of the instrumented region.
+    """
+    restore = [_patch_target(profiler, module, *target) for target in _HOST_STAGES]
+    try:
+        yield
+    finally:
+        for entry in reversed(restore):
+            if entry is None:
+                continue
+            holder, attr, original = entry
+            if isinstance(holder, dict):
+                holder[attr] = original
+            else:
+                setattr(holder, attr, original)
+
+
 def _aggregate_engine_metrics(engines: list[dict]) -> dict:
     # Fold every Engine created during the run into one set of totals
     construction = {"wkb_decode_ns": 0, "statistics_ns": 0}
@@ -124,7 +188,7 @@ def profile_payload(
             the Engine section empty rather than genuinely zero.
 
     Returns:
-        A JSON-serialisable payload with time, memory, and Engine sections.
+        A JSON-serialisable payload with time, memory, host stage and Engine sections.
     """
     engine = _aggregate_engine_metrics(engines)
     engine_ns = sum(engine["construction"].values())
@@ -139,6 +203,7 @@ def profile_payload(
             "baseline": profiler.baseline,
             "peak": profiler.peak,
         },
+        "stages": dict(profiler.times),
         "engine": engine,
         "metrics": metrics,
     }

bench/spatial_bench/report_utils.py

@@ -337,6 +337,11 @@ def _section(query_id: str, result: dict) -> str:
         f"{construction['wkb_decode_ns'] / NS_PER_SECOND:7.3f}s  statistics "
         f"{construction['statistics_ns'] / NS_PER_SECOND:7.3f}s",
     ]
+    host = _host_times(result)
+    if host:
+        lines.append(
+            "host stages   " + "  ".join(f"{name} {value:.3f}s" for name, value in host.items())
+        )
     if engine["index_builds"]:
         lines.append("index builds")
         for metric in engine["index_builds"]:
@@ -457,6 +462,49 @@ def _stage_times(result: dict) -> dict[str, float]:
     return stages
 
 
+# Host stage order in pipeline sequence with the derived remainder last
+_HOST_ORDER = ("read", "wkb decode", "frame build", "polars compute")
+
+
+def _host_times(result: dict) -> dict[str, float]:
+    # Host wall per stage with everything unattributed folded into polars compute
+    if result.get("status") != "ok":
+        return {}
+    profile = result["profile"]
+    stages = dict(profile.get("stages", {}))
+    if not stages:
+        return {}
+    engine = profile["engine"]
+    outside = sum(metric["elapsed_compute_ns"] for metric in engine["index_builds"])
+    outside += sum(metric["elapsed_compute_ns"] for metric in engine["operations"])
+    accounted = sum(stages.values()) + outside / NS_PER_SECOND
+    stages["polars compute"] = max(profile["time"]["total"] - accounted, 0.0)
+    return {name: stages[name] for name in _HOST_ORDER if name in stages}
+
+
+def _host_table(results: dict[str, dict], query_ids: list[str], labels: list[str]) -> str:
+    # Per-query host wall by stage, covering the time no Engine metric reports
+    head = f"{'query':<7}{'stage':<54}{labels[0]:>12}{labels[1]:>13}{'delta':>14}"
+    lines = [_WIDE_SEP, "Host time by stage, seconds", _WIDE_SUB, head, "-" * len(head)]
+    for query_id in query_ids:
+        stages = [_host_times(results[label].get(query_id, {})) for label in labels]
+        names = _stage_order(stages)
+        if not names:
+            continue
+        for position, name in enumerate(names):
+            new, old = stages[0].get(name, 0.0), stages[1].get(name, 0.0)
+            lines.append(
+                f"{query_id if position == 0 else '':<7}{name:<54}{new:>12.4f}"
+                f"{old:>13.4f}{_delta(new, old):>14}"
+            )
+    lines.append(
+        "\nHost stages are wall clock around the read and decode calls. frame build contains "
+        "the Engine's own WKB decode and statistics. polars compute is the unattributed "
+        "remainder after every host stage and Engine metric is subtracted."
+    )
+    return "\n".join(lines)
+
+
 def _reports_metrics(results: dict) -> bool:
     # A wheel without the private metrics hook still profiles wall time and memory
     return any(
@@ -590,6 +638,7 @@ def write_profile_comparison(transports: dict[str, dict], out_path: Path) -> Non
     if len(order) == 2:
         parts.append(_comparison_table(results, query_ids, order))
         parts.append(_stage_table(results, query_ids, order))
+        parts.append(_host_table(results, query_ids, order))
         parts.append(f"{_WIDE_SEP}\nPer-query detail, {order[0]} build")
     parts.extend(
         _section(query_id, primary[query_id]) for query_id in query_ids if query_id in primary

bench/spatial_bench/run_query.py

@@ -57,15 +57,18 @@ def _run_profiled(query_id: str, data_dir: str, scale_factor: int) -> None:
     # Deferred because an ordinary run of another engine installs neither PyCanopy nor Polars
     from bench.spatial_bench.profiler_utils import (  # noqa: PLC0415
         StageProfiler,
+        instrument_host_stages,
         profile_payload,
         verify_output,
     )
+    from bench.spatial_bench.queries import pycanopy as queries  # noqa: PLC0415
 
     runner = load_runner("pycanopy")
     profiler = StageProfiler()
     try:
         runner.prepare(data_dir)
-        with _capture_metrics() as capture:
+        module = queries.BY_ID[query_id]
+        with _capture_metrics() as capture, instrument_host_stages(profiler, module):
             started = time.perf_counter()
             result = runner.execute(query_id)
             result = _materialize(result)

docs/contributing.md

@@ -29,17 +29,23 @@ uv run maturin develop --release
 
 | Command | What it does |
 |:--------|:-------------|
-| `make check` | fmt + build + lint + test |
-| `make test` | Run the Python test suite |
+| `make check` | Format, build, lint and run every test. The default target |
 | `make build` | Debug build |
-| `make build-prod` | Release build |
+| `make build-prod` | Release build, needed for accurate benchmark numbers |
+| `make tune-engine` | Calibrate planner costs and update the bundled profile |
+| `make profile` | Two-build SF1 profile, writes `assets/profile.txt` |
+| `make sf1`, `make sf10` | SpatialBench at that scale factor, all four engines |
 | `make clean` | Remove build artifacts |
 
+`profile`, `sf1` and `sf10` launch EC2 instances and need AWS credentials. Narrow the engine
+list with `make sf1 engines=pycanopy`.
+
 ## Running tests
 
 ```bash
-make test
+make check
 # or directly
+cargo nextest run
 uv run pytest tests/python -x -q
 ```
 
@@ -52,11 +58,31 @@ uv run ruff format && uv run ruff check
 cargo fmt && cargo clippy
 ```
 
+`scripts/check_comments.py` enforces the rules below that ruff and clippy cannot express. Run
+`uv run python scripts/check_comments.py` to list violations, or with `--fix` to strip trailing
+periods from single-line comments.
+
 To avoid a slopocolypse, I recommend using these guidelines:
 
+**Comments**
+
+- Comments annotate code in one line. Use a multi-line block only when one line truly cannot carry it.
+- Comments use minimal commas. The checker allows one. Say the one thing the reader needs and stop.
+- A comment states its fact and stops. Never trail a justification clause off a comma, such as "so this is exact" or "the way the old code did".
 - No em dashes, no semicolons in comments or docstrings.
-- All Python imports at module level.
-- Public Python functions use Google-style docstrings (`Args:`, `Returns:`).
-- Private Python functions use a `#` comment as the first line in the body.
-- Rust `pub` items require `///` doc comments; every module file requires `//!`.
 - Single-line comments have no trailing period, multi-line comment blocks end each sentence with a period.
+- Write a TODO as `// TODO(name):` or `# TODO(name):`.
+
+**Python**
+
+- All imports at module level.
+- Public functions use Google-style docstrings with `Args:`, `Returns:` and `Yields:` as applicable.
+- Docstrings carry no `Raises:` section and no line for a `None` return or input.
+- Private functions carry no docstring. They use a `#` comment as the first line in the body.
+
+**Rust**
+
+- `pub` items require a one-line `///` doc comment. Private `fn` usually carry none.
+- `///` docs are free prose. `Args:` and `Returns:` headings are Python-only.
+- Every module file requires a single-line `//!` module doc.
+- Every `unsafe` block requires a `// SAFETY:` comment.