fix: remove hard-coded type from pin_metrics, add tests

Author: machow

vetiver/monitor.py

@@ -73,8 +73,13 @@ def _rolling_df(df: pd.DataFrame, td: timedelta):
 
 
 def pin_metrics(
-    board, df_metrics, metrics_pin_name: str, index_name="index", overwrite=False
-):
+    board,
+    df_metrics: pd.DataFrame,
+    metrics_pin_name: str,
+    pin_type: "str | None" = None,
+    index_name: str = "index",
+    overwrite: bool = False,
+) -> pd.DataFrame:
     """
     Update an existing pin storing model metrics over time
 
@@ -97,21 +102,18 @@ def pin_metrics(
         the existing pin.
     """
 
-    new_dates = df_metrics[index_name]
+    old_metrics_raw = board.pin_read(metrics_pin_name)
 
-    old_metrics = board.pin_read(metrics_pin_name)
-    old_dates = old_metrics[index_name]
+    # need to coerce date index to a datetime, since pandas does not infer
+    # date columns from CSV (but note that formats like arrow do)
+    old_metrics = old_metrics_raw.copy()
+    old_metrics[index_name] = pd.to_datetime(old_metrics[index_name])
 
     # handle overlapping dates ----
-    if new_dates.dtype != old_dates.dtype:
-        raise TypeError(
-            f"index_name column ({repr(index_name)}) in old and new metrics "
-            "must have the same dtype. "
-            f"\nOld dtype: {old_dates.dtype}"
-            f"\nNew dtype: {new_dates.dtype}"
-        )
+    dt_new = pd.to_datetime(df_metrics[index_name])
+    dt_old = old_metrics[index_name]
 
-    indx_old_overlap = old_metrics[index_name].isin(new_dates)
+    indx_old_overlap = dt_old.isin(dt_new)
 
     if overwrite:
         # get only rows specific to old metrics, so when we concat below
@@ -128,7 +130,14 @@ def pin_metrics(
     combined_metrics = pd.concat([old_metrics, df_metrics], ignore_index=True)
     sorted_metrics = combined_metrics.sort_values(index_name)
 
-    board.pin_write(sorted_metrics, metrics_pin_name, type="arrow")
+    if pin_type is None:
+        meta = board.pin_meta(metrics_pin_name)
+
+        final_pin_type = meta.type
+    else:
+        final_pin_type = pin_type
+
+    board.pin_write(sorted_metrics, metrics_pin_name, type=final_pin_type)
 
     return sorted_metrics
 

vetiver/tests/test_monitor.py

@@ -6,6 +6,8 @@
 import numpy
 import vetiver
 
+import pytest
+
 rng = pd.date_range("1/1/2012", periods=10, freq="S")
 new = dict(x=range(len(rng)), y=range(len(rng)))
 df = pd.DataFrame(new, index=rng)
@@ -41,17 +43,22 @@ def test_monitor(snapshot):
     snapshot.assert_match(m.to_json(), "test_monitor.json")
 
 
-def test_vetiver_pin_metrics():
-    board = pins.board_temp()
-    df_metrics_old = pd.DataFrame(
+@pytest.fixture
+def df_metrics_old():
+    return pd.DataFrame(
         {
             "index": pd.to_datetime(["2021-01-01", "2021-01-02"]),
             "n": [1, 2],
             "metric": ["x", "x"],
-            "estimate": [0.6, 0.7],
+            "estimate": [0.1, 0.2],
         }
     )
 
+
+def test_vetiver_pin_metrics_simple(df_metrics_old):
+    board = pins.board_temp()
+    board.pin_write(df_metrics_old, "test_metrics", type="csv")
+
     df_metrics_new = pd.DataFrame(
         {
             "index": pd.to_datetime(["2021-01-03", "2021-01-04"]),
@@ -61,8 +68,53 @@ def test_vetiver_pin_metrics():
         }
     )
 
-    board.pin_write(df_metrics_old, "test_metrics", type="arrow")
-
     df_res = vetiver.pin_metrics(board, df_metrics_new, "test_metrics")
 
     assert len(df_res) == 4
+    assert df_res.equals(pd.concat([df_metrics_old, df_metrics_new], ignore_index=True))
+
+
+def test_vetiver_pin_metrics_overlap_error(df_metrics_old):
+    board = pins.board_temp()
+    board.pin_write(df_metrics_old, "test_metrics", type="csv")
+
+    with pytest.raises(ValueError) as exc_info:
+        vetiver.pin_metrics(board, df_metrics_old, "test_metrics")
+
+    assert "The new metrics overlap" in exc_info.value.args[0]
+
+
+def test_vetiver_pin_metrics_overwrite(df_metrics_old):
+    board = pins.board_temp()
+    board.pin_write(df_metrics_old, "test_metrics", type="csv")
+
+    # first row should update existing metrics
+    df_metrics_new = pd.DataFrame(
+        {
+            "index": pd.to_datetime(["2021-01-01", "2021-01-03"]),
+            "n": [200, 201],
+            "metric": ["y", "y"],
+            "estimate": [0.8, 0.9],
+        }
+    )
+
+    df_res = vetiver.pin_metrics(board, df_metrics_new, "test_metrics", overwrite=True)
+    assert len(df_res) == 3
+
+    df_dst = pd.concat([df_metrics_old.iloc[[1], :], df_metrics_new], ignore_index=True)
+    assert df_res.equals(df_dst.sort_values("index"))
+
+
+def test_vetiver_pin_metrics_manual_pin_type(df_metrics_old):
+    board = pins.board_temp()
+    board.pin_write(df_metrics_old, "test_metrics", type="csv")
+
+    df_res = vetiver.pin_metrics(
+        board, df_metrics_old, "test_metrics", overwrite=True, pin_type="joblib"
+    )
+
+    assert len(df_res) == 2
+
+    meta = board.pin_meta("test_metrics")
+
+    assert meta.type == "joblib"