perfect-postcode/server-rs/src/data/poi.rs
Andras Schmelczer e2b85fe819
Some checks failed
CI / Check (push) Failing after 2m47s
Build and publish Docker image / build-and-push (push) Successful in 5m37s
lgtm
2026-06-26 16:48:20 +01:00

678 lines
24 KiB
Rust

use std::collections::{HashMap, HashSet};
use std::path::Path;
use anyhow::{bail, Context};
use polars::frame::DataFrame;
use polars::lazy::frame::LazyFrame;
use polars::prelude::*;
use rustc_hash::FxHashSet;
use serde::Serialize;
use tracing::info;
use crate::features::POI_GROUP_ORDER;
use crate::utils::{generate_priorities, InternedColumn};
#[derive(Serialize, Clone)]
pub struct POICategoryGroup {
pub name: String,
pub categories: Vec<String>,
}
const GROCERY_DASHBOARD_CATEGORIES: &[&str] = &[
"Supermarket",
"Convenience Store",
"Bakery",
"Greengrocer",
"Aldi",
"Amazon",
"Asda",
"Booths",
"Budgens",
"Centra",
"Co-op",
"COOK",
"Costco",
"Dunnes Stores",
"Farmfoods",
"Heron Foods",
"Iceland",
"Lidl",
"Makro",
"M&S",
"Morrisons",
"Planet Organic",
"Sainsbury's",
"Spar",
"Tesco",
"The Food Warehouse",
"Waitrose",
"Whole Foods Market",
];
const DASHBOARD_POI_GROUPS: &[(&str, &[&str])] = &[
(
"Public Transport",
&[
"Rail station",
"Tube station",
"DLR station",
"Tram & Metro stop",
"Bus station",
"Bus stop",
"Airport",
],
),
("Groceries", GROCERY_DASHBOARD_CATEGORIES),
("Food & Drink", &["Café", "Restaurant", "Pub", "Fast Food"]),
("Green Space", &["Park", "Playground"]),
(
"Education",
&[
"Nursery school",
"Primary school",
"Secondary school",
"All-through school",
"Sixth form",
"Further education college",
"University",
"Special school",
"School",
],
),
(
"Health",
&["GP Surgery", "Pharmacy", "Dentist", "Hospital", "Clinic"],
),
(
"Leisure",
&[
"Gym & Fitness",
"Sports Centre",
"Cinema",
"Theatre",
"Library",
],
),
(
"Practical",
&["Post Office", "Bank", "EV Charging", "Fuel Station"],
),
];
fn add_category_filter_index(
category_values: &[String],
category: &str,
selected: &mut FxHashSet<u16>,
) {
if let Some(pos) = category_values.iter().position(|value| value == category) {
selected.insert(pos as u16);
}
}
fn canonical_poi_category(category: &str) -> &str {
match category {
"Allendale Co-operative Society"
| "Central England Co-operative"
| "Channel Islands Co-operative Society"
| "Chelmsford Star Co-operative Society"
| "Clydebank Co-operative"
| "Coniston Co-operative Society"
| "Co-op Food"
| "East of England Co-operative"
| "Heart of England Co-operative"
| "Langdale Co-operative Society"
| "Lincolnshire Co-operative"
| "Midcounties Co-operative"
| "Scottish Midland Co-operative"
| "Tamworth Co-operative Society"
| "The Co-operative Food"
| "The Co-operative Food PFS"
| "The Co-operative Group"
| "The Radstock Co-operative Society"
| "The Southern Co-operative" => "Co-op",
_ => category,
}
}
/// Categories the pipeline emits for the GIAS-derived school POIs. A bare
/// `poi=School` URL (predating the per-phase split) is expanded to all of these
/// so bookmarked links keep showing schools.
const SCHOOL_CATEGORY_ALIASES: &[&str] = &[
"Nursery school",
"Primary school",
"Secondary school",
"All-through school",
"Sixth form",
"Further education college",
"University",
"Special school",
"School",
];
pub fn resolve_poi_category_filter(category_values: &[String], categories: &str) -> FxHashSet<u16> {
let mut selected = FxHashSet::default();
for part in categories.split(',') {
let category = canonical_poi_category(part.trim());
if category.is_empty() {
continue;
}
if category == "School" {
for alias in SCHOOL_CATEGORY_ALIASES {
add_category_filter_index(category_values, alias, &mut selected);
}
continue;
}
add_category_filter_index(category_values, category, &mut selected);
}
selected
}
/// Metadata for state-funded school POIs (sourced from the DfE GIAS register).
/// Every field is optional because GIAS does not populate every column for every
/// establishment type (e.g. nurseries have no sixth form, FE colleges no FSM).
#[derive(Serialize, Clone, Default)]
pub struct SchoolMetadata {
#[serde(skip_serializing_if = "Option::is_none")]
pub phase: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub r#type: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub type_group: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub age_range: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub gender: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub religious_character: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub admissions_policy: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub nursery_provision: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub sixth_form: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub capacity: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub pupils: Option<u32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub fsm_percent: Option<f32>,
#[serde(skip_serializing_if = "Option::is_none")]
pub trust: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub address: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub postcode: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub local_authority: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub website: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub telephone: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub head_name: Option<String>,
#[serde(skip_serializing_if = "Option::is_none")]
pub ofsted_rating: Option<String>,
}
pub struct POIData {
/// Contiguous buffer holding all POI ID strings end-to-end.
id_buffer: String,
/// Byte offset into `id_buffer` where each row's ID starts.
id_offsets: Vec<u32>,
/// Length in bytes of each row's ID.
id_lengths: Vec<u16>,
pub group: InternedColumn,
pub category: InternedColumn,
pub icon_category: InternedColumn,
pub name: Vec<String>,
pub lat: Vec<f32>,
pub lng: Vec<f32>,
pub emoji: InternedColumn,
/// Deterministic pseudo-random priority per row, used to select a spatially
/// uniform subset when the POI count exceeds the per-request limit.
/// Computed once at load time so the same POIs are always chosen for a given viewport.
pub priority: Vec<u32>,
/// Indirection table: row idx → index into `school_meta`, or u32::MAX when
/// the POI is not a school. Keeps the per-row overhead at 4 bytes regardless
/// of how many school metadata fields we carry.
school_meta_idx: Vec<u32>,
school_meta: Vec<SchoolMetadata>,
}
impl POIData {
/// Get the ID string for a given row.
pub fn id(&self, row: usize) -> &str {
let offset = self.id_offsets[row] as usize;
let length = self.id_lengths[row] as usize;
&self.id_buffer[offset..offset + length]
}
/// Get the school metadata for a given row, or None if not a school.
pub fn school(&self, row: usize) -> Option<&SchoolMetadata> {
let idx = self.school_meta_idx[row];
if idx == u32::MAX {
None
} else {
Some(&self.school_meta[idx as usize])
}
}
}
fn extract_str_col(df: &DataFrame, name: &str) -> anyhow::Result<Vec<String>> {
let column = df
.column(name)
.with_context(|| format!("Missing column '{name}' in POI data"))?;
let string_column = column
.str()
.with_context(|| format!("Column '{name}' is not a string column"))?;
string_column
.into_iter()
.enumerate()
.map(|(row, value)| {
value
.map(ToString::to_string)
.with_context(|| format!("Column '{name}' has null at row {row}"))
})
.collect()
}
fn extract_f32_col(df: &DataFrame, name: &str) -> anyhow::Result<Vec<f32>> {
let column = df
.column(name)
.with_context(|| format!("Missing column '{name}' in POI data"))?;
let cast = column
.cast(&DataType::Float32)
.with_context(|| format!("Failed to cast column '{name}' to Float32"))?;
let float_column = cast
.f32()
.with_context(|| format!("Column '{name}' is not a float32 column"))?;
float_column
.into_iter()
.enumerate()
.map(|(row, value)| value.with_context(|| format!("Column '{name}' has null at row {row}")))
.collect()
}
/// Read an optional string column. Returns None when the column itself is missing
/// (older POI parquets without the school_* extension); returns Some(vec) of
/// length row_count where each entry is None for null cells.
fn extract_optional_str_col(
df: &DataFrame,
name: &str,
) -> anyhow::Result<Option<Vec<Option<String>>>> {
let column = match df.column(name) {
Ok(column) => column,
Err(_) => return Ok(None),
};
let string_column = column
.str()
.with_context(|| format!("Column '{name}' is not a string column"))?;
Ok(Some(
string_column
.into_iter()
.map(|value| value.map(ToString::to_string))
.collect(),
))
}
fn extract_optional_u32_col(
df: &DataFrame,
name: &str,
) -> anyhow::Result<Option<Vec<Option<u32>>>> {
let column = match df.column(name) {
Ok(column) => column,
Err(_) => return Ok(None),
};
let cast = column
.cast(&DataType::Int64)
.with_context(|| format!("Failed to cast column '{name}' to Int64"))?;
let int_column = cast
.i64()
.with_context(|| format!("Column '{name}' is not an integer column"))?;
Ok(Some(
int_column
.into_iter()
.map(|value| value.and_then(|v| if v < 0 { None } else { Some(v as u32) }))
.collect(),
))
}
fn extract_optional_f32_col(
df: &DataFrame,
name: &str,
) -> anyhow::Result<Option<Vec<Option<f32>>>> {
let column = match df.column(name) {
Ok(column) => column,
Err(_) => return Ok(None),
};
let cast = column
.cast(&DataType::Float32)
.with_context(|| format!("Failed to cast column '{name}' to Float32"))?;
let float_column = cast
.f32()
.with_context(|| format!("Column '{name}' is not a float32 column"))?;
Ok(Some(float_column.into_iter().collect()))
}
fn build_school_meta(
row_count: usize,
df: &DataFrame,
) -> anyhow::Result<(Vec<u32>, Vec<SchoolMetadata>)> {
let phase = extract_optional_str_col(df, "school_phase")?;
if phase.is_none() {
// POI parquet predates the school metadata extension — record an empty
// table and a sentinel-filled index, so callers transparently see None.
return Ok((vec![u32::MAX; row_count], Vec::new()));
}
let phase = phase.unwrap();
let r#type = extract_optional_str_col(df, "school_type")?.unwrap_or_default();
let type_group = extract_optional_str_col(df, "school_type_group")?.unwrap_or_default();
let age_range = extract_optional_str_col(df, "school_age_range")?.unwrap_or_default();
let gender = extract_optional_str_col(df, "school_gender")?.unwrap_or_default();
let religious_character =
extract_optional_str_col(df, "school_religious_character")?.unwrap_or_default();
let admissions_policy =
extract_optional_str_col(df, "school_admissions_policy")?.unwrap_or_default();
let nursery_provision =
extract_optional_str_col(df, "school_nursery_provision")?.unwrap_or_default();
let sixth_form = extract_optional_str_col(df, "school_sixth_form")?.unwrap_or_default();
let capacity = extract_optional_u32_col(df, "school_capacity")?.unwrap_or_default();
let pupils = extract_optional_u32_col(df, "school_pupils")?.unwrap_or_default();
let fsm_percent = extract_optional_f32_col(df, "school_fsm_percent")?.unwrap_or_default();
let trust = extract_optional_str_col(df, "school_trust")?.unwrap_or_default();
let address = extract_optional_str_col(df, "school_address")?.unwrap_or_default();
let postcode = extract_optional_str_col(df, "school_postcode")?.unwrap_or_default();
let local_authority =
extract_optional_str_col(df, "school_local_authority")?.unwrap_or_default();
let website = extract_optional_str_col(df, "school_website")?.unwrap_or_default();
let telephone = extract_optional_str_col(df, "school_telephone")?.unwrap_or_default();
let head_name = extract_optional_str_col(df, "school_head_name")?.unwrap_or_default();
let ofsted_rating = extract_optional_str_col(df, "school_ofsted_rating")?.unwrap_or_default();
let fetch_str = |col: &Vec<Option<String>>, row: usize| -> Option<String> {
col.get(row).cloned().flatten()
};
let fetch_u32 =
|col: &Vec<Option<u32>>, row: usize| -> Option<u32> { col.get(row).copied().flatten() };
let fetch_f32 =
|col: &Vec<Option<f32>>, row: usize| -> Option<f32> { col.get(row).copied().flatten() };
let mut idx = vec![u32::MAX; row_count];
let mut meta = Vec::new();
for (row, meta_idx) in idx.iter_mut().enumerate().take(row_count) {
let type_group_val = fetch_str(&type_group, row);
let type_val = fetch_str(&r#type, row);
// type_group is present for every GIAS row, so use it as the sentinel
// for "this POI is a school" — matches the pipeline guarantee.
if type_group_val.is_none() && type_val.is_none() {
continue;
}
*meta_idx = meta.len() as u32;
meta.push(SchoolMetadata {
phase: fetch_str(&phase, row),
r#type: type_val,
type_group: type_group_val,
age_range: fetch_str(&age_range, row),
gender: fetch_str(&gender, row),
religious_character: fetch_str(&religious_character, row),
admissions_policy: fetch_str(&admissions_policy, row),
nursery_provision: fetch_str(&nursery_provision, row),
sixth_form: fetch_str(&sixth_form, row),
capacity: fetch_u32(&capacity, row),
pupils: fetch_u32(&pupils, row),
fsm_percent: fetch_f32(&fsm_percent, row),
trust: fetch_str(&trust, row),
address: fetch_str(&address, row),
postcode: fetch_str(&postcode, row),
local_authority: fetch_str(&local_authority, row),
website: fetch_str(&website, row),
telephone: fetch_str(&telephone, row),
head_name: fetch_str(&head_name, row),
ofsted_rating: fetch_str(&ofsted_rating, row),
});
}
Ok((idx, meta))
}
impl POIData {
pub fn load(parquet_path: &Path) -> anyhow::Result<Self> {
super::run_polars_io(|| Self::load_inner(parquet_path))
}
fn load_inner(parquet_path: &Path) -> anyhow::Result<Self> {
info!("Loading POI data from {:?}...", parquet_path);
let parquet_path = PlRefPath::try_from_path(parquet_path)
.context("Failed to normalize POI parquet path")?;
let df = LazyFrame::scan_parquet(parquet_path, Default::default())
.context("Failed to scan POI parquet")?
.collect()
.context("Failed to read POI parquet")?;
let row_count = df.height();
info!("Loaded {} POIs", row_count);
let id_raw: Vec<String> = extract_str_col(&df, "id")?;
let name = extract_str_col(&df, "name")?;
let category_raw: Vec<String> = extract_str_col(&df, "category")?
.into_iter()
.map(|category| canonical_poi_category(&category).to_string())
.collect();
let group_raw = extract_str_col(&df, "group")?;
let lat = extract_f32_col(&df, "lat")?;
let lng = extract_f32_col(&df, "lng")?;
let emoji_raw = extract_str_col(&df, "emoji")?;
let icon_category_raw: Vec<String> = extract_str_col(&df, "icon_category")?
.into_iter()
.map(|category| canonical_poi_category(&category).to_string())
.collect();
// Pack POI IDs into a contiguous buffer
let total_id_bytes: usize = id_raw.iter().map(|s| s.len()).sum();
let mut id_buffer = String::with_capacity(total_id_bytes);
let mut id_offsets = Vec::with_capacity(row_count);
let mut id_lengths = Vec::with_capacity(row_count);
for s in &id_raw {
let offset = id_buffer.len() as u32;
let length = s.len().min(u16::MAX as usize) as u16;
id_offsets.push(offset);
id_lengths.push(length);
id_buffer.push_str(&s[..length as usize]);
}
let category = InternedColumn::build(&category_raw);
let icon_category = InternedColumn::build(&icon_category_raw);
let group = InternedColumn::build(&group_raw);
let emoji = InternedColumn::build(&emoji_raw);
info!(
category_unique = category.values.len(),
icon_category_unique = icon_category.values.len(),
group_unique = group.values.len(),
emoji_unique = emoji.values.len(),
"POI string columns interned"
);
// Assign a deterministic pseudo-random priority to each row.
// This ensures the same POIs are selected across requests,
// preventing visual "shuffling" when panning the map.
let priority = generate_priorities(row_count);
let (school_meta_idx, school_meta) = build_school_meta(row_count, &df)?;
info!(schools = school_meta.len(), "Loaded GIAS school metadata");
info!("POI data loading complete.");
Ok(POIData {
id_buffer,
id_offsets,
id_lengths,
name,
category,
icon_category,
group,
lat,
lng,
emoji,
priority,
school_meta_idx,
school_meta,
})
}
/// Build dashboard category groups from every category present in the loaded POI data.
pub fn category_groups(&self) -> anyhow::Result<Vec<POICategoryGroup>> {
let mut group_cats: HashMap<String, HashSet<String>> = HashMap::new();
let num_pois = self.category.indices.len();
for row in 0..num_pois {
let category = self.category.get(row).to_string();
let group = self.group.get(row).to_string();
group_cats.entry(group).or_default().insert(category);
}
// Validate that data groups match the hardcoded order exactly
let expected: HashSet<&str> = POI_GROUP_ORDER.iter().copied().collect();
let actual: HashSet<&str> = group_cats.keys().map(|key| key.as_str()).collect();
let missing_from_data: Vec<&&str> = expected.difference(&actual).collect();
let missing_from_order: Vec<&&str> = actual.difference(&expected).collect();
if !missing_from_data.is_empty() || !missing_from_order.is_empty() {
bail!(
"POI group mismatch!\n In POI_GROUP_ORDER but not in data: {:?}\n In data but not in POI_GROUP_ORDER: {:?}",
missing_from_data, missing_from_order
);
}
let preferred_order: HashMap<&str, HashMap<&str, usize>> = DASHBOARD_POI_GROUPS
.iter()
.map(|(group, categories)| {
(
*group,
categories
.iter()
.enumerate()
.map(|(idx, category)| (*category, idx))
.collect(),
)
})
.collect();
let groups: Vec<POICategoryGroup> = POI_GROUP_ORDER
.iter()
.filter_map(|group_name| {
let mut categories: Vec<String> = group_cats
.get(*group_name)
.map(|categories| categories.iter().cloned().collect())
.unwrap_or_default();
if categories.is_empty() {
return None;
}
let group_order = preferred_order.get(*group_name);
categories.sort_by(|a, b| {
let a_order = group_order.and_then(|order| order.get(a.as_str())).copied();
let b_order = group_order.and_then(|order| order.get(b.as_str())).copied();
match (a_order, b_order) {
(Some(left), Some(right)) => left.cmp(&right),
(Some(_), None) => std::cmp::Ordering::Less,
(None, Some(_)) => std::cmp::Ordering::Greater,
(None, None) => a.cmp(b),
}
});
Some(POICategoryGroup {
name: (*group_name).to_string(),
categories,
})
})
.collect();
Ok(groups)
}
}
#[cfg(test)]
impl POIData {
/// Minimal empty instance for integration tests that need an `AppState`
/// but never touch POI data.
pub(crate) fn empty_for_tests() -> Self {
POIData {
id_buffer: String::new(),
id_offsets: Vec::new(),
id_lengths: Vec::new(),
group: InternedColumn::build(&[]),
category: InternedColumn::build(&[]),
icon_category: InternedColumn::build(&[]),
name: Vec::new(),
lat: Vec::new(),
lng: Vec::new(),
emoji: InternedColumn::build(&[]),
priority: Vec::new(),
school_meta_idx: Vec::new(),
school_meta: Vec::new(),
}
}
}
#[cfg(test)]
mod tests {
use super::*;
#[test]
fn category_filter_matches_exact_present_categories() {
let values = vec![
"Supermarket".to_string(),
"Tesco".to_string(),
"Aldi".to_string(),
"Rail station".to_string(),
];
let selected = resolve_poi_category_filter(&values, "Supermarket,Rail station");
assert!(selected.contains(&0));
assert!(selected.contains(&3));
assert_eq!(selected.len(), 2);
}
#[test]
fn unknown_category_filter_matches_nothing() {
let values = vec!["Supermarket".to_string()];
let selected = resolve_poi_category_filter(&values, "Unknown");
assert!(selected.is_empty());
}
#[test]
fn legacy_school_filter_expands_to_all_school_categories() {
// Bookmarked URLs from before the per-phase split sent `poi=School`;
// they should still match every school category that's loaded.
let values = vec![
"Primary school".to_string(),
"Secondary school".to_string(),
"University".to_string(),
"Tesco".to_string(),
];
let selected = resolve_poi_category_filter(&values, "School");
assert!(selected.contains(&0));
assert!(selected.contains(&1));
assert!(selected.contains(&2));
assert!(!selected.contains(&3));
assert_eq!(selected.len(), 3);
}
#[test]
fn coop_category_aliases_resolve_to_single_category() {
let values = vec!["Co-op".to_string(), "Tesco".to_string()];
let selected = resolve_poi_category_filter(
&values,
"Central England Co-operative,The Southern Co-operative",
);
assert!(selected.contains(&0));
assert_eq!(selected.len(), 1);
assert_eq!(canonical_poi_category("Lincolnshire Co-operative"), "Co-op");
assert_eq!(canonical_poi_category("Tesco"), "Tesco");
}
}