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.

Sign up for GitHub

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

Concatenate small input chunks before P2P rechunking #8832

Merged
merged 22 commits into from
Aug 23, 2024
Merged

Concatenate small input chunks before P2P rechunking #8832

Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
6fd92b8
Use task-based rechunking to prechunk along partial boundaries
hendrikmakait Aug 13, 2024
5b85a65
Move slicing_is_necessary
hendrikmakait Aug 13, 2024
dac50a9
Add tests
hendrikmakait Aug 13, 2024
d8a9c8d
Minor
hendrikmakait Aug 13, 2024
c23fd8f
REVERT ME: Adjust CI repository
hendrikmakait Aug 15, 2024
81760fa
Minor
hendrikmakait Aug 15, 2024
6f7a107
Concatenate using task-based shuffling before P2P
hendrikmakait Aug 14, 2024
2ae27fd
Merge branch 'main' into use-rechunking-to-prerechunk
hendrikmakait Aug 16, 2024
1b7941d
Merge branch 'main' into use-rechunking-to-prerechunk
hendrikmakait Aug 20, 2024
d269e03
Revert "REVERT ME: Adjust CI repository"
hendrikmakait Aug 20, 2024
24bafbe
Fix merge (WIP)
hendrikmakait Aug 20, 2024
1db2284
Fix merge (2/2)
hendrikmakait Aug 20, 2024
100fa43
Fix concat logic
hendrikmakait Aug 20, 2024
8035a4a
Adjust tests
hendrikmakait Aug 20, 2024
2c79a6e
Add tests
hendrikmakait Aug 21, 2024
73dfd99
Add test
hendrikmakait Aug 21, 2024
3a119a5
Add tests and fix ordering
hendrikmakait Aug 22, 2024
54f234e
Adopt algorithm from task-based rechunking
hendrikmakait Aug 22, 2024
e4fdfdb
formatting
hendrikmakait Aug 23, 2024
34f6ee1
Update distributed/shuffle/_rechunk.py
hendrikmakait Aug 23, 2024
d86ae35
Minor
hendrikmakait Aug 23, 2024
5487c23
Add docs
hendrikmakait Aug 23, 2024
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
Prev Previous commit
Next Next commit
Fix merge (2/2)
  • Loading branch information
@hendrikmakait
hendrikmakait committed Aug 20, 2024
commit 1db2284beb6663b53a589980d8c383a45ac92db0
6 changes: 4 additions & 2 deletions distributed/shuffle/_rechunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -451,8 +451,8 @@ def _calculate_prechunking(
old_to_new_axis = _old_to_new[axis_index]
old_axis = old_chunks[axis_index]
split_axis = []
partial_chunks = []
for slice_ in slices:
partial_chunks = []
first_new_chunk = slice_.start
first_old_chunk, first_old_slice = old_to_new_axis[first_new_chunk][0]
last_new_chunk = slice_.stop - 1
Expand Down Expand Up @@ -484,7 +484,9 @@ def _calculate_prechunking(

partial_chunks.append(chunk_size)
split_axis.append(partial_chunks)

partial_chunks = []
if partial_chunks:
split_axis.append(partial_chunks)
split_axes.append(split_axis)

has_nans = (any(math.isnan(y) for y in x) for x in old_chunks)
Expand Down
4 changes: 2 additions & 2 deletions distributed/shuffle/tests/test_rechunk.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -1355,7 +1355,7 @@ async def test_partial_rechunk_taskgroups(c, s):
],
)
def test_calculate_prechunking_1d(old, new, expected):
actual = _calculate_prechunking(old, new)
actual = _calculate_prechunking(old, new, np.dtype)
assert actual == expected


Expand All @@ -1375,5 +1375,5 @@ def test_calculate_prechunking_1d(old, new, expected):
],
)
def test_calculate_prechunking_2d(old, new, expected):
actual = _calculate_prechunking(old, new)
actual = _calculate_prechunking(old, new, np.dtype(np.int16))
assert actual == expected
Loading