Enable flox in GroupBy and resample

[andersy005]

This PR supersedes #4540. This is still a very rough draft :)

• Closes
  • Closes Wrap numpy-groupies to speed up Xarray's groupby aggregations #4473
  • Closes Resample is ~100x slower than Pandas resample; Speed is related to resample period (unlike Pandas) #4498
  • Closes groupby very slow compared to pandas #659
  • Closes why time grouping doesn't preserve chunks #2237
  • xref calculating climatologies efficiently pangeo-data/pangeo#271
• Tests added
• Passes pre-commit run --all-files
• User visible changes (including notable bug fixes) are documented in whats-new.rst

TODO:

• flox vs legacy tests with nans
• test passing flox-specific kwargs like engine and method.

[github-actions]

Unit Test Results

         6 files           6 suites   57m 10s ⏱️
16 290 tests 14 551 ✔️ 1 739 💤 0 ❌
90 936 runs  82 738 ✔️ 8 198 💤 0 ❌

Results for commit a2168df.

♻️ This comment has been updated with latest results.

[max-sixty]

Very excited for this! Thank you @andersy005 !

If you've already run them — how do the benchmarks look? I imagine they're going to do really well for things like resample.

[dcherian]

We don't have any asv benchmarks for groupby currently. It would be good to add some!

[Illviljan]

Looking at those unit errors it appears numpy_groupies is forcing the duck arrays to numpy arrays. Perhaps adding a like in the np.asanyarray() will do the trick?

C:\Miniconda\envs\xarray-tests\lib\site-packages\numpy_groupies\utils_numpy.py:199: in input_validation
    a = np.asanyarray(a)

https://github.com/ml31415/numpy-groupies/blob/7a31d1f9bbd51111b4a4d01cf1df01a5b4827e85/numpy_groupies/utils_numpy.py#L199

____________________ TestDataset.test_resample[int32-data] ____________________
[gw1] win32 -- Python 3.9.7 C:\Miniconda\envs\xarray-tests\python.exe

self = <xarray.tests.test_units.TestDataset object at 0x0000028489DCFC70>
variant = 'data', dtype = dtype('int32')

    @pytest.mark.parametrize(
        "variant",
        (
            "data",
            pytest.param(
                "dims", marks=pytest.mark.skip(reason="indexes don't support units")
            ),
            "coords",
        ),
    )
    def test_resample(self, variant, dtype):
        # TODO: move this to test_computation_objects
        variants = {
            "data": ((unit_registry.degK, unit_registry.Pa), 1, 1),
            "dims": ((1, 1), unit_registry.m, 1),
            "coords": ((1, 1), 1, unit_registry.m),
        }
        (unit1, unit2), dim_unit, coord_unit = variants.get(variant)
    
        array1 = np.linspace(-5, 5, 10 * 5).reshape(10, 5).astype(dtype) * unit1
        array2 = np.linspace(10, 20, 10 * 8).reshape(10, 8).astype(dtype) * unit2
    
        t = pd.date_range("10-09-2010", periods=array1.shape[0], freq="1y")
        y = np.arange(5) * dim_unit
        z = np.arange(8) * dim_unit
    
        u = np.linspace(-1, 0, 5) * coord_unit
    
        ds = xr.Dataset(
            data_vars={"a": (("time", "y"), array1), "b": (("time", "z"), array2)},
            coords={"time": t, "y": y, "z": z, "u": ("y", u)},
        )
        units = extract_units(ds)
    
        func = method("resample", time="6m")
    
        expected = attach_units(func(strip_units(ds)).mean(), units)
>       actual = func(ds).mean()

D:\a\xarray\xarray\xarray\tests\test_units.py:5366: 
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
D:\a\xarray\xarray\xarray\core\groupby.py:602: in wrapped_func
    result = xarray_reduce(
C:\Miniconda\envs\xarray-tests\lib\site-packages\dask_groupby\xarray.py:259: in xarray_reduce
    actual = xr.apply_ufunc(
D:\a\xarray\xarray\xarray\core\computation.py:1153: in apply_ufunc
    return apply_dataset_vfunc(
D:\a\xarray\xarray\xarray\core\computation.py:447: in apply_dataset_vfunc
    result_vars = apply_dict_of_variables_vfunc(
D:\a\xarray\xarray\xarray\core\computation.py:391: in apply_dict_of_variables_vfunc
    result_vars[name] = func(*variable_args)
D:\a\xarray\xarray\xarray\core\computation.py:733: in apply_variable_ufunc
    result_data = func(*input_data)
C:\Miniconda\envs\xarray-tests\lib\site-packages\dask_groupby\xarray.py:232: in wrapper
    result, groups = groupby_reduce(*args, **kwargs)
C:\Miniconda\envs\xarray-tests\lib\site-packages\dask_groupby\core.py:1119: in groupby_reduce
    results = chunk_reduce(
C:\Miniconda\envs\xarray-tests\lib\site-packages\dask_groupby\core.py:521: in chunk_reduce
    result = _get_aggregate(backend)(
C:\Miniconda\envs\xarray-tests\lib\site-packages\numpy_groupies\aggregate_numpy.py:291: in aggregate
    return _aggregate_base(group_idx, a, size=size, fill_value=fill_value,
C:\Miniconda\envs\xarray-tests\lib\site-packages\numpy_groupies\aggregate_numpy.py:256: in _aggregate_base
    group_idx, a, flat_size, ndim_idx, size = input_validation(group_idx, a,
C:\Miniconda\envs\xarray-tests\lib\site-packages\numpy_groupies\utils_numpy.py:199: in input_validation
    a = np.asanyarray(a)
C:\Miniconda\envs\xarray-tests\lib\site-packages\numpy\core\_asarray.py:171: in asanyarray
    return array(a, dtype, copy=False, order=order, subok=True)
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <Quantity([[-5 -3 -2 -1  0  0  1  2  3  4]
 [-4 -3 -2 -1  0  0  1  2  3  4]
 [-4 -3 -2 -1  0  0  1  2  3  4]
 [-4 -3 -2 -1  0  0  1  2  3  4]
 [-4 -3 -2 -1  0  0  1  2  3  5]], 'kelvin')>
t = None

    def __array__(self, t=None):
>       warnings.warn(
            "The unit of the quantity is stripped when downcasting to ndarray.",
            UnitStrippedWarning,
            stacklevel=2,
        )
E       pint.errors.UnitStrippedWarning: The unit of the quantity is stripped when downcasting to ndarray.

C:\Miniconda\envs\xarray-tests\lib\site-packages\pint\quantity.py:1700: UnitStrippedWarning

[Illviljan]

Copy/pasteing the Resample-class removes the rest of the errors. Not especially elegant though, maybe there's a better way?

xarray/core/dask_array_compat.py:8: error: Incompatible types in assignment (expression has type "None", variable has type Module)  [assignment]
xarray/backends/locks.py:10: error: Cannot assign to a type  [misc]
xarray/backends/locks.py:10: error: Incompatible types in assignment (expression has type "Type[Lock]", variable has type "Type[SerializableLock]")  [assignment]
Installing missing stub packages:
xarray/backends/locks.py:15: error: Cannot assign to a type  [misc]
/usr/share/miniconda/envs/xarray-tests/bin/python -m pip install types-PyYAML types-paramiko types-python-dateutil types-pytz types-setuptools
xarray/backends/locks.py:15: error: Incompatible types in assignment (expression has type "None", variable has type "Type[Lock]")  [assignment]

xarray/core/types.py:19: error: Cannot assign to a type  [misc]

xarray/core/types.py:19: error: Incompatible types in assignment (expression has type "Type[ndarray[Any, Any]]", variable has type "Type[Array]")  [assignment]
Found 25 errors in 14 files (checked 140 source files)
xarray/core/duck_array_ops.py:33: error: Incompatible types in assignment (expression has type "None", variable has type Module)  [assignment]
xarray/core/nanops.py:14: error: Incompatible types in assignment (expression has type "None", variable has type Module)  [assignment]
xarray/core/_reductions.py:15: error: Skipping analyzing "flox": module is installed, but missing library stubs or py.typed marker  [import]
xarray/core/_reductions.py:15: note: See https://mypy.readthedocs.io/en/stable/running_mypy.html#missing-imports
xarray/core/groupby.py:584: error: Skipping analyzing "flox.xarray": module is installed, but missing library stubs or py.typed marker  [import]
xarray/core/dataset.py:113: error: Cannot assign to a type  [misc]
xarray/core/dataset.py:113: error: Incompatible types in assignment (expression has type "None", variable has type "Type[Delayed]")  [assignment]
xarray/core/dataset.py:1761: error: Incompatible return value type (got "Union[Tuple[ArrayWriter, AbstractDataStore], bytes, Delayed, None]", expected "Union[bytes, Delayed, None]")  [return-value]
xarray/core/dataarray.py:70: error: Cannot assign to a type  [misc]
xarray/core/dataarray.py:70: error: Incompatible types in assignment (expression has type "None", variable has type "Type[Delayed]")  [assignment]
xarray/core/computation.py:1864: error: Overloaded function implementation cannot satisfy signature 2 due to inconsistencies in how they use type variables  [misc]
xarray/core/computation.py:1864: error: Overloaded function implementation cannot satisfy signature 3 due to inconsistencies in how they use type variables  [misc]
xarray/core/_typed_ops.pyi:24: error: Cannot assign to a type  [misc]
xarray/core/_typed_ops.pyi:24: error: Incompatible types in assignment (expression has type "Type[ndarray[Any, Any]]", variable has type "Type[Array]")  [assignment]
xarray/backends/api.py:38: error: Cannot assign to a type  [misc]
xarray/backends/api.py:38: error: Incompatible types in assignment (expression has type "None", variable has type "Type[Delayed]")  [assignment]
xarray/tests/test_computation.py:2019: error: No overload variant of "polyval" matches argument types "Any", "Any"  [call-overload]
xarray/tests/test_computation.py:2019: note: Possible overload variants:
xarray/tests/test_computation.py:2019: note:     def polyval(coord: DataArray, coeffs: DataArray, degree_dim: Hashable) -> DataArray
xarray/tests/test_computation.py:2019: note:     def [T_Xarray in (DataArray, Dataset)] polyval(coord: T_Xarray, coeffs: Dataset, degree_dim: Hashable) -> Dataset
xarray/tests/test_computation.py:2019: note:     def [T_Xarray in (DataArray, Dataset)] polyval(coord: Dataset, coeffs: T_Xarray, degree_dim: Hashable) -> Dataset
xarray/tests/test_testing.py:13: error: Cannot infer type of lambda  [misc]
xarray/tests/test_testing.py:13: error: Incompatible types in assignment (expression has type "Callable[[Any], Any]", variable has type "Callable[[Any, Any, Any, Any, Any, Any, Any, Any, Any], Any]")  [assignment]

[Illviljan]

Yay, mypy is passing now after flox updates. :)

[aulemahal]

Hi @dcherian and @andersy005, thanks a lot for this PR and for flox. I'm currently trying this with our climate indices library (xclim) and it works well! I get a clear improvement in performance,

In xclim's indicators, all groupby operations are time resampling ops and the best performance is achieved with method='blockwise', by far. Is there a plan to let users control the default method through something like xr.set_options ?

Like, use_flox could accept False, True or any method name? (The default still being map-reduce) ? What do you think?

[dcherian]

Thanks for testing it out! I was going to ping xclim when this finally got merged. Presumably you haven't found any bugs?

---

You can pass method as .mean(..., method=...). Clearly this needs docs :)

We could actually consider adding flox_kwargs to the groupby constructor since a method is really only dependent on the distribution of group labels across the chunks. Right now, I'd just like this to get merged :)

For resampling-type, we are using cohorts by default which generalizes to blockwise when applicable but is slower at graph-construction time. Note you can only blockwise if all members of a group are in a single block. So if you are resampling to yearly but a year of data occupies multiple chunks, you want "cohorts", not "blockwise".

[aulemahal]

Presumably you haven't found any bugs?

Not for now,, but I haven't done much and it wasn't very complex.

You can pass method as .mean(..., method=...)

Yes! I saw that, but I did search a bit. In fact, my idea was triggered partly because I'm lazy. My suggestion was to do something alike what xarray does with keep_attrs : a global/context option that changes the default value of method, so it doesn't need to be explicitly added everywhere.

Note you can only blockwise if all members of a group are in a single block.

Oh! I didn't catch that. My test data was already appropriately chunked.

Right now, I'd just like this to get merged :)

Haha, I understand. My idea, if pertinent, could be done in another PR.

[dcherian]

a global/context option that changes the default value of method

Unfortunately the optimal method depends on distribution of group labels across chunks, so a global option doesn't make sense. It would make sense to create a method="auto" and use that but it doesn't exist yet ("cohorts" is closest)

[andersy005]

👏🏽 Glad to see this brought to completion! In my opinion, this is ready to be merged into main to allow folks to try it out. thank you, @dcherian, @Illviljan.

Does anyone know how to fix the benchmark CI? It appears that it's breaking during the setup

[shoyer]

👍 this looks great to me!

[max-sixty]

Congratulations @dcherian & @Illviljan & [edit] @andersy005 !!!

[dcherian]

and @andersy005 !