Refactor the server

server-rs/src/routes/hexagons.rs

@@ -0,0 +1,257 @@
+use std::fmt::Write;
+use std::sync::Arc;
+
+use axum::extract::Query;
+use axum::http::StatusCode;
+use axum::response::IntoResponse;
+use rustc_hash::FxHashMap;
+use serde::Deserialize;
+use tracing::{info, warn};
+
+use crate::consts::{H3_PRECOMPUTE_MAX, H3_PRECOMPUTE_MIN};
+use crate::filter::parse_filters;
+use crate::state::AppState;
+
+const BOUNDS_BUFFER_PERCENT: f64 = 0.2;
+
+#[derive(Deserialize)]
+pub struct HexagonParams {
+    resolution: u8,
+    bounds: Option<String>,
+    /// Comma-separated filters: `name:min:max,...`
+    /// Rows must have non-NaN values within [min,max] for each filter.
+    filters: Option<String>,
+}
+
+/// Per-cell accumulator for aggregating features
+struct CellAgg {
+    count: u32,
+    mins: Vec<f64>,
+    maxs: Vec<f64>,
+}
+
+impl CellAgg {
+    fn new(num_features: usize) -> Self {
+        CellAgg {
+            count: 0,
+            mins: vec![f64::INFINITY; num_features],
+            maxs: vec![f64::NEG_INFINITY; num_features],
+        }
+    }
+
+    /// Add a row using row-major feature_data layout.
+    /// feature_data[row * num_features + feat_idx] — all features for one row
+    /// are contiguous, so this reads a single cache line per ~8 features.
+    #[inline]
+    fn add_row(&mut self, feature_data: &[f64], row: usize, num_features: usize) {
+        self.count += 1;
+        let base = row * num_features;
+        let row_slice = &feature_data[base..base + num_features];
+        for (i, &v) in row_slice.iter().enumerate() {
+            if v.is_finite() {
+                if v < self.mins[i] {
+                    self.mins[i] = v;
+                }
+                if v > self.maxs[i] {
+                    self.maxs[i] = v;
+                }
+            }
+        }
+    }
+}
+
+/// Write the hexagons JSON response directly to a String buffer,
+/// avoiding serde_json::Value allocations entirely.
+fn write_hexagons_json(
+    buf: &mut String,
+    groups: &FxHashMap<u64, CellAgg>,
+    min_keys: &[String],
+    max_keys: &[String],
+    num_features: usize,
+) {
+    buf.push_str("{\"features\":[");
+    let mut first = true;
+    for (&cell_id, agg) in groups {
+        if !first {
+            buf.push(',');
+        }
+        first = false;
+
+        let cell = h3o::CellIndex::try_from(cell_id).unwrap();
+        write!(buf, "{{\"h3\":\"{}\",\"count\":{}", cell, agg.count).unwrap();
+
+        for i in 0..num_features {
+            if agg.mins[i] != f64::INFINITY {
+                write!(
+                    buf,
+                    ",\"{}\":{},\"{}\":{}",
+                    min_keys[i], agg.mins[i], max_keys[i], agg.maxs[i]
+                )
+                .unwrap();
+            }
+        }
+        buf.push('}');
+    }
+    buf.push_str("]}");
+}
+
+pub async fn get_hexagons(
+    state: Arc<AppState>,
+    Query(params): Query<HexagonParams>,
+) -> Result<impl IntoResponse, (StatusCode, String)> {
+    let resolution = params.resolution;
+    if resolution < H3_PRECOMPUTE_MIN || resolution > H3_PRECOMPUTE_MAX {
+        warn!(resolution, "Resolution out of range [{}, {}]", H3_PRECOMPUTE_MIN, H3_PRECOMPUTE_MAX);
+        return Err((
+            StatusCode::BAD_REQUEST,
+            format!(
+                "resolution must be between {} and {}",
+                H3_PRECOMPUTE_MIN, H3_PRECOMPUTE_MAX
+            ),
+        ));
+    }
+
+    let bounds_str = params.bounds.ok_or((
+        StatusCode::BAD_REQUEST,
+        "bounds parameter is required".into(),
+    ))?;
+
+    let parts: Vec<f64> = bounds_str
+        .split(',')
+        .map(|s| s.trim().parse::<f64>())
+        .collect::<Result<Vec<_>, _>>()
+        .map_err(|_| {
+            (
+                StatusCode::BAD_REQUEST,
+                "Invalid bounds format. Use: south,west,north,east".into(),
+            )
+        })?;
+
+    if parts.len() != 4 {
+        return Err((
+            StatusCode::BAD_REQUEST,
+            "Invalid bounds format. Use: south,west,north,east".into(),
+        ));
+    }
+
+    let (mut south, mut west, mut north, mut east) = (parts[0], parts[1], parts[2], parts[3]);
+
+    let lat_range = north - south;
+    let lng_range = east - west;
+    south -= lat_range * BOUNDS_BUFFER_PERCENT;
+    north += lat_range * BOUNDS_BUFFER_PERCENT;
+    west -= lng_range * BOUNDS_BUFFER_PERCENT;
+    east += lng_range * BOUNDS_BUFFER_PERCENT;
+
+    let precision = 0.01;
+    south = (south / precision).floor() * precision;
+    west = (west / precision).floor() * precision;
+    north = (north / precision).ceil() * precision;
+    east = (east / precision).ceil() * precision;
+
+    let filters_str = params.filters.clone();
+    let (parsed_filters, parsed_enum_filters) = parse_filters(
+        params.filters.as_deref(),
+        &state.data.feature_names,
+        &state.data.enum_features,
+    );
+    let num_filters = parsed_filters.len() + parsed_enum_filters.len();
+
+    let json_body = tokio::task::spawn_blocking(move || {
+        let t0 = std::time::Instant::now();
+
+        let num_features = state.data.num_features;
+        let feature_data = &state.data.feature_data;
+
+        let min_keys: Vec<String> = state
+            .data
+            .feature_names
+            .iter()
+            .map(|n| format!("min_{}", n))
+            .collect();
+        let max_keys: Vec<String> = state
+            .data
+            .feature_names
+            .iter()
+            .map(|n| format!("max_{}", n))
+            .collect();
+
+        let h3_cells_for_res: Option<&[u64]> = state
+            .h3_cells
+            .get(resolution as usize)
+            .filter(|v| !v.is_empty())
+            .map(|v| v.as_slice());
+
+        let mut groups: FxHashMap<u64, CellAgg> = FxHashMap::default();
+
+        let enum_features = &state.data.enum_features;
+
+        // 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(|f| {
+                let v = feature_data[row * num_features + f.feat_idx];
+                v.is_finite() && v >= f.min && v <= f.max
+            }) && parsed_enum_filters.iter().all(|ef| {
+                let v = enum_features[ef.enum_idx].data[row];
+                v != 255 && ef.allowed.contains(&v)
+            })
+        };
+
+        if let Some(precomputed) = h3_cells_for_res {
+            state
+                .grid
+                .for_each_in_bounds(south, west, north, east, |row_idx| {
+                    let row = row_idx as usize;
+                    if !row_passes(row) {
+                        return;
+                    }
+                    let cell_id = precomputed[row];
+                    groups
+                        .entry(cell_id)
+                        .or_insert_with(|| CellAgg::new(num_features))
+                        .add_row(feature_data, row, num_features);
+                });
+        } else {
+            let h3_res = h3o::Resolution::try_from(resolution).unwrap();
+            state
+                .grid
+                .for_each_in_bounds(south, west, north, east, |row_idx| {
+                    let row = row_idx as usize;
+                    if !row_passes(row) {
+                        return;
+                    }
+                    let cell_id = h3o::LatLng::new(state.data.lat[row], state.data.lon[row])
+                        .map(|c| u64::from(c.to_cell(h3_res)))
+                        .unwrap_or(0);
+                    groups
+                        .entry(cell_id)
+                        .or_insert_with(|| CellAgg::new(num_features))
+                        .add_row(feature_data, row, num_features);
+                });
+        }
+
+        let t_agg = t0.elapsed();
+
+        let mut json_buf = String::with_capacity(groups.len() * 128);
+        write_hexagons_json(&mut json_buf, &groups, &min_keys, &max_keys, num_features);
+
+        let t_total = t0.elapsed();
+        info!(
+            resolution,
+            cells = groups.len(),
+            filters = num_filters,
+            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),
+            bytes = json_buf.len(),
+            "GET /api/hexagons"
+        );
+
+        json_buf
+    })
+    .await
+    .unwrap();
+
+    Ok(([("content-type", "application/json")], json_body))
+}