Fix filter

CDCgov · Feb 5, 2025 · 1b7b106 · 1b7b106
1 parent b9bb9cb
commit 1b7b106
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/pipelines/prep_data.py b/pipelines/prep_data.py
@@ -146,7 +146,7 @@ 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 +207,7 @@ 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,