[Data] Add Expression Support & with_columns API #54322

goutamvenkat-anyscale · 2025-07-03T18:57:03Z

Why are these changes needed?

Adding with_columns API to enable users to pass in multiple expressions, that can get bundled into a single Project call.

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

Signed-off-by: Goutam V <[email protected]>

python/ray/data/dataset.py

Signed-off-by: Goutam V <[email protected]>

python/ray/data/dataset.py

python/ray/data/expressions.py

Signed-off-by: Goutam V <[email protected]>

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

python/ray/data/expressions.py

python/ray/data/dataset.py

iamjustinhsu · 2025-07-09T18:24:48Z

python/ray/data/expressions.py

+# ──────────────────────────────────────
+# This is used by Dataset.with_columns – kept here so it can be re-used by
+# future optimised executors.
+_PANDAS_OPS: Dict[str, Callable[[Any, Any], Any]] = {


if these pandas, numpy, and pyarrow all need to have keys for "add", "sub", etc.. maybe better to define a base class for each key, and then implement __call__ for each subclass (pandas, numpy, pyarrow)?

python/ray/data/expressions.py

Signed-off-by: Goutam V <[email protected]>

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

python/ray/data/dataset.py

python/ray/data/_expression_evaluator.py

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

python/ray/data/dataset.py

Signed-off-by: Goutam V <[email protected]>

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

python/ray/data/_expression_evaluator.py

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin

Please address comments before landing

python/ray/data/_expression_evaluator.py

python/ray/data/_internal/block_builder.py

python/ray/data/expressions.py

Signed-off-by: Goutam V <[email protected]>

python/ray/data/expressions.py

Signed-off-by: Goutam V <[email protected]>

python/ray/data/_expression_evaluator.py

python/ray/data/expressions.py

Signed-off-by: Goutam V <[email protected]>

python/ray/data/expressions.py

Signed-off-by: Goutam V <[email protected]>

Signed-off-by: jugalshah291 <[email protected]>

Signed-off-by: Douglas Strodtman <[email protected]>

goutamvenkat-anyscale requested a review from a team as a code owner July 3, 2025 18:57

goutamvenkat-anyscale added 2 commits July 3, 2025 11:57

Add Expression Support & with_column API

50a74b3

Signed-off-by: Goutam V <[email protected]>

Rename to use_columns, use list[expr]

9a5086f

Signed-off-by: Goutam V <[email protected]>

goutamvenkat-anyscale force-pushed the goutam/expressions branch from a8807e8 to 9a5086f Compare July 3, 2025 18:57

alexeykudinkin reviewed Jul 3, 2025

View reviewed changes

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

Use project operator & update doc

9804716

Signed-off-by: Goutam V <[email protected]>

goutamvenkat-anyscale added the go add ONLY when ready to merge, run all tests label Jul 7, 2025

goutamvenkat-anyscale added 2 commits July 7, 2025 00:53

Fix linting issue

16cccff

Signed-off-by: Goutam V <[email protected]>

Doc linter

a309598

Signed-off-by: Goutam V <[email protected]>

goutamvenkat-anyscale requested a review from alexeykudinkin July 7, 2025 16:25

goutamvenkat-anyscale changed the title ~~Add Expression Support & with_columns API~~ [Data] Add Expression Support & with_columns API Jul 7, 2025

goutamvenkat-anyscale added the data Ray Data-related issues label Jul 7, 2025

doctest

176b444

Signed-off-by: Goutam V <[email protected]>

raulchen reviewed Jul 7, 2025

View reviewed changes

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

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

python/ray/data/expressions.py Show resolved Hide resolved

Address comment

c17570f

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin reviewed Jul 9, 2025

View reviewed changes

iamjustinhsu reviewed Jul 9, 2025

View reviewed changes

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

iamjustinhsu reviewed Jul 9, 2025

View reviewed changes

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

goutamvenkat-anyscale added 2 commits July 9, 2025 13:39

Address comments

eea0d55

Signed-off-by: Goutam V <[email protected]>

Linter & remove dataclass for operations

ed251a1

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin reviewed Jul 9, 2025

View reviewed changes

goutamvenkat-anyscale added 5 commits July 9, 2025 15:27

Address comments

e36a87b

Signed-off-by: Goutam V <[email protected]>

revert old change

86bc9fb

Signed-off-by: Goutam V <[email protected]>

Remove unnecessary arg

3053b95

Signed-off-by: Goutam V <[email protected]>

Merge branch 'master' into goutam/expressions

b09ae8f

doctest + pytest skip if version is not met

fb3c6a1

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin reviewed Jul 10, 2025

View reviewed changes

python/ray/data/_internal/planner/plan_udf_map_op.py Show resolved Hide resolved

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

goutamvenkat-anyscale added 2 commits July 10, 2025 10:18

Remove circular dep

bd7bc77

Signed-off-by: Goutam V <[email protected]>

Address comments

6d443f8

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin approved these changes Jul 10, 2025

View reviewed changes

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

python/ray/data/_internal/block_builder.py Outdated Show resolved Hide resolved

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

goutamvenkat-anyscale added 5 commits July 10, 2025 11:20

remove change in block builder

fc034ec

Signed-off-by: Goutam V <[email protected]>

Remove block builder change

1a3941f

Signed-off-by: Goutam V <[email protected]>

Make pre-commit happy

3f30cbb

Signed-off-by: Goutam V <[email protected]>

Address comment on Expr AST comparison

9b8de87

Signed-off-by: Goutam V <[email protected]>

Add expressions test to bazel build

c13f679

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin reviewed Jul 10, 2025

View reviewed changes

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

goutamvenkat-anyscale added 3 commits July 10, 2025 15:15

Remove match expression

8d61562

Signed-off-by: Goutam V <[email protected]>

Comments

d8890fd

Signed-off-by: Goutam V <[email protected]>

Merge branch 'master' into goutam/expressions

49e3ccb

alexeykudinkin reviewed Jul 11, 2025

View reviewed changes

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

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

python/ray/data/expressions.py Show resolved Hide resolved

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

goutamvenkat-anyscale added 2 commits July 11, 2025 10:01

Address comments

164cbd3

Signed-off-by: Goutam V <[email protected]>

Add comments back

b64beef

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin reviewed Jul 11, 2025

View reviewed changes

python/ray/data/expressions.py Show resolved Hide resolved

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

goutamvenkat-anyscale added 2 commits July 11, 2025 10:30

Make expression classes dev api

a3f3050

Signed-off-by: Goutam V <[email protected]>

Add stability to DeveloperAPIs

821b73e

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin approved these changes Jul 11, 2025

View reviewed changes

alexeykudinkin enabled auto-merge (squash) July 11, 2025 20:01

Add .rst files

f5b08eb

Signed-off-by: Goutam V <[email protected]>

github-actions bot disabled auto-merge July 11, 2025 23:23

goutamvenkat-anyscale added 3 commits July 11, 2025 17:06

idk rst

2741b12

Signed-off-by: Goutam V <[email protected]>

Merge branch 'master' into goutam/expressions

f4f620c

Remove code snippet

c7f0424

Signed-off-by: Goutam V <[email protected]>

alexeykudinkin merged commit 0cebaa1 into ray-project:master Jul 12, 2025
5 checks passed

jugalshah291 pushed a commit to jugalshah291/ray_fork that referenced this pull request Sep 11, 2025

[Data] Add Expression Support & with_columns API (ray-project#54322)

0a52939

Signed-off-by: jugalshah291 <[email protected]>

dstrodtman pushed a commit to dstrodtman/ray that referenced this pull request Oct 6, 2025

[Data] Add Expression Support & with_columns API (ray-project#54322)

ccce219

Signed-off-by: Douglas Strodtman <[email protected]>

[Data] Add Expression Support & with_columns API #54322

[Data] Add Expression Support & with_columns API #54322

Uh oh!

Conversation

goutamvenkat-anyscale commented Jul 3, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Why are these changes needed?

Related issue number

Checks

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

iamjustinhsu Jul 9, 2025

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

alexeykudinkin left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

4 participants

goutamvenkat-anyscale commented Jul 3, 2025 •

edited

Loading