feat(wikidata): Count OSM features and names per brand

docs/TECHNICAL_DESIGN.md

@@ -211,6 +211,8 @@ graph TD
     ATP_PARQUET --> COLLECT_WIKI(collect_wikidata_ids)
     OSM_INDEX --> COLLECT_WIKI
     COLLECT_WIKI --> WIKIDATA[wikidata-ids]
+    COLLECT_WIKI --> OSM_BRANDS[osm-brands.tsv]
+    COLLECT_WIKI --> BRAND_NAMES[brand-names.tsv]
     WIKIDATA --> QRANK_QUANTILES(compute_qrank_quantiles)
     QRANK_CSV --> QRANK_QUANTILES
     QRANK_QUANTILES --> QRANK_Q[qrank-quantiles]
@@ -395,6 +397,14 @@ fetch away.
   ([#713](https://github.com/brawer/osmdiffs/issues/713)). Plain
   `wikidata` is left out, since on OSM it mostly points at cities,
   rivers and the like. Visits every feature in `osm-features.index`.
+  On the side, it writes `osm-brands.tsv`, the number of OSM features
+  per Wikidata ID, country and main category (such as
+  `shop=supermarket`), and `brand-names.tsv`, one `brand`,
+  `operator` or `network` name per ID and key: the most common one in
+  Latin script, or else the most common one, from OSM and ATP
+  ([`brand_totals.rs`](../src/pipeline/brand_totals.rs)). The spider
+  stats compare ATP with the first, and show the second so that
+  readers needn't look up Wikidata IDs.
 - **`compute_qrank_quantiles`**
   ([`src/pipeline/qrank/mod.rs`](../src/pipeline/qrank/mod.rs)) — gives
   each ID in `wikidata-ids` that has pageviews its QRank quantile among

src/pipeline/brand_totals.rs

@@ -0,0 +1,307 @@
+//! For each brand, operator and network that features refer to by
+//! Wikidata ID: how many OpenStreetMap features it has per country and
+//! main category, and a name for it, from OSM or AllThePlaces. The
+//! spider stats compare these totals with AllThePlaces, and show the
+//! names so that their readers needn't look up Wikidata IDs.
+//!
+//! [super::wikidata_ids::collect_wikidata_ids] fills them in, since it
+//! visits the same features anyway.
+
+use crate::matchers::{contact::main_category, parse_wikidata_ids};
+use crate::utils::countries::Country;
+use anyhow::{Context, Result};
+use std::collections::HashMap;
+use std::io::{BufWriter, Write};
+use std::path::{Path, PathBuf};
+
+use super::wikidata_ids::WIKIDATA_KEYS;
+
+/// Filename, within `workdir`, of the OSM feature totals.
+const FEATURES_FILENAME: &str = "osm-brands.tsv";
+
+/// Filename, within `workdir`, of the names.
+const NAMES_FILENAME: &str = "brand-names.tsv";
+
+pub(super) fn features_path(workdir: &Path) -> PathBuf {
+    workdir.join(FEATURES_FILENAME)
+}
+
+pub(super) fn names_path(workdir: &Path) -> PathBuf {
+    workdir.join(NAMES_FILENAME)
+}
+
+#[derive(Clone, Copy, Debug, PartialEq, Eq)]
+pub(super) enum Source {
+    Osm,
+    Atp,
+}
+
+/// Wikidata ID, and the key of the name, such as `brand`.
+type NameKey = (u64, &'static str);
+
+#[derive(Debug, Default)]
+pub(super) struct BrandTotals {
+    /// OSM features per Wikidata ID, country and main category (see
+    /// [main_category]). A feature that refers to the same ID twice,
+    /// such as by `brand:wikidata` and `operator:wikidata`, counts once.
+    features: HashMap<(u64, Option<Country>, Option<String>), u64>,
+
+    /// How often OSM and ATP features use each name, in that order.
+    names: HashMap<NameKey, HashMap<String, [u64; 2]>>,
+}
+
+impl BrandTotals {
+    /// Adds a feature with `tags` that lies in `country`, which only
+    /// matters for OSM features. Names count only where the Wikidata
+    /// tag names a single ID: with several, there's no telling which
+    /// name goes with which.
+    pub fn add<'a>(
+        &mut self,
+        source: Source,
+        tags: impl Iterator<Item = (&'a str, &'a str)> + Clone,
+        country: Option<Country>,
+    ) {
+        let mut ids: Vec<u64> = Vec::new();
+        for key in WIKIDATA_KEYS {
+            let Some((_, value)) = tags.clone().find(|(k, _)| *k == key) else {
+                continue;
+            };
+            let found: Vec<u64> = parse_wikidata_ids(value).collect();
+            let name_key = key.strip_suffix(":wikidata").expect("Wikidata key");
+            if let [id] = found[..]
+                && let Some((_, name)) = tags.clone().find(|(k, _)| *k == name_key)
+            {
+                let name = clean_name(name);
+                if !name.is_empty() {
+                    let counts = self
+                        .names
+                        .entry((id, name_key))
+                        .or_default()
+                        .entry(name)
+                        .or_default();
+                    counts[source as usize] += 1;
+                }
+            }
+            ids.extend(found);
+        }
+        if source != Source::Osm || ids.is_empty() {
+            return;
+        }
+        ids.sort_unstable();
+        ids.dedup();
+        let category = main_category(tags);
+        for id in ids {
+            *self
+                .features
+                .entry((id, country, category.clone()))
+                .or_default() += 1;
+        }
+    }
+
+    pub fn merge(mut self, other: BrandTotals) -> BrandTotals {
+        if self.features.len() + self.names.len() < other.features.len() + other.names.len() {
+            return other.merge(self);
+        }
+        for (key, n) in other.features {
+            *self.features.entry(key).or_default() += n;
+        }
+        for (key, names) in other.names {
+            let mine = self.names.entry(key).or_default();
+            for (name, [osm, atp]) in names {
+                let counts = mine.entry(name).or_default();
+                counts[0] += osm;
+                counts[1] += atp;
+            }
+        }
+        self
+    }
+
+    /// Writes the feature totals to [features_path] and a name per
+    /// Wikidata ID and key to [names_path], both sorted, and returns
+    /// how many lines each has. The name is the most common one in
+    /// Latin script, or if there is none, the most common one: enough
+    /// to tell brands apart for readers of the spider stats.
+    pub fn write(&self, workdir: &Path) -> Result<(usize, usize)> {
+        let mut features: Vec<_> = self.features.iter().collect();
+        features.sort_by_cached_key(|((id, country, category), _)| {
+            (*id, code(*country), category.clone())
+        });
+        let num_features = features.len();
+        write_tsv(
+            &features_path(workdir),
+            "wikidata_id\tcountry\tcategory\tfeatures",
+            features.into_iter().map(|((id, country, category), n)| {
+                format!(
+                    "Q{id}\t{}\t{}\t{n}",
+                    code(*country),
+                    category.as_deref().unwrap_or("")
+                )
+            }),
+        )?;
+
+        let mut keys: Vec<&NameKey> = self.names.keys().collect();
+        keys.sort();
+        let lines: Vec<String> = keys
+            .into_iter()
+            .filter_map(|k| {
+                let (name, [osm, atp]) = display_name(&self.names[k])?;
+                let (id, key) = k;
+                Some(format!("Q{id}\t{key}\t{name}\t{osm}\t{atp}"))
+            })
+            .collect();
+        let num_names = lines.len();
+        write_tsv(
+            &names_path(workdir),
+            "wikidata_id\tkey\tname\tosm\tatp",
+            lines.into_iter(),
+        )?;
+        Ok((num_features, num_names))
+    }
+}
+
+/// The name to show, out of `names` with their counts: the most
+/// common in Latin script, or else the most common, with ties broken
+/// by the name itself.
+fn display_name(names: &HashMap<String, [u64; 2]>) -> Option<(&String, [u64; 2])> {
+    let rank = |(name, counts): &(&String, &[u64; 2])| (is_latin(name), counts[0] + counts[1]);
+    names
+        .iter()
+        .max_by(|a, b| rank(a).cmp(&rank(b)).then(b.0.cmp(a.0)))
+        .map(|(name, counts)| (name, *counts))
+}
+
+/// Whether every letter of `name` is in Latin script, such as
+/// `Migros` or `Żabka`, unlike `ミニストップ`.
+fn is_latin(name: &str) -> bool {
+    name.chars().filter(|c| c.is_alphabetic()).all(
+        |c| matches!(c, 'A'..='Z' | 'a'..='z' | '\u{00C0}'..='\u{024F}' | '\u{1E00}'..='\u{1EFF}'),
+    )
+}
+
+fn code(country: Option<Country>) -> String {
+    country.map(|c| c.code().to_string()).unwrap_or_default()
+}
+
+/// `name` without surrounding spaces, and with tabs and line breaks
+/// replaced, so it fits into one TSV column.
+fn clean_name(name: &str) -> String {
+    name.trim()
+        .chars()
+        .map(|c| if c.is_control() { ' ' } else { c })
+        .collect()
+}
+
+fn write_tsv(out: &Path, header: &str, lines: impl Iterator<Item = String>) -> Result<()> {
+    let mut tmp = out.to_path_buf();
+    tmp.add_extension("tmp");
+    let mut w = BufWriter::new(
+        std::fs::File::create(&tmp).with_context(|| format!("could not create {tmp:?}"))?,
+    );
+    writeln!(w, "{header}")?;
+    for line in lines {
+        writeln!(w, "{line}")?;
+    }
+    w.into_inner()?.sync_all()?;
+    std::fs::rename(&tmp, out)?;
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    const CH: Option<Country> = Some(Country::SWITZERLAND);
+    const DE: Option<Country> = Some(Country::GERMANY);
+
+    #[test]
+    fn test_brand_totals() -> Result<()> {
+        let mut a = BrandTotals::default();
+        let mut b = BrandTotals::default();
+        let migros = [
+            ("shop", "supermarket"),
+            ("brand", "Migros"),
+            ("brand:wikidata", "Q680727"),
+            ("operator", "Genossenschaft Migros Zürich"),
+            ("operator:wikidata", "Q1"),
+        ];
+        a.add(Source::Osm, migros.iter().copied(), CH);
+        b.add(Source::Osm, migros.iter().copied(), CH);
+        let typo = [
+            ("shop", "supermarket"),
+            ("brand", "Migors"),
+            ("brand:wikidata", "Q680727"),
+        ];
+        b.add(Source::Osm, typo.iter().copied(), CH);
+        // ATP names count, but ATP features don't.
+        let atp = [
+            ("shop", "supermarket"),
+            ("brand", "Migros\t"),
+            ("brand:wikidata", "Q680727"),
+        ];
+        b.add(Source::Atp, atp.iter().copied(), CH);
+        // Two IDs: no name.
+        let two = [
+            ("amenity", "fuel"),
+            ("brand", "Migrol;Migros"),
+            ("brand:wikidata", "Q1;Q680727"),
+        ];
+        b.add(Source::Osm, two.iter().copied(), DE);
+        // The same ID twice counts once.
+        let same = [
+            ("brand", "Migros"),
+            ("brand:wikidata", "Q680727"),
+            ("operator:wikidata", "Q680727"),
+        ];
+        b.add(Source::Osm, same.iter().copied(), None);
+        let totals = a.merge(b);
+
+        let dir = tempfile::tempdir()?;
+        assert_eq!(totals.write(dir.path())?, (5, 2));
+        assert_eq!(
+            std::fs::read_to_string(features_path(dir.path()))?,
+            "wikidata_id\tcountry\tcategory\tfeatures\n\
+             Q1\tCH\tshop=supermarket\t2\n\
+             Q1\tDE\tamenity=fuel\t1\n\
+             Q680727\t\t\t1\n\
+             Q680727\tCH\tshop=supermarket\t3\n\
+             Q680727\tDE\tamenity=fuel\t1\n"
+        );
+        assert_eq!(
+            std::fs::read_to_string(names_path(dir.path()))?,
+            "wikidata_id\tkey\tname\tosm\tatp\n\
+             Q1\toperator\tGenossenschaft Migros Zürich\t2\t0\n\
+             Q680727\tbrand\tMigros\t3\t1\n"
+        );
+        Ok(())
+    }
+
+    #[test]
+    fn test_display_name() {
+        let shown = |list: &[(&str, u64)]| {
+            let names: HashMap<String, [u64; 2]> =
+                list.iter().map(|(n, c)| (n.to_string(), [*c, 0])).collect();
+            display_name(&names).map(|(n, _)| n.clone())
+        };
+        assert_eq!(shown(&[("A", 1), ("B", 4), ("C", 4)]), Some("B".into()));
+        assert_eq!(
+            shown(&[("ミニストップ", 9), ("Ministop", 2)]),
+            Some("Ministop".into())
+        );
+        assert_eq!(
+            shown(&[("ミニストップ", 9), ("ローソン", 2)]),
+            Some("ミニストップ".into())
+        );
+        assert_eq!(shown(&[]), None);
+    }
+
+    #[test]
+    fn test_is_latin() {
+        assert!(is_latin("Migros"));
+        assert!(is_latin("Coop Pronto 24/7"));
+        assert!(is_latin("Żabka"));
+        assert!(is_latin("Phở Hòa"));
+        assert!(!is_latin("ミニストップ"));
+        assert!(!is_latin("Пятёрочка"));
+        assert!(!is_latin("7-Eleven セブン"));
+    }
+}

src/pipeline/mod.rs

@@ -31,6 +31,7 @@ use time::UtcDateTime;
 pub(crate) const EXTERNAL_SORT_CHUNK_BYTES: usize = 512 * 1024 * 1024;
 
 mod atp;
+mod brand_totals;
 mod conflate;
 mod conflated_tiles;
 mod datapackage;