1046 lines
35 KiB
Rust
1046 lines
35 KiB
Rust
use std::path::Path;
|
||
|
||
use anyhow::Context;
|
||
use polars::frame::DataFrame;
|
||
use polars::lazy::frame::LazyFrame;
|
||
use polars::prelude::*;
|
||
use rustc_hash::FxHashMap;
|
||
use tracing::info;
|
||
|
||
use crate::utils::InternedColumn;
|
||
|
||
/// Upper bound on place rows scored per query (candidate sets are normally far smaller).
|
||
const PLACE_CANDIDATE_LIMIT: usize = 50_000;
|
||
const PLACE_PREFIX_MIN_LEN: usize = 2;
|
||
const PLACE_PREFIX_MAX_LEN: usize = 6;
|
||
|
||
pub struct PlaceData {
|
||
pub name: Vec<String>,
|
||
pub name_lower: Vec<String>,
|
||
pub name_search: Vec<String>,
|
||
pub place_type: InternedColumn,
|
||
pub type_rank: Vec<u8>,
|
||
pub population: Vec<u32>,
|
||
pub lat: Vec<f32>,
|
||
pub lon: Vec<f32>,
|
||
pub city: Vec<Option<String>>,
|
||
pub travel_destination: Vec<bool>,
|
||
/// Inverted index from an alias token to the (ascending) place rows containing it. Lets place
|
||
/// search gather candidates instead of scanning all ~1M+ rows per keystroke.
|
||
token_index: FxHashMap<String, Vec<u32>>,
|
||
/// Prefix → indexed tokens, for matching a partially-typed final word.
|
||
token_prefix_index: FxHashMap<String, Vec<String>>,
|
||
/// Trigram → fuzzy-eligible rows (settlements/stations only), for bounded typo matching.
|
||
fuzzy_trigram_index: FxHashMap<u32, Vec<u32>>,
|
||
}
|
||
|
||
#[derive(Clone, Copy)]
|
||
pub(super) struct CityCandidate<'a> {
|
||
name: &'a str,
|
||
lat: f32,
|
||
lon: f32,
|
||
population: u32,
|
||
max_dist_sq: f32,
|
||
}
|
||
|
||
const PARENT_CITY_MAX_DIST_SQ: f32 = 0.81;
|
||
const LONDON_DISPLAY_MAX_DEGREES: f32 = 30.0 / 111.0;
|
||
const LONDON_DISPLAY_MAX_DIST_SQ: f32 = LONDON_DISPLAY_MAX_DEGREES * LONDON_DISPLAY_MAX_DEGREES;
|
||
const SUBSUMED_CITY_MAX_DEGREES: f32 = 5.0 / 111.0;
|
||
const SUBSUMED_CITY_MAX_DIST_SQ: f32 = SUBSUMED_CITY_MAX_DEGREES * SUBSUMED_CITY_MAX_DEGREES;
|
||
const SUBSUMED_CITY_MIN_POPULATION_RATIO: u32 = 10;
|
||
|
||
fn type_rank(place_type: &str) -> u8 {
|
||
match place_type {
|
||
"city" => 0,
|
||
"town" => 1,
|
||
"village" => 2,
|
||
"suburb" | "neighbourhood" | "quarter" | "borough" | "locality" => 3,
|
||
"station" | "university" => 4,
|
||
"hamlet" | "isolated_dwelling" | "island" => 5,
|
||
_ => 6,
|
||
}
|
||
}
|
||
|
||
pub fn is_travel_destination_type(place_type: &str) -> bool {
|
||
matches!(place_type, "city" | "station" | "university")
|
||
}
|
||
|
||
impl<'a> CityCandidate<'a> {
|
||
fn from_place(name: &'a str, lat: f32, lon: f32, population: u32) -> Self {
|
||
let max_dist_sq = if name == "London" {
|
||
LONDON_DISPLAY_MAX_DIST_SQ
|
||
} else {
|
||
PARENT_CITY_MAX_DIST_SQ
|
||
};
|
||
|
||
Self {
|
||
name,
|
||
lat,
|
||
lon,
|
||
population,
|
||
max_dist_sq,
|
||
}
|
||
}
|
||
|
||
fn distance_sq(&self, lat: f32, lon: f32, cos_lat: f32) -> f32 {
|
||
let dlat = self.lat - lat;
|
||
let dlon = (self.lon - lon) * cos_lat;
|
||
dlat * dlat + dlon * dlon
|
||
}
|
||
|
||
fn is_subsumed_by(&self, other: &Self) -> bool {
|
||
if self.population == 0 {
|
||
return false;
|
||
}
|
||
|
||
let min_parent_population =
|
||
u64::from(self.population) * u64::from(SUBSUMED_CITY_MIN_POPULATION_RATIO);
|
||
if u64::from(other.population) < min_parent_population {
|
||
return false;
|
||
}
|
||
|
||
other.distance_sq(self.lat, self.lon, self.lat.to_radians().cos())
|
||
< SUBSUMED_CITY_MAX_DIST_SQ
|
||
}
|
||
}
|
||
|
||
pub(super) fn display_city_candidates<'a>(
|
||
names: &'a [String],
|
||
type_rank: &[u8],
|
||
population: &[u32],
|
||
lat: &[f32],
|
||
lon: &[f32],
|
||
) -> Vec<CityCandidate<'a>> {
|
||
let cities: Vec<CityCandidate<'_>> = type_rank
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(idx, &rank)| {
|
||
if rank == 0 {
|
||
Some(CityCandidate::from_place(
|
||
&names[idx],
|
||
lat[idx],
|
||
lon[idx],
|
||
population[idx],
|
||
))
|
||
} else {
|
||
None
|
||
}
|
||
})
|
||
.collect();
|
||
|
||
cities
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(idx, city)| {
|
||
let is_subsumed = cities
|
||
.iter()
|
||
.enumerate()
|
||
.any(|(other_idx, other)| other_idx != idx && city.is_subsumed_by(other));
|
||
(!is_subsumed).then_some(*city)
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
pub(super) fn nearest_display_city<'a>(
|
||
lat: f32,
|
||
lon: f32,
|
||
cities: &'a [CityCandidate<'a>],
|
||
) -> Option<&'a str> {
|
||
let cos_lat = lat.to_radians().cos();
|
||
let (best_city, best_dist_sq) = cities
|
||
.iter()
|
||
.map(|city| (city, city.distance_sq(lat, lon, cos_lat)))
|
||
.min_by(|(_, lhs), (_, rhs)| lhs.total_cmp(rhs))?;
|
||
|
||
(best_dist_sq < best_city.max_dist_sq).then_some(best_city.name)
|
||
}
|
||
|
||
pub fn normalize_search_text(text: &str) -> String {
|
||
let mut result = String::with_capacity(text.len());
|
||
let mut last_was_space = true;
|
||
|
||
for ch in text.chars() {
|
||
if super::is_apostrophe(ch) {
|
||
continue;
|
||
}
|
||
|
||
let lower = ch.to_ascii_lowercase();
|
||
if lower.is_ascii_alphanumeric() {
|
||
result.push(lower);
|
||
last_was_space = false;
|
||
} else if !last_was_space {
|
||
result.push(' ');
|
||
last_was_space = true;
|
||
}
|
||
}
|
||
|
||
if result.ends_with(' ') {
|
||
result.pop();
|
||
}
|
||
result
|
||
}
|
||
|
||
/// Tokens across all of a place's search aliases (split on word and alias separators),
|
||
/// for token-AND matching where every query word must prefix-match some place token.
|
||
pub fn place_alias_tokens(search_text: &str) -> impl Iterator<Item = &str> {
|
||
search_text
|
||
.split([' ', '|'])
|
||
.filter(|token| !token.is_empty())
|
||
}
|
||
|
||
fn trigram_hash(first: char, second: char, third: char) -> u32 {
|
||
let mut hash = 2_166_136_261u32;
|
||
for ch in [first, second, third] {
|
||
hash = (hash ^ (ch as u32)).wrapping_mul(16_777_619);
|
||
}
|
||
hash
|
||
}
|
||
|
||
/// Sorted, de-duplicated padded character trigrams of `text`, for Jaccard fuzzy matching.
|
||
pub fn compute_trigrams(text: &str) -> Vec<u32> {
|
||
let norm = normalize_search_text(text);
|
||
if norm.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
let chars: Vec<char> = [' ', ' ']
|
||
.into_iter()
|
||
.chain(norm.chars())
|
||
.chain(std::iter::once(' '))
|
||
.collect();
|
||
let mut grams: Vec<u32> = chars
|
||
.windows(3)
|
||
.map(|window| trigram_hash(window[0], window[1], window[2]))
|
||
.collect();
|
||
grams.sort_unstable();
|
||
grams.dedup();
|
||
grams
|
||
}
|
||
|
||
/// Intersect two ascending-sorted row-id slices.
|
||
fn intersect_sorted(left: &[u32], right: &[u32]) -> Vec<u32> {
|
||
let mut out = Vec::new();
|
||
let (mut i, mut j) = (0, 0);
|
||
while i < left.len() && j < right.len() {
|
||
match left[i].cmp(&right[j]) {
|
||
std::cmp::Ordering::Less => i += 1,
|
||
std::cmp::Ordering::Greater => j += 1,
|
||
std::cmp::Ordering::Equal => {
|
||
out.push(left[i]);
|
||
i += 1;
|
||
j += 1;
|
||
}
|
||
}
|
||
}
|
||
out
|
||
}
|
||
|
||
/// Union two ascending-sorted row-id slices (deduplicated, stays sorted).
|
||
fn union_sorted(left: &[u32], right: &[u32]) -> Vec<u32> {
|
||
let mut out = Vec::with_capacity(left.len() + right.len());
|
||
let (mut i, mut j) = (0, 0);
|
||
while i < left.len() && j < right.len() {
|
||
match left[i].cmp(&right[j]) {
|
||
std::cmp::Ordering::Less => {
|
||
out.push(left[i]);
|
||
i += 1;
|
||
}
|
||
std::cmp::Ordering::Greater => {
|
||
out.push(right[j]);
|
||
j += 1;
|
||
}
|
||
std::cmp::Ordering::Equal => {
|
||
out.push(left[i]);
|
||
i += 1;
|
||
j += 1;
|
||
}
|
||
}
|
||
}
|
||
out.extend_from_slice(&left[i..]);
|
||
out.extend_from_slice(&right[j..]);
|
||
out
|
||
}
|
||
|
||
/// Distinct indexable tokens (len ≥ 2) across all of a place's search aliases. ASCII because
|
||
/// `normalize_search_text` already dropped non-alphanumerics, so prefix byte-slicing is safe.
|
||
fn place_index_tokens(search_text: &str) -> Vec<String> {
|
||
let mut tokens: Vec<String> = place_alias_tokens(search_text)
|
||
.filter(|token| token.len() >= 2)
|
||
.map(ToString::to_string)
|
||
.collect();
|
||
tokens.sort_unstable();
|
||
tokens.dedup();
|
||
tokens
|
||
}
|
||
|
||
fn build_place_prefix_index(
|
||
token_index: &FxHashMap<String, Vec<u32>>,
|
||
) -> FxHashMap<String, Vec<String>> {
|
||
let mut prefix_index: FxHashMap<String, Vec<String>> = FxHashMap::default();
|
||
for token in token_index.keys() {
|
||
let max_len = token.len().min(PLACE_PREFIX_MAX_LEN);
|
||
for len in PLACE_PREFIX_MIN_LEN..=max_len {
|
||
prefix_index
|
||
.entry(token[..len].to_string())
|
||
.or_default()
|
||
.push(token.clone());
|
||
}
|
||
}
|
||
for tokens in prefix_index.values_mut() {
|
||
tokens.sort_unstable();
|
||
tokens.dedup();
|
||
}
|
||
prefix_index
|
||
}
|
||
|
||
/// Whether a place type participates in fuzzy (typo) matching. Settlements/stations/universities
|
||
/// do; the ~1M streets and POIs do not (people rarely misspell a road and it keeps fuzzy bounded).
|
||
fn is_fuzzy_eligible_type(place_type: &str) -> bool {
|
||
!matches!(
|
||
place_type,
|
||
"street" | "park" | "attraction" | "hospital" | "retail"
|
||
)
|
||
}
|
||
|
||
/// Jaccard similarity between two sorted trigram sets (0.0–1.0).
|
||
pub fn trigram_similarity(left: &[u32], right: &[u32]) -> f32 {
|
||
if left.is_empty() || right.is_empty() {
|
||
return 0.0;
|
||
}
|
||
let (mut i, mut j, mut intersection) = (0, 0, 0usize);
|
||
while i < left.len() && j < right.len() {
|
||
match left[i].cmp(&right[j]) {
|
||
std::cmp::Ordering::Less => i += 1,
|
||
std::cmp::Ordering::Greater => j += 1,
|
||
std::cmp::Ordering::Equal => {
|
||
intersection += 1;
|
||
i += 1;
|
||
j += 1;
|
||
}
|
||
}
|
||
}
|
||
let union = left.len() + right.len() - intersection;
|
||
intersection as f32 / union as f32
|
||
}
|
||
|
||
fn replace_token(text: &str, from: &str, to: &str) -> Option<String> {
|
||
let mut changed = false;
|
||
let replaced: Vec<&str> = text
|
||
.split_whitespace()
|
||
.map(|token| {
|
||
if token == from {
|
||
changed = true;
|
||
to
|
||
} else {
|
||
token
|
||
}
|
||
})
|
||
.collect();
|
||
|
||
changed.then(|| replaced.join(" "))
|
||
}
|
||
|
||
fn push_alias(aliases: &mut Vec<String>, alias: String) {
|
||
if !alias.is_empty() && !aliases.iter().any(|existing| existing == &alias) {
|
||
aliases.push(alias);
|
||
}
|
||
}
|
||
|
||
/// Bidirectional token abbreviations expanded into search aliases so a query typed either
|
||
/// way matches (e.g. "gt missenden" ↔ "Great Missenden", "mt" ↔ "Mount").
|
||
const PLACE_TOKEN_ALIASES: &[(&str, &str)] = &[
|
||
("st", "saint"),
|
||
("saint", "st"),
|
||
("mt", "mount"),
|
||
("mount", "mt"),
|
||
("gt", "great"),
|
||
("great", "gt"),
|
||
("lt", "little"),
|
||
("little", "lt"),
|
||
("upr", "upper"),
|
||
("upper", "upr"),
|
||
("lwr", "lower"),
|
||
("lower", "lwr"),
|
||
];
|
||
|
||
fn build_search_text(name: &str, place_type: &str) -> String {
|
||
let primary = normalize_search_text(name);
|
||
let mut aliases = vec![primary.clone()];
|
||
|
||
for (from, to) in PLACE_TOKEN_ALIASES {
|
||
if let Some(alias) = replace_token(&primary, from, to) {
|
||
push_alias(&mut aliases, alias);
|
||
}
|
||
}
|
||
|
||
if place_type == "station" {
|
||
let suffix_aliases: [(&str, &[&str]); 6] = [
|
||
(
|
||
" tube station",
|
||
&[" underground station", " station", " tube", " underground"],
|
||
),
|
||
(
|
||
" underground station",
|
||
&[" tube station", " station", " tube", " underground"],
|
||
),
|
||
(
|
||
" railway station",
|
||
&[" rail station", " station", " railway", " rail"],
|
||
),
|
||
(
|
||
" overground station",
|
||
&[" station", " overground", " railway station"],
|
||
),
|
||
(
|
||
" elizabeth line station",
|
||
&[" station", " elizabeth line", " crossrail station"],
|
||
),
|
||
(" dlr station", &[" station", " dlr"]),
|
||
];
|
||
|
||
for (suffix, replacements) in suffix_aliases {
|
||
if let Some(stem) = primary.strip_suffix(suffix) {
|
||
for replacement in replacements {
|
||
push_alias(&mut aliases, format!("{stem}{replacement}"));
|
||
}
|
||
}
|
||
}
|
||
}
|
||
|
||
aliases.join(" | ")
|
||
}
|
||
|
||
fn extract_str_col(df: &DataFrame, name: &str) -> anyhow::Result<Vec<String>> {
|
||
let column = df
|
||
.column(name)
|
||
.with_context(|| format!("Missing column '{name}' in places data"))?;
|
||
let string_column = column
|
||
.str()
|
||
.with_context(|| format!("Column '{name}' is not a string column"))?;
|
||
string_column
|
||
.into_iter()
|
||
.enumerate()
|
||
.map(|(row, value)| {
|
||
value
|
||
.map(ToString::to_string)
|
||
.with_context(|| format!("Column '{name}' has null at row {row}"))
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn extract_f32_col(df: &DataFrame, name: &str) -> anyhow::Result<Vec<f32>> {
|
||
let column = df
|
||
.column(name)
|
||
.with_context(|| format!("Missing column '{name}' in places data"))?;
|
||
let cast = column
|
||
.cast(&DataType::Float32)
|
||
.with_context(|| format!("Failed to cast column '{name}' to Float32"))?;
|
||
let float_column = cast
|
||
.f32()
|
||
.with_context(|| format!("Column '{name}' is not a float32 column"))?;
|
||
float_column
|
||
.into_iter()
|
||
.enumerate()
|
||
.map(|(row, value)| value.with_context(|| format!("Column '{name}' has null at row {row}")))
|
||
.collect()
|
||
}
|
||
|
||
fn extract_bool_col(df: &DataFrame, name: &str) -> anyhow::Result<Vec<bool>> {
|
||
let column = df
|
||
.column(name)
|
||
.with_context(|| format!("Missing column '{name}' in places data"))?;
|
||
let bool_column = column
|
||
.bool()
|
||
.with_context(|| format!("Column '{name}' is not a boolean column"))?;
|
||
bool_column
|
||
.into_iter()
|
||
.enumerate()
|
||
.map(|(row, value)| value.with_context(|| format!("Column '{name}' has null at row {row}")))
|
||
.collect()
|
||
}
|
||
|
||
fn extract_optional_str_col(
|
||
df: &DataFrame,
|
||
name: &str,
|
||
) -> anyhow::Result<Option<Vec<Option<String>>>> {
|
||
let column = match df.column(name) {
|
||
Ok(column) => column,
|
||
Err(_) => return Ok(None),
|
||
};
|
||
let string_column = column
|
||
.str()
|
||
.with_context(|| format!("Column '{name}' is not a string column"))?;
|
||
Ok(Some(
|
||
string_column
|
||
.into_iter()
|
||
.map(|value| value.map(ToString::to_string))
|
||
.collect(),
|
||
))
|
||
}
|
||
|
||
impl PlaceData {
|
||
pub fn load(parquet_path: &Path) -> anyhow::Result<Self> {
|
||
super::run_polars_io(|| Self::load_inner(parquet_path))
|
||
}
|
||
|
||
fn load_inner(parquet_path: &Path) -> anyhow::Result<Self> {
|
||
info!("Loading place data from {:?}...", parquet_path);
|
||
|
||
let parquet_path = PlRefPath::try_from_path(parquet_path)
|
||
.context("Failed to normalize places parquet path")?;
|
||
let df = LazyFrame::scan_parquet(parquet_path, Default::default())
|
||
.context("Failed to scan places parquet")?
|
||
.collect()
|
||
.context("Failed to read places parquet")?;
|
||
|
||
let row_count = df.height();
|
||
info!("Loaded {} places", row_count);
|
||
|
||
let name = extract_str_col(&df, "name")?;
|
||
let place_type_raw = extract_str_col(&df, "place_type")?;
|
||
let lat = extract_f32_col(&df, "lat")?;
|
||
let lon = extract_f32_col(&df, "lon")?;
|
||
let population: Vec<u32> = if df.column("population").is_ok() {
|
||
let pop_f32 = extract_f32_col(&df, "population")?;
|
||
pop_f32
|
||
.iter()
|
||
.map(|&val| val.max(0.0).min(u32::MAX as f32) as u32)
|
||
.collect()
|
||
} else {
|
||
vec![0; row_count]
|
||
};
|
||
|
||
let name_lower: Vec<String> = name.iter().map(|nm| nm.to_lowercase()).collect();
|
||
let name_search: Vec<String> = name
|
||
.iter()
|
||
.zip(&place_type_raw)
|
||
.map(|(nm, pt)| build_search_text(nm, pt))
|
||
.collect();
|
||
let type_rank_vec: Vec<u8> = place_type_raw.iter().map(|pt| type_rank(pt)).collect();
|
||
let place_type = InternedColumn::build(&place_type_raw);
|
||
let travel_destination = if df.column("travel_destination").is_ok() {
|
||
extract_bool_col(&df, "travel_destination")?
|
||
} else {
|
||
place_type_raw
|
||
.iter()
|
||
.map(|place_type| is_travel_destination_type(place_type))
|
||
.collect()
|
||
};
|
||
let display_city_override = extract_optional_str_col(&df, "display_city")?;
|
||
|
||
// Precompute nearest city for each non-city place
|
||
let city_candidates =
|
||
display_city_candidates(&name, &type_rank_vec, &population, &lat, &lon);
|
||
|
||
let fallback_city: Vec<Option<String>> = (0..row_count)
|
||
.map(|idx| {
|
||
if type_rank_vec[idx] == 0 {
|
||
return None; // Cities don't need a city label
|
||
}
|
||
nearest_display_city(lat[idx], lon[idx], &city_candidates).map(str::to_string)
|
||
})
|
||
.collect();
|
||
|
||
let city: Vec<Option<String>> = if let Some(display_city_override) = display_city_override {
|
||
fallback_city
|
||
.into_iter()
|
||
.zip(display_city_override)
|
||
.enumerate()
|
||
.map(|(idx, (fallback, override_city))| {
|
||
if type_rank_vec[idx] == 0 {
|
||
return None;
|
||
}
|
||
override_city
|
||
.and_then(|value| {
|
||
let trimmed = value.trim();
|
||
(!trimmed.is_empty()).then(|| trimmed.to_string())
|
||
})
|
||
.or(fallback)
|
||
})
|
||
.collect()
|
||
} else {
|
||
fallback_city
|
||
};
|
||
|
||
// Build the place search index: an inverted token index over all rows (so the per-query
|
||
// cost scales with matched candidates, not the ~1M-row corpus), plus a trigram index over
|
||
// only fuzzy-eligible rows for bounded typo matching.
|
||
let mut token_index: FxHashMap<String, Vec<u32>> = FxHashMap::default();
|
||
let mut fuzzy_trigram_index: FxHashMap<u32, Vec<u32>> = FxHashMap::default();
|
||
for idx in 0..row_count {
|
||
for token in place_index_tokens(&name_search[idx]) {
|
||
token_index.entry(token).or_default().push(idx as u32);
|
||
}
|
||
if is_fuzzy_eligible_type(&place_type_raw[idx]) {
|
||
for trigram in compute_trigrams(&name[idx]) {
|
||
fuzzy_trigram_index
|
||
.entry(trigram)
|
||
.or_default()
|
||
.push(idx as u32);
|
||
}
|
||
}
|
||
}
|
||
let token_prefix_index = build_place_prefix_index(&token_index);
|
||
|
||
let with_pop = population.iter().filter(|&&pop| pop > 0).count();
|
||
let with_city = city.iter().filter(|c| c.is_some()).count();
|
||
info!(
|
||
places = row_count,
|
||
types = place_type.values.len(),
|
||
with_population = with_pop,
|
||
with_city = with_city,
|
||
tokens = token_index.len(),
|
||
fuzzy_trigrams = fuzzy_trigram_index.len(),
|
||
"Place data loaded"
|
||
);
|
||
|
||
Ok(PlaceData {
|
||
name,
|
||
name_lower,
|
||
name_search,
|
||
place_type,
|
||
type_rank: type_rank_vec,
|
||
population,
|
||
lat,
|
||
lon,
|
||
city,
|
||
travel_destination,
|
||
token_index,
|
||
token_prefix_index,
|
||
fuzzy_trigram_index,
|
||
})
|
||
}
|
||
|
||
/// Candidate place rows for the query content tokens: intersect the posting lists of words
|
||
/// typed in full; if none matched an indexed token exactly, seed from the smallest
|
||
/// prefix-expanded list (so a partially-typed final word still works). Bounded by
|
||
/// `PLACE_CANDIDATE_LIMIT`.
|
||
pub fn place_candidate_rows(&self, tokens: &[&str]) -> Vec<u32> {
|
||
let mut exact: Vec<&[u32]> = tokens
|
||
.iter()
|
||
.filter_map(|token| self.token_index.get(*token).map(Vec::as_slice))
|
||
.collect();
|
||
|
||
let mut rows = if exact.is_empty() {
|
||
self.place_prefix_seed(tokens)
|
||
} else {
|
||
exact.sort_by_key(|posting| posting.len());
|
||
let mut acc = exact[0].to_vec();
|
||
for posting in &exact[1..] {
|
||
if acc.is_empty() {
|
||
break;
|
||
}
|
||
acc = intersect_sorted(&acc, posting);
|
||
}
|
||
acc
|
||
};
|
||
rows.truncate(PLACE_CANDIDATE_LIMIT);
|
||
rows
|
||
}
|
||
|
||
fn place_prefix_seed(&self, tokens: &[&str]) -> Vec<u32> {
|
||
let mut best: Option<Vec<u32>> = None;
|
||
for token in tokens {
|
||
if token.len() < PLACE_PREFIX_MIN_LEN {
|
||
continue;
|
||
}
|
||
let key = &token[..token.len().min(PLACE_PREFIX_MAX_LEN)];
|
||
let Some(indexed) = self.token_prefix_index.get(key) else {
|
||
continue;
|
||
};
|
||
let mut union: Vec<u32> = Vec::new();
|
||
for indexed_token in indexed {
|
||
if !indexed_token.starts_with(token) {
|
||
continue;
|
||
}
|
||
if let Some(rows) = self.token_index.get(indexed_token) {
|
||
union = if union.is_empty() {
|
||
rows.clone()
|
||
} else {
|
||
union_sorted(&union, rows)
|
||
};
|
||
}
|
||
}
|
||
if !union.is_empty()
|
||
&& best
|
||
.as_ref()
|
||
.is_none_or(|current| union.len() < current.len())
|
||
{
|
||
best = Some(union);
|
||
}
|
||
}
|
||
best.unwrap_or_default()
|
||
}
|
||
|
||
/// Fuzzy-eligible rows sharing enough trigrams with the query to be worth Jaccard scoring.
|
||
/// Bounded by the (small) fuzzy trigram index rather than scanning every place.
|
||
pub fn fuzzy_candidate_rows(&self, query_trigrams: &[u32]) -> Vec<u32> {
|
||
if query_trigrams.is_empty() {
|
||
return Vec::new();
|
||
}
|
||
let mut counts: FxHashMap<u32, u16> = FxHashMap::default();
|
||
for trigram in query_trigrams {
|
||
if let Some(rows) = self.fuzzy_trigram_index.get(trigram) {
|
||
for &row in rows {
|
||
*counts.entry(row).or_default() += 1;
|
||
}
|
||
}
|
||
}
|
||
let min_shared = (((query_trigrams.len() as f32) * 0.4).ceil() as u16).max(1);
|
||
counts
|
||
.into_iter()
|
||
.filter_map(|(row, shared)| (shared >= min_shared).then_some(row))
|
||
.collect()
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
impl PlaceData {
|
||
/// Build a minimal PlaceData from (name, place_type) pairs for index tests.
|
||
fn from_names<S: AsRef<str>>(rows: &[(S, S)]) -> Self {
|
||
let name: Vec<String> = rows.iter().map(|(nm, _)| nm.as_ref().to_string()).collect();
|
||
let place_type_raw: Vec<String> =
|
||
rows.iter().map(|(_, pt)| pt.as_ref().to_string()).collect();
|
||
let name_lower: Vec<String> = name.iter().map(|nm| nm.to_lowercase()).collect();
|
||
let name_search: Vec<String> = name
|
||
.iter()
|
||
.zip(&place_type_raw)
|
||
.map(|(nm, pt)| build_search_text(nm, pt))
|
||
.collect();
|
||
let mut token_index: FxHashMap<String, Vec<u32>> = FxHashMap::default();
|
||
let mut fuzzy_trigram_index: FxHashMap<u32, Vec<u32>> = FxHashMap::default();
|
||
for idx in 0..name.len() {
|
||
for token in place_index_tokens(&name_search[idx]) {
|
||
token_index.entry(token).or_default().push(idx as u32);
|
||
}
|
||
if is_fuzzy_eligible_type(&place_type_raw[idx]) {
|
||
for trigram in compute_trigrams(&name[idx]) {
|
||
fuzzy_trigram_index
|
||
.entry(trigram)
|
||
.or_default()
|
||
.push(idx as u32);
|
||
}
|
||
}
|
||
}
|
||
let token_prefix_index = build_place_prefix_index(&token_index);
|
||
let len = name.len();
|
||
PlaceData {
|
||
name,
|
||
name_lower,
|
||
name_search,
|
||
place_type: InternedColumn::build(&place_type_raw),
|
||
type_rank: place_type_raw.iter().map(|pt| type_rank(pt)).collect(),
|
||
population: vec![0; len],
|
||
lat: vec![0.0; len],
|
||
lon: vec![0.0; len],
|
||
city: vec![None; len],
|
||
travel_destination: vec![false; len],
|
||
token_index,
|
||
token_prefix_index,
|
||
fuzzy_trigram_index,
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
impl PlaceData {
|
||
/// Minimal empty instance for integration tests that need an `AppState`
|
||
/// but never touch place data.
|
||
pub(crate) fn empty_for_tests() -> Self {
|
||
PlaceData {
|
||
name: Vec::new(),
|
||
name_lower: Vec::new(),
|
||
name_search: Vec::new(),
|
||
place_type: InternedColumn::build(&[]),
|
||
type_rank: Vec::new(),
|
||
population: Vec::new(),
|
||
lat: Vec::new(),
|
||
lon: Vec::new(),
|
||
city: Vec::new(),
|
||
travel_destination: Vec::new(),
|
||
token_index: FxHashMap::default(),
|
||
token_prefix_index: FxHashMap::default(),
|
||
fuzzy_trigram_index: FxHashMap::default(),
|
||
}
|
||
}
|
||
}
|
||
|
||
#[cfg(test)]
|
||
mod tests {
|
||
use super::*;
|
||
|
||
#[test]
|
||
fn normalize_search_text_elides_every_apostrophe_variant() {
|
||
// Whatever glyph the keyboard / autocorrect / paste produced for the apostrophe, the
|
||
// normalized form must be identical — otherwise "King's Cross" tokenizes as `king s cross`
|
||
// and matching breaks. See is_apostrophe for the full set.
|
||
let expected = "kings cross";
|
||
for q in [
|
||
"King's Cross", // U+0027 straight
|
||
"King’s Cross", // U+2019 right single quote
|
||
"King‘s Cross", // U+2018 left single quote
|
||
"King´s Cross", // U+00B4 acute accent
|
||
"Kingʼs Cross", // U+02BC modifier letter apostrophe
|
||
"King′s Cross", // U+2032 prime
|
||
"King`s Cross", // U+0060 grave accent
|
||
"Kingʻs Cross", // U+02BB modifier letter turned comma
|
||
] {
|
||
assert_eq!(normalize_search_text(q), expected, "failed for {q:?}");
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn place_index_tokens_dedup_and_min_length() {
|
||
// "a" is too short; aliases split on " | ".
|
||
assert_eq!(
|
||
place_index_tokens("st albans | saint albans"),
|
||
vec!["albans".to_string(), "saint".to_string(), "st".to_string()]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn place_candidate_rows_intersect_and_prefix_seed() {
|
||
let pd = PlaceData::from_names(&[
|
||
("Camden", "suburb"),
|
||
("Camden Town", "suburb"),
|
||
("Camden Market", "attraction"),
|
||
("Manchester", "city"),
|
||
("Manchester Piccadilly", "station"),
|
||
]);
|
||
|
||
// Full word → posting list (Camden, Camden Town, Camden Market).
|
||
let camden = pd.place_candidate_rows(&["camden"]);
|
||
assert_eq!(camden, vec![0, 1, 2]);
|
||
|
||
// Two full words intersect to rows containing BOTH (Camden Town only).
|
||
let camden_town = pd.place_candidate_rows(&["camden", "town"]);
|
||
assert_eq!(camden_town, vec![1]);
|
||
|
||
// A partially-typed final word with no exact token seeds from the prefix index.
|
||
let piccad = pd.place_candidate_rows(&["piccad"]);
|
||
assert_eq!(piccad, vec![4]);
|
||
|
||
// No match → empty.
|
||
assert!(pd.place_candidate_rows(&["zzzz"]).is_empty());
|
||
}
|
||
|
||
// Run with: cargo test --release bench_place_search -- --ignored --nocapture
|
||
#[test]
|
||
#[ignore]
|
||
fn bench_place_search_at_one_million_rows() {
|
||
let roads = [
|
||
"High Street",
|
||
"Station Road",
|
||
"Church Lane",
|
||
"Victoria Road",
|
||
"Mill Lane",
|
||
"Park Avenue",
|
||
"Queens Road",
|
||
"Kings Road",
|
||
];
|
||
let mut rows: Vec<(String, String)> = Vec::with_capacity(1_000_000);
|
||
for i in 0..1_000_000usize {
|
||
// Vary the name so the index resembles ~1M distinct (street, area) rows.
|
||
rows.push((
|
||
format!("{} {}", roads[i % roads.len()], i % 4000),
|
||
"street".into(),
|
||
));
|
||
}
|
||
rows.push(("London".into(), "city".into()));
|
||
let pd = PlaceData::from_names(&rows);
|
||
|
||
let start = std::time::Instant::now();
|
||
let mut hits = 0usize;
|
||
for _ in 0..50 {
|
||
let candidates = pd.place_candidate_rows(&["high", "street"]);
|
||
for row in candidates {
|
||
let idx = row as usize;
|
||
if place_search_test_score(&pd, idx, "high street", &["high", "street"]).is_some() {
|
||
hits += 1;
|
||
}
|
||
}
|
||
}
|
||
let per_query = start.elapsed() / 50;
|
||
println!(
|
||
"indexed place search over {} rows: {:?}/query ({} hits)",
|
||
pd.name.len(),
|
||
per_query,
|
||
hits / 50
|
||
);
|
||
// The old full O(N) scan measured ~36ms here; candidate-based must be far under that.
|
||
assert!(per_query.as_millis() < 10, "per_query was {per_query:?}");
|
||
}
|
||
|
||
/// Mirrors the route's per-candidate match check for the bench.
|
||
fn place_search_test_score(
|
||
pd: &PlaceData,
|
||
idx: usize,
|
||
query_search: &str,
|
||
query_tokens: &[&str],
|
||
) -> Option<f32> {
|
||
let search_text = &pd.name_search[idx];
|
||
if query_tokens.iter().all(|qt| {
|
||
place_alias_tokens(search_text)
|
||
.any(|t| t == *qt || (qt.len() >= 2 && t.starts_with(qt)))
|
||
}) {
|
||
Some(640.0)
|
||
} else if pd.name_lower[idx] == query_search {
|
||
Some(1000.0)
|
||
} else {
|
||
None
|
||
}
|
||
}
|
||
|
||
#[test]
|
||
fn fuzzy_candidate_rows_finds_typos_only_for_eligible_rows() {
|
||
let pd = PlaceData::from_names(&[
|
||
("London", "city"),
|
||
("Baker Street", "street"), // not fuzzy-eligible
|
||
]);
|
||
let typo = compute_trigrams("Londn");
|
||
let candidates = pd.fuzzy_candidate_rows(&typo);
|
||
assert!(candidates.contains(&0)); // London (city) is reachable by fuzzy
|
||
assert!(!candidates.contains(&1)); // streets are excluded from the fuzzy index
|
||
}
|
||
|
||
fn test_city_rows() -> [(&'static str, f32, f32, u32); 5] {
|
||
[
|
||
("London", 51.507_446, -0.1277653, 8_908_083),
|
||
("Westminster", 51.497_322, -0.137149, 211_365),
|
||
("City of London", 51.515_617, -0.0919983, 10_847),
|
||
("Cambridge", 52.205_532, 0.1186637, 145_818),
|
||
("Oxford", 51.752_014, -1.2578499, 165_000),
|
||
]
|
||
}
|
||
|
||
fn all_test_city_candidates() -> Vec<CityCandidate<'static>> {
|
||
test_city_rows()
|
||
.into_iter()
|
||
.map(|(name, lat, lon, population)| {
|
||
CityCandidate::from_place(name, lat, lon, population)
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
fn test_city_candidates() -> Vec<CityCandidate<'static>> {
|
||
let cities = all_test_city_candidates();
|
||
|
||
cities
|
||
.iter()
|
||
.enumerate()
|
||
.filter_map(|(idx, city)| {
|
||
let is_subsumed = cities
|
||
.iter()
|
||
.enumerate()
|
||
.any(|(other_idx, other)| other_idx != idx && city.is_subsumed_by(other));
|
||
(!is_subsumed).then_some(*city)
|
||
})
|
||
.collect()
|
||
}
|
||
|
||
#[test]
|
||
fn type_rank_ordering() {
|
||
assert!(type_rank("city") < type_rank("town"));
|
||
assert!(type_rank("town") < type_rank("station"));
|
||
assert!(type_rank("station") < type_rank("unknown"));
|
||
}
|
||
|
||
#[test]
|
||
fn search_text_handles_common_address_variants() {
|
||
assert!(build_search_text("King's Cross tube station", "station")
|
||
.contains("kings cross underground"));
|
||
assert!(build_search_text("St Albans", "city").contains("saint albans"));
|
||
assert!(build_search_text("Shadwell DLR station", "station").contains("shadwell station"));
|
||
}
|
||
|
||
#[test]
|
||
fn search_text_expands_directional_and_size_abbreviations() {
|
||
assert!(build_search_text("Great Missenden", "village").contains("gt missenden"));
|
||
assert!(build_search_text("Mount Pleasant", "suburb").contains("mt pleasant"));
|
||
assert!(build_search_text("Little Venice", "suburb").contains("lt venice"));
|
||
}
|
||
|
||
#[test]
|
||
fn trigram_similarity_is_high_for_typos_and_low_for_unrelated() {
|
||
let london = compute_trigrams("London");
|
||
let typo = compute_trigrams("Londn");
|
||
let other = compute_trigrams("Manchester");
|
||
assert!(trigram_similarity(&london, &typo) >= 0.4);
|
||
assert!(trigram_similarity(&london, &other) < 0.2);
|
||
assert!((trigram_similarity(&london, &london) - 1.0).abs() < 1e-6);
|
||
}
|
||
|
||
#[test]
|
||
fn place_alias_tokens_split_across_aliases() {
|
||
let tokens: Vec<&str> = place_alias_tokens("kings cross | kings x").collect();
|
||
assert_eq!(tokens, vec!["kings", "cross", "kings", "x"]);
|
||
}
|
||
|
||
#[test]
|
||
fn travel_destination_types_match_legacy_places() {
|
||
assert!(is_travel_destination_type("city"));
|
||
assert!(is_travel_destination_type("station"));
|
||
assert!(!is_travel_destination_type("town"));
|
||
assert!(!is_travel_destination_type("suburb"));
|
||
}
|
||
|
||
#[test]
|
||
fn display_city_candidates_drop_city_nodes_subsumed_by_much_larger_nearby_city() {
|
||
let rows = test_city_rows();
|
||
let names: Vec<String> = rows
|
||
.iter()
|
||
.map(|(name, _, _, _)| name.to_string())
|
||
.collect();
|
||
let type_rank: Vec<u8> = vec![0; rows.len()];
|
||
let population: Vec<u32> = rows
|
||
.iter()
|
||
.map(|(_, _, _, population)| *population)
|
||
.collect();
|
||
let lat: Vec<f32> = rows.iter().map(|(_, lat, _, _)| *lat).collect();
|
||
let lon: Vec<f32> = rows.iter().map(|(_, _, lon, _)| *lon).collect();
|
||
|
||
let cities = display_city_candidates(&names, &type_rank, &population, &lat, &lon);
|
||
|
||
assert_eq!(
|
||
cities.iter().map(|city| city.name).collect::<Vec<_>>(),
|
||
["London", "Cambridge", "Oxford"]
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn nearest_display_city_labels_inner_greater_london_from_london_candidate() {
|
||
let cities = test_city_candidates();
|
||
|
||
assert_eq!(
|
||
nearest_display_city(51.371_304, -0.101957, &cities),
|
||
Some("London")
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn nearest_display_city_preserves_non_london_duplicates() {
|
||
let cities = test_city_candidates();
|
||
|
||
assert_eq!(
|
||
nearest_display_city(52.127_77, -0.0813098, &cities),
|
||
Some("Cambridge")
|
||
);
|
||
}
|
||
|
||
#[test]
|
||
fn nearest_display_city_does_not_extend_london_past_its_display_radius() {
|
||
let cities = test_city_candidates();
|
||
|
||
assert_eq!(nearest_display_city(51.5093, -0.5954, &cities), None);
|
||
}
|
||
|
||
#[test]
|
||
fn nearest_display_city_keeps_normal_non_london_city() {
|
||
let cities = test_city_candidates();
|
||
|
||
assert_eq!(
|
||
nearest_display_city(51.456659, -0.969651, &cities),
|
||
Some("Oxford")
|
||
);
|
||
}
|
||
}
|