Update contains_cftime_datetimes to avoid loading entire variable array

[agoodm]

• Closes Opening datasets with large object dtype arrays is very slow #7484
• User visible changes (including notable bug fixes) are documented in whats-new.rst

This PR greatly improves the performance for opening datasets with large arrays of object type (typically string arrays) since contains_cftime_datetimes was triggering the entire array to be read from the file just to check the very first element in the entire array.

@Illviljan continuing our discussion from the issue thread, I did try to pass in var._data to _contains_cftime_datetimes, but I had a lot of trouble finding a way to generalize how to index the first array element. The best I could do was var._data.array.get_array(), but I don't think get_array is implemented for every backend. So for now I am leaving my original proposed solution.

[mathause]

Thanks for the PR. However, does that actually make a difference? To me it looks like _contains_cftime_datetimes also only considers one element of the array.

xarray/xarray/core/common.py

Lines 1779 to 1780 in b451558

	if array.dtype == np.dtype("O") and array.size > 0:
	sample = np.asarray(array).flat[0]

[agoodm]

Thanks for the PR. However, does that actually make a difference? To me it looks like _contains_cftime_datetimes also only considers one element of the array.

xarray/xarray/core/common.py

Lines 1779 to 1780 in b451558

	if array.dtype == np.dtype("O") and array.size > 0:
	sample = np.asarray(array).flat[0]

This isn't actually the line of code that's causing the performance bottleneck, it's the access to var.data in the function call that is actually problematic as I explained in the issue thread. You can verify this yourself running this simple example before and after applying the changes in this PR:

import numpy as np
import xarray as xr

str_array = np.arange(100000000).astype(str)
ds = xr.DataArray(dims=('x',), data=str_array).to_dataset(name='str_array')
ds = ds.chunk(x=10000)
ds['str_array'] = ds.str_array.astype('O') # Needs to actually be object dtype to show the problem
ds.to_zarr('str_array.zarr')

%time xr.open_zarr('str_array.zarr')

[Illviljan]

@agoodm, what you think of this version? Using xr.Variable directly seems a little easier to work with than trying to guess which type of array (cupy, dask, pint, backendarray, etc) is in the variable.

[agoodm]

@Illviljan I gave your update a quick test, it seems to work well enough and still maintains the performance improvement. It looks fine to me though I guess it looks like you still need to fix this failing mypy stuff now?

[mathause]

I thought da.data passes the array along as is - but you can learn something everyday. Thanks @Illviljan for taking over and sorry @agoodm for not properly reading the issue...

[dcherian]

very clean. It'd be nice to add some sort of test like DuckBackendArrayWrapper in https://github.com/pydata/xarray/pull/6874/files . __getitem__ should raise if it will return more than one value.

[agoodm]

Thanks for taking a look at this. I am a little confused for what you are suggesting here. Are you looking for a simple test in test_variable.py that applies the same logic in this block to extract the very first element via Variable.__getitem__ here and check that it returns one value, a more general contains_cftime_datetimes test, or both?

[dcherian]

Sotry that was a bit complicated and intended for IIlviljan.

I pushed a commit with a test. I also changed the code to account for those lazily indexed backend arrays explicitly.

[headtr1ck]

LGTM.

Some simple test of the new functionality would be nice.

If you are really motivated, you can think about adding an asv benchmark here: https://github.com/pydata/xarray/blob/main/asv_bench/benchmarks/dataset_io.py

[dcherian]

Seems like we're passing a DataArray instead of a Variable somewhere.

[mathause]

I think that's in the tests themselves

xarray/xarray/tests/test_coding_times.py

Line 778 in 6531b57

@pytest.fixture()

[dcherian]

Thanks @agoodm this work prompted a bunch of internal cleanup!

[agoodm]

Thanks @Illviljan and @dcherian for helping to see this through.