Add step to df if not set on dataframe

polyaxon · Sep 28, 2023 · 6aabcd3 · 6aabcd3
1 parent 23e3fde
commit 6aabcd3
Showing 1 changed file with 2 additions and 0 deletions.
diff --git a/traceml/traceml/events/schemas.py b/traceml/traceml/events/schemas.py
@@ -417,6 +417,8 @@ def _read_jsonl(
             # Pyarrow automatically converts timestamp fields
             if "timestamp" in df.columns:
                 df["timestamp"] = df["timestamp"].astype(str)
+        if not hasattr(df, 'step'):
+            df['step'] = None
         return df
 
     @classmethod