Repository: dask/dask

Issue:
Column assignments fails with Float64Dtype type
**What happened**:
When using pandas' new Float64 nullable type (with pandas >= 1.2), column assignment fails with 
```
TypeError: Cannot interpret 'Float64Dtype()' as a data type
```

**Minimal Complete Verifiable Example**:

Fails at least with pandas version 1.2.0 (below that, the new extension type was not introduced).

```python
import pandas as pd
import dask.dataframe as dd

# some example data. Important is only the Float64, the new pandas extension type
df = dd.from_pandas(pd.DataFrame({"a": [1.1]}, dtype="Float64"), npartitions=1)

df.assign(new_col=df["a"])
# TypeError: Cannot interpret 'Float64Dtype()' as a data type
```

Full stacktrace
```
Traceback (most recent call last):
  File "bla.py", line 7, in <module>
    df.assign(new_col=df["a"])
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/core.py", line 4004, in assign
    df2 = self._meta_nonempty.assign(**_extract_meta(kwargs, nonempty=True))
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/core.py", line 360, in _meta_nonempty
    return meta_nonempty(self._meta)
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/utils.py", line 509, in __call__
    return meth(arg, *args, **kwargs)
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/utils.py", line 391, in meta_nonempty_dataframe
    dt_s_dict[dt] = _nonempty_series(x.iloc[:, i], idx=idx)
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/utils.py", line 578, in _nonempty_series
    data = np.array([entry, entry], dtype=dtype)
TypeError: Cannot interpret 'Float64Dtype()' as a data type
```

As a fix, I guess one wants to introduce a function similar to https://github.com/dask/dask/blob/master/dask/dataframe/utils.py#L43 (but for floats) and use it in https://github.com/dask/dask/blob/master/dask/dataframe/utils.py#L540. I am happy to open a PR to fix this if this solution is ok.

**Environment**:

- Dask version: 2021.1.1
- Pandas version: 1.2.0
- Python version:
- Operating System: ubuntu
- Install method (conda, pip, source): conda

Column assignments fails with Float64Dtype type
**What happened**:
When using pandas' new Float64 nullable type (with pandas >= 1.2), column assignment fails with 
```
TypeError: Cannot interpret 'Float64Dtype()' as a data type
```

**Minimal Complete Verifiable Example**:

Fails at least with pandas version 1.2.0 (below that, the new extension type was not introduced).

```python
import pandas as pd
import dask.dataframe as dd

# some example data. Important is only the Float64, the new pandas extension type
df = dd.from_pandas(pd.DataFrame({"a": [1.1]}, dtype="Float64"), npartitions=1)

df.assign(new_col=df["a"])
# TypeError: Cannot interpret 'Float64Dtype()' as a data type
```

Full stacktrace
```
Traceback (most recent call last):
  File "bla.py", line 7, in <module>
    df.assign(new_col=df["a"])
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/core.py", line 4004, in assign
    df2 = self._meta_nonempty.assign(**_extract_meta(kwargs, nonempty=True))
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/core.py", line 360, in _meta_nonempty
    return meta_nonempty(self._meta)
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/utils.py", line 509, in __call__
    return meth(arg, *args, **kwargs)
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/utils.py", line 391, in meta_nonempty_dataframe
    dt_s_dict[dt] = _nonempty_series(x.iloc[:, i], idx=idx)
  File "/home/nils/anaconda3/envs/dask-sql/lib/python3.8/site-packages/dask/dataframe/utils.py", line 578, in _nonempty_series
    data = np.array([entry, entry], dtype=dtype)
TypeError: Cannot interpret 'Float64Dtype()' as a data type
```

As a fix, I guess one wants to introduce a function similar to https://github.com/dask/dask/blob/master/dask/dataframe/utils.py#L43 (but for floats) and use it in https://github.com/dask/dask/blob/master/dask/dataframe/utils.py#L540. I am happy to open a PR to fix this if this solution is ok.

**Environment**:

- Dask version: 2021.1.1
- Pandas version: 1.2.0
- Python version:
- Operating System: ubuntu
- Install method (conda, pip, source): conda



Discussion:
I think the solution that you propose is very reasonable and if you are willing to open a PR even better! Just as a heads up there is a release planned for friday (tomorrow) so it'd be great to get this in.
I think the solution that you propose is very reasonable and if you are willing to open a PR even better! Just as a heads up there is a release planned for friday (tomorrow) so it'd be great to get this in.

Resolving patch:
diff --git a/dask/dataframe/utils.py b/dask/dataframe/utils.py
--- a/dask/dataframe/utils.py
+++ b/dask/dataframe/utils.py
@@ -23,6 +23,7 @@
 from ._compat import (  # noqa: F401
     PANDAS_GT_100,
     PANDAS_GT_110,
+    PANDAS_GT_120,
     tm,
 )
 
@@ -55,6 +56,18 @@ def is_integer_na_dtype(t):
     return isinstance(dtype, types)
 
 
+def is_float_na_dtype(t):
+    if not PANDAS_GT_120:
+        return False
+
+    dtype = getattr(t, "dtype", t)
+    types = (
+        pd.Float32Dtype,
+        pd.Float64Dtype,
+    )
+    return isinstance(dtype, types)
+
+
 def shard_df_on_index(df, divisions):
     """Shard a DataFrame by ranges on its index
 
@@ -558,6 +571,8 @@ def _nonempty_series(s, idx=None):
         data = pd.Categorical(data, categories=cats, ordered=s.cat.ordered)
     elif is_integer_na_dtype(dtype):
         data = pd.array([1, None], dtype=dtype)
+    elif is_float_na_dtype(dtype):
+        data = pd.array([1.0, None], dtype=dtype)
     elif is_period_dtype(dtype):
         # pandas 0.24.0+ should infer this to be Series[Period[freq]]
         freq = dtype.freq


Test patch:
diff --git a/dask/dataframe/tests/test_dataframe.py b/dask/dataframe/tests/test_dataframe.py
--- a/dask/dataframe/tests/test_dataframe.py
+++ b/dask/dataframe/tests/test_dataframe.py
@@ -4468,3 +4468,17 @@ def test_join_series():
     expected_df = dd.from_pandas(df.join(df["x"], lsuffix="_"), npartitions=1)
     actual_df = ddf.join(ddf["x"], lsuffix="_")
     assert_eq(actual_df, expected_df)
+
+
+@pytest.mark.skipif(
+    not dd._compat.PANDAS_GT_120, reason="Float64 was introduced in pandas>=1.2"
+)
+def test_assign_na_float_columns():
+    # See https://github.com/dask/dask/issues/7156
+    df_pandas = pd.DataFrame({"a": [1.1]}, dtype="Float64")
+    df = dd.from_pandas(df_pandas, npartitions=1)
+
+    df = df.assign(new_col=df["a"])
+
+    assert df.compute()["a"].dtypes == "Float64"
+    assert df.compute()["new_col"].dtypes == "Float64"
diff --git a/dask/dataframe/tests/test_utils_dataframe.py b/dask/dataframe/tests/test_utils_dataframe.py
--- a/dask/dataframe/tests/test_utils_dataframe.py
+++ b/dask/dataframe/tests/test_utils_dataframe.py
@@ -17,6 +17,7 @@
     is_series_like,
     is_index_like,
     PANDAS_GT_100,
+    PANDAS_GT_120,
 )
 
 import pytest
@@ -470,3 +471,10 @@ def test_nonempty_series_sparse():
         dd.utils._nonempty_series(ser)
 
     assert len(w) == 0
+
+
+@pytest.mark.skipif(not PANDAS_GT_120, reason="Float64 was introduced in pandas>=1.2")
+def test_nonempty_series_nullable_float():
+    ser = pd.Series([], dtype="Float64")
+    non_empty = dd.utils._nonempty_series(ser)
+    assert non_empty.dtype == "Float64"