unnest

ray.data.expressions.unnest(expr: Expr) → UnnestExpr

Expand a struct-typed expression into one output column per struct field.

Use this with Dataset.with_columns to let a single expression — typically a UDF that computes several related values and returns them bundled as a struct — produce multiple output columns. The output column names and order come from the struct’s fields. The wrapped expression is evaluated once per block, not once per field.

The struct type must be known when the plan is built: either the wrapped expression declares it (a UDF’s return_dtype), or it is a reference to a struct column of a dataset whose schema is known.

Expansion is one level deep: a field that is itself a struct comes out as a single struct-typed column, not flattened further. Chaining a second with_columns(unnest(col(...))) flattens it, provided the intermediate schema is known at plan time: it is when the struct type came from a declared return_dtype; for a plain struct column, call materialize() between the two steps. unnest() cannot wrap another unnest() — the inner one already denotes multiple columns, so there is no single struct value left to expand — and raises TypeError if you try.

Parameters:

expr (Expr) – An expression that resolves to a PyArrow struct type.

Returns:

An UnnestExpr suitable for passing positionally to Dataset.with_columns.

Return type:

UnnestExpr

Example

>>> import pyarrow as pa
>>> import pyarrow.compute as pc
>>> import ray
>>> from ray.data.datatype import DataType
>>> from ray.data.expressions import col, udf, unnest
>>>
>>> @udf(return_dtype=DataType.struct([
...     ("sum_ab", DataType.int64()),
...     ("product_ab", DataType.int64()),
... ]))
... def make_features(a: pa.Array, b: pa.Array) -> pa.StructArray:
...     return pa.StructArray.from_arrays(
...         [pc.add(a, b).combine_chunks(), pc.multiply(a, b).combine_chunks()],
...         names=["sum_ab", "product_ab"],
...     )
>>>
>>> ds = ray.data.from_items([{"a": 2, "b": 10}, {"a": 3, "b": 20}])
>>> ds.with_columns(unnest(make_features(col("a"), col("b")))).show(1)
{'a': 2, 'b': 10, 'sum_ab': 12, 'product_ab': 20}

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