Test Empire with NANs

pipeline/s3_timeseries/pipeline.py

@@ -1,3 +1,4 @@
+import logging
 from datetime import date, datetime
 from textwrap import dedent
 from typing import Annotated
@@ -199,6 +200,21 @@ def monthly_ds(
                     f"Column name collision after renaming for data on {df_date}, trying to squish duplicates",
                 )
                 df = df.groupby(df.columns, axis=1).first()
+
+            # Avoid attempting to convert the time column to numeric inside
+            # clean_up_dtypes_and_nas, which causes unnecessary exceptions.
+            if "time" in df.columns:
+                time_col = df["time"]
+                df_wo_time = df.drop(columns=["time"])
+                df_wo_time = clean_up_dtypes_and_nas(
+                    df_wo_time,
+                    na_values="NAN",
+                    logger=context.log,
+                )
+                df_wo_time["time"] = time_col
+                df = df_wo_time
+            else:
+                df = clean_up_dtypes_and_nas(df, na_values="NAN", logger=context.log)
             daily_dfs.append(df)
 
         df = pd.concat(daily_dfs, ignore_index=True)
@@ -367,3 +383,25 @@ def build_defs() -> dg.Definitions:
             defs = dg.Definitions.merge(defs, dataset_defs)
 
         return defs
+
+
+def clean_up_dtypes_and_nas(
+    df: pd.DataFrame,
+    na_values: None | str | list[str] = None,
+    logger: logging.Logger | None = None,
+) -> pd.DataFrame:
+    """Clean up data types and NA values in a dataframe"""
+    df = df.copy()
+    if not logger:
+        logger = logging.getLogger(__name__)
+    if na_values is not None:
+        df = df.replace(na_values, pd.NA).dropna()
+
+    for c in df.columns:
+        try:
+            df[c] = pd.to_numeric(df[c])
+        except (ValueError, TypeError) as e:
+            if logger:
+                logger.warning(f"Could not convert column {c} to numeric: {e}")
+
+    return df