Method with_column

python/ray/data/dataset.py:873–969

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, R

(
        self,
        column_name: str,
        expr: Expr,
        *,
        compute: Optional[ComputeStrategy] = None,
        **ray_remote_args,
    )

    @PublicAPI(api_group=EXPRESSION_API_GROUP, stability="alpha")
    def with_column(
        self,
        column_name: str,
        expr: Expr,
        *,
        compute: Optional[ComputeStrategy] = None,
        **ray_remote_args,
    ) -> "Dataset":
        """
        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}

        Args:
            column_name: The name of the new column.
            expr: An expression that defines the new column values.
            compute: The compute strategy to use for the projection operation.