[Data] UDF Expression Support for with_column

[goutamvenkat-anyscale]

Why are these changes needed?

This adds support for UDF as expressions into with_column.

Since this UDFExpr is designed for batches of data each parameter represents a pyarrow.Array.

Example usage:

import pyarrow as pa
import pyarrow.compute as pc
from ray.data.expressions import col


@udf()
def add_one(x: pa.Array) -> pa.Array:
       return pc.add(x, 1)

ds = ds.with_column('id_plus_1', add_one(col("id")))

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

[gemini-code-assist]

Code Review

This pull request introduces a significant enhancement to Ray Data by adding support for User-Defined Functions (UDFs) within with_column expressions. The implementation is well-designed, introducing a UDFExpr for representing UDFs and a DataType class for better type handling. The logic to dynamically switch between a Project operator and map_batches based on the presence of a UDF and a batch_size is a smart approach. The tests are thorough and cover a wide range of scenarios.

My review includes a few suggestions for improvement:

• The docstring for with_column should be updated to document the new functionality.
• A potential correctness issue in the __hash__ implementation of the new DataType class needs to be addressed.
• An enhancement is proposed for the DataType class to improve Python-to-Arrow type conversion.

Overall, this is a valuable contribution that greatly increases the power and flexibility of Ray Data's expression API.

[gemini-code-assist]

The logic to apply the UDF based on column names inside test_with_column_udf_multi_column can be simplified. Consider moving the UDF application logic into the test_scenario parametrization. For example, you could add a columns key to your test_scenario dicts (e.g., "columns": ["a", "b"]) and then simplify this block to cols_to_use = [col(c) for c in test_scenario["columns"]]; ds_with_udf = ds.with_column(column_name, udf_fn(*cols_to_use)). This would make the test body cleaner and the different test cases more explicit.

[omatthew98]

Small comment, overall lgtm.

[srinathk10]

Nice test coverage!