From 072c643d8c1d3b50ef41bf3eade0eb68fd029980 Mon Sep 17 00:00:00 2001 From: Markus Zimmerberger Date: Sat, 12 Sep 2026 08:07:42 +0200 Subject: [PATCH] Reduce RAM consumtion during import --- stops_import.py | 152 ++++++++++++++++++++++++++++++++---------------- 1 file changed, 102 insertions(+), 50 deletions(-) diff --git a/stops_import.py b/stops_import.py index 72970f2..4e09b08 100644 --- a/stops_import.py +++ b/stops_import.py @@ -4,6 +4,13 @@ The snapshot (data/stops_dach.csv.gz) is produced by scripts/fetch_stops.py from OpenStreetMap data (© OpenStreetMap contributors, ODbL) for Austria, Germany and Switzerland. Importing needs no network access. +At ~307k rows, processing this in one shot - the whole table loaded as ORM +objects, ~307k new Stop() instances held until a single final commit - is +enough to OOM a small host (confirmed: it did, on a 512 MB production LXC). +Everything below is batched (default 2000 rows) with a commit + expunge_all() +after each batch, so peak memory stays roughly flat regardless of table size +instead of growing with it. + flask import-stops # import the checked-in snapshot flask import-stops --file X # import an alternative CSV(.gz) """ @@ -16,12 +23,12 @@ import unicodedata from pathlib import Path import click -from flask import current_app from flask.cli import with_appcontext from models import Stop, Trip, db SNAPSHOT_PATH = Path(__file__).resolve().parent / "data" / "stops_dach.csv.gz" +BATCH_SIZE = 2000 def normalize_name(value: str) -> str: @@ -40,6 +47,78 @@ def _open_csv(path: Path) -> io.TextIOBase: return open(path, "rt", newline="", encoding="utf-8") +def _row_fields(row: dict) -> dict: + name = row["name"].strip() + return dict( + name=name, + name_normalized=normalize_name(name), + stop_type=row["stop_type"] or "other", + country=(row.get("country") or "AT").strip().upper(), + municipality=(row.get("municipality") or "").strip() or None, + latitude=float(row["latitude"]), + longitude=float(row["longitude"]), + lines=(row.get("lines") or "").strip(), + ) + + +def _import_batch(batch: list[dict]) -> tuple[int, int]: + """Upsert one batch. Only this batch's rows are ever loaded as ORM + objects, so memory stays bounded by BATCH_SIZE, not table size.""" + ids = [row["osm_id"] for row in batch] + existing = { + (s.osm_type, s.osm_id): s + for s in Stop.query.filter(Stop.osm_id.in_(ids)) + } + inserted = updated = 0 + for row in batch: + key = (row["osm_type"], row["osm_id"]) + fields = row["fields"] + stop = existing.get(key) + if stop is None: + db.session.add(Stop(osm_type=key[0], osm_id=key[1], **fields)) + inserted += 1 + else: + changed = False + for attr, value in fields.items(): + if getattr(stop, attr) != value: + setattr(stop, attr, value) + changed = True + updated += changed + db.session.commit() + db.session.expunge_all() # drop the identity map's refs so GC can reclaim them + return inserted, updated + + +def _remove_stale(seen: set[tuple[str, int]]) -> int: + """Delete stops absent from the new snapshot, unless a trip still + references them. Streams the id/key columns only (yield_per) rather than + loading full Stop rows for the whole table.""" + referenced = { + sid + for (sid,) in db.session.query(Trip.origin_stop_id).distinct() + if sid is not None + } | { + sid + for (sid,) in db.session.query(Trip.destination_stop_id).distinct() + if sid is not None + } + + stale_ids = [ + sid + for sid, osm_type, osm_id in db.session.query( + Stop.id, Stop.osm_type, Stop.osm_id + ).yield_per(5000) + if (osm_type, osm_id) not in seen and sid not in referenced + ] + + removed = 0 + for i in range(0, len(stale_ids), BATCH_SIZE): + chunk = stale_ids[i : i + BATCH_SIZE] + removed += Stop.query.filter(Stop.id.in_(chunk)).delete(synchronize_session=False) + db.session.commit() + return removed + + def import_stops(path: Path | str = SNAPSHOT_PATH) -> dict[str, int]: """Upsert every row of the snapshot into ``stops`` by (osm_type, osm_id). @@ -52,59 +131,32 @@ def import_stops(path: Path | str = SNAPSHOT_PATH) -> dict[str, int]: f"Snapshot not found: {path}\nRun scripts/fetch_stops.py first." ) + seen: set[tuple[str, int]] = set() + inserted = updated = total = 0 + batch: list[dict] = [] + with _open_csv(path) as fh: - rows = list(csv.DictReader(fh)) - if not rows: + for row in csv.DictReader(fh): + key = (row["osm_type"], int(row["osm_id"])) + seen.add(key) + total += 1 + batch.append({"osm_type": key[0], "osm_id": key[1], "fields": _row_fields(row)}) + if len(batch) >= BATCH_SIZE: + i, u = _import_batch(batch) + inserted += i + updated += u + batch = [] + if batch: + i, u = _import_batch(batch) + inserted += i + updated += u + + if total == 0: raise click.ClickException(f"{path} contains no rows.") - existing = {(s.osm_type, s.osm_id): s for s in Stop.query.all()} - seen: set[tuple[str, int]] = set() - inserted = updated = 0 - - for row in rows: - key = (row["osm_type"], int(row["osm_id"])) - seen.add(key) - name = row["name"].strip() - fields = dict( - name=name, - name_normalized=normalize_name(name), - stop_type=row["stop_type"] or "other", - country=(row.get("country") or "AT").strip().upper(), - municipality=(row.get("municipality") or "").strip() or None, - latitude=float(row["latitude"]), - longitude=float(row["longitude"]), - lines=(row.get("lines") or "").strip(), - ) - stop = existing.get(key) - if stop is None: - db.session.add(Stop(osm_type=key[0], osm_id=key[1], **fields)) - inserted += 1 - else: - changed = False - for attr, value in fields.items(): - if getattr(stop, attr) != value: - setattr(stop, attr, value) - changed = True - updated += changed - - referenced = { - sid - for (sid,) in db.session.query(Trip.origin_stop_id).distinct() - if sid is not None - } | { - sid - for (sid,) in db.session.query(Trip.destination_stop_id).distinct() - if sid is not None - } - removed = 0 - for key, stop in existing.items(): - if key not in seen and stop.id not in referenced: - db.session.delete(stop) - removed += 1 - - db.session.commit() + removed = _remove_stale(seen) return { - "total": len(rows), + "total": total, "inserted": inserted, "updated": updated, "removed": removed,