fix(chore): cleanup and adds progress bar in the hf download

Author: kshitijrajsharma

README.md

@@ -181,6 +181,9 @@ Any flag works as a YAML key. See [docs/Manual.md](./docs/Manual.md) for the ful
 Every run writes `stats.duckdb` (or `<--name>.duckdb`) plus the formats you ask for via
 `-f parquet|csv|json|markdown|psql`. Parquet is the default. Open it with duckdb, polars, pandas, anything.
 
+Rerunning the same query with a different `-f` re-exports from the existing `<name>.duckdb` instead of
+refetching, so adding a format is instant. Pass `--overwrite` to force a fresh recompute.
+
 ## Configuration
 
 Every meaningful flag has a matching `OSMSG_*` env var so the CLI, a `.env` file, and a
@@ -196,6 +199,7 @@ docker-compose `environment:` block all reach the same setting. CLI flag wins ov
 | `--cache-dir` | `OSMSG_CACHE_DIR` | platform cache | Where downloaded OSM files are kept across runs. |
 | `--output-dir` | `OSMSG_OUTPUT_DIR` | `.` | Where `<name>.duckdb` and exports are written. |
 | `--format` / `-f` | `OSMSG_FORMAT` | `parquet` | Repeat for multiple. Comma-separated when set via env. |
+| `--overwrite` | (none) | off | Recompute even if `<name>.duckdb` already holds this exact query. |
 | `--psql-dsn` | `OSMSG_PSQL_DSN` | unset | libpq DSN for `-f psql`. |
 | `--psql-bulk` | `OSMSG_PSQL_BULK` | off | Faster first full load to Postgres. |
 | `--history` / `--no-history` | `OSMSG_HISTORY` | on | Read covered months from the published dataset. |

docs/Manual.md

@@ -77,6 +77,10 @@ osmsg --last day -f psql --psql-dsn "host=localhost dbname=osm user=osm"
 > breakdown is just a query over the four base tables, so consumers derive it on demand instead of
 > duplicating data.
 
+`<name>.duckdb` is stamped with the query that built it. Rerunning the same query with a different `-f`
+re-exports from that store instead of refetching, so adding a format is instant. Changing any query
+parameter (window, hashtags, tags, boundary) recomputes; `--overwrite` forces a fresh recompute.
+
 ## Config file
 
 Long invocations are easier to maintain in YAML. Keys mirror the CLI flag names.

osmsg/cli.py

@@ -271,6 +271,14 @@ def main(
         str | None,
         typer.Option("--changeset-file", help="Changeset dump (.osm.bz2) paired with --osh-file."),
     ] = None,
+    overwrite: Annotated[
+        bool,
+        typer.Option(
+            "--overwrite",
+            help="Recompute even if <name>.duckdb already holds this exact query; otherwise a rerun "
+            "that only changes the output format re-exports from the existing store.",
+        ),
+    ] = False,
 ) -> None:
     """Run osmsg. With no subcommand this generates stats (or loads history with --insert)."""
     if ctx.invoked_subcommand is not None:
@@ -338,6 +346,7 @@ def main(
         insert=insert,
         osh_file=osh_file,
         changeset_file=changeset_file,
+        overwrite=overwrite,
     )
 
     if last is not None:

osmsg/export/psql.py

@@ -5,10 +5,6 @@
 from ..exceptions import OsmsgError
 from ..pg_schema import PG_SCHEMA
 
-# Secondary indexes and foreign keys that make a row-by-row insert slow. For a one-time bulk load
-# they are dropped before the COPY and rebuilt once after (one index build + one FK validation,
-# instead of maintaining them per row). Primary keys stay, because the ON CONFLICT upserts need them.
-# Indexes are (name, create-sql); foreign keys are (table, name, add-clause).
 _BULK_INDEXES = [
     ("idx_changesets_created_at", "CREATE INDEX idx_changesets_created_at ON changesets (created_at)"),
     ("idx_changesets_geom", "CREATE INDEX idx_changesets_geom ON changesets USING GIST (geom)"),
@@ -25,8 +21,6 @@
 ]
 
 
-# Bulk loads push the big tables in this many changeset_id ranges, each its own statement and so its
-# own commit, so a failure costs one range instead of rolling back the whole multi-GB load.
 _BULK_COMMIT_CHUNKS = 32
 
 
@@ -102,9 +96,6 @@ def to_psql(conn: duckdb.DuckDBPyConnection, dsn: str, *, bulk_load: bool = Fals
             )
 
         if bulk_load:
-            # Stream rows instead of buffering them to preserve order; buffering 180M+ JSON-bearing
-            # rows is what exhausts memory in a single INSERT. Then drop the secondary indexes and
-            # foreign keys so the load does not maintain them per row.
             conn.execute("SET preserve_insertion_order = false")
             for table, name, _add in _BULK_FKS:
                 _pg(conn, f"ALTER TABLE {table} DROP CONSTRAINT IF EXISTS {name}")
@@ -114,8 +105,6 @@ def to_psql(conn: duckdb.DuckDBPyConnection, dsn: str, *, bulk_load: bool = Fals
             _push_chunked(conn, "changesets", _push_changesets)
             _push_chunked(conn, "changeset_stats", _push_changeset_stats)
         elif _pg_has_history(conn):
-            # The history layer (seq_id=0) is already in PG from the bulk load and never changes, so an
-            # incremental --update pushes only the live layer and its parents, not the 180M history rows.
             live_ids = "changeset_id IN (SELECT changeset_id FROM changeset_stats WHERE seq_id <> 0)"
             conn.execute(
                 "INSERT INTO pg_target.users SELECT * FROM users "
@@ -124,7 +113,6 @@ def to_psql(conn: duckdb.DuckDBPyConnection, dsn: str, *, bulk_load: bool = Fals
             _push_changesets(conn, f"WHERE {live_ids}")
             _push_changeset_stats(conn, "WHERE seq_id <> 0")
         else:
-            # No history in PG (a plain live target): push everything (live rows are all seq_id<>0).
             conn.execute("INSERT INTO pg_target.users SELECT * FROM users ON CONFLICT DO NOTHING")
             _push_changesets(conn)
             _push_changeset_stats(conn)
@@ -141,7 +129,6 @@ def to_psql(conn: duckdb.DuckDBPyConnection, dsn: str, *, bulk_load: bool = Fals
         )
 
         if bulk_load:
-            # Rebuild once, with more memory for the sort-based index builds, then refresh planner stats.
             for table, name, add in _BULK_FKS:
                 _pg(conn, f"ALTER TABLE {table} ADD CONSTRAINT {name} {add}")
             for _name, create in _BULK_INDEXES:

osmsg/history.py

@@ -10,26 +10,26 @@
 import duckdb
 import requests
 
-from .ui import info, warn
+from .ui import info, progress_bar, warn
 
 UTC = dt.UTC
 SCHEMA_VERSION = 1
 DEFAULT_HISTORY_URL = "hf://datasets/kshitijrajsharma/osmsg-history"
-HISTORY_SEQ_ID = 0  # sentinel seq_id for rows sourced from the history backfill (no replication seq)
+HISTORY_SEQ_ID = 0
 
 
 @dataclass
 class Manifest:
     schema_version: int
-    min_month: dt.datetime  # first day of the earliest covered month (UTC)
-    frontier: dt.datetime  # first day of the month AFTER the latest covered month (exclusive bound)
+    min_month: dt.datetime
+    frontier: dt.datetime
 
 
 @dataclass
 class WindowSplit:
     remote_start: dt.datetime | None
-    remote_end: dt.datetime | None  # exclusive
-    live_start: dt.datetime  # the live diff path handles [live_start, end]
+    remote_end: dt.datetime | None
+    live_start: dt.datetime
 
     @property
     def has_remote(self) -> bool:
@@ -51,7 +51,6 @@ def has_metadata_filter(self) -> bool:
 
 
 def _manifest_http_url(history_url: str) -> str:
-    # hf://datasets/<repo> -> https://huggingface.co/datasets/<repo>/resolve/main/manifest.json
     if history_url.startswith("hf://datasets/"):
         repo = history_url[len("hf://datasets/") :]
         return f"https://huggingface.co/datasets/{repo}/resolve/main/manifest.json"
@@ -78,7 +77,7 @@ def fetch_manifest(history_url: str, timeout: int = 15) -> Manifest | None:
                 return None
             payload = response.json()
         else:
-            with open(url) as handle:  # local path (testing / self-hosted mirror)
+            with open(url) as handle:
                 payload = json.load(handle)
     except (requests.RequestException, OSError, ValueError) as exc:
         warn(f"history: manifest unreachable ({type(exc).__name__}); using live path.")
@@ -124,9 +123,8 @@ def _months(start: dt.datetime, end: dt.datetime) -> list[tuple[int, int]]:
 
 
 def _partition_list(base: str, dataset: str, months: list[tuple[int, int]]) -> str | None:
-    """Direct read_parquet() over the dataset's month partitions, or None when none exist. A glob would
-    make DuckDB list every partition over the HF API. Local bases are filtered to files that exist,
-    since a converted slice may lack a partition (e.g. a month with metadata but no counted edits)."""
+    """Direct read_parquet() over the given month partitions (local bases filtered to existing files),
+    or None when none exist."""
     root = base.rstrip("/")
     remote = root.startswith(("hf://", "http://", "https://", "s3://"))
     files = [f"{root}/{dataset}/year={year}/month={month}/data.parquet" for (year, month) in months]
@@ -138,9 +136,7 @@ def _partition_list(base: str, dataset: str, months: list[tuple[int, int]]) -> s
 
 
 def _hashtag_predicate(hashtags: list[str], exact_lookup: bool) -> str:
-    """SQL predicate over the changesets `hashtags` list, matching the live ChangesetHandler.
-    Whole-token (case-insensitive) with exact_lookup, otherwise substring. hashtags are already
-    canonicalised to a leading '#'."""
+    """SQL predicate matching the changesets `hashtags` list: whole-token with exact_lookup, else substring."""
     needles = [h.lower() for h in hashtags]
     if exact_lookup:
         terms = ", ".join(f"'{n}'" for n in needles)
@@ -160,10 +156,6 @@ def ingest_remote(
     if split.remote_start is None or split.remote_end is None:
         return 0
     months = _months(split.remote_start, split.remote_end)
-    changesets_src = _partition_list(history_url, "changesets", months)
-    changefiles_src = _partition_list(history_url, "changefiles", months)
-    if changesets_src is None and changefiles_src is None:
-        return 0
     start_iso = split.remote_start.astimezone(UTC).isoformat()
     end_iso = split.remote_end.astimezone(UTC).isoformat()
     in_window = f"created_at >= TIMESTAMPTZ '{start_iso}' AND created_at < TIMESTAMPTZ '{end_iso}'"
@@ -172,21 +164,8 @@ def ingest_remote(
     conn.execute("INSTALL spatial; LOAD spatial;")
     if history_url.startswith(("hf://", "http://", "https://", "s3://")):
         conn.execute("INSTALL httpfs; LOAD httpfs;")
-        # Ride out HF rate-limits on multi-partition reads instead of failing the run.
         conn.execute("SET http_retries=10; SET http_retry_wait_ms=2000; SET http_retry_backoff=1.5;")
 
-    info(f"history: remote ingest {start_iso} -> {end_iso} ({len(months)} month partitions) from {history_url}")
-
-    if changesets_src is not None:
-        # Names for everyone in the window; every changeset_stats uid has a changeset row here.
-        conn.execute(
-            f"""INSERT INTO users
-                SELECT uid, any_value(username) FROM {changesets_src}
-                WHERE {in_window} AND username IS NOT NULL
-                GROUP BY uid
-                ON CONFLICT (uid) DO NOTHING"""
-        )
-
     changeset_preds = [in_window]
     if filters.hashtags:
         changeset_preds.append(_hashtag_predicate(filters.hashtags, filters.exact_lookup))
@@ -199,51 +178,55 @@ def ingest_remote(
         changeset_preds.append(f"uid IN (SELECT uid FROM users WHERE username IN ({names}))")
     changeset_where = " AND ".join(changeset_preds)
 
-    # Always populate changesets: every changeset_stats row needs a parent row (the live path keeps
-    # this invariant via stubs, and Postgres enforces it as a foreign key). A metadata filter narrows
-    # which changesets (and thus which stats) are kept; a plain run keeps all in the window.
-    if changesets_src is not None:
-        conn.execute(
-            f"""INSERT INTO changesets
-                SELECT changeset_id, uid, created_at, hashtags, editor,
-                       CASE WHEN min_lon IS NOT NULL
-                            THEN ST_MakeEnvelope(min_lon, min_lat, max_lon, max_lat) END
-                FROM {changesets_src} WHERE {changeset_where}
-                ON CONFLICT (changeset_id) DO NOTHING"""
-        )
-
     stats_preds = [in_window]
     if filters.has_metadata_filter:
-        # Keep element stats only for changesets that passed the metadata filter above.
         stats_preds.append("changeset_id IN (SELECT changeset_id FROM changesets)")
     stats_where = " AND ".join(stats_preds)
 
-    if changefiles_src is not None:
-        conn.execute(
-            f"""INSERT INTO changeset_stats
-                SELECT changeset_id, {HISTORY_SEQ_ID} AS seq_id, uid,
-                       nodes_created, nodes_modified, nodes_deleted,
-                       ways_created, ways_modified, ways_deleted,
-                       rels_created, rels_modified, rels_deleted,
-                       poi_created, poi_modified, tag_stats
-                FROM {changefiles_src} WHERE {stats_where}
-                ON CONFLICT (seq_id, changeset_id) DO NOTHING"""
-        )
+    info(f"history: remote ingest {start_iso} -> {end_iso} ({len(months)} month partitions) from {history_url}")
+
+    with progress_bar(len(months), unit="months", description="Reading history") as advance:
+        for month in months:
+            changesets_src = _partition_list(history_url, "changesets", [month])
+            changefiles_src = _partition_list(history_url, "changefiles", [month])
+            if changesets_src is not None:
+                conn.execute(
+                    f"""INSERT INTO users
+                        SELECT uid, any_value(username) FROM {changesets_src}
+                        WHERE {in_window} AND username IS NOT NULL
+                        GROUP BY uid ON CONFLICT (uid) DO NOTHING"""
+                )
+                conn.execute(
+                    f"""INSERT INTO changesets
+                        SELECT changeset_id, uid, created_at, hashtags, editor,
+                               CASE WHEN min_lon IS NOT NULL
+                                    THEN ST_MakeEnvelope(min_lon, min_lat, max_lon, max_lat) END
+                        FROM {changesets_src} WHERE {changeset_where}
+                        ON CONFLICT (changeset_id) DO NOTHING"""
+                )
+            if changefiles_src is not None:
+                conn.execute(
+                    f"""INSERT INTO changeset_stats
+                        SELECT changeset_id, {HISTORY_SEQ_ID} AS seq_id, uid,
+                               nodes_created, nodes_modified, nodes_deleted,
+                               ways_created, ways_modified, ways_deleted,
+                               rels_created, rels_modified, rels_deleted,
+                               poi_created, poi_modified, tag_stats
+                        FROM {changefiles_src} WHERE {stats_where}
+                        ON CONFLICT (seq_id, changeset_id) DO NOTHING"""
+                )
+            advance()
+
     row = conn.execute(f"SELECT count(*) FROM changeset_stats WHERE seq_id = {HISTORY_SEQ_ID}").fetchone()
     return row[0] if row else 0
 
 
-# Resume one day before the frontier, not at it. A changeset can stay open for up to 24h, so its
-# edits can straddle the frontier, and converting a date to a replication sequence is not exact. The
-# re-scanned day overlaps the history layer, which the seq_id=0 dedup removes, so this never misses an
-# edit and never double counts.
 RESUME_SAFETY = dt.timedelta(days=1)
 
 
 def seed_resume_at(conn: duckdb.DuckDBPyConnection, resume_at: dt.datetime, replication_url: str) -> dt.datetime | None:
-    """Seed the `state` table so `osmsg --update` resumes at `resume_at` on `replication_url`. Derives
-    the replication sequence from the timestamp, so the caller never picks a seq by hand. Returns the
-    resume timestamp, or None if no sequence resolves at that time."""
+    """Seed `state` so `osmsg --update` resumes at `resume_at` on `replication_url`. Returns resume_at,
+    or None if no sequence resolves."""
     from osmium.replication.server import ReplicationServer
 
     from .db.schema import upsert_state

osmsg/maintain/convert.py

@@ -1,7 +1,5 @@
 """Convert a planet .osh history plus a changeset dump into the changefiles/changesets parquet
-datasets, out of core via osmsg's own DuckDB tables. Streams raw per-edit rows to parquet in bounded
-batches, then aggregates and joins in DuckDB (a changeset's edits are scattered across the .osh, so an
-in-memory pass OOMs at planet scale)."""
+datasets, out of core via osmsg's own DuckDB tables."""
 
 import concurrent.futures as cf
 import datetime as dt
@@ -20,11 +18,8 @@
 
 BATCH = 1_000_000
 CREATE, MODIFY, DELETE = 0, 1, 2
-# Out-of-core settings for planet-scale aggregation. Leave headroom below physical RAM; spill to disk.
 DUCKDB_MEMORY_LIMIT = "40GB"
 DUCKDB_THREADS = 24
-# A global GROUP BY over all string-keyed tag rows OOMs even with spill, and json_group_object does
-# not spill. Shard raw tags to disk by changeset_id % K, then aggregate each shard independently.
 TAG_SHARDS = 64
 
 ELEM_SCHEMA = pa.schema(
@@ -162,9 +157,7 @@ def stream_changesets(dump: str, start: dt.datetime, end: dt.datetime, work: pat
 
 
 def build_tables(con: duckdb.DuckDBPyConnection, work: pathlib.Path) -> None:
-    """Populate osmsg's tables (users, changesets, changeset_stats) from the streamed raw rows. Globs
-    raw_elements_*/raw_tags_* so single-process and split-parallel runs both work: one global GROUP BY
-    recombines each changeset's edits across parts."""
+    """Populate osmsg's tables (users, changesets, changeset_stats) from the streamed raw rows."""
     con.execute("INSTALL json; LOAD json;")
     work = pathlib.Path(work)
     cs = (work / "raw_changesets.parquet").as_posix()
@@ -209,8 +202,6 @@ def build_tables(con: duckdb.DuckDBPyConnection, work: pathlib.Path) -> None:
               a.rels_created, a.rels_modified, a.rels_deleted,
               a.poi_created, a.poi_modified"""
     for b in range(TAG_SHARDS):
-        # Insert this shard's agg changesets; attach tag_stats only if the shard has tags (tiny inputs
-        # and edit-only changesets carry none).
         shard_dir = shards / f"shard={b}"
         if shard_dir.is_dir():
             shard_glob = (shard_dir / "*.parquet").as_posix()
@@ -244,11 +235,8 @@ def build_tables(con: duckdb.DuckDBPyConnection, work: pathlib.Path) -> None:
 
 
 def export_parquet(con: duckdb.DuckDBPyConnection, out: pathlib.Path) -> None:
-    """Materialise the two datasets as persisted tables (a view would re-run the planet-scale joins per
-    partition; a TEMP table would hold 180M JSON rows in RAM), then write Morton-sorted partitions."""
+    """Materialise the two datasets as persisted tables, then write Morton-sorted partitions."""
     con.execute(MORTON_MACROS)
-    # changefiles created_at falls back to the element edit time when the changeset predates the window,
-    # so in-window edits are never dropped.
     con.execute(
         f"""CREATE TABLE changefiles_all AS
             SELECT s.* EXCLUDE (seq_id),
@@ -292,9 +280,8 @@ def aggregate(work: pathlib.Path, out: pathlib.Path) -> pathlib.Path:
 def convert(
     osh: str, changesets: str, start: dt.datetime, end: dt.datetime, work_dir: pathlib.Path, parts: int = 1
 ) -> pathlib.Path:
-    """Convert one .osh history + changeset dump to the two parquet datasets under `work_dir/out`.
-    With parts>1 the history is split at blob boundaries and streamed concurrently. Returns the out
-    directory holding changefiles/, changesets/, and stats.duckdb."""
+    """Convert one .osh history + changeset dump to the two parquet datasets under `work_dir/out`,
+    returned as a path. With parts>1 the history is split and streamed concurrently."""
     work = pathlib.Path(work_dir)
     raw = work / "raw"
     raw.mkdir(parents=True, exist_ok=True)

osmsg/maintain/month.py

@@ -12,9 +12,6 @@
 from .parquet import GEOM_COLS, MORTON_MACROS, write_partitions
 
 UTC = dt.UTC
-# Planet-wide edits are continuous, so a complete month reaches within minutes of its end. A larger
-# shortfall means the source day diffs did not cover the whole month (a mid-day snapshot or lagging
-# replication), so the partition would be published short, the exact gap the read-side backstep masks.
 COMPLETENESS_TOLERANCE = dt.timedelta(hours=1)