flat_map#

Apply the given function to each row and then flatten results.

Use this method if your transformation returns multiple rows for each input row.

You can use either a function or a callable class to perform the transformation. For functions, Ray Data uses stateless Ray tasks. For classes, Ray Data uses stateful Ray actors. For more information, see Stateful Transforms.

Tip

map_batches() can also modify the number of rows. If your transformation is vectorized like most NumPy and pandas operations, it might be faster.

Warning

Specifying both num_cpus and num_gpus for map tasks is experimental, and may result in scheduling or stability issues. Please report any issues to the Ray team.

Examples

>>> from typing import Any, Dict, List
>>> import ray
>>> def duplicate_row(row: Dict[str, Any]) -> List[Dict[str, Any]]:
...     return [row] * 2
>>> print(
...     ray.data.range(3)
...     .flat_map(duplicate_row)
...     .take_all()
... )
[{'id': 0}, {'id': 0}, {'id': 1}, {'id': 1}, {'id': 2}, {'id': 2}]

Time complexity: O(dataset size / parallelism)

Parameters:

fn (Callable[[Dict[str, Any]], List[Dict[str, Any]] | Dict[str, Any]] | Callable[[Dict[str, Any]], Iterator[List[Dict[str, Any]] | Dict[str, Any]]] | type[_CallableClassProtocol]) – The function or generator to apply to each record, or a class type that can be instantiated to create such a callable.
compute (ComputeStrategy | None) –
The compute strategy to use for the map operation.
- If compute is not specified for a function, will use ray.data.TaskPoolStrategy() to launch concurrent tasks based on the available resources and number of input blocks.
- Use ray.data.TaskPoolStrategy(size=n) to launch at most n concurrent Ray tasks.
- If compute is not specified for a callable class, will use ray.data.ActorPoolStrategy(min_size=1, max_size=None) to launch an autoscaling actor pool from 1 to unlimited workers.
- 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.
- Use ray.data.ActorPoolStrategy(min_size=m, max_size=n, initial_size=initial) to use an autoscaling actor pool from m to n workers, with an initial size of initial.
fn_args (Iterable[Any] | None) – Positional arguments to pass to fn after the first argument. These arguments are top-level arguments to the underlying Ray task.
fn_kwargs (Dict[str, Any] | None) – Keyword arguments to pass to fn. These arguments are top-level arguments to the underlying Ray task.
fn_constructor_args (Iterable[Any] | None) – Positional arguments to pass to fn’s constructor. You can only provide this if fn is a callable class. These arguments are top-level arguments in the underlying Ray actor construction task.
fn_constructor_kwargs (Dict[str, Any] | None) – Keyword arguments to pass to fn’s constructor. This can only be provided if fn is a callable class. These arguments are top-level arguments in the underlying Ray actor construction task.
num_cpus (float | None) – The number of CPUs to reserve for each parallel map worker.
num_gpus (float | None) – The number of GPUs to reserve for each parallel map worker. For example, specify num_gpus=1 to request 1 GPU for each parallel map worker.
memory (float | None) – The heap memory in bytes to reserve for each parallel map worker.
concurrency (int | Tuple[int, int] | Tuple[int, int, int] | None) – This argument is deprecated. Use compute argument.
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_fn (Callable[[], Dict[str, Any]] | None) – A function that returns a dictionary of remote args passed to each map worker. The purpose of this argument is to generate dynamic arguments for each actor/task, and will be called each time prior to initializing the worker. Args returned from this dict will always override the args in ray_remote_args. Note: this is an advanced, experimental feature. This argument is deprecated and will be removed in Ray 2.64.
**ray_remote_args – Additional resource requirements to request from Ray for each map worker. See ray.remote() for details. This argument is deprecated and will be removed in Ray 2.64.