perfect-postcode/server-rs/src/data/crime_records.rs
2026-06-25 22:29:52 +01:00

534 lines
22 KiB
Rust

//! Individual police.uk crime records (last 7 years) backing the right pane's
//! "individual crimes" list and the `/api/crime-records` endpoint.
//!
//! This table is enormous — ~500M rows, because each incident is replicated to
//! every postcode whose buffer covers it (see [`gather`](CrimeRecords::gather)),
//! so it is NOT held as a `Vec<struct>`: each field is a flat columnar
//! [`SpillVec`] (mmap-backed and kernel-reclaimable when `--spill-dir` is set),
//! small string fields are dictionary-encoded, and the parquet is pre-sorted by
//! postcode so each postcode's records are a contiguous `[start, start+count)`
//! slice located via a CSR-style offset index. Resident RSS is ~0 until records
//! are actually read.
//!
//! At ~500M rows the parquet's string columns (postcode/type/location/outcome)
//! decode to tens of GB if read whole, so the loader never materialises the
//! whole `DataFrame`: it streams the file in bounded row-count chunks (only the
//! row groups overlapping each slice are decoded) and writes each column
//! straight into its (optionally spilled) backing store via [`SpillVecBuilder`],
//! keeping the transient footprint to one chunk plus the index maps.
use std::fs::File;
use std::path::Path;
use anyhow::{bail, Context};
use lasso::{Rodeo, RodeoReader, Spur};
use polars::prelude::*;
use rustc_hash::FxHashMap;
use tracing::info;
use super::run_polars_io;
use super::spill::{SpillVec, SpillVecBuilder};
/// Rows decoded per streaming slice. `with_slice` decodes only the row groups
/// overlapping the slice, so the transient decode is roughly one chunk's worth
/// (~tens of MB at the writer's ~123k-row groups) instead of the tens-of-GB
/// whole-file `DataFrame`. The bound's only dependency on file layout is a
/// reasonable input row-group size, which our pipeline writer produces.
const CHUNK_ROWS: usize = 2_000_000;
/// A resolved view of one record (strings dereferenced from the dictionaries).
pub struct CrimeRecordView<'a> {
/// `year * 12 + (month - 1)`.
pub month_index: u32,
pub crime_type: &'a str,
pub outcome: Option<&'a str>,
pub location: Option<&'a str>,
pub lat: f32,
pub lon: f32,
}
pub struct CrimeRecords {
month: SpillVec<u32>,
ctype: SpillVec<u8>,
outcome: SpillVec<u8>,
location: SpillVec<Spur>,
lat: SpillVec<f32>,
lon: SpillVec<f32>,
/// Dictionary for `ctype` (bare crime type names, e.g. "Burglary").
crime_type_dict: Vec<String>,
/// Dictionary for `outcome`; index 0 is the empty/unknown sentinel.
outcome_dict: Vec<String>,
/// Resolver for the interned `location` strings (`""` means withheld).
location_resolver: RodeoReader,
/// Postcode → `(start, count)` into the columnar arrays (records for a
/// postcode are contiguous because the parquet is sorted by postcode).
by_postcode: FxHashMap<String, (u32, u32)>,
}
impl CrimeRecords {
#[cfg(test)]
pub fn empty() -> Self {
Self {
month: SpillVec::owned(Vec::new()),
ctype: SpillVec::owned(Vec::new()),
outcome: SpillVec::owned(Vec::new()),
location: SpillVec::owned(Vec::new()),
lat: SpillVec::owned(Vec::new()),
lon: SpillVec::owned(Vec::new()),
crime_type_dict: Vec::new(),
outcome_dict: vec![String::new()],
location_resolver: Rodeo::default().into_reader(),
by_postcode: FxHashMap::default(),
}
}
/// Number of records stored for a postcode (0 if none).
pub fn total_for(&self, postcode: &str) -> u32 {
self.by_postcode.get(postcode).map_or(0, |&(_, c)| c)
}
/// Resolve a record index to a borrowing view.
pub fn view(&self, idx: u32) -> CrimeRecordView<'_> {
let i = idx as usize;
let outcome_idx = self.outcome[i] as usize;
let outcome = self
.outcome_dict
.get(outcome_idx)
.filter(|s| !s.is_empty())
.map(String::as_str);
let location = self.location_resolver.resolve(&self.location[i]);
CrimeRecordView {
month_index: self.month[i],
crime_type: self
.crime_type_dict
.get(self.ctype[i] as usize)
.map_or("", String::as_str),
outcome,
location: (!location.is_empty()).then_some(location),
lat: self.lat[i],
lon: self.lon[i],
}
}
/// Record indices across `postcodes`, newest first, optionally restricted to
/// months `>= since_month`. These are exactly the incidents counted for the
/// selected postcodes — for a single postcode that is its precise incident
/// list; for a multi-postcode selection a boundary incident counted for
/// several postcodes appears once per postcode, matching the count. We do not
/// de-duplicate because police.uk snaps many genuinely distinct incidents
/// (especially anti-social behaviour) to the same point/month and provides no
/// per-incident id to tell a true duplicate from two real incidents apart.
pub fn gather(&self, postcodes: &[&str], since_month: Option<u32>) -> Vec<u32> {
let month = self.month.as_slice();
let mut out: Vec<u32> = Vec::new();
for pc in postcodes {
let Some(&(start, count)) = self.by_postcode.get(*pc) else {
continue;
};
for i in start..start + count {
if since_month.map_or(true, |s| month[i as usize] >= s) {
out.push(i);
}
}
}
out.sort_unstable_by(|&a, &b| month[b as usize].cmp(&month[a as usize]));
out
}
pub fn load(path: &Path, spill_dir: Option<&Path>) -> anyhow::Result<Self> {
run_polars_io(|| Self::load_inner(path, spill_dir, CHUNK_ROWS))
}
fn load_inner(
path: &Path,
spill_dir: Option<&Path>,
chunk_rows: usize,
) -> anyhow::Result<Self> {
// A zero chunk size would loop forever; the public entry point passes the
// const, but tests parameterise this to exercise chunk boundaries.
let chunk_rows = chunk_rows.max(1);
info!("Loading crime records from {}", path.display());
// Read the footer once for the row count, and keep it to hand to every
// per-chunk reader so the 2.9MB metadata is never re-parsed.
let metadata = {
let file = File::open(path).with_context(|| {
format!("Failed to open crime-records parquet at {}", path.display())
})?;
ParquetReader::new(file)
.get_metadata()
.with_context(|| {
format!("Failed to read crime-records parquet metadata at {}", path.display())
})?
.clone()
};
let n = metadata.num_rows;
// Record indices are stored as `u32` (and `by_postcode` holds `(start,
// count)` as `u32`), so the table must fit in that index space.
if n > u32::MAX as usize {
bail!("crime-records parquet has {n} rows, exceeding the u32 record-index limit");
}
// Columns that, when spilling, are written straight into mmap-backed files
// as we stream — so the ~9GB of columnar data never lands on the heap.
let mut month = SpillVecBuilder::<u32>::with_len(n, spill_dir, "crime_month")?;
let mut ctype = SpillVecBuilder::<u8>::with_len(n, spill_dir, "crime_ctype")?;
let mut outcome = SpillVecBuilder::<u8>::with_len(n, spill_dir, "crime_outcome")?;
let mut location = SpillVecBuilder::<Spur>::with_len(n, spill_dir, "crime_location")?;
let mut lat = SpillVecBuilder::<f32>::with_len(n, spill_dir, "crime_lat")?;
let mut lon = SpillVecBuilder::<f32>::with_len(n, spill_dir, "crime_lon")?;
let mut crime_type_dict: Vec<String> = Vec::new();
let mut type_index: FxHashMap<String, u8> = FxHashMap::default();
// Outcome index 0 is the empty/unknown sentinel.
let mut outcome_dict: Vec<String> = vec![String::new()];
let mut outcome_index: FxHashMap<String, u8> = FxHashMap::default();
let mut rodeo = Rodeo::default();
let empty_spur = rodeo.get_or_intern("");
let mut by_postcode: FxHashMap<String, (u32, u32)> = FxHashMap::default();
let mut cur_pc: Option<String> = None;
let mut cur_start: u32 = 0;
// Absolute record index across all chunks; drives both the CSR index and
// the column builders' write order.
let mut global_row: u32 = 0;
let columns: Vec<String> = ["postcode", "month_index", "crime_type", "location", "outcome", "lat", "lon"]
.iter()
.map(|s| s.to_string())
.collect();
let mut offset = 0usize;
while offset < n {
let len = chunk_rows.min(n - offset);
let file = File::open(path).with_context(|| {
format!("Failed to open crime-records parquet at {}", path.display())
})?;
let mut reader = ParquetReader::new(file);
reader.set_metadata(metadata.clone());
// `with_slice` only decodes the row groups overlapping `[offset,
// offset+len)` (the file is memory-mapped, so untouched groups are
// never faulted in), capping the transient decode to one chunk.
let df = reader
.with_columns(Some(columns.clone()))
.with_slice(Some((offset, len)))
.finish()
.with_context(|| {
format!(
"Failed to read crime-records rows [{offset}, {}) from {}",
offset + len,
path.display()
)
})?;
let postcode_col = df
.column("postcode")
.context("crime-records parquet missing 'postcode'")?
.str()
.context("'postcode' is not a string")?;
let month_col = df
.column("month_index")
.context("crime-records parquet missing 'month_index'")?
.cast(&DataType::Int32)
.context("'month_index' not castable to i32")?;
let month_ca = month_col.i32().context("'month_index' is not i32")?;
// Months are `year*12 + month0` (~24_000), always positive. A null or
// non-positive value means a corrupt parquet; fail loudly rather than
// silently clamping it to 0 and later rendering it as "0000-01".
if month_ca.null_count() > 0 {
bail!("crime-records 'month_index' has null values (corrupt parquet)");
}
match month_ca.min() {
Some(m) if m > 0 => {}
_ => bail!("crime-records 'month_index' must be a positive year*12+month index"),
}
let type_col = df
.column("crime_type")
.context("crime-records parquet missing 'crime_type'")?
.str()
.context("'crime_type' is not a string")?;
let location_col = df
.column("location")
.context("crime-records parquet missing 'location'")?
.str()
.context("'location' is not a string")?;
let outcome_col = df
.column("outcome")
.context("crime-records parquet missing 'outcome'")?
.str()
.context("'outcome' is not a string")?;
let lat_col = df
.column("lat")
.context("crime-records parquet missing 'lat'")?
.cast(&DataType::Float32)?;
let lat_ca = lat_col.f32().context("'lat' is not f32")?;
let lon_col = df
.column("lon")
.context("crime-records parquet missing 'lon'")?
.cast(&DataType::Float32)?;
let lon_ca = lon_col.f32().context("'lon' is not f32")?;
let height = df.height();
for row in 0..height {
// CSR index: the parquet is sorted by postcode, so a change in the
// postcode value (across chunk boundaries too) closes the previous
// run and opens a new one.
let pc = postcode_col
.get(row)
.with_context(|| {
format!("crime-records row {} has null postcode", offset + row)
})?
.trim();
if cur_pc.as_deref() != Some(pc) {
if let Some(prev) = cur_pc.take() {
by_postcode.insert(prev, (cur_start, global_row - cur_start));
}
cur_pc = Some(pc.to_string());
cur_start = global_row;
}
month.push(month_ca.get(row).unwrap() as u32);
let ty = type_col.get(row).unwrap_or("");
let ty_id = match type_index.get(ty) {
Some(&id) => id,
None => {
let id = u8::try_from(crime_type_dict.len())
.context("more than 256 distinct crime types")?;
crime_type_dict.push(ty.to_string());
type_index.insert(ty.to_string(), id);
id
}
};
ctype.push(ty_id);
let oc = outcome_col.get(row).unwrap_or("");
let oc_id = if oc.is_empty() {
0
} else {
match outcome_index.get(oc) {
Some(&id) => id,
None => {
let id = u8::try_from(outcome_dict.len())
.context("more than 256 distinct outcomes")?;
outcome_dict.push(oc.to_string());
outcome_index.insert(oc.to_string(), id);
id
}
}
};
outcome.push(oc_id);
let loc = location_col.get(row).unwrap_or("");
location.push(if loc.is_empty() {
empty_spur
} else {
rodeo.get_or_intern(loc)
});
lat.push(lat_ca.get(row).unwrap_or(f32::NAN));
lon.push(lon_ca.get(row).unwrap_or(f32::NAN));
global_row += 1;
}
offset += len;
}
if let Some(prev) = cur_pc.take() {
by_postcode.insert(prev, (cur_start, global_row - cur_start));
}
debug_assert_eq!(global_row as usize, n, "streamed fewer rows than the parquet declares");
let records = Self {
month: month.finish()?,
ctype: ctype.finish()?,
outcome: outcome.finish()?,
location: location.finish()?,
lat: lat.finish()?,
lon: lon.finish()?,
crime_type_dict,
outcome_dict,
location_resolver: rodeo.into_reader(),
by_postcode,
};
info!(
records = n,
postcodes = records.by_postcode.len(),
crime_types = records.crime_type_dict.len(),
outcomes = records.outcome_dict.len(),
"Crime records loaded"
);
Ok(records)
}
}
#[cfg(test)]
mod tests {
use super::*;
fn write_fixture(path: &Path) {
// Two postcodes, postcode-sorted. AA1 1AA has 3 records across two
// months (a null outcome and a null location), BB2 2BB has 1.
let mut df = df!(
"postcode" => ["AA1 1AA", "AA1 1AA", "AA1 1AA", "BB2 2BB"],
"month_index" => [24300i32, 24300, 24290, 24305],
"crime_type" => ["Burglary", "Burglary", "Robbery", "Drugs"],
"location" => [Some("On or near A St"), Some("On or near A St"), None, Some("On or near B Rd")],
"outcome" => [Some("Under investigation"), None, Some("Court result"), None],
"lat" => [51.5f32, 51.5, 51.6, 52.0],
"lon" => [-0.1f32, -0.1, -0.2, -1.0],
)
.unwrap();
let mut file = std::fs::File::create(path).unwrap();
ParquetWriter::new(&mut file).finish(&mut df).unwrap();
}
#[test]
fn loads_indexes_and_gathers() {
let dir = std::env::temp_dir().join(format!("crimerec-{}", std::process::id()));
std::fs::create_dir_all(&dir).unwrap();
let path = dir.join("records.parquet");
write_fixture(&path);
let recs = CrimeRecords::load(&path, None).unwrap();
assert_eq!(recs.total_for("AA1 1AA"), 3);
assert_eq!(recs.total_for("BB2 2BB"), 1);
assert_eq!(recs.total_for("ZZ9 9ZZ"), 0);
// Newest-first across the two postcodes.
let all = recs.gather(&["AA1 1AA", "BB2 2BB"], None);
assert_eq!(all.len(), 4);
let months: Vec<u32> = all.iter().map(|&i| recs.view(i).month_index).collect();
assert_eq!(months, vec![24305, 24300, 24300, 24290]);
// `since` window filter (keep months >= 24300).
assert_eq!(recs.gather(&["AA1 1AA"], Some(24300)).len(), 2);
// String resolution + null handling.
let robbery = all
.iter()
.map(|&i| recs.view(i))
.find(|v| v.crime_type == "Robbery")
.unwrap();
assert_eq!(robbery.outcome, Some("Court result"));
assert_eq!(robbery.location, None); // null location → None
// Two records have a null outcome (an AA1 Burglary and the BB2 Drugs).
let null_outcomes = all
.iter()
.map(|&i| recs.view(i))
.filter(|v| v.outcome.is_none())
.count();
assert_eq!(null_outcomes, 2);
std::fs::remove_dir_all(&dir).ok();
}
/// The CSR per-postcode index and the column builders must compose correctly
/// across streaming chunk boundaries — including a postcode run split between
/// two chunks. Forces `chunk_rows = 2` over the 4-row fixture so AA1 1AA's
/// three records straddle the boundary (rows 0,1 in chunk 0; row 2 in chunk 1)
/// and is exercised both heap-backed (no spill) and mmap-backed (spill).
#[test]
fn streams_across_chunk_boundaries() {
let base = std::env::temp_dir().join(format!("crimerec-chunk-{}", std::process::id()));
std::fs::create_dir_all(&base).unwrap();
let path = base.join("records.parquet");
write_fixture(&path);
let spill = base.join("spill");
std::fs::create_dir_all(&spill).unwrap();
for spill_dir in [None, Some(spill.as_path())] {
let recs = CrimeRecords::load_inner(&path, spill_dir, 2).unwrap();
// Counts match regardless of how the runs were split across chunks.
assert_eq!(recs.total_for("AA1 1AA"), 3);
assert_eq!(recs.total_for("BB2 2BB"), 1);
assert_eq!(recs.total_for("ZZ9 9ZZ"), 0);
// Full gather, newest-first, identical to the single-chunk load.
let all = recs.gather(&["AA1 1AA", "BB2 2BB"], None);
assert_eq!(all.len(), 4);
let months: Vec<u32> = all.iter().map(|&i| recs.view(i).month_index).collect();
assert_eq!(months, vec![24305, 24300, 24300, 24290]);
// The run that straddled the boundary still resolves its strings.
let robbery = all
.iter()
.map(|&i| recs.view(i))
.find(|v| v.crime_type == "Robbery")
.unwrap();
assert_eq!(robbery.outcome, Some("Court result"));
assert_eq!(robbery.location, None);
}
std::fs::remove_dir_all(&base).ok();
}
/// Peak/resident RSS in MiB from `/proc/self/status` (Linux only).
fn rss_mib() -> (f64, f64) {
let status = std::fs::read_to_string("/proc/self/status").unwrap_or_default();
let field = |key: &str| -> f64 {
status
.lines()
.find(|l| l.starts_with(key))
.and_then(|l| l.split_whitespace().nth(1))
.and_then(|kb| kb.parse::<f64>().ok())
.map_or(0.0, |kb| kb / 1024.0)
};
(field("VmHWM:"), field("VmRSS:"))
}
/// Manual, real-data smoke test: load the actual ~500M-row parquet and report
/// peak RSS, proving the streaming + spill load completes without the
/// tens-of-GB `DataFrame` materialisation that OOMed the old `.collect()`.
///
/// Run with:
/// PPC_REAL_CRIME_RECORDS=/path/to/crime_records.parquet \
/// cargo test --bins -- --ignored --nocapture real_crime_records_load_is_bounded
#[test]
#[ignore = "needs the full crime_records.parquet; run manually"]
fn real_crime_records_load_is_bounded() {
let path = std::env::var("PPC_REAL_CRIME_RECORDS")
.unwrap_or_else(|_| "../property-data/crime_records.parquet".to_string());
let path = Path::new(&path);
if !path.exists() {
eprintln!("skipping: {} not found", path.display());
return;
}
let spill = std::env::var("PPC_REAL_SPILL")
.unwrap_or_else(|_| "../.tmp/crime-spill-realtest".to_string());
let spill = Path::new(&spill);
std::fs::create_dir_all(spill).unwrap();
let (hwm_before, _rss_before) = rss_mib();
let start = std::time::Instant::now();
let recs = CrimeRecords::load(path, Some(spill)).unwrap();
let elapsed = start.elapsed();
let (hwm_after, rss_after) = rss_mib();
let total: u64 = recs.by_postcode.values().map(|&(_, c)| c as u64).sum();
eprintln!(
"loaded {} records across {} postcodes in {:.1}s | RSS peak {:.0}->{:.0} MiB (Δ{:.0}) resident now {:.0} MiB",
total,
recs.by_postcode.len(),
elapsed.as_secs_f64(),
hwm_before,
hwm_after,
hwm_after - hwm_before,
rss_after,
);
assert!(recs.by_postcode.len() > 0, "expected at least one postcode");
assert!(total > 0, "expected at least one record");
// The old `.collect()` decoded all rows' string columns at once (tens of
// GB). Streaming must keep the peak growth far below that; a generous 20GiB
// ceiling still proves we never materialise the whole file.
assert!(
hwm_after - hwm_before < 20_480.0,
"peak RSS grew by {:.0} MiB during load — streaming/spill not bounding memory",
hwm_after - hwm_before
);
std::fs::remove_dir_all(spill).ok();
}
}