Support spilling
This commit is contained in:
parent
3dd5cde626
commit
8d0ecdab19
18 changed files with 150 additions and 65 deletions
1
server-rs/Cargo.lock
generated
1
server-rs/Cargo.lock
generated
|
|
@ -3887,6 +3887,7 @@ dependencies = [
|
|||
"hmac 0.13.0",
|
||||
"lasso",
|
||||
"libc",
|
||||
"memmap2",
|
||||
"metrics",
|
||||
"metrics-exporter-prometheus",
|
||||
"parking_lot",
|
||||
|
|
|
|||
|
|
@ -27,6 +27,9 @@ urlencoding = "2"
|
|||
url = "2"
|
||||
rust_xlsxwriter = "0.94"
|
||||
pmtiles = { version = "0.23", features = ["mmap-async-tokio"] }
|
||||
# Read-only mmap backing for the large property arrays under the `--spill-dir`
|
||||
# dev flag (see data/spill.rs); lets a low-memory dev box page them from disk.
|
||||
memmap2 = "0.9"
|
||||
rand = "0.10"
|
||||
hmac = "0.13"
|
||||
sha2 = "0.11"
|
||||
|
|
|
|||
|
|
@ -4,6 +4,7 @@ mod places;
|
|||
mod poi;
|
||||
mod postcodes;
|
||||
mod property;
|
||||
pub mod spill;
|
||||
pub mod travel_time;
|
||||
|
||||
/// Apostrophe-like code points that should be elided (not treated as word breaks) when
|
||||
|
|
|
|||
|
|
@ -19,6 +19,7 @@ use super::address_search::{
|
|||
use super::poi_metrics::{PostcodePoiMetrics, NO_POI_METRIC_ROW};
|
||||
use super::stats::{column_to_f32_vec, compute_feature_stats, FeatureStats, Histogram};
|
||||
use super::{HistoricalPrice, PropertyData, RenovationEvent, TenureEvent};
|
||||
use crate::data::spill::SpillVec;
|
||||
|
||||
const MISSING_COORDINATE_SAMPLE_LIMIT: usize = 10;
|
||||
const COUNTRY_COLUMN_CANDIDATES: &[&str] = &[
|
||||
|
|
@ -160,11 +161,25 @@ fn validate_no_england_rows_missing_coordinates(
|
|||
}
|
||||
|
||||
impl PropertyData {
|
||||
pub fn load(properties_path: &Path, postcode_features_path: &Path) -> anyhow::Result<Self> {
|
||||
crate::data::run_polars_io(|| Self::load_inner(properties_path, postcode_features_path))
|
||||
/// Load the property data. When `spill` is `Some(dir)`, the large flat arrays
|
||||
/// (the feature matrix and the address-search index) are written to anonymous
|
||||
/// files in `dir` and memory-mapped read-only instead of held on the heap —
|
||||
/// the `--spill-dir` dev flag, which lets a low-memory box page them from disk.
|
||||
pub fn load(
|
||||
properties_path: &Path,
|
||||
postcode_features_path: &Path,
|
||||
spill: Option<&Path>,
|
||||
) -> anyhow::Result<Self> {
|
||||
crate::data::run_polars_io(|| {
|
||||
Self::load_inner(properties_path, postcode_features_path, spill)
|
||||
})
|
||||
}
|
||||
|
||||
fn load_inner(properties_path: &Path, postcode_features_path: &Path) -> anyhow::Result<Self> {
|
||||
fn load_inner(
|
||||
properties_path: &Path,
|
||||
postcode_features_path: &Path,
|
||||
spill: Option<&Path>,
|
||||
) -> anyhow::Result<Self> {
|
||||
// Load postcode.parquet
|
||||
tracing::info!(
|
||||
"Loading postcode features from {:?}",
|
||||
|
|
@ -1011,37 +1026,71 @@ impl PropertyData {
|
|||
// Transpose to row-major AND apply spatial permutation in one pass.
|
||||
// Combines numeric and enum features into a single feature_data array, quantized to u16.
|
||||
tracing::info!("Transposing to row-major layout (spatially sorted, quantized to u16)");
|
||||
let mut feature_data = vec![NAN_U16; row_count * num_features];
|
||||
feature_data
|
||||
.par_chunks_mut(num_features)
|
||||
.enumerate()
|
||||
.for_each(|(new_row, row_slice)| {
|
||||
let old_index = perm[new_row] as usize;
|
||||
// Numeric features: quantize to u16
|
||||
for (feat_idx, col_vec) in numeric_col_major.iter().enumerate() {
|
||||
let value = col_vec[old_index];
|
||||
row_slice[feat_idx] = if value.is_finite() {
|
||||
let range = quant_range[feat_idx];
|
||||
if range > 0.0 {
|
||||
let normalized = (value - quant_min[feat_idx]) / range;
|
||||
(normalized * QUANT_SCALE).round().clamp(0.0, QUANT_SCALE) as u16
|
||||
} else {
|
||||
0
|
||||
// Built in place so that under `--spill-dir` the 3GB matrix lands straight
|
||||
// in the mmap-backed file and never also exists on the heap.
|
||||
let feature_data = SpillVec::build_u16(
|
||||
row_count * num_features,
|
||||
NAN_U16,
|
||||
spill,
|
||||
"feature_data",
|
||||
|feature_data| {
|
||||
feature_data
|
||||
.par_chunks_mut(num_features)
|
||||
.enumerate()
|
||||
.for_each(|(new_row, row_slice)| {
|
||||
let old_index = perm[new_row] as usize;
|
||||
// Numeric features: quantize to u16
|
||||
for (feat_idx, col_vec) in numeric_col_major.iter().enumerate() {
|
||||
let value = col_vec[old_index];
|
||||
row_slice[feat_idx] = if value.is_finite() {
|
||||
let range = quant_range[feat_idx];
|
||||
if range > 0.0 {
|
||||
let normalized = (value - quant_min[feat_idx]) / range;
|
||||
(normalized * QUANT_SCALE).round().clamp(0.0, QUANT_SCALE)
|
||||
as u16
|
||||
} else {
|
||||
0
|
||||
}
|
||||
} else {
|
||||
NAN_U16
|
||||
};
|
||||
}
|
||||
} else {
|
||||
NAN_U16
|
||||
};
|
||||
}
|
||||
// Enum features: store as u16 directly
|
||||
for (enum_idx, (_, encoded)) in enum_col_major.iter().enumerate() {
|
||||
let value = encoded[old_index];
|
||||
row_slice[num_numeric + enum_idx] = if value.is_finite() {
|
||||
value as u16
|
||||
} else {
|
||||
NAN_U16
|
||||
};
|
||||
}
|
||||
});
|
||||
// Enum features: store as u16 directly
|
||||
for (enum_idx, (_, encoded)) in enum_col_major.iter().enumerate() {
|
||||
let value = encoded[old_index];
|
||||
row_slice[num_numeric + enum_idx] = if value.is_finite() {
|
||||
value as u16
|
||||
} else {
|
||||
NAN_U16
|
||||
};
|
||||
}
|
||||
});
|
||||
},
|
||||
)?;
|
||||
|
||||
// Spill the remaining large flat arrays (the address-search index) to disk
|
||||
// when configured. The posting-list hashmaps and interners stay on the heap
|
||||
// — they aren't flat arrays, and they're a small fraction of the footprint.
|
||||
let address_buffer =
|
||||
SpillVec::maybe_spill(address_buffer.into_bytes(), spill, "address_buffer")?;
|
||||
let address_offsets = SpillVec::maybe_spill(address_offsets, spill, "address_offsets")?;
|
||||
let address_lengths = SpillVec::maybe_spill(address_lengths, spill, "address_lengths")?;
|
||||
let address_search_token_keys = SpillVec::maybe_spill(
|
||||
address_search_token_keys,
|
||||
spill,
|
||||
"address_search_token_keys",
|
||||
)?;
|
||||
let address_search_token_offsets = SpillVec::maybe_spill(
|
||||
address_search_token_offsets,
|
||||
spill,
|
||||
"address_search_token_offsets",
|
||||
)?;
|
||||
let address_search_token_lengths = SpillVec::maybe_spill(
|
||||
address_search_token_lengths,
|
||||
spill,
|
||||
"address_search_token_lengths",
|
||||
)?;
|
||||
let postcode_keys = SpillVec::maybe_spill(postcode_keys, spill, "postcode_keys")?;
|
||||
|
||||
tracing::info!("Data loading complete");
|
||||
|
||||
|
|
|
|||
|
|
@ -26,6 +26,7 @@ use rustc_hash::FxHashMap;
|
|||
use serde::Serialize;
|
||||
|
||||
use crate::consts::NAN_U16;
|
||||
use crate::data::spill::SpillVec;
|
||||
|
||||
#[derive(Serialize, Clone)]
|
||||
pub struct RenovationEvent {
|
||||
|
|
@ -61,7 +62,9 @@ pub struct PropertyData {
|
|||
/// Quantized to u16. NaN sentinel = u16::MAX (65535).
|
||||
/// Numeric features: encoded via (val - min) / range * 65534.
|
||||
/// Enum features: stored directly as u16 cast of the f32 index.
|
||||
pub feature_data: Vec<u16>,
|
||||
/// Heap-resident in prod; mmap-backed under `--spill-dir` (the largest array,
|
||||
/// so the biggest dev-memory win). Derefs to `[u16]`, so reads are unchanged.
|
||||
pub feature_data: SpillVec<u16>,
|
||||
/// Per-feature: range / QUANT_SCALE for fast decode.
|
||||
dequant_a: Vec<f32>,
|
||||
/// Per-feature: minimum value (offset for dequantization).
|
||||
|
|
@ -72,15 +75,16 @@ pub struct PropertyData {
|
|||
pub poi_metrics: PostcodePoiMetrics,
|
||||
/// Unquantized last sale price used by the price-history chart.
|
||||
last_known_price_raw: Vec<f32>,
|
||||
/// Contiguous buffer holding all address strings end-to-end.
|
||||
address_buffer: String,
|
||||
/// Contiguous buffer holding all address strings end-to-end (valid UTF-8 by
|
||||
/// construction; spillable, so stored as bytes rather than `String`).
|
||||
address_buffer: SpillVec<u8>,
|
||||
/// Byte offset into `address_buffer` where each row's address starts.
|
||||
address_offsets: Vec<u32>,
|
||||
address_offsets: SpillVec<u32>,
|
||||
/// Length in bytes of each row's address.
|
||||
address_lengths: Vec<u16>,
|
||||
address_lengths: SpillVec<u16>,
|
||||
/// Interned postcodes: reader is thread-safe, keys index into it.
|
||||
postcode_interner: lasso::RodeoReader,
|
||||
postcode_keys: Vec<lasso::Spur>,
|
||||
postcode_keys: SpillVec<lasso::Spur>,
|
||||
/// Rows for each postcode, keyed by the interned postcode key.
|
||||
postcode_row_index: FxHashMap<lasso::Spur, Vec<u32>>,
|
||||
/// Inverted index from address tokens to property rows.
|
||||
|
|
@ -91,11 +95,11 @@ pub struct PropertyData {
|
|||
/// Resolve-only (no string->key reverse map): scoring only resolves keys.
|
||||
address_search_interner: lasso::RodeoResolver,
|
||||
/// Flat per-row normalized address-search token keys.
|
||||
address_search_token_keys: Vec<lasso::Spur>,
|
||||
address_search_token_keys: SpillVec<lasso::Spur>,
|
||||
/// Offset into `address_search_token_keys` for each row.
|
||||
address_search_token_offsets: Vec<u32>,
|
||||
address_search_token_offsets: SpillVec<u32>,
|
||||
/// Number of normalized address-search token keys for each row.
|
||||
address_search_token_lengths: Vec<u16>,
|
||||
address_search_token_lengths: SpillVec<u16>,
|
||||
/// For enum features: maps feature index to list of possible string values.
|
||||
/// Index in values list corresponds to the u16 value stored in feature_data.
|
||||
pub enum_values: rustc_hash::FxHashMap<usize, Vec<String>>,
|
||||
|
|
@ -123,7 +127,11 @@ impl PropertyData {
|
|||
pub fn address(&self, row: usize) -> &str {
|
||||
let offset = self.address_offsets[row] as usize;
|
||||
let length = self.address_lengths[row] as usize;
|
||||
&self.address_buffer[offset..offset + length]
|
||||
let bytes = &self.address_buffer[offset..offset + length];
|
||||
// SAFETY: `address_buffer` is built by concatenating `&str` slices on UTF-8
|
||||
// char boundaries, so every recorded `[offset, offset + length)` range is
|
||||
// valid UTF-8.
|
||||
unsafe { std::str::from_utf8_unchecked(bytes) }
|
||||
}
|
||||
|
||||
/// Get the postcode string for a given row.
|
||||
|
|
@ -133,7 +141,7 @@ impl PropertyData {
|
|||
|
||||
/// Get postcode components for field-level borrowing (avoids conflicting borrows with feature_data).
|
||||
pub fn postcode_parts(&self) -> (&lasso::RodeoReader, &[lasso::Spur]) {
|
||||
(&self.postcode_interner, &self.postcode_keys)
|
||||
(&self.postcode_interner, self.postcode_keys.as_slice())
|
||||
}
|
||||
|
||||
/// Property rows for a given postcode string, or empty if unknown.
|
||||
|
|
@ -229,25 +237,25 @@ impl PropertyData {
|
|||
feature_names: Vec::new(),
|
||||
num_features: 0,
|
||||
num_numeric: 0,
|
||||
feature_data: Vec::new(),
|
||||
feature_data: SpillVec::owned(Vec::new()),
|
||||
dequant_a: Vec::new(),
|
||||
quant_min: Vec::new(),
|
||||
quant_range: Vec::new(),
|
||||
feature_stats: Vec::new(),
|
||||
poi_metrics: PostcodePoiMetrics::empty(0),
|
||||
last_known_price_raw: Vec::new(),
|
||||
address_buffer: String::new(),
|
||||
address_offsets: Vec::new(),
|
||||
address_lengths: Vec::new(),
|
||||
address_buffer: SpillVec::owned(Vec::new()),
|
||||
address_offsets: SpillVec::owned(Vec::new()),
|
||||
address_lengths: SpillVec::owned(Vec::new()),
|
||||
postcode_interner: lasso::Rodeo::default().into_reader(),
|
||||
postcode_keys: Vec::new(),
|
||||
postcode_keys: SpillVec::owned(Vec::new()),
|
||||
postcode_row_index: FxHashMap::default(),
|
||||
address_token_index: FxHashMap::default(),
|
||||
address_prefix_index: FxHashMap::default(),
|
||||
address_search_interner: lasso::Rodeo::default().into_resolver(),
|
||||
address_search_token_keys: Vec::new(),
|
||||
address_search_token_offsets: Vec::new(),
|
||||
address_search_token_lengths: Vec::new(),
|
||||
address_search_token_keys: SpillVec::owned(Vec::new()),
|
||||
address_search_token_offsets: SpillVec::owned(Vec::new()),
|
||||
address_search_token_lengths: SpillVec::owned(Vec::new()),
|
||||
enum_values: rustc_hash::FxHashMap::default(),
|
||||
enum_counts: rustc_hash::FxHashMap::default(),
|
||||
approx_build_date_bits: Vec::new(),
|
||||
|
|
|
|||
|
|
@ -281,6 +281,14 @@ struct Cli {
|
|||
#[arg(long)]
|
||||
dist: Option<PathBuf>,
|
||||
|
||||
/// Dev only: spill the large property arrays (feature matrix + address-search
|
||||
/// index) to anonymous files in this directory and memory-map them read-only,
|
||||
/// instead of holding them on the heap. Trades a little speed for a much
|
||||
/// smaller resident set so a low-memory dev box can run the full dataset; the
|
||||
/// mapped pages are file-backed and reclaimable under pressure. Omit in prod.
|
||||
#[arg(long, env = "SPILL_DIR")]
|
||||
spill_dir: Option<PathBuf>,
|
||||
|
||||
/// URL of the screenshot service (e.g. http://screenshot:8002)
|
||||
#[arg(long, env = "SCREENSHOT_URL")]
|
||||
screenshot_url: String,
|
||||
|
|
@ -465,7 +473,15 @@ async fn main() -> anyhow::Result<()> {
|
|||
cli.properties.display(),
|
||||
cli.postcode_features.display(),
|
||||
);
|
||||
let property_data = data::PropertyData::load(&cli.properties, &cli.postcode_features)?;
|
||||
let spill_dir = cli.spill_dir.as_deref();
|
||||
if let Some(dir) = spill_dir {
|
||||
info!(
|
||||
"Spill-to-disk enabled: large property arrays will be memory-mapped from {}",
|
||||
dir.display()
|
||||
);
|
||||
}
|
||||
let property_data =
|
||||
data::PropertyData::load(&cli.properties, &cli.postcode_features, spill_dir)?;
|
||||
trim_allocator("property data load");
|
||||
info!(
|
||||
rows = property_data.lat.len(),
|
||||
|
|
|
|||
|
|
@ -48,7 +48,7 @@ const KEEP_UNKNOWN_LISTING_FILTER_FEATURES: &[&str] = &[
|
|||
"Construction year",
|
||||
"Interior height (m)",
|
||||
"Current energy rating",
|
||||
"Potential energy rating"
|
||||
"Potential energy rating",
|
||||
];
|
||||
const LISTING_BOUNDS_EPSILON_DEGREES: f64 = 0.00001;
|
||||
|
||||
|
|
|
|||
|
|
@ -105,7 +105,7 @@ pub(super) fn count_matching_rows(
|
|||
.collect();
|
||||
let has_travel = !travel_data.is_empty();
|
||||
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
let num_features = state.data.num_features;
|
||||
let num_rows = state.data.lat.len();
|
||||
let (pc_interner, pc_keys) = state.data.postcode_parts();
|
||||
|
|
|
|||
|
|
@ -632,7 +632,7 @@ pub async fn get_export(
|
|||
let bytes = tokio::task::spawn_blocking(move || -> Result<Vec<u8>, String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
let num_features = state.data.num_features;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
let quant = state.data.quant_ref();
|
||||
let feature_names = &state.data.feature_names;
|
||||
let enum_values = &state.data.enum_values;
|
||||
|
|
|
|||
|
|
@ -88,7 +88,7 @@ pub async fn get_filter_counts(
|
|||
let response = tokio::task::spawn_blocking(move || -> Result<FilterCountsResponse, String> {
|
||||
let t0 = std::time::Instant::now();
|
||||
let num_features = state.data.num_features;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
|
||||
// Load travel time data
|
||||
let travel_data: Vec<TravelData> = travel_entries
|
||||
|
|
|
|||
|
|
@ -226,7 +226,7 @@ pub(super) fn top_filter_exclusions(
|
|||
return Vec::new();
|
||||
}
|
||||
|
||||
let feature_data = &data.feature_data;
|
||||
let feature_data: &[u16] = &data.feature_data;
|
||||
let num_features = data.num_features;
|
||||
let quant = data.quant_ref();
|
||||
let poi_quant = data.poi_metrics.quant_ref();
|
||||
|
|
@ -494,7 +494,7 @@ pub async fn get_hexagon_stats(
|
|||
.map_err(|err| format!("Invalid H3 resolution {}: {}", resolution, err))?;
|
||||
let need_parent = needs_parent(resolution);
|
||||
let num_features = state.data.num_features;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
let travel_data = load_travel_data(&state.travel_time_store, &travel_entries)?;
|
||||
let has_travel = !travel_entries.is_empty();
|
||||
|
||||
|
|
|
|||
|
|
@ -332,7 +332,7 @@ pub async fn get_hexagons(
|
|||
.collect();
|
||||
|
||||
let num_features = state.data.num_features;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
let quant = state.data.quant_ref();
|
||||
let (pc_interner, pc_keys) = state.data.postcode_parts();
|
||||
let min_keys = &state.min_keys;
|
||||
|
|
|
|||
|
|
@ -108,7 +108,7 @@ pub async fn get_postcode_properties(
|
|||
let result = 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 feature_data: &[u16] = &state.data.feature_data;
|
||||
let feature_names = &state.data.feature_names;
|
||||
let feature_name_to_index = &state.feature_name_to_index;
|
||||
let enum_values = &state.data.enum_values;
|
||||
|
|
|
|||
|
|
@ -92,7 +92,7 @@ pub async fn get_postcode_stats(
|
|||
let response = tokio::task::spawn_blocking(move || {
|
||||
let start_time = std::time::Instant::now();
|
||||
let num_features = state.data.num_features;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
let travel_data = load_travel_data(&state.travel_time_store, &travel_entries)?;
|
||||
let has_travel = !travel_entries.is_empty();
|
||||
|
||||
|
|
|
|||
|
|
@ -131,7 +131,7 @@ pub async fn get_postcodes(
|
|||
.collect();
|
||||
|
||||
let num_features = state.data.num_features;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
let quant = state.data.quant_ref();
|
||||
let min_keys = &state.min_keys;
|
||||
let max_keys = &state.max_keys;
|
||||
|
|
|
|||
|
|
@ -332,7 +332,7 @@ pub async fn get_hexagon_properties(
|
|||
.map_err(|err| format!("Invalid H3 resolution {}: {}", resolution, err))?;
|
||||
let need_parent = needs_parent(resolution);
|
||||
let num_features = state.data.num_features;
|
||||
let feature_data = &state.data.feature_data;
|
||||
let feature_data: &[u16] = &state.data.feature_data;
|
||||
let feature_names = &state.data.feature_names;
|
||||
let feature_name_to_index = &state.feature_name_to_index;
|
||||
let enum_values = &state.data.enum_values;
|
||||
|
|
|
|||
|
|
@ -482,12 +482,13 @@ pub fn compute_poi_feature_stats(
|
|||
mod tests {
|
||||
use super::*;
|
||||
use crate::consts::NAN_U16;
|
||||
use crate::data::spill::SpillVec;
|
||||
|
||||
fn enum_data(values: &[u16]) -> PropertyData {
|
||||
let mut data = PropertyData::empty_for_tests();
|
||||
data.num_features = 1;
|
||||
data.num_numeric = 0; // single enum feature at index 0
|
||||
data.feature_data = values.to_vec();
|
||||
data.feature_data = SpillVec::owned(values.to_vec());
|
||||
data.enum_values
|
||||
.insert(0, vec!["Yes".to_string(), "No".to_string()]);
|
||||
data
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue