Flatten metrics column in parquet representation (#106)

* When exporting to parquet, flatten the `metrics` column into multiple
  columns of Parquet.
* When importing, do the reverse (if the parquet file
  doesn't already have a `metrics` column).
* If the parquet file has a metrics column (legacy format), it will skip
  the unflattening on import, so it should continue to work.
* Adds "metrics" to the reserved metric names so that it's not possible
  to log a metric name that will get this logic confused.

Fixes #105

Author: znation

trackio/sqlite_storage.py

@@ -93,6 +93,14 @@ def export_to_parquet():
             ):
                 with sqlite3.connect(db_path) as conn:
                     df = pd.read_sql("SELECT * from metrics", conn)
+                # break out the single JSON metrics column into individual columns
+                metrics = df["metrics"].copy()
+                metrics = pd.DataFrame(
+                    metrics.apply(json.loads).values.tolist(), index=df.index
+                )
+                del df["metrics"]
+                for col in metrics.columns:
+                    df[col] = metrics[col]
                 df.to_parquet(parquet_path)
 
     @staticmethod
@@ -107,6 +115,17 @@ def import_from_parquet():
             db_path = parquet_path.with_suffix(".db")
             df = pd.read_parquet(parquet_path)
             with sqlite3.connect(db_path) as conn:
+                # fix up df to have a single JSON metrics column
+                if "metrics" not in df.columns:
+                    # separate other columns from metrics
+                    metrics = df.copy()
+                    other_cols = ["id", "timestamp", "run_name", "step"]
+                    df = df[other_cols]
+                    for col in other_cols:
+                        del metrics[col]
+                    # combine them all into a single metrics col
+                    metrics = json.loads(metrics.to_json(orient="records"))
+                    df["metrics"] = [json.dumps(row) for row in metrics]
                 df.to_sql("metrics", conn, if_exists="replace", index=False)
 
     @staticmethod

trackio/utils.py

@@ -8,7 +8,7 @@
 import pandas as pd
 from huggingface_hub.constants import HF_HOME
 
-RESERVED_KEYS = ["project", "run", "timestamp", "step", "time"]
+RESERVED_KEYS = ["project", "run", "timestamp", "step", "time", "metrics"]
 TRACKIO_DIR = Path(HF_HOME) / "trackio"
 
 TRACKIO_LOGO_DIR = Path(__file__).parent / "assets"