[Data] ArrowInvalid error when you backfill missing fields from map tasks by machichima · Pull Request #60643 · ray-project/ray

machichima · 2026-02-01T11:07:15Z

Description

Try type casting if struct field types mismatch when backfilling missing fields

Related issues

Closes #60628

Additional information

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

Signed-off-by: machichima <nary12321@gmail.com>

gemini-code-assist

Code Review

This pull request addresses an ArrowInvalid error that occurs when backfilling missing fields in struct columns, particularly from map tasks. The issue arises from type mismatches between struct fields in different blocks (e.g., int64 vs. float64) that are not handled after schema unification. The proposed change correctly identifies these type discrepancies and explicitly casts the array to the unified field type. The implementation is clean, includes robust error handling with an informative ValueError, and effectively resolves the bug. The changes look good to me.

Signed-off-by: machichima <nary12321@gmail.com>

machichima · 2026-02-01T11:16:20Z

@bveeramani PTAL, thank you!

Signed-off-by: machichima <nary12321@gmail.com>

alexeykudinkin · 2026-02-02T21:15:50Z

@machichima can you please take a look at build failures?

bveeramani · 2026-02-02T21:25:21Z

python/ray/data/tests/test_map_batches.py

+
+    ds = ray.data.range(4, override_num_blocks=1)
+    ds = ds.map_batches(generator_fn, batch_size=4)
+    result = ds.materialize()


Nit: I think this materialize is redundant. take_all() already materializes the dataset

bveeramani · 2026-02-02T21:26:34Z

python/ray/data/tests/test_map_batches.py

+    # Rows 0 and 2 should have int cast to float, with c=None
+    assert rows[0]["data"] == {"a": 1.0, "b": "hello", "c": None}
+    assert rows[2]["data"] == {"a": 1.0, "b": "hello", "c": None}
+
+    # Rows 1 and 3 should have float a, with b=None
+    assert rows[1]["data"] == {"a": 1.5, "b": None, "c": 100}
+    assert rows[3]["data"] == {"a": 1.5, "b": None, "c": 100}


Though I don't think this is possible given the current Ray Data implementation, this row ordering isn't guaranteed by the interface. This has historically been the cause of a lot of Ray Data's flaky tests.

Could you refactor this test so that it doesn't depend on a particular test ordering?

Updated in 773cd0d

bveeramani · 2026-02-02T21:27:57Z

python/ray/data/tests/test_map_batches.py

    assert len(output) == len(expected_output), (len(output), len(expected_output))


+def test_map_batches_struct_field_type_divergence(shutdown_only):


In addition to this E2E test, I think we should add a unit test for this bug as well (maybe at the concat function layer of abstraction). Unit tests are not only much faster to run, but also serve as documentation

Added unit test in: 6a35990

python/ray/data/tests/test_map_batches.py

Signed-off-by: machichima <nary12321@gmail.com>

bveeramani

Nice

bveeramani · 2026-02-03T21:44:46Z

@machichima looks like there's a minor import error. Would you mind fixing it? Think we should be good to land after

Signed-off-by: machichima <nary12321@gmail.com>

cursor

Cursor Bugbot has reviewed your changes and found 1 potential issue.

cursor · 2026-02-04T00:34:06Z

python/ray/data/tests/test_map_batches.py

+    ds = ds.map_batches(generator_fn, batch_size=4)
+    result = ds.materialize()
+
+    rows = result.take_all()


Redundant materialize() call before take_all()

Low Severity

The materialize() call on line 824 is redundant because take_all() on line 826 already materializes the dataset. This was noted in the PR review comments. The extra call doesn't cause incorrect behavior but adds unnecessary overhead and clutters the test code.

machichima · 2026-02-04T00:50:12Z

@machichima looks like there's a minor import error. Would you mind fixing it? Think we should be good to land after

Done!

…asks (ray-project#60643) ## Description Try type casting if struct field types mismatch when backfilling missing fields ## Related issues Closes ray-project#60628 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: machichima <nary12321@gmail.com> Co-authored-by: Balaji Veeramani <balaji@anyscale.com> Signed-off-by: tiennguyentony <46289799+tiennguyentony@users.noreply.github.com>

…asks (ray-project#60643) ## Description Try type casting if struct field types mismatch when backfilling missing fields ## Related issues Closes ray-project#60628 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: machichima <nary12321@gmail.com> Co-authored-by: Balaji Veeramani <balaji@anyscale.com>

…asks (#60643) ## Description Try type casting if struct field types mismatch when backfilling missing fields ## Related issues Closes #60628 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: machichima <nary12321@gmail.com> Co-authored-by: Balaji Veeramani <balaji@anyscale.com> Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>

…asks (#60643) ## Description Try type casting if struct field types mismatch when backfilling missing fields ## Related issues Closes #60628 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: machichima <nary12321@gmail.com> Co-authored-by: Balaji Veeramani <balaji@anyscale.com>

…asks (ray-project#60643) ## Description Try type casting if struct field types mismatch when backfilling missing fields ## Related issues Closes ray-project#60628 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: machichima <nary12321@gmail.com> Co-authored-by: Balaji Veeramani <balaji@anyscale.com> Signed-off-by: Adel Nour <ans9868@nyu.edu>

…asks (ray-project#60643) ## Description Try type casting if struct field types mismatch when backfilling missing fields ## Related issues Closes ray-project#60628 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: machichima <nary12321@gmail.com> Co-authored-by: Balaji Veeramani <balaji@anyscale.com>

…asks (ray-project#60643) ## Description Try type casting if struct field types mismatch when backfilling missing fields ## Related issues Closes ray-project#60628 ## Additional information > Optional: Add implementation details, API changes, usage examples, screenshots, etc. --------- Signed-off-by: machichima <nary12321@gmail.com> Co-authored-by: Balaji Veeramani <balaji@anyscale.com> Signed-off-by: peterxcli <peterxcli@gmail.com>

fix: try type casting if field type mismatch

4dccb02

Signed-off-by: machichima <nary12321@gmail.com>

machichima requested a review from a team as a code owner February 1, 2026 11:07

gemini-code-assist bot reviewed Feb 1, 2026

View reviewed changes

test: check type casting

861d061

Signed-off-by: machichima <nary12321@gmail.com>

Merge branch 'master' into 60628-arrow-invalid-err

0a741ae

ray-gardener bot added data Ray Data-related issues community-contribution Contributed by the community labels Feb 1, 2026

refactor: precommit

a65c8d8

Signed-off-by: machichima <nary12321@gmail.com>

alexeykudinkin approved these changes Feb 2, 2026

View reviewed changes

alexeykudinkin added the go add ONLY when ready to merge, run all tests label Feb 2, 2026

bveeramani reviewed Feb 2, 2026

View reviewed changes

Merge branch 'master' into 60628-arrow-invalid-err

8c2cda3

cursor bot reviewed Feb 3, 2026

View reviewed changes

python/ray/data/tests/test_map_batches.py Show resolved Hide resolved

machichima and others added 4 commits February 3, 2026 20:06

test: make test result deterministic to prevent flaky

773cd0d

Signed-off-by: machichima <nary12321@gmail.com>

Merge branch 'master' into 60628-arrow-invalid-err

8bdfa20

test: skip test for old pyarrow version

dd8a662

Signed-off-by: machichima <nary12321@gmail.com>

test: add unit test using concat

6a35990

Signed-off-by: machichima <nary12321@gmail.com>

bveeramani approved these changes Feb 3, 2026

View reviewed changes

fix: import error

ec2a812

Signed-off-by: machichima <nary12321@gmail.com>

cursor bot reviewed Feb 4, 2026

View reviewed changes

bveeramani enabled auto-merge (squash) February 4, 2026 00:52

Merge branch 'master' into 60628-arrow-invalid-err

4c52650

github-actions bot disabled auto-merge February 4, 2026 02:12

bveeramani merged commit 4cf7794 into ray-project:master Feb 4, 2026
6 checks passed

		assert len(output) == len(expected_output), (len(output), len(expected_output))


		def test_map_batches_struct_field_type_divergence(shutdown_only):

Conversation

machichima commented Feb 1, 2026

Description

Related issues

Additional information

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

machichima commented Feb 1, 2026

Uh oh!

alexeykudinkin commented Feb 2, 2026

Uh oh!

bveeramani Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

bveeramani Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

machichima Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

bveeramani Feb 2, 2026

Choose a reason for hiding this comment

Uh oh!

machichima Feb 3, 2026

Choose a reason for hiding this comment

Uh oh!

Uh oh!

bveeramani left a comment

Choose a reason for hiding this comment

Uh oh!

bveeramani commented Feb 3, 2026

Uh oh!

cursor bot left a comment

Choose a reason for hiding this comment

Uh oh!

cursor bot Feb 4, 2026

Choose a reason for hiding this comment

Redundant materialize() call before take_all()

Uh oh!

machichima commented Feb 4, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants

Redundant `materialize()` call before `take_all()`