feat(pems_data): initial station data helper classes

- metadata and imputed aggregate 5min URLs
- get district metadata
- get station imputed aggregate 5min data

Author: thekaveman

pems_data/src/pems_data/stations.py

@@ -0,0 +1,55 @@
+import pandas as pd
+
+from pems_data.s3 import S3Bucket
+
+
+class StationsBucket(S3Bucket):
+    """Station-specific bucket config."""
+
+    imputation_detector_agg_5min = "imputation/detector_imputed_agg_five_minutes"
+    metadata_file = "geo/current_stations.parquet"
+
+
+class Stations:
+    """Station-specific data."""
+
+    def __init__(self, bucket: StationsBucket = StationsBucket()):
+        self.bucket = bucket
+
+    def get_district_metadata(self, district_number: str) -> pd.DataFrame:
+        """Loads metadata for all stations in the selected District from S3."""
+
+        columns = [
+            "STATION_ID",
+            "NAME",
+            "PHYSICAL_LANES",
+            "STATE_POSTMILE",
+            "ABSOLUTE_POSTMILE",
+            "LATITUDE",
+            "LONGITUDE",
+            "LENGTH",
+            "STATION_TYPE",
+            "DISTRICT",
+            "FREEWAY",
+            "DIRECTION",
+            "COUNTY_NAME",
+            "CITY_NAME",
+        ]
+        filters = [("DISTRICT", "=", district_number)]
+
+        return self.bucket.read_parquet(self.bucket.metadata_file, columns=columns, filters=filters)
+
+    def get_imputed_agg_5min(self, station_id: str) -> pd.DataFrame:
+        """Loads imputed aggregate 5 minute data for a specific station."""
+
+        columns = [
+            "STATION_ID",
+            "LANE",
+            "SAMPLE_TIMESTAMP",
+            "VOLUME_SUM",
+            "SPEED_FIVE_MINS",
+            "OCCUPANCY_AVG",
+        ]
+        filters = [("STATION_ID", "=", station_id)]
+
+        return self.bucket.read_parquet(self.bucket.imputation_detector_agg_5min, columns=columns, filters=filters)

tests/pytest/pems_data/test_stations.py

@@ -0,0 +1,74 @@
+import pandas as pd
+
+from pems_data.stations import Stations, StationsBucket
+import pytest
+
+
+class TestStationsBucket:
+    def test_imputation_detector_agg_5min(self):
+        assert StationsBucket.imputation_detector_agg_5min == "imputation/detector_imputed_agg_five_minutes"
+
+    def test_metadata_file(self):
+        assert StationsBucket.metadata_file == "geo/current_stations.parquet"
+
+
+class TestStations:
+    @pytest.fixture
+    def df(self):
+        return pd.DataFrame({"STATION_ID": [1]})
+
+    @pytest.fixture
+    def stations(self):
+        return Stations()
+
+    @pytest.fixture
+    def mock_read_parquet(self, df, mocker):
+        return mocker.patch("pems_data.stations.StationsBucket.read_parquet", return_value=df)
+
+    def test_bucket(self, stations: Stations):
+        assert isinstance(stations.bucket, StationsBucket)
+
+    def test_get_district_metadata(self, stations: Stations, df, mock_read_parquet):
+        district_number = "7"
+        result = stations.get_district_metadata(district_number)
+
+        mock_read_parquet.assert_called_once_with(
+            StationsBucket.metadata_file,
+            columns=[
+                "STATION_ID",
+                "NAME",
+                "PHYSICAL_LANES",
+                "STATE_POSTMILE",
+                "ABSOLUTE_POSTMILE",
+                "LATITUDE",
+                "LONGITUDE",
+                "LENGTH",
+                "STATION_TYPE",
+                "DISTRICT",
+                "FREEWAY",
+                "DIRECTION",
+                "COUNTY_NAME",
+                "CITY_NAME",
+            ],
+            filters=[("DISTRICT", "=", district_number)],
+        )
+        assert result.equals(df)
+
+    def test_get_imputed_agg_5min(self, stations: Stations, df, mock_read_parquet):
+        station_id = "123"
+
+        result = stations.get_imputed_agg_5min(station_id)
+
+        mock_read_parquet.assert_called_once_with(
+            StationsBucket.imputation_detector_agg_5min,
+            columns=[
+                "STATION_ID",
+                "LANE",
+                "SAMPLE_TIMESTAMP",
+                "VOLUME_SUM",
+                "SPEED_FIVE_MINS",
+                "OCCUPANCY_AVG",
+            ],
+            filters=[("STATION_ID", "=", station_id)],
+        )
+        assert result.equals(df)