(feat): Support for pandas ExtensionArray

doc/whats-new.rst

@@ -39,6 +39,10 @@ New Features
   This is currently limited to the linear interpolation method (`method='linear'`).
   (:issue:`7377`, :pull:`8684`) By `Marco Wolsza <https://github.com/maawoo>`_.
 
+- Xarray now makes a best attempt not to coerce :py:class:`pandas.api.extensions.ExtensionArray` to a numpy array
+by supporting 1D `ExtensionArray` objects internally where possible.  Thus, `Dataset`s initialized with a `pd.Catgeorical`, for example,
+will retain the object.  However, one cannot do operations that are not possible on the `ExtensionArray` then, such as broadcasting.
+
 Breaking changes
 ~~~~~~~~~~~~~~~~
 

pyproject.toml

@@ -30,7 +30,7 @@ dependencies = [
 
 [project.optional-dependencies]
 accel = ["scipy", "bottleneck", "numbagg", "flox", "opt_einsum"]
-complete = ["xarray[accel,io,parallel,viz,dev]"]
+complete = ["xarray[accel,io,parallel,viz,dev,extension_arrays]"]
 dev = [
   "hypothesis",
   "pre-commit",
@@ -45,6 +45,7 @@ dev = [
 io = ["netCDF4", "h5netcdf", "scipy", 'pydap; python_version<"3.10"', "zarr", "fsspec", "cftime", "pooch"]
 parallel = ["dask[complete]"]
 viz = ["matplotlib", "seaborn", "nc-time-axis"]
+extension-arrays = ["plum-dispatch"]
 
 [project.urls]
 Documentation = "https://docs.xarray.dev"
@@ -124,6 +125,7 @@ module = [
   "opt_einsum.*",
   "pandas.*",
   "pooch.*",
+  "pyarrow.*",
   "pydap.*",
   "pytest.*",
   "scipy.*",

xarray/coding/strings.py

@@ -5,6 +5,7 @@
 from functools import partial
 
 import numpy as np
+from pandas.api.types import is_extension_array_dtype
 
 from xarray.coding.variables import (
     VariableCoder,
@@ -27,11 +28,10 @@ def create_vlen_dtype(element_type):
 
 
 def check_vlen_dtype(dtype):
-    if dtype.kind != "O" or dtype.metadata is None:
+    if is_extension_array_dtype(dtype) or dtype.kind != "O" or dtype.metadata is None:
         return None
-    else:
-        # check xarray (element_type) as well as h5py (vlen)
-        return dtype.metadata.get("element_type", dtype.metadata.get("vlen"))
+    # check xarray (element_type) as well as h5py (vlen)
+    return dtype.metadata.get("element_type", dtype.metadata.get("vlen"))
 
 
 def is_unicode_dtype(dtype):

xarray/coding/times.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import pandas as pd
+from pandas.api.types import is_extension_array_dtype
 from pandas.errors import OutOfBoundsDatetime, OutOfBoundsTimedelta
 
 from xarray.coding.variables import (
@@ -967,9 +968,10 @@ def __init__(self, use_cftime: bool | None = None) -> None:
         self.use_cftime = use_cftime
 
     def encode(self, variable: Variable, name: T_Name = None) -> Variable:
-        if np.issubdtype(
-            variable.data.dtype, np.datetime64
-        ) or contains_cftime_datetimes(variable):
+        if (not is_extension_array_dtype(variable.data)) and (
+            np.issubdtype(variable.data.dtype, np.datetime64)
+            or contains_cftime_datetimes(variable)
+        ):
             dims, data, attrs, encoding = unpack_for_encoding(variable)
 
             units = encoding.pop("units", None)
@@ -1007,7 +1009,9 @@ def decode(self, variable: Variable, name: T_Name = None) -> Variable:
 
 class CFTimedeltaCoder(VariableCoder):
     def encode(self, variable: Variable, name: T_Name = None) -> Variable:
-        if np.issubdtype(variable.data.dtype, np.timedelta64):
+        if (not is_extension_array_dtype(variable.data)) and np.issubdtype(
+            variable.data.dtype, np.timedelta64
+        ):
             dims, data, attrs, encoding = unpack_for_encoding(variable)
 
             data, units = encode_cf_timedelta(

xarray/coding/variables.py

@@ -9,6 +9,7 @@
 
 import numpy as np
 import pandas as pd
+from pandas.api.types import is_extension_array_dtype
 
 from xarray.core import dtypes, duck_array_ops, indexing
 from xarray.core.parallelcompat import get_chunked_array_type
@@ -250,7 +251,6 @@ class CFMaskCoder(VariableCoder):
     def encode(self, variable: Variable, name: T_Name = None):
         dims, data, attrs, encoding = unpack_for_encoding(variable)
 
-        dtype = np.dtype(encoding.get("dtype", data.dtype))
         fv = encoding.get("_FillValue")
         mv = encoding.get("missing_value")
 
@@ -268,6 +268,8 @@ def encode(self, variable: Variable, name: T_Name = None):
         # special case DateTime to properly handle NaT
         is_time_like = _is_time_like(attrs.get("units"))
 
+        dtype = np.dtype(encoding.get("dtype", data.dtype))
+
         if fv_exists:
             # Ensure _FillValue is cast to same dtype as data's
             encoding["_FillValue"] = dtype.type(fv)
@@ -472,16 +474,18 @@ class DefaultFillvalueCoder(VariableCoder):
 
     def encode(self, variable: Variable, name: T_Name = None) -> Variable:
         dims, data, attrs, encoding = unpack_for_encoding(variable)
+        has_no_fill = "_FillValue" not in attrs and "_FillValue" not in encoding
         # make NaN the fill value for float types
-        if (
-            "_FillValue" not in attrs
-            and "_FillValue" not in encoding
-            and np.issubdtype(variable.dtype, np.floating)
-        ):
+        if is_extension_array_dtype(data):
+            if not has_no_fill:
+                raise ValueError(
+                    "Found _FillValue encoding or attr on extension array."
+                )
+            return variable
+        if has_no_fill and np.issubdtype(variable.dtype, np.floating):
             attrs["_FillValue"] = variable.dtype.type(np.nan)
             return Variable(dims, data, attrs, encoding, fastpath=True)
-        else:
-            return variable
+        return variable
 
     def decode(self, variable: Variable, name: T_Name = None) -> Variable:
         raise NotImplementedError()

xarray/conventions.py

@@ -6,6 +6,7 @@
 
 import numpy as np
 import pandas as pd
+from pandas.api.types import is_extension_array_dtype
 
 from xarray.coding import strings, times, variables
 from xarray.coding.variables import SerializationWarning, pop_to
@@ -114,7 +115,10 @@ def ensure_dtype_not_object(var: Variable, name: T_Name = None) -> Variable:
         dims, data, attrs, encoding = variables.unpack_for_encoding(var)
 
         # leave vlen dtypes unchanged
-        if strings.check_vlen_dtype(data.dtype) is not None:
+        if (
+            is_extension_array_dtype(data)
+            or strings.check_vlen_dtype(data.dtype) is not None
+        ):
             return var
 
         if is_duck_dask_array(data):
@@ -356,9 +360,9 @@ def _update_bounds_encoding(variables: T_Variables) -> None:
         attrs = v.attrs
         encoding = v.encoding
         has_date_units = "units" in encoding and "since" in encoding["units"]
-        is_datetime_type = np.issubdtype(
-            v.dtype, np.datetime64
-        ) or contains_cftime_datetimes(v)
+        is_datetime_type = (not is_extension_array_dtype(v)) and (
+            contains_cftime_datetimes(v) or np.issubdtype(v.dtype, np.datetime64)
+        )
 
         if (
             is_datetime_type

xarray/core/dataset.py

@@ -24,6 +24,7 @@
 from typing import IO, TYPE_CHECKING, Any, Callable, Generic, Literal, cast, overload
 
 import numpy as np
+from pandas.api.types import is_extension_array_dtype
 
 # remove once numpy 2.0 is the oldest supported version
 try:
@@ -6835,6 +6836,7 @@ def reduce(
                     # that don't have the reduce dims: PR5393
                     not reduce_dims
                     or not numeric_only
+                    or not is_extension_array_dtype(var.dtype)
                     or np.issubdtype(var.dtype, np.number)
                     or (var.dtype == np.bool_)
                 ):
@@ -7149,13 +7151,33 @@ def to_pandas(self) -> pd.Series | pd.DataFrame:
         )
 
     def _to_dataframe(self, ordered_dims: Mapping[Any, int]):
-        columns = [k for k in self.variables if k not in self.dims]
+        columns = [
+            k
+            for k in self.variables
+            if k not in self.dims
+            and not is_extension_array_dtype(self.variables[k].data)
+        ]
+        extension_array_columns = [
+            k
+            for k in self.variables
+            if k not in self.dims and is_extension_array_dtype(self.variables[k].data)
+        ]
         data = [
             self._variables[k].set_dims(ordered_dims).values.reshape(-1)
             for k in columns
         ]
         index = self.coords.to_index([*ordered_dims])
-        return pd.DataFrame(dict(zip(columns, data)), index=index)
+        broadcasted_df = pd.DataFrame(dict(zip(columns, data)), index=index)
+        for extension_array_column in extension_array_columns:
+            extension_array = self.variables[extension_array_column].data.array
+            index = self[self.variables[extension_array_column].dims[0]].data
+            cat_df = pd.DataFrame(
+                {extension_array_column: extension_array},
+                index=self[self.variables[extension_array_column].dims[0]].data,
+            )
+            cat_df.index.name = self.variables[extension_array_column].dims[0]
+            broadcasted_df = broadcasted_df.join(cat_df)
+        return broadcasted_df
 
     def to_dataframe(self, dim_order: Sequence[Hashable] | None = None) -> pd.DataFrame:
         """Convert this dataset into a pandas.DataFrame.
@@ -7301,11 +7323,14 @@ def from_dataframe(cls, dataframe: pd.DataFrame, sparse: bool = False) -> Self:
                 "cannot convert a DataFrame with a non-unique MultiIndex into xarray"
             )
 
-        # Cast to a NumPy array first, in case the Series is a pandas Extension
-        # array (which doesn't have a valid NumPy dtype)
-        # TODO: allow users to control how this casting happens, e.g., by
-        # forwarding arguments to pandas.Series.to_numpy?
-        arrays = [(k, np.asarray(v)) for k, v in dataframe.items()]
+        arrays = [
+            (k, np.asarray(v))
+            for k, v in dataframe.items()
+            if not is_extension_array_dtype(v)
+        ]
+        extension_arrays = [
+            (k, v) for k, v in dataframe.items() if is_extension_array_dtype(v)
+        ]
 
         indexes: dict[Hashable, Index] = {}
         index_vars: dict[Hashable, Variable] = {}
@@ -7319,6 +7344,8 @@ def from_dataframe(cls, dataframe: pd.DataFrame, sparse: bool = False) -> Self:
                 xr_idx = PandasIndex(lev, dim)
                 indexes[dim] = xr_idx
                 index_vars.update(xr_idx.create_variables())
+            arrays += [(k, np.asarray(v)) for k, v in extension_arrays]
+            extension_arrays = []
         else:
             index_name = idx.name if idx.name is not None else "index"
             dims = (index_name,)
@@ -7332,6 +7359,8 @@ def from_dataframe(cls, dataframe: pd.DataFrame, sparse: bool = False) -> Self:
             obj._set_sparse_data_from_dataframe(idx, arrays, dims)
         else:
             obj._set_numpy_data_from_dataframe(idx, arrays, dims)
+        for name, extension_array in extension_arrays:
+            obj[name] = (dims, extension_array)
         return obj
 
     def to_dask_dataframe(

xarray/core/duck_array_ops.py

@@ -10,8 +10,10 @@
 import datetime
 import inspect
 import warnings
+from collections.abc import Sequence
 from functools import partial
 from importlib import import_module
+from typing import Callable
 
 import numpy as np
 import pandas as pd
@@ -32,11 +34,21 @@
 from numpy import concatenate as _concatenate
 from numpy.lib.stride_tricks import sliding_window_view  # noqa
 from packaging.version import Version
+from pandas.api.types import is_extension_array_dtype
+
+try:
+    from plum import dispatch  # type: ignore[import-not-found]
+except ImportError:
+
+    def dispatch(func):
+        return func
+
 
 from xarray.core import dask_array_ops, dtypes, nputils, pycompat
 from xarray.core.options import OPTIONS
 from xarray.core.parallelcompat import get_chunked_array_type, is_chunked_array
 from xarray.core.pycompat import array_type, is_duck_dask_array
+from xarray.core.types import DTypeLikeSave, T_ExtensionArray
 from xarray.core.utils import is_duck_array, module_available
 
 # remove once numpy 2.0 is the oldest supported version
@@ -53,6 +65,64 @@
 dask_available = module_available("dask")
 
 
+HANDLED_EXTENSION_ARRAY_FUNCTIONS: dict[Callable, Callable] = {}
+
+
+def implements(numpy_function):
+    """Register an __array_function__ implementation for MyArray objects."""
+
+    def decorator(func):
+        HANDLED_EXTENSION_ARRAY_FUNCTIONS[numpy_function] = func
+        return func
+
+    return decorator
+
+
+@implements(np.issubdtype)
+@dispatch
+def __extension_duck_array__issubdtype(
+    extension_array_dtype: T_ExtensionArray, other_dtype: DTypeLikeSave
+) -> bool:
+    return False  # never want a function to think a pandas extension dtype is a subtype of numpy
+
+
+@implements(np.broadcast_to)
+@dispatch
+def __extension_duck_array__broadcast(arr: T_ExtensionArray, shape: tuple):
+    if shape[0] == len(arr) and len(shape) == 1:
+        return arr
+    raise NotImplementedError("Cannot broadcast 1d-only pandas categorical array.")
+
+
+@implements(np.stack)
+@dispatch
+def __extension_duck_array__stack(arr: T_ExtensionArray, axis: int):
+    raise NotImplementedError("Cannot stack 1d-only pandas categorical array.")
+
+
+@implements(np.concatenate)
+@dispatch
+def __extension_duck_array__concatenate(
+    arrays: Sequence[T_ExtensionArray], axis: int = 0, out=None
+) -> T_ExtensionArray:
+    return type(arrays[0])._concat_same_type(arrays)
+
+
+@implements(np.where)
+@dispatch
+def __extension_duck_array__where(
+    condition: np.ndarray, x: T_ExtensionArray, y: T_ExtensionArray
+) -> T_ExtensionArray:
+    if (
+        isinstance(x, pd.Categorical)
+        and isinstance(y, pd.Categorical)
+        and x.dtype != y.dtype
+    ):
+        x = x.add_categories(set(y.categories).difference(set(x.categories)))
+        y = y.add_categories(set(x.categories).difference(set(y.categories)))
+    return pd.Series(x).where(condition, pd.Series(y)).array
+
+
 def get_array_namespace(x):
     if hasattr(x, "__array_namespace__"):
         return x.__array_namespace__()
@@ -155,7 +225,7 @@ def isnull(data):
         return full_like(data, dtype=bool, fill_value=False)
     else:
         # at this point, array should have dtype=object
-        if isinstance(data, np.ndarray):
+        if isinstance(data, np.ndarray) or is_extension_array_dtype(data):
             return pandas_isnull(data)
         else:
             # Not reachable yet, but intended for use with other duck array
@@ -220,9 +290,17 @@ def asarray(data, xp=np):
 
 def as_shared_dtype(scalars_or_arrays, xp=np):
     """Cast a arrays to a shared dtype using xarray's type promotion rules."""
-    array_type_cupy = array_type("cupy")
-    if array_type_cupy and any(
-        isinstance(x, array_type_cupy) for x in scalars_or_arrays
+    if any(is_extension_array_dtype(x) for x in scalars_or_arrays):
+        extension_array_types = [
+            x.dtype for x in scalars_or_arrays if is_extension_array_dtype(x)
+        ]
+        if len(extension_array_types) == len(scalars_or_arrays) and all(
+            isinstance(x, type(extension_array_types[0])) for x in extension_array_types
+        ):
+            return scalars_or_arrays
+        arrays = [asarray(np.array(x), xp=xp) for x in scalars_or_arrays]
+    elif array_type_cupy := array_type("cupy") and any(  # noqa: F841
+        isinstance(x, array_type_cupy) for x in scalars_or_arrays  # noqa: F821
     ):
         import cupy as cp