Xarray: DataArray.sel extremely slow

[mschrimpf]

Problem description

.sel is an xarray method I use a lot and I would have expected it to fairly efficient.
However, even on tiny DataArrays, it takes seconds.

Code Sample, a copy-pastable example if possible

import timeit

setup = """
import itertools
import numpy as np
import xarray as xr
import string

a = list(string.printable)
b = list(string.ascii_lowercase)
d = xr.DataArray(np.random.rand(len(a), len(b)), coords={'a': a, 'b': b}, dims=['a', 'b'])
d.load()
"""

run = """
for _a, _b in itertools.product(a, b):
    d.sel(a=_a, b=_b)
"""
running_times = timeit.repeat(run, setup, repeat=3, number=10)
print("xarray", running_times)  # e.g. [14.792144000064582, 15.19372400001157, 15.345327000017278]

Expected Output

I would have expected the above code to run in milliseconds.
However, it takes over 10 seconds!
Adding an additional d = d.stack(aa=['a'], bb=['b']) makes it even slower, about twice as slow.

For reference, a naive dict-indexing implementation in Python takes 0.01 seconds:

setup = """
import itertools
import numpy as np
import string

a = list(string.printable)
b = list(string.ascii_lowercase)

d = np.random.rand(len(a), len(b))
indexers = {'a': {coord: index for (index, coord) in enumerate(a)},
            'b': {coord: index for (index, coord) in enumerate(b)}}
"""

run = """
for _a, _b in itertools.product(a, b):
    index_a, index_b = indexers['a'][_a], indexers['b'][_b]
    item = d[index_a][index_b]
"""
running_times = timeit.repeat(run, setup, repeat=3, number=10)
print("dicts", running_times)  # e.g. [0.015355999930761755, 0.01466800004709512, 0.014295000000856817]

Output of xr.show_versions()

INSTALLED VERSIONS

commit: None
python: 3.7.0.final.0
python-bits: 64
OS: Linux
OS-release: 4.4.0-17134-Microsoft
machine: x86_64
processor: x86_64
byteorder: little
LC_ALL: None
LANG: None
LOCALE: en_US.UTF-8
xarray: 0.10.8
pandas: 0.23.4
numpy: 1.15.1
scipy: 1.1.0
netCDF4: 1.4.1
h5netcdf: None
h5py: None
Nio: None
zarr: None
bottleneck: None
cyordereddict: None
dask: None
distributed: None
matplotlib: 2.2.3
cartopy: None
seaborn: None
setuptools: 40.2.0
pip: 10.0.1
conda: None
pytest: 3.7.4
IPython: 6.5.0
sphinx: None

this is a follow-up from #2438

[max-sixty]

Thanks for the issue @mschrimpf

.sel is slow per operation, mainly because it's a python function call (although not the only reason - it's also doing a set of checks / potential alignments / etc). When I say 'slow', I mean about 0.5ms:

In [6]: %timeit d.sel(a='a', b='a')
1000 loops, best of 3: 522 µs per loop

While there's an overhead, the time is fairly consistent regardless of the number of items it's selecting. For example:

In [11]: %timeit d.sel(a=d['a'], b=d['b'])
1000 loops, best of 3: 1 ms per loop

So, as is often the case in the pandas / python ecosystem, if you can write code in a vectorized way, without using python in the tight loops, it's fast. If you need to run python in each loop, it's much slower.

Does that resonate?

---

While I think not the main point here, there might be some optimizations on sel. It runs isinstance 144 times! And initializes a collection 13 times? Here's the %prun of the 0.5ms command:

        1077 function calls (1066 primitive calls) in 0.002 seconds

   Ordered by: internal time

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
        6    0.000    0.000    0.000    0.000 coordinates.py:169(<genexpr>)
       13    0.000    0.000    0.000    0.000 collections.py:50(__init__)
       14    0.000    0.000    0.000    0.000 _abcoll.py:548(update)
       33    0.000    0.000    0.000    0.000 _weakrefset.py:70(__contains__)
        2    0.000    0.000    0.001    0.000 dataset.py:881(_construct_dataarray)
      144    0.000    0.000    0.000    0.000 {isinstance}
        1    0.000    0.000    0.001    0.001 dataset.py:1496(isel)
       18    0.000    0.000    0.000    0.000 {numpy.core.multiarray.array}
        3    0.000    0.000    0.000    0.000 dataset.py:92(calculate_dimensions)
       13    0.000    0.000    0.000    0.000 abc.py:128(__instancecheck__)
       36    0.000    0.000    0.000    0.000 common.py:183(__setattr__)
        2    0.000    0.000    0.000    0.000 coordinates.py:167(variables)
        2    0.000    0.000    0.000    0.000 {method 'get_loc' of 'pandas._libs.index.IndexEngine' objects}
       26    0.000    0.000    0.000    0.000 variable.py:271(shape)
       65    0.000    0.000    0.000    0.000 collections.py:90(__iter__)
        5    0.000    0.000    0.000    0.000 variable.py:136(as_compatible_data)
        3    0.000    0.000    0.000    0.000 dataarray.py:165(__init__)
        2    0.000    0.000    0.000    0.000 indexing.py:1255(__getitem__)
        3    0.000    0.000    0.000    0.000 variable.py:880(isel)
       14    0.000    0.000    0.000    0.000 collections.py:71(__setitem__)
        1    0.000    0.000    0.000    0.000 dataset.py:1414(_validate_indexers)
        6    0.000    0.000    0.000    0.000 coordinates.py:38(__iter__)
        3    0.000    0.000    0.000    0.000 variable.py:433(_broadcast_indexes)
        2    0.000    0.000    0.000    0.000 variable.py:1826(to_index)
        3    0.000    0.000    0.000    0.000 dataset.py:636(_construct_direct)
        2    0.000    0.000    0.000    0.000 indexing.py:122(convert_label_indexer)
       15    0.000    0.000    0.000    0.000 utils.py:306(__init__)
        3    0.000    0.000    0.000    0.000 indexing.py:17(expanded_indexer)
       28    0.000    0.000    0.000    0.000 collections.py:138(iteritems)
        1    0.000    0.000    0.001    0.001 indexing.py:226(remap_label_indexers)
       15    0.000    0.000    0.000    0.000 numeric.py:424(asarray)
        1    0.000    0.000    0.001    0.001 indexing.py:193(get_dim_indexers)
    80/70    0.000    0.000    0.000    0.000 {len}

[mschrimpf]

Thanks @max-sixty,
the checks per call make sense, although I still find 0.5 ms insane for a single-value lookup (python dict-indexing takes about a 50th to index every single item in the array).

The reason I'm looking into this is actually multi-dimensional grouping (#2438) which is unfortunately not implemented (the above code is essentially a step towards trying to implement that). Is there a way of vectorizing these calls with that in mind? I.e. apply a method for each group.

[max-sixty]

Is there a way of vectorizing these calls with that in mind? I.e. apply a method for each group.

I can't think of anything immediately, and doubt there's an easy way given it doesn't exist yet (though that logic can be a trap!). There's some hacky pandas reshaping you may be able to do to solve this as a one-off. Otherwise it does likely require a concerted effort with numbagg.

I occasionally hit this issue too, so as keen as you are to find a solution. Thanks for giving it a try.

[mschrimpf]

I posted a manual solution to the multi-dimensional grouping in the stackoverflow thread.
Hopefully, .sel can be made more efficient though, it's such an everyday method.

[max-sixty]

Thanks @mschrimpf. Hopefully we can get multi-dimensional groupbys, too.