Repository: dask/dask

Issue:
Pandas 1.2.0 compatibility - `std()` for the `datetime64[ns]` dtype
This doesn't work yet in dask because here the implementation of `std` generally depends on `var`, which doesn't work for datetimes. 

Some tests are skipped for this in `dask/dataframe/tests/test_arithmetics_reduction.py`


Discussion:
I thought I'd split up #7100 to make it easier to keep track of the different parts. 
Does this seem like an okay **first issue**? I'd like to take this on if so! I just don't want to start an issue someone else might be able to do much quicker if I get stuck. Thanks!
It's a pretty good first issue in that you probably don't need a ton of dask-specific knowledge. But it might be a good bit of work. Please feel free to take it on and ping me if you get stuck.
I'll give it a shot then! Thanks for such a quick response.

Resolving patch:
diff --git a/dask/dataframe/core.py b/dask/dataframe/core.py
--- a/dask/dataframe/core.py
+++ b/dask/dataframe/core.py
@@ -2178,11 +2178,29 @@ def std(
     ):
         axis = self._validate_axis(axis)
         _raise_if_object_series(self, "std")
+        _raise_if_not_series_or_dataframe(self, "std")
+
         meta = self._meta_nonempty.std(axis=axis, skipna=skipna)
+        is_df_like = is_dataframe_like(self._meta)
+        needs_time_conversion = False
+        numeric_dd = self
+
+        if PANDAS_GT_120 and is_df_like:
+            time_cols = self._meta.select_dtypes(include="datetime").columns
+            if len(time_cols) > 0:
+                (
+                    numeric_dd,
+                    needs_time_conversion,
+                ) = self._convert_time_cols_to_numeric(time_cols, axis, meta, skipna)
+        elif PANDAS_GT_120 and not is_df_like:
+            needs_time_conversion = is_datetime64_any_dtype(self._meta)
+            if needs_time_conversion:
+                numeric_dd = _convert_to_numeric(self, skipna)
+
         if axis == 1:
             result = map_partitions(
-                M.std,
-                self,
+                M.std if not needs_time_conversion else _sqrt_and_convert_to_timedelta,
+                numeric_dd,
                 meta=meta,
                 token=self._token_prefix + "std",
                 axis=axis,
@@ -2192,18 +2210,67 @@ def std(
                 parent_meta=self._meta,
             )
             return handle_out(out, result)
+
+        # Case where axis=0 or axis=None
+        v = numeric_dd.var(skipna=skipna, ddof=ddof, split_every=split_every)
+        name = self._token_prefix + "std"
+
+        if needs_time_conversion:
+            sqrt_func_kwargs = {
+                "is_df_like": is_df_like,
+                "time_cols": time_cols if is_df_like else None,
+                "axis": axis,
+            }
+            sqrt_func = _sqrt_and_convert_to_timedelta
         else:
-            v = self.var(skipna=skipna, ddof=ddof, split_every=split_every)
-            name = self._token_prefix + "std"
-            result = map_partitions(
-                np.sqrt,
-                v,
-                meta=meta,
-                token=name,
-                enforce_metadata=False,
-                parent_meta=self._meta,
+            sqrt_func_kwargs = {}
+            sqrt_func = np.sqrt
+
+        result = map_partitions(
+            sqrt_func,
+            v,
+            meta=meta,
+            token=name,
+            enforce_metadata=False,
+            parent_meta=self._meta,
+            **sqrt_func_kwargs,
+        )
+
+        # Try to match the Pandas result dtype
+        if is_df_like and hasattr(meta, "dtype"):
+            result = result.astype(meta.dtype)
+
+        return handle_out(out, result)
+
+    def _convert_time_cols_to_numeric(self, time_cols, axis, meta, skipna):
+        from .io import from_pandas
+
+        needs_time_conversion = True
+
+        # Ensure all columns are correct type. Need to shallow copy since cols will be modified
+        if axis == 0:
+            numeric_dd = self[meta.index].copy()
+        else:
+            numeric_dd = self.copy()
+
+        # Mix of datetimes with other numeric types produces NaNs for each value in std() series
+        if axis == 1 and len(time_cols) != len(self.columns):
+            # This is faster than converting each column to numeric when it's not necessary
+            # since each standard deviation will just be NaN
+            needs_time_conversion = False
+            numeric_dd = from_pandas(
+                pd.DataFrame(
+                    {"_": pd.Series([np.nan])},
+                    index=self.index,
+                ),
+                npartitions=self.npartitions,
             )
-            return handle_out(out, result)
+        else:
+            # Convert timedelta and datetime columns to integer types so we can use var
+            for col in time_cols:
+                numeric_dd[col] = _convert_to_numeric(numeric_dd[col], skipna)
+
+        return numeric_dd, needs_time_conversion
 
     @_numeric_only
     @derived_from(pd.DataFrame)
@@ -7494,3 +7561,41 @@ def series_map(base_series, map_series):
     divisions = list(base_series.divisions)
 
     return new_dd_object(graph, final_prefix, meta, divisions)
+
+
+def _convert_to_numeric(series, skipna):
+    if skipna:
+        return series.dropna().view("i8")
+
+    # series.view("i8") with pd.NaT produces -9223372036854775808 is why we need to do this
+    return series.view("i8").mask(series.isnull(), np.nan)
+
+
+def _sqrt_and_convert_to_timedelta(partition, axis, *args, **kwargs):
+    if axis == 1:
+        return pd.to_timedelta(M.std(partition, axis=axis, *args, **kwargs))
+
+    is_df_like, time_cols = kwargs["is_df_like"], kwargs["time_cols"]
+
+    sqrt = np.sqrt(partition)
+
+    if not is_df_like:
+        return pd.to_timedelta(sqrt)
+
+    time_col_mask = sqrt.index.isin(time_cols)
+    matching_vals = sqrt[time_col_mask]
+    for time_col, matching_val in zip(time_cols, matching_vals):
+        sqrt[time_col] = pd.to_timedelta(matching_val)
+
+    return sqrt
+
+
+def _raise_if_not_series_or_dataframe(x, funcname):
+    """
+    Utility function to raise an error if an object is not a Series or DataFrame
+    """
+    if not is_series_like(x) and not is_dataframe_like(x):
+        raise NotImplementedError(
+            "`%s` is only supported with objects that are Dataframes or Series"
+            % funcname
+        )


Test patch:
diff --git a/dask/dataframe/tests/test_arithmetics_reduction.py b/dask/dataframe/tests/test_arithmetics_reduction.py
--- a/dask/dataframe/tests/test_arithmetics_reduction.py
+++ b/dask/dataframe/tests/test_arithmetics_reduction.py
@@ -4,6 +4,7 @@
 import numpy as np
 import pandas as pd
 import pytest
+from pandas.api.types import is_scalar
 
 import dask.dataframe as dd
 from dask.dataframe._compat import PANDAS_GT_120, PANDAS_VERSION
@@ -1184,33 +1185,20 @@ def test_reductions_frame_dtypes():
     assert_eq(df.min(), ddf.min())
     assert_eq(df.max(), ddf.max())
     assert_eq(df.count(), ddf.count())
-    if PANDAS_GT_120:
-        # std is implemented for datetimes in pandas 1.2.0, but dask
-        # implementation depends on var which isn't
-        assert_eq(
-            df_no_timedelta.drop("dt", axis=1).std(),
-            ddf_no_timedelta.drop("dt", axis=1).std(),
-        )
-    else:
-        assert_eq(df_no_timedelta.std(), ddf_no_timedelta.std())
-    assert_eq(df_no_timedelta.var(), ddf_no_timedelta.var())
-
-    df2 = df.drop("timedelta", axis=1)
-    ddf2 = ddf.drop("timedelta", axis=1)
-
-    assert_eq(df2.var(skipna=False), ddf2.var(skipna=False))
     assert_eq(df.sem(), ddf.sem())
-    if PANDAS_GT_120:
-        assert_eq(
-            df_no_timedelta.drop("dt", axis=1).std(ddof=0),
-            ddf_no_timedelta.drop("dt", axis=1).std(ddof=0),
-        )
-    else:
-        assert_eq(df_no_timedelta.std(ddof=0), ddf_no_timedelta.std(ddof=0))
-    assert_eq(df2.var(ddof=0), ddf2.var(ddof=0))
-    assert_eq(df2.var(ddof=0, skipna=False), ddf2.var(ddof=0, skipna=False))
     assert_eq(df.sem(ddof=0), ddf.sem(ddof=0))
 
+    assert_eq(df_no_timedelta.std(), ddf_no_timedelta.std())
+    assert_eq(df_no_timedelta.std(skipna=False), ddf_no_timedelta.std(skipna=False))
+    assert_eq(df_no_timedelta.std(ddof=0), ddf_no_timedelta.std(ddof=0))
+    assert_eq(df_no_timedelta.var(), ddf_no_timedelta.var())
+    assert_eq(df_no_timedelta.var(skipna=False), ddf_no_timedelta.var(skipna=False))
+    assert_eq(df_no_timedelta.var(ddof=0), ddf_no_timedelta.var(ddof=0))
+    assert_eq(
+        df_no_timedelta.var(ddof=0, skipna=False),
+        ddf_no_timedelta.var(ddof=0, skipna=False),
+    )
+
     assert_eq(df._get_numeric_data(), ddf._get_numeric_data())
 
     numerics = ddf[["int", "float"]]
@@ -1477,3 +1465,154 @@ def test_series_agg_with_min_count(method, min_count):
         assert result == 1
     else:
         assert result is np.nan
+
+
+# Default absolute tolerance of 2000 nanoseconds
+def assert_near_timedeltas(t1, t2, atol=2000):
+    if is_scalar(t1):
+        t1 = pd.Series([t1])
+    if is_scalar(t2):
+        t2 = pd.Series([t2])
+
+    assert t1.dtype == t2.dtype
+    assert_eq(pd.to_numeric(t1), pd.to_numeric(t2), atol=atol)
+
+
+@pytest.mark.skipif(
+    not PANDAS_GT_120, reason="std() for datetime only added in pandas>=1.2"
+)
+@pytest.mark.parametrize("axis", [0, 1])
+def test_datetime_std_creates_copy_cols(axis):
+    pdf = pd.DataFrame(
+        {
+            "dt1": [
+                datetime.fromtimestamp(1636426700 + (i * 250000)) for i in range(10)
+            ],
+            "dt2": [
+                datetime.fromtimestamp(1636426700 + (i * 300000)) for i in range(10)
+            ],
+        }
+    )
+
+    ddf = dd.from_pandas(pdf, 3)
+
+    # Series test (same line twice to make sure data structure wasn't mutated)
+    assert_eq(ddf["dt1"].std(), pdf["dt1"].std())
+    assert_eq(ddf["dt1"].std(), pdf["dt1"].std())
+
+    # DataFrame test (same line twice to make sure data structure wasn't mutated)
+    assert_near_timedeltas(ddf.std(axis=axis).compute(), pdf.std(axis=axis))
+    assert_near_timedeltas(ddf.std(axis=axis).compute(), pdf.std(axis=axis))
+
+
+@pytest.mark.skipif(
+    not PANDAS_GT_120, reason="std() for datetime only added in pandas>=1.2"
+)
+@pytest.mark.parametrize("axis", [0, 1])
+@pytest.mark.parametrize("skipna", [False, True])
+def test_datetime_std_with_larger_dataset(axis, skipna):
+    num_rows = 250
+
+    dt1 = pd.concat(
+        [
+            pd.Series([pd.NaT] * 15, index=range(15)),
+            pd.to_datetime(
+                pd.Series(
+                    [
+                        datetime.fromtimestamp(1636426704 + (i * 250000))
+                        for i in range(num_rows - 15)
+                    ],
+                    index=range(15, 250),
+                )
+            ),
+        ],
+        ignore_index=False,
+    )
+
+    base_numbers = [
+        (1638290040706793300 + (i * 69527182702409)) for i in range(num_rows)
+    ]
+
+    pdf = pd.DataFrame(
+        {"dt1": dt1, "dt2": pd.to_datetime(pd.Series(base_numbers))}, index=range(250)
+    )
+
+    for i in range(3, 8):
+        pdf[f"dt{i}"] = pd.to_datetime(
+            pd.Series([int(x + (0.12 * i)) for x in base_numbers])
+        )
+
+    ddf = dd.from_pandas(pdf, 8)
+
+    assert_near_timedeltas(
+        ddf[["dt1"]].std(axis=axis, skipna=skipna).compute(),
+        pdf[["dt1"]].std(axis=axis, skipna=skipna),
+    )
+
+    # Same thing but as Series. No axis, since axis=1 raises error
+    assert_near_timedeltas(
+        ddf["dt1"].std(skipna=skipna).compute(), pdf["dt1"].std(skipna=skipna)
+    )
+
+    # Computation on full dataset
+    assert_near_timedeltas(
+        ddf.std(axis=axis, skipna=skipna).compute(), pdf.std(axis=axis, skipna=skipna)
+    )
+
+
+@pytest.mark.skipif(
+    not PANDAS_GT_120, reason="std() for datetime only added in pandas>=1.2"
+)
+@pytest.mark.filterwarnings(
+    "ignore:Dropping of nuisance columns:FutureWarning"
+)  # https://github.com/dask/dask/issues/7714
+@pytest.mark.parametrize("skipna", [False, True])
+def test_datetime_std_across_axis1_null_results(skipna):
+    pdf = pd.DataFrame(
+        {
+            "dt1": [
+                datetime.fromtimestamp(1636426704 + (i * 250000)) for i in range(10)
+            ],
+            "dt2": [
+                datetime.fromtimestamp(1636426704 + (i * 217790)) for i in range(10)
+            ],
+            "nums": [i for i in range(10)],
+        }
+    )
+
+    ddf = dd.from_pandas(pdf, 3)
+
+    # Single column always results in NaT
+    assert_eq(
+        ddf[["dt1"]].std(axis=1, skipna=skipna), pdf[["dt1"]].std(axis=1, skipna=skipna)
+    )
+
+    # Mix of datetimes with other numeric types produces NaNs
+    assert_eq(ddf.std(axis=1, skipna=skipna), pdf.std(axis=1, skipna=skipna))
+
+    # Test with mix of na and truthy datetimes
+    pdf2 = pd.DataFrame(
+        {
+            "dt1": [pd.NaT]
+            + [datetime.fromtimestamp(1636426704 + (i * 250000)) for i in range(10)]
+            + [pd.NaT],
+            "dt2": [
+                datetime.fromtimestamp(1636426704 + (i * 250000)) for i in range(12)
+            ],
+            "dt3": [
+                datetime.fromtimestamp(1636426704 + (i * 282616)) for i in range(12)
+            ],
+        }
+    )
+
+    ddf2 = dd.from_pandas(pdf2, 3)
+
+    assert_eq(ddf2.std(axis=1, skipna=skipna), pdf2.std(axis=1, skipna=skipna))
+
+
+def test_std_raises_on_index():
+    with pytest.raises(
+        NotImplementedError,
+        match="`std` is only supported with objects that are Dataframes or Series",
+    ):
+        dd.from_pandas(pd.DataFrame({"test": [1, 2]}), npartitions=2).index.std()