Enable flox in GroupBy and resample (#5734)

Closes #5734
Closes #4473
Closes #4498
Closes #659
Closes #2237
xref pangeo-data/pangeo#271

Co-authored-by: Anderson Banihirwe <[email protected]>
Co-authored-by: pre-commit-ci[bot] <66853113+pre-commit-ci[bot]@users.noreply.github.com>
Co-authored-by: Illviljan <[email protected]>
Co-authored-by: Mathias Hauser <[email protected]>
Co-authored-by: Stephan Hoyer <[email protected]>

asv_bench/asv.conf.json

@@ -67,6 +67,8 @@
  "bottleneck": [""],
  "dask": [""],
  "distributed": [""],
+ "flox": [""],
+ "numpy_groupies": [""],
  "sparse": [""]
  },
 

asv_bench/benchmarks/groupby.py

@@ -13,6 +13,7 @@ def setup(self, *args, **kwargs):
  {
  "a": xr.DataArray(np.r_[np.repeat(1, self.n), np.repeat(2, self.n)]),
  "b": xr.DataArray(np.arange(2 * self.n)),
+ "c": xr.DataArray(np.arange(2 * self.n)),
  }
  )
  self.ds2d = self.ds1d.expand_dims(z=10)
@@ -50,10 +51,11 @@ class GroupByDask(GroupBy):
  def setup(self, *args, **kwargs):
  requires_dask()
  super().setup(**kwargs)
- self.ds1d = self.ds1d.sel(dim_0=slice(None, None, 2)).chunk({"dim_0": 50})
- self.ds2d = self.ds2d.sel(dim_0=slice(None, None, 2)).chunk(
- {"dim_0": 50, "z": 5}
- )
+
+ self.ds1d = self.ds1d.sel(dim_0=slice(None, None, 2))
+ self.ds1d["c"] = self.ds1d["c"].chunk({"dim_0": 50})
+ self.ds2d = self.ds2d.sel(dim_0=slice(None, None, 2))
+ self.ds2d["c"] = self.ds2d["c"].chunk({"dim_0": 50, "z": 5})
  self.ds1d_mean = self.ds1d.groupby("b").mean()
  self.ds2d_mean = self.ds2d.groupby("b").mean()
 

ci/install-upstream-wheels.sh

@@ -15,6 +15,7 @@ conda uninstall -y --force \
  pint \
  bottleneck \
  sparse \
+ flox \
  h5netcdf \
  xarray
 # to limit the runtime of Upstream CI
@@ -47,4 +48,5 @@ python -m pip install \
  git+https://github.com/pydata/sparse \
  git+https://github.com/intake/filesystem_spec \
  git+https://github.com/SciTools/nc-time-axis \
+ git+https://github.com/dcherian/flox \
  git+https://github.com/h5netcdf/h5netcdf

ci/requirements/all-but-dask.yml

@@ -13,6 +13,7 @@ dependencies:
  - cfgrib
  - cftime
  - coveralls
+ - flox
  - h5netcdf
  - h5py
  - hdf5

ci/requirements/environment-windows.yml

@@ -10,6 +10,7 @@ dependencies:
  - cftime
  - dask-core
  - distributed
+ - flox
  - fsspec!=2021.7.0
  - h5netcdf
  - h5py

ci/requirements/environment.yml

@@ -12,6 +12,7 @@ dependencies:
  - cftime
  - dask-core
  - distributed
+ - flox
  - fsspec!=2021.7.0
  - h5netcdf
  - h5py

ci/requirements/min-all-deps.yml

@@ -17,6 +17,7 @@ dependencies:
  - coveralls
  - dask-core=2021.04
  - distributed=2021.04
+ - flox=0.5
  - h5netcdf=0.11
  - h5py=3.1
  # hdf5 1.12 conflicts with h5py=3.1

doc/whats-new.rst

@@ -141,6 +141,11 @@ Performance
 - GroupBy binary operations are now vectorized.
  Previously this involved looping over all groups. (:issue:`5804`,:pull:`6160`)
  By `Deepak Cherian <https://github.com/dcherian>`_.
+- Substantially improved GroupBy operations using `flox <https://flox.readthedocs.io/en/latest/>`_.
+ This is auto-enabled when ``flox`` is installed. Use ``xr.set_options(use_flox=False)`` to use
+ the old algorithm. (:issue:`4473`, :issue:`4498`, :issue:`659`, :issue:`2237`, :pull:`271`).
+ By `Deepak Cherian <https://github.com/dcherian>`_,`Anderson Banihirwe <https://github.com/andersy005>`_,
+ `Jimmy Westling <https://github.com/illviljan>`_.
 
 Internal Changes
 ~~~~~~~~~~~~~~~~

setup.cfg

@@ -98,6 +98,7 @@ accel =
  scipy
  bottleneck
  numbagg
+ flox
 
 parallel =
  dask[complete]