Skip to content

[Data] Improve handling of mismatched columns#53861

Merged
bveeramani merged 4 commits intomasterfrom
column-handling
Jun 17, 2025
Merged

[Data] Improve handling of mismatched columns#53861
bveeramani merged 4 commits intomasterfrom
column-handling

Conversation

@bveeramani
Copy link
Copy Markdown
Member

@bveeramani bveeramani commented Jun 16, 2025 •
edited
Loading

Why are these changes needed?

If a map transform function produces outputs with different column (e.g., you're reading a JSONL file where lines contain different keys), then Ray Data errors.

This PR updates the implementation to make it more robust to variation in column names.

Related issue number

Checks

  • I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
  • I've run scripts/format.sh to lint the changes in this PR.
  • I've included any doc changes needed for https://docs.ray.io/en/master/.
    • I've added any new APIs to the API Reference. For example, if I added a
      method in Tune, I've added it in doc/source/tune/api/ under the
      corresponding .rst file.
  • I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
  • Testing Strategy
    • Unit tests
    • Release tests
    • This PR is not tested :(

@bveeramani
Iminitl commit
027821e 
Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
Copilot AI review requested due to automatic review settings June 16, 2025 19:43
@bveeramani bveeramani requested a review from a team as a code owner June 16, 2025 19:43
Copilot AI reviewed Jun 16, 2025
View reviewed changes
Copy link
Copy Markdown
Contributor

Copilot AI left a comment

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Pull Request Overview

This PR improves the handling of mismatched columns by refining error reporting for filtering, extending tests for Arrow block building, and refactoring the internal TableBlockBuilder logic to allow heterogeneous row schemas.

  • Updates error message normalization in filtering tests
  • Adds tests for handling different column names in ArrowBlockBuilder operations
  • Refactors TableBlockBuilder to fill missing columns with None rather than raising an error

Reviewed Changes

Copilot reviewed 3 out of 3 changed files in this pull request and generated no comments.

File Description
python/ray/data/tests/test_map.py Updated error handling and normalization for filtered dataset error messages
python/ray/data/tests/test_arrow_block.py Added tests verifying correct behavior when merging rows/blocks with mismatched columns
python/ray/data/_internal/table_block.py Removed strict schema validation in favor of filling missing columns with None
Comments suppressed due to low confidence (3)

python/ray/data/tests/test_map.py:805

  • Ensure that using RuntimeError for filter expression failures is consistent with the overall error handling strategy. If this change in exception type is intentional, consider updating related documentation to clarify the API behavior.
with pytest.raises(RuntimeError) as exc_info:

python/ray/data/_internal/table_block.py:87

  • Consider adding a comment explaining that allowing mismatched columns and populating missing values with None is intentional and documenting how downstream operations should handle such cases.
# Fill in missing columns with None.

python/ray/data/tests/test_arrow_block.py:529

  • [nitpick] Consider extending this test with multiple rows per column to further validate the behavior of merging blocks with different column names.
expected_table = pa.Table.from_pydict({"col1": ["spam", None], "col2": [None, "foo"]})

@bveeramani
Fxi stuff:
0a54c1b 
Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
@bveeramani
Remove stuff
858e604 
Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
@bveeramani
Update stuff
4d0fadf 
Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
@bveeramani bveeramani enabled auto-merge (squash) June 16, 2025 22:11
@github-actions github-actions bot added the go add ONLY when ready to merge, run all tests label Jun 16, 2025
@bveeramani bveeramani merged commit eed7e02 into master Jun 17, 2025
7 checks passed
@bveeramani bveeramani deleted the column-handling branch June 17, 2025 01:13
elliot-barn pushed a commit that referenced this pull request Jun 18, 2025
@bveeramani @elliot-barn
[Data] Improve handling of mismatched columns (#53861)
9aef295 
<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this
solves. -->

If a map transform function produces outputs with different column
(e.g., you're reading a JSONL file where lines contain different keys),
then Ray Data errors.

This PR updates the implementation to make it more robust to variation
in column names.

## Related issue number

<!-- For example: "Closes #1234" -->

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>
minerharry pushed a commit to minerharry/ray that referenced this pull request Jun 27, 2025
@bveeramani @minerharry
[Data] Improve handling of mismatched columns (ray-project#53861)
fc9a4b1 
<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this
solves. -->

If a map transform function produces outputs with different column
(e.g., you're reading a JSONL file where lines contain different keys),
then Ray Data errors.

This PR updates the implementation to make it more robust to variation
in column names.

## Related issue number

<!-- For example: "Closes ray-project#1234" -->

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
elliot-barn pushed a commit that referenced this pull request Jul 2, 2025
@bveeramani @elliot-barn
[Data] Improve handling of mismatched columns (#53861)
574c9bd 
<!-- Thank you for your contribution! Please review
https://github.com/ray-project/ray/blob/master/CONTRIBUTING.rst before
opening a pull request. -->

<!-- Please add a reviewer to the assignee section when you create a PR.
If you don't have the access to it, we will shortly find a reviewer and
assign them to your PR. -->

## Why are these changes needed?

<!-- Please give a short summary of the change and the problem this
solves. -->

If a map transform function produces outputs with different column
(e.g., you're reading a JSONL file where lines contain different keys),
then Ray Data errors.

This PR updates the implementation to make it more robust to variation
in column names.

## Related issue number

<!-- For example: "Closes #1234" -->

## Checks

- [ ] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [ ] I've run `scripts/format.sh` to lint the changes in this PR.
- [ ] I've included any doc changes needed for
https://docs.ray.io/en/master/.
- [ ] I've added any new APIs to the API Reference. For example, if I
added a
method in Tune, I've added it in `doc/source/tune/api/` under the
           corresponding `.rst` file.
- [ ] I've made sure the tests are passing. Note that there might be a
few flaky tests, see the recent failures at https://flakey-tests.ray.io/
- Testing Strategy
   - [ ] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Balaji Veeramani <bveeramani@berkeley.edu>
Signed-off-by: elliot-barn <elliot.barnwell@anyscale.com>
@Oblynx
Copy link
Copy Markdown

Oblynx commented Jul 20, 2025

Very happy to see this! I've been hitting this bug!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

Copilot code review Copilot Copilot left review comments

@raulchen raulchen raulchen approved these changes

Assignees

@raulchen raulchen

Labels

go add ONLY when ready to merge, run all tests

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

4 participants

@bveeramani @Oblynx @raulchen