with_column#

Dataset.with_column(column_name: str, expr: Expr, *, compute: ComputeStrategy | None = None, num_cpus: float | None = None, num_gpus: float | None = None, memory: float | None = None, label_selector: Dict[str, str] | None = None, fallback_strategy: List[Dict[str, Any]] | None = None, max_calls: int | None = None, resources: Dict[str, float] | None = None, accelerator_type: str | None = None, runtime_env: Dict[str, Any] | None = None, **ray_remote_args) → Dataset[source]#

Add a new column to the dataset via an expression.

This method allows you to add a new column to a dataset by applying an expression. The expression can be composed of existing columns, literals, and user-defined functions (UDFs).

For callable class UDFs, Ray Data automatically uses actor semantics to maintain state across batches. You can customize the compute strategy to control parallelism and resource allocation.

Examples

>>> import ray
>>> from ray.data.expressions import col
>>> ds = ray.data.range(100)
>>> # Add a new column 'id_2' by multiplying 'id' by 2.
>>> ds.with_column("id_2", col("id") * 2).show(2)
{'id': 0, 'id_2': 0}
{'id': 1, 'id_2': 2}

>>> # Using a UDF with with_column
>>> from ray.data.datatype import DataType
>>> from ray.data.expressions import udf
>>> import pyarrow.compute as pc
>>>
>>> @udf(return_dtype=DataType.int32())
... def add_one(column):
...     return pc.add(column, 1)
>>>
>>> ds.with_column("id_plus_one", add_one(col("id"))).show(2)
{'id': 0, 'id_plus_one': 1}
{'id': 1, 'id_plus_one': 2}

>>> # Using a callable class UDF (automatically uses actors)
>>> @udf(return_dtype=DataType.int32())
... class AddOffset:
...     def __init__(self, offset):
...         self.offset = offset
...     def __call__(self, x):
...         return pc.add(x, self.offset)
>>>
>>> add_five = AddOffset(5)
>>> ds.with_column("id_plus_five", add_five(col("id"))).show(2)
{'id': 0, 'id_plus_five': 5}
{'id': 1, 'id_plus_five': 6}
Parameters:

  • column_name (str) – The name of the new column.

  • expr (Expr) – An expression that defines the new column values.

  • compute (ComputeStrategy | None) –

    The compute strategy to use for the projection operation. If not specified and the expression contains callable class UDFs, Ray Data automatically uses ActorPoolStrategy for actor semantics. Otherwise, uses TaskPoolStrategy.

    • Use ray.data.ActorPoolStrategy(size=n) to use a fixed size actor pool of n workers.

    • Use ray.data.ActorPoolStrategy(min_size=m, max_size=n) to use an autoscaling actor pool from m to n workers.

  • num_cpus (float | None) – The number of CPUs to reserve for each worker.

  • num_gpus (float | None) – The number of GPUs to reserve for each worker.

  • memory (float | None) – The heap memory in bytes to reserve for each worker.

  • label_selector (Dict[str, str] | None) – Labels required on the node where each worker runs.

  • fallback_strategy (List[Dict[str, Any]] | None) – Alternative label requirements that Ray tries in order when label_selector can’t be satisfied.

  • max_calls (int | None) – The maximum number of calls a task worker handles before exiting. This option only applies to task workers.

  • resources (Dict[str, float] | None) – Custom resources to reserve for each worker, expressed as a mapping from resource name to quantity.

  • accelerator_type (str | None) – The accelerator type required on the node where each worker runs.

  • runtime_env (Dict[str, Any] | None) – The runtime environment to use for each worker.

  • **ray_remote_args – Additional resource requirements to request from Ray for the map tasks (e.g., num_gpus=1). This argument is deprecated and will be removed in Ray 2.64.

Returns:

A new dataset with the added column evaluated via the expression.

Return type:

Dataset

PublicAPI (alpha): This API is in alpha and may change before becoming stable.

On this page