Skip to content

Fix regression: IndexVariable.copy(deep=True) casts dtype=U to object #3095

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.

Sign up for GitHub

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

Merged
merged 7 commits into from
Aug 2, 2019
Merged

Fix regression: IndexVariable.copy(deep=True) casts dtype=U to object #3095

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
7 commits
Select commit Hold shift + click to select a range
262581f
Fix regression: IndexVariable.copy(deep=True) casts dtype=U to object
Jul 11, 2019
ab6960f
What's New
Jul 11, 2019
678e82b
Simplify copy()
crusaderky Jul 13, 2019
3e7825c
Merge remote-tracking branch 'upstream/master' into object_index
Jul 24, 2019
ae212a0
Cosmetic
crusaderky Jul 24, 2019
16e229d
Merge remote-tracking branch 'upstream/master' into object_index
Aug 2, 2019
da45c3b
Code review
Aug 2, 2019
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
6 changes: 4 additions & 2 deletions doc/whats-new.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -46,8 +46,10 @@ Enhancements

Bug fixes
~~~~~~~~~

- Improved error handling and documentation for `.expand_dims()`
- Fix regression introduced in v0.12.2 where ``copy(deep=True)`` would convert
unicode indices to dtype=object (:issue:`3094`).
By `Guido Imperiale <https://github.com/crusaderky>`_.
- Improved error handling and documentation for `.expand_dims()`
read-only view.
- Fix tests for big-endian systems (:issue:`3125`).
By `Graham Inggs <https://github.com/ginggs>`_.
Expand Down
40 changes: 31 additions & 9 deletions xarray/core/indexing.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,12 +3,13 @@
from collections import defaultdict
from contextlib import suppress
from datetime import timedelta
from typing import Sequence
from typing import Any, Tuple, Sequence, Union

import numpy as np
import pandas as pd

from . import duck_array_ops, nputils, utils
from .npcompat import DTypeLike
from .pycompat import dask_array_type, integer_types
from .utils import is_dict_like

Expand Down Expand Up @@ -1227,9 +1228,10 @@ def transpose(self, order):


class PandasIndexAdapter(ExplicitlyIndexedNDArrayMixin):
"""Wrap a pandas.Index to preserve dtypes and handle explicit indexing."""
"""Wrap a pandas.Index to preserve dtypes and handle explicit indexing.
"""

def __init__(self, array, dtype=None):
def __init__(self, array: Any, dtype: DTypeLike = None):
self.array = utils.safe_cast_to_index(array)
if dtype is None:
if isinstance(array, pd.PeriodIndex):
Expand All @@ -1241,13 +1243,15 @@ def __init__(self, array, dtype=None):
dtype = np.dtype('O')
else:
dtype = array.dtype
else:
dtype = np.dtype(dtype)
self._dtype = dtype

@property
def dtype(self):
def dtype(self) -> np.dtype:
return self._dtype

def __array__(self, dtype=None):
def __array__(self, dtype: DTypeLike = None) -> np.ndarray:
if dtype is None:
dtype = self.dtype
array = self.array
Expand All @@ -1258,11 +1262,18 @@ def __array__(self, dtype=None):
return np.asarray(array.values, dtype=dtype)

@property
def shape(self):
def shape(self) -> Tuple[int]:
# .shape is broken on pandas prior to v0.15.2
return (len(self.array),)

def __getitem__(self, indexer):
def __getitem__(
self, indexer
) -> Union[
NumpyIndexingAdapter,
np.ndarray,
np.datetime64,
np.timedelta64,
]:
key = indexer.tuple
if isinstance(key, tuple) and len(key) == 1:
# unpack key so it can index a pandas.Index object (pandas.Index
Expand Down Expand Up @@ -1299,9 +1310,20 @@ def __getitem__(self, indexer):

return result

def transpose(self, order):
def transpose(self, order) -> pd.Index:
return self.array # self.array should be always one-dimensional

def __repr__(self):
def __repr__(self) -> str:
return ('%s(array=%r, dtype=%r)'
% (type(self).__name__, self.array, self.dtype))

def copy(self, deep: bool = True) -> 'PandasIndexAdapter':
# Not the same as just writing `self.array.copy(deep=deep)`, as
# shallow copies of the underlying numpy.ndarrays become deep ones
# upon pickling
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I don't quite follow this comment -- how does pickling relate to this new method?

Copy link
Contributor Author

@crusaderky crusaderky Jul 15, 2019 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

pandas.Index.copy(deep=False) creates new identical views of the underlying numpy arrays.
A numpy view becomes a real, deep-copied numpy array upon pickling. This is by design to prevent people from accidentally sending a huge numpy array over network or IPC when they just want to send a slice of it. Crucially though, this (IMHO improvable) design makes no distinction when the base array is the same size or smaller than the view.

Back to the method at hand: if two xarray.DataArray/Dataset share the same underlying pandas.Index, which is the whole intent of copy(deep=False), when they are pickled together they continue sharing the same data. This is already true when you shallow-copy xarray.Variable.
But if instead of

array = self.array.copy(deep=True) if deep else self.array

we wrote

array = self.array.copy(deep=deep)

which would be tempting as it is much more readable, everything would be the same initially, until you pickle the original index and its shallow copy together, at which point the copy would automatically and silently become a deep one, causing double RAM/disk usage upon storing/unpickling.

>>> idx = xarray.IndexVariable('x', list(range(500000)))._data.array
>>> len(pickle.dumps((idx, idx)))                                                                                                       
4000281
>>> len(pickle.dumps((idx, idx.copy(deep=False))))                                                                                     
8000341

This was a real problem I faced a couple of years ago where I had to dump to disk (for forensic purposes) about 10,000 intermediate steps of a Monte Carlo simulation, each step being a DataArray or Dataset. The data variables were all dask-backed, so they were fine. But among the indices I had a 'scenario' dimension with 500,000 points, dtype='<U13'.

Before pickling: 13 * 500,000 = 6.2 MB; the 50,000 shallow copies of it being just views
After pickling: 13 * 500,000 * 10,000 = I needed a new hard disk!

This issue was from a couple of years ago and has since been fixed. My (apparently overcomplicated) code above prevents it from showing up again.

P.S. this, by the way, is a big problem with xarray.align when you are dealing with mixed numpy+dask arrays

>>> a = xarray.DataArray(da.ones(500000, chunks=10000), dims=['x'])                                                                     
>>> b = xarray.concat([a] + [xarray.DataArray(i) for i in range(1000)], dim='y')                                                        
>>> c = b.sum()
>>> c.compute()
<xarray.DataArray ()>
array(2.497505e+11)
# Everything is fine so far (not very fast due to the dask chunks of size 1 but still)...
>>> client = distributed.Client()
>>> c.compute()

Watch the RAM usage of your dask client rocket to 8 GB, and several GBs worker-side too.
This will take a while because, client side, you just created 3.8 GB (500000 * 8 * 1000) worth of pickle file and are now sending it over the network.

[EDIT] nvm this last example. The initial concat() is already sending the RAM usage to 4 GB, meaning it's not calling numpy.ndarray.broadcast_to under the hood as it used to do...

Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Thanks for the explanation, this makes complete sense.

# >>> len(pickle.dumps((self.array, self.array)))
# 4000281
# >>> len(pickle.dumps((self.array, self.array.copy(deep=False))))
# 8000341
array = self.array.copy(deep=True) if deep else self.array
return PandasIndexAdapter(array, self._dtype)
9 changes: 1 addition & 8 deletions xarray/core/variable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1942,14 +1942,7 @@ def copy(self, deep=True, data=None):
data copied from original.
"""
if data is None:
if deep:
# self._data should be a `PandasIndexAdapter` instance at this
# point, which doesn't have a copy method, so make a deep copy
# of the underlying `pandas.MultiIndex` and create a new
# `PandasIndexAdapter` instance with it.
data = PandasIndexAdapter(self._data.array.copy(deep=True))
else:
data = self._data
data = self._data.copy(deep=deep)
else:
data = as_compatible_data(data)
if self.shape != data.shape:
Expand Down
5 changes: 3 additions & 2 deletions xarray/tests/test_variable.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -477,8 +477,9 @@ def test_concat_mixed_dtypes(self):
assert actual.dtype == object

@pytest.mark.parametrize('deep', [True, False])
def test_copy(self, deep):
v = self.cls('x', 0.5 * np.arange(10), {'foo': 'bar'})
@pytest.mark.parametrize('astype', [float, int, str])
def test_copy(self, deep, astype):
v = self.cls('x', (0.5 * np.arange(10)).astype(astype), {'foo': 'bar'})
w = v.copy(deep=deep)
assert type(v) is type(w)
assert_identical(v, w)
Expand Down