[Data] Enable expressions for grouped with_column in Ray Data #58231
Conversation
Signed-off-by: YoussefEssDS <oyoussefesseddiq@gmail.com>
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request introduces GroupedData.with_column, which allows applying expressions to grouped data. The implementation is clean, reusing existing components like map_groups and eval_projection. The added tests cover both UDF-based and arithmetic expressions, ensuring the new feature is well-tested. My feedback includes a minor suggestion to improve code clarity in the tests by removing a redundant dataset creation.
| ds = ray.data.from_items( | ||
| [ | ||
| {"group": 1, "value": 1}, | ||
| {"group": 1, "value": 2}, | ||
| {"group": 2, "value": 3}, | ||
| {"group": 2, "value": 4}, | ||
| ] | ||
| ) |
There was a problem hiding this comment.
To improve code clarity and avoid redundancy, you can remove this duplicate dataset creation. The ds variable is already defined with the same data earlier in the test and can be reused for the second assertion.
ray-gardener
bot
added
data
Signed-off-by: YoussefEssDS <oyoussefesseddiq@gmail.com>
|
Hi @alexeykudinkin do you have any further suggestions? Thanks! |
1 participant
Description
This PR will:
1- Introduce
GroupedData.with_column, allowing grouped datasets to evaluate Ray Data expressions per group while preserving existing columns.2- Validate the supplied expression type (reject non‑Expr and DownloadExpr since the expression evaluator can’t visit downloads as far as I understand) and reuse the projection engine so grouping flows stay aligned with the dataset-level expression API.
3- Add tests for grouped expression usage through udf-based and arithmetic expressions.
Related issues
Closes #57907