Skip to content

Commit

Permalink
fix
Browse files Browse the repository at this point in the history
  • Loading branch information
dcherian committed Sep 15, 2024
1 parent 3ce5183 commit ba87150
Showing 1 changed file with 3 additions and 1 deletion.
4 changes: 3 additions & 1 deletion flox/core.py
Original file line number Diff line number Diff line change
Expand Up @@ -642,6 +642,7 @@ def rechunk_for_blockwise(array: DaskArray, axis: T_Axis, labels: np.ndarray) ->
DaskArray
Rechunked array
"""
# TODO: this should be unnecessary?
labels = factorize_((labels,), axes=())[0]
chunks = array.chunks[axis]
newchunks = _get_optimal_chunks_for_groups(chunks, labels)
Expand Down Expand Up @@ -2623,7 +2624,8 @@ def groupby_reduce(

partial_agg = partial(dask_groupby_agg, **kwargs)

if method == "blockwise" and by_.ndim == 1:
# if preferred method is already blockwise, no need to rechunk
if preferred_method != "blockwise" and method == "blockwise" and by_.ndim == 1:
array = rechunk_for_blockwise(array, axis=-1, labels=by_)

result, groups = partial_agg(
Expand Down

0 comments on commit ba87150

Please sign in to comment.