[Data] `map_batches` support limit_pushdown by owenowenisme · Pull Request #57880 · ray-project/ray

owenowenisme · 2025-10-18T16:21:44Z

Description

This PR adds a ‎preserve_row option to ‎map_batches. When ‎preserve_row is true, the limit operator can be pushed down through this ‎map_batches call for optimization.

Note: ‎map_group is built on ‎map_batches, but limit pushdown support for ‎map_group is out of scope for this PR, so ‎preserve_row_count is set to false for it.

Related issues

Additional information

Signed-off-by: You-Cheng Lin <[email protected]>

goutamvenkat-anyscale · 2025-10-21T00:45:55Z

python/ray/data/_internal/logical/operators/map_operator.py

        fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
        min_rows_per_bundled_input: Optional[int] = None,
        compute: Optional[ComputeStrategy] = None,
+        preserve_row_count: bool = False,


preserves_row_count

goutamvenkat-anyscale · 2025-10-21T00:46:18Z

python/ray/data/tests/test_execution_optimizer_limit_pushdown.py

    assert result_with == expected


+def test_limit_pushdown_preserve_row_count_with_map_batches(


Can we make this a parameterized test?

goutamvenkat-anyscale · 2025-10-21T00:51:03Z

python/ray/data/dataset.py

                worker.
            memory: The heap memory in bytes to reserve for each parallel map worker.
            concurrency: This argument is deprecated. Use ``compute`` argument.
+            preserve_row_count: Set to True only if the UDF always emits the same number of records it receives (no drops or duplicates). When true, the optimizer can push downstream limits past this transform for better pruning.


For the 2nd sentence: When set to True, the logical optimizer, in the presence of a limit(limit=k), will only scan k rows prior to executing the UDF, thereby saving on compute resources.

goutamvenkat-anyscale

lgtm. Just a few cosmetic changes

Signed-off-by: You-Cheng Lin <[email protected]>

python/ray/data/_internal/logical/operators/map_operator.py

python/ray/data/dataset.py

…support-limit

Signed-off-by: You-Cheng Lin <[email protected]>

## Description This PR adds a ‎`preserve_row` option to ‎`map_batches`. When ‎`preserve_row` is true, the limit operator can be pushed down through this ‎`map_batches` call for optimization. Note: ‎`map_group` is built on ‎`map_batches`, but limit pushdown support for ‎`map_group` is out of scope for this PR, so ‎`preserve_row_count` is set to false for it. ## Related issues ## Additional information --------- Signed-off-by: You-Cheng Lin <[email protected]> Signed-off-by: You-Cheng Lin <[email protected]> Co-authored-by: You-Cheng Lin <[email protected]> Signed-off-by: xgui <[email protected]>

## Description This PR adds a ‎`preserve_row` option to ‎`map_batches`. When ‎`preserve_row` is true, the limit operator can be pushed down through this ‎`map_batches` call for optimization. Note: ‎`map_group` is built on ‎`map_batches`, but limit pushdown support for ‎`map_group` is out of scope for this PR, so ‎`preserve_row_count` is set to false for it. ## Related issues ## Additional information --------- Signed-off-by: You-Cheng Lin <[email protected]> Signed-off-by: You-Cheng Lin <[email protected]> Co-authored-by: You-Cheng Lin <[email protected]>

## Description This PR adds a ‎`preserve_row` option to ‎`map_batches`. When ‎`preserve_row` is true, the limit operator can be pushed down through this ‎`map_batches` call for optimization. Note: ‎`map_group` is built on ‎`map_batches`, but limit pushdown support for ‎`map_group` is out of scope for this PR, so ‎`preserve_row_count` is set to false for it. ## Related issues ## Additional information --------- Signed-off-by: You-Cheng Lin <[email protected]> Signed-off-by: You-Cheng Lin <[email protected]> Co-authored-by: You-Cheng Lin <[email protected]> Signed-off-by: Aydin Abiar <[email protected]>

## Description This PR adds a ‎`preserve_row` option to ‎`map_batches`. When ‎`preserve_row` is true, the limit operator can be pushed down through this ‎`map_batches` call for optimization. Note: ‎`map_group` is built on ‎`map_batches`, but limit pushdown support for ‎`map_group` is out of scope for this PR, so ‎`preserve_row_count` is set to false for it. ## Related issues ## Additional information --------- Signed-off-by: You-Cheng Lin <[email protected]> Signed-off-by: You-Cheng Lin <[email protected]> Co-authored-by: You-Cheng Lin <[email protected]> Signed-off-by: Future-Outlier <[email protected]>

This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([#39486](#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([#57880](#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR #60448](#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]>

…ct#60756) This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (ray-project#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (ray-project#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (ray-project#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([ray-project#39486](ray-project#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([ray-project#57880](ray-project#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR ray-project#60448](ray-project#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]> Signed-off-by: tiennguyentony <[email protected]>

…ct#60756) This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (ray-project#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (ray-project#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (ray-project#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([ray-project#39486](ray-project#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([ray-project#57880](ray-project#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR ray-project#60448](ray-project#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]>

This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([#39486](#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([#57880](#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR #60448](#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]> Signed-off-by: elliot-barn <[email protected]>

This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([#39486](#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([#57880](#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR #60448](#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]>

…ct#60756) This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (ray-project#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (ray-project#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (ray-project#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([ray-project#39486](ray-project#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([ray-project#57880](ray-project#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR ray-project#60448](ray-project#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]>

…ct#60756) This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (ray-project#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (ray-project#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (ray-project#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([ray-project#39486](ray-project#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([ray-project#57880](ray-project#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR ray-project#60448](ray-project#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]> Signed-off-by: Adel Nour <[email protected]>

…ct#60756) This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (ray-project#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (ray-project#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (ray-project#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([ray-project#39486](ray-project#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([ray-project#57880](ray-project#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR ray-project#60448](ray-project#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]>

## Description This PR adds a ‎`preserve_row` option to ‎`map_batches`. When ‎`preserve_row` is true, the limit operator can be pushed down through this ‎`map_batches` call for optimization. Note: ‎`map_group` is built on ‎`map_batches`, but limit pushdown support for ‎`map_group` is out of scope for this PR, so ‎`preserve_row_count` is set to false for it. ## Related issues ## Additional information --------- Signed-off-by: You-Cheng Lin <[email protected]> Signed-off-by: You-Cheng Lin <[email protected]> Co-authored-by: You-Cheng Lin <[email protected]> Signed-off-by: peterxcli <[email protected]>

…ct#60756) This PR updates the operator fusion rule to fuse `MapBatches` even if they modify the row counts. The intention of this PR is to preserve the historical operator fusion behavior and avoid introducing regressions. For more details, see the timeline below. --- ### Timeline of Changes | Date | Event | Description | | :--- | :--- | :--- | | **June 8, 2023** | **Limit pushdown added** | Added limit pushdown and a property to `MapBatches` incorrectly stating it doesn't modify row counts. (ray-project#35950) | | **June 27, 2023** | **Limit pushdown disabled** | Rule disabled because it incorrectly pushed limits past UDFs that modified row counts. (ray-project#36831) | | **April 28, 2025** | **Fusion restricted** | Added logic to stop fusing operators that modify row counts when the downstream has a batch size. `MapBatches` stayed fused only because of its incorrect property (ray-project#52570). | | **July 8, 2025** | **Limit pushdown re-enabled with special case** | Re-enabled with a special case to prevent pushing limits past `MapBatches`. ([ray-project#39486](ray-project#39486)) | | **Oct 24, 2025** | **Special case removed** | Special case removed, re-introducing the bug where limits are pushed past `MapBatches`. ([ray-project#57880](ray-project#57880)) | | **Feb 2, 2026** | **Property Fix** | Updated `MapBatches` to correctly report it modifies rows by default. This fixed the pushdown bug but broke fusion logic. ([PR ray-project#60448](ray-project#60448)) | | **Feb 4, 2026** | (This PR) | Add a special-case to preserve the historical `MapBatches` fusion behavior | --- <a href="proxy.php?url=https://cursor.com/dashboard?tab=bugbot">Cursor Bugbot</a> reviewed your changes and found no issues for commit d99e7b1 --------- Signed-off-by: Balaji Veeramani <[email protected]> Signed-off-by: peterxcli <[email protected]>

update

f97b26e

Signed-off-by: You-Cheng Lin <[email protected]>

owenowenisme added the go add ONLY when ready to merge, run all tests label Oct 18, 2025

owenowenisme changed the title ~~[Data map_batches support limit_pushdown~~ [Data] map_batches support limit_pushdown Oct 18, 2025

set map_group presrve row to false

771122f

Signed-off-by: You-Cheng Lin <[email protected]>

owenowenisme marked this pull request as ready for review October 20, 2025 05:53

owenowenisme requested a review from a team as a code owner October 20, 2025 05:53

ray-gardener bot added the data Ray Data-related issues label Oct 20, 2025

goutamvenkat-anyscale reviewed Oct 21, 2025

View reviewed changes

goutamvenkat-anyscale requested changes Oct 21, 2025

View reviewed changes

resolve comments

169d4d2

Signed-off-by: You-Cheng Lin <[email protected]>

owenowenisme requested a review from goutamvenkat-anyscale October 21, 2025 04:11

goutamvenkat-anyscale approved these changes Oct 21, 2025

View reviewed changes

alexeykudinkin reviewed Oct 24, 2025

View reviewed changes

python/ray/data/_internal/logical/operators/map_operator.py Outdated Show resolved Hide resolved

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

owenowenisme added 2 commits October 24, 2025 08:46

Merge remote-tracking branch 'upstream/master' into data/map-batches-…

f173103

…support-limit

rename preserve_row_count to udf_modifying_row_count

8d0416c

Signed-off-by: You-Cheng Lin <[email protected]>

This comment was marked as outdated.

Sign in to view

make map batches fusion use udf_modifying_row_count

93e8111

Signed-off-by: You-Cheng Lin <[email protected]>

owenowenisme force-pushed the data/map-batches-support-limit branch from 2bc26a1 to 93e8111 Compare October 24, 2025 04:26

update fusion test

3a8c54d

Signed-off-by: You-Cheng Lin <[email protected]>

owenowenisme requested a review from alexeykudinkin October 24, 2025 16:22

alexeykudinkin approved these changes Oct 24, 2025

View reviewed changes

alexeykudinkin merged commit 978ca10 into ray-project:master Oct 24, 2025
6 checks passed

bveeramani mentioned this pull request Feb 4, 2026

[Data] Fuse MapBatches even if they modify the row count #60756

Merged

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

[Data] `map_batches` support limit_pushdown#57880

[Data] `map_batches` support limit_pushdown#57880
alexeykudinkin merged 7 commits intoray-project:masterfrom
owenowenisme:data/map-batches-support-limit

owenowenisme commented Oct 18, 2025 •

edited

Loading

Uh oh!

goutamvenkat-anyscale Oct 21, 2025

Uh oh!

goutamvenkat-anyscale Oct 21, 2025

Uh oh!

goutamvenkat-anyscale Oct 21, 2025

Uh oh!

goutamvenkat-anyscale left a comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

		assert result_with == expected


		def test_limit_pushdown_preserve_row_count_with_map_batches(

Conversation

owenowenisme commented Oct 18, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Description

Related issues

Additional information

Uh oh!

goutamvenkat-anyscale Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale Oct 21, 2025

Choose a reason for hiding this comment

Uh oh!

goutamvenkat-anyscale left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

owenowenisme commented Oct 18, 2025 •

edited

Loading