Handle dtypes.NA properly for datetime/timedelta

Author: dcherian

flox/aggregations.py

@@ -158,12 +158,12 @@ def _get_fill_value(dtype, fill_value):
             return np.nan
         # This is madness, but npg checks that fill_value is compatible
         # with array dtype even if the fill_value is never used.
-        elif (
-            np.issubdtype(dtype, np.integer)
-            or np.issubdtype(dtype, np.timedelta64)
-            or np.issubdtype(dtype, np.datetime64)
-        ):
+        elif np.issubdtype(dtype, np.integer):
             return dtypes.get_neg_infinity(dtype, min_for_int=True)
+        elif np.issubdtype(dtype, np.timedelta64):
+            return np.timedelta64("NaT")
+        elif np.issubdtype(dtype, np.datetime64):
+            return np.datetime64("NaT")
         else:
             return None
     return fill_value
@@ -435,9 +435,9 @@ def _std_finalize(sumsq, sum_, count, ddof=0):
 
 
 min_ = Aggregation("min", chunk="min", combine="min", fill_value=dtypes.INF)
-nanmin = Aggregation("nanmin", chunk="nanmin", combine="nanmin", fill_value=np.nan)
+nanmin = Aggregation("nanmin", chunk="nanmin", combine="nanmin", fill_value=dtypes.NA)
 max_ = Aggregation("max", chunk="max", combine="max", fill_value=dtypes.NINF)
-nanmax = Aggregation("nanmax", chunk="nanmax", combine="nanmax", fill_value=np.nan)
+nanmax = Aggregation("nanmax", chunk="nanmax", combine="nanmax", fill_value=dtypes.NA)
 
 
 def argreduce_preprocess(array, axis):
@@ -741,15 +741,15 @@ def scan_binary_op(left_state: ScanState, right_state: ScanState, *, agg: Scan)
     binary_op=None,
     reduction="nanlast",
     scan="ffill",
-    identity=np.nan,
+    identity=dtypes.NA,
     mode="concat_then_scan",
 )
 bfill = Scan(
     "bfill",
     binary_op=None,
     reduction="nanlast",
     scan="ffill",
-    identity=np.nan,
+    identity=dtypes.NA,
     mode="concat_then_scan",
     preprocess=reverse,
     finalize=reverse,

tests/test_properties.py

@@ -218,10 +218,11 @@ def test_first_last(data, array: dask.array.Array, func: str) -> None:
 @settings(report_multiple_bugs=False)
 @given(data=st.data(), array=chunked_arrays())
 def test_topk_max_min(data, array):
-    "top 1 == max; top -1 == min"
+    "top 1 == nanmax; top -1 == nanmin"
     size = array.shape[-1]
+    note(array.compute())
     by = data.draw(by_arrays(shape=(size,)))
-    k, npfunc = data.draw(st.sampled_from([(1, "max"), (-1, "min")]))
+    k, npfunc = data.draw(st.sampled_from([(1, "nanmax"), (-1, "nanmin")]))
 
     for a in (array, array.compute()):
         actual, _ = groupby_reduce(a, by, func="topk", finalize_kwargs={"k": k})