drop extra columns

openclimatefix · Nov 18, 2024 · 3a004f3 · 3a004f3
1 parent 6ce150f
commit 3a004f3
Showing 1 changed file with 8 additions and 0 deletions.
diff --git a/database-cleanup/database_cleanup/app.py b/database-cleanup/database_cleanup/app.py
@@ -101,6 +101,14 @@ def save_forecast_and_values(
         forecasts_sql = query.all()
         forecasts_df = pd.DataFrame([f.__dict__ for f in forecasts_sql])
 
+        # drop column _sa_instance_state if it is there
+        if "_sa_instance_state" in forecasts_df.columns:
+            forecasts_df = forecasts_df.drop(columns="_sa_instance_state")
+
+        # drop forecast_value_uuid as we dont need it
+        if table == "forecast_value":
+            forecasts_df = forecasts_df.drop(columns="forecast_value_uuid")
+
         # save to csv
         _log.info(f"saving to {directory}, Saving {len(forecasts_df)} rows to {table}.csv")
         forecasts_df.to_csv(f"{directory}/{table}_{index}.csv", index=False)