Support spilling
This commit is contained in:
parent
3dd5cde626
commit
8d0ecdab19
18 changed files with 150 additions and 65 deletions
|
|
@ -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");
|
||||
|
||||
|
|
|
|||
Loading…
Add table
Add a link
Reference in a new issue