Support spilling

This commit is contained in:
Andras Schmelczer 2026-06-20 11:03:04 +01:00
parent 3dd5cde626
commit 8d0ecdab19
18 changed files with 150 additions and 65 deletions

View file

@ -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");