[data] fix map groups don't break down blocks #58988

iamjustinhsu · 2025-11-25T23:29:38Z

Description

Fixes #58603. map_group assumes that all partitions fix in one block. However, we my PR broke this behavior by breaking down blocks. To address this, I added a data context variable that is implicitly set to False to breakdown blocks. If True, block size will be preserved

Related issues

Link related issues: "Fixes #1234", "Closes #1234", or "Related to #1234".

Additional information

Optional: Add implementation details, API changes, usage examples, screenshots, etc.

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

python/ray/data/grouped_data.py

gemini-code-assist

Code Review

The PR introduces a new context flag _preserve_hash_shuffle_finalize_blocks to prevent block breakdown during hash shuffle finalization, which is necessary for map_groups. The overall approach is sound, but I've found a critical logic issue in the implementation and some inconsistencies in the documentation and tests that need to be addressed.

There's a logical error in hash_shuffle.py where an or is used instead of an and, which could cause the fix to not work as intended when target_max_block_size is set.
The documentation and comments for the new flag in context.py contradict the implementation's behavior for map_groups.
The new test for map_groups has a misleading name and comments.

I've left specific comments with suggestions for these points. Once these are addressed, the PR should be in good shape.

python/ray/data/_internal/execution/operators/hash_shuffle.py

python/ray/data/context.py

python/ray/data/tests/test_preserve_hash_shuffle_blocks.py

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

python/ray/data/_internal/planner/plan_all_to_all_op.py

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin · 2025-11-26T01:21:39Z

python/ray/data/_internal/execution/operators/hash_shuffle.py

        aggregator_ray_remote_args_override: Optional[Dict[str, Any]] = None,
        shuffle_progress_bar_name: Optional[str] = None,
        finalize_progress_bar_name: Optional[str] = None,
+        preserve_finalize_blocks: bool = False,


disallow_block_splitting

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

srinathk10

LGTM

cursor · 2025-11-26T01:28:02Z

python/ray/data/tests/test_preserve_hash_shuffle_blocks.py

+        ctx = DataContext.get_current()
+        # Very small to force splitting if enabled
+        ctx.target_max_block_size = 1
+        yield


Bug: Test fixture doesn't restore modified DataContext value

The setup fixture modifies ctx.target_max_block_size = 1 but never restores the original value after the test. The established pattern in conftest.py for such fixtures is to save the original value before modifying, then restore it after yield. Without restoration, the modified value of 1 persists and could cause test pollution affecting subsequent tests that depend on the default target_max_block_size value.

…/preserve-block-size

python/ray/data/tests/test_preserve_hash_shuffle_blocks.py

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin · 2025-11-26T18:52:37Z

python/ray/data/_internal/execution/operators/hash_shuffle.py

        target_max_block_size = self._data_context.target_max_block_size
        # None means the user wants to preserve the block distribution,
        # so we do not break the block down further.
-        if target_max_block_size is not None:
+        # Also check _disallow_block_splitting parameter.
+        if target_max_block_size is not None and not self._disallow_block_splitting:


Instead pass target_max_block_size as None when block splitting not allowed

alexeykudinkin · 2025-11-26T18:54:02Z

python/ray/data/_internal/execution/operators/hash_shuffle.py

+            # NOTE: This is set to True because num_partitions (aka, # of output blocks)
+            # must be preserved.
+            disallow_block_splitting=True,


Suggested change

# NOTE: This is set to True because num_partitions (aka, # of output blocks)

# must be preserved.

disallow_block_splitting=True,

# NOTE: In cases like ``groupby`` blocks can't be split as this might violate an invariant that all rows

# with the same key are in the same group (block)

disallow_block_splitting=True,

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

…/preserve-block-size

alexeykudinkin · 2025-12-02T18:04:02Z

python/ray/data/_internal/execution/operators/hash_shuffle.py

+        self._data_context = data_context.copy()
+        if disallow_block_splitting:
+            self._data_context.target_max_block_size = None


Sorry, i should have been more clear:

Don't patch the DC

Instead pass target_max_block_size directly into the aggregator (to avoid passing to overlapping configs DC and disallowing flag)

…llow_block_splitting=True Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

## Description `test_preserve_hash_shuffle_blocks` has been flaking consistently. To mitigate the flakiness, this PR bumps the test size from "small" to "medium". ``` [2025-12-06T07:06:32Z] //python/ray/data:test_preserve_hash_shuffle_blocks TIMEOUT in 3 out of 3 in 63.4s -- [2025-12-06T07:06:32Z] Stats over 3 runs: max = 63.4s, min = 60.1s, avg = 62.3s, dev = 1.6s [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test.log [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test_attempts/attempt_1.log [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test_attempts/attempt_2.log [2025-12-06T07:06:32Z] ``` See also #58988 Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>

## Description `test_preserve_hash_shuffle_blocks` has been flaking consistently. To mitigate the flakiness, this PR bumps the test size from "small" to "medium". This is a follow-up to #59256, which accidentally bumped the wrong test. ``` [2025-12-06T07:06:32Z] //python/ray/data:test_preserve_hash_shuffle_blocks TIMEOUT in 3 out of 3 in 63.4s -- [2025-12-06T07:06:32Z] Stats over 3 runs: max = 63.4s, min = 60.1s, avg = 62.3s, dev = 1.6s [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test.log [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test_attempts/attempt_1.log [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test_attempts/attempt_2.log [2025-12-06T07:06:32Z] ``` See also #58988 Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>

…ect#59399) ## Description `test_preserve_hash_shuffle_blocks` has been flaking consistently. To mitigate the flakiness, this PR bumps the test size from "small" to "medium". This is a follow-up to ray-project#59256, which accidentally bumped the wrong test. ``` [2025-12-06T07:06:32Z] //python/ray/data:test_preserve_hash_shuffle_blocks TIMEOUT in 3 out of 3 in 63.4s -- [2025-12-06T07:06:32Z] Stats over 3 runs: max = 63.4s, min = 60.1s, avg = 62.3s, dev = 1.6s [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test.log [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test_attempts/attempt_1.log [2025-12-06T07:06:32Z] /root/.cache/bazel/_bazel_root/1df605deb6d24fc8068f6e25793ec703/execroot/io_ray/bazel-out/k8-opt/testlogs/python/ray/data/test_preserve_hash_shuffle_blocks/test_attempts/attempt_2.log [2025-12-06T07:06:32Z] ``` See also ray-project#58988 Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>

[data] fix map groups don't break down blocks

c4a9046

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

iamjustinhsu requested a review from a team as a code owner November 25, 2025 23:29

fix

3d5ab27

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

cursor bot reviewed Nov 25, 2025

View reviewed changes

python/ray/data/grouped_data.py Outdated Show resolved Hide resolved

gemini-code-assist bot reviewed Nov 25, 2025

View reviewed changes

iamjustinhsu added 4 commits November 25, 2025 16:05

fix

15e0905

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

comments

49d74ab

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

update

e0d81c0

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

fix

7c03899

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

cursor bot reviewed Nov 26, 2025

View reviewed changes

python/ray/data/_internal/planner/plan_all_to_all_op.py Outdated Show resolved Hide resolved

iamjustinhsu added 5 commits November 25, 2025 17:15

fix

e2178bc

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

remove

f20a181

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

comment

339d3c2

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

fix

9465fce

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

fix

bb5623e

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin reviewed Nov 26, 2025

View reviewed changes

done

2818ee1

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

srinathk10 approved these changes Nov 26, 2025

View reviewed changes

cursor bot reviewed Nov 26, 2025

View reviewed changes

ray-gardener bot added the data Ray Data-related issues label Nov 26, 2025

Merge branch 'master' of https://github.com/ray-project/ray into jhsu…

0870f93

…/preserve-block-size

cursor bot reviewed Nov 26, 2025

View reviewed changes

python/ray/data/tests/test_preserve_hash_shuffle_blocks.py Outdated Show resolved Hide resolved

fix

ad3d7fb

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

srinathk10 added the go add ONLY when ready to merge, run all tests label Nov 26, 2025

alexeykudinkin reviewed Nov 26, 2025

View reviewed changes

iamjustinhsu added 2 commits November 26, 2025 10:59

fix

c8a2f0c

Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

Merge branch 'master' of https://github.com/ray-project/ray into jhsu…

8c19656

…/preserve-block-size

alexeykudinkin reviewed Dec 2, 2025

View reviewed changes

prevent overlapping concepts: target_max_block_size is None, and disa…

20d9685

…llow_block_splitting=True Signed-off-by: iamjustinhsu <jhsu@anyscale.com>

alexeykudinkin approved these changes Dec 2, 2025

View reviewed changes

alexeykudinkin merged commit 94ef5ff into ray-project:master Dec 2, 2025
6 checks passed

iamjustinhsu deleted the jhsu/preserve-block-size branch December 2, 2025 21:02

bveeramani mentioned this pull request Dec 8, 2025

[Data] Increase test_preserve_hash_shuffle_blocks timeout #59256

Merged

bveeramani mentioned this pull request Dec 12, 2025

[Data] Increase test_preserve_hash_shuffle_blocks timeout #59399

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

[data] fix map groups don't break down blocks #58988

[data] fix map groups don't break down blocks #58988

Uh oh!

iamjustinhsu commented Nov 25, 2025

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alexeykudinkin Nov 26, 2025

Uh oh!

srinathk10 left a comment

Uh oh!

cursor bot Nov 26, 2025

Uh oh!

Uh oh!

alexeykudinkin Nov 26, 2025

Uh oh!

alexeykudinkin Nov 26, 2025

Uh oh!

alexeykudinkin Dec 2, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

[data] fix map groups don't break down blocks #58988

[data] fix map groups don't break down blocks #58988

Uh oh!

Conversation

iamjustinhsu commented Nov 25, 2025

Description

Related issues

Additional information

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alexeykudinkin Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

srinathk10 left a comment

Choose a reason for hiding this comment

Uh oh!

cursor bot Nov 26, 2025

Choose a reason for hiding this comment

Bug: Test fixture doesn't restore modified DataContext value

Uh oh!

Uh oh!

alexeykudinkin Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

alexeykudinkin Nov 26, 2025

Choose a reason for hiding this comment

Uh oh!

alexeykudinkin Dec 2, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants