python/ray/data/dataset.py
def map_batches(
        self,
        fn: Union[CallableClass, Callable[[BatchType], BatchType]],
        *,
        batch_size: Optional[int] = 4096,
        compute: Union[str, ComputeStrategy] = None,
        batch_format: str = "native",
        **ray_remote_args,
    ) -> "Dataset[Any]":
        """Apply the given function to batches of records of this dataset.

        This is a blocking operation.

        Examples:
            >>> import ray
            >>> # Transform python objects.
            >>> ds = ray.data.range(1000) # doctest: +SKIP
            >>> # Transform batches in parallel.
            >>> ds.map_batches(lambda batch: [v * 2 for v in batch]) # doctest: +SKIP
            >>> # Define a callable class that persists state across
            >>> # function invocations for efficiency.
            >>> init_model = ... # doctest: +SKIP
            >>> class CachedModel:
            ...    def __init__(self):
            ...        self.model = init_model()
            ...    def __call__(self, item):
            ...        return self.model(item)
            >>> # Apply the transform in parallel on GPUs. Since
            >>> # compute=ActorPoolStrategy(2, 8) the transform will be applied on an
            >>> # autoscaling pool of 2-8 Ray actors, each allocated 1 GPU by Ray.
            >>> from ray.data.impl.compute import ActorPoolStrategy
            >>> ds.map_batches( # doctest: +SKIP
            ...     CachedModel, # doctest: +SKIP
            ...     batch_size=256, # doctest: +SKIP
            ...     compute=ActorPoolStrategy(2, 8), # doctest: +SKIP
            ...     num_gpus=1) # doctest: +SKIP

        Time complexity: O(dataset size / parallelism)

        Args:
            fn: The function to apply to each record batch, or a class type
                that can be instantiated to create such a callable. Callable classes are
                only supported for the actor compute strategy.
            batch_size: Request a specific batch size, or None to use entire
                blocks as batches. Defaults to a system-chosen batch size.
            compute: The compute strategy, either "tasks" (default) to use Ray
                tasks, or ActorPoolStrategy(min, max) to use an autoscaling actor pool.
            batch_format: Specify "native" to use the native block format
                (promotes Arrow to pandas), "pandas" to select
                ``pandas.DataFrame`` as the batch format,
                or "pyarrow" to select ``pyarrow.Table``.
            ray_remote_args: Additional resource requirements to request from
                ray (e.g., num_gpus=1 to request GPUs for the map tasks).
        """
        import pyarrow as pa
        import pandas as pd

        if batch_size is not None and batch_size < 1:
            raise ValueError("Batch size cannot be negative or 0")

        fn = cache_wrapper(fn, compute)
        context = DatasetContext.get_current()

        def transform(block: Block) -> Iterable[Block]:
            DatasetContext._set_current(context)
            output_buffer = BlockOutputBuffer(None, context.target_max_block_size)
            block = BlockAccessor.for_block(block)
            total_rows = block.num_rows()
            max_batch_size = batch_size
            if max_batch_size is None:
                max_batch_size = max(total_rows, 1)

            for start in range(0, total_rows, max_batch_size):
                # Build a block for each batch.
                end = min(total_rows, start + max_batch_size)
                # Make sure to copy if slicing to avoid the Arrow serialization
                # bug where we include the entire base view on serialization.
                view = block.slice(start, end, copy=batch_size is not None)
                if batch_format == "native":
                    # Always promote Arrow blocks to pandas for consistency.
                    if isinstance(view, pa.Table) or isinstance(view, bytes):
                        view = BlockAccessor.for_block(view).to_pandas()
                elif batch_format == "pandas":
                    view = BlockAccessor.for_block(view).to_pandas()
                elif batch_format == "pyarrow":
                    view = BlockAccessor.for_block(view).to_arrow()
                else:
                    raise ValueError(
                        "The batch format must be one of 'native', 'pandas', "
                        "or 'pyarrow', got: {}".format(batch_format)
                    )

                applied = fn(view)
                if not (
                    isinstance(applied, list)
                    or isinstance(applied, pa.Table)
                    or isinstance(applied, pd.core.frame.DataFrame)
                ):
                    raise ValueError(
                        "The map batches UDF returned the value "
                        f"{applied}, which is not allowed. "
                        "The return type must be either list, "
                        "pandas.DataFrame, or pyarrow.Table"
                    )
                output_buffer.add_block(applied)
                if output_buffer.has_next():
                    yield output_buffer.next()

            output_buffer.finalize()
            if output_buffer.has_next():
                yield output_buffer.next()

        plan = self._plan.with_stage(
            OneToOneStage("map_batches", transform, compute, ray_remote_args)
        )
        return Dataset(plan, self._epoch, self._lazy)
Apply the given function to batches of records of this dataset.

        This is a blocking operation.

        Examples:
            >>> import ray
            >>> # Transform python objects.
            >>> ds = ray.data.range(1000) # doctest: +SKIP
            >>> # Transform batches in parallel.
            >>> ds.map_batches(lambda batch: [v * 2 for v in batch]) # doctest: +SKIP
            >>> # Define a callable class that persists state across
            >>> # function invocations for efficiency.
            >>> init_model = ... # doctest: +SKIP
            >>> class CachedModel:
            ...    def __init__(self):
            ...        self.model = init_model()
            ...    def __call__(self, item):
            ...        return self.model(item)
            >>> # Apply the transform in parallel on GPUs. Since
            >>> # compute=ActorPoolStrategy(2, 8) the transform will be applied on an
            >>> # autoscaling pool of 2-8 Ray actors, each allocated 1 GPU by Ray.
            >>> from ray.data.impl.compute import ActorPoolStrategy
            >>> ds.map_batches( # doctest: +SKIP
            ...     CachedModel, # doctest: +SKIP
            ...     batch_size=256, # doctest: +SKIP
            ...     compute=ActorPoolStrategy(2, 8), # doctest: +SKIP
            ...     num_gpus=1) # doctest: +SKIP

        Time complexity: O(dataset size / parallelism)

        Args:
            fn: The function to apply to each record batch, or a class type
                that can be instantiated to create such a callable. Callable classes are
                only supported for the actor compute strategy.
            batch_size: Request a specific batch size, or None to use entire
                blocks as batches. Defaults to a system-chosen batch size.
            compute: The compute strategy, either "tasks" (default) to use Ray
                tasks, or ActorPoolStrategy(min, max) to use an autoscaling actor pool.
            batch_format: Specify "native" to use the native block format
                (promotes Arrow to pandas), "pandas" to select
                ``pandas.DataFrame`` as the batch format,
                or "pyarrow" to select ``pyarrow.Table``.
            ray_remote_args: Additional resource requirements to request from
                ray (e.g., num_gpus=1 to request GPUs for the map tasks).
        
Apply the given function to batches of records of this dataset.
This is a blocking operation.
Apply the given function to batches of records of this dataset . This is a blocking operation .
def map_batches(
        self,
        fn: Union[CallableClass, Callable[[BatchType], BatchType]],
        *,
        batch_size: Optional[int] = 4096,
        compute: Union[str, ComputeStrategy] = None,
        batch_format: str = "native",
        **ray_remote_args,
    ) -> "Dataset[Any]":
        import pyarrow as pa
        import pandas as pd
        if batch_size is not None and batch_size < 1:
            raise ValueError("Batch size cannot be negative or 0")
        fn = cache_wrapper(fn, compute)
        context = DatasetContext.get_current()
        def transform(block: Block) -> Iterable[Block]:
            DatasetContext._set_current(context)
            output_buffer = BlockOutputBuffer(None, context.target_max_block_size)
            block = BlockAccessor.for_block(block)
            total_rows = block.num_rows()
            max_batch_size = batch_size
            if max_batch_size is None:
                max_batch_size = max(total_rows, 1)
            for start in range(0, total_rows, max_batch_size):
                end = min(total_rows, start + max_batch_size)
                view = block.slice(start, end, copy=batch_size is not None)
                if batch_format == "native":
                    if isinstance(view, pa.Table) or isinstance(view, bytes):
                        view = BlockAccessor.for_block(view).to_pandas()
                elif batch_format == "pandas":
                    view = BlockAccessor.for_block(view).to_pandas()
                elif batch_format == "pyarrow":
                    view = BlockAccessor.for_block(view).to_arrow()
                else:
                    raise ValueError(
                        "The batch format must be one of 'native', 'pandas', "
                        "or 'pyarrow', got: {}".format(batch_format)
                    )
                applied = fn(view)
                if not (
                    isinstance(applied, list)
                    or isinstance(applied, pa.Table)
                    or isinstance(applied, pd.core.frame.DataFrame)
                ):
                    raise ValueError(
                        "The map batches UDF returned the value "
                        f"{applied}, which is not allowed. "
                        "The return type must be either list, "
                        "pandas.DataFrame, or pyarrow.Table"
                    )
                output_buffer.add_block(applied)
                if output_buffer.has_next():
                    yield output_buffer.next()
            output_buffer.finalize()
            if output_buffer.has_next():
                yield output_buffer.next()
        plan = self._plan.with_stage(
            OneToOneStage("map_batches", transform, compute, ray_remote_args)
        )
        return Dataset(plan, self._epoch, self._lazy)
def map_batches ( self , fn : Union [ CallableClass , Callable [ [ BatchType ] , BatchType ] ] , * , batch_size : Optional [ int ] = 4096 , compute : Union [ str , ComputeStrategy ] = None , batch_format : str = "native" , ** ray_remote_args , ) -> "Dataset[Any]" : import pyarrow as pa import pandas as pd if batch_size is not None and batch_size < 1 : raise ValueError ( "Batch size cannot be negative or 0" ) fn = cache_wrapper ( fn , compute ) context = DatasetContext . get_current ( ) def transform ( block : Block ) -> Iterable [ Block ] : DatasetContext . _set_current ( context ) output_buffer = BlockOutputBuffer ( None , context . target_max_block_size ) block = BlockAccessor . for_block ( block ) total_rows = block . num_rows ( ) max_batch_size = batch_size if max_batch_size is None : max_batch_size = max ( total_rows , 1 ) for start in range ( 0 , total_rows , max_batch_size ) : end = min ( total_rows , start + max_batch_size ) view = block . slice ( start , end , copy = batch_size is not None ) if batch_format == "native" : if isinstance ( view , pa . Table ) or isinstance ( view , bytes ) : view = BlockAccessor . for_block ( view ) . to_pandas ( ) elif batch_format == "pandas" : view = BlockAccessor . for_block ( view ) . to_pandas ( ) elif batch_format == "pyarrow" : view = BlockAccessor . for_block ( view ) . to_arrow ( ) else : raise ValueError ( "The batch format must be one of 'native', 'pandas', " "or 'pyarrow', got: {}" . format ( batch_format ) ) applied = fn ( view ) if not ( isinstance ( applied , list ) or isinstance ( applied , pa . Table ) or isinstance ( applied , pd . core . frame . DataFrame ) ) : raise ValueError ( "The map batches UDF returned the value " f"{applied}, which is not allowed. " "The return type must be either list, " "pandas.DataFrame, or pyarrow.Table" ) output_buffer . add_block ( applied ) if output_buffer . has_next ( ) : yield output_buffer . next ( ) output_buffer . finalize ( ) if output_buffer . has_next ( ) : yield output_buffer . next ( ) plan = self . _plan . with_stage ( OneToOneStage ( "map_batches" , transform , compute , ray_remote_args ) ) return Dataset ( plan , self . _epoch , self . _lazy )
Apply the given function to batches of records of this dataset.
Apply the given function to batches of records of this dataset .
"""Apply the given function to batches of records of this dataset.

        This is a blocking operation.

        Examples:
            >>> import ray
            >>> # Transform python objects.
            >>> ds = ray.data.range(1000) # doctest: +SKIP
            >>> # Transform batches in parallel.
            >>> ds.map_batches(lambda batch: [v * 2 for v in batch]) # doctest: +SKIP
            >>> # Define a callable class that persists state across
            >>> # function invocations for efficiency.
            >>> init_model = ... # doctest: +SKIP
            >>> class CachedModel:
            ...    def __init__(self):
            ...        self.model = init_model()
            ...    def __call__(self, item):
            ...        return self.model(item)
            >>> # Apply the transform in parallel on GPUs. Since
            >>> # compute=ActorPoolStrategy(2, 8) the transform will be applied on an
            >>> # autoscaling pool of 2-8 Ray actors, each allocated 1 GPU by Ray.
            >>> from ray.data.impl.compute import ActorPoolStrategy
            >>> ds.map_batches( # doctest: +SKIP
            ...     CachedModel, # doctest: +SKIP
            ...     batch_size=256, # doctest: +SKIP
            ...     compute=ActorPoolStrategy(2, 8), # doctest: +SKIP
            ...     num_gpus=1) # doctest: +SKIP

        Time complexity: O(dataset size / parallelism)

        Args:
            fn: The function to apply to each record batch, or a class type
                that can be instantiated to create such a callable. Callable classes are
                only supported for the actor compute strategy.
            batch_size: Request a specific batch size, or None to use entire
                blocks as batches. Defaults to a system-chosen batch size.
            compute: The compute strategy, either "tasks" (default) to use Ray
                tasks, or ActorPoolStrategy(min, max) to use an autoscaling actor pool.
            batch_format: Specify "native" to use the native block format
                (promotes Arrow to pandas), "pandas" to select
                ``pandas.DataFrame`` as the batch format,
                or "pyarrow" to select ``pyarrow.Table``.
            ray_remote_args: Additional resource requirements to request from
                ray (e.g., num_gpus=1 to request GPUs for the map tasks).
        """ # Build a block for each batch. # Make sure to copy if slicing to avoid the Arrow serialization # bug where we include the entire base view on serialization. # Always promote Arrow blocks to pandas for consistency.
self
fn
Union[CallableClass, Callable[[BatchType], BatchType]]
The function to apply to each record batch, or a class type
that can be instantiated to create such a callable. Callable classes are
only supported for the actor compute strategy.
The function to apply to each record batch or a class type that can be instantiated to create such a callable . Callable classes are only supported for the actor compute strategy .
batch_size
Optional[int]
Request a specific batch size, or None to use entire
blocks as batches. Defaults to a system-chosen batch size.
Request a specific batch size or None to use entire blocks as batches . Defaults to a system - chosen batch size .
compute
Union[str, ComputeStrategy]
The compute strategy, either "tasks" (default) to use Ray
tasks, or ActorPoolStrategy(min, max) to use an autoscaling actor pool.
The compute strategy either " tasks " ( default ) to use Ray tasks or ActorPoolStrategy ( min max ) to use an autoscaling actor pool .
batch_format
str
Specify "native" to use the native block format
(promotes Arrow to pandas), "pandas" to select
``pandas.DataFrame`` as the batch format,
or "pyarrow" to select ``pyarrow.Table``.
Specify " native " to use the native block format ( promotes Arrow to pandas ) " pandas " to select ` ` pandas . DataFrame ` ` as the batch format or " pyarrow " to select ` ` pyarrow . Table ` ` .
ray_remote_args
Additional resource requirements to request from
ray .
Additional resource requirements to request from ray .
examples
>>> import ray
>>> # Transform python objects.
>>> ds = ray.data.range(1000) # doctest: +SKIP
>>> # Transform batches in parallel.
>>> ds.map_batches(lambda batch: [v * 2 for v in batch]) # doctest: +SKIP
>>> # Define a callable class that persists state across
>>> # function invocations for efficiency.
>>> import ray >>> # Transform python objects . >>> ds = ray . data . range ( 1000 ) # doctest : + SKIP >>> # Transform batches in parallel . >>> ds . map_batches ( lambda batch : [ v * 2 for v in batch ] ) # doctest : + SKIP >>> # Define a callable class that persists state across >>> # function invocations for efficiency .