Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Pull xarray's nbytes from nbytes attribute on arrays #6797

Merged
merged 10 commits into from
Jul 22, 2022
1 change: 0 additions & 1 deletion doc/api.rst
Original file line number Diff line number Diff line change
Expand Up @@ -282,7 +282,6 @@ ndarray attributes
DataArray.shape
DataArray.size
DataArray.dtype
DataArray.nbytes
DataArray.chunks


Expand Down
2 changes: 2 additions & 0 deletions doc/whats-new.rst
Original file line number Diff line number Diff line change
Expand Up @@ -34,6 +34,8 @@ Deprecations
Bug fixes
~~~~~~~~~

- :py:attr:`DataArray.nbytes` now uses the ``nbytes`` property of the underlying array if available.
By `Max Jones <https://github.com/maxrjones>`_.

Documentation
~~~~~~~~~~~~~
Expand Down
6 changes: 6 additions & 0 deletions xarray/core/dataarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -646,6 +646,12 @@ def size(self) -> int:

@property
def nbytes(self) -> int:
"""
Total bytes consumed by the elements of this DataArray's data.

If the backend data array does not include ``nbytes``, estimates
the bytes consumed based on the ``size`` and ``dtype``.
"""
return self.variable.nbytes

@property
Expand Down
6 changes: 6 additions & 0 deletions xarray/core/dataset.py
Original file line number Diff line number Diff line change
Expand Up @@ -1378,6 +1378,12 @@ def __array__(self, dtype=None):

@property
def nbytes(self) -> int:
"""
Total bytes consumed by the data arrays of all variables in this dataset.

If the backend array for any variable does not include ``nbytes``, estimates
the total bytes for that array based on the ``size`` and ``dtype``.
"""
return sum(v.nbytes for v in self.variables.values())

@property
Expand Down
10 changes: 8 additions & 2 deletions xarray/core/variable.py
Original file line number Diff line number Diff line change
Expand Up @@ -334,8 +334,14 @@ def shape(self):
return self._data.shape

@property
def nbytes(self):
return self.size * self.dtype.itemsize
def nbytes(self) -> int:
"""
Total bytes consumed by the elements of the data array.
"""
if hasattr(self.data, "nbytes"):
return self.data.nbytes
else:
return self.size * self.dtype.itemsize

@property
def _in_memory(self):
Expand Down
6 changes: 6 additions & 0 deletions xarray/tests/test_array_api.py
Original file line number Diff line number Diff line change
Expand Up @@ -43,6 +43,12 @@ def test_indexing(arrays) -> None:
assert_equal(actual, expected)


def test_properties(arrays) -> None:
np_arr, xp_arr = arrays
assert np_arr.nbytes == np_arr.data.nbytes
assert xp_arr.nbytes == np_arr.data.nbytes


def test_reorganizing_operation(arrays) -> None:
np_arr, xp_arr = arrays
expected = np_arr.transpose()
Expand Down
3 changes: 3 additions & 0 deletions xarray/tests/test_sparse.py
Original file line number Diff line number Diff line change
Expand Up @@ -274,6 +274,9 @@ def setUp(self):
self.data = sparse.random((4, 6), random_state=0, density=0.5)
self.var = xr.Variable(("x", "y"), self.data)

def test_nbytes(self):
assert self.var.nbytes == self.data.nbytes

def test_unary_op(self):
assert_sparse_equal(-self.var.data, -self.data)
assert_sparse_equal(abs(self.var).data, abs(self.data))
Expand Down