Method flat_map

python/ray/data/dataset.py:1408–1548

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 Ra

(
        self,
        fn: UserDefinedFunction[
            Dict[str, Any], Union[List[Dict[str, Any]], Dict[str, Any]]
        ],
        *,
        compute: Optional[ComputeStrategy] = None,
        fn_args: Optional[Iterable[Any]] = None,
        fn_kwargs: Optional[Dict[str, Any]] = None,
        fn_constructor_args: Optional[Iterable[Any]] = None,
        fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
        num_cpus: Optional[float] = None,
        num_gpus: Optional[float] = None,
        memory: Optional[float] = None,
        concurrency: Optional[Union[int, Tuple[int, int], Tuple[int, int, int]]] = None,
        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
        **ray_remote_args,
    )

    @PublicAPI(api_group=BT_API_GROUP)
    def flat_map(
        self,
        fn: UserDefinedFunction[
            Dict[str, Any], Union[List[Dict[str, Any]], Dict[str, Any]]
        ],
        *,
        compute: Optional[ComputeStrategy] = None,
        fn_args: Optional[Iterable[Any]] = None,
        fn_kwargs: Optional[Dict[str, Any]] = None,
        fn_constructor_args: Optional[Iterable[Any]] = None,
        fn_constructor_kwargs: Optional[Dict[str, Any]] = None,
        num_cpus: Optional[float] = None,
        num_gpus: Optional[float] = None,
        memory: Optional[float] = None,
        concurrency: Optional[Union[int, Tuple[int, int], Tuple[int, int, int]]] = None,
        ray_remote_args_fn: Optional[Callable[[], Dict[str, Any]]] = None,
        **ray_remote_args,
    ) -> "Dataset":
        """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
        :ref:`Stateful Transforms <stateful_transforms>`.

        .. tip::
            :meth:`~Dataset.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 <https://github.com/ray-project/ray/issues/new/choose>`_
            to the Ray team.

        Examples:

            .. testcode::

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

            .. testoutput::

                [{'id': 0}, {'id': 0}, {'id': 1}, {'id': 1}, {'id': 2}, {'id': 2}]