diff --git a/docs/TECHNICAL_DESIGN.md b/docs/TECHNICAL_DESIGN.md index 5eced07..ec064d5 100644 --- a/docs/TECHNICAL_DESIGN.md +++ b/docs/TECHNICAL_DESIGN.md @@ -370,6 +370,10 @@ fetch away. spatial index (`OsmFeatureIndex`), queryable by S2 cell range without decoding every candidate (`index.rs`, [`src/tables/feature_index.rs`](../src/tables/feature_index.rs)). + The index also keeps the country of each feature’s centroid, one + byte per feature, so that later steps needn’t decode geometry to + find it; looking them up adds about 30 seconds on the planet. An + index in an older format is built again. **~4h48m** for the whole step on the #665 bare-VM run, back when the planet download itself went over BitTorrent and dominated that figure. On the newer, HTTPS-based cpx42 run, the download itself diff --git a/src/pipeline/conflate/deletion.rs b/src/pipeline/conflate/deletion.rs index 7dec05f..3a32644 100644 --- a/src/pipeline/conflate/deletion.rs +++ b/src/pipeline/conflate/deletion.rs @@ -508,13 +508,15 @@ fn swiss_uics(value: &str) -> impl Iterator { split_values(value).filter(|code| code.starts_with("85")) } -/// Why the unmatched OSM feature with `tags` is a candidate for -/// deletion, or `None` if nothing suggests so. `location` gives the -/// feature's (longitude, latitude); it is only called when needed, -/// and at most once, since it means decoding the geometry. +/// Why the unmatched OSM feature with `tags` in `country` is a +/// candidate for deletion, or `None` if nothing suggests so. +/// `location` gives the feature's (longitude, latitude); it is only +/// called when needed, and at most once, since it means decoding the +/// geometry. pub(super) fn deletion_reason<'a>( tags: impl Iterator + Clone, atp: &AtpCoverage, + country: Option, location: impl FnOnce() -> Option<(f64, f64)>, ) -> Option { let mut location = Some(location); @@ -551,7 +553,7 @@ pub(super) fn deletion_reason<'a>( // stops that the register has since given their own // country's codes, such as `uic_ref=8581987` in Nauders, // Austria. Neither Swiss signal is about such a stop. - if in_swiss_home(country(lng, lat), lng, lat) { + if in_swiss_home(country, lng, lat) { if has_swiss_id { return Some(Candidate::swiss(Reason::SwissTransitId)); } @@ -569,7 +571,6 @@ pub(super) fn deletion_reason<'a>( let brand = brand_tags(tags)?; let local_kind = brand.local_kind(); - let country = locate().and_then(|(lng, lat)| country(lng, lat)); let keys: Vec = brand .ids .iter() @@ -792,7 +793,8 @@ mod tests { } fn reason(atp: &AtpCoverage, tags: &[(&str, &str)], at: (f64, f64)) -> Option { - deletion_reason(tags.iter().copied(), atp, || Some(at)).map(|c| c.reason) + deletion_reason(tags.iter().copied(), atp, country(at.0, at.1), || Some(at)) + .map(|c| c.reason) } #[test] @@ -1081,7 +1083,10 @@ mod tests { // Without a location. let tags = [("highway", "bus_stop")]; - assert_eq!(deletion_reason(tags.iter().copied(), &atp, || None), None); + assert_eq!( + deletion_reason(tags.iter().copied(), &atp, None, || None), + None + ); // ATP lacks the Swiss register this week. assert_eq!(reason(&coverage(false), &tags, ZURICH_NORTH), None); @@ -1190,7 +1195,7 @@ mod tests { assert_eq!(reason(&atp, &tags, SINGEN), Some(Reason::Wikidata)); // Without a location, nothing says where the stop is. assert_eq!( - deletion_reason(old_code.iter().copied(), &atp, || None), + deletion_reason(old_code.iter().copied(), &atp, None, || None), None ); } @@ -1199,7 +1204,13 @@ mod tests { fn test_candidate_keys() { let atp = coverage(false); let tags = [("shop", "supermarket"), ("brand:wikidata", "Q9;Q1")]; - let c = deletion_reason(tags.iter().copied(), &atp, || Some(SINGEN)).expect("candidate"); + let c = deletion_reason( + tags.iter().copied(), + &atp, + country(SINGEN.0, SINGEN.1), + || Some(SINGEN), + ) + .expect("candidate"); assert_eq!(c.reason, Reason::Wikidata); // Only Q1 is covered by ATP, so only Q1 counts locally. assert_eq!(c.keys, vec![LocalKey::wikidata(1, None)]); diff --git a/src/pipeline/conflate/mod.rs b/src/pipeline/conflate/mod.rs index b0b4d0f..257a32f 100644 --- a/src/pipeline/conflate/mod.rs +++ b/src/pipeline/conflate/mod.rs @@ -542,7 +542,8 @@ fn send_deletion_candidates( strings: &osm.strings, }; let location = || centroid(&candidate); - let Some(found) = deletion_reason(candidate.tags(), coverage, location) else { + let country = osm.index.country(r); + let Some(found) = deletion_reason(candidate.tags(), coverage, country, location) else { return Ok(None); }; // A feature without a location can't be put on a map, nor diff --git a/src/pipeline/inputs.rs b/src/pipeline/inputs.rs index 373160c..e09f77c 100644 --- a/src/pipeline/inputs.rs +++ b/src/pipeline/inputs.rs @@ -237,14 +237,12 @@ mod tests { ] { symlink(fixture(name), workdir.path().join(name))?; } - for name in [ - "alltheplaces.parquet", - "osm-features.index", - "osm-features.index.inverted", - "osm-assemble.strings", - ] { + for name in ["alltheplaces.parquet", "osm-assemble.strings"] { std::fs::write(workdir.path().join(name), b"")?; } + crate::tables::OsmFeatureIndex::write_stub_for_test( + &workdir.path().join("osm-features.index"), + )?; let inputs = fetch_inputs( &test_client(), diff --git a/src/pipeline/osm/fetch.rs b/src/pipeline/osm/fetch.rs index 3e3a768..c59c572 100644 --- a/src/pipeline/osm/fetch.rs +++ b/src/pipeline/osm/fetch.rs @@ -322,13 +322,10 @@ mod tests { .path() .join(format!("{PLANET_PBF_FILENAME}.meta.json")), )?; - for name in [ - "osm-features.index", - "osm-features.index.inverted", - "osm-assemble.strings", - ] { - std::fs::write(workdir.path().join(name), b"")?; - } + std::fs::write(workdir.path().join("osm-assemble.strings"), b"")?; + crate::tables::OsmFeatureIndex::write_stub_for_test( + &workdir.path().join("osm-features.index"), + )?; let input = fetch_planet( UNREACHABLE_URL, diff --git a/src/pipeline/wikidata_ids.rs b/src/pipeline/wikidata_ids.rs index d7f1633..3d82de9 100644 --- a/src/pipeline/wikidata_ids.rs +++ b/src/pipeline/wikidata_ids.rs @@ -27,9 +27,7 @@ use crate::make_progress_bar; use crate::matchers::{OsmCandidate, parse_wikidata_ids}; use crate::places::PlaceReader; use crate::tables::{OsmFeatures, U64Set}; -use crate::utils::countries::country; use anyhow::Result; -use geo::Centroid; use indicatif::MultiProgress; use rayon::prelude::*; use std::path::{Path, PathBuf}; @@ -141,8 +139,8 @@ fn collect_from_atp(atp: &Path) -> Result<(Vec, BrandTotals)> { /// The Wikidata IDs in `osm`'s brand/operator/network tags, with /// duplicates, and their [BrandTotals]. Visits every feature in the /// index, in parallel; tag keys are compared as string pool indices, so -/// only features with matching tags get their tags looked up as strings -/// and their geometry decoded, for the country. +/// only features with matching tags get their tags looked up as +/// strings. fn collect_from_osm( osm: &OsmFeatures, progress: &MultiProgress, @@ -185,12 +183,7 @@ fn collect_from_osm( feature: &feature, strings: &osm.strings, }; - let country = candidate - .geometry() - .ok() - .and_then(|g| g.centroid()) - .and_then(|c| country(c.x(), c.y())); - totals.add(Source::Osm, candidate.tags(), country); + totals.add(Source::Osm, candidate.tags(), osm.index.country(r)); } Ok::<_, anyhow::Error>((ids, totals)) }, diff --git a/src/tables/feature_index.rs b/src/tables/feature_index.rs index c402453..e9dce7d 100644 --- a/src/tables/feature_index.rs +++ b/src/tables/feature_index.rs @@ -31,12 +31,14 @@ //! Feature storage (`out`): //! //! ```text -//! byte 0..8: magic "featstg0" +//! byte 0..8: magic "featstg1" //! byte 8..16: entry count, u64 little-endian //! byte 16..24: offset of the `id` array, u64 little-endian //! byte 24..32: offset of the `starts` array, u64 little-endian //! byte 32..40: offset of the `blobs` region, u64 little-endian -//! byte 40..64: reserved, zero-filled +//! byte 40..48: offset of the `countries` array, u64 little-endian +//! byte 48..56: offset of the `country_codes` table, u64 little-endian +//! byte 56..64: number of country codes, u64 little-endian //! //! id array: `entry count` OSM feature ids, u64 little-endian each //! starts array: `entry count + 1` byte offsets into `blobs`, u64 @@ -45,6 +47,14 @@ //! range is `starts[i]..starts[i+1]` without special-casing //! the final entry //! blobs region: concatenated `Feature.encode_to_vec()` bytes +//! countries: `entry count` bytes, the country of each feature's +//! centroid: 0 for none, else 1 + its position in +//! `country_codes` +//! country_codes: two-letter ISO 3166-1 codes, two ASCII bytes each, +//! as `crate::utils::countries::all` gave them when the +//! index was built. Stored, rather than taken from the +//! running binary, so that an index stays readable after +//! the boundary data changes. //! ``` //! //! Inverted index (`out` with an added `.inverted` extension): @@ -68,6 +78,7 @@ use crate::matchers::MatchMask; use crate::pipeline::EXTERNAL_SORT_CHUNK_BYTES; use crate::tables::{Feature, FeatureToIndex, StringPool}; +use crate::utils::countries::{self, Country}; use anyhow::{Context, Result}; use deepsize::DeepSizeOf; use ext_sort::{ @@ -91,7 +102,7 @@ use std::{ /// Size of each file's header, in bytes: see the "File format" section above. const HEADER_SIZE: usize = 8 * 8; -const FEATURES_FILE_SIGNATURE: &[u8; 8] = b"featstg0"; +const FEATURES_FILE_SIGNATURE: &[u8; 8] = b"featstg1"; const INVERTED_INDEX_FILE_SIGNATURE: &[u8; 8] = b"featinv0"; /// An opaque reference to a feature's position in an [OsmFeatureIndex]'s @@ -146,6 +157,8 @@ pub struct OsmFeatureIndex<'a> { id: &'a [u64], starts: &'a [u64], blobs: &'a [u8], + countries: &'a [u8], + country_codes: Vec>, inverted_index_file: File, _inverted_index_mmap: Mmap, @@ -187,11 +200,27 @@ impl<'a> OsmFeatureIndex<'a> { p } + /// Writes files at `path` that [OsmFeatureIndex::exists] accepts + /// but that can't be opened, for tests that only check whether an + /// index is there. + #[cfg(test)] + pub(crate) fn write_stub_for_test(path: &Path) -> Result<()> { + std::fs::write(path, FEATURES_FILE_SIGNATURE)?; + std::fs::write(Self::inverted_index_path(path), b"")?; + Ok(()) + } + /// Whether both files an `OsmFeatureIndex` at `path` needs already - /// exist -- for callers that memoize an expensive build step behind - /// an existence check rather than reopening unconditionally. + /// exist, in the current format -- for callers that memoize an + /// expensive build step behind an existence check rather than + /// reopening unconditionally. An index in an older format counts as + /// missing, so that it gets built again. pub fn exists(path: &Path) -> bool { - path.exists() && Self::inverted_index_path(path).exists() + let mut signature = [0_u8; 8]; + let current = File::open(path) + .and_then(|mut f| f.read_exact(&mut signature)) + .is_ok_and(|_| &signature == FEATURES_FILE_SIGNATURE); + current && Self::inverted_index_path(path).exists() } /// Builds an `OsmFeatureIndex` from `fti`, which may be in any order. @@ -254,7 +283,8 @@ impl<'a> OsmFeatureIndex<'a> { let fti = FeatureToIndex::decode(item.fti_bytes.as_slice()) .context("failed to decode a FeatureToIndex spooled during pass 1")?; let feature = fti.feature.expect("feature"); - feature_storage_writer.write(feature.id, &feature.encode_to_vec())?; + let country = centroid_country(fti.centroid_s2_cell_id); + feature_storage_writer.write(feature.id, country, &feature.encode_to_vec())?; for cell_id in &fti.coverage_s2_cell_id { pass2_writer.write_all(&cell_id.to_le_bytes())?; pass2_writer.write_all(&fti.match_mask.to_le_bytes())?; @@ -357,18 +387,28 @@ impl<'a> OsmFeatureIndex<'a> { path, "starts", )?; - let blobs = { - let blobs_offset = usize::try_from(u64::from_le(header[4]))?; - if blobs_offset <= features_mmap.len() { - // SAFETY: Verified length; no alignment constraints on &[u8]. - unsafe { - let ptr = features_mmap.as_ptr().add(blobs_offset); - std::slice::from_raw_parts(ptr, features_mmap.len() - blobs_offset) - } - } else { - anyhow::bail!("misplaced blobs in OsmFeatureIndex: {}", path.display()); - } + let blobs_offset = usize::try_from(u64::from_le(header[4]))?; + let countries_offset = usize::try_from(u64::from_le(header[5]))?; + let codes_offset = usize::try_from(u64::from_le(header[6]))?; + let num_codes = usize::try_from(u64::from_le(header[7]))?; + if blobs_offset > countries_offset + || countries_offset + entries_count != codes_offset + || codes_offset + 2 * num_codes != features_mmap.len() + { + anyhow::bail!("misplaced arrays in OsmFeatureIndex: {}", path.display()); + } + // SAFETY: Verified length; no alignment constraints on &[u8]. + let bytes = |offset: usize, len: usize| -> &'a [u8] { + unsafe { std::slice::from_raw_parts(features_mmap.as_ptr().add(offset), len) } }; + let blobs = bytes(blobs_offset, countries_offset - blobs_offset); + let countries = bytes(countries_offset, entries_count); + let country_codes = bytes(codes_offset, 2 * num_codes) + .as_chunks::<2>() + .0 + .iter() + .map(|code| Country::from_code(*code)) + .collect(); let inverted_index_path = Self::inverted_index_path(path); let inverted_index_file = File::open(&inverted_index_path).with_context(|| { @@ -418,6 +458,8 @@ impl<'a> OsmFeatureIndex<'a> { id, starts, blobs, + countries, + country_codes, inverted_index_file, _inverted_index_mmap: inverted_index_mmap, @@ -499,6 +541,18 @@ impl<'a> OsmFeatureIndex<'a> { (0..count).into_par_iter().map(LocalFeatureRef) } + /// The country of the feature `r` refers to, by its centroid, or + /// `None` if it lies in none, such as in the open sea. No decode; + /// `None` too for a country that the running binary's boundary data + /// doesn't know, which can only happen with an index built by + /// another version. + pub fn country(&self, r: LocalFeatureRef) -> Option { + match self.countries[r.0 as usize] { + 0 => None, + i => self.country_codes.get(i as usize - 1).copied().flatten(), + } + } + /// Fully decodes the feature `r` refers to. The only operation on /// this index that costs a protobuf decode. pub fn get_feature(&self, r: LocalFeatureRef) -> Result { @@ -602,6 +656,28 @@ impl Iterator for QueryIter<'_> { /// rather than the `&Mmap` reference -- otherwise the borrow checker /// would tie the slice to this function's local borrow of `mmap`, not to /// the `Mmap` value itself once it's moved into the returned struct. +/// The country at the centroid of a feature, given as an S2 cell id; +/// 0 means the feature has no centroid. +fn centroid_country(cell_id: u64) -> Option { + if cell_id == 0 { + return None; + } + let ll = s2::latlng::LatLng::from(CellID(cell_id)); + countries::country(ll.lng.deg(), ll.lat.deg()) +} + +/// How `country` is stored: 0 for none, else 1 + its position in +/// [countries::all]. +fn country_byte(country: Option) -> u8 { + let Some(country) = country else { + return 0; + }; + match countries::all().binary_search(&country) { + std::result::Result::Ok(i) => u8::try_from(i + 1).expect("fewer than 255 countries"), + Err(_) => 0, + } +} + fn read_u64_array<'a>( mmap: &Mmap, offset: usize, @@ -646,6 +722,9 @@ struct FeatureStorageWriter { blobs_path: PathBuf, blobs_writer: BufWriter, + countries_path: PathBuf, + countries_writer: BufWriter, + entries_count: u64, } @@ -664,6 +743,10 @@ impl FeatureStorageWriter { blobs_path.add_extension("blobs.tmp"); let blobs_writer = BufWriter::with_capacity(32 * 1024, File::create(&blobs_path)?); + let mut countries_path = PathBuf::from(path); + countries_path.add_extension("countries.tmp"); + let countries_writer = BufWriter::with_capacity(32 * 1024, File::create(&countries_path)?); + Ok(FeatureStorageWriter { path: PathBuf::from(path), tmp_path, @@ -672,16 +755,19 @@ impl FeatureStorageWriter { starts_writer, blobs_path, blobs_writer, + countries_path, + countries_writer, entries_count: 0, }) } /// Appends `feature_bytes` (a `Feature.encode_to_vec()`), recording - /// `id` alongside it at the same position. - fn write(&mut self, id: u64, feature_bytes: &[u8]) -> Result<()> { + /// `id` and `country` alongside it at the same position. + fn write(&mut self, id: u64, country: Option, feature_bytes: &[u8]) -> Result<()> { let start: u64 = self.blobs_writer.stream_position()?; self.starts_writer.write_all(&start.to_le_bytes())?; self.blobs_writer.write_all(feature_bytes)?; + self.countries_writer.write_all(&[country_byte(country)])?; self.writer.write_all(&id.to_le_bytes())?; self.entries_count += 1; @@ -716,12 +802,28 @@ impl FeatureStorageWriter { std::io::copy(&mut File::open(&self.blobs_path)?, &mut self.writer)?; remove_file(&self.blobs_path)?; + let countries_offset: u64 = self.writer.stream_position()?; + self.countries_writer.flush()?; // flush() returns errors + drop(self.countries_writer); // drop() does not return errors + std::io::copy(&mut File::open(&self.countries_path)?, &mut self.writer)?; + remove_file(&self.countries_path)?; + + let codes_offset: u64 = self.writer.stream_position()?; + assert_eq!(codes_offset, countries_offset + self.entries_count); + let codes = countries::all(); + for country in codes { + self.writer.write_all(&country.code_bytes())?; + } + self.writer.seek(SeekFrom::Start(0))?; self.writer.write_all(FEATURES_FILE_SIGNATURE)?; self.writer.write_all(&self.entries_count.to_le_bytes())?; self.writer.write_all(&ids_offset.to_le_bytes())?; self.writer.write_all(&starts_offset.to_le_bytes())?; self.writer.write_all(&blobs_offset.to_le_bytes())?; + self.writer.write_all(&countries_offset.to_le_bytes())?; + self.writer.write_all(&codes_offset.to_le_bytes())?; + self.writer.write_all(&(codes.len() as u64).to_le_bytes())?; assert!(self.writer.stream_position()? <= HEADER_SIZE as u64); self.writer.seek(SeekFrom::End(0))?; @@ -1023,6 +1125,52 @@ mod tests { assert!(OsmFeatureIndex::open(&bad).is_err()); } + #[test] + fn country() { + let cell = |lng: f64, lat: f64| CellID::from(s2::latlng::LatLng::from_degrees(lat, lng)).0; + let (_dir, index) = build(vec![ + fti(11, cell(8.5402, 47.3782), &[1], MatchMask::SHOP.0), // Zürich + fti(21, cell(-30.0, 40.0), &[2], MatchMask::SHOP.0), // Atlantic + fti(31, 0, &[3], MatchMask::SHOP.0), // no centroid + fti(41, cell(121.5654, 25.0330), &[4], MatchMask::SHOP.0), // Taipei + ]); + let mut found: Vec<(u64, Option)> = (0..index.len() as u32) + .map(LocalFeatureRef) + .map(|r| { + ( + index.feature_id(r), + index.country(r).map(|c| c.code().to_string()), + ) + }) + .collect(); + found.sort(); + assert_eq!( + found, + vec![ + (11, Some("CH".into())), + (21, None), + (31, None), + (41, Some("TW".into())) + ] + ); + } + + #[test] + fn exists_rejects_older_format() { + let dir = TempDir::new().expect("tempdir"); + let out = dir.path().join("osm-features.index"); + OsmFeatureIndex::create( + vec![fti(1, 1, &[1], MatchMask::SHOP.0)].into_iter(), + dir.path(), + &out, + ) + .expect("create"); + let mut bytes = std::fs::read(&out).expect("read"); + bytes[..8].copy_from_slice(b"featstg0"); + std::fs::write(&out, bytes).expect("write"); + assert!(!OsmFeatureIndex::exists(&out)); + } + #[test] fn exists_requires_both_files() { let dir = TempDir::new().expect("tempdir"); diff --git a/src/utils/countries.rs b/src/utils/countries.rs index dac1eb1..86dae4b 100644 --- a/src/utils/countries.rs +++ b/src/utils/countries.rs @@ -5,7 +5,7 @@ //! crate, whose boundaries are derived from OpenStreetMap //! (© OpenStreetMap contributors, ODbL) via JOSM's `boundaries.osm`. -use country_boundaries::{BOUNDARIES_ODBL_360X180, CountryBoundaries, LatLon}; +use country_boundaries::{BOUNDARIES_ODBL_360X180, BoundingBox, CountryBoundaries, LatLon}; use std::sync::LazyLock; /// The version of the `country-boundaries` crate, which ships the @@ -25,7 +25,7 @@ static BOUNDARIES: LazyLock = LazyLock::new(|| { }); /// A country, by its two-letter ISO 3166-1 code. -#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash)] +#[derive(Clone, Copy, Debug, PartialEq, Eq, Hash, PartialOrd, Ord)] pub struct Country([u8; 2]); impl Country { @@ -40,6 +40,34 @@ impl Country { pub fn code(&self) -> &str { std::str::from_utf8(&self.0).expect("ASCII") } + + /// The two-letter code as bytes, such as `*b"CH"`. + pub fn code_bytes(&self) -> [u8; 2] { + self.0 + } + + /// The country with the two-letter code `code`, if it is one that + /// [country] can return. + pub fn from_code(code: [u8; 2]) -> Option { + all().binary_search(&Country(code)).ok().map(|i| all()[i]) + } +} + +static ALL: LazyLock> = LazyLock::new(|| { + let world = BoundingBox::new(-90.0, -180.0, 90.0, 179.999_999).expect("valid bounds"); + let mut all: Vec = BOUNDARIES + .intersecting_ids(world) + .into_iter() + .filter_map(|id| <[u8; 2]>::try_from(id.as_bytes()).ok()) + .map(Country) + .collect(); + all.sort(); + all +}); + +/// Every country that [country] can return, sorted by code. +pub fn all() -> &'static [Country] { + &ALL } /// The country that (`lng`, `lat`), in degrees, lies in, or `None` if @@ -94,4 +122,27 @@ mod tests { assert_eq!(code(-30.0, 40.0), None); // Atlantic assert_eq!(code(0.0, 91.0), None); // invalid } + + #[test] + fn test_all() { + let all = all(); + // Fits into a byte, with 0 for none; see OsmFeatureIndex. As of + // country-boundaries 1.2.0, there are exactly 255. If a newer + // version has more, OsmFeatureIndex needs two bytes per feature. + assert!(all.len() <= 255, "{}", all.len()); + assert!(all.len() > 190, "{}", all.len()); + assert!(all.windows(2).all(|w| w[0] < w[1])); + assert!(all.contains(&Country::SWITZERLAND)); + assert!(all.contains(&Country(*b"HK"))); + assert_eq!(Country::from_code(*b"CH"), Some(Country::SWITZERLAND)); + assert_eq!(Country::from_code(*b"ZZ"), None); + // Every country found on a grid over the world is in `all`. + for lat in (-89..90).step_by(1) { + for lng in (-179..180).step_by(1) { + if let Some(c) = country(lng as f64 + 0.5, lat as f64 + 0.5) { + assert!(all.binary_search(&c).is_ok(), "{c:?}"); + } + } + } + } }