Refactor and improve
This commit is contained in:
parent
1f148b2185
commit
242acff987
22 changed files with 754 additions and 1053 deletions
|
|
@ -1,35 +1,22 @@
|
|||
use std::fmt::{self, Write};
|
||||
use std::sync::Arc;
|
||||
|
||||
use axum::extract::Query;
|
||||
use axum::http::StatusCode;
|
||||
use axum::response::IntoResponse;
|
||||
use axum::response::Json;
|
||||
use rustc_hash::FxHashMap;
|
||||
use serde::Deserialize;
|
||||
use serde::{Deserialize, Serialize};
|
||||
use serde_json::{Map, Value};
|
||||
use tracing::{info, warn};
|
||||
|
||||
use crate::consts::{
|
||||
BOUNDS_BUFFER_PERCENT, BOUNDS_QUANTIZATION, ENUM_NULL, H3_PRECOMPUTE_MAX, H3_REQUEST_MAX,
|
||||
H3_REQUEST_MIN, POSTCODE_MIN_RESOLUTION,
|
||||
BOUNDS_BUFFER_PERCENT, BOUNDS_QUANTIZATION, H3_PRECOMPUTE_MAX, H3_REQUEST_MAX, H3_REQUEST_MIN,
|
||||
};
|
||||
use crate::filter::parse_filters;
|
||||
use crate::parsing::{parse_bounds, parse_filters, row_passes_filters};
|
||||
use crate::state::AppState;
|
||||
|
||||
use super::parse::parse_bounds;
|
||||
|
||||
struct HumanBytes(usize);
|
||||
|
||||
impl fmt::Display for HumanBytes {
|
||||
fn fmt(&self, formatter: &mut fmt::Formatter<'_>) -> fmt::Result {
|
||||
let bytes = self.0;
|
||||
if bytes >= 1_000_000 {
|
||||
write!(formatter, "{:.1} MB", bytes as f64 / 1_000_000.0)
|
||||
} else if bytes >= 1_000 {
|
||||
write!(formatter, "{:.1} KB", bytes as f64 / 1_000.0)
|
||||
} else {
|
||||
write!(formatter, "{} B", bytes)
|
||||
}
|
||||
}
|
||||
#[derive(Serialize)]
|
||||
pub struct HexagonsResponse {
|
||||
features: Vec<Map<String, Value>>,
|
||||
}
|
||||
|
||||
#[derive(Deserialize)]
|
||||
|
|
@ -51,28 +38,14 @@ struct CellAgg {
|
|||
count: u32,
|
||||
mins: Box<[f32]>,
|
||||
maxs: Box<[f32]>,
|
||||
/// Min/max ordinal indices for enum features (255 = no data yet)
|
||||
enum_mins: Box<[u8]>,
|
||||
enum_maxs: Box<[u8]>,
|
||||
/// Most common postcode in this cell (only tracked at high resolutions)
|
||||
postcode: Option<String>,
|
||||
postcode_count: u32,
|
||||
lat_sum: f64,
|
||||
lon_sum: f64,
|
||||
}
|
||||
|
||||
impl CellAgg {
|
||||
fn new(num_features: usize, num_enums: usize) -> Self {
|
||||
fn new(num_features: usize) -> Self {
|
||||
CellAgg {
|
||||
count: 0,
|
||||
mins: vec![f32::INFINITY; num_features].into_boxed_slice(),
|
||||
maxs: vec![f32::NEG_INFINITY; num_features].into_boxed_slice(),
|
||||
enum_mins: vec![ENUM_NULL; num_enums].into_boxed_slice(),
|
||||
enum_maxs: vec![0; num_enums].into_boxed_slice(),
|
||||
postcode: None,
|
||||
postcode_count: 0,
|
||||
lat_sum: 0.0,
|
||||
lon_sum: 0.0,
|
||||
}
|
||||
}
|
||||
|
||||
|
|
@ -96,23 +69,6 @@ impl CellAgg {
|
|||
}
|
||||
}
|
||||
|
||||
/// Track min/max ordinal index for each enum feature in this cell.
|
||||
#[inline]
|
||||
fn add_enums(&mut self, enum_data: &[u8], row: usize, num_enums: usize) {
|
||||
let base = row * num_enums;
|
||||
let row_slice = &enum_data[base..base + num_enums];
|
||||
for (enum_index, &value) in row_slice.iter().enumerate() {
|
||||
if value != ENUM_NULL {
|
||||
if self.enum_mins[enum_index] == ENUM_NULL || value < self.enum_mins[enum_index] {
|
||||
self.enum_mins[enum_index] = value;
|
||||
}
|
||||
if value > self.enum_maxs[enum_index] {
|
||||
self.enum_maxs[enum_index] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Add a row, only aggregating the features at the given indices.
|
||||
#[inline]
|
||||
fn add_row_selective(
|
||||
|
|
@ -136,178 +92,57 @@ impl CellAgg {
|
|||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Track min/max ordinal index for selected enum features only.
|
||||
#[inline]
|
||||
fn add_enums_selective(
|
||||
&mut self,
|
||||
enum_data: &[u8],
|
||||
row: usize,
|
||||
num_enums: usize,
|
||||
indices: &[usize],
|
||||
) {
|
||||
let base = row * num_enums;
|
||||
for &enum_index in indices {
|
||||
let value = enum_data[base + enum_index];
|
||||
if value != ENUM_NULL {
|
||||
if self.enum_mins[enum_index] == ENUM_NULL || value < self.enum_mins[enum_index] {
|
||||
self.enum_mins[enum_index] = value;
|
||||
}
|
||||
if value > self.enum_maxs[enum_index] {
|
||||
self.enum_maxs[enum_index] = value;
|
||||
}
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Track postcode and centroid for high-resolution cells.
|
||||
/// Uses simple "first seen" approach — at res 11/12, most rows in a cell share a postcode.
|
||||
#[inline]
|
||||
fn add_postcode(&mut self, postcode: &str, lat: f32, lon: f32) {
|
||||
self.lat_sum += lat as f64;
|
||||
self.lon_sum += lon as f64;
|
||||
if postcode.is_empty() {
|
||||
return;
|
||||
}
|
||||
if self.postcode.is_none() {
|
||||
self.postcode = Some(postcode.to_string());
|
||||
self.postcode_count = 1;
|
||||
} else if self.postcode.as_deref() == Some(postcode) {
|
||||
self.postcode_count += 1;
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Escape a string for inclusion in a JSON string literal.
|
||||
pub(crate) fn write_json_escaped(buf: &mut String, text: &str) {
|
||||
for character in text.chars() {
|
||||
match character {
|
||||
'"' => buf.push_str("\\\""),
|
||||
'\\' => buf.push_str("\\\\"),
|
||||
'\n' => buf.push_str("\\n"),
|
||||
'\r' => buf.push_str("\\r"),
|
||||
'\t' => buf.push_str("\\t"),
|
||||
ctrl if ctrl < '\x20' => {
|
||||
let _ = write!(buf, "\\u{:04x}", ctrl as u32);
|
||||
}
|
||||
other => buf.push(other),
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
/// Write the hexagons JSON response directly to a String buffer,
|
||||
/// avoiding serde_json::Value allocations entirely.
|
||||
/// When `numeric_indices` / `enum_indices` are Some, only those features are written.
|
||||
#[allow(clippy::too_many_arguments)]
|
||||
fn write_hexagons_json(
|
||||
buf: &mut String,
|
||||
/// Build feature maps from aggregated cell data.
|
||||
fn build_feature_maps(
|
||||
groups: &FxHashMap<u64, CellAgg>,
|
||||
min_keys: &[String],
|
||||
max_keys: &[String],
|
||||
num_features: usize,
|
||||
enum_min_keys: &[String],
|
||||
enum_max_keys: &[String],
|
||||
num_enums: usize,
|
||||
include_postcode: bool,
|
||||
numeric_indices: Option<&[usize]>,
|
||||
enum_indices: Option<&[usize]>,
|
||||
) {
|
||||
buf.push_str("{\"features\":[");
|
||||
let mut first = true;
|
||||
indices: Option<&[usize]>,
|
||||
) -> Vec<Map<String, Value>> {
|
||||
let mut features = Vec::with_capacity(groups.len());
|
||||
|
||||
for (&cell_id, aggregation) in groups {
|
||||
let Some(cell) = h3o::CellIndex::try_from(cell_id).ok() else {
|
||||
continue;
|
||||
};
|
||||
|
||||
if !first {
|
||||
buf.push(',');
|
||||
}
|
||||
first = false;
|
||||
let mut map = Map::new();
|
||||
map.insert("h3".into(), Value::String(cell.to_string()));
|
||||
map.insert("count".into(), Value::Number(aggregation.count.into()));
|
||||
|
||||
let _ = write!(buf, "{{\"h3\":\"{}\",\"count\":{}", cell, aggregation.count);
|
||||
|
||||
if let Some(indices) = numeric_indices {
|
||||
for &feat_index in indices {
|
||||
if aggregation.mins[feat_index].is_finite()
|
||||
&& aggregation.maxs[feat_index].is_finite()
|
||||
{
|
||||
let _ = write!(
|
||||
buf,
|
||||
",\"{}\":{},\"{}\":{}",
|
||||
min_keys[feat_index],
|
||||
aggregation.mins[feat_index],
|
||||
max_keys[feat_index],
|
||||
aggregation.maxs[feat_index]
|
||||
);
|
||||
}
|
||||
}
|
||||
let iter: Box<dyn Iterator<Item = usize>> = if let Some(idx) = indices {
|
||||
Box::new(idx.iter().copied())
|
||||
} else {
|
||||
for feat_index in 0..num_features {
|
||||
if aggregation.mins[feat_index].is_finite()
|
||||
&& aggregation.maxs[feat_index].is_finite()
|
||||
{
|
||||
let _ = write!(
|
||||
buf,
|
||||
",\"{}\":{},\"{}\":{}",
|
||||
min_keys[feat_index],
|
||||
aggregation.mins[feat_index],
|
||||
max_keys[feat_index],
|
||||
aggregation.maxs[feat_index]
|
||||
);
|
||||
Box::new(0..num_features)
|
||||
};
|
||||
|
||||
for feat_index in iter {
|
||||
if aggregation.mins[feat_index].is_finite()
|
||||
&& aggregation.maxs[feat_index].is_finite()
|
||||
{
|
||||
if let (Some(min_num), Some(max_num)) = (
|
||||
serde_json::Number::from_f64(aggregation.mins[feat_index] as f64),
|
||||
serde_json::Number::from_f64(aggregation.maxs[feat_index] as f64),
|
||||
) {
|
||||
map.insert(min_keys[feat_index].clone(), Value::Number(min_num));
|
||||
map.insert(max_keys[feat_index].clone(), Value::Number(max_num));
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if let Some(indices) = enum_indices {
|
||||
for &enum_index in indices {
|
||||
if aggregation.enum_mins[enum_index] != ENUM_NULL {
|
||||
let _ = write!(
|
||||
buf,
|
||||
",\"{}\":{},\"{}\":{}",
|
||||
enum_min_keys[enum_index],
|
||||
aggregation.enum_mins[enum_index],
|
||||
enum_max_keys[enum_index],
|
||||
aggregation.enum_maxs[enum_index]
|
||||
);
|
||||
}
|
||||
}
|
||||
} else {
|
||||
for enum_index in 0..num_enums {
|
||||
if aggregation.enum_mins[enum_index] != ENUM_NULL {
|
||||
let _ = write!(
|
||||
buf,
|
||||
",\"{}\":{},\"{}\":{}",
|
||||
enum_min_keys[enum_index],
|
||||
aggregation.enum_mins[enum_index],
|
||||
enum_max_keys[enum_index],
|
||||
aggregation.enum_maxs[enum_index]
|
||||
);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
if include_postcode {
|
||||
if let Some(ref postcode) = aggregation.postcode {
|
||||
let total = aggregation.count as f64;
|
||||
let centroid_lat = aggregation.lat_sum / total;
|
||||
let centroid_lon = aggregation.lon_sum / total;
|
||||
if centroid_lat.is_finite() && centroid_lon.is_finite() {
|
||||
buf.push_str(",\"postcode\":\"");
|
||||
write_json_escaped(buf, postcode);
|
||||
let _ = write!(buf, "\",\"lat\":{},\"lon\":{}", centroid_lat, centroid_lon);
|
||||
}
|
||||
}
|
||||
}
|
||||
|
||||
buf.push('}');
|
||||
features.push(map);
|
||||
}
|
||||
buf.push_str("]}");
|
||||
|
||||
features
|
||||
}
|
||||
|
||||
pub async fn get_hexagons(
|
||||
state: Arc<AppState>,
|
||||
Query(params): Query<HexagonParams>,
|
||||
) -> Result<impl IntoResponse, (StatusCode, String)> {
|
||||
) -> Result<Json<HexagonsResponse>, (StatusCode, String)> {
|
||||
let resolution = params.resolution;
|
||||
if !(H3_REQUEST_MIN..=H3_REQUEST_MAX).contains(&resolution) {
|
||||
warn!(
|
||||
|
|
@ -346,50 +181,40 @@ pub async fn get_hexagons(
|
|||
let (parsed_filters, parsed_enum_filters) = parse_filters(
|
||||
params.filters.as_deref(),
|
||||
&state.data.feature_names,
|
||||
&state.data.enum_features,
|
||||
&state.data.enum_values,
|
||||
);
|
||||
let num_filters = parsed_filters.len() + parsed_enum_filters.len();
|
||||
|
||||
// Parse optional `fields` param into numeric and enum index sets.
|
||||
// Parse optional `fields` param into feature indices.
|
||||
// If `fields` is absent (None), all features are included.
|
||||
// If `fields` is present (even empty string), only listed features are included.
|
||||
let field_indices: Option<(Vec<usize>, Vec<usize>)> =
|
||||
params.fields.as_ref().map(|fields_str| {
|
||||
let mut numeric_indices = Vec::new();
|
||||
let mut enum_indices = Vec::new();
|
||||
if !fields_str.is_empty() {
|
||||
for name in fields_str.split(',') {
|
||||
let name = name.trim();
|
||||
if name.is_empty() {
|
||||
continue;
|
||||
}
|
||||
if let Some(idx) = state
|
||||
.data
|
||||
.feature_names
|
||||
.iter()
|
||||
.position(|feat| feat == name)
|
||||
{
|
||||
numeric_indices.push(idx);
|
||||
} else if let Some(&idx) = state.enum_name_to_idx.get(name) {
|
||||
enum_indices.push(idx);
|
||||
}
|
||||
let field_indices: Option<Vec<usize>> = params.fields.as_ref().map(|fields_str| {
|
||||
if fields_str.is_empty() {
|
||||
return Vec::new();
|
||||
}
|
||||
fields_str
|
||||
.split(',')
|
||||
.filter_map(|name| {
|
||||
let name = name.trim();
|
||||
if name.is_empty() {
|
||||
return None;
|
||||
}
|
||||
}
|
||||
(numeric_indices, enum_indices)
|
||||
});
|
||||
state
|
||||
.data
|
||||
.feature_names
|
||||
.iter()
|
||||
.position(|feat| feat == name)
|
||||
})
|
||||
.collect()
|
||||
});
|
||||
|
||||
let json_body = tokio::task::spawn_blocking(move || -> Result<String, String> {
|
||||
let response = tokio::task::spawn_blocking(move || -> Result<HexagonsResponse, String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
|
||||
let num_features = state.data.num_features;
|
||||
let num_enums = state.data.num_enums;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let enum_data = &state.data.enum_data;
|
||||
|
||||
let min_keys = &state.min_keys;
|
||||
let max_keys = &state.max_keys;
|
||||
let enum_min_keys = &state.enum_min_keys;
|
||||
let enum_max_keys = &state.enum_max_keys;
|
||||
|
||||
let h3_res = h3o::Resolution::try_from(resolution)
|
||||
.map_err(|error| format!("Invalid H3 resolution {}: {}", resolution, error))?;
|
||||
|
|
@ -398,50 +223,20 @@ pub async fn get_hexagons(
|
|||
|
||||
let mut groups: FxHashMap<u64, CellAgg> = FxHashMap::default();
|
||||
|
||||
let include_postcode = resolution >= POSTCODE_MIN_RESOLUTION;
|
||||
|
||||
// Row-level filter check: numeric must be non-NaN and within [min, max],
|
||||
// enum must have value index in the allowed set
|
||||
let row_passes = |row: usize| -> bool {
|
||||
parsed_filters.iter().all(|filter| {
|
||||
let value = feature_data[row * num_features + filter.feat_idx];
|
||||
value.is_finite() && value >= filter.min && value <= filter.max
|
||||
}) && parsed_enum_filters.iter().all(|enum_filter| {
|
||||
let value = enum_data[row * num_enums + enum_filter.enum_idx];
|
||||
value != ENUM_NULL && enum_filter.allowed.contains(&value)
|
||||
})
|
||||
};
|
||||
|
||||
// Choose aggregation strategy based on whether fields are specified
|
||||
let has_selective = field_indices.is_some();
|
||||
let (sel_numeric, sel_enum) = field_indices
|
||||
.as_ref()
|
||||
.map_or((&[][..], &[][..]), |(ni, ei)| {
|
||||
(ni.as_slice(), ei.as_slice())
|
||||
});
|
||||
let sel_indices = field_indices.as_deref().unwrap_or(&[]);
|
||||
|
||||
let aggregate_row = |groups: &mut FxHashMap<u64, CellAgg>, cell_id: u64, row: usize| {
|
||||
let aggregation = groups
|
||||
.entry(cell_id)
|
||||
.or_insert_with(|| CellAgg::new(num_features, num_enums));
|
||||
.or_insert_with(|| CellAgg::new(num_features));
|
||||
if has_selective {
|
||||
aggregation.add_row_selective(feature_data, row, num_features, sel_numeric);
|
||||
aggregation.add_enums_selective(enum_data, row, num_enums, sel_enum);
|
||||
aggregation.add_row_selective(feature_data, row, num_features, sel_indices);
|
||||
} else {
|
||||
aggregation.add_row(feature_data, row, num_features);
|
||||
aggregation.add_enums(enum_data, row, num_enums);
|
||||
}
|
||||
if include_postcode {
|
||||
aggregation.add_postcode(
|
||||
state.data.postcode(row),
|
||||
state.data.lat[row],
|
||||
state.data.lon[row],
|
||||
);
|
||||
}
|
||||
};
|
||||
|
||||
// Resolve cell at requested resolution from precomputed max-resolution cell.
|
||||
// For max resolution, use directly; for lower resolutions, derive parent.
|
||||
let cell_for_row = |row: usize| -> u64 {
|
||||
let max_cell = precomputed[row];
|
||||
if !need_parent || max_cell == 0 {
|
||||
|
|
@ -458,7 +253,13 @@ pub async fn get_hexagons(
|
|||
.grid
|
||||
.for_each_in_bounds(south, west, north, east, |row_idx| {
|
||||
let row = row_idx as usize;
|
||||
if !row_passes(row) {
|
||||
if !row_passes_filters(
|
||||
row,
|
||||
&parsed_filters,
|
||||
&parsed_enum_filters,
|
||||
feature_data,
|
||||
num_features,
|
||||
) {
|
||||
return;
|
||||
}
|
||||
aggregate_row(&mut groups, cell_for_row(row), row);
|
||||
|
|
@ -466,19 +267,12 @@ pub async fn get_hexagons(
|
|||
|
||||
let t_agg = t0.elapsed();
|
||||
|
||||
let mut json_buf = String::with_capacity(groups.len() * 128);
|
||||
write_hexagons_json(
|
||||
&mut json_buf,
|
||||
let features = build_feature_maps(
|
||||
&groups,
|
||||
min_keys,
|
||||
max_keys,
|
||||
num_features,
|
||||
enum_min_keys,
|
||||
enum_max_keys,
|
||||
num_enums,
|
||||
include_postcode,
|
||||
field_indices.as_ref().map(|(ni, _)| ni.as_slice()),
|
||||
field_indices.as_ref().map(|(_, ei)| ei.as_slice()),
|
||||
field_indices.as_deref(),
|
||||
);
|
||||
|
||||
let t_total = t0.elapsed();
|
||||
|
|
@ -489,15 +283,14 @@ pub async fn get_hexagons(
|
|||
filters_raw = filters_str.as_deref().unwrap_or("-"),
|
||||
agg_ms = format_args!("{:.1}", t_agg.as_secs_f64() * 1000.0),
|
||||
total_ms = format_args!("{:.1}", t_total.as_secs_f64() * 1000.0),
|
||||
size = format_args!("{}", HumanBytes(json_buf.len())),
|
||||
"GET /api/hexagons"
|
||||
);
|
||||
|
||||
Ok(json_buf)
|
||||
Ok(HexagonsResponse { features })
|
||||
})
|
||||
.await
|
||||
.map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error.to_string()))?
|
||||
.map_err(|error| (StatusCode::INTERNAL_SERVER_ERROR, error))?;
|
||||
|
||||
Ok(([("content-type", "application/json")], json_body))
|
||||
Ok(Json(response))
|
||||
}
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue