[Data] - Unify Project Operator to use Expressions #57076

goutamvenkat-anyscale · 2025-10-01T03:52:21Z

Why are these changes needed?

Makes the following changes:

Move Ray data's expression evaluator to an internal module
Use the visitor pattern for native expressions
Use expressions for rename_columns, with_column, select_columns and remove cols and cols_rename in Project
Modify Projection Pushdown to work with combinations of the above operators correctly

Related issue number

Closes #56878, #57700

Checks

I've signed off every commit(by using the -s flag, i.e., git commit -s) in this PR.
I've run pre-commit jobs to lint the changes in this PR. (pre-commit setup)
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 :(

Note

Switch Project/column ops to named expression lists with preserve_existing, move/evolve expression evaluation to internal native Block visitor, and implement robust Project fusion/pushdown; plus Arrow/Pandas block fixes and test updates.

Expressions & Evaluation:
- Move expression evaluator to _internal/planner/plan_expression/expression_evaluator.py; add native Block/BlockColumn visitor (NativeExpressionEvaluator) and centralize op maps.
- Update all call sites to use new internal eval_expr (Arrow/Pandas blocks, UDF planner, tests).
Project Operator & APIs:
- Redesign Project to accept exprs: List[Expr] with enforced names and preserve_existing flag; remove cols/cols_rename paths.
- Rework planner (plan_project_op) to compute expressions, fill/upsert columns, and finalize schema based on preserve_existing.
- Migrate Dataset APIs: with_column, select_columns, rename_columns to build named expressions; count uses empty exprs with preserve_existing=False.
Projection Pushdown/Fusion:
- Replace rule with expression-aware fusion of consecutive Projects (handles renames, substitution, and visibility via preserve_existing).
- Push down required columns into reads only for selection-style projects.
Arrow/Pandas Block & Shuffle:
- Arrow concat: preserve rows for 0-column tables via stub column; fill_column uses upsert_column; num_rows returns raw num_rows.
- Remove row-count assert in shuffle map; minor filter import path changes.
Expressions minor:
- Fix UnaryExpr default data_type initialization.
Tests/Plans:
- Update tests to new expression APIs and plan expectations; add comprehensive projection fusion tests; adjust stats/explain outputs.

^{Written by Cursor Bugbot for commit dbf609a. This will update automatically on new commits. Configure here.}

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

goutamvenkat-anyscale · 2025-10-01T03:53:40Z

/gemini review

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

gemini-code-assist

Code Review

This pull request is a significant refactoring to unify projection-like operations (select_columns, with_column, rename_columns) under a single Project logical operator that uses expressions. This is a great improvement for consistency and maintainability. The changes include moving the expression evaluator, implementing a visitor pattern for evaluation, and rewriting the projection pushdown logic to correctly fuse chains of projection operations. The new logic is more robust and handles complex cases like column swaps and chained renames correctly. I've found a few minor issues with type hints and docstrings that could be improved for clarity. Overall, this is a high-quality contribution.

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

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

gemini-code-assist

Code Review

This pull request is a significant and well-executed refactoring that unifies the select_columns, rename_columns, and with_column operations into a single expression-based Project operator. This change enhances the planner's optimization capabilities. The implementation is solid, introducing a visitor pattern for expression evaluation and rewriting the projection pushdown logic to correctly handle complex chains of operations. My review focuses on the new, complex logic in the planner and expression evaluator, with suggestions to improve type hint accuracy and simplify some logic for better long-term maintainability.

python/ray/data/_internal/logical/rules/projection_pushdown.py

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

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

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

goutamvenkat-anyscale · 2025-10-01T06:45:47Z

/gemini review

gemini-code-assist

Code Review

This pull request is an excellent and significant refactoring that unifies the Project operator to be fully expression-driven. The changes, including moving the expression evaluator to an internal module, adopting the visitor pattern for evaluation, and overhauling the projection pushdown logic, greatly improve the code's structure, maintainability, and correctness. The new ProjectionPushdown implementation is particularly impressive in how it handles complex chains of select, with_column, and rename operations. The addition of comprehensive tests is also a major plus. I've found one minor issue related to error handling that could be improved.

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

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

goutamvenkat-anyscale · 2025-10-01T17:54:07Z

/gemini review

gemini-code-assist

Code Review

This pull request is a major and well-executed refactoring that unifies select_columns, with_column, and rename_columns under a single expression-based Project logical operator. This is a significant improvement that simplifies the logical plan and enables more powerful and robust optimizations, particularly for projection fusion and pushdown. The changes are extensive, touching the logical operators, the physical planner, the expression evaluation logic, and the public Dataset API. The new implementation of projection pushdown is much cleaner and correctly handles complex chains of operations. The introduction of a visitor-based expression evaluator is also a solid design choice. I've found one minor issue with error handling in a helper function, but overall this is an excellent refactoring.

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

python/ray/data/_internal/arrow_block.py

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

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

python/ray/data/_internal/arrow_block.py

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

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

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

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

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

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

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

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

python/ray/data/expressions.py

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

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

goutamvenkat-anyscale · 2025-10-14T02:07:06Z

/gemini review

gemini-code-assist

Code Review

This pull request is a significant and well-executed refactoring that unifies various column operations (with_column, select_columns, rename_columns) under a single expression-based Project operator. Moving the expression evaluation logic to an internal module and adopting the visitor pattern is a great architectural improvement. The introduction of StarColumnsExpr and star() provides a clean mechanism for additive column operations. The rewrite of the projection pushdown rule is particularly impressive, enabling more powerful and robust fusion of consecutive Project operators. Overall, these changes greatly improve the clarity, consistency, and maintainability of the Ray Data API and its underlying implementation. I've included a few suggestions for further improvement.

python/ray/data/expressions.py

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

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

cursor · 2025-10-14T02:21:49Z

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

+    elif _is_pa_string_like(x) and isinstance(x, (pa.Array, pa.ChunkedArray)):
+        x = _pa_decode_dict_string_array(x)
+    else:
+        raise


Bug: Unconditional Raising Causes Runtime Error

A bare raise statement without an explicit exception type is present. This attempts to re-raise the last active exception, but if no exception is currently being handled, it will cause a RuntimeError.

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

alexeykudinkin · 2025-10-14T19:40:04Z

python/ray/data/_internal/logical/rules/projection_pushdown.py

+    # If any expression is star(), we need all columns
+    if any(isinstance(expr, StarColumnsExpr) for expr in exprs):
+        return None


This encoding (none -> all) is implicit -- let's make it explicit.

To make it explicit we'd need to add an expression resolution (against schema) step. We don't need to do in this PR, but let's cut a ticket and link it here

alexeykudinkin · 2025-10-14T19:40:35Z

python/ray/data/_internal/logical/rules/projection_pushdown.py

+    """
+    referenced_columns: Set[str] = set()
+
+    def visit_expr(expr: Expr) -> None:


+1

Make all traversals a corresponding Visitor (and make visitors w/o state singletons)

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

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

alexeykudinkin · 2025-10-16T00:29:30Z

python/ray/data/_internal/logical/rules/projection_pushdown.py


-logger = logging.getLogger(__name__)

+class _ColumnReferenceCollector(_ExprVisitor):


This file is impossible to review now

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py

alexeykudinkin · 2025-10-16T00:37:03Z

python/ray/data/expressions.py

    operand: Expr

-    data_type: DataType = field(init=False)
+    data_type: DataType = field(default_factory=lambda: DataType(object), init=False)


What does this default mean?

Added comment in new PR. It's the default return dtype. For UnaryExpr it should actually be bool not object.

Without this defined, further expressions can't be chained. If I have expr1.is_not_null().alias('x') then we need to ensure that AliasExpr gets the datatype propagated from the UnaryExpr

python/ray/data/expressions.py

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

alexeykudinkin · 2025-10-16T00:43:27Z

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

+    if len(exprs) == 1 and isinstance(exprs[0], StarColumnsExpr):
+        return block


Expand stars before continuing processing (will later move this into resolving stage)

Do you mean creating a list of col expressions?

That requires the schema

If I have 100k+ columns, isn't that slow?

Addressed this in new PR

alexeykudinkin · 2025-10-16T00:43:36Z

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py

+        if isinstance(expr, StarColumnsExpr):
+            continue


See comment above

Addressed in new PR

goutamvenkat-anyscale · 2025-10-20T16:40:54Z

Closing in favor of #57855

goutamvenkat-anyscale added 2 commits September 29, 2025 12:27

[Data][1/n] - Use visitor pattern for Ray Data Expressions

fb4cb18

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

Fix projection pushdown

f99ae1c

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

goutamvenkat-anyscale requested a review from a team as a code owner October 1, 2025 03:52

goutamvenkat-anyscale added the data Ray Data-related issues label Oct 1, 2025

goutamvenkat-anyscale commented Oct 1, 2025

View reviewed changes

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

gemini-code-assist bot reviewed Oct 1, 2025

View reviewed changes

This comment was marked as outdated.

Sign in to view

gemini-code-assist bot reviewed Oct 1, 2025

View reviewed changes

python/ray/data/_internal/logical/rules/projection_pushdown.py Outdated Show resolved Hide resolved

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py Show resolved Hide resolved

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

goutamvenkat-anyscale added 4 commits September 30, 2025 22:53

Fix count

275d919

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

gemini comments

4c273e8

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

Clean up

7aed3af

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

Merge from master

40f724e

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

gemini-code-assist bot reviewed Oct 1, 2025

View reviewed changes

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py Show resolved Hide resolved

This comment was marked as outdated.

Sign in to view

Fix filter + cleanup

c879ad6

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

This comment was marked as outdated.

Sign in to view

goutamvenkat-anyscale mentioned this pull request Oct 1, 2025

[Ray Data] Add a force overwrite option to the rename_columns() #56878

Closed

goutamvenkat-anyscale added 2 commits October 1, 2025 01:12

Move test to fusion + fix projection fusion tests

c140ad7

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

Fix operator fusion

f35cfd0

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

This comment was marked as outdated.

Sign in to view

goutamvenkat-anyscale added 3 commits October 1, 2025 10:34

Fix parquet tests

30a7ab5

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

Spelling mistake

6a626e8

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

one more cleanup

f1b5d97

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

gemini-code-assist bot reviewed Oct 1, 2025

View reviewed changes

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py Show resolved Hide resolved

goutamvenkat-anyscale commented Oct 1, 2025

View reviewed changes

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

Handle invalid orders in projection pushdown

73cb01b

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

This comment was marked as outdated.

Sign in to view

goutamvenkat-anyscale added 2 commits October 7, 2025 15:46

Merge from master

282c3d5

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

remove old test suite

eb52554

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

cursor bot reviewed Oct 7, 2025

View reviewed changes

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py Outdated Show resolved Hide resolved

python/ray/data/_internal/arrow_block.py Show resolved Hide resolved

python/ray/data/_internal/planner/plan_expression/expression_evaluator.py Show resolved Hide resolved

Fix download test + comments

ebe2441

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

cursor bot reviewed Oct 8, 2025

View reviewed changes

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py Outdated Show resolved Hide resolved

Add all expression

5f814f9

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

cursor bot reviewed Oct 8, 2025

View reviewed changes

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py Outdated Show resolved Hide resolved

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

doc fix

fb9d380

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

cursor bot reviewed Oct 8, 2025

View reviewed changes

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py Outdated Show resolved Hide resolved

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

goutamvenkat-anyscale added 2 commits October 8, 2025 00:17

Fix test + api_policy_check

f46a40f

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

merge from master

cd84a20

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

cursor bot reviewed Oct 8, 2025

View reviewed changes

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py Outdated Show resolved Hide resolved

alexeykudinkin reviewed Oct 10, 2025

View reviewed changes

goutamvenkat-anyscale added 3 commits October 13, 2025 16:17

Merge branch 'master' into goutam/visitor_expr

8867156

Address comments

c675dd6

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

doc fix

aa2a9b5

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

gemini-code-assist bot reviewed Oct 14, 2025

View reviewed changes

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

python/ray/data/_internal/planner/exchange/shuffle_task_spec.py Outdated Show resolved Hide resolved

zero_div error

29feb0f

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

cursor bot reviewed Oct 14, 2025

View reviewed changes

goutamvenkat-anyscale added 2 commits October 13, 2025 20:11

Fix test

638b9c6

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

Add stub column back and remove arrow block change

62f354b

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

alexeykudinkin reviewed Oct 14, 2025

View reviewed changes

ray-project deleted a comment from alexeykudinkin Oct 15, 2025

goutamvenkat-anyscale added 3 commits October 14, 2025 17:32

use visitor in pushdown

b869cfe

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

Merge branch 'master' into goutam/visitor_expr

3df14b3

one more visitor

6ded04a

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

alexeykudinkin reviewed Oct 16, 2025

View reviewed changes

goutamvenkat-anyscale closed this Oct 20, 2025


		logger = logging.getLogger(__name__)

		class _ColumnReferenceCollector(_ExprVisitor):

		if len(exprs) == 1 and isinstance(exprs[0], StarColumnsExpr):
		return block

[Data] - Unify Project Operator to use Expressions #57076

[Data] - Unify Project Operator to use Expressions #57076

Uh oh!

Conversation

goutamvenkat-anyscale commented Oct 1, 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!

goutamvenkat-anyscale commented Oct 1, 2025

Uh oh!

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

goutamvenkat-anyscale commented Oct 1, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

This comment was marked as outdated.

Uh oh!

goutamvenkat-anyscale commented Oct 1, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

This comment was marked as outdated.

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!

Uh oh!

Uh oh!

goutamvenkat-anyscale commented Oct 14, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

cursor bot Oct 14, 2025

Choose a reason for hiding this comment

Bug: Unconditional Raising Causes Runtime Error

goutamvenkat-anyscale commented Oct 1, 2025 •

edited

Loading

goutamvenkat-anyscale Oct 16, 2025 •

edited

Loading