Repository: dask/dask

Issue:
When using PyArrow dtypes, aggregations create NaNs of unexpected type
<!-- Please include a self-contained copy-pastable example that generates the issue if possible.

Please be concise with code posted. See guidelines below on how to provide a good bug report:

- Craft Minimal Bug Reports http://matthewrocklin.com/blog/work/2018/02/28/minimal-bug-reports
- Minimal Complete Verifiable Examples https://stackoverflow.com/help/mcve

Bug reports that follow these guidelines are easier to diagnose, and so are often handled much more quickly.
-->

**Describe the issue**:

Under certain conditions—aggregations of groups that have only missing values, when using PyArrow floats—Dask creates NaNs of an unexpected type. PyArrow types are nullable, so it should create `<NA>` but it creates `NaN`, which are then incorrectly not treated as missing values. 

**Minimal Complete Verifiable Example**:

```python
import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame({
    "x": range(15),
    "y": [pd.NA] * 10 + [1.0] * 5,
    "group": ["a"] * 5 + ["b"] * 5 + ["c"] * 5
})

ddf = dd.from_pandas(df, npartitions=2).astype({
    "y": "double[pyarrow]"
})
print(ddf.compute())
#      x     y group
# 0    0  <NA>     a
# 1    1  <NA>     a
# 2    2  <NA>     a
# 3    3  <NA>     a
# 4    4  <NA>     a
# 5    5  <NA>     b
# 6    6  <NA>     b
# 7    7  <NA>     b
# 8    8  <NA>     b
# 9    9  <NA>     b
# 10  10   1.0     c
# 11  11   1.0     c
# 12  12   1.0     c
# 13  13   1.0     c
# 14  14   1.0     c

ddf = ddf.groupby("group").agg(x=("x", "mean"), y=("y", "mean"))
print(ddf.compute())
#           x    y
# group           
# a       2.0  NaN
# b       7.0  NaN
# c      12.0  1.0
# (Note that groups with <NA> resulted in NaN rather than <NA>)

# These NaNs are then not recognized as missing values
print(ddf["y"].isnull().sum().compute())
# 0
# (While the expected result is 2).

# This breaks other functions, example further aggregations in Pandas
print(ddf["y"].compute().mean(skipna=True))
# nan
# (While the exected result is 1.0).

# And also further aggreations in Dask, albeit with a different unexpected result
print(ddf["y"].mean(skipna=True).compute())
# 0.0
# (While the exected result is 1.0)
```

**Anything else we need to know?**:

If we use Numpy dtypes, i.e. convert `y` to `float64` rather than `double[pyarrow]` everything works as expected. 



**Environment**:

- Dask version: 2024.4.1
- Python version: 3.10.13
- Operating System: MacOS on the client, dask running inside Docker based on ghcr.io/dask/dask image on the cluster. 
- Install method (conda, pip, source): pip on the client, using ghcr.io/dask/dask as the base image on the cluster. 


Discussion:

Patch:
diff --git a/dask/dataframe/groupby.py b/dask/dataframe/groupby.py
index a1350694f..5e6c23a9d 100644
--- a/dask/dataframe/groupby.py
+++ b/dask/dataframe/groupby.py
@@ -1227,10 +1227,23 @@ def _apply_func_to_columns(df_like, prefix, func):
 
 
 def _finalize_mean(df, sum_column, count_column):
-    return df[sum_column] / df[count_column]
+    result = df[sum_column] / df[count_column]
+    return _adjust_for_arrow_na(result, df[count_column])
 
 
-def _finalize_var(df, count_column, sum_column, sum2_column, **kwargs):
+def _adjust_for_arrow_na(result, df, check_for_isna=False):
+    if isinstance(result.dtype, pd.ArrowDtype):
+        # Our mean computation results in np.nan here but pandas doesn't
+        if check_for_isna:
+            result[df.isna()] = pd.NA
+        else:
+            result[df == 0] = pd.NA
+    return result
+
+
+def _finalize_var(
+    df, count_column, sum_column, sum2_column, adjust_arrow=True, **kwargs
+):
     # arguments are being checked when building the finalizer. As of this moment,
     # we're only using ddof, and raising an error on other keyword args.
     ddof = kwargs.get("ddof", 1)
@@ -1243,13 +1256,20 @@ def _finalize_var(df, count_column, sum_column, sum2_column, **kwargs):
     div[div < 0] = 0
     result /= div
     result[(n - ddof) == 0] = np.nan
-
-    return result
+    if adjust_arrow:
+        return _adjust_for_arrow_na(result, div)
+    else:
+        return result
 
 
 def _finalize_std(df, count_column, sum_column, sum2_column, **kwargs):
-    result = _finalize_var(df, count_column, sum_column, sum2_column, **kwargs)
-    return np.sqrt(result)
+    result = _finalize_var(
+        df, count_column, sum_column, sum2_column, adjust_arrow=False, **kwargs
+    )
+    res = np.sqrt(result)
+    if res.dtype != result.dtype:
+        res = res.astype(result.dtype)
+    return _adjust_for_arrow_na(res, result, check_for_isna=True)
 
 
 def _cum_agg_aligned(part, cum_last, index, columns, func, initial):


Test patch:
diff --git a/dask/dataframe/tests/test_groupby.py b/dask/dataframe/tests/test_groupby.py
index 11aaeecff..df5c0b45e 100644
--- a/dask/dataframe/tests/test_groupby.py
+++ b/dask/dataframe/tests/test_groupby.py
@@ -3755,3 +3755,29 @@ def test_groupby_value_counts_all_na_partitions():
         ddf.groupby("A")["B"].value_counts(),
         df.groupby("A")["B"].value_counts(),
     )
+
+
+def test_agg_pyarrow_casts():
+    pytest.importorskip("pyarrow")
+    df = pd.DataFrame(
+        {
+            "x": range(15),
+            "y": pd.Series([pd.NA] * 10 + [1.0] * 5, dtype="double[pyarrow]"),
+            "group": ["a"] * 5 + ["b"] * 5 + ["c"] * 5,
+        }
+    )
+
+    ddf = dd.from_pandas(df, npartitions=2)
+    if PANDAS_GE_220:
+        # np.sqrt doesn't work before 2.2
+        additional_aggs = {"z": ("y", "std")}
+    else:
+        additional_aggs = {}
+
+    result = ddf.groupby("group").agg(
+        x=("y", "var"), y=("y", "mean"), **additional_aggs
+    )
+    expected = df.groupby("group").agg(
+        x=("y", "var"), y=("y", "mean"), **additional_aggs
+    )
+    assert_eq(result, expected)