(fix): no fill_value on reindex

properties/test_pandas_roundtrip.py

@@ -134,10 +134,21 @@ def test_roundtrip_pandas_dataframe_datetime(df) -> None:
     xr.testing.assert_identical(dataset, roundtripped.to_xarray())
 
 
-def test_roundtrip_1d_pandas_extension_array() -> None:
-    df = pd.DataFrame({"cat": pd.Categorical(["a", "b", "c"])})
-    arr = xr.Dataset.from_dataframe(df)["cat"]
+@pytest.mark.parametrize(
+    "extension_array",
+    [
+        pd.Categorical(["a", "b", "c"]),
+        pd.array([1, 2, 3], dtype="int64"),
+        pd.array(["a", "b", "c"], dtype="string"),
+        pd.arrays.IntervalArray(
+            [pd.Interval(0, 1), pd.Interval(1, 5), pd.Interval(2, 6)]
+        ),
+    ],
+)
+def test_roundtrip_1d_pandas_extension_array(extension_array) -> None:
+    df = pd.DataFrame({"arr": extension_array})
+    arr = xr.Dataset.from_dataframe(df)["arr"]
     roundtripped = arr.to_pandas()
-    assert (df["cat"] == roundtripped).all()
-    assert df["cat"].dtype == roundtripped.dtype
+    assert (df["arr"] == roundtripped).all()
+    assert df["arr"].dtype == roundtripped.dtype
     xr.testing.assert_identical(arr, roundtripped.to_xarray())

xarray/core/dtypes.py

@@ -4,6 +4,7 @@
 from typing import Any
 
 import numpy as np
+import pandas as pd
 from pandas.api.types import is_extension_array_dtype
 
 from xarray.compat import array_api_compat, npcompat
@@ -63,7 +64,9 @@ def maybe_promote(dtype: np.dtype) -> tuple[np.dtype, Any]:
     # N.B. these casting rules should match pandas
     dtype_: np.typing.DTypeLike
     fill_value: Any
-    if HAS_STRING_DTYPE and np.issubdtype(dtype, np.dtypes.StringDType()):
+    if pd.api.types.is_extension_array_dtype(dtype):
+        return dtype, pd.NA
+    elif HAS_STRING_DTYPE and np.issubdtype(dtype, np.dtypes.StringDType()):
         # for now, we always promote string dtypes to object for consistency with existing behavior
         # TODO: refactor this once we have a better way to handle numpy vlen-string dtypes
         dtype_ = object

xarray/core/duck_array_ops.py

@@ -251,7 +251,7 @@
 
     if xp == np:
         # numpy currently doesn't have a astype:
         return data.astype(dtype, **kwargs)
     return xp.astype(data, dtype, **kwargs)
 
 
@@ -273,18 +273,35 @@
         extension_array_types = [
             x.dtype for x in scalars_or_arrays if is_extension_array_dtype(x)
         ]
-        non_nans = [x for x in scalars_or_arrays if not isna(x)]
-        if len(extension_array_types) == len(non_nans) and all(
+        non_nans_or_scalar = [
+            x for x in scalars_or_arrays if not (isna(x) or np.isscalar(x))
+        ]
+        if len(extension_array_types) == len(non_nans_or_scalar) and all(
             isinstance(x, type(extension_array_types[0])) for x in extension_array_types
         ):
-            return [
+            # Get the extension array class of the first element, guaranteed to be the same
+            # as the others thanks to the anove check.
+            extension_array_class = type(
+                non_nans_or_scalar[0].array
+                if isinstance(non_nans_or_scalar[0], PandasExtensionArray)
+                else non_nans_or_scalar[0]
+            )
+            # Cast scalars/nans to extension array class
+            arrays_with_nan_to_sequence = [
                 x
-                if not isna(x)
-                else PandasExtensionArray(
-                    type(non_nans[0].array)._from_sequence([x], dtype=non_nans[0].dtype)
+                if not (isna(x) or np.isscalar(x))
+                else extension_array_class._from_sequence(
+                    [x], dtype=non_nans_or_scalar[0].dtype
                 )
                 for x in scalars_or_arrays
             ]
+            # Wrap the output if necessary
+            return [
+                PandasExtensionArray(x)
+                if not isinstance(x, PandasExtensionArray)
+                else x
+                for x in arrays_with_nan_to_sequence
+            ]
         raise ValueError(
             f"Cannot cast values to shared type, found values: {scalars_or_arrays}"
         )
@@ -407,7 +424,6 @@
         condition = asarray(condition, dtype=dtype, xp=xp)
     else:
         condition = astype(condition, dtype=dtype, xp=xp)
-
     return xp.where(condition, *as_shared_dtype([x, y], xp=xp))
 
 

xarray/tests/__init__.py

@@ -363,6 +363,16 @@ def create_test_data(
                 )
             ),
         )
+        obj["var5"] = (
+            "dim1",
+            pd.array(
+                rs.integers(1, 10, size=dim_sizes[0]).tolist(), dtype=pd.Int64Dtype()
+            ),
+        )
+        obj["var6"] = (
+            "dim1",
+            pd.array(list(string.ascii_lowercase[: dim_sizes[0]]), dtype="string"),
+        )
     if dim_sizes == _DEFAULT_TEST_DIM_SIZES:
         numbers_values = np.array([0, 1, 2, 0, 0, 1, 1, 2, 2, 3], dtype="int64")
     else:

xarray/tests/test_concat.py

@@ -154,19 +154,20 @@ def test_concat_missing_var() -> None:
     assert_identical(actual, expected)
 
 
-def test_concat_categorical() -> None:
+def test_concat_extension_array() -> None:
     data1 = create_test_data(use_extension_array=True)
     data2 = create_test_data(use_extension_array=True)
     concatenated = concat([data1, data2], dim="dim1")
-    assert (
-        concatenated["var4"]
-        == type(data2["var4"].variable.data)._concat_same_type(
-            [
-                data1["var4"].variable.data,
-                data2["var4"].variable.data,
-            ]
-        )
-    ).all()
+    for var in ["var4", "var5"]:
+        assert (
+            concatenated[var]
+            == type(data2[var].variable.data)._concat_same_type(
+                [
+                    data1[var].variable.data,
+                    data2[var].variable.data,
+                ]
+            )
+        ).all()
 
 
 def test_concat_missing_multiple_consecutive_var() -> None:

xarray/tests/test_dataarray.py

@@ -3075,6 +3075,24 @@ def test_propagate_attrs(self, func) -> None:
         with set_options(keep_attrs=True):
             assert func(da).attrs == da.attrs
 
+    def test_fillna_extension_array_int(self) -> None:
+        srs: pd.Series = pd.Series(
+            index=np.array([1, 2, 3]), data=pd.array([pd.NA, 1, 1])
+        )
+        da = srs.to_xarray()
+        filled = da.fillna(0)
+        assert filled.dtype == pd.Int64Dtype()
+        assert (filled.values == np.array([0, 1, 1])).all()
+
+    def test_dropna_extension_array_int(self) -> None:
+        srs: pd.Series = pd.Series(
+            index=np.array([1, 2, 3]), data=pd.array([pd.NA, 1, 1])
+        )
+        da = srs.to_xarray()
+        filled = da.dropna("index")
+        assert filled.dtype == pd.Int64Dtype()
+        assert (filled.values == np.array([1, 1])).all()
+
     def test_fillna(self) -> None:
         a = DataArray([np.nan, 1, np.nan, 3], coords={"x": range(4)}, dims="x")
         actual = a.fillna(-1)

xarray/tests/test_dataset.py

@@ -298,6 +298,8 @@ def test_repr(self) -> None:
                 var2     (dim1, dim2) float64 576B 0.953 1.52 1.704 ... 0.1347 -0.6423
                 var3     (dim3, dim1) float64 640B 0.4107 0.9941 0.1665 ... 0.716 1.555
                 var4     (dim1) category 32B 'b' 'c' 'b' 'a' 'c' 'a' 'c' 'a'
+                var5     (dim1) Int64 72B 5 9 7 2 6 2 8 1
+                var6     (dim1) string 64B 'a' 'b' 'c' 'd' 'e' 'f' 'g' 'h'
             Attributes:
                 foo:      bar""".format(
                 data["dim3"].dtype,
@@ -1827,25 +1829,36 @@ def test_categorical_index_reindex(self) -> None:
         actual = ds.reindex(cat=["foo"])["cat"].values
         assert (actual == np.array(["foo"])).all()
 
-    @pytest.mark.parametrize("fill_value", [np.nan, pd.NA])
-    def test_extensionarray_negative_reindex(self, fill_value) -> None:
-        cat = pd.Categorical(
-            ["foo", "bar", "baz"],
-            categories=["foo", "bar", "baz", "qux", "quux", "corge"],
-        )
+    @pytest.mark.parametrize("fill_value", [np.nan, pd.NA, None])
+    @pytest.mark.parametrize(
+        "extension_array",
+        [
+            pd.Categorical(
+                ["foo", "bar", "baz"],
+                categories=["foo", "bar", "baz", "qux"],
+            ),
+            pd.array([1, 2, 3], dtype=pd.Int32Dtype()),
+            pd.array(["a", "b", "c"], dtype="string"),
+        ],
+    )
+    def test_extensionarray_negative_reindex(self, fill_value, extension_array) -> None:
         ds = xr.Dataset(
-            {"cat": ("index", cat)},
+            {"arr": ("index", extension_array)},
             coords={"index": ("index", np.arange(3))},
         )
+        kwargs = {}
+        if fill_value is not None:
+            kwargs["fill_value"] = fill_value
         reindexed_cat = cast(
             pd.api.extensions.ExtensionArray,
-            (
-                ds.reindex(index=[-1, 1, 1], fill_value=fill_value)["cat"]
-                .to_pandas()
-                .values
-            ),
+            (ds.reindex(index=[-1, 1, 1], **kwargs)["arr"].to_pandas().values),
         )
-        assert reindexed_cat.equals(pd.array([pd.NA, "bar", "bar"], dtype=cat.dtype))  # type: ignore[attr-defined]
+        assert reindexed_cat.equals(
+            pd.array(
+                [pd.NA, extension_array[1], extension_array[1]],
+                dtype=extension_array.dtype,
+            )
+        )  # type: ignore[attr-defined]
 
     def test_extension_array_reindex_same(self) -> None:
         series = pd.Series([1, 2, pd.NA, 3], dtype=pd.Int32Dtype())
@@ -5443,6 +5456,20 @@ def test_dropna(self) -> None:
         with pytest.raises(TypeError, match=r"must specify how or thresh"):
             ds.dropna("a", how=None)  # type: ignore[arg-type]
 
+    def test_fillna_extension_array_int(self) -> None:
+        srs = pd.DataFrame({"data": pd.array([pd.NA, 1, 1])}, index=np.array([1, 2, 3]))
+        ds = srs.to_xarray()
+        filled = ds.fillna(0)
+        assert filled["data"].dtype == pd.Int64Dtype()
+        assert (filled["data"].values == np.array([0, 1, 1])).all()
+
+    def test_dropna_extension_array_int(self) -> None:
+        srs = pd.DataFrame({"data": pd.array([pd.NA, 1, 1])}, index=np.array([1, 2, 3]))
+        ds = srs.to_xarray()
+        dropped = ds.dropna("index")
+        assert dropped["data"].dtype == pd.Int64Dtype()
+        assert (dropped["data"].values == np.array([1, 1])).all()
+
     def test_fillna(self) -> None:
         ds = Dataset({"a": ("x", [np.nan, 1, np.nan, 3])}, {"x": [0, 1, 2, 3]})