Skip to content

[Data] - Add DataType to Expression class#55915

Merged
alexeykudinkin merged 24 commits intoray-project:masterfrom
goutamvenkat-anyscale:goutam/datatype_expr
Sep 4, 2025
Merged

[Data] - Add DataType to Expression class#55915
alexeykudinkin merged 24 commits intoray-project:masterfrom
goutamvenkat-anyscale:goutam/datatype_expr

Conversation

@goutamvenkat-anyscale
Copy link
Contributor

@goutamvenkat-anyscale goutamvenkat-anyscale commented Aug 25, 2025

Why are these changes needed?

Add return_dtype to all Expr types.

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 :(

@goutamvenkat-anyscale goutamvenkat-anyscale requested a review from a team as a code owner August 25, 2025 23:05
gemini-code-assist[bot]
gemini-code-assist bot reviewed Aug 25, 2025
View reviewed changes
Copy link
Contributor

@gemini-code-assist gemini-code-assist bot left a comment

Choose a reason for hiding this comment

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

Code Review

This pull request introduces a new DataType class to represent data types within Ray Data expressions, and integrates it into the Expr class hierarchy. The new DataType class provides a unified abstraction over PyArrow, NumPy, and Python types, with methods for conversion and type inference.

The implementation is a good first step. I've identified a few areas for improvement in the new DataType class to enhance its robustness and correctness, particularly around type conversion fallbacks and exception handling. My comments focus on ensuring the new DataType class is as reliable and performant as possible.

@goutamvenkat-anyscale goutamvenkat-anyscale added data Ray Data-related issues go add ONLY when ready to merge, run all tests labels Aug 25, 2025
@goutamvenkat-anyscale
address comments
81051d7 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Fix doclint
c268bba 
Signed-off-by: Goutam V <[email protected]>
bveeramani
bveeramani reviewed Aug 26, 2025
View reviewed changes
@goutamvenkat-anyscale
Address comments
493bccb 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Fix docstring
b7e5e9b 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Fix doclint again
d657418 
Signed-off-by: Goutam V <[email protected]>
bveeramani
bveeramani approved these changes Aug 26, 2025
View reviewed changes
Copy link
Member

@bveeramani bveeramani left a comment

Choose a reason for hiding this comment

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

Minor comment but looks good to me https://github.com/ray-project/ray/pull/55915/files#r2302280395

omatthew98
omatthew98 reviewed Aug 29, 2025
View reviewed changes
python/ray/data/datatype.py Outdated
)

# Type checking methods
def is_arrow_type(self) -> bool:
Copy link
Contributor

@omatthew98 omatthew98 Aug 29, 2025

Choose a reason for hiding this comment

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

Nit: Could use a field like self._internal_type_kind or similar to store the kind of type on construction and add some type of enum to check against rather than having to constantly call these functions / the underlying isinstance function.

srinathk10
srinathk10 reviewed Aug 29, 2025
View reviewed changes
alexeykudinkin
alexeykudinkin approved these changes Aug 29, 2025
View reviewed changes
python/ray/data/datatype.py Outdated

def __post_init__(self):
"""Validate the _internal_type after initialization."""
if not isinstance(self._internal_type, (pa.DataType, np.dtype, type)):
Copy link
Contributor

@alexeykudinkin alexeykudinkin Aug 29, 2025

Choose a reason for hiding this comment

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

We need to match support of batch UDFs: ie we'd have to support PA, Numpy, Pandas, Python (not a fan of the latter but since we already support it)

Copy link
Contributor Author

@goutamvenkat-anyscale goutamvenkat-anyscale Aug 29, 2025 •
edited
Loading

Choose a reason for hiding this comment

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

Do you mind if I add Pandas support in a separate PR? I've added a TODO for this.

@goutamvenkat-anyscale
Merge from master
dd4340a 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Address comments
833de7c 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Add unit tests
c346317 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Doc lint
9e17680 
Signed-off-by: Goutam V <[email protected]>
alexeykudinkin
alexeykudinkin reviewed Aug 30, 2025
View reviewed changes
@goutamvenkat-anyscale
colocate classes
eca2f9e 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Annotation
0aa65b7 
Signed-off-by: Goutam V <[email protected]>
@goutamvenkat-anyscale
Make linter happy
7b212c7 
Signed-off-by: Goutam V <[email protected]>
alexeykudinkin
alexeykudinkin reviewed Sep 2, 2025
View reviewed changes
iamjustinhsu
iamjustinhsu reviewed Sep 2, 2025
View reviewed changes
alexeykudinkin
alexeykudinkin reviewed Sep 3, 2025
View reviewed changes
@alexeykudinkin alexeykudinkin enabled auto-merge (squash) September 4, 2025 01:23
@github-actions github-actions bot disabled auto-merge September 4, 2025 03:28
@alexeykudinkin alexeykudinkin enabled auto-merge (squash) September 4, 2025 04:12
@goutamvenkat-anyscale
Remove from rst file
ab6ef27 
Signed-off-by: Goutam V. <[email protected]>
@github-actions github-actions bot disabled auto-merge September 4, 2025 17:56
@alexeykudinkin alexeykudinkin merged commit 5022884 into ray-project:master Sep 4, 2025
5 checks passed
@goutamvenkat-anyscale goutamvenkat-anyscale deleted the goutam/datatype_expr branch September 4, 2025 19:17
sampan-s-nayak pushed a commit to sampan-s-nayak/ray that referenced this pull request Sep 8, 2025
@goutamvenkat-anyscale
[Data] - Add DataType to Expression class (ray-project#55915)
16e8ec6 
<!-- 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?

Add return_dtype to all Expr types.

## Related issue number

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

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] 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
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Goutam V <[email protected]>
Signed-off-by: Goutam V. <[email protected]>
Signed-off-by: sampan <[email protected]>
jugalshah291 pushed a commit to jugalshah291/ray_fork that referenced this pull request Sep 11, 2025
@goutamvenkat-anyscale @jugalshah291
[Data] - Add DataType to Expression class (ray-project#55915)
acef136 
<!-- 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?

Add return_dtype to all Expr types.

## Related issue number

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

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] 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
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Goutam V <[email protected]>
Signed-off-by: Goutam V. <[email protected]>
Signed-off-by: jugalshah291 <[email protected]>
wyhong3103 pushed a commit to wyhong3103/ray that referenced this pull request Sep 12, 2025
@goutamvenkat-anyscale
[Data] - Add DataType to Expression class (ray-project#55915)
ce4a19e 
<!-- 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?

Add return_dtype to all Expr types.

## Related issue number

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

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] 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
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Goutam V <[email protected]>
Signed-off-by: Goutam V. <[email protected]>
Signed-off-by: yenhong.wong <[email protected]>
dstrodtman pushed a commit that referenced this pull request Oct 6, 2025
@goutamvenkat-anyscale @dstrodtman
[Data] - Add DataType to Expression class (#55915)
42f4810 
<!-- 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?

Add return_dtype to all Expr types.

## Related issue number

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

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] 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
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Goutam V <[email protected]>
Signed-off-by: Goutam V. <[email protected]>
Signed-off-by: Douglas Strodtman <[email protected]>
snorkelopstesting4-web pushed a commit to snorkel-marlin-repos/ray-project_ray_pr_55915_6b89ff95-a69e-45de-bced-b95a4351578a that referenced this pull request Oct 22, 2025
[Data] - Add DataType to Expression class
8037f70 
Original PR #55915 by goutamvenkat-anyscale
Original: ray-project/ray#55915
@snorkelopstesting3-bot snorkelopstesting3-bot mentioned this pull request Oct 22, 2025
Merged
1 task
snorkelopstesting3-bot added a commit to snorkel-marlin-repos/ray-project_ray_pr_55915_6b89ff95-a69e-45de-bced-b95a4351578a that referenced this pull request Oct 22, 2025
@snorkelopstesting3-bot
Merge pull request #1: [Data] - Add DataType to Expression class
3c46ebc 
Merged from original PR #55915
Original: ray-project/ray#55915
landscapepainter pushed a commit to landscapepainter/ray that referenced this pull request Nov 17, 2025
@goutamvenkat-anyscale @landscapepainter
[Data] - Add DataType to Expression class (ray-project#55915)
6fd4cf6 
<!-- 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?

Add return_dtype to all Expr types.


## Related issue number

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

## Checks

- [x] I've signed off every commit(by using the -s flag, i.e., `git
commit -s`) in this PR.
- [x] 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
   - [x] Unit tests
   - [ ] Release tests
   - [ ] This PR is not tested :(

---------

Signed-off-by: Goutam V <[email protected]>
Signed-off-by: Goutam V. <[email protected]>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@omatthew98 omatthew98 omatthew98 left review comments

@srinathk10 srinathk10 srinathk10 left review comments

@iamjustinhsu iamjustinhsu iamjustinhsu left review comments

@alexeykudinkin alexeykudinkin alexeykudinkin approved these changes

@bveeramani bveeramani bveeramani approved these changes

+1 more reviewer

@gemini-code-assist gemini-code-assist[bot] gemini-code-assist[bot] left review comments

Reviewers whose approvals may not affect merge requirements

Assignees

No one assigned

Labels

data Ray Data-related issues 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.

6 participants

@goutamvenkat-anyscale @alexeykudinkin @omatthew98 @bveeramani @srinathk10 @iamjustinhsu