Fix Polars temporal column stats JSON serialization for tracker

ui/sdk/src/hamilton_sdk/tracking/polars_col_stats.py

@@ -19,6 +19,7 @@
 
 import polars as pl
 from polars.exceptions import InvalidOperationError
+
 from hamilton_sdk.tracking import dataframe_stats as dfs
 
 
@@ -130,6 +131,16 @@ def numeric_column_stats(
     )
 
 
+def _temporal_to_jsonable(value: object) -> object:
+    # datetime.datetime is a subclass of datetime.date. Time/Duration also arrive via
+    # selectors.temporal() and must be JSON-safe for the UI (same base_data_type="datetime").
+    if isinstance(value, (datetime.date, datetime.time)):
+        return value.isoformat()
+    if isinstance(value, datetime.timedelta):
+        return str(value)
+    return value
+
+
 def datetime_column_stats(
     name: str,
     position: int,
@@ -144,13 +155,10 @@ def datetime_column_stats(
     histogram: dict[str, int],
 ) -> dfs.DatetimeColumnStatistics:
     # TODO: push these conversions into Hamilton functions.
-    # Note: datetime.datetime is a subclass of datetime.date, so checking datetime.date catches both
-    min = min.isoformat() if isinstance(min, datetime.date) else min
-    max = max.isoformat() if isinstance(max, datetime.date) else max
-    mean = mean.isoformat() if isinstance(mean, datetime.date) else mean
-    quantiles = {
-        q: v.isoformat() if isinstance(v, datetime.date) else v for q, v in quantiles.items()
-    }
+    min = _temporal_to_jsonable(min)
+    max = _temporal_to_jsonable(max)
+    mean = _temporal_to_jsonable(mean)
+    quantiles = {q: _temporal_to_jsonable(v) for q, v in quantiles.items()}
     return dfs.DatetimeColumnStatistics(
         name=name,
         pos=position,

ui/sdk/tests/tracking/test_polars_col_stats.py

@@ -17,8 +17,11 @@
 
 """Module for testing pandas column stats."""
 
+from datetime import date, datetime, time, timedelta
+
 import polars as pl
 import pytest
+
 from hamilton_sdk.tracking import polars_col_stats as pcs
 
 
@@ -129,3 +132,67 @@ def test_max_string(example_df_string):
 
 def test_str_len(example_df_string):
     assert pcs.str_len(example_df_string["a"]).to_list() == [1, 1, 1, 1, 1]
+
+
+def test_temporal_to_jsonable():
+    assert pcs._temporal_to_jsonable(datetime(2021, 1, 2, 3, 4, 5)) == "2021-01-02T03:04:05"
+    assert pcs._temporal_to_jsonable(date(2021, 1, 2)) == "2021-01-02"
+    assert pcs._temporal_to_jsonable(time(1, 2, 3)) == "01:02:03"
+    assert pcs._temporal_to_jsonable(timedelta(days=1, hours=2)) == "1 day, 2:00:00"
+    assert pcs._temporal_to_jsonable(3.14) == 3.14
+
+
+def test_datetime_column_stats_serializes_temporal_values():
+    stats = pcs.datetime_column_stats(
+        name="ts",
+        position=0,
+        data_type="Datetime(time_unit='us', time_zone=None)",
+        count=3,
+        missing=0,
+        zeros=0,
+        min=datetime(2021, 1, 1),
+        max=datetime(2021, 1, 3),
+        mean=datetime(2021, 1, 2),
+        quantiles={0.5: datetime(2021, 1, 2)},
+        histogram={},
+    )
+    assert stats.min == "2021-01-01T00:00:00"
+    assert stats.max == "2021-01-03T00:00:00"
+    assert stats.mean == "2021-01-02T00:00:00"
+    assert stats.std == 0.0
+    assert stats.quantiles[0.5] == "2021-01-02T00:00:00"
+    assert stats.base_data_type == "datetime"
+
+    time_stats = pcs.datetime_column_stats(
+        name="t",
+        position=1,
+        data_type="Time",
+        count=2,
+        missing=0,
+        zeros=0,
+        min=time(1, 0),
+        max=time(2, 0),
+        mean=time(1, 30),
+        quantiles={0.5: time(2, 0)},
+        histogram={},
+    )
+    assert time_stats.min == "01:00:00"
+    assert time_stats.max == "02:00:00"
+    assert time_stats.mean == "01:30:00"
+
+    duration_stats = pcs.datetime_column_stats(
+        name="d",
+        position=2,
+        data_type="Duration(time_unit='us')",
+        count=2,
+        missing=0,
+        zeros=0,
+        min=timedelta(days=1),
+        max=timedelta(days=2),
+        mean=timedelta(days=1, hours=12),
+        quantiles={0.5: timedelta(days=2)},
+        histogram={},
+    )
+    assert duration_stats.min == "1 day, 0:00:00"
+    assert duration_stats.max == "2 days, 0:00:00"
+    assert duration_stats.mean == "1 day, 12:00:00"

ui/sdk/tests/tracking/test_polars_stats.py

@@ -15,7 +15,8 @@
 # specific language governing permissions and limitations
 # under the License.
 
-from datetime import date
+import json
+from datetime import date, datetime, time, timedelta
 
 import polars as pl
 
@@ -213,3 +214,47 @@ def test_compute_stats_df():
         actual["observability_value"][col].pop("quantiles", None)
         expected_stats["observability_value"][col].pop("quantiles", None)
     assert actual == expected_stats
+
+
+def test_compute_stats_datetime_series_regression():
+    # Regression for #1127: Datetime columns must not error via std() and must stay trackable.
+    series = pl.Series(
+        "timestamp",
+        [
+            datetime(2021, 1, 1),
+            datetime(2021, 1, 2),
+            datetime(2021, 1, 3),
+        ],
+    )
+    actual = ps.compute_stats_series(series, "df", {})
+    column_stats = actual["observability_value"]["df"]
+    assert column_stats["base_data_type"] == "datetime"
+    assert column_stats["data_type"].startswith("Datetime")
+    assert column_stats["std"] == 0.0
+    assert column_stats["min"] == "2021-01-01T00:00:00"
+    assert column_stats["max"] == "2021-01-03T00:00:00"
+    assert column_stats["mean"] == "2021-01-02T00:00:00"
+    json.dumps(actual)
+
+
+def test_compute_stats_time_and_duration_columns_are_json_safe():
+    df = pl.DataFrame(
+        {
+            "t": pl.Series([time(1, 0), time(2, 0), time(3, 0)]),
+            "d": pl.Series([timedelta(days=1), timedelta(days=2), timedelta(days=3)]),
+        }
+    )
+    actual = ps.compute_stats_df(df, "test", {})
+    time_stats = actual["observability_value"]["t"]
+    duration_stats = actual["observability_value"]["d"]
+    assert time_stats["base_data_type"] == "datetime"
+    assert duration_stats["base_data_type"] == "datetime"
+    assert time_stats["std"] == 0.0
+    assert duration_stats["std"] == 0.0
+    assert isinstance(time_stats["min"], str)
+    assert isinstance(time_stats["max"], str)
+    assert isinstance(time_stats["mean"], str)
+    assert isinstance(duration_stats["min"], str)
+    assert isinstance(duration_stats["max"], str)
+    assert isinstance(duration_stats["mean"], str)
+    json.dumps(actual)