Merge branch '323-use-forecasttoolslocation_table-instead-of-download…

…ing-populations' into dhm-run-pyrenew-h
CDCgov · Feb 5, 2025 · 8a8aac9 · 8a8aac9
2 parents d2e218c + 59cbd33
commit 8a8aac9
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/pipelines/prep_data.py b/pipelines/prep_data.py
@@ -146,7 +146,9 @@ def process_state_level_data(
 
     if state_abb == "US":
         locations_to_aggregate = (
-            state_pop_df["abb"].filter(pl.col("abb") != "US").unique()
+            state_pop_df.filter(pl.col("abb") != "US")
+            .get_column("abb")
+            .unique()
         )
         logger.info("Aggregating state-level data to national")
         state_level_nssp_data = aggregate_to_national(
@@ -207,7 +209,9 @@ def aggregate_facility_level_nssp_to_state(
     if state_abb == "US":
         logger.info("Aggregating facility-level data to national")
         locations_to_aggregate = (
-            state_pop_df["abb"].filter(pl.col("abb") != "US").unique()
+            state_pop_df.filter(pl.col("abb") != "US")
+            .get_column("abb")
+            .unique()
         )
         facility_level_nssp_data = aggregate_to_national(
             facility_level_nssp_data,