perf: Reduce memory overhead in trip augmentation during ingest

Replace the bucket-and-sort approach in get_trip_rows_with_extra_time_fields
with a single streaming pass over stop_times that tracks min/max stop
sequence per trip, avoiding materializing full trip/stop_time lists in memory.

Author: ankoure

.github/workflows/ci.yml

@@ -0,0 +1,25 @@
+name: CI
+
+on:
+  push:
+    branches: [main]
+  pull_request:
+
+jobs:
+  install-and-lint:
+    runs-on: ubuntu-latest
+    steps:
+      - uses: actions/checkout@v4
+
+      - uses: actions/setup-python@v5
+        with:
+          python-version: "3.12"
+
+      - name: Install Poetry
+        run: pip install poetry
+
+      - name: Install dependencies
+        run: poetry install
+
+      - name: Check formatting with black
+        run: poetry run black --check .

mbta_gtfs_sqlite/feed.py

@@ -13,6 +13,7 @@
 ALL_DB_FILES = [DB_FILE, DB_COMPACT_FILE]
 DEFAULT_INGEST_BATCH_SIZE = 300000
 
+
 @dataclass
 class GtfsFeed(object):
     archive: "MbtaGtfsArchive"

mbta_gtfs_sqlite/ingest.py

@@ -1,12 +1,10 @@
 from sqlalchemy.orm import Session
-from typing import Dict, Any, Callable, List, Type, Iterable, Union
+from typing import Dict, Any, Callable, Type, Iterable, Union
 from more_itertools import ichunked
 
 from .build import GtfsFeedDownloadResult
 from .reader import GtfsReader
 from .utils.time import date_from_string, seconds_from_string
-from .utils.decorators import listify
-from .utils.indexes import bucket_by
 from .models.agency import Agency
 from .models.base import Base
 from .models.calendar_attributes import CalendarAttribute
@@ -47,23 +45,32 @@ def transform_row_dict(
     }
 
 
-@listify
-def get_trip_rows_with_extra_time_fields(
-    trip_rows: List[Dict[str, str]],
-    stop_time_rows: List[Dict[str, str]],
-) -> List[Dict[str, str]]:
-    stop_times_by_trip_id = bucket_by(stop_time_rows, "trip_id")
-    for trip_row in trip_rows:
-        stop_times_for_trip = sorted(
-            stop_times_by_trip_id[trip_row["trip_id"]],
-            key=lambda stop_time: int(stop_time["stop_sequence"]),
-        )
-        yield {
-            **trip_row,
-            "start_time": stop_times_for_trip[0]["arrival_time"],
-            "end_time": stop_times_for_trip[-1]["arrival_time"],
-            "stop_count": len(stop_times_for_trip),
-        }
+def get_trip_extra_fields_by_trip_id(
+    stop_time_rows: Iterable[Dict[str, str]],
+) -> Dict[str, Dict[str, Any]]:
+    extra_fields_by_trip_id: Dict[str, Dict[str, Any]] = {}
+    for stop_time_row in stop_time_rows:
+        trip_id = stop_time_row["trip_id"]
+        stop_sequence = int(stop_time_row["stop_sequence"])
+        arrival_time = stop_time_row["arrival_time"]
+        extra_fields = extra_fields_by_trip_id.get(trip_id)
+        if extra_fields is None:
+            extra_fields_by_trip_id[trip_id] = {
+                "min_stop_sequence": stop_sequence,
+                "start_time": arrival_time,
+                "max_stop_sequence": stop_sequence,
+                "end_time": arrival_time,
+                "stop_count": 1,
+            }
+            continue
+        extra_fields["stop_count"] += 1
+        if stop_sequence < extra_fields["min_stop_sequence"]:
+            extra_fields["min_stop_sequence"] = stop_sequence
+            extra_fields["start_time"] = arrival_time
+        if stop_sequence > extra_fields["max_stop_sequence"]:
+            extra_fields["max_stop_sequence"] = stop_sequence
+            extra_fields["end_time"] = arrival_time
+    return extra_fields_by_trip_id
 
 
 def ingest_feed_info(
@@ -117,11 +124,16 @@ def ingest_rows(
         session.bulk_insert_mappings(model, mappings)
 
 
-def get_augmented_trip_rows(reader: GtfsReader):
-    stop_times = list(reader.read_stop_times())
-    trips = list(reader.read_trips())
-    trip_rows = get_trip_rows_with_extra_time_fields(trips, stop_times)
-    return trip_rows
+def get_augmented_trip_rows(reader: GtfsReader) -> Iterable[Dict[str, str]]:
+    extra_fields_by_trip_id = get_trip_extra_fields_by_trip_id(reader.read_stop_times())
+    for trip_row in reader.read_trips():
+        extra_fields = extra_fields_by_trip_id[trip_row["trip_id"]]
+        yield {
+            **trip_row,
+            "start_time": extra_fields["start_time"],
+            "end_time": extra_fields["end_time"],
+            "stop_count": extra_fields["stop_count"],
+        }
 
 
 def ingest_gtfs_csv_into_db(