pydata / pydata/xarray

[PERFORMANCE]: `isin` on `CFTimeIndex`-backed `Coordinate` slow

Open
#6,230 5 comments 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

enhancement topic-cftime topic-performance
Dominant language
Python
Stars
4.2k
Forks
1.4k
Avg merge
2d 15h
Merged PRs (30d)
14

Description

Is your feature request related to a problem?

I want to do coord1.isin.coord2 and it is quite slow when coords are large and of object type CFTimeIndex.

import xarray as xr
import numpy as np

n=1000
coord1 = xr.cftime_range(start='2000', freq='MS', periods=n)
coord2 = xr.cftime_range(start='2000', freq='3MS', periods=n)

# cftimeindex: very fast
%timeit coord1.isin(coord2) # 743 µs ± 1.33 µs

# np.isin on index.asi8
%timeit np.isin(coord1.asi8,coord2.asi8) # 7.83 ms ± 14.1 µs

da = xr.DataArray(np.random.random((n,n)),dims=['a','b'],coords={'a':coord1,'b':coord2})

# when xr.DataArray coordinate slow
%timeit da.a.isin(da.b) # 94.9 ms ± 959 µs

# when converting xr.DataArray coordinate back to index slow
%timeit np.isin(da.a.to_index(), da.b.to_index()) # 97.4 ms ± 819 µs

# when converting xr.DataArray coordinate back to index asi
%timeit np.isin(da.a.to_index().asi8, da.b.to_index().asi8) # 7.89 ms ± 15.2 µs
Describe the solution you'd like

faster coord1.isin.coord2 by default. could we re-route here, e.g. to the alternative?

conversion from coordinate to_index() is costly I guess

Describe alternatives you've considered

np.isin(coord1.to_index().asi8, coord2.to_index().asi8 brings me nice speedups in https://github.com/pangeo-data/climpred/pull/724

Additional context

unsure whether this issue should go here on in cftime

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start at the Coordinate .isin and .to_index entry points, using the provided CFTimeIndex-backed benchmark to compare them with np.isin on .asi8. The work is done when coordinate membership tests preserve their current behavior while avoiding the costly conversion path and show a meaningful speedup for large object-type CFTimeIndex coordinates.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
data, performance
Issue type
Feature
Difficulty
4/5
Estimated time
3-5 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
42/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.