diff --git a/examples/osm/.gitignore b/examples/osm/.gitignore new file mode 100644 index 00000000..5dab2dc1 --- /dev/null +++ b/examples/osm/.gitignore @@ -0,0 +1,10 @@ +# Rust build artifacts for the bundled native parser/sorter +osmium-rs/target/ + +# Local OSM data + derived caches (can be hundreds of GB — never commit) +*.pbf +*.f64 +*.f32 +*.idx +*.xyzones +osm-data/ diff --git a/examples/osm/README.md b/examples/osm/README.md new file mode 100644 index 00000000..1295bbd4 --- /dev/null +++ b/examples/osm/README.md @@ -0,0 +1,99 @@ +# Every OpenStreetMap node, out-of-core + +Render **all ~10.7 billion OpenStreetMap planet nodes** as an interactive +density scatter in `xy`, straight off disk. This is the end-to-end proof of the +design dossier's §27 *"mmap (native)"* canonical-store row and the §2 *"100M+ / +out-of-core — interactive via viewport tiling, bounded RAM"* target, at true +planet scale: **172 GB of lon/lat lives on disk, resident memory stays +screen-bounded** (the density pyramid), never data-bounded. + +The dataset used here: `planet-latest.osm.pbf` decoded to +**10,742,674,832 nodes** → two f64 columns (`osm_lon.f64`, `osm_lat.f64`, +86 GB each). + +## What's in here + +| File | Role | +|---|---| +| `osmium-rs/` | Native Rust crate: `osm-nodes` (PBF → f64 columns, ~250× pyosmium) and `osm-sort` (builds the Tier-3 spatial index). | +| `ingest.py` | Parse a planet `.pbf` into f64 columns and render the out-of-core density scatter (with timings). | +| `viewer.py` | Interactive browser viewer — serves `xy`'s real WebGL client over plain HTTP; pan/zoom re-aggregate every viewport from disk. | +| `_pbf_split.py` | Standalone PBF blob-framing scanner (byte-range enumeration for parallel decode); reference/legacy helper. | + +## Prerequisites + +- **Rust** (for the native parser/sorter): `cd examples/osm/osmium-rs && cargo build --release` +- **The render client**, built once from the repo root (it is not committed — see #214): `npm ci && node js/build.mjs` +- **A planet PBF**: download `planet-latest.osm.pbf` from a mirror (~80 GB). *(Any `.osm.pbf` extract works — a city/country extract is a fast way to try the pipeline.)* +- Disk: ~260 GB free for the full planet (172 GB columns + 84 GB spatial index). + +All commands below are run **from the repo root**. + +## 1. Parse: PBF → canonical f64 columns + +```bash +python examples/osm/ingest.py \ + --pbf /path/to/planet-latest.osm.pbf \ + --out /path/to/osm-data +``` + +`osm-nodes` decodes DenseNodes (delta + zigzag varint) across all CPU cores +with a zlib-ng inflate backend, writing `osm_lon.f64` / `osm_lat.f64` straight +to disk. `ingest.py` then builds the figure and prints the memory report and a +few zoom timings — proving `canonical_bytes` (RAM-resident) stays **0** while +`canonical_mapped_bytes` (disk) holds all 172 GB. + +Re-run against columns already on disk with `--reuse` (skips parsing). + +## 2. (Optional but recommended) Build the deep-zoom spatial index + +The density pyramid alone answers zoomed-out views crisply, but its finest +level over the whole planet is ~2.4 km/cell — so city/street zoom would only +upsample (blurry). `osm-sort` spatially sorts every point into a grid-bucketed +layout so a viewport reads only its in-window points — deep zoom then gets +**sharper and cheaper the further you go**. + +```bash +examples/osm/osmium-rs/target/release/osm-sort \ + /path/to/osm-data/osm_lon.f64 \ + /path/to/osm-data/osm_lat.f64 \ + /path/to/osm-data/osm_spatial \ + --grid 8192 --partitions 512 +``` + +Produces `osm_spatial.lon.f32` / `osm_spatial.lat.f32` (sorted by cell) and +`osm_spatial.idx` (header + cumulative cell offsets). It's an external counting +sort — sequential I/O only, peak RAM ≈ one partition. `viewer.py` attaches it +automatically if `osm_spatial.idx` sits next to the columns. + +## 3. View + +```bash +python examples/osm/viewer.py --out /path/to/osm-data --port 8777 +# open http://localhost:8777/ +``` + +Drag to pan, scroll to zoom. Every viewport is re-aggregated through the same +`channel.handle_message` kernel protocol the notebook widget uses. The tier +served rides each update as `binning`: + +- `pyramid-L` / `-upsampled` — zoomed-out aggregate from the pre-binned pyramid. +- `spatial-exact` (`filter: nearest`) — deep zoom re-binned exactly from the + in-window points at full screen resolution (crisp streets, no interpolation blur). +- `spatial-points` — deep enough that the window fits the direct budget: the + real points ship as individual marks. + +## Rough timings (10.7B nodes, NVMe SSD, 16 cores) + +| Stage | Cost | Notes | +|---|---|---| +| Parse (`osm-nodes`) | minutes | ~250× the pyosmium binding; multithreaded + zlib-ng | +| Spatial sort (`osm-sort`) | ~8 min (~21 M pts/s) | 3-pass parallel external counting sort | +| Figure build (zone maps) | **~51 s first time, then instant** | one 172 GB scan; result cached to a `.xyzones` sidecar next to each column, so every later run (e.g. viewer restart) loads it in milliseconds | +| Pan / zoom | single-digit ms | screen-bounded per viewport, independent of N | + +The `.xyzones` zone-map cache is a general `xy` out-of-core feature (see +`python/xy/columns.py`): the first `figure()` over a memmapped column persists +its per-chunk statistics, and reopening validates the sidecar against the +file's size/mtime before reusing it — so the 51 s domain scan is paid once, not +on every restart. diff --git a/examples/osm/_pbf_split.py b/examples/osm/_pbf_split.py new file mode 100644 index 00000000..c5b28700 --- /dev/null +++ b/examples/osm/_pbf_split.py @@ -0,0 +1,93 @@ +"""Minimal PBF *framing* scanner — enumerate blob byte-ranges for parallel +parsing. This reads only the fixed 4-byte length prefixes and the small +`BlobHeader` protobuf (type + datasize) to walk from one blob to the next; it +never touches the compressed OSM payload. pyosmium still performs the actual +(trusted) coordinate decode on each byte range via `osmium.io.FileBuffer`. + +A .osm.pbf file is a sequence of: + [4-byte big-endian BlobHeader length][BlobHeader protobuf][Blob protobuf] +The first blob is an `OSMHeader`; the rest are `OSMData`. A byte range made of +`OSMHeader` bytes followed by any run of whole `OSMData` blobs is itself a valid +PBF stream, so it can be parsed independently in another process. +""" + +from __future__ import annotations + +import struct +from dataclasses import dataclass + + +def _read_varint(buf: memoryview, pos: int) -> tuple[int, int]: + """Decode a protobuf base-128 varint at ``pos``; return (value, new_pos).""" + result = 0 + shift = 0 + while True: + b = buf[pos] + pos += 1 + result |= (b & 0x7F) << shift + if not (b & 0x80): + return result, pos + shift += 7 + + +def _blob_header_fields(bh: memoryview) -> tuple[str, int]: + """Extract (type, datasize) from a BlobHeader protobuf (fields 1 and 3).""" + pos = 0 + btype = "" + datasize = -1 + n = len(bh) + while pos < n: + tag, pos = _read_varint(bh, pos) + field, wire = tag >> 3, tag & 0x7 + if wire == 2: # length-delimited + ln, pos = _read_varint(bh, pos) + if field == 1: # type (string) + btype = bytes(bh[pos : pos + ln]).decode("ascii") + pos += ln + elif wire == 0: # varint + val, pos = _read_varint(bh, pos) + if field == 3: # datasize + datasize = val + else: # 32/64-bit — not used by BlobHeader, skip defensively + pos += 4 if wire == 5 else 8 + if datasize < 0: + raise ValueError("BlobHeader missing datasize") + return btype, datasize + + +@dataclass +class Blob: + start: int # byte offset of the 4-byte length prefix + end: int # byte offset just past the blob body + btype: str + + +def scan_blobs(mm: memoryview) -> list[Blob]: + """Walk every blob in the mapped file, returning their byte ranges.""" + blobs: list[Blob] = [] + pos = 0 + total = len(mm) + while pos + 4 <= total: + (hlen,) = struct.unpack_from(">I", mm, pos) + hstart = pos + 4 + btype, datasize = _blob_header_fields(mm[hstart : hstart + hlen]) + end = hstart + hlen + datasize + if end > total: + break # truncated tail (e.g. partial download) — stop cleanly + blobs.append(Blob(pos, end, btype)) + pos = end + return blobs + + +def partition(blobs: list[Blob], n_parts: int) -> list[tuple[int, int]]: + """Split the OSMData blobs into ``n_parts`` contiguous (start,end) byte + ranges. The OSMHeader blob (blobs[0]) is excluded — callers prepend it.""" + data = [b for b in blobs if b.btype == "OSMData"] + if not data: + return [] + per = -(-len(data) // n_parts) # ceil + parts: list[tuple[int, int]] = [] + for i in range(0, len(data), per): + group = data[i : i + per] + parts.append((group[0].start, group[-1].end)) + return parts diff --git a/examples/osm/ingest.py b/examples/osm/ingest.py new file mode 100644 index 00000000..f3b42953 --- /dev/null +++ b/examples/osm/ingest.py @@ -0,0 +1,147 @@ +"""Ingest **every OpenStreetMap node** (~9 billion, from `planet-latest.osm.pbf`) +into xy's out-of-core canonical store and render it as a density scatter. + +This is the end-to-end proof of the §27 "mmap (native)" canonical-store row and +the §2 "100M+ / out-of-core — interactive via viewport tiling, bounded RAM" +target, at true planet scale. Peak resident memory stays screen-bounded (the +density pyramid), not data-bounded (172 GB of lon/lat on disk). + +Pipeline: + planet.pbf ──osm-nodes (native)──▶ osm_lon.f64 / osm_lat.f64 (disk f64) + ──▶ xy.scatter(density=True) ──▶ build_payload / density_view + +Usage (from the repo root): + python examples/osm/ingest.py \ + --pbf /path/to/planet-latest.osm.pbf \ + --out /path/to/osm-data \ + [--reuse] # skip parsing; reuse osm_lon.f64/osm_lat.f64 on disk + +The native parser (examples/osm/osmium-rs) is built with: + cd examples/osm/osmium-rs && cargo build --release +""" + +from __future__ import annotations + +import argparse +import os +import resource +import subprocess +import sys +import time + +import numpy as np + +# examples/osm/ → repo root is two levels up; the package lives in python/. +_HERE = os.path.dirname(__file__) +_REPO_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +sys.path.insert(0, os.path.join(_REPO_ROOT, "python")) +import xy # noqa: E402 +from xy._ooc import open_f64 # noqa: E402 + + +def _rss_gb() -> float: + return resource.getrusage(resource.RUSAGE_SELF).ru_maxrss / (1024 * 1024) # GB on Linux + + +def parse_planet(pbf: str, out_dir: str) -> tuple[np.memmap, np.memmap]: + """Decode all node coordinates via the native `osm-nodes` parser + (`osmium-rs/`, ~250x faster than the Python binding), then memmap the + resulting f64 columns. The parser writes canonical columns straight to + disk; xy never holds the data in RAM.""" + px = os.path.join(out_dir, "osm_lon.f64") + py = os.path.join(out_dir, "osm_lat.f64") + binary = os.path.join(_HERE, "osmium-rs", "target", "release", "osm-nodes") + if not os.path.exists(binary): + raise SystemExit( + f"native parser not built: {binary}\n" + "build it with: cd examples/osm/osmium-rs && cargo build --release" + ) + t0 = time.perf_counter() + subprocess.run([binary, pbf, px, py], check=True) + print(f"native parse wall-clock: {time.perf_counter() - t0:.0f}s", flush=True) + xcol, ycol = open_f64(px), open_f64(py) + print( + f"canonical on disk: {xcol.nbytes + ycol.nbytes:,} bytes across 2 columns " + f"({len(xcol):,} nodes)", + flush=True, + ) + return xcol, ycol + + +def bench(xcol: np.memmap, ycol: np.memmap) -> None: + n = len(xcol) + print(f"\n=== rendering {n:,} points (out-of-core density scatter) ===", flush=True) + + t0 = time.perf_counter() + fig = xy.chart(xy.scatter(x=xcol, y=ycol, density=True)).figure() + print( + f"figure build (ingest + zone maps, one disk scan): {time.perf_counter() - t0:.0f}s", + flush=True, + ) + + rep = fig.store.memory_report() + print(f" canonical_bytes (RAM-resident): {rep['canonical_bytes']:,}") + print(f" canonical_mapped_bytes (disk): {rep['canonical_mapped_bytes']:,}") + assert rep["canonical_bytes"] == 0, "canonical data leaked into RAM" + + t0 = time.perf_counter() + spec, blob = fig.build_payload(2048) + tr = spec["traces"][0] + print( + f"first paint: {time.perf_counter() - t0:.0f}s | wire blob = {len(blob):,} B " + f"| tier={tr.get('tier')}", + flush=True, + ) + + # Zoom into a continent-scale window; should serve from the pyramid. + for label, (x0, x1, y0, y1) in { + "world": (-180.0, 180.0, -85.0, 85.0), + "europe": (-12.0, 32.0, 35.0, 60.0), + "london": (-0.5, 0.3, 51.3, 51.7), + }.items(): + t0 = time.perf_counter() + s2, bufs2 = fig.density_view(0, x0, x1, y0, y1, 1000, 800) + tr2 = s2["traces"][0] if s2.get("traces") else {} + nbytes = sum(len(b) for b in bufs2) + print( + f" zoom [{label}]: {(time.perf_counter() - t0) * 1e3:.0f} ms " + f"| {len(bufs2)} bufs / {nbytes:,} B " + f"| tier={tr2.get('tier')} | binning={tr2.get('binning')} | visible={tr2.get('visible')}", + flush=True, + ) + + full = fig.memory_report() + print( + f"\nresident_array_bytes: {full['resident_array_bytes']:,} " + f"(canonical RAM {full['canonical_bytes']:,} + channels {full['channel_bytes']:,} " + f"+ pyramid {full['pyramid_bytes']:,})" + ) + print(f"canonical on disk (mapped): {full['canonical_mapped_bytes']:,} bytes") + print(f"peak process RSS: {_rss_gb():.2f} GB") + print("\nOK — every OSM node rendered with screen-bounded resident memory.") + + +def main() -> None: + ap = argparse.ArgumentParser() + ap.add_argument("--pbf", help="planet-latest.osm.pbf (omit with --reuse)") + ap.add_argument("--out", required=True, help="output dir for the f64 columns") + ap.add_argument( + "--reuse", + action="store_true", + help="skip parsing; reuse osm_lon.f64/osm_lat.f64 already on disk", + ) + args = ap.parse_args() + if not args.reuse and not args.pbf: + ap.error("--pbf is required unless --reuse is given") + + if args.reuse: + xcol = open_f64(os.path.join(args.out, "osm_lon.f64")) + ycol = open_f64(os.path.join(args.out, "osm_lat.f64")) + print(f"reusing {len(xcol):,} nodes from disk", flush=True) + else: + xcol, ycol = parse_planet(args.pbf, args.out) + bench(xcol, ycol) + + +if __name__ == "__main__": + main() diff --git a/examples/osm/osmium-rs/.gitignore b/examples/osm/osmium-rs/.gitignore new file mode 100644 index 00000000..ea8c4bf7 --- /dev/null +++ b/examples/osm/osmium-rs/.gitignore @@ -0,0 +1 @@ +/target diff --git a/examples/osm/osmium-rs/Cargo.lock b/examples/osm/osmium-rs/Cargo.lock new file mode 100644 index 00000000..d073966e --- /dev/null +++ b/examples/osm/osmium-rs/Cargo.lock @@ -0,0 +1,105 @@ +# This file is automatically @generated by Cargo. +# It is not intended for manual editing. +version = 4 + +[[package]] +name = "adler2" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "320119579fcad9c21884f5c4861d16174d0e06250625266f50fe6898340abefa" + +[[package]] +name = "cc" +version = "1.2.64" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "dad887fd958be91b5098c0248def011f4523ab786cd411be668777e55063501f" +dependencies = [ + "find-msvc-tools", + "shlex", +] + +[[package]] +name = "cfg-if" +version = "1.0.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9330f8b2ff13f34540b44e946ef35111825727b38d33286ef986142615121801" + +[[package]] +name = "cmake" +version = "0.1.58" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "c0f78a02292a74a88ac736019ab962ece0bc380e3f977bf72e376c5d78ff0678" +dependencies = [ + "cc", +] + +[[package]] +name = "crc32fast" +version = "1.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9481c1c90cbf2ac953f07c8d4a58aa3945c425b7185c9154d67a65e4230da511" +dependencies = [ + "cfg-if", +] + +[[package]] +name = "find-msvc-tools" +version = "0.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "5baebc0774151f905a1a2cc41989300b1e6fbb29aff0ceffa1064fdd3088d582" + +[[package]] +name = "flate2" +version = "1.1.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "843fba2746e448b37e26a819579957415c8cef339bf08564fe8b7ddbd959573c" +dependencies = [ + "crc32fast", + "libz-ng-sys", + "miniz_oxide", +] + +[[package]] +name = "libc" +version = "0.2.186" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "68ab91017fe16c622486840e4c83c9a37afeff978bd239b5293d61ece587de66" + +[[package]] +name = "libz-ng-sys" +version = "1.1.29" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "879917b256f6317769b9f374b435805a8697013098aacce5a38ac106cd6a9469" +dependencies = [ + "cmake", + "libc", +] + +[[package]] +name = "miniz_oxide" +version = "0.8.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "1fa76a2c86f704bdb222d66965fb3d63269ce38518b83cb0575fca855ebb6316" +dependencies = [ + "adler2", + "simd-adler32", +] + +[[package]] +name = "osmpbf-nodes" +version = "0.1.0" +dependencies = [ + "flate2", +] + +[[package]] +name = "shlex" +version = "2.0.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "f8fadd59c855ef2080decdef8ff161eb6661b86933c9d82e5ba29dc602a55aba" + +[[package]] +name = "simd-adler32" +version = "0.3.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "703d5c7ef118737c72f1af64ad2f6f8c5e1921f818cdcb97b8fe6fc69bf66214" diff --git a/examples/osm/osmium-rs/Cargo.toml b/examples/osm/osmium-rs/Cargo.toml new file mode 100644 index 00000000..9e0ac980 --- /dev/null +++ b/examples/osm/osmium-rs/Cargo.toml @@ -0,0 +1,30 @@ +[package] +name = "osmpbf-nodes" +version = "0.1.0" +edition = "2021" +description = "Fast, multithreaded extraction of node coordinates from OpenStreetMap .osm.pbf into raw f64 columns." +license = "Apache-2.0" +repository = "https://github.com/reflex-dev/osmpbf-nodes" +keywords = ["openstreetmap", "osm", "pbf", "geo", "parser"] +categories = ["parser-implementations", "science::geo"] + +[lib] +name = "osmpbf_nodes" + +[[bin]] +name = "osm-nodes" +path = "src/main.rs" + +[[bin]] +name = "osm-sort" +path = "src/sort_main.rs" + +# zlib-ng backend: ~2-4x faster inflate than the pure-Rust miniz_oxide default, +# which dominates PBF decode. Same flate2 API. (Falls back to `flate2 = "1"` +# with default features if zlib-ng is unavailable on a target.) +[dependencies] +flate2 = { version = "1", default-features = false, features = ["zlib-ng"] } + +[profile.release] +lto = "thin" +codegen-units = 1 diff --git a/examples/osm/osmium-rs/README.md b/examples/osm/osmium-rs/README.md new file mode 100644 index 00000000..a7615674 --- /dev/null +++ b/examples/osm/osmium-rs/README.md @@ -0,0 +1,66 @@ +# osmpbf-nodes + +Fast, multithreaded extraction of **node coordinates** from an OpenStreetMap +`.osm.pbf` file into two flat little-endian `f64` column files (`lon`, `lat`). + +Built for feeding large-scale point-cloud / plotting engines that memory-map +their canonical columns: the output is exactly what such an engine wants on +disk, and the whole planet (~9 billion nodes) decodes in **minutes** instead of +the hours a per-object Python binding takes. + +## Why + +The reference Python binding (pyosmium) materializes a Python object per node, +which caps throughput at a few hundred thousand nodes/second. This crate decodes +the PBF protobuf by hand, decompresses blobs with pure-Rust zlib +(`flate2`/`miniz_oxide`), and fans blob decoding across a thread pool. On a +16-core desktop it sustains **~50–60 M nodes/second**, validated to produce +bit-for-bit identical coordinates to pyosmium. + +## Install / build + +```sh +cargo build --release +``` + +## Use (CLI) + +```sh +osm-nodes planet-latest.osm.pbf lon.f64 lat.f64 --threads 16 +``` + +- `lon.f64` / `lat.f64`: native-endian `f64`, one value per node, `lon[i]` + paired with `lat[i]`. Load with `numpy.memmap(path, dtype='. +//! +//! ## Output +//! +//! Two files of native-endian `f64`, one value per node, `lon[i]`/`lat[i]` +//! paired. Node order is not preserved across threads — irrelevant for an +//! unordered point cloud, and the whole point of the parallel write. +//! +//! Unix only (uses positioned writes); the decoder itself is portable. + +pub mod sort; + +use std::fs::File; +use std::io::{self, BufReader, Read}; +use std::os::unix::fs::FileExt; +use std::path::Path; +use std::sync::atomic::{AtomicU64, Ordering}; +use std::sync::mpsc::sync_channel; +use std::sync::Arc; + +/// Summary of a decode run. +#[derive(Debug, Clone, Copy, Default)] +pub struct Stats { + /// Node coordinates written to the output columns. + pub nodes: u64, + /// Non-dense (`Node`) entries encountered — 0 for planet dumps; reported, + /// never silently dropped. + pub sparse_nodes: u64, + /// `OSMData` blocks decoded. + pub blocks: u64, +} + +// ---- protobuf primitives (trusted, well-formed input) ---------------------- + +#[inline] +fn read_varint(buf: &[u8], pos: &mut usize) -> u64 { + let mut result: u64 = 0; + let mut shift = 0u32; + loop { + let b = buf[*pos]; + *pos += 1; + result |= ((b & 0x7F) as u64) << shift; + if b & 0x80 == 0 { + return result; + } + shift += 7; + } +} + +#[inline] +fn zigzag(n: u64) -> i64 { + ((n >> 1) as i64) ^ -((n & 1) as i64) +} + +/// Advance `pos` past one field's value given its wire type. +#[inline] +fn skip_field(buf: &[u8], pos: &mut usize, wire: u64) { + match wire { + 0 => { + read_varint(buf, pos); + } + 1 => *pos += 8, + 2 => { + let len = read_varint(buf, pos) as usize; + *pos += len; + } + 5 => *pos += 4, + _ => panic!("unsupported protobuf wire type {wire}"), + } +} + +// ---- blob framing ---------------------------------------------------------- + +/// `(type, datasize)` from a `BlobHeader` (fields 1 = type, 3 = datasize). +fn parse_blob_header(bh: &[u8]) -> io::Result<(String, usize)> { + let mut pos = 0; + let mut btype = String::new(); + let mut datasize: i64 = -1; + while pos < bh.len() { + let tag = read_varint(bh, &mut pos); + let (field, wire) = (tag >> 3, tag & 0x7); + match (field, wire) { + (1, 2) => { + let len = read_varint(bh, &mut pos) as usize; + btype = String::from_utf8_lossy(&bh[pos..pos + len]).into_owned(); + pos += len; + } + (3, 0) => datasize = read_varint(bh, &mut pos) as i64, + (_, w) => skip_field(bh, &mut pos, w), + } + } + if datasize < 0 { + return Err(io::Error::new( + io::ErrorKind::InvalidData, + "BlobHeader has no datasize", + )); + } + Ok((btype, datasize as usize)) +} + +/// Decompress a `Blob` (field 1 = raw, field 3 = zlib_data, field 2 = raw_size) +/// into its `PrimitiveBlock` bytes. +fn inflate_blob(blob: &[u8]) -> io::Result> { + let mut pos = 0; + let mut raw_size = 0usize; + let mut raw: Option<(usize, usize)> = None; + let mut zlib: Option<(usize, usize)> = None; + while pos < blob.len() { + let tag = read_varint(blob, &mut pos); + let (field, wire) = (tag >> 3, tag & 0x7); + match (field, wire) { + (2, 0) => raw_size = read_varint(blob, &mut pos) as usize, + (1, 2) => { + let len = read_varint(blob, &mut pos) as usize; + raw = Some((pos, pos + len)); + pos += len; + } + (3, 2) => { + let len = read_varint(blob, &mut pos) as usize; + zlib = Some((pos, pos + len)); + pos += len; + } + (_, w) => skip_field(blob, &mut pos, w), + } + } + if let Some((s, e)) = raw { + return Ok(blob[s..e].to_vec()); + } + if let Some((s, e)) = zlib { + let mut out = Vec::with_capacity(raw_size); + flate2::read::ZlibDecoder::new(&blob[s..e]).read_to_end(&mut out)?; + return Ok(out); + } + Err(io::Error::new( + io::ErrorKind::InvalidData, + "Blob uses an unsupported compression (only raw/zlib handled)", + )) +} + +// ---- PrimitiveBlock / DenseNodes decode ------------------------------------ + +/// Decode every node coordinate in a `PrimitiveBlock`, appending degrees to +/// `lon`/`lat`. Returns the count of non-dense `Node` entries seen (also +/// decoded, not skipped) — reported for information only. +fn decode_primitive_block(buf: &[u8], lon: &mut Vec, lat: &mut Vec) -> u64 { + // Pass 1: coordinate scaling (granularity=17 default 100, lat_offset=19, + // lon_offset=20, all int, appear *after* the groups in the byte stream) and + // the byte ranges of each primitivegroup (field 2). + let mut pos = 0; + let mut granularity: i64 = 100; + let mut lat_off: i64 = 0; + let mut lon_off: i64 = 0; + let mut groups: Vec<(usize, usize)> = Vec::new(); + while pos < buf.len() { + let tag = read_varint(buf, &mut pos); + let (field, wire) = (tag >> 3, tag & 0x7); + match (field, wire) { + (2, 2) => { + let len = read_varint(buf, &mut pos) as usize; + groups.push((pos, pos + len)); + pos += len; + } + (17, 0) => granularity = read_varint(buf, &mut pos) as i64, + (19, 0) => lat_off = read_varint(buf, &mut pos) as i64, + (20, 0) => lon_off = read_varint(buf, &mut pos) as i64, + (_, w) => skip_field(buf, &mut pos, w), + } + } + let scale = granularity as f64; + let mut sparse = 0u64; + for (s, e) in groups { + sparse += decode_group(&buf[s..e], scale, lat_off as f64, lon_off as f64, lon, lat); + } + sparse +} + +fn decode_group( + buf: &[u8], + scale: f64, + lat_off: f64, + lon_off: f64, + lon: &mut Vec, + lat: &mut Vec, +) -> u64 { + let mut pos = 0; + let mut sparse = 0u64; + while pos < buf.len() { + let tag = read_varint(buf, &mut pos); + let (field, wire) = (tag >> 3, tag & 0x7); + match (field, wire) { + (2, 2) => { + // DenseNodes + let len = read_varint(buf, &mut pos) as usize; + decode_dense(&buf[pos..pos + len], scale, lat_off, lon_off, lon, lat); + pos += len; + } + (1, 2) => { + // repeated Node (sparse encoding). The planet file uses + // DenseNodes exclusively, but hand-made extracts may carry + // these — decode them too so "every node" actually holds. + let len = read_varint(buf, &mut pos) as usize; + decode_node(&buf[pos..pos + len], scale, lat_off, lon_off, lon, lat); + sparse += 1; + pos += len; + } + (_, w) => skip_field(buf, &mut pos, w), + } + } + sparse +} + +/// Decode one sparse `Node` message: `lat` (field 8) and `lon` (field 9) are +/// single zigzag sint64 varints, scaled exactly like DenseNodes. A node +/// missing either coordinate is skipped (malformed). +fn decode_node( + buf: &[u8], + scale: f64, + lat_off: f64, + lon_off: f64, + lon: &mut Vec, + lat: &mut Vec, +) { + let mut pos = 0; + let (mut lat_v, mut lon_v) = (None, None); + while pos < buf.len() { + let tag = read_varint(buf, &mut pos); + let (field, wire) = (tag >> 3, tag & 0x7); + match (field, wire) { + (8, 0) => lat_v = Some(zigzag(read_varint(buf, &mut pos))), + (9, 0) => lon_v = Some(zigzag(read_varint(buf, &mut pos))), + (_, w) => skip_field(buf, &mut pos, w), + } + } + if let (Some(la), Some(lo)) = (lat_v, lon_v) { + lat.push(1e-9 * (lat_off + scale * la as f64)); + lon.push(1e-9 * (lon_off + scale * lo as f64)); + } +} + +fn decode_dense( + buf: &[u8], + scale: f64, + lat_off: f64, + lon_off: f64, + lon: &mut Vec, + lat: &mut Vec, +) { + // DenseNodes: lat=8, lon=9 are packed, delta-encoded sint64 arrays of equal + // length (one entry per node). id=1 is ignored. + let mut pos = 0; + let mut lat_r: Option<(usize, usize)> = None; + let mut lon_r: Option<(usize, usize)> = None; + while pos < buf.len() { + let tag = read_varint(buf, &mut pos); + let (field, wire) = (tag >> 3, tag & 0x7); + match (field, wire) { + (8, 2) => { + let len = read_varint(buf, &mut pos) as usize; + lat_r = Some((pos, pos + len)); + pos += len; + } + (9, 2) => { + let len = read_varint(buf, &mut pos) as usize; + lon_r = Some((pos, pos + len)); + pos += len; + } + (_, w) => skip_field(buf, &mut pos, w), + } + } + let (Some((mut p, pe)), Some((mut q, qe))) = (lat_r, lon_r) else { + return; + }; + let (mut lat_v, mut lon_v) = (0i64, 0i64); + while p < pe && q < qe { + lat_v += zigzag(read_varint(buf, &mut p)); + lon_v += zigzag(read_varint(buf, &mut q)); + lat.push(1e-9 * (lat_off + scale * lat_v as f64)); + lon.push(1e-9 * (lon_off + scale * lon_v as f64)); + } +} + +// ---- driver ---------------------------------------------------------------- + +/// Decode all node coordinates from `pbf` into the `lon`/`lat` output files. +/// +/// The outputs are pre-sized to `capacity` values (a sparse file until +/// written) and truncated to the exact node count on completion, so an +/// over-estimate is cheap. `n_threads` blob decoders run in parallel while a +/// reader thread streams blobs sequentially. +pub fn decode_pbf_nodes( + pbf: &Path, + lon_path: &Path, + lat_path: &Path, + capacity: u64, + n_threads: usize, +) -> io::Result { + let lon_file = Arc::new(File::create(lon_path)?); + let lat_file = Arc::new(File::create(lat_path)?); + lon_file.set_len(capacity * 8)?; + lat_file.set_len(capacity * 8)?; + + let cursor = Arc::new(AtomicU64::new(0)); + let sparse = Arc::new(AtomicU64::new(0)); + let blocks = Arc::new(AtomicU64::new(0)); + let n_threads = n_threads.max(1); + + let (tx, rx) = sync_channel::>(n_threads * 2); + let rx = Arc::new(std::sync::Mutex::new(rx)); + + let mut workers = Vec::with_capacity(n_threads); + for _ in 0..n_threads { + let rx = Arc::clone(&rx); + let lon_file = Arc::clone(&lon_file); + let lat_file = Arc::clone(&lat_file); + let cursor = Arc::clone(&cursor); + let sparse = Arc::clone(&sparse); + let blocks = Arc::clone(&blocks); + workers.push(std::thread::spawn(move || -> io::Result<()> { + let mut lon: Vec = Vec::new(); + let mut lat: Vec = Vec::new(); + loop { + let blob = { + let guard = rx.lock().unwrap(); + match guard.recv() { + Ok(b) => b, + Err(_) => break, + } + }; + lon.clear(); + lat.clear(); + let pb = inflate_blob(&blob)?; + let s = decode_primitive_block(&pb, &mut lon, &mut lat); + if s > 0 { + sparse.fetch_add(s, Ordering::Relaxed); + } + blocks.fetch_add(1, Ordering::Relaxed); + let count = lon.len() as u64; + if count == 0 { + continue; + } + let off = cursor.fetch_add(count, Ordering::Relaxed); + if off + count > capacity { + return Err(io::Error::other( + "node count exceeded capacity; raise --capacity", + )); + } + let byte_off = off * 8; + lon_file.write_all_at(as_bytes(&lon), byte_off)?; + lat_file.write_all_at(as_bytes(&lat), byte_off)?; + } + Ok(()) + })); + } + + // Reader: stream blobs sequentially; hand OSMData blobs to the pool. + let read_result = (|| -> io::Result<()> { + let mut reader = BufReader::with_capacity(1 << 22, File::open(pbf)?); + // A clean stream ends exactly at a blob boundary — EOF *before* the next + // 4-byte length prefix. EOF anywhere inside a length prefix, header, or + // blob means the file is truncated (e.g. a partial download) and is a + // hard error: incomplete data must never masquerade as a complete decode. + loop { + let hlen = match read_len_prefix(&mut reader)? { + None => break, // clean end at a boundary + Some(hlen) => hlen, + }; + let mut hbuf = vec![0u8; hlen]; + read_exact_or_truncated(&mut reader, &mut hbuf)?; + let (btype, datasize) = parse_blob_header(&hbuf)?; + let mut blob = vec![0u8; datasize]; + read_exact_or_truncated(&mut reader, &mut blob)?; + if btype == "OSMData" { + // Send fails only if all workers died (their error surfaces on join). + if tx.send(blob).is_err() { + break; + } + } + } + Ok(()) + })(); + drop(tx); + + let mut first_err = read_result.err(); + for w in workers { + let joined = match w.join() { + Ok(inner) => inner, + Err(_) => Err(io::Error::other("worker panicked")), + }; + if let Err(e) = joined { + first_err.get_or_insert(e); + } + } + if let Some(e) = first_err { + return Err(e); + } + + let nodes = cursor.load(Ordering::Relaxed); + lon_file.set_len(nodes * 8)?; + lat_file.set_len(nodes * 8)?; + Ok(Stats { + nodes, + sparse_nodes: sparse.load(Ordering::Relaxed), + blocks: blocks.load(Ordering::Relaxed), + }) +} + +/// Read exactly `buf.len()` bytes, or return `None` if EOF arrives first +/// (a truncated final blob). Distinguishes clean truncation from a real error. +fn read_full_or_eof(r: &mut R, buf: &mut [u8]) -> io::Result> { + let mut filled = 0; + while filled < buf.len() { + match r.read(&mut buf[filled..]) { + Ok(0) => return Ok(None), + Ok(n) => filled += n, + Err(e) if e.kind() == io::ErrorKind::Interrupted => {} + Err(e) => return Err(e), + } + } + Ok(Some(())) +} + +fn truncated() -> io::Error { + io::Error::new( + io::ErrorKind::UnexpectedEof, + "truncated PBF: unexpected end of file mid-blob (incomplete download?)", + ) +} + +/// Read the 4-byte big-endian blob-length prefix. `Ok(None)` only when the +/// stream ends cleanly *at* the boundary (0 bytes available); a 1–3 byte +/// partial read is a truncated file, not a clean end. +fn read_len_prefix(r: &mut R) -> io::Result> { + let mut buf = [0u8; 4]; + let mut filled = 0; + while filled < 4 { + match r.read(&mut buf[filled..]) { + Ok(0) if filled == 0 => return Ok(None), + Ok(0) => return Err(truncated()), + Ok(n) => filled += n, + Err(e) if e.kind() == io::ErrorKind::Interrupted => {} + Err(e) => return Err(e), + } + } + Ok(Some(u32::from_be_bytes(buf) as usize)) +} + +/// Fill `buf` fully; an EOF before it is full means the file is truncated. +fn read_exact_or_truncated(r: &mut R, buf: &mut [u8]) -> io::Result<()> { + match read_full_or_eof(r, buf)? { + Some(()) => Ok(()), + None => Err(truncated()), + } +} + +#[inline] +fn as_bytes(v: &[f64]) -> &[u8] { + // f64 is plain-old-data; a read-only reinterpret for positioned writes. + unsafe { std::slice::from_raw_parts(v.as_ptr() as *const u8, std::mem::size_of_val(v)) } +} + +#[cfg(test)] +mod tests { + use super::*; + + #[test] + fn varint_roundtrip() { + // Encode a few varints back-to-back and read them. + let mut buf = Vec::new(); + for v in [0u64, 1, 127, 128, 300, 16384, u64::MAX] { + let mut x = v; + loop { + let mut b = (x & 0x7F) as u8; + x >>= 7; + if x != 0 { + b |= 0x80; + } + buf.push(b); + if x == 0 { + break; + } + } + } + let mut pos = 0; + for v in [0u64, 1, 127, 128, 300, 16384, u64::MAX] { + assert_eq!(read_varint(&buf, &mut pos), v); + } + assert_eq!(pos, buf.len()); + } + + #[test] + fn zigzag_known_values() { + assert_eq!(zigzag(0), 0); + assert_eq!(zigzag(1), -1); + assert_eq!(zigzag(2), 1); + assert_eq!(zigzag(3), -2); + assert_eq!(zigzag(4), 2); + // A delta chain like DenseNodes stores: +1, +1, -1 → 0,1,2,1. + let deltas = [zigzag(2), zigzag(2), zigzag(1)]; // 1, 1, -1 + let mut acc = 0i64; + let cum: Vec = deltas + .iter() + .map(|d| { + acc += d; + acc + }) + .collect(); + assert_eq!(cum, vec![1, 2, 1]); + } + + fn put_varint(buf: &mut Vec, mut x: u64) { + loop { + let mut b = (x & 0x7F) as u8; + x >>= 7; + if x != 0 { + b |= 0x80; + } + buf.push(b); + if x == 0 { + break; + } + } + } + + #[test] + fn read_len_prefix_distinguishes_clean_end_from_truncation() { + // 0 bytes at a boundary → clean end. + assert!(read_len_prefix(&mut &[][..]).unwrap().is_none()); + // A full 4-byte prefix decodes big-endian. + assert_eq!(read_len_prefix(&mut &[0, 0, 1, 0][..]).unwrap(), Some(256)); + // 1–3 bytes then EOF → truncation error, never a silent stop. + let err = read_len_prefix(&mut &[0, 0][..]).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + } + + #[test] + fn read_exact_or_truncated_errors_on_short_blob() { + let mut buf = [0u8; 8]; + let err = read_exact_or_truncated(&mut &[1, 2, 3][..], &mut buf).unwrap_err(); + assert_eq!(err.kind(), io::ErrorKind::UnexpectedEof); + } + + #[test] + fn decode_node_extracts_sparse_coordinates() { + // A sparse Node message with lat (field 8) and lon (field 9) as zigzag + // sint64 varints — the encoding planet files omit but extracts may use. + let (scale, lat_off, lon_off) = (100.0, 0.0, 0.0); + let (lat_raw, lon_raw) = (10_000_000i64, -20_000_000i64); + let mut msg = Vec::new(); + msg.push((8 << 3) | 0); // field 8, varint + put_varint(&mut msg, ((lat_raw << 1) ^ (lat_raw >> 63)) as u64); + msg.push((9 << 3) | 0); // field 9, varint + put_varint(&mut msg, ((lon_raw << 1) ^ (lon_raw >> 63)) as u64); + let (mut lon, mut lat) = (Vec::new(), Vec::new()); + decode_node(&msg, scale, lat_off, lon_off, &mut lon, &mut lat); + assert_eq!(lat, vec![1e-9 * scale * lat_raw as f64]); + assert_eq!(lon, vec![1e-9 * scale * lon_raw as f64]); + } +} diff --git a/examples/osm/osmium-rs/src/main.rs b/examples/osm/osmium-rs/src/main.rs new file mode 100644 index 00000000..b4bb99fa --- /dev/null +++ b/examples/osm/osmium-rs/src/main.rs @@ -0,0 +1,88 @@ +//! `osm-nodes` — extract node coordinates from an `.osm.pbf` into two flat +//! `f64` column files. +//! +//! ```text +//! osm-nodes [--threads N] [--capacity N] +//! ``` + +use std::path::Path; +use std::process::ExitCode; +use std::time::Instant; + +use osmpbf_nodes::decode_pbf_nodes; + +// Generous upper bound on planet node count (~9.5B as of 2026). The output +// files are sparse until written and truncated to the true count, so +// over-estimating costs nothing. +const DEFAULT_CAPACITY: u64 = 12_000_000_000; + +fn main() -> ExitCode { + let args: Vec = std::env::args().collect(); + let mut positional: Vec<&str> = Vec::new(); + let mut threads = std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(4); + let mut capacity = DEFAULT_CAPACITY; + + let mut i = 1; + while i < args.len() { + match args[i].as_str() { + "--threads" | "-t" => { + i += 1; + threads = args.get(i).and_then(|s| s.parse().ok()).unwrap_or(threads); + } + "--capacity" | "-c" => { + i += 1; + capacity = args.get(i).and_then(|s| s.parse().ok()).unwrap_or(capacity); + } + "-h" | "--help" => { + eprintln!( + "usage: osm-nodes \ + [--threads N] [--capacity N]" + ); + return ExitCode::SUCCESS; + } + other => positional.push(other), + } + i += 1; + } + + if positional.len() != 3 { + eprintln!("error: expected 3 paths (pbf, out_lon, out_lat); see --help"); + return ExitCode::FAILURE; + } + + let (pbf, lon, lat) = (positional[0], positional[1], positional[2]); + eprintln!("decoding {pbf} with {threads} threads (capacity {capacity})…"); + let t0 = Instant::now(); + match decode_pbf_nodes( + Path::new(pbf), + Path::new(lon), + Path::new(lat), + capacity, + threads, + ) { + Ok(stats) => { + let secs = t0.elapsed().as_secs_f64(); + eprintln!( + "done: {} nodes in {:.1}s ({:.1} M nodes/s) across {} blocks{}", + stats.nodes, + secs, + stats.nodes as f64 / secs / 1e6, + stats.blocks, + if stats.sparse_nodes > 0 { + format!(" (incl. {} non-dense nodes)", stats.sparse_nodes) + } else { + String::new() + }, + ); + // machine-readable last line + println!("{}", stats.nodes); + ExitCode::SUCCESS + } + Err(e) => { + eprintln!("error: {e}"); + ExitCode::FAILURE + } + } +} diff --git a/examples/osm/osmium-rs/src/sort.rs b/examples/osm/osmium-rs/src/sort.rs new file mode 100644 index 00000000..1987f2fe --- /dev/null +++ b/examples/osm/osmium-rs/src/sort.rs @@ -0,0 +1,369 @@ +//! Spatial sort: reorder canonical (lon, lat) f64 columns into a grid-bucketed +//! layout so a viewport query reads only its in-window points, not the whole +//! column. This is the Tier-3 out-of-core index (xy dossier §28/§32b): the +//! per-viewport cost drops from O(N) to O(points in window), so zoom gets +//! *sharper and faster* the deeper you go. +//! +//! Output is a derived f32 cache (half the bytes; ~1 m geo precision is plenty +//! for rendering — the f64 canonical store remains the source of truth): +//! - `_lon.f32`, `_lat.f32`: points sorted by row-major grid cell +//! - `.idx`: header + (g*g + 1) u64 cumulative offsets (prefix sum), so +//! cell b owns points [off[b], off[b+1]); a window's cells are contiguous +//! per grid row → one slice read per row. +//! +//! Algorithm: an external counting sort that never holds all N in RAM and does +//! sequential I/O only. Pass 1 histograms cell counts. Points are then split +//! into P equal-count partitions (skew-proof) written as sequential f32 files; +//! each partition (sized to RAM) is counting-sorted in memory and appended to +//! the final columns in cell order. Peak RAM ≈ one partition. + +use std::fs::File; +use std::io::{self, BufWriter, Read, Write}; +use std::os::unix::fs::FileExt; +use std::path::Path; +use std::sync::atomic::{AtomicU64, AtomicUsize, Ordering}; +use std::time::Instant; + +const MAGIC: &[u8; 8] = b"XYSPIDX1"; + +fn nthreads() -> usize { + std::thread::available_parallelism() + .map(|n| n.get()) + .unwrap_or(4) +} + +/// Read exactly `buf.len()` f64 at byte offset `off*8` via a positional read +/// (thread-safe, no shared cursor — many threads scan disjoint slices at once). +fn read_f64_at(f: &File, buf: &mut [f64], point_off: u64) -> io::Result<()> { + let bytes = + unsafe { std::slice::from_raw_parts_mut(buf.as_mut_ptr() as *mut u8, buf.len() * 8) }; + f.read_exact_at(bytes, point_off * 8) +} + +#[derive(Debug)] +pub struct SortStats { + pub n: u64, + pub g: usize, + pub partitions: usize, + pub max_partition: u64, +} + +#[inline] +fn cell_of(v: f64, lo: f64, inv_span_g: f64, g: usize) -> usize { + // Bucket by the *f32* value the index stores, not the f64 source: passes 1/2 + // see f64 and pass 3 sees the round-tripped f32, so rounding to f32 here + // makes the cell identical across all passes (else a boundary point lands in + // a different partition than its file → out of range). Clamp to [0, g-1]. + let vf = (v as f32) as f64; + let c = ((vf - lo) * inv_span_g) as isize; + c.max(0).min(g as isize - 1) as usize +} + +/// Append a thread-local scatter buffer to partition `q`'s file at a +/// reserved, non-overlapping offset. `pwrite` to disjoint ranges is safe from +/// many threads without a lock, so pass 2 fans out across cores while still +/// producing exactly `used_parts` files (not threads×parts — that would blow +/// the open-fd limit). Order within a partition file is arbitrary; pass 3 +/// counting-sorts by cell regardless. +fn flush_part(files: &[File], tails: &[AtomicU64], q: usize, buf: &mut Vec) -> io::Result<()> { + if buf.is_empty() { + return Ok(()); + } + let at = tails[q].fetch_add(buf.len() as u64, Ordering::Relaxed); + files[q].write_all_at(buf, at)?; + buf.clear(); + Ok(()) +} + +#[allow(clippy::too_many_arguments)] +pub fn spatial_sort( + lon_path: &Path, + lat_path: &Path, + out_prefix: &Path, + g: usize, + x0: f64, + x1: f64, + y0: f64, + y1: f64, + n_partitions: usize, +) -> io::Result { + // Validate the caller/CLI options up front so a bad value returns an error + // through the normal path instead of panicking deep in a worker: g==0 would + // index an empty histogram, n_partitions==0 would divide by zero, and g must + // fit the u32 the index header stores (with g*g not overflowing usize). + let bad = |m: &str| io::Error::new(io::ErrorKind::InvalidInput, m.to_string()); + if g == 0 { + return Err(bad("--grid must be >= 1")); + } + if g > u32::MAX as usize { + return Err(bad("--grid too large (must fit in a u32)")); + } + if n_partitions == 0 { + return Err(bad("--partitions must be >= 1")); + } + let cells = g + .checked_mul(g) + .ok_or_else(|| bad("--grid too large (g*g overflows)"))?; + let n = File::open(lon_path)?.metadata()?.len() / 8; + let inv_x = g as f64 / (x1 - x0); + let inv_y = g as f64 / (y1 - y0); + let block = 1usize << 20; // 1M values per I/O chunk + + // ---- Pass 1: histogram cell counts, parallel over disjoint slices -------- + // Each thread positionally reads its slice and folds a private histogram; + // the T partials are then summed. Turns a read-bound O(N) scan into a + // parallel one that saturates NVMe bandwidth. + let threads = nthreads(); + let mut counts = vec![0u64; cells]; + let t_p1 = Instant::now(); + { + let lon_f = File::open(lon_path)?; + let lat_f = File::open(lat_path)?; + let per = n.div_ceil(threads as u64); + let partials: Vec>> = std::thread::scope(|s| { + let hs: Vec<_> = (0..threads) + .map(|ti| { + let (lon_f, lat_f) = (&lon_f, &lat_f); + s.spawn(move || -> io::Result> { + let mut local = vec![0u64; cells]; + let (start, end) = (ti as u64 * per, ((ti as u64 + 1) * per).min(n)); + let (mut lb, mut tb) = (vec![0f64; block], vec![0f64; block]); + let mut off = start; + while off < end { + let k = (block as u64).min(end - off) as usize; + read_f64_at(lon_f, &mut lb[..k], off)?; + read_f64_at(lat_f, &mut tb[..k], off)?; + for i in 0..k { + let cx = cell_of(lb[i], x0, inv_x, g); + let cy = cell_of(tb[i], y0, inv_y, g); + local[cy * g + cx] += 1; + } + off += k as u64; + } + Ok(local) + }) + }) + .collect(); + hs.into_iter().map(|h| h.join().unwrap()).collect() + }); + for pres in partials { + let local = pres?; + for b in 0..cells { + counts[b] += local[b]; + } + } + } + + eprintln!(" pass 1 (histogram): {:.1}s", t_p1.elapsed().as_secs_f64()); + + // ---- Prefix sum → cell start offsets; equal-count partition boundaries --- + let mut offsets = vec![0u64; cells + 1]; + for b in 0..cells { + offsets[b + 1] = offsets[b] + counts[b]; + } + let total = offsets[cells]; + // Partition p owns cell range [cell_lo[p], cell_lo[p+1]); boundaries chosen + // so each partition holds ≈ total/P points. A cell→partition LUT makes the + // scatter O(1)/point. + let target = total.div_ceil(n_partitions as u64).max(1); + let mut part_of = vec![0u32; cells]; + let mut cell_lo = vec![0usize; n_partitions + 1]; + let mut p = 0usize; + for b in 0..cells { + while p + 1 < n_partitions && offsets[b] >= (p as u64 + 1) * target { + p += 1; + cell_lo[p] = b; + } + part_of[b] = p as u32; + } + for slot in &mut cell_lo[p + 1..=n_partitions] { + *slot = cells; + } + let used_parts = p + 1; + + // ---- Pass 2: scatter points into per-partition f32 files (sequential) ---- + let t_p2 = Instant::now(); + let tmp: Vec<_> = (0..used_parts) + .map(|q| out_prefix.with_extension(format!("part{q}"))) + .collect(); + { + // T threads each scatter a disjoint input slice into per-thread-local + // buffers keyed by partition, appending to the shared partition files + // via reserved-offset positional writes (see `flush_part`). Fully + // parallel — pass 2 is no longer the serial floor. + let files: Vec = tmp.iter().map(File::create).collect::>()?; + let tails: Vec = (0..used_parts).map(|_| AtomicU64::new(0)).collect(); + let lon_f = File::open(lon_path)?; + let lat_f = File::open(lat_path)?; + let per = n.div_ceil(threads as u64); + // Per-partition local buffer size before a positional write; 128 KiB + // keeps write syscalls coarse while the working set stays ≈ threads × + // used_parts × 128 KiB (~1 GiB at 16×512). + const FLUSH: usize = 1 << 17; + let (files, tails, part_of) = (&files, &tails, &part_of); + let results: Vec> = std::thread::scope(|s| { + let hs: Vec<_> = (0..threads) + .map(|ti| { + let (lon_f, lat_f) = (&lon_f, &lat_f); + s.spawn(move || -> io::Result<()> { + let (start, end) = (ti as u64 * per, ((ti as u64 + 1) * per).min(n)); + let (mut lb, mut tb) = (vec![0f64; block], vec![0f64; block]); + let mut bufs: Vec> = (0..used_parts) + .map(|_| Vec::with_capacity(FLUSH + 8)) + .collect(); + let mut off = start; + while off < end { + let k = (block as u64).min(end - off) as usize; + read_f64_at(lon_f, &mut lb[..k], off)?; + read_f64_at(lat_f, &mut tb[..k], off)?; + for i in 0..k { + let cx = cell_of(lb[i], x0, inv_x, g); + let cy = cell_of(tb[i], y0, inv_y, g); + let q = part_of[cy * g + cx] as usize; + let b = &mut bufs[q]; + b.extend_from_slice(&(lb[i] as f32).to_le_bytes()); + b.extend_from_slice(&(tb[i] as f32).to_le_bytes()); + if b.len() >= FLUSH { + flush_part(files, tails, q, b)?; + } + } + off += k as u64; + } + for (q, buf) in bufs.iter_mut().enumerate() { + flush_part(files, tails, q, buf)?; + } + Ok(()) + }) + }) + .collect(); + hs.into_iter().map(|h| h.join().unwrap()).collect() + }); + for r in results { + r?; + } + } + eprintln!(" pass 2 (scatter): {:.1}s", t_p2.elapsed().as_secs_f64()); + + // ---- Pass 3: counting-sort each partition in RAM → final columns --------- + let t_p3 = Instant::now(); + // Partitions are independent and land in disjoint, known output ranges + // (offsets[cell_lo[q]]), so workers sort them in parallel and write each + // straight to its slice via a positional write — no ordering barrier. + let out_lon = File::create(out_prefix.with_extension("lon.f32"))?; + let out_lat = File::create(out_prefix.with_extension("lat.f32"))?; + out_lon.set_len(total * 4)?; + out_lat.set_len(total * 4)?; + let next = AtomicUsize::new(0); + let max_partition = AtomicUsize::new(0); + let results: Vec> = std::thread::scope(|s| { + let hs: Vec<_> = (0..threads) + .map(|_| { + let (next, max_partition) = (&next, &max_partition); + let (out_lon, out_lat, tmp, offsets, cell_lo) = + (&out_lon, &out_lat, &tmp, &offsets, &cell_lo); + s.spawn(move || -> io::Result<()> { + loop { + let q = next.fetch_add(1, Ordering::Relaxed); + if q >= used_parts { + return Ok(()); + } + let (b_lo, b_hi) = (cell_lo[q], cell_lo[q + 1]); + let base = offsets[b_lo]; + let m = (offsets[b_hi] - base) as usize; + max_partition.fetch_max(m, Ordering::Relaxed); + let mut raw = vec![0u8; m * 8]; + File::open(&tmp[q])?.read_exact(&mut raw)?; + let pairs = unsafe { + std::slice::from_raw_parts(raw.as_ptr() as *const f32, m * 2) + }; + let span = b_hi - b_lo; + let mut cur = vec![0u32; span + 1]; + for i in 0..m { + let cx = cell_of(pairs[2 * i] as f64, x0, inv_x, g); + let cy = cell_of(pairs[2 * i + 1] as f64, y0, inv_y, g); + cur[(cy * g + cx) - b_lo + 1] += 1; + } + for k in 0..span { + cur[k + 1] += cur[k]; + } + let (mut slon, mut slat) = (vec![0f32; m], vec![0f32; m]); + for i in 0..m { + let cx = cell_of(pairs[2 * i] as f64, x0, inv_x, g); + let cy = cell_of(pairs[2 * i + 1] as f64, y0, inv_y, g); + let slot = &mut cur[(cy * g + cx) - b_lo]; + slon[*slot as usize] = pairs[2 * i]; + slat[*slot as usize] = pairs[2 * i + 1]; + *slot += 1; + } + out_lon.write_all_at( + unsafe { + std::slice::from_raw_parts(slon.as_ptr() as *const u8, m * 4) + }, + base * 4, + )?; + out_lat.write_all_at( + unsafe { + std::slice::from_raw_parts(slat.as_ptr() as *const u8, m * 4) + }, + base * 4, + )?; + std::fs::remove_file(&tmp[q])?; + } + }) + }) + .collect(); + hs.into_iter().map(|h| h.join().unwrap()).collect() + }); + for r in results { + r?; + } + let max_partition = max_partition.load(Ordering::Relaxed) as u64; + eprintln!(" pass 3 (sort+write):{:.1}s", t_p3.elapsed().as_secs_f64()); + + // ---- Index file: header + cumulative offsets ---------------------------- + let mut idx = BufWriter::new(File::create(out_prefix.with_extension("idx"))?); + idx.write_all(MAGIC)?; + idx.write_all(&(g as u32).to_le_bytes())?; + idx.write_all(&0u32.to_le_bytes())?; + for v in [x0, x1, y0, y1] { + idx.write_all(&v.to_le_bytes())?; + } + idx.write_all(&total.to_le_bytes())?; + let off_bytes = + unsafe { std::slice::from_raw_parts(offsets.as_ptr() as *const u8, offsets.len() * 8) }; + idx.write_all(off_bytes)?; + idx.flush()?; + + Ok(SortStats { + n, + g, + partitions: used_parts, + max_partition, + }) +} + +#[cfg(test)] +mod tests { + use super::*; + use std::io::Write; + + fn tmp_f64(name: &str) -> std::path::PathBuf { + let p = std::env::temp_dir().join(format!("osmsort_{}_{}", std::process::id(), name)); + File::create(&p) + .unwrap() + .write_all(&1.0f64.to_le_bytes()) + .unwrap(); + p + } + + #[test] + fn invalid_grid_and_partitions_error_not_panic() { + let (lon, lat) = (tmp_f64("lon"), tmp_f64("lat")); + let out = std::env::temp_dir().join(format!("osmsort_{}_idx", std::process::id())); + let call = |g, p| spatial_sort(&lon, &lat, &out, g, -180.0, 180.0, -90.0, 90.0, p); + assert_eq!(call(0, 4).unwrap_err().kind(), io::ErrorKind::InvalidInput); + assert_eq!(call(2, 0).unwrap_err().kind(), io::ErrorKind::InvalidInput); + // A sane call succeeds. + assert!(call(2, 4).is_ok()); + } +} diff --git a/examples/osm/osmium-rs/src/sort_main.rs b/examples/osm/osmium-rs/src/sort_main.rs new file mode 100644 index 00000000..d2e464e7 --- /dev/null +++ b/examples/osm/osmium-rs/src/sort_main.rs @@ -0,0 +1,105 @@ +//! `osm-sort` — build the spatial index (grid-bucketed f32 columns + offsets) +//! from canonical f64 lon/lat columns. +//! +//! ```text +//! osm-sort [--grid G] [--partitions P] \ +//! [--extent x0 x1 y0 y1] +//! ``` +//! Writes .lon.f32, .lat.f32, .idx. + +use std::path::Path; +use std::process::ExitCode; +use std::time::Instant; + +use osmpbf_nodes::sort::spatial_sort; + +const USAGE: &str = + "usage: osm-sort [--grid G] [--partitions P] [--extent x0 x1 y0 y1]"; + +fn usage(msg: &str) -> ExitCode { + eprintln!("error: {msg}\n{USAGE}"); + ExitCode::FAILURE +} + +fn main() -> ExitCode { + let args: Vec = std::env::args().collect(); + let mut pos: Vec<&str> = Vec::new(); + let mut g = 8192usize; + let mut parts = 512usize; + // Default to the full planet extent; nodes fall within it. + let (mut x0, mut x1, mut y0, mut y1) = (-180.0, 180.0, -90.0, 90.0); + let mut i = 1; + // Every option reads its values through `args.get(..).parse()`, so a flag at + // the end of argv or a non-numeric value is a clean usage error — never an + // out-of-bounds index or `unwrap` panic. + while i < args.len() { + match args[i].as_str() { + "--grid" => match args.get(i + 1).and_then(|s| s.parse().ok()) { + Some(v) => { + g = v; + i += 2; + } + None => return usage("--grid needs a positive integer"), + }, + "--partitions" => match args.get(i + 1).and_then(|s| s.parse().ok()) { + Some(v) => { + parts = v; + i += 2; + } + None => return usage("--partitions needs a positive integer"), + }, + "--extent" => { + let vals: Option> = (1..=4) + .map(|k| args.get(i + k).and_then(|s| s.parse().ok())) + .collect(); + match vals { + Some(v) => { + (x0, x1, y0, y1) = (v[0], v[1], v[2], v[3]); + i += 5; + } + None => return usage("--extent needs four numbers: x0 x1 y0 y1"), + } + } + other => { + pos.push(other); + i += 1; + } + } + } + if pos.len() != 3 { + return usage("expected exactly three positional args: "); + } + eprintln!("sorting into {g}x{g} grid, {parts} partitions, extent ({x0},{x1},{y0},{y1})…"); + let t0 = Instant::now(); + match spatial_sort( + Path::new(pos[0]), + Path::new(pos[1]), + Path::new(pos[2]), + g, + x0, + x1, + y0, + y1, + parts, + ) { + Ok(s) => { + let secs = t0.elapsed().as_secs_f64(); + eprintln!( + "done: {} points, {}x{} grid, {} partitions (max {}) in {:.0}s ({:.1} M/s)", + s.n, + s.g, + s.g, + s.partitions, + s.max_partition, + secs, + s.n as f64 / secs / 1e6, + ); + println!("{}", s.n); + ExitCode::SUCCESS + } + Err(e) => { + eprintln!("error: {e}"); + ExitCode::FAILURE + } + } +} diff --git a/examples/osm/viewer.py b/examples/osm/viewer.py new file mode 100644 index 00000000..94089e4d --- /dev/null +++ b/examples/osm/viewer.py @@ -0,0 +1,202 @@ +"""Interactive browser viewer for an out-of-core OSM node scatter. + +Serves xy's real WebGL render client (`python/xy/static/index.js`) against a +disk-backed density-scatter Figure. Pan/zoom drive the same kernel protocol the +notebook widget uses (`channel.handle_message`), so every viewport is +re-aggregated from the density pyramid — screen-bounded, at 1B+ points. + +Transport is plain HTTP (the client correlates replies by `seq`, so a +request/response POST per message is sufficient — no websocket, stdlib only): + + GET / → the page (imports the client, builds a ChartView) + GET /index.js → the render client bundle + GET /init → framed initial payload (spec + buffers) + POST /msg → one client message → framed kernel reply + +Run (from the repo root, after building the client once — `node js/build.mjs`): + python examples/osm/viewer.py --out /path/to/osm-data [--port 8777] +then open http://localhost:8777/. +""" + +from __future__ import annotations + +import argparse +import json +import os +import struct +import sys +import threading +import time +from http.server import BaseHTTPRequestHandler, ThreadingHTTPServer + +# examples/osm/ → repo root is two levels up; the package lives in python/. +_REPO_ROOT = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "..")) +sys.path.insert(0, os.path.join(_REPO_ROOT, "python")) +import xy # noqa: E402 +from xy import channel # noqa: E402 +from xy._ooc import open_f64 # noqa: E402 + +# The anywidget ESM bundle. It is built (not committed — see #214) by +# `node js/build.mjs`; the viewer errors clearly below if it is missing. +STATIC = os.path.join(_REPO_ROOT, "python", "xy", "static", "index.js") + +# Client /msg bodies are small JSON; reject anything larger (defensive bound). +_MAX_MSG_BYTES = 1 << 20 + +PAGE = """ +OSM nodes — xy out-of-core scatter + + +
Loading OpenStreetMap nodes…
drag to pan · scroll to zoom
+
+ + +""" + + +def build_figure(out_dir: str): + xcol = open_f64(os.path.join(out_dir, "osm_lon.f64")) + ycol = open_f64(os.path.join(out_dir, "osm_lat.f64")) + print(f"loaded {len(xcol):,} nodes from disk (out-of-core)", flush=True) + fig = xy.chart(xy.scatter(x=xcol, y=ycol, density=True)).figure() + spec, bufs = fig.build_payload_split() + spec.setdefault("interaction", {})["_transport_view_change"] = True + # Attach the Tier-3 spatial index if built (osm-sort): deep zoom then serves + # exact street-level detail from just the in-window points. + idx_prefix = os.path.join(out_dir, "osm_spatial") + if os.path.exists(idx_prefix + ".idx"): + from xy._spatial import SpatialIndex + + fig.traces[0]._spatial_index = SpatialIndex.load(idx_prefix) + print(f"spatial index attached ({idx_prefix}.idx)", flush=True) + # Warm the density pyramid now (one O(N) disk scan) so the *first* pan/zoom + # is instant instead of stalling on the lazy build. Full-domain view. + t0 = time.perf_counter() + tr = fig.traces[0] + x0, x1, y0, y1 = tr.x.min, tr.x.max, tr.y.min, tr.y.max + fig.density_view(0, x0, x1, y0, y1, 1200, 900) + print(f"pyramid warmed in {time.perf_counter() - t0:.0f}s", flush=True) + rep = fig.store.memory_report() + print( + f"canonical RAM-resident: {rep['canonical_bytes']:,} B | on disk: " + f"{rep['canonical_mapped_bytes']:,} B", + flush=True, + ) + return fig, spec, bufs + + +def frame(content: dict, buffers) -> bytes: + jb = json.dumps(content).encode() + out = bytearray(struct.pack(" None: + ap = argparse.ArgumentParser() + ap.add_argument( + "--out", required=True, help="dir holding osm_lon.f64 / osm_lat.f64 (see README)" + ) + ap.add_argument("--port", type=int, default=8777) + args = ap.parse_args() + + if not os.path.exists(STATIC): + raise SystemExit( + f"render client not built: {STATIC}\n" + "build it once from the repo root: npm ci && node js/build.mjs" + ) + + fig, spec, bufs = build_figure(args.out) + init_frame = frame(spec, bufs) + lock = threading.Lock() + with open(STATIC, "rb") as f: + client_js = f.read() + + class Handler(BaseHTTPRequestHandler): + def log_message(self, format, *args): # noqa: A002 — quiet + pass + + def _send(self, body: bytes, ctype: str) -> None: + self.send_response(200) + self.send_header("Content-Type", ctype) + self.send_header("Content-Length", str(len(body))) + self.end_headers() + self.wfile.write(body) + + def do_GET(self): + if self.path == "/": + self._send(PAGE.encode(), "text/html; charset=utf-8") + elif self.path == "/index.js": + self._send(client_js, "text/javascript") + elif self.path == "/init": + self._send(init_frame, "application/octet-stream") + else: + self.send_error(404) + + def do_POST(self): + if self.path != "/msg": + self.send_error(404) + return + n = int(self.headers.get("Content-Length", 0)) + # Client messages are tiny JSON (a viewport + a few scalars); cap the + # body so a stray/hostile request can't make the demo allocate wildly. + if n > _MAX_MSG_BYTES: + self.send_error(413) + return + content = json.loads(self.rfile.read(n) or b"{}") + with lock: + reply = channel.handle_message(fig, content) + if reply is None: + self._send(b"", "application/octet-stream") + else: + msg, out = reply + self._send(frame(msg, out), "application/octet-stream") + + srv = ThreadingHTTPServer(("0.0.0.0", args.port), Handler) + print(f"\n ▶ http://localhost:{args.port}/\n", flush=True) + srv.serve_forever() + + +if __name__ == "__main__": + main()