fix: Normalize Dask timestamps to UTC without a row-wise apply

_normalize_timestamp made timestamp columns tz-aware with a row-wise apply declaring meta="datetime64[ns, UTC]". meta only declares the dtype; the partition gets whatever the lambda returns. With zero rows the lambda never runs, so the computed column stayed tz-naive while meta claimed UTC, and the tz-aware comparison in _filter_ttl then raised: TypeError: Invalid comparison between dtype=datetime64[ns] and DatetimeArray A zero-row entity_df is a normal case in batch scoring, when the upstream query matched nothing for that run, so get_historical_features crashed rather than returning an empty result. A column carrying a non-UTC timezone diverged from the declared meta the same way, because the lambda returned any value that already had tzinfo untouched. Replace both applies with dd.to_datetime(..., utc=True), which localizes tz-naive values, converts tz-aware ones, and yields the correct dtype for an empty frame. It is also vectorized, so it avoids a Python-level call per row on non-empty frames. Signed-off-by: Daksha1611 <mehtadaksha1611@gmail.com>
Daksha1611 · Sep 6, 2026 · Sep 11, 2026 · Oct 6, 2026 · Sep 6, 2026 · 042397e7b2ace8052c5a7b641f9d651dfcb06b77
commit 042397e7b2ace8052c5a7b641f9d651dfcb06b77
@@ -1136,6 +1136,22 @@ def _merge(
     return df_to_join
 
 
+def _to_utc(series: dd.Series) -> dd.Series:
+    """Return *series* as tz-aware UTC, correctly for an empty partition.
+
+    The row-wise ``apply`` this replaces declared ``meta="datetime64[ns, UTC]"``
+    but produced whatever the lambda returned. With zero rows the lambda never
+    runs, so the computed partition stayed tz-naive while meta claimed
+    otherwise, and the later tz-aware comparison in ``_filter_ttl`` raised
+    ``TypeError: Invalid comparison``. A column carrying a non-UTC tz hit the
+    same mismatch, since the lambda passed those values through untouched.
+
+    ``to_datetime(..., utc=True)`` localizes tz-naive values, converts tz-aware
+    ones, and yields the right dtype for an empty frame.
+    """
+    return dd.to_datetime(series, utc=True)
+
+
 def _normalize_timestamp(
     df_to_join: dd.DataFrame,
     timestamp_field: str,
@@ -1155,10 +1171,7 @@ def _normalize_timestamp(
             df_to_join = df_to_join.drop(columns=dups)
 
         # Make sure all timestamp fields are tz-aware. We default tz-naive fields to UTC
-        df_to_join[timestamp_field] = df_to_join[timestamp_field].apply(
-            lambda x: x if x.tzinfo else x.replace(tzinfo=timezone.utc),
-            meta=(timestamp_field, "datetime64[ns, UTC]"),
-        )
+        df_to_join[timestamp_field] = _to_utc(df_to_join[timestamp_field])
 
     # TODO: need to figure out why the value of created_timestamp_column_type.tz is pytz.UTC
     if created_timestamp_column and (
@@ -1170,11 +1183,8 @@ def _normalize_timestamp(
             df_to_join, dups = _df_column_uniquify(df_to_join)
             df_to_join = df_to_join.drop(columns=dups)
 
-        df_to_join[created_timestamp_column] = df_to_join[
-            created_timestamp_column
-        ].apply(
-            lambda x: x if x.tzinfo else x.replace(tzinfo=timezone.utc),
-            meta=(timestamp_field, "datetime64[ns, UTC]"),
+        df_to_join[created_timestamp_column] = _to_utc(
+            df_to_join[created_timestamp_column]
         )
 
     return df_to_join.persist()

@@ -0,0 +1,58 @@
+"""Timestamp normalization must produce tz-aware UTC even for an empty frame.
+
+A zero-row ``entity_df`` is a normal degenerate case in batch scoring, when the
+upstream query matched nothing for that run.
+"""
+
+from datetime import datetime, timedelta
+from unittest.mock import MagicMock
+
+import dask.dataframe as dd
+import pandas as pd
+import pytest
+
+from feast.infra.offline_stores.dask import _filter_ttl, _normalize_timestamp
+
+ENTITY_TS = "entity_timestamp"
+EVENT_TS = "event_timestamp"
+
+
+def _frame(n: int, tz: str | None = None) -> dd.DataFrame:
+    stamps = pd.to_datetime([datetime(2026, 2, 1)] * n)
+    if tz is not None:
+        stamps = stamps.tz_localize(tz)
+    return dd.from_pandas(
+        pd.DataFrame(
+            {
+                EVENT_TS: stamps,
+                ENTITY_TS: pd.to_datetime([datetime(2026, 2, 1)] * n, utc=True),
+                "conv_rate": [0.5] * n,
+            }
+        ),
+        npartitions=1,
+    )
+
+
+@pytest.mark.parametrize("rows", [0, 1])
+def test_normalize_timestamp_is_utc_aware_regardless_of_row_count(rows):
+    normalized = _normalize_timestamp(_frame(rows), EVENT_TS).compute()
+    assert isinstance(normalized[EVENT_TS].dtype, pd.DatetimeTZDtype)
+    assert str(normalized[EVENT_TS].dtype.tz) == "UTC"
+
+
+@pytest.mark.parametrize("rows", [0, 1])
+def test_filter_ttl_on_empty_frame_does_not_raise(rows):
+    """The tz-naive/tz-aware mismatch used to surface here as a TypeError."""
+    feature_view = MagicMock()
+    feature_view.ttl = timedelta(days=3650)
+
+    normalized = _normalize_timestamp(_frame(rows), EVENT_TS)
+    result = _filter_ttl(normalized, feature_view, ENTITY_TS, EVENT_TS).compute()
+
+    assert len(result) == rows
+
+
+def test_non_utc_timezone_is_converted_not_passed_through():
+    """A non-UTC column also diverged from the declared UTC meta."""
+    normalized = _normalize_timestamp(_frame(1, tz="US/Eastern"), EVENT_TS).compute()
+    assert str(normalized[EVENT_TS].dtype.tz) == "UTC"