Optimise

server-rs/src/routes/hexagon_stats.rs

@@ -8,7 +8,7 @@ use axum::response::IntoResponse;
 use serde::Deserialize;
 use tracing::{info, warn};
 
-use crate::consts::{ENUM_NULL, H3_REQUEST_MAX, H3_REQUEST_MIN, HISTOGRAM_BINS};
+use crate::consts::{ENUM_NULL, H3_PRECOMPUTE_MAX, H3_REQUEST_MAX, H3_REQUEST_MIN, HISTOGRAM_BINS};
 use crate::filter::{parse_filters, row_passes_filters};
 use crate::state::AppState;
 
@@ -19,6 +19,10 @@ pub struct HexagonStatsParams {
     pub h3: String,
     pub resolution: u8,
     pub filters: Option<String>,
+    /// Comma-separated feature names to include in stats response.
+    /// When present (even if empty), only listed features are computed.
+    /// When absent, all features are returned (backward compatible).
+    pub fields: Option<String>,
 }
 
 pub async fn get_hexagon_stats(
@@ -45,8 +49,6 @@ pub async fn get_hexagon_stats(
             ),
         ));
     }
-    let resolution_idx = resolution as usize;
-
     let h3_str = params.h3.clone();
     let filters_str = params.filters.clone();
     let (parsed_filters, parsed_enum_filters) = parse_filters(
@@ -56,42 +58,58 @@ pub async fn get_hexagon_stats(
     );
     let num_filters = parsed_filters.len() + parsed_enum_filters.len();
 
+    // Parse optional `fields` param into sets of feature names.
+    // None = include all, Some = only include listed features.
+    let field_set: Option<std::collections::HashSet<String>> = params.fields.as_ref().map(|fields_str| {
+        fields_str
+            .split(',')
+            .map(|field| field.trim().to_string())
+            .filter(|field| !field.is_empty())
+            .collect()
+    });
+
     let result = tokio::task::spawn_blocking(move || {
         let start_time = std::time::Instant::now();
-        let precomputed: Option<&[u64]> = state
-            .h3_cells
-            .get(resolution_idx)
-            .filter(|cells| !cells.is_empty())
-            .map(|cells| cells.as_slice());
+        let precomputed = &state.h3_cells;
         let h3_res = h3o::Resolution::try_from(resolution)
             .map_err(|err| format!("Invalid H3 resolution {}: {}", resolution, err))?;
+        let need_parent = resolution < H3_PRECOMPUTE_MAX;
         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 enum_features = &state.data.enum_features;
 
         let (min_lat, min_lon, max_lat, max_lon) = h3_cell_bounds(cell, 0.001);
 
+        // Resolve cell at requested resolution from precomputed max-resolution cell
+        let cell_for_row = |row: usize| -> u64 {
+            let max_cell = precomputed[row];
+            if !need_parent || max_cell == 0 {
+                return max_cell;
+            }
+            h3o::CellIndex::try_from(max_cell)
+                .ok()
+                .and_then(|ci| ci.parent(h3_res))
+                .map(u64::from)
+                .unwrap_or(0)
+        };
+
         // Collect matching rows
         let mut matching_rows: Vec<usize> = Vec::new();
         state
             .grid
             .for_each_in_bounds(min_lat, min_lon, max_lat, max_lon, |row_idx| {
                 let row = row_idx as usize;
-                let row_cell = if let Some(h3_data) = precomputed {
-                    h3_data[row]
-                } else {
-                    h3o::LatLng::new(state.data.lat[row] as f64, state.data.lon[row] as f64)
-                        .map(|coord| u64::from(coord.to_cell(h3_res)))
-                        .unwrap_or(0)
-                };
-                if row_cell == cell_u64
+                if cell_for_row(row) == cell_u64
                     && row_passes_filters(
                         row,
                         &parsed_filters,
                         &parsed_enum_filters,
                         feature_data,
                         num_features,
-                        enum_features,
+                        enum_data,
+                        num_enums,
                     )
                 {
                     matching_rows.push(row);
@@ -109,6 +127,12 @@ pub async fn get_hexagon_stats(
         output.push_str(",\"numeric_features\":[");
         let mut first_numeric = true;
         for (feature_index, feature_name) in state.data.feature_names.iter().enumerate() {
+            // Skip features not in the requested set (when fields param is present)
+            if let Some(ref set) = field_set {
+                if !set.contains(feature_name.as_str()) {
+                    continue;
+                }
+            }
             let global_stats = &state.data.feature_stats[feature_index];
             let histogram_min = global_stats.histogram.min;
             let histogram_max = global_stats.histogram.max;
@@ -178,15 +202,20 @@ pub async fn get_hexagon_stats(
         output.push_str("],\"enum_features\":[");
         let mut first_enum = true;
         for enum_feature in enum_features {
+            // Skip enum features not in the requested set
+            if let Some(ref set) = field_set {
+                if !set.contains(enum_feature.name.as_str()) {
+                    continue;
+                }
+            }
             let enum_index = match state.enum_name_to_idx.get(&enum_feature.name) {
                 Some(&index) => index,
                 None => continue,
             };
-            let enum_data = &state.data.enum_features[enum_index];
 
-            let mut value_counts = vec![0u64; enum_data.values.len()];
+            let mut value_counts = vec![0u64; enum_feature.values.len()];
             for &row in &matching_rows {
-                let value = enum_data.data[row];
+                let value = enum_data[row * num_enums + enum_index];
                 if value != ENUM_NULL && (value as usize) < value_counts.len() {
                     value_counts[value as usize] += 1;
                 }
@@ -215,7 +244,7 @@ pub async fn get_hexagon_stats(
                     output.push(',');
                 }
                 first_value = false;
-                write_json_string(&mut output, &enum_data.values[value_index]);
+                write_json_string(&mut output, &enum_feature.values[value_index]);
                 write!(output, ":{}", count).unwrap();
             }
             output.push_str("}}");