apply bugfix

nextstrain · Dec 10, 2021 · aba5515 · aba5515
1 parent 4e3e155
commit aba5515
Showing 1 changed file with 5 additions and 2 deletions.
diff --git a/augur/filter.py b/augur/filter.py
@@ -957,8 +957,11 @@ def expand_date_col(metadata: pd.DataFrame, group_by_set: set) -> Tuple[pd.DataF
     skipped_strains = []
     # replace date with year/month/day as nullable ints
     date_cols = ['year', 'month', 'day']
-    df_dates = (metadata_new['date'].str.split('-', n=2, expand=True)
-                                    .set_axis(date_cols, axis=1))
+    df_dates = metadata['date'].str.split('-', n=2, expand=True)
+    df_dates = df_dates.set_axis(date_cols[:len(df_dates.columns)], axis=1)
+    missing_date_cols = set(date_cols) - set(df_dates.columns)
+    for col in missing_date_cols:
+        df_dates[col] = pd.NA
     for col in date_cols:
         df_dates[col] = pd.to_numeric(df_dates[col], errors='coerce').astype(pd.Int64Dtype())
     metadata_new = pd.concat([metadata_new.drop('date', axis=1), df_dates], axis=1)