diff --git a/ui/sdk/src/hamilton_sdk/tracking/polars_col_stats.py b/ui/sdk/src/hamilton_sdk/tracking/polars_col_stats.py index 04586d3c0..aab11f679 100644 --- a/ui/sdk/src/hamilton_sdk/tracking/polars_col_stats.py +++ b/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, diff --git a/ui/sdk/tests/tracking/test_polars_col_stats.py b/ui/sdk/tests/tracking/test_polars_col_stats.py index f80c12b86..4f2fd8104 100644 --- a/ui/sdk/tests/tracking/test_polars_col_stats.py +++ b/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" diff --git a/ui/sdk/tests/tracking/test_polars_stats.py b/ui/sdk/tests/tracking/test_polars_stats.py index 6c8129a59..98d10c3b3 100644 --- a/ui/sdk/tests/tracking/test_polars_stats.py +++ b/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)