Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

use engine flox for ordered groups #266

Merged
merged 34 commits into from
Oct 15, 2023
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
34 commits
Select commit Hold shift + click to select a range
ef6bbc4
use engine flox for ordered groups
mathause Sep 29, 2023
e0ea569
Add issorted helper func
dcherian Oct 7, 2023
d4e30d8
Some fixes
dcherian Oct 7, 2023
c8000e3
In xarray too
dcherian Oct 7, 2023
68edd74
Merge remote-tracking branch 'upstream/main' into engine_none
dcherian Oct 7, 2023
c483a1a
formatting
dcherian Oct 7, 2023
3e4bae9
simplify
dcherian Oct 7, 2023
d14299a
Merge remote-tracking branch 'upstream/main' into engine_none
dcherian Oct 7, 2023
92eaf2c
retry
dcherian Oct 10, 2023
1a230eb
flox
dcherian Oct 10, 2023
22bd20c
minversion numabgg
dcherian Oct 10, 2023
699ecd9
cleanup
dcherian Oct 10, 2023
3d6413f
fix type
dcherian Oct 10, 2023
7061fda
Merge branch 'main' into engine_none
dcherian Oct 11, 2023
443fbc2
update gitignore
dcherian Oct 11, 2023
eedeb8e
add types
dcherian Oct 11, 2023
cc90509
Fix env?
dcherian Oct 11, 2023
56957a1
fix
dcherian Oct 11, 2023
16d4393
fix merge
dcherian Oct 11, 2023
f6262d2
cleanup
dcherian Oct 11, 2023
12dc816
Merge branch 'main' into engine_none
dcherian Oct 11, 2023
e13c2d9
[skip-ci] bench
dcherian Oct 11, 2023
a7cd95c
Merge branch 'main' into engine_none
dcherian Oct 11, 2023
46262f0
temporarily disable numbagg
dcherian Oct 12, 2023
ed8bd51
don't cache env
dcherian Oct 12, 2023
db0a6cc
Finally!
dcherian Oct 12, 2023
90cecf8
bugfix
dcherian Oct 14, 2023
f2e0aa6
Fix doctest
dcherian Oct 14, 2023
39c19b8
more fixes
dcherian Oct 14, 2023
7126e60
Fix CI
dcherian Oct 14, 2023
1243903
readd numbagg
dcherian Oct 14, 2023
524f540
Merge remote-tracking branch 'upstream/main' into engine_none
dcherian Oct 14, 2023
b70b0e6
Merge branch 'main' into engine_none
dcherian Oct 14, 2023
3534766
Fix.
dcherian Oct 15, 2023
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 16 additions & 6 deletions flox/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -1755,7 +1755,7 @@ def groupby_reduce(
dtype: np.typing.DTypeLike = None,
min_count: int | None = None,
method: T_Method = "map-reduce",
engine: T_Engine = "numpy",
engine: T_Engine = None,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Needs a docstring update.

reindex: bool | None = None,
finalize_kwargs: dict[Any, Any] | None = None,
) -> tuple[DaskArray, Unpack[tuple[np.ndarray | DaskArray, ...]]]: # type: ignore[misc] # Unpack not in mypy yet
Expand Down Expand Up @@ -1851,17 +1851,27 @@ def groupby_reduce(
xarray.xarray_reduce
"""

bys: T_Bys = tuple(np.asarray(b) if not is_duck_array(b) else b for b in by)
nby = len(bys)
by_is_dask = tuple(is_duck_dask_array(b) for b in bys)
any_by_dask = any(by_is_dask)

if engine is None:
dcherian marked this conversation as resolved.
Show resolved Hide resolved
# choose numpy per default
engine = "numpy"

if nby == 1 and not any_by_dask and bys[0].ndim == 1:
# maybe move to helper function
issorted = lambda arr: (arr[:-1] <= arr[1:]).all()
dcherian marked this conversation as resolved.
Show resolved Hide resolved
if not _is_arg_reduction(func) and issorted(bys[0]):
engine = "flox"

if engine == "flox" and _is_arg_reduction(func):
raise NotImplementedError(
"argreductions not supported for engine='flox' yet."
"Try engine='numpy' or engine='numba' instead."
)

bys: T_Bys = tuple(np.asarray(b) if not is_duck_array(b) else b for b in by)
nby = len(bys)
by_is_dask = tuple(is_duck_dask_array(b) for b in bys)
any_by_dask = any(by_is_dask)

if method in ["split-reduce", "cohorts"] and any_by_dask:
raise ValueError(f"method={method!r} can only be used when grouping by numpy arrays.")

Expand Down
2 changes: 1 addition & 1 deletion flox/xarray.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,7 +72,7 @@ def xarray_reduce(
fill_value=None,
dtype: np.typing.DTypeLike = None,
method: str = "map-reduce",
engine: str = "numpy",
engine: str = None,
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Again needs docstring update.

keep_attrs: bool | None = True,
skipna: bool | None = None,
min_count: int | None = None,
Expand Down
2 changes: 1 addition & 1 deletion tests/conftest.py
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
import pytest


@pytest.fixture(scope="module", params=["flox", "numpy", "numba"])
@pytest.fixture(scope="module", params=[None, "flox", "numpy", "numba"])
dcherian marked this conversation as resolved.
Show resolved Hide resolved
def engine(request):
if request.param == "numba":
try:
Expand Down
Loading