Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -2595,6 +2595,28 @@ def argmax(
"""
return self._unravel_argminmax("argmax", dim, axis, keep_attrs, skipna)

def _as_sparse(self, sparse_format=_default, fill_value=_default) -> Variable:
"""
Use sparse-array as backend.
"""
from xarray.namedarray.utils import _default as _default_named

if sparse_format is _default:
sparse_format = _default_named

if fill_value is _default:
fill_value = _default_named

out = super()._as_sparse(sparse_format, fill_value)
return cast("Variable", out)

def _to_dense(self) -> Variable:
"""
Change backend from sparse to np.array.
"""
out = super()._to_dense()
return cast("Variable", out)


class IndexVariable(Variable):
"""Wrapper for accommodating a pandas.Index in an xarray.Variable.
Expand Down
9 changes: 7 additions & 2 deletions xarray/namedarray/_typing.py
Original file line number Diff line number Diff line change
Expand Up @@ -248,7 +248,7 @@ class _sparsearrayfunction(
Corresponds to np.ndarray.
"""

def todense(self) -> NDArray[_ScalarType_co]:
def todense(self) -> np.ndarray[Any, _DType_co]:
...


Expand All @@ -262,9 +262,14 @@ class _sparsearrayapi(
Corresponds to np.ndarray.
"""

def todense(self) -> NDArray[_ScalarType_co]:
def todense(self) -> np.ndarray[Any, _DType_co]:
...


# NamedArray can most likely use both __array_function__ and __array_namespace__:
_sparsearrayfunction_or_api = (_sparsearrayfunction, _sparsearrayapi)

sparseduckarray = Union[
_sparsearrayfunction[_ShapeType_co, _DType_co],
_sparsearrayapi[_ShapeType_co, _DType_co],
]
18 changes: 8 additions & 10 deletions xarray/namedarray/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@
_DType_co,
_ScalarType_co,
_ShapeType_co,
_sparsearrayfunction_or_api,
_SupportsImag,
_SupportsReal,
)
Expand Down Expand Up @@ -810,9 +811,9 @@ def _as_sparse(
self,
sparse_format: Literal["coo"] | Default = _default,
fill_value: ArrayLike | Default = _default,
) -> Self:
) -> NamedArray[Any, _DType_co]:
"""
use sparse-array as backend.
Use sparse-array as backend.
"""
import sparse

Expand All @@ -832,18 +833,15 @@ def _as_sparse(
raise ValueError(f"{sparse_format} is not a valid sparse format") from exc

data = as_sparse(astype(self, dtype).data, fill_value=fill_value)
return self._replace(data=data)
return self._new(data=data)

def _to_dense(self) -> Self:
def _to_dense(self) -> NamedArray[Any, _DType_co]:
"""
Change backend from sparse to np.array
Change backend from sparse to np.array.
"""
from xarray.namedarray._typing import _sparsearrayfunction_or_api

if isinstance(self._data, _sparsearrayfunction_or_api):
# return self._replace(data=self._data.todense())
data_: np.ndarray[Any, Any] = self._data.todense()
return self._replace(data=data_)
data_dense: np.ndarray[Any, _DType_co] = self._data.todense()
return self._new(data=data_dense)
else:
raise TypeError("self.data is not a sparse array")

Expand Down