Skip to content

Commit 670e8f4

Browse files
Digvijay  Waghelacursoragent
authored andcommitted
Make Polars temporal column stats JSON-safe for the tracker
Datetime tracking for #1127 already avoids std(); extend the same datetime stats path so Time and Duration values serialize cleanly, and add a Datetime regression test matching the reported failure. Co-authored-by: Cursor <cursoragent@cursor.com>
1 parent da58cad commit 670e8f4

3 files changed

Lines changed: 128 additions & 8 deletions

File tree

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

Lines changed: 15 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919

2020
import polars as pl
2121
from polars.exceptions import InvalidOperationError
22+
2223
from hamilton_sdk.tracking import dataframe_stats as dfs
2324

2425

@@ -130,6 +131,16 @@ def numeric_column_stats(
130131
)
131132

132133

134+
def _temporal_to_jsonable(value: object) -> object:
135+
# datetime.datetime is a subclass of datetime.date. Time/Duration also arrive via
136+
# selectors.temporal() and must be JSON-safe for the UI (same base_data_type="datetime").
137+
if isinstance(value, (datetime.date, datetime.time)):
138+
return value.isoformat()
139+
if isinstance(value, datetime.timedelta):
140+
return str(value)
141+
return value
142+
143+
133144
def datetime_column_stats(
134145
name: str,
135146
position: int,
@@ -144,13 +155,10 @@ def datetime_column_stats(
144155
histogram: dict[str, int],
145156
) -> dfs.DatetimeColumnStatistics:
146157
# TODO: push these conversions into Hamilton functions.
147-
# Note: datetime.datetime is a subclass of datetime.date, so checking datetime.date catches both
148-
min = min.isoformat() if isinstance(min, datetime.date) else min
149-
max = max.isoformat() if isinstance(max, datetime.date) else max
150-
mean = mean.isoformat() if isinstance(mean, datetime.date) else mean
151-
quantiles = {
152-
q: v.isoformat() if isinstance(v, datetime.date) else v for q, v in quantiles.items()
153-
}
158+
min = _temporal_to_jsonable(min)
159+
max = _temporal_to_jsonable(max)
160+
mean = _temporal_to_jsonable(mean)
161+
quantiles = {q: _temporal_to_jsonable(v) for q, v in quantiles.items()}
154162
return dfs.DatetimeColumnStatistics(
155163
name=name,
156164
pos=position,

ui/sdk/tests/tracking/test_polars_col_stats.py

Lines changed: 67 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -17,8 +17,11 @@
1717

1818
"""Module for testing pandas column stats."""
1919

20+
from datetime import date, datetime, time, timedelta
21+
2022
import polars as pl
2123
import pytest
24+
2225
from hamilton_sdk.tracking import polars_col_stats as pcs
2326

2427

@@ -129,3 +132,67 @@ def test_max_string(example_df_string):
129132

130133
def test_str_len(example_df_string):
131134
assert pcs.str_len(example_df_string["a"]).to_list() == [1, 1, 1, 1, 1]
135+
136+
137+
def test_temporal_to_jsonable():
138+
assert pcs._temporal_to_jsonable(datetime(2021, 1, 2, 3, 4, 5)) == "2021-01-02T03:04:05"
139+
assert pcs._temporal_to_jsonable(date(2021, 1, 2)) == "2021-01-02"
140+
assert pcs._temporal_to_jsonable(time(1, 2, 3)) == "01:02:03"
141+
assert pcs._temporal_to_jsonable(timedelta(days=1, hours=2)) == "1 day, 2:00:00"
142+
assert pcs._temporal_to_jsonable(3.14) == 3.14
143+
144+
145+
def test_datetime_column_stats_serializes_temporal_values():
146+
stats = pcs.datetime_column_stats(
147+
name="ts",
148+
position=0,
149+
data_type="Datetime(time_unit='us', time_zone=None)",
150+
count=3,
151+
missing=0,
152+
zeros=0,
153+
min=datetime(2021, 1, 1),
154+
max=datetime(2021, 1, 3),
155+
mean=datetime(2021, 1, 2),
156+
quantiles={0.5: datetime(2021, 1, 2)},
157+
histogram={},
158+
)
159+
assert stats.min == "2021-01-01T00:00:00"
160+
assert stats.max == "2021-01-03T00:00:00"
161+
assert stats.mean == "2021-01-02T00:00:00"
162+
assert stats.std == 0.0
163+
assert stats.quantiles[0.5] == "2021-01-02T00:00:00"
164+
assert stats.base_data_type == "datetime"
165+
166+
time_stats = pcs.datetime_column_stats(
167+
name="t",
168+
position=1,
169+
data_type="Time",
170+
count=2,
171+
missing=0,
172+
zeros=0,
173+
min=time(1, 0),
174+
max=time(2, 0),
175+
mean=time(1, 30),
176+
quantiles={0.5: time(2, 0)},
177+
histogram={},
178+
)
179+
assert time_stats.min == "01:00:00"
180+
assert time_stats.max == "02:00:00"
181+
assert time_stats.mean == "01:30:00"
182+
183+
duration_stats = pcs.datetime_column_stats(
184+
name="d",
185+
position=2,
186+
data_type="Duration(time_unit='us')",
187+
count=2,
188+
missing=0,
189+
zeros=0,
190+
min=timedelta(days=1),
191+
max=timedelta(days=2),
192+
mean=timedelta(days=1, hours=12),
193+
quantiles={0.5: timedelta(days=2)},
194+
histogram={},
195+
)
196+
assert duration_stats.min == "1 day, 0:00:00"
197+
assert duration_stats.max == "2 days, 0:00:00"
198+
assert duration_stats.mean == "1 day, 12:00:00"

ui/sdk/tests/tracking/test_polars_stats.py

Lines changed: 46 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -15,7 +15,8 @@
1515
# specific language governing permissions and limitations
1616
# under the License.
1717

18-
from datetime import date
18+
import json
19+
from datetime import date, datetime, time, timedelta
1920

2021
import polars as pl
2122

@@ -213,3 +214,47 @@ def test_compute_stats_df():
213214
actual["observability_value"][col].pop("quantiles", None)
214215
expected_stats["observability_value"][col].pop("quantiles", None)
215216
assert actual == expected_stats
217+
218+
219+
def test_compute_stats_datetime_series_regression():
220+
# Regression for #1127: Datetime columns must not error via std() and must stay trackable.
221+
series = pl.Series(
222+
"timestamp",
223+
[
224+
datetime(2021, 1, 1),
225+
datetime(2021, 1, 2),
226+
datetime(2021, 1, 3),
227+
],
228+
)
229+
actual = ps.compute_stats_series(series, "df", {})
230+
column_stats = actual["observability_value"]["df"]
231+
assert column_stats["base_data_type"] == "datetime"
232+
assert column_stats["data_type"].startswith("Datetime")
233+
assert column_stats["std"] == 0.0
234+
assert column_stats["min"] == "2021-01-01T00:00:00"
235+
assert column_stats["max"] == "2021-01-03T00:00:00"
236+
assert column_stats["mean"] == "2021-01-02T00:00:00"
237+
json.dumps(actual)
238+
239+
240+
def test_compute_stats_time_and_duration_columns_are_json_safe():
241+
df = pl.DataFrame(
242+
{
243+
"t": pl.Series([time(1, 0), time(2, 0), time(3, 0)]),
244+
"d": pl.Series([timedelta(days=1), timedelta(days=2), timedelta(days=3)]),
245+
}
246+
)
247+
actual = ps.compute_stats_df(df, "test", {})
248+
time_stats = actual["observability_value"]["t"]
249+
duration_stats = actual["observability_value"]["d"]
250+
assert time_stats["base_data_type"] == "datetime"
251+
assert duration_stats["base_data_type"] == "datetime"
252+
assert time_stats["std"] == 0.0
253+
assert duration_stats["std"] == 0.0
254+
assert isinstance(time_stats["min"], str)
255+
assert isinstance(time_stats["max"], str)
256+
assert isinstance(time_stats["mean"], str)
257+
assert isinstance(duration_stats["min"], str)
258+
assert isinstance(duration_stats["max"], str)
259+
assert isinstance(duration_stats["mean"], str)
260+
json.dumps(actual)

0 commit comments

Comments
 (0)