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
 
 

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

@@ -58,7 +58,8 @@ def _compute_stats(df: pl.DataFrame) -> dict[str, dict[str, Any]]:
     numeric_types = df.select(selectors.numeric())
     bool_types = df.select([pl.col(pl.Boolean)])
     # df.select([pl.col(pl.Object)])
-    date_types = df.select(selectors.temporal())
+    # Time/Duration are not valid JS Dates in the UI; leave them for unhandled stats.
+    date_types = df.select(selectors.date() | selectors.datetime())
     # get all other columns that have not been selected
     # df.select(
     #     ~cs.by_dtype([pl.Categorical, pl.Utf8, pl.Boolean, pl.Object])

ui/sdk/tests/tracking/test_polars_col_stats.py

@@ -17,8 +17,11 @@
 
 """Module for testing pandas column stats."""
 
+from datetime import date, datetime
+
 import polars as pl
 import pytest
+
 from hamilton_sdk.tracking import polars_col_stats as pcs
 
 
@@ -129,3 +132,44 @@ 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_datetime_column_stats_serializes_date_and_datetime_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"
+
+    date_stats = pcs.datetime_column_stats(
+        name="d",
+        position=1,
+        data_type="Date",
+        count=2,
+        missing=0,
+        zeros=0,
+        min=date(2021, 1, 1),
+        max=date(2021, 1, 3),
+        mean=datetime(2021, 1, 2),
+        quantiles={0.5: date(2021, 1, 2)},
+        histogram={},
+    )
+    assert date_stats.min == "2021-01-01"
+    assert date_stats.max == "2021-01-03"
+    assert date_stats.mean == "2021-01-02T00:00:00"
+    assert date_stats.quantiles[0.5] == "2021-01-02"
+    assert date_stats.base_data_type == "datetime"

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,48 @@ 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_unhandled():
+    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)]),
+            "ts": pl.Series([datetime(2021, 1, 1), datetime(2021, 1, 2), datetime(2021, 1, 3)]),
+        }
+    )
+    actual = ps.compute_stats_df(df, "test", {})
+    time_stats = actual["observability_value"]["t"]
+    duration_stats = actual["observability_value"]["d"]
+    datetime_stats = actual["observability_value"]["ts"]
+    assert time_stats["base_data_type"] == "unhandled"
+    assert duration_stats["base_data_type"] == "unhandled"
+    assert "min" not in time_stats
+    assert "max" not in time_stats
+    assert "min" not in duration_stats
+    assert "max" not in duration_stats
+    assert datetime_stats["base_data_type"] == "datetime"
+    assert datetime_stats["min"] == "2021-01-01T00:00:00"
+    assert datetime_stats["max"] == "2021-01-03T00:00:00"
+    json.dumps(actual)