chore: import upstream snapshot with attribution

doc/source/data/aggregating-data.rst

@@ -0,0 +1,168 @@
+.. _aggregations:
+
+Aggregating Data
+================
+
+Ray Data provides a flexible and performant API for performing aggregations on :class:`~ray.data.dataset.Dataset`. 
+
+Basic Aggregations
+------------------
+
+Ray Data provides several built-in aggregation functions like :class:`~ray.data.Dataset.max`,
+:class:`~ray.data.Dataset.min`, :class:`~ray.data.Dataset.sum`.
+
+These can be used directly on a Dataset or a GroupedData object, as shown below:
+
+.. testcode::
+
+    import ray
+
+    # Create a sample dataset
+    ds = ray.data.range(100)
+    ds = ds.add_column("group_key", lambda x: x["id"].to_numpy() % 3)
+    # Schema: {'id': int64, 'group_key': int64}
+
+    # Find the max
+    result = ds.max("id")
+    # result: 99
+
+    # Find the minimum value per group
+    result = ds.groupby("group_key").min("id")
+    # result: [{'group_key': 0, 'min(id)': 0}, {'group_key': 1, 'min(id)': 1}, {'group_key': 2, 'min(id)': 2}]
+
+The full list of built-in aggregation functions is available in the :ref:`Dataset API reference <dataset-api>`.
+
+Each of the preceding methods also has a corresponding :ref:`AggregateFnV2 <aggregations_api_ref>` object. These objects can be used in :meth:`~ray.data.Dataset.aggregate()` or :meth:`Dataset.groupby().aggregate() <ray.data.grouped_data.GroupedData.aggregate>`.
+
+Aggregation objects can be used directly with a Dataset like shown below:
+
+.. testcode::
+
+    import ray
+    from ray.data.aggregate import Count, Mean, Quantile
+
+    # Create a sample dataset
+    ds = ray.data.range(100)
+    ds = ds.add_column("group_key", lambda x: x["id"].to_numpy() % 3)
+
+    # Count all rows
+    result = ds.aggregate(Count())
+    # result: {'count()': 100}
+
+    # Calculate mean per group
+    result = ds.groupby("group_key").aggregate(Mean(on="id")).take_all()
+    # result: [{'group_key': 0, 'mean(id)': ...},
+    #          {'group_key': 1, 'mean(id)': ...},
+    #          {'group_key': 2, 'mean(id)': ...}]
+
+    # Calculate 75th percentile
+    result = ds.aggregate(Quantile(on="id", q=0.75))
+    # result: {'quantile(id)': 75.0}
+
+Multiple aggregations can also be computed at once:
+
+.. testcode::
+
+    import ray
+    from ray.data.aggregate import Count, Mean, Min, Max, Std
+
+    ds = ray.data.range(100)
+    ds = ds.add_column("group_key", lambda x: x["id"].to_numpy() % 3)
+
+    # Compute multiple aggregations at once
+    result = ds.groupby("group_key").aggregate(
+        Count(on="id"),
+        Mean(on="id"),
+        Min(on="id"),
+        Max(on="id"),
+        Std(on="id")
+    ).take_all()
+    # result: [{'group_key': 0, 'count(id)': 34, 'mean(id)': ..., 'min(id)': ..., 'max(id)': ..., 'std(id)': ...},
+    #          {'group_key': 1, 'count(id)': 33, 'mean(id)': ..., 'min(id)': ..., 'max(id)': ..., 'std(id)': ...},
+    #          {'group_key': 2, 'count(id)': 33, 'mean(id)': ..., 'min(id)': ..., 'max(id)': ..., 'std(id)': ...}]
+    
+
+Custom Aggregations
+--------------------
+
+You can create custom aggregations by implementing the :class:`~ray.data.aggregate.AggregateFnV2` interface. The AggregateFnV2 interface has three key methods to implement:
+
+1. `aggregate_block`: Processes a single block of data and returns a partial aggregation result
+2. `combine`: Merges two partial aggregation results into a single result
+3. `finalize`: Transforms the final accumulated result into the desired output format
+
+The aggregation process follows these steps:
+
+1. **Initialization**: For each group (if grouping) or for the entire dataset, an initial accumulator is created using `zero_factory`
+2. **Block Aggregation**: The `aggregate_block` method is applied to each block independently
+3. **Combination**: The `combine` method merges partial results into a single accumulator
+4. **Finalization**: The `finalize` method transforms the final accumulator into the desired output
+
+Example: Creating a Custom Mean Aggregator
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Here's an example of creating a custom aggregator that calculates the Mean of values in a column:
+
+.. testcode::
+
+    import numpy as np
+    from ray.data.aggregate import AggregateFnV2
+    from ray.data._internal.util import is_null
+    from ray.data.block import Block, BlockAccessor, AggType, U
+    import pyarrow.compute as pc
+    from typing import List, Optional
+
+    class Mean(AggregateFnV2):
+        """Defines mean aggregation."""
+
+        def __init__(
+            self,
+            on: Optional[str] = None,
+            ignore_nulls: bool = True,
+            alias_name: Optional[str] = None,
+        ):
+            super().__init__(
+                alias_name if alias_name else f"mean({str(on)})",
+                on=on,
+                ignore_nulls=ignore_nulls,
+                # NOTE: We've to copy returned list here, as some
+                #       aggregations might be modifying elements in-place
+                zero_factory=lambda: list([0, 0]),  # noqa: C410
+            )
+
+        def aggregate_block(self, block: Block) -> AggType:
+            block_acc = BlockAccessor.for_block(block)
+            count = block_acc.count(self._target_col_name, self._ignore_nulls)
+
+            if count == 0 or count is None:
+                # Empty or all null.
+                return None
+
+            sum_ = block_acc.sum(self._target_col_name, self._ignore_nulls)
+
+            if is_null(sum_):
+                # In case of ignore_nulls=False and column containing 'null'
+                # return as is (to prevent unnecessary type conversions, when, for ex,
+                # using Pandas and returning None)
+                return sum_
+
+            return [sum_, count]
+
+        def combine(self, current_accumulator: AggType, new: AggType) -> AggType:
+            return [current_accumulator[0] + new[0], current_accumulator[1] + new[1]]
+
+        def finalize(self, accumulator: AggType) -> Optional[U]:
+            if accumulator[1] == 0:
+                return np.nan
+
+            return accumulator[0] / accumulator[1]
+
+
+.. note::
+    Internally, aggregations support both the :ref:`hash-shuffle backend <hash-shuffle>` and the :ref:`range based backend <range-partitioning-shuffle>`.
+
+    Hash-shuffling can provide better performance for aggregations in certain cases. For more information see `comparison between hash based shuffling and Range Based shuffling approach <https://www.anyscale.com/blog/ray-data-joins-hash-shuffle#performance-benchmarks/>`_ .
+
+    To use the hash-shuffle algorithm for aggregations, you need to set the shuffle strategy explicitly:    
+    ``ray.data.DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE`` before creating a ``Dataset``
+    

doc/source/data/api/_autogen.rst

@@ -0,0 +1,24 @@
+:orphan:
+
+.. # This file is only used to auto-generate API docs.
+.. # It should not be included in the toctree.
+.. #
+.. # For any classes that you want to include in the
+.. # API docs, add them to the list of autosummary
+.. # below, then include the generated ray.data.<class>.rst
+.. # file in your top level rst file.
+
+.. currentmodule:: ray.data
+
+.. autosummary::
+    :nosignatures:
+    :template: autosummary/class_v2.rst
+    :toctree:
+
+    DataIterator
+    Dataset
+    Schema
+    stats.DatasetSummary
+    grouped_data.GroupedData
+    aggregate.AggregateFn
+    aggregate.AggregateFnV2

doc/source/data/api/aggregate.rst

@@ -0,0 +1,34 @@
+.. _aggregations_api_ref:
+
+Aggregation API
+===============
+
+Pass :class:`AggregateFnV2 <ray.data.aggregate.AggregateFnV2>` objects to
+:meth:`Dataset.aggregate() <ray.data.Dataset.aggregate>` or 
+:meth:`Dataset.groupby().aggregate() <ray.data.grouped_data.GroupedData.aggregate>` to 
+compute aggregations.
+
+.. currentmodule:: ray.data.aggregate
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    AggregateFnV2
+    AggregateFn
+    Count
+    Sum
+    Min
+    Max
+    Mean
+    Std
+    AbsMax
+    Quantile
+    Unique
+    AsList
+    CountDistinct
+    ValueCounter
+    MissingValuePercentage
+    ZeroPercentage
+    ApproximateQuantile
+    ApproximateTopK

doc/source/data/api/api.rst

@@ -0,0 +1,22 @@
+.. _data-api:
+
+Ray Data API
+================
+
+.. toctree::
+    :maxdepth: 2
+
+    loading_data.rst
+    saving_data.rst
+    dataset.rst
+    data_iterator.rst
+    execution_options.rst
+    checkpoint.rst
+    aggregate.rst
+    grouped_data.rst
+    expressions.rst
+    datatype.rst
+    data_context.rst
+    preprocessor.rst
+    llm.rst
+    from_other_data_libs.rst

doc/source/data/api/checkpoint.rst

@@ -0,0 +1,18 @@
+.. _checkpoint-api:
+
+Checkpoint API
+==============
+
+.. currentmodule:: ray.data.checkpoint.interfaces
+
+Configuration
+-------------
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+   :template: autosummary/class_without_autosummary.rst
+
+   CheckpointConfig
+   CheckpointBackend
+

doc/source/data/api/data_context.rst

@@ -0,0 +1,17 @@
+.. _data-context-api:
+
+Global configuration
+====================
+
+.. currentmodule:: ray.data.context
+
+.. autoclass:: DataContext
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   DataContext.get_current
+
+
+.. autoclass:: AutoscalingConfig

doc/source/data/api/data_iterator.rst

@@ -0,0 +1,6 @@
+.. _dataset-iterator-api:
+
+DataIterator API
+================
+
+.. include:: ray.data.DataIterator.rst

doc/source/data/api/dataset.rst

@@ -0,0 +1,68 @@
+.. _dataset-api:
+
+Dataset API
+==============
+
+.. include:: ray.data.Dataset.rst
+
+Compute Strategy API
+--------------------
+.. currentmodule:: ray.data
+.. autosummary::
+  :nosignatures:
+  :toctree: doc/
+
+  ActorPoolStrategy
+  TaskPoolStrategy
+
+Mixing API
+----------
+.. currentmodule:: ray.data
+
+.. autoclass:: MixStoppingCondition
+    :members:
+    :undoc-members:
+
+Schema
+------
+.. currentmodule:: ray.data
+
+.. autoclass:: Schema
+    :members:
+
+DatasetSummary
+--------------
+.. currentmodule:: ray.data.stats
+
+.. autoclass:: DatasetSummary
+    :members:
+
+Developer API
+-------------
+
+.. currentmodule:: ray.data
+
+.. autosummary::
+  :nosignatures:
+  :toctree: doc/
+
+  Dataset.to_pandas_refs
+  Dataset.to_numpy_refs
+  Dataset.to_arrow_refs
+  Dataset.iter_internal_ref_bundles
+  block.Block
+  block.BlockExecStats
+  block.BlockMetadata
+  block.BlockAccessor
+  RandomSeedConfig
+
+Deprecated API
+--------------
+
+.. currentmodule:: ray.data
+
+.. autosummary::
+  :nosignatures:
+  :toctree: doc/
+
+  Dataset.iter_tf_batches

doc/source/data/api/datatype.rst

@@ -0,0 +1,18 @@
+.. _datatype-api:
+
+Data types
+==========
+
+.. currentmodule:: ray.data.datatype
+
+Class
+-----
+
+.. autoclass:: DataType
+    :members:
+
+Enumeration
+-----------
+
+.. autoclass:: TypeCategory
+    :members:

doc/source/data/api/execution_options.rst

@@ -0,0 +1,26 @@
+.. _execution-options-api:
+
+ExecutionOptions API
+====================
+
+.. currentmodule:: ray.data
+
+Constructor
+-----------
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+   :template: autosummary/class_without_autosummary.rst
+
+   ExecutionOptions
+
+Resource Options
+----------------
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+   :template: autosummary/class_without_autosummary.rst
+
+   ExecutionResources

doc/source/data/api/expressions.rst

@@ -0,0 +1,218 @@
+.. _expressions-api:
+
+Expressions API
+================
+
+.. currentmodule:: ray.data.expressions
+
+Expressions provide a way to specify column-based operations on datasets.
+Use :func:`col` to reference columns and :func:`lit` to create literal values.
+You can combine these with operators to create complex expressions for filtering,
+transformations, and computations.
+
+Public API
+----------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    star
+    col
+    lit
+    udf
+    pyarrow_udf
+    download
+    monotonically_increasing_id
+    random
+    uuid
+
+Expression Classes
+------------------
+
+These classes represent the structure of expressions. You typically don't need to
+instantiate them directly, but you may encounter them when working with expressions.
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    Expr
+    ColumnExpr
+    LiteralExpr
+    BinaryExpr
+    UnaryExpr
+    UDFExpr
+    StarExpr
+    DownloadExpr
+    MonotonicallyIncreasingIdExpr
+    RandomExpr
+    UUIDExpr
+
+Expression namespaces
+------------------------------------
+
+These namespace classes provide specialized operations for list, string, struct, array, and
+`datetime` columns. You access them through properties on expressions: ``.list``, ``.str``,
+``.struct``, ``.arr``, and ``.dt``.
+
+The following example shows how to use the string namespace to transform text columns:
+
+.. testcode::
+
+    import ray
+    from ray.data.expressions import col
+
+    # Create a dataset with a text column
+    ds = ray.data.from_items([
+        {"name": "alice"},
+        {"name": "bob"},
+        {"name": "charlie"}
+    ])
+
+    # Use the string namespace to uppercase the names
+    ds = ds.with_column("upper_name", col("name").str.upper())
+    ds.show()
+
+.. testoutput::
+
+    {'name': 'alice', 'upper_name': 'ALICE'}
+    {'name': 'bob', 'upper_name': 'BOB'}
+    {'name': 'charlie', 'upper_name': 'CHARLIE'}
+
+The following example demonstrates using the list namespace to work with array columns:
+
+.. testcode::
+
+    import ray
+    from ray.data.expressions import col
+
+    # Create a dataset with list columns
+    ds = ray.data.from_items([
+        {"scores": [85, 90, 78]},
+        {"scores": [92, 88]},
+        {"scores": [76, 82, 88, 91]}
+    ])
+
+    # Use the list namespace to get the length of each list
+    ds = ds.with_column("num_scores", col("scores").list.len())
+    ds.show()
+
+.. testoutput::
+
+    {'scores': [85, 90, 78], 'num_scores': 3}
+    {'scores': [92, 88], 'num_scores': 2}
+    {'scores': [76, 82, 88, 91], 'num_scores': 4}
+
+You can also perform list-specific transformations like sorting and flattening:
+
+.. testcode::
+
+    import ray
+    from ray.data.expressions import col
+
+    ds = ray.data.from_items([
+        {"values": [3, 1, 2], "nested": [[1, 2], [3]]},
+        {"values": [2, None, 5], "nested": [[4], []]}
+    ])
+
+    ds = ds.with_column(
+        "sorted_values", col("values").list.sort(order="descending")
+    )
+    ds = ds.with_column(
+        "flattened_nested", col("nested").list.flatten()
+    )
+    ds.show()
+
+.. testoutput::
+
+    {'values': [3, 1, 2], 'nested': [[1, 2], [3]], 'sorted_values': [3, 2, 1], 'flattened_nested': [1, 2, 3]}
+    {'values': [2, None, 5], 'nested': [[4], []], 'sorted_values': [5, 2, None], 'flattened_nested': [4]}
+
+The following example shows how to use the struct namespace to access nested fields:
+
+.. testcode::
+
+    import ray
+    from ray.data.expressions import col
+
+    # Create a dataset with struct columns
+    ds = ray.data.from_items([
+        {"user": {"name": "alice", "age": 25}},
+        {"user": {"name": "bob", "age": 30}},
+        {"user": {"name": "charlie", "age": 35}}
+    ])
+
+    # Use the struct namespace to extract a specific field
+    ds = ds.with_column("user_name", col("user").struct.field("name"))
+    ds.show()
+
+.. testoutput::
+
+    {'user': {'name': 'alice', 'age': 25}, 'user_name': 'alice'}
+    {'user': {'name': 'bob', 'age': 30}, 'user_name': 'bob'}
+    {'user': {'name': 'charlie', 'age': 35}, 'user_name': 'charlie'}
+
+The following example shows how to use the array namespace to convert fixed-size
+list columns to variable-length lists:
+
+.. testcode::
+
+    import pyarrow as pa
+    import ray
+    from ray.data.expressions import col
+
+    values = pa.array([1, 2, 3, 4])
+    fixed = pa.FixedSizeListArray.from_arrays(values, 2)
+    table = pa.table({"features": fixed})
+
+    ds = ray.data.from_arrow(table)
+    ds = ds.with_column("features_list", col("features").arr.to_list())
+    ds.show()
+
+.. testoutput::
+
+    {'features': [1, 2], 'features_list': [1, 2]}
+    {'features': [3, 4], 'features_list': [3, 4]}
+
+The following example shows how to use the `datetime` namespace to extract components:
+
+.. testcode::
+
+    import datetime
+    import pandas as pd
+    import ray
+    from ray.data.expressions import col
+
+    ds = ray.data.from_items([
+        {"ts": pd.Timestamp("2024-01-02 03:04:05")},
+        {"ts": pd.Timestamp("2024-02-03 04:05:06")}
+    ])
+
+    ds = ds.with_column("year", col("ts").dt.year())
+    ds.show()
+
+.. testoutput::
+
+    {'ts': datetime.datetime(2024, 1, 2, 3, 4, 5), 'year': 2024}
+    {'ts': datetime.datetime(2024, 2, 3, 4, 5, 6), 'year': 2024}
+
+.. autoclass:: _ListNamespace
+    :members:
+    :exclude-members: _expr
+
+.. autoclass:: _StringNamespace
+    :members:
+    :exclude-members: _expr
+
+.. autoclass:: _StructNamespace
+    :members:
+    :exclude-members: _expr
+
+.. autoclass:: _ArrayNamespace
+    :members:
+    :exclude-members: _expr
+
+.. autoclass:: _DatetimeNamespace
+    :members:
+    :exclude-members: _expr

doc/source/data/api/from_other_data_libs.rst

@@ -0,0 +1,93 @@
+.. _api-guide-for-users-from-other-data-libs:
+
+API Guide for Users from Other Data Libraries
+=============================================
+
+Ray Data is a data loading and preprocessing library for ML. It shares certain
+similarities with other ETL data processing libraries, but also has its own focus.
+This guide provides API mappings for users who come from those data
+libraries, so you can quickly map what you may already know to Ray Data APIs.
+
+.. note::
+
+  - This is meant to map APIs that perform comparable but not necessarily identical operations.
+    Select the API reference for exact semantics and usage.
+  - This list may not be exhaustive: It focuses on common APIs or APIs that are less obvious to see a connection.
+
+.. _api-guide-for-pandas-users:
+
+For Pandas Users
+----------------
+
+.. list-table:: Pandas DataFrame vs. Ray Data APIs
+   :header-rows: 1
+
+   * - Pandas DataFrame API
+     - Ray Data API
+   * - df.head()
+     - :meth:`ds.show() <ray.data.Dataset.show>`, :meth:`ds.take() <ray.data.Dataset.take>`, or :meth:`ds.take_batch() <ray.data.Dataset.take_batch>`
+   * - df.dtypes
+     - :meth:`ds.schema() <ray.data.Dataset.schema>`
+   * - len(df) or df.shape[0]
+     - :meth:`ds.count() <ray.data.Dataset.count>`
+   * - df.truncate()
+     - :meth:`ds.limit() <ray.data.Dataset.limit>`
+   * - df.iterrows()
+     - :meth:`ds.iter_rows() <ray.data.Dataset.iter_rows>`
+   * - df.drop()
+     - :meth:`ds.drop_columns() <ray.data.Dataset.drop_columns>`
+   * - df.transform()
+     - :meth:`ds.map_batches() <ray.data.Dataset.map_batches>` or :meth:`ds.map() <ray.data.Dataset.map>`
+   * - df.groupby()
+     - :meth:`ds.groupby() <ray.data.Dataset.groupby>`
+   * - df.groupby().apply()
+     - :meth:`ds.groupby().map_groups() <ray.data.grouped_data.GroupedData.map_groups>`
+   * - df.sample()
+     - :meth:`ds.random_sample() <ray.data.Dataset.random_sample>`
+   * - df.sort_values()
+     - :meth:`ds.sort() <ray.data.Dataset.sort>`
+   * - df.append()
+     - :meth:`ds.union() <ray.data.Dataset.union>`
+   * - df.aggregate()
+     - :meth:`ds.aggregate() <ray.data.Dataset.aggregate>`
+   * - df.min()
+     - :meth:`ds.min() <ray.data.Dataset.min>`
+   * - df.max()
+     - :meth:`ds.max() <ray.data.Dataset.max>`
+   * - df.sum()
+     - :meth:`ds.sum() <ray.data.Dataset.sum>`
+   * - df.mean()
+     - :meth:`ds.mean() <ray.data.Dataset.mean>`
+   * - df.std()
+     - :meth:`ds.std() <ray.data.Dataset.std>`
+
+.. _api-guide-for-pyarrow-users:
+
+For PyArrow Users
+-----------------
+
+.. list-table:: PyArrow Table vs. Ray Data APIs
+   :header-rows: 1
+
+   * - PyArrow Table API
+     - Ray Data API
+   * - ``pa.Table.schema``
+     - :meth:`ds.schema() <ray.data.Dataset.schema>`
+   * - ``pa.Table.num_rows``
+     - :meth:`ds.count() <ray.data.Dataset.count>`
+   * - ``pa.Table.filter()``
+     - :meth:`ds.filter() <ray.data.Dataset.filter>`
+   * - ``pa.Table.drop()``
+     - :meth:`ds.drop_columns() <ray.data.Dataset.drop_columns>`
+   * - ``pa.Table.add_column()``
+     - :meth:`ds.with_column() <ray.data.Dataset.with_column>`
+   * - ``pa.Table.groupby()``
+     - :meth:`ds.groupby() <ray.data.Dataset.groupby>`
+   * - ``pa.Table.sort_by()``
+     - :meth:`ds.sort() <ray.data.Dataset.sort>`
+
+
+For PyTorch Dataset & DataLoader Users
+--------------------------------------
+
+For more details, see the :ref:`Migrating from PyTorch to Ray Data <migrate_pytorch>`.

doc/source/data/api/grouped_data.rst

@@ -0,0 +1,11 @@
+.. _grouped-dataset-api:
+
+GroupedData API
+===============
+
+.. currentmodule:: ray.data
+
+The groupby call returns GroupedData objects:
+:meth:`Dataset.groupby() <ray.data.Dataset.groupby>`.
+
+.. include:: ray.data.grouped_data.GroupedData.rst

doc/source/data/api/llm.rst

@@ -0,0 +1,52 @@
+.. _llm-ref:
+
+Large Language Model (LLM) API
+==============================
+
+.. currentmodule:: ray.data.llm
+
+LLM processor builder
+---------------------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~build_processor
+
+Processor
+---------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~Processor
+
+Processor configs
+-----------------
+
+.. autosummary::
+    :nosignatures:
+    :template: autosummary/class_without_autosummary_noinheritance.rst
+    :toctree: doc/
+
+    ~ProcessorConfig
+    ~HttpRequestProcessorConfig
+    ~vLLMEngineProcessorConfig
+    ~SGLangEngineProcessorConfig
+
+.. _stage-configs-ref:
+
+Stage configs
+-------------
+
+.. autosummary::
+    :nosignatures:
+    :template: autosummary/class_without_autosummary_noinheritance.rst
+    :toctree: doc/
+
+    ~ChatTemplateStageConfig
+    ~TokenizerStageConfig
+    ~DetokenizeStageConfig
+    ~PrepareMultimodalStageConfig

doc/source/data/api/loading_data.rst

@@ -0,0 +1,467 @@
+.. _loading-data-api:
+
+Loading Data API
+================
+
+.. currentmodule:: ray.data
+
+Public APIs
+-----------
+
+Arrow
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_arrow
+
+Audio
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_audio
+
+Avro
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_avro
+
+BigQuery
+^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_bigquery
+
+Binary
+^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_binary_files
+
+Catalog
+^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Catalog
+
+   DatabricksUnityCatalog
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+   :template: autosummary/class_without_autosummary.rst
+
+   ReaderFormat
+
+CSV
+^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_csv
+
+ClickHouse
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_clickhouse
+
+Daft
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_daft
+
+Dask
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_dask
+
+Databricks
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_databricks_tables
+
+Delta Lake
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_delta
+
+Delta Sharing
+^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_delta_sharing_tables
+
+Hudi
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_hudi
+
+Hugging Face
+^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_huggingface
+
+Iceberg
+^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_iceberg
+
+Images
+^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_images
+
+JSON
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_json
+
+Kafka
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_kafka
+
+Lance
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_lance
+
+MCAP (Message Capture)
+^^^^^^^^^^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_mcap
+
+Mars
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_mars
+
+Modin
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_modin
+
+MongoDB
+^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_mongo
+
+NumPy
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_numpy
+   read_numpy
+
+Pandas
+^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_pandas
+
+Parquet
+^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_parquet
+
+Python Objects
+^^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_items
+
+SQL Databases
+^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_sql
+
+Snowflake
+^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_snowflake
+
+Spark
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_spark
+
+Synthetic Data
+^^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   range
+   range_tensor
+
+TFRecords
+^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_tfrecords
+
+TensorFlow
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_tf
+
+Text
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_text
+
+Torch
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_torch
+
+Unity Catalog
+^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_unity_catalog
+
+Video
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_videos
+
+WebDataset
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_webdataset
+
+Zarr
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   read_zarr
+
+Partitioning API
+^^^^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   datasource.Partitioning
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+   :template: autosummary/class_without_autosummary.rst
+
+   datasource.PartitionStyle
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   datasource.PathPartitionFilter
+   datasource.PathPartitionParser
+
+Shuffling API
+^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   FileShuffleConfig
+
+Developer APIs
+--------------
+
+Arrow refs
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_arrow_refs
+
+Datasource API
+^^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Datasource
+   datasource.FileBasedDatasource
+   read_datasource
+   ReadTask
+
+.. _metadata_provider:
+
+MetadataProvider API
+^^^^^^^^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   datasource.BaseFileMetadataProvider
+   datasource.DefaultFileMetadataProvider
+   datasource.FileMetadataProvider
+
+NumPy refs
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_numpy_refs
+
+Pandas refs
+^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   from_pandas_refs

doc/source/data/api/preprocessor.rst

@@ -0,0 +1,100 @@
+.. _preprocessor-ref:
+
+Preprocessor
+============
+
+Preprocessor Interface
+------------------------
+
+.. currentmodule:: ray.data
+
+Constructor
+~~~~~~~~~~~
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessor.Preprocessor
+
+Fit/Transform APIs
+~~~~~~~~~~~~~~~~~~
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessor.Preprocessor.fit
+    ~preprocessor.Preprocessor.fit_transform
+    ~preprocessor.Preprocessor.transform
+    ~preprocessor.Preprocessor.transform_batch
+    ~preprocessor.PreprocessorNotFittedException
+
+
+Generic Preprocessors
+---------------------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessors.Concatenator
+    ~preprocessors.SimpleImputer
+    ~preprocessors.Chain
+
+Categorical Encoders
+--------------------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessors.Categorizer
+    ~preprocessors.LabelEncoder
+    ~preprocessors.MultiHotEncoder
+    ~preprocessors.OneHotEncoder
+    ~preprocessors.OrdinalEncoder
+
+Feature Scalers
+---------------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessors.MaxAbsScaler
+    ~preprocessors.MinMaxScaler
+    ~preprocessors.Normalizer
+    ~preprocessors.PowerTransformer
+    ~preprocessors.RobustScaler
+    ~preprocessors.StandardScaler
+
+K-Bins Discretizers
+-------------------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessors.CustomKBinsDiscretizer
+    ~preprocessors.UniformKBinsDiscretizer
+
+Feature Hashers and Vectorizers
+-------------------------------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessors.FeatureHasher
+    ~preprocessors.CountVectorizer
+    ~preprocessors.HashingVectorizer
+
+Specialized Preprocessors
+-------------------------
+
+.. autosummary::
+    :nosignatures:
+    :toctree: doc/
+
+    ~preprocessors.TorchVisionPreprocessor

doc/source/data/api/saving_data.rst

@@ -0,0 +1,234 @@
+.. _saving-data-api:
+
+Saving Data API
+===============
+
+.. currentmodule:: ray.data
+
+Public APIs
+-----------
+
+BigQuery
+^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_bigquery
+
+CSV
+^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_csv
+
+ClickHouse
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_clickhouse
+
+Daft
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_daft
+
+Dask
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_dask
+
+Iceberg
+^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_iceberg
+
+Images
+^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_images
+
+JSON
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_json
+
+Lance
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_lance
+
+Mars
+^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_mars
+
+Modin
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_modin
+
+MongoDB
+^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_mongo
+
+NumPy
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_numpy
+
+Pandas
+^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_pandas
+
+Parquet
+^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_parquet
+
+SQL Databases
+^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_sql
+
+Snowflake
+^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_snowflake
+
+Spark
+^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_spark
+
+TFRecords
+^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.write_tfrecords
+
+Developer APIs
+--------------
+
+Arrow refs
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_arrow_refs
+
+Datasink API
+^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Datasink
+   Dataset.write_datasink
+   datasource.RowBasedFileDatasink
+   datasource.BlockBasedFileDatasink
+   datasource.WriteResult
+   datasource.WriteReturnType
+
+FilenameProvider
+^^^^^^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   datasource.FilenameProvider
+
+NumPy refs
+^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_numpy_refs
+
+Pandas refs
+^^^^^^^^^^^
+
+.. autosummary::
+   :nosignatures:
+   :toctree: doc/
+
+   Dataset.to_pandas_refs
+

doc/source/data/batch_inference.rst

@@ -0,0 +1,491 @@
+.. _batch_inference_home:
+
+End-to-end: Offline Batch Inference
+===================================
+
+Offline batch inference is a process for generating model predictions on a fixed set of input data. Ray Data offers an efficient and scalable solution for batch inference, providing faster execution and cost-effectiveness for deep learning applications.
+
+..
+ https://docs.google.com/presentation/d/1l03C1-4jsujvEFZUM4JVNy8Ju8jnY5Lc_3q7MBWi2PQ/edit#slide=id.g230eb261ad2_0_0
+
+.. image:: images/stream-example.png
+   :width: 650px
+   :align: center
+
+.. note::
+    This guide is primarily focused on batch inference with deep learning frameworks.
+    For more information on batch inference with LLMs, see :ref:`Working with LLMs <working-with-llms>`.
+
+.. _batch_inference_quickstart:
+
+Quickstart
+----------
+To start, install Ray Data:
+
+.. code-block:: bash
+
+    pip install -U "ray[data]"
+
+Using Ray Data for offline inference involves four basic steps:
+
+- **Step 1:** Load your data into a Ray Dataset. Ray Data supports many different datasources and formats. For more details, see :ref:`Loading Data <loading_data>`.
+- **Step 2:** Define a Python class to load the pre-trained model.
+- **Step 3:** Transform your dataset using the pre-trained model by calling :meth:`ds.map_batches() <ray.data.Dataset.map_batches>`. For more details, see :ref:`Transforming Data <transforming_data>`.
+- **Step 4:** Get the final predictions by either iterating through the output or saving the results. For more details, see the :ref:`Iterating over data <iterating-over-data>` and :ref:`Saving data <saving-data>` user guides.
+
+For more in-depth examples for your use case, see :doc:`the batch inference examples</data/examples>`.
+For how to configure batch inference, see :ref:`the configuration guide<batch_inference_configuration>`.
+
+.. tab-set::
+
+    .. tab-item:: HuggingFace
+        :sync: HuggingFace
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+
+            import ray
+
+            # Step 1: Create a Ray Dataset from in-memory Numpy arrays.
+            # You can also create a Ray Dataset from many other sources and file
+            # formats.
+            ds = ray.data.from_numpy(np.asarray(["Complete this", "for me"]))
+
+            # Step 2: Define a Predictor class for inference.
+            # Use a class to initialize the model just once in `__init__`
+            # and reuse it for inference across multiple batches.
+            class HuggingFacePredictor:
+                def __init__(self):
+                    from transformers import pipeline
+                    # Initialize a pre-trained GPT2 Huggingface pipeline.
+                    self.model = pipeline("text-generation", model="gpt2")
+
+                # Logic for inference on 1 batch of data.
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, list]:
+                    # Get the predictions from the input batch.
+                    predictions = self.model(list(batch["data"]), max_length=20, num_return_sequences=1)
+                    # `predictions` is a list of length-one lists. For example:
+                    # [[{'generated_text': 'output_1'}], ..., [{'generated_text': 'output_2'}]]
+                    # Modify the output to get it into the following format instead:
+                    # ['output_1', 'output_2']
+                    batch["output"] = [sequences[0]["generated_text"] for sequences in predictions]
+                    return batch
+
+            # Step 2: Map the Predictor over the Dataset to get predictions.
+            # Use 2 parallel actors for inference. Each actor predicts on a
+            # different partition of data.
+            predictions = ds.map_batches(
+                HuggingFacePredictor, 
+                compute=ray.data.ActorPoolStrategy(size=2), 
+                batch_size="auto"
+            )
+            # Step 3: Show one prediction output.
+            predictions.show(limit=1)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'data': 'Complete this', 'output': 'Complete this information or purchase any item from this site.\n\nAll purchases are final and non-'}
+
+
+    .. tab-item:: PyTorch
+        :sync: PyTorch
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+            import torch
+            import torch.nn as nn
+
+            import ray
+
+            # Step 1: Create a Ray Dataset from in-memory Numpy arrays.
+            # You can also create a Ray Dataset from many other sources and file
+            # formats.
+            ds = ray.data.from_numpy(np.ones((1, 100)))
+
+            # Step 2: Define a Predictor class for inference.
+            # Use a class to initialize the model just once in `__init__`
+            # and reuse it for inference across multiple batches.
+            class TorchPredictor:
+                def __init__(self):
+                    # Load a dummy neural network.
+                    # Set `self.model` to your pre-trained PyTorch model.
+                    self.model = nn.Sequential(
+                        nn.Linear(in_features=100, out_features=1),
+                        nn.Sigmoid(),
+                    )
+                    self.model.eval()
+
+                # Logic for inference on 1 batch of data.
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                    tensor = torch.as_tensor(batch["data"], dtype=torch.float32)
+                    with torch.inference_mode():
+                        # Get the predictions from the input batch.
+                        return {"output": self.model(tensor).numpy()}
+
+            # Step 2: Map the Predictor over the Dataset to get predictions.
+            # Use 2 parallel actors for inference. Each actor predicts on a
+            # different partition of data.
+            predictions = ds.map_batches(TorchPredictor, compute=ray.data.ActorPoolStrategy(size=2), batch_size="auto")
+            # Step 3: Show one prediction output.
+            predictions.show(limit=1)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'output': array([0.5590901], dtype=float32)}
+
+    .. tab-item:: TensorFlow
+        :sync: TensorFlow
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+
+            import ray
+
+            # Step 1: Create a Ray Dataset from in-memory Numpy arrays.
+            # You can also create a Ray Dataset from many other sources and file
+            # formats.
+            ds = ray.data.from_numpy(np.ones((1, 100)))
+
+            # Step 2: Define a Predictor class for inference.
+            # Use a class to initialize the model just once in `__init__`
+            # and reuse it for inference across multiple batches.
+            class TFPredictor:
+                def __init__(self):
+                    from tensorflow import keras
+
+                    # Load a dummy neural network.
+                    # Set `self.model` to your pre-trained Keras model.
+                    input_layer = keras.Input(shape=(100,))
+                    output_layer = keras.layers.Dense(1, activation="sigmoid")
+                    self.model = keras.Sequential([input_layer, output_layer])
+
+                # Logic for inference on 1 batch of data.
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                    # Get the predictions from the input batch.
+                    return {"output": self.model(batch["data"]).numpy()}
+
+            # Step 2: Map the Predictor over the Dataset to get predictions.
+            # Use 2 parallel actors for inference. Each actor predicts on a
+            # different partition of data.
+            predictions = ds.map_batches(TFPredictor, compute=ray.data.ActorPoolStrategy(size=2), batch_size="auto")
+             # Step 3: Show one prediction output.
+            predictions.show(limit=1)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'output': array([0.625576], dtype=float32)}
+
+    .. tab-item:: LLM Inference
+        :sync: vLLM
+
+        Ray Data offers native integration with vLLM, a high-performance inference engine for large language models (LLMs).
+
+        .. testcode::
+            :skipif: True
+
+            import ray
+            from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+            import numpy as np
+
+            config = vLLMEngineProcessorConfig(
+                model="unsloth/Llama-3.1-8B-Instruct",
+                engine_kwargs={
+                    "enable_chunked_prefill": True,
+                    "max_num_batched_tokens": 4096,
+                    "max_model_len": 16384,
+                },
+                concurrency=1,
+                batch_size=64,
+            )
+            processor = build_processor(
+                config,
+                preprocess=lambda row: dict(
+                    messages=[
+                        {"role": "system", "content": "You are a bot that responds with haikus."},
+                        {"role": "user", "content": row["item"]}
+                    ],
+                    sampling_params=dict(
+                        temperature=0.3,
+                        max_tokens=250,
+                    )
+                ),
+                postprocess=lambda row: dict(
+                    answer=row["generated_text"]
+                ),
+            )
+
+            ds = ray.data.from_items(["Start of the haiku is: Complete this for me..."])
+
+            ds = processor(ds)
+            ds.show(limit=1)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'answer': 'Snowflakes gently fall\nBlanketing the winter scene\nFrozen peaceful hush'}
+
+.. _batch_inference_configuration:
+
+Configuration and troubleshooting
+---------------------------------
+
+.. _batch_inference_gpu:
+
+Job-level Checkpointing
+~~~~~~~~~~~~~~~~~~~~~~~
+
+Use job-level checkpointing to make offline batch inference jobs resilient to failures
+like node restarts or transient execution errors.
+
+When enabled, Ray Data records progress during execution. If a batch inference
+job fails partway through processing, rerunning the same pipeline with the same
+checkpoint configuration resumes by skipping already-processed records instead
+of reprocessing the entire dataset.
+
+This is especially useful for large batch inference workloads where restarting
+from the beginning would be expensive.
+
+To enable job-level checkpointing, configure a
+:class:`~ray.data.checkpoint.CheckpointConfig` on the current
+:class:`~ray.data.DataContext`. See the
+:ref:`Execution Configurations <execution_configurations>` guide for details.
+
+Using GPUs for inference
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+To use GPUs for inference, make the following changes to your code:
+
+1. Update the class implementation to move the model and data to and from GPU.
+2. Specify ``num_gpus=1`` in the :meth:`ds.map_batches() <ray.data.Dataset.map_batches>` call to indicate that each actor should use 1 GPU.
+3. Specify a ``batch_size`` for inference. For more details on how to configure the batch size, see :ref:`Configuring Batch Size <batch_inference_batch_size>`.
+
+The remaining is the same as the :ref:`Quickstart <batch_inference_quickstart>`.
+
+.. tab-set::
+
+    .. tab-item:: HuggingFace
+        :sync: HuggingFace
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+
+            import ray
+
+            ds = ray.data.from_numpy(np.asarray(["Complete this", "for me"]))
+
+            class HuggingFacePredictor:
+                def __init__(self):
+                    from transformers import pipeline
+                    # Set "cuda:0" as the device so the Huggingface pipeline uses GPU.
+                    self.model = pipeline("text-generation", model="gpt2", device="cuda:0")
+
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, list]:
+                    predictions = self.model(list(batch["data"]), max_length=20, num_return_sequences=1)
+                    batch["output"] = [sequences[0]["generated_text"] for sequences in predictions]
+                    return batch
+
+            # Use 2 actors, each actor using 1 GPU. 2 GPUs total.
+            predictions = ds.map_batches(
+                HuggingFacePredictor,
+                num_gpus=1,
+                # Specify the batch size for inference.
+                # Increase this for larger datasets.
+                batch_size=1,
+                # Set the concurrency to the number of GPUs in your cluster.
+                compute=ray.data.ActorPoolStrategy(size=2),
+                )
+            predictions.show(limit=1)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'data': 'Complete this', 'output': 'Complete this poll. Which one do you think holds the most promise for you?\n\nThank you'}
+
+
+    .. tab-item:: PyTorch
+        :sync: PyTorch
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+            import torch
+            import torch.nn as nn
+
+            import ray
+
+            ds = ray.data.from_numpy(np.ones((1, 100)))
+
+            class TorchPredictor:
+                def __init__(self):
+                    # Move the neural network to GPU device by specifying "cuda".
+                    self.model = nn.Sequential(
+                        nn.Linear(in_features=100, out_features=1),
+                        nn.Sigmoid(),
+                    ).cuda()
+                    self.model.eval()
+
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                    # Move the input batch to GPU device by specifying "cuda".
+                    tensor = torch.as_tensor(batch["data"], dtype=torch.float32, device="cuda")
+                    with torch.inference_mode():
+                        # Move the prediction output back to CPU before returning.
+                        return {"output": self.model(tensor).cpu().numpy()}
+
+            # Use 2 actors, each actor using 1 GPU. 2 GPUs total.
+            predictions = ds.map_batches(
+                TorchPredictor,
+                num_gpus=1,
+                # Specify the batch size for inference.
+                # Increase this for larger datasets.
+                batch_size=1,
+                # Set the concurrency to the number of GPUs in your cluster.
+                compute=ray.data.ActorPoolStrategy(size=2),
+                )
+            predictions.show(limit=1)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'output': array([0.5590901], dtype=float32)}
+
+    .. tab-item:: TensorFlow
+        :sync: TensorFlow
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+
+            import ray
+
+            ds = ray.data.from_numpy(np.ones((1, 100)))
+
+            class TFPredictor:
+                def __init__(self):
+                    import tensorflow as tf
+                    from tensorflow import keras
+
+                    # Move the neural network to GPU by specifying the GPU device.
+                    with tf.device("GPU:0"):
+                        input_layer = keras.Input(shape=(100,))
+                        output_layer = keras.layers.Dense(1, activation="sigmoid")
+                        self.model = keras.Sequential([input_layer, output_layer])
+
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                    import tensorflow as tf
+
+                    # Move the input batch to GPU by specifying GPU device.
+                    with tf.device("GPU:0"):
+                        return {"output": self.model(batch["data"]).numpy()}
+
+            # Use 2 actors, each actor using 1 GPU. 2 GPUs total.
+            predictions = ds.map_batches(
+                TFPredictor,
+                num_gpus=1,
+                # Specify the batch size for inference.
+                # Increase this for larger datasets.
+                batch_size=1,
+                # Set the concurrency to the number of GPUs in your cluster.
+                compute=ray.data.ActorPoolStrategy(size=2),
+            )
+            predictions.show(limit=1)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'output': array([0.625576], dtype=float32)}
+
+.. _batch_inference_batch_size:
+
+Configuring Batch Size
+~~~~~~~~~~~~~~~~~~~~~~
+
+Configure the size of the input batch that's passed to ``__call__`` by setting the ``batch_size`` argument for :meth:`ds.map_batches() <ray.data.Dataset.map_batches>`.
+
+Increasing batch size results in faster execution because inference is a vectorized operation. For GPU inference, increasing batch size increases GPU utilization.
+
+For **CPU inference**, use ``batch_size="auto"`` to let Ray Data automatically determine an appropriate batch size based on your data. For **GPU inference**, specify an explicit integer ``batch_size`` as large as possible without running out of GPU memory. If you encounter out-of-memory errors, decrease ``batch_size``.
+
+.. testcode::
+
+    import numpy as np
+
+    import ray
+
+    ds = ray.data.from_numpy(np.ones((10, 100)))
+
+    def assert_batch(batch: Dict[str, np.ndarray]):
+        assert len(batch) == 2
+        return batch
+
+    # Specify that each input batch should be of size 2.
+    ds.map_batches(assert_batch, batch_size=2)
+
+Handling GPU out-of-memory failures
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If you run into CUDA out-of-memory issues, your batch size is likely too large. Decrease
+the batch size by following :ref:`these steps <batch_inference_batch_size>`. If your
+batch size is already set to 1, then use either a smaller model or GPU devices with more
+memory.
+
+For advanced users working with large models, you can use model parallelism to shard the model across multiple GPUs.
+
+Optimizing expensive CPU preprocessing
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If your workload involves expensive CPU preprocessing in addition to model inference, you can optimize throughput by separating the preprocessing and inference logic into separate operations. This separation allows inference on batch :math:`N` to execute concurrently with preprocessing on batch :math:`N+1`.
+
+For an example where preprocessing is done in a separate `map` call, see :doc:`Image Classification Batch Inference with PyTorch ResNet18 </data/examples/pytorch_resnet_batch_prediction>`.
+
+Handling CPU out-of-memory failures
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If you run out of CPU RAM, you likely have too many model replicas that are running concurrently on the same node. For example, if a model
+uses 5 GB of RAM when created / run, and a machine has 16 GB of RAM total, then no more
+than three of these models can be run at the same time. The default resource assignments
+of one CPU per task/actor might lead to `OutOfMemoryError` from Ray in this situation.
+
+Suppose your cluster has 4 nodes, each with 16 CPUs. To limit to at most
+3 of these actors per node, you can override the CPU or memory:
+
+.. testcode::
+    :skipif: True
+
+    from typing import Dict
+    import numpy as np
+
+    import ray
+
+    ds = ray.data.from_numpy(np.asarray(["Complete this", "for me"]))
+
+    class HuggingFacePredictor:
+        def __init__(self):
+            from transformers import pipeline
+            self.model = pipeline("text-generation", model="gpt2")
+
+        def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, list]:
+            predictions = self.model(list(batch["data"]), max_length=20, num_return_sequences=1)
+            batch["output"] = [sequences[0]["generated_text"] for sequences in predictions]
+            return batch
+
+    predictions = ds.map_batches(
+        HuggingFacePredictor,
+        # Require 5 CPUs per actor (so at most 3 can fit per 16 CPU node).
+        num_cpus=5,
+        # 3 actors per node, with 4 nodes in the cluster means concurrency of 12.
+        compute=ray.data.ActorPoolStrategy(size=12),
+        )
+    predictions.show(limit=1)

doc/source/data/benchmark.md

@@ -0,0 +1,136 @@
+# Ray Data Benchmarks
+
+This page documents benchmark results and methodologies for evaluating Ray Data performance across a variety of data modalities and workloads.
+
+---
+
+## Workload Summary
+
+- **Image Classification**: Processing 800k ImageNet images using ResNet18. The pipeline downloads images, deserializes them, applies transformations, runs ResNet18 inference on GPU, and outputs predicted labels.
+- **Document Embedding**: Processing 10k PDF documents from Digital Corpora. The pipeline reads PDF documents, extracts text page-by-page, splits into chunks with overlap, embeds using a `all-MiniLM-L6-v2` model on GPU, and outputs embeddings with metadata.
+- **Audio Transcription**: Transcribing 113,800 audio files from Mozilla Common Voice 17 dataset using a Whisper-tiny model. The pipeline loads FLAC audio files, resamples to 16kHz, extracts features using Whisper's processor, runs GPU-accelerated batch inference with the model, and outputs transcriptions with metadata.
+- **Video Object Detection**: Processing 10k video frames from Hollywood2 action videos dataset using YOLOv11n for object detection. The pipeline loads video frames, resizes them to 640x640, runs batch inference with YOLO to detect objects, extracts individual object crops, and outputs object metadata and cropped images in Parquet format.
+- **Large-scale Image Embedding**: Processing 4TiB of base64-encoded images from a Parquet dataset using ViT for image embedding. The pipeline decodes base64 images, converts to RGB, preprocesses using ViTImageProcessor (resizing, normalization), runs GPU-accelerated batch inference with ViT to generate embeddings, and outputs results to Parquet format.
+
+Ray Data 2.50 is compared with Daft 0.6.2, an open source multimodal data processing library built on Ray.
+
+---
+
+## Results Summary
+
+![Multimodal Inference Benchmark Results](/data/images/multimodal_inference_results.png)
+
+```{list-table}
+:header-rows: 1
+:name: benchmark-results-summary
+-   - Workload
+    - **Daft (s)**
+    - **Ray Data (s)**
+-   - **Image Classification**
+    - 195.3 ± 2.5
+    - **111.2 ± 1.2**
+-   - **Document Embedding**
+    - 51.3 ± 1.3
+    - **29.4 ± 0.8**
+-   - **Audio Transcription**
+    - 510.5 ± 10.4
+    - **312.6 ± 3.1**
+-   - **Video Object Detection**
+    - 735.3 ± 7.6
+    - **623 ± 1.4**
+-   - **Large Scale Image Embedding**
+    - 752.75 ± 5.5
+    - **105.81 ± 0.79**
+```
+
+
+All benchmark results are taken from an average/std across 4 runs. A warmup was also run to download the model and remove any startup overheads that would affect the result.
+
+
+## Workload Configuration
+
+
+```{list-table}
+:header-rows: 1
+:name: workload-configuration
+-   - Workload
+    - Dataset
+    - Data Path
+    - Cluster Configuration
+    - Code
+-   - **Image Classification**
+    - 800k images from ImageNet
+    - s3://ray-example-data/imagenet/metadata_file.parquet
+    - 1 head / 8 workers of varying instance types
+    - [Link](https://github.com/ray-project/ray/tree/master/release/nightly_tests/multimodal_inference_benchmarks/image_classification)
+-   - **Document Embedding**
+    - 10k PDFs from Digital Corpora
+    - s3://ray-example-data/digitalcorpora/metadata
+    - g6.xlarge head, 8 g6.xlarge workers
+    - [Link](https://github.com/ray-project/ray/tree/master/release/nightly_tests/multimodal_inference_benchmarks/document_embedding)
+-   - **Audio Transcription**
+    - 113,800 audio files from Mozilla Common Voice 17 en dataset
+    - s3://air-example-data/common_voice_17/parquet/
+    - g6.xlarge head,  8 g6.xlarge workers
+    - [Link](https://github.com/ray-project/ray/tree/master/release/nightly_tests/multimodal_inference_benchmarks/audio_transcription)
+-   - **Video Object Detection**
+    - 1,000 videos from Hollywood-2 Human Actions dataset
+    - s3://ray-example-data/videos/Hollywood2-actions-videos/Hollywood2/AVIClips/
+    - 1 head, 8 workers of varying instance types
+    - [Link](https://github.com/ray-project/ray/tree/master/release/nightly_tests/multimodal_inference_benchmarks/video_object_detection)
+-   - **Large-scale Image Embedding**
+    - 4 TiB of Parquet files containing base64 encoded images
+    - s3://ray-example-data/image-datasets/10TiB-b64encoded-images-in-parquet-v3/
+    - m5.24xlarge (head), 40 g6e.xlarge (gpu workers), 64 r6i.8xlarge (cpu workers)
+    - [Link](https://github.com/ray-project/ray/tree/master/release/nightly_tests/multimodal_inference_benchmarks/large_image_embedding)
+```
+
+## Image Classification across different instance types
+
+This experiment compares the performance of Ray Data with Daft on the image classification workload across a variety of instance types. Each run is an average/std across 3 runs. A warmup was also run to download the model and remove any startup overheads that would affect the result.
+
+```{list-table}
+:header-rows: 1
+:name: image-classification-results
+-   -
+    - g6.xlarge (4 CPUs)
+    - g6.2xlarge (8 CPUs)
+    - g6.4xlarge (16 CPUs)
+    - g6.8xlarge (32 CPUs)
+-   - **Ray Data (s)**
+    - 456.2 ± 39.9
+    - **195.5 ± 7.6**
+    - **144.8 ± 1.9**
+    - **111.2 ± 1.2**
+-   - **Daft (s)**
+    - **315.0 ± 31.2**
+    - 202.0 ± 2.2
+    - 195.0 ± 6.6
+    - 195.3 ± 2.5
+```
+
+## Video Object Detection across different instance types
+
+This experiment compares the performance of Ray Data with Daft on the video object detection workload across a variety of instance types. Each run is an average/std across 4 runs. A warmup was also run to download the model and remove any startup overheads that would affect the result.
+
+
+
+```{list-table}
+:header-rows: 1
+:name: video-object-detection-results
+-   -
+    - g6.xlarge (4 CPUs)
+    - g6.2xlarge (8 CPUs)
+    - g6.4xlarge (16 CPUs)
+    - g6.8xlarge (32 CPUs)
+-   - **Ray Data (s)**
+    - 922 ± 13.8
+    - **704.8 ± 25.0**
+    - **629 ± 1.8**
+    - **623 ± 1.4**
+-   - **Daft (s)**
+    - **758.8 ± 10.4**
+    - 735.3 ± 7.6
+    - 747.5 ± 13.4
+    - 771.3 ± 25.6
+```

doc/source/data/comparisons.rst

@@ -0,0 +1,61 @@
+Comparing Ray Data to other systems
+===================================
+
+How does Ray Data compare to other solutions for offline inference?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. dropdown:: Batch Services: AWS Batch, GCP Batch
+
+    Cloud providers such as AWS, GCP, and Azure provide batch services to manage compute infrastructure for you. Each service uses the same process: you provide the code, and the service runs your code on each node in a cluster. However, while infrastructure management is necessary, it is often not enough. These services have limitations, such as a lack of software libraries to address optimized parallelization, efficient data transfer, and easy debugging. These solutions are suitable only for experienced users who can write their own optimized batch inference code.
+
+    Ray Data abstracts away not only the infrastructure management, but also the sharding of your dataset, the parallelization of the inference over these shards, and the transfer of data from storage to CPU to GPU.
+
+
+.. dropdown:: Online inference solutions: Bento ML, Sagemaker Batch Transform
+
+    Solutions like `Bento ML <https://www.bentoml.com/>`_, `Sagemaker Batch Transform <https://docs.aws.amazon.com/sagemaker/latest/dg/batch-transform.html>`_, or :ref:`Ray Serve <rayserve>` provide APIs to make it easy to write performant inference code and can abstract away infrastructure complexities. But they are designed for online inference rather than offline batch inference, which are two different problems with different sets of requirements. These solutions introduce additional complexity like HTTP, and cannot effectively handle large datasets leading inference service providers like `Bento ML to integrating with Apache Spark <https://www.youtube.com/watch?v=HcT0lZ4U1EM>`_ for offline inference.
+
+    Ray Data is built for offline batch jobs, without all the extra complexities of starting servers or sending HTTP requests.
+
+    For a more detailed performance comparison between Ray Data and Sagemaker Batch Transform, see `Offline Batch Inference: Comparing Ray, Apache Spark, and SageMaker <https://www.anyscale.com/blog/offline-batch-inference-comparing-ray-apache-spark-and-sagemaker>`_.
+
+.. dropdown:: Distributed Data Processing Frameworks: Apache Spark and Daft
+
+    Ray Data handles many of the same batch processing workloads as `Apache Spark <https://spark.apache.org/>`_ and `Daft <https://www.daft.ai>`_, but with a streaming paradigm that is better suited for GPU workloads for deep learning inference.
+
+    However, Ray Data doesn't have a SQL interface unlike Spark and Daft.
+
+    For a more detailed performance comparison between Ray Data and Apache Spark, see `Offline Batch Inference: Comparing Ray, Apache Spark, and SageMaker <https://www.anyscale.com/blog/offline-batch-inference-comparing-ray-apache-spark-and-sagemaker>`_.
+
+
+
+How does Ray Data compare to other solutions for ML training ingest?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. dropdown:: PyTorch Dataset and DataLoader
+
+    * **Framework-agnostic:** Datasets is framework-agnostic and portable between different distributed training frameworks, while `Torch datasets <https://pytorch.org/docs/stable/data.html>`__ are specific to Torch.
+    * **No built-in IO layer:** Torch datasets do not have an I/O layer for common file formats or in-memory exchange with other frameworks; users need to bring in other libraries and roll this integration themselves.
+    * **Generic distributed data processing:** Datasets is more general: it can handle generic distributed operations, including global per-epoch shuffling, which would otherwise have to be implemented by stitching together two separate systems. Torch datasets would require such stitching for anything more involved than batch-based preprocessing, and does not natively support shuffling across worker shards. See our `blog post <https://www.anyscale.com/blog/deep-dive-data-ingest-in-a-third-generation-ml-architecture>`__ on why this shared infrastructure is important for 3rd generation ML architectures.
+    * **Lower overhead:** Datasets is lower overhead: it supports zero-copy exchange between processes, in contrast to the multi-processing-based pipelines of Torch datasets.
+
+
+.. dropdown:: TensorFlow Dataset
+
+    * **Framework-agnostic:** Datasets is framework-agnostic and portable between different distributed training frameworks, while `TensorFlow datasets <https://www.tensorflow.org/api_docs/python/tf/data/Dataset>`__ is specific to TensorFlow.
+    * **Unified single-node and distributed:** Datasets unifies single and multi-node training under the same abstraction. TensorFlow datasets presents `separate concepts <https://www.tensorflow.org/api_docs/python/tf/distribute/DistributedDataset>`__ for distributed data loading and prevents code from being seamlessly scaled to larger clusters.
+    * **Generic distributed data processing:** Datasets is more general: it can handle generic distributed operations, including global per-epoch shuffling, which would otherwise have to be implemented by stitching together two separate systems. TensorFlow datasets would require such stitching for anything more involved than basic preprocessing, and does not natively support full-shuffling across worker shards; only file interleaving is supported. See our `blog post <https://www.anyscale.com/blog/deep-dive-data-ingest-in-a-third-generation-ml-architecture>`__ on why this shared infrastructure is important for 3rd generation ML architectures.
+    * **Lower overhead:** Datasets is lower overhead: it supports zero-copy exchange between processes, in contrast to the multi-processing-based pipelines of TensorFlow datasets.
+
+.. dropdown:: Petastorm
+
+    * **Supported data types:** `Petastorm <https://github.com/uber/petastorm>`__ only supports Parquet data, while Ray Data supports many file formats.
+    * **Lower overhead:** Datasets is lower overhead: it supports zero-copy exchange between processes, in contrast to the multi-processing-based pipelines used by Petastorm.
+    * **No data processing:** Petastorm does not expose any data processing APIs.
+
+
+.. dropdown:: NVTabular
+
+    * **Supported data types:** `NVTabular <https://github.com/NVIDIA-Merlin/NVTabular>`__ only supports tabular (Parquet, CSV, Avro) data, while Ray Data supports many other file formats.
+    * **Lower overhead:** Datasets is lower overhead: it supports zero-copy exchange between processes, in contrast to the multi-processing-based pipelines used by NVTabular.
+    * **Heterogeneous compute:** NVTabular doesn't support mixing heterogeneous resources in dataset transforms (e.g. both CPU and GPU transformations), while Ray Data supports this.

doc/source/data/concurrent-dataset-execution.md

@@ -0,0 +1,106 @@
+(data_concurrent_execution)=
+
+# Run multiple Datasets in one cluster
+
+When two or more Ray Data Datasets share a single Ray cluster, they compete for the same pool of nodes by default. That competition can cause unwanted contention — one Dataset's reads can starve a second Dataset's GPU stage, autoscaling decisions get muddled, and runtime becomes a function of whatever else happens to be running.
+
+Ray Data lets you assign each Dataset to its own **subcluster** — a labeled subset of nodes that only that Dataset uses. Subclusters give you smooth, predictable execution for concurrent Datasets and a natural way to express "this Dataset runs here, that one runs there."
+
+Common use cases:
+
+- **Asynchronous validation during training.** A training Dataset feeds the trainer. A validation Dataset feeds a separate validation task on different hardware. See {ref}`train-validating-checkpoints` for the Ray Train integration.
+- **Multitenancy on a shared workspace.** Several Datasets — different users, different pipelines, or different stages of one workflow — share one Anyscale workspace and don't disturb each other.
+
+## How it works
+
+Each Dataset carries an `ExecutionOptions.label_selector` (a `Dict[str, str]`) that Ray Data attaches to every task and actor the Dataset launches. The autoscaling coordinator buckets nodes by the value at the reserved label key `"ray-subcluster"` and only places a Dataset's work on nodes whose label matches.
+
+## Configuration
+
+There are two steps.
+
+### 1. Label your worker nodes
+
+Label each worker node with the reserved key `ray-subcluster` to mark which subcluster it belongs to. See {ref}`labels` for how to configure labels. The mechanism used depends on your deployment (cluster YAML, KubeRay, or `ray start --labels`).
+
+For example, in a Ray cluster YAML config:
+
+```yaml
+available_node_types:
+  train_workers:
+    min_workers: 2
+    max_workers: 4
+    labels:
+      ray-subcluster: training
+    node_config:
+      InstanceType: g5.xlarge
+  validation_workers:
+    min_workers: 0
+    max_workers: 2
+    labels:
+      ray-subcluster: validation
+    node_config:
+      InstanceType: g4dn.xlarge
+```
+
+Subcluster values are arbitrary strings (`"training"`, `"validation"`, `"tenant_a"`, `"team-blue"`) — pick whatever makes sense for your workload.
+
+### 2. Tag each Dataset with a `label_selector`
+
+Copy the current `DataContext`, set the selector on the copy, and apply the copy temporarily with the `DataContext.current()` context manager. Construct your Dataset inside the `with` block:
+
+```python
+import ray
+
+ctx = ray.data.DataContext.get_current().copy()
+ctx.execution_options.label_selector = {"ray-subcluster": "tenant_a"}
+
+with ray.data.DataContext.current(ctx):
+    # Tasks launched during construction (reads, schema inference) read
+    # the temporary context. ``Dataset.context`` is a deep copy of the
+    # current context, so the new Dataset keeps the selector after the
+    # ``with`` block exits.
+    dataset = ray.data.read_parquet("s3://my-bucket/tenant_a/")
+```
+
+:::{important}
+Mutating `ray.data.DataContext.get_current()` in place permanently affects every subsequent Dataset in the same driver process. Use the `DataContext.current()` context manager to scope each Dataset's selector to its own construction block.
+
+Set the selector *before* creating the Dataset, not after. Tasks Ray Data spawns during construction (for example, the parquet read tasks that infer the schema) read the current context, so setting `dataset.context.execution_options.label_selector` after the fact doesn't retroactively re-route them.
+:::
+
+## Example: two Datasets, two subclusters
+
+```python
+import ray
+import threading
+
+
+def make_dataset(subcluster: str, path: str) -> ray.data.Dataset:
+    ctx = ray.data.DataContext.get_current().copy()
+    ctx.execution_options.label_selector = {"ray-subcluster": subcluster}
+    with ray.data.DataContext.current(ctx):
+        return ray.data.read_parquet(path)
+
+
+# Construct each Dataset in the main thread so the temporary contexts
+# don't race on the process-global ``_default_context``.
+ds_a = make_dataset("tenant_a", "s3://my-bucket/tenant_a/")
+ds_b = make_dataset("tenant_b", "s3://my-bucket/tenant_b/")
+
+# Then run them concurrently. ds_a's tasks only land on
+# ray-subcluster=tenant_a nodes; ds_b's only on
+# ray-subcluster=tenant_b nodes.
+threading.Thread(target=lambda: ds_a.materialize()).start()
+threading.Thread(target=lambda: ds_b.materialize()).start()
+```
+
+## Ray Train integration
+
+When you wire the Datasets into a `TorchTrainer` (or any `DataParallelTrainer`), `ray.train.DataConfig` is the more ergonomic entry point — it takes a per-dataset `ExecutionOptions` map. See {ref}`train-validating-checkpoints` for the full pattern, including how to set the training-side selector through `DataConfig` and the validation-side selector inside your `validation_fn`.
+
+## API reference
+
+- {class}`ray.data.ExecutionOptions` — see the `label_selector` parameter.
+- {class}`ray.data.DataContext` — the per-process Ray Data configuration the `execution_options` live on.
+- {class}`ray.train.DataConfig` — accepts a `Dict[str, ExecutionOptions]` so each Train dataset can carry its own selector.

doc/source/data/contributing/contributing-guide.md

@@ -0,0 +1,62 @@
+# Contributing Guide
+
+If you want your changes to be reviewed and merged quickly, following a few key practices makes a big difference. Clear, focused, and well-structured contributions help reviewers understand your intent and ensure your improvements land smoothly.
+
+:::{seealso}
+This guide covers contributing to Ray Data in specific. For information on contributing to the Ray project in general, see the {ref}`general Ray contributing guide<getting-involved>`.
+:::
+
+## Find something to work on
+
+Start by solving a problem you encounter, like fixing a bug or adding a missing feature. If you're unsure where to start:
+* Browse the issue tracker for problems you understand.
+* Look for labels like ["good first issue"](https://github.com/ray-project/ray/issues?q=is%3Aissue%20state%3Aopen%20label%3Agood-first-issue%20label%3Adata) for approachable tasks.
+* [Join the Ray Slack](https://www.ray.io/join-slack) and post in #data-contributors.
+
+## Get early feedback
+
+If you’re adding a new public API or making a substantial refactor, **share your plan early**. Discussing changes before you invest a lot of work can save time and align your work with the project’s direction.
+
+You can open a draft PR, discuss on an Issue, or post in Slack for early feedback. It won’t affect acceptance and often improves the final design.
+
+## Write good tests
+
+Most changes to Ray Data require tests. For tips on how to write good tests, see {ref}`How to write tests <how-to-write-tests>`.
+
+## Write simple, clear code
+
+Ray Data values **readable, maintainable, and extendable** code over clever tricks. For guidance on how to write code that aligns with Ray Data's design taste, see [A Philosophy of Software Design](https://web.stanford.edu/~ouster/cgi-bin/aposd2ndEdExtract.pdf).
+
+## Test your changes locally
+
+To test your changes locally, build [Ray from source](https://docs.ray.io/en/latest/ray-contribute/development.html). For Ray Data development, you typically only need the Python environment—you can skip the C++ build unless you’re also contributing to Ray Core.
+
+Before submitting a PR, run `pre-commit` to lint your changes and `pytest` to execute your tests.
+
+Note that the full Ray Data test suite can be heavy to run locally, start with tests directly related to your changes. For example, if you modified `map`, from `python/ray/data/tests` run: `pytest test_map.py`.
+
+
+## Open a pull request
+
+### Write a clear pull request description
+
+Explain **why the change exists and what it achieves**. Clear descriptions reduce back-and-forth and speed up reviews.
+
+Here's an example of a PR with a good description: [[Data] Refactor PhysicalOperator.completed to fix side effects ](https://github.com/ray-project/ray/pull/58915).
+
+### Keep pull requests small
+
+Review difficulty scales non-linearly with PR size.
+
+For fast reviews, do the following:
+* **Keep PRs under ~200 lines** of change when possible.
+* **Split large PRs** into multiple incremental PRs.
+* Avoid mixing refactors and new features in the same PR.
+
+Here's an example of a PR that keeps its scope small: [[Data] Support Non-String Items for ApproximateTopK Aggregator](https://github.com/ray-project/ray/pull/58659). While the broader effort focuses on optimizing preprocessors, this change was deliberately split out as a small, incremental PR, which made it much easier to review.
+
+### Make CI pass
+
+Ray's CI runs lint and a small set of tests first in the `buildkite/microcheck` check. Start by making that pass.
+
+Once it’s green, tag your reviewer. They can add the go label to trigger the full test suite.

doc/source/data/contributing/contributing.rst

@@ -0,0 +1,9 @@
+========================
+Contributing to Ray Data
+========================
+
+.. toctree::
+    :maxdepth: 2
+
+    contributing-guide
+    how-to-write-tests

doc/source/data/contributing/how-to-write-tests.md

@@ -0,0 +1,176 @@
+(how-to-write-tests)=
+# How to write tests
+
+:::{note}
+**Disclaimer**: There are no hard rules in software engineering. Use your judgment when applying these.
+:::
+
+Flaky or brittle tests (the kind that break when assumptions shift) slow development. Nobody likes getting stuck on a PR because a test failed for reasons unrelated to their change.
+
+This guide is a collection of practices to help you write tests that support the Ray Data project, not slow it down.
+
+## General good practices
+
+### Prefer unit tests over integration tests
+
+Unit tests give faster feedback and make it easier to pinpoint failures. They run in milliseconds, not seconds, and don’t depend on Ray clusters, external systems, or timing. This keeps the test suite fast, reliable, and easy to maintain.
+
+:::{note}
+Put unit tests in `python/ray/data/tests/unit`.
+:::
+
+### Use fixtures, skip try-finally
+
+Fixtures make tests cleaner, more reusable, and better isolated. They’re the right tool for setup and teardown, especially for things like `monkeypatch`.
+
+`try-finally` works, but fixtures make intent clearer and avoid boilerplate.
+
+**Original code**
+```python
+def test_dynamic_block_split(ray_start_regular_shared):
+    ctx = ray.data.context.DataContext.get_current()
+    original_target_max_block_size = ctx.target_max_block_size
+
+    ctx.target_max_block_size = 1
+    try: 
+        ...
+    finally:
+        ctx.target_max_block_size = original_target_max_block_size
+```
+
+**Better**
+```python
+def test_dynamic_block_split(ray_start_regular_shared, restore_data_context):
+    ctx = ray.data.context.DataContext.get_current()
+    target_max_block_size = ctx.target_max_block_size
+    ... # No need for try-finally
+```
+
+## Ray-specific practices
+
+### Don't assume Datasets produce outputs in a specific order
+
+Unless you set `preserve_order=True` in the `DataContext`, Ray Data doesn’t guarantee an output order. If your test relies on order without explicitly asking for it, you’re setting yourself up for brittle failures.
+
+**Original code**
+```python
+ds_dfs = []
+for path in os.listdir(out_path):
+    assert path.startswith("data_") and path.endswith(".parquet")
+    ds_dfs.append(pd.read_parquet(os.path.join(out_path, path)))
+
+ds_df = pd.concat(ds_dfs).reset_index(drop=True)
+df = pd.concat([df1, df2]).reset_index(drop=True)
+assert ds_df.equals(df)
+```
+
+**Better**
+```python
+from ray.data._internal.util import rows_same
+
+actual_data = pd.read_parquet(out_path)
+expected_data = pd.concat([df1, df2]
+assert rows_same(actual_data, expected_data)
+```
+
+:::{tip}
+Use the `ray.data._internal.util.rows_same` utility function to compare pandas DataFrames for equality while ignoring indices and order.
+:::
+
+### Prefer shared cluster fixtures
+
+Prefer shared cluster fixtures like `ray_start_regular_shared` over isolated cluster fixtures like `shutdown_only` and `ray_start_regular`.
+
+`shutdown_only` and `ray_start_regular` restart the Ray cluster after each test finishes. Starting and stopping Ray can take over a second — which sounds small, but across thousands of tests (plus parameterizations) it adds up fast.
+
+Only use isolated clusters when your test truly needs a fresh cluster.
+
+:::{note}
+There's an inherent tradeoff between isolation and speed here. For this specific case, choose to prioritize speed.
+:::
+
+**Original code**
+```python
+@pytest.mark.parametrize("concurrency", [-1, 1.5], ids=["negative", "float"])
+def test_invalid_concurrency_raises(shutdown_only, concurrency):
+    ds = ray.data.range(1)  # Each parametrization restarts the Ray cluster!
+    with pytest.raises(ValueError):
+        ds.map(lambda row: row, concurrency=concurrency)
+```
+
+**Better**
+```python
+@pytest.mark.parametrize("concurrency", [-1, 1.5], ids=["negative", "float"])
+def test_invalid_concurrency_raises(ray_start_regular_shared, concurrency):
+    ds = ray.data.range(1)  # Each parametrization reuses the same Ray cluster.
+    with pytest.raises(ValueError):
+        ds.map(lambda row: row, concurrency=concurrency)
+```
+
+### Avoid testing against repr outputs to validate specific data
+
+`repr` output isn’t part of any interface contract — it can change at any time. Besides, tests that assert against repr often hide the real intent: are you trying to check the data, or just how it happens to print? Be explicit about what you care about.
+
+
+
+**Original code**
+```python
+assert str(ds) == "Dataset(num_rows=6, schema={one: int64, two: string})", ds
+```
+
+**Better**
+```python
+assert ds.schema() == Schema(pa.schema({"one": pa.int64(), "two": pa.string()}))
+assert ds.count() == 6
+```
+
+### Avoid assumptions about the number or size of blocks
+
+Unless you’re testing an API like `repartition`, don’t lock your test to a specific number or size of blocks. Both can change depending on the implementation or the cluster config — and that’s usually fine.
+
+**Original code**
+
+```python
+ds = ray.data.read_parquet(paths + [txt_path], filesystem=fs)
+assert ds._plan.initial_num_blocks() == 2  # Where does 2 come from?
+assert rows_same(ds.to_pandas(), expected_data)
+```
+
+**Better**
+
+```python
+ds = ray.data.read_parquet(paths + [txt_path], filesystem=fs)
+# Assertion about number of blocks has been removed.
+assert rows_same(ds.to_pandas(), expected_data)
+```
+
+**Original code**
+```python
+ds2 = ds.repartition(5)
+assert ds2._plan.initial_num_blocks() == 5
+assert ds2._block_num_rows() == [10, 10, 0, 0, 0]  # Magic numbers?
+```
+
+**Better**
+```python
+ds2 = ds.repartition(5)
+assert sum(len(bundle.blocks) for bundle in ds.iter_internal_ref_bundles()) == 5
+# Assertion about the number of rows in each block has been removed.
+```
+
+### Avoid testing that the DAG looks a particular way
+
+The operators in the execution plan can shift over time as the implementation evolves. Unless you’re specifically testing optimization rules or working at the operator level, tests shouldn’t expect a particular DAG structure.
+
+**Original code**
+```python
+# Check that metadata fetch is included in stats.
+assert "FromArrow" in ds.stats()
+# Underlying implementation uses `FromArrow` operator
+assert ds._plan._logical_plan.dag.name == "FromArrow"
+```
+
+**Better**
+```python
+# (Assertions removed). 
+```

doc/source/data/custom-datasource-example.rst

@@ -0,0 +1,132 @@
+.. _custom_datasource:
+
+Advanced: Read and Write Custom File Types
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. vale off
+
+.. Ignoring Vale because of future tense.
+
+This guide shows you how to extend Ray Data to read and write file types that aren't
+natively supported. This is an advanced guide, and you'll use unstable internal APIs.
+
+.. vale on
+
+Images are already supported with the :func:`~ray.data.read_images`
+and :meth:`~ray.data.Dataset.write_images` APIs, but this example shows you how to
+implement them for illustrative purposes.
+
+Read data from files
+--------------------
+
+.. tip::
+    If you're not contributing to Ray Data, you don't need to create a
+    :class:`~ray.data.Datasource`. Instead, you can call
+    :func:`~ray.data.read_binary_files` and decode files with
+    :meth:`~ray.data.Dataset.map`.
+
+The core abstraction for reading files is :class:`~ray.data.datasource.FileBasedDatasource`.
+It provides file-specific functionality on top of the
+:class:`~ray.data.Datasource` interface.
+
+To subclass :class:`~ray.data.datasource.FileBasedDatasource`, implement the constructor
+and ``_read_stream``.
+
+Implement the constructor
+=========================
+
+Call the superclass constructor and specify the files you want to read.
+Optionally, specify valid file extensions. Ray Data ignores files with other extensions.
+
+.. literalinclude:: doc_code/custom_datasource_example.py
+    :language: python
+    :start-after: __datasource_constructor_start__
+    :end-before: __datasource_constructor_end__
+
+Implement ``_read_stream``
+==========================
+
+``_read_stream`` is a generator that yields one or more blocks of data from a file.
+
+`Blocks <https://github.com/ray-project/ray/blob/23d3bfcb9dd97ea666b7b4b389f29b9cc0810121/python/ray/data/block.py#L54>`_
+are a Data-internal abstraction for a collection of rows. They can be PyArrow tables,
+pandas DataFrames, or dictionaries of NumPy arrays.
+
+Don't create a block directly. Instead, add rows of data to a
+`DelegatingBlockBuilder <https://github.com/ray-project/ray/blob/23d3bfcb9dd97ea666b7b4b389f29b9cc0810121/python/ray/data/_internal/delegating_block_builder.py#L10>`_.
+
+.. literalinclude:: doc_code/custom_datasource_example.py
+    :language: python
+    :start-after: __read_stream_start__
+    :end-before: __read_stream_end__
+
+Read your data
+==============
+
+Once you've implemented ``ImageDatasource``, call :func:`~ray.data.read_datasource` to
+read images into a :class:`~ray.data.Dataset`. Ray Data reads your files in parallel.
+
+.. literalinclude:: doc_code/custom_datasource_example.py
+    :language: python
+    :start-after: __read_datasource_start__
+    :end-before: __read_datasource_end__
+
+Write data to files
+-------------------
+
+.. note::
+    The write interface is under active development and might change in the future. If
+    you have feature requests,
+    `open a GitHub Issue <https://github.com/ray-project/ray/issues/new?assignees=&labels=enhancement%2Ctriage&projects=&template=feature-request.yml&title=%5B%3CRay+component%3A+Core%7CRLlib%7Cetc...%3E%5D+>`_.
+
+The core abstractions for writing data to files are :class:`~ray.data.datasource.RowBasedFileDatasink` and
+:class:`~ray.data.datasource.BlockBasedFileDatasink`. They provide file-specific functionality on top of the
+:class:`~ray.data.Datasink` interface.
+
+If you want to write one row per file, subclass :class:`~ray.data.datasource.RowBasedFileDatasink`.
+Otherwise, subclass :class:`~ray.data.datasource.BlockBasedFileDatasink`.
+
+.. vale off
+
+.. Ignoring Vale because of future tense.
+
+In this example, you'll write one image per file, so you'll subclass
+:class:`~ray.data.datasource.RowBasedFileDatasink`. To subclass
+:class:`~ray.data.datasource.RowBasedFileDatasink`, implement the constructor and
+:meth:`~ray.data.datasource.RowBasedFileDatasink.write_row_to_file`.
+
+.. vale on
+
+Implement the constructor
+=========================
+
+Call the superclass constructor and specify the folder to write to. Optionally, specify
+a string representing the file format (for example, ``"png"``). Ray Data uses the
+file format as the file extension.
+
+.. literalinclude:: doc_code/custom_datasource_example.py
+    :language: python
+    :start-after: __datasink_constructor_start__
+    :end-before: __datasink_constructor_end__
+
+Implement ``write_row_to_file``
+===============================
+
+``write_row_to_file`` writes a row of data to a file. Each row is a dictionary that maps
+column names to values.
+
+.. literalinclude:: doc_code/custom_datasource_example.py
+    :language: python
+    :start-after: __write_row_to_file_start__
+    :end-before: __write_row_to_file_end__
+
+Write your data
+===============
+
+Once you've implemented ``ImageDatasink``, call :meth:`~ray.data.Dataset.write_datasink`
+to write images to files. Ray Data writes to multiple files in parallel.
+
+.. literalinclude:: doc_code/custom_datasource_example.py
+    :language: python
+    :start-after: __write_datasink_start__
+    :end-before: __write_datasink_end__

doc/source/data/data-internals.rst

@@ -0,0 +1,356 @@
+.. _datasets_scheduling:
+
+==================
+Ray Data Internals
+==================
+
+This guide describes the implementation of Ray Data. The intended audience is advanced
+users and Ray Data developers.
+
+For a gentler introduction to Ray Data, see :ref:`Quickstart <data_quickstart>`.
+
+.. _dataset_concept:
+
+Key concepts
+============
+
+Datasets and blocks
+-------------------
+
+Datasets
+~~~~~~~~
+
+:class:`Dataset <ray.data.Dataset>` is the main user-facing Python API. It represents a
+distributed data collection, and defines data loading and processing operations. You
+typically use the API in this way:
+
+1. Create a Ray Dataset from external storage or in-memory data.
+2. Apply transformations to the data.
+3. Write the outputs to external storage or feed the outputs to training workers.
+
+Blocks
+~~~~~~
+
+A *block* is the basic unit of data bulk that Ray Data stores in the object store and
+transfers over the network. Each block contains a disjoint subset of rows, and Ray Data
+loads and transforms these blocks in parallel.
+
+The following figure visualizes a dataset with three blocks, each holding 1000 rows.
+Ray Data holds the :class:`~ray.data.Dataset` on the process that triggers execution
+(which is usually the driver) and stores the blocks as objects in Ray's shared-memory
+:ref:`object store <objects-in-ray>`.
+
+.. image:: images/dataset-arch.svg
+
+..
+  https://docs.google.com/drawings/d/1PmbDvHRfVthme9XD7EYM-LIHPXtHdOfjCbc1SCsM64k/edit
+
+Block formats
+~~~~~~~~~~~~~
+
+Blocks are Arrow tables or `pandas` DataFrames. Generally, blocks are Arrow tables
+unless Arrow can’t represent your data.
+
+The block format doesn’t affect the type of data returned by APIs like
+:meth:`~ray.data.Dataset.iter_batches`.
+
+Block size limiting
+~~~~~~~~~~~~~~~~~~~
+
+Ray Data bounds block sizes to avoid excessive communication overhead and prevent
+out-of-memory errors. Small blocks are good for latency and more streamed execution,
+while large blocks reduce scheduler and communication overhead. The default range
+attempts to make a good tradeoff for most jobs.
+
+Ray Data attempts to bound block sizes between 1 MiB and 128 MiB. To change the block
+size range, configure the ``target_min_block_size`` and  ``target_max_block_size``
+attributes of :class:`~ray.data.context.DataContext`.
+
+.. testcode::
+
+    import ray
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.target_min_block_size = 1 * 1024 * 1024
+    ctx.target_max_block_size = 128 * 1024 * 1024
+
+Dynamic block splitting
+~~~~~~~~~~~~~~~~~~~~~~~
+
+If a block is larger than 192 MiB (50% more than the target max size), Ray Data
+dynamically splits the block into smaller blocks.
+
+To change the size at which Ray Data splits blocks, configure
+``MAX_SAFE_BLOCK_SIZE_FACTOR``. The default value is 1.5.
+
+.. testcode::
+
+    import ray
+
+    ray.data.context.MAX_SAFE_BLOCK_SIZE_FACTOR = 1.5
+
+Ray Data can’t split rows. So, if your dataset contains large rows (for example, large
+images), then Ray Data can’t bound the block size.
+
+
+Shuffle Algorithms
+------------------
+
+In data processing, shuffling refers to the process of redistributing individual dataset's partitions (that in Ray Data are
+called :ref:`blocks <data_key_concepts>`).
+
+Ray Data implements two main shuffle algorithms:
+
+.. _hash-shuffle:
+
+Hash-shuffling
+~~~~~~~~~~~~~~
+
+.. note:: Hash-shuffling is available in Ray 2.46
+
+Hash-shuffling is a classical hash-partitioning based shuffling where:
+
+1. **Partition phase:** rows in every block are hash-partitioned based on values in the *key columns* into a specified number of partitions, following a simple residual formula of ``hash(key-values) % N`` (used in hash-tables and pretty much everywhere).
+2. **Push phase:** partition's shards from individual blocks are then pushed into corresponding aggregating actors (called ``HashShuffleAggregator``) handling respective partitions.
+3. **Reduce phase:** aggregators combine received individual partition's shards back into blocks optionally applying additional transformations before producing the resulting blocks.
+
+Hash-shuffling is particularly useful for operations that require deterministic partitioning based on keys, such as joins, group-by operations, and key-based repartitioning, by
+ensuring that rows with the same key-values are being placed into the same partition.
+
+.. note:: To use hash-shuffling in your aggregations and repartitioning operations, you need to currently specify
+    ``ray.data.DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE`` before creating a ``Dataset``.
+
+.. _range-partitioning-shuffle:
+
+Range-partitioning shuffle
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Range-partitioning based shuffle also is a classical algorithm, based on the dataset being split into target number of ranges as determined by boundaries approximating
+the real ranges of the totally ordered (sorted) dataset.
+
+1. **Sampling phase:** every input block is randomly sampled for (10) rows. Samples are combined into a single dataset, which is then sorted and split into
+   target number of partitions defining approximate *range boundaries*.
+2. **Partition phase:** every block is sorted and split into partitions based on the *range boundaries* derived in the previous step.
+3. **Reduce phase:** individual partitions within the same range are then recombined to produce the resulting block.
+
+.. note:: Range-partitioning shuffle is a default shuffling strategy. To set it explicitly specify
+    ``ray.data.DataContext.get_current().shuffle_strategy = ShuffleStrategy.SORT_SHUFFLE_PULL_BASED`` before creating a ``Dataset``.
+
+
+Operators, plans, and planning
+------------------------------
+
+Operators
+~~~~~~~~~
+
+There are two types of operators: *logical operators* and *physical operators*. Logical
+operators are stateless objects that describe “what” to do. Physical operators are
+stateful objects that describe “how” to do it. An example of a logical operator is
+``ReadOp``, and an example of a physical operator is ``TaskPoolMapOperator``.
+
+Plans
+~~~~~
+
+A *logical plan* is a series of logical operators, and a *physical plan* is a series of
+physical operators. When you call APIs like :func:`ray.data.read_images` and
+:meth:`ray.data.Dataset.map_batches`, Ray Data produces a logical plan. When execution
+starts, the planner generates a corresponding physical plan.
+
+The planner
+~~~~~~~~~~~
+
+The Ray Data planner translates logical operators to one or more physical operators. For
+example, the planner translates the ``ReadOp`` logical operator into two physical
+operators: an ``InputDataBuffer`` and ``TaskPoolMapOperator``. Whereas the ``ReadOp``
+logical operator only describes the input data, the ``TaskPoolMapOperator`` physical
+operator actually launches tasks to read the data.
+
+Plan optimization
+~~~~~~~~~~~~~~~~~
+
+Ray Data applies optimizations to both logical and physical plans. For example, the
+``OperatorFusionRule`` combines a chain of physical map operators into a single map
+operator. This prevents unnecessary serialization between map operators.
+
+To add custom optimization rules, implement a class that extends ``Rule`` and configure
+``DEFAULT_LOGICAL_RULES`` or ``DEFAULT_PHYSICAL_RULES``.
+
+.. testcode::
+
+    import ray
+    from ray.data._internal.logical.interfaces import Rule
+    from ray.data._internal.logical.optimizers import get_logical_ruleset
+
+    class CustomRule(Rule):
+        def apply(self, plan):
+            ...
+
+    logical_ruleset = get_logical_ruleset()
+    logical_ruleset.add(CustomRule)
+
+.. testcode::
+    :hide:
+
+    logical_ruleset.remove(CustomRule)
+
+Types of physical operators
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Physical operators take in a stream of block references and output another stream of
+block references. Some physical operators launch Ray Tasks and Actors to transform
+the blocks, and others only manipulate the references.
+
+``MapOperator`` is the most common operator. All read, transform, and write operations
+are implemented with it. To process data, ``MapOperator`` implementations use either Ray
+Tasks or Ray Actors.
+
+Non-map operators include ``OutputSplitter`` and ``LimitOperator``. These two operators
+manipulate references to data, but don’t launch tasks or modify the underlying data.
+
+Execution
+---------
+
+The executor
+~~~~~~~~~~~~
+
+The *executor* schedules tasks and moves data between physical operators.
+
+The executor and operators are located on the process where dataset execution starts.
+For batch inference jobs, this process is usually the driver. For training jobs, the
+executor runs on a special actor called ``SplitCoordinator`` which handles
+:meth:`~ray.data.Dataset.streaming_split`.
+
+Tasks and actors launched by operators are scheduled across the cluster, and outputs are
+stored in Ray’s distributed object store. The executor manipulates references to
+objects, and doesn’t fetch the underlying data itself to the executor.
+
+Out queues
+~~~~~~~~~~
+
+Each physical operator has an associated *out queue*. When a physical operator produces
+outputs, the executor moves the outputs to the operator’s out queue.
+
+.. _streaming_execution:
+
+Streaming execution
+~~~~~~~~~~~~~~~~~~~
+
+In contrast to bulk synchronous execution, Ray Data’s streaming execution doesn’t wait
+for one operator to complete to start the next. Each operator takes in and outputs a
+stream of blocks. This approach allows you to process datasets that are too large to fit
+in your cluster’s memory.
+
+The scheduling loop
+~~~~~~~~~~~~~~~~~~~
+
+The executor runs a loop. Each step works like this:
+
+1. Wait until running tasks and actors have new outputs.
+2. Move new outputs into the appropriate operator out queues.
+3. Choose some operators and assign new inputs to them. These operator process the new
+   inputs either by launching new tasks or manipulating metadata.
+
+Choosing the best operator to assign inputs is one of the most important decisions in
+Ray Data. This decision is critical to the performance, stability, and scalability of a
+Ray Data job. The executor can schedule an operator if the operator satisfies the
+following conditions:
+
+* The operator has inputs.
+* There are adequate resources available.
+* The operator isn’t backpressured.
+
+If there are multiple viable operators, the executor chooses the operator with the
+smallest out queue.
+
+Scheduling
+==========
+
+Ray Data uses Ray Core for execution. Below is a summary of the :ref:`scheduling strategy <ray-scheduling-strategies>` for Ray Data:
+
+* The ``SPREAD`` scheduling strategy ensures that data blocks and map tasks are evenly balanced across the cluster.
+* Dataset tasks ignore placement groups by default, see :ref:`Ray Data and Placement Groups <datasets_pg>`.
+* Map operations use the ``SPREAD`` scheduling strategy if the total argument size is less than 50 MB; otherwise, they use the ``DEFAULT`` scheduling strategy.
+* Read operations use the ``SPREAD`` scheduling strategy.
+* All other operations, such as split, sort, and shuffle, use the ``DEFAULT`` scheduling strategy.
+
+.. _datasets_pg:
+
+Ray Data and placement groups
+-----------------------------
+
+By default, Ray Data configures its tasks and actors to use the cluster-default scheduling strategy (``"DEFAULT"``). You can inspect this configuration variable here:
+:class:`ray.data.DataContext.get_current().scheduling_strategy <ray.data.DataContext>`. This scheduling strategy schedules these Tasks and Actors outside any present
+placement group. To use current placement group resources specifically for Ray Data, set ``ray.data.DataContext.get_current().scheduling_strategy = None``.
+
+Consider this override only for advanced use cases to improve performance predictability. The general recommendation is to let Ray Data run outside placement groups.
+
+.. _datasets_tune:
+
+Ray Data and Tune
+-----------------
+
+When using Ray Data in conjunction with :ref:`Ray Tune <tune-main>`, it's important to ensure there are enough free CPUs for Ray Data to run on. By default, Tune tries to fully utilize cluster CPUs. This can prevent Ray Data from scheduling tasks, reducing performance or causing workloads to hang.
+
+To ensure CPU resources are always available for Ray Data execution, limit the number of concurrent Tune trials with the ``max_concurrent_trials`` Tune option.
+
+.. literalinclude:: ./doc_code/key_concepts.py
+  :language: python
+  :start-after: __resource_allocation_1_begin__
+  :end-before: __resource_allocation_1_end__
+
+.. _data_memory_management:
+
+Memory Model
+============
+
+This section describes how Ray Data manages execution and object store memory.
+
+Ray divides each node's memory into three pools. By default, it reserves 30% for the
+object store and 10% for system overhead, and treats the remaining as logical memory.
+
+.. image:: ./data-memory-model-1.svg
+   :width: 300
+   :align: center
+
+Each pool serves a different purpose:
+
+- **Logical memory** is what's available for the heap of UDFs and built-in
+  transformations like reads.
+- **Object store** holds buffered blocks.
+- **System memory** is what's left for Ray Core (the raylet) and other processes outside
+  your tasks. 
+
+.. note::
+
+  Zero-copy deserializable objects are an exception. They're used in the UDF but
+  accounted for only in the object store, so they serve as both the buffer and the
+  working memory.
+
+.. image:: ./data-memory-model-2.svg
+   :width: 360
+   :align: center
+
+When a UDF processes data, it uses heap memory to do the work. For example, a UDF that
+calls a Torch preprocessor holds the tensors on the heap. As the UDF produces output 
+rows or batches, Ray Data serializes them into PyArrow tables and stores them in the 
+shared object store.
+
+.. image:: ./data-memory-model-3.svg
+   :width: 550
+   :align: center
+
+To limit object store use, Ray Data applies backpressure and stops launching tasks once
+enough data is buffered. If Ray Data produces more data than fits, Ray Core *spills* 
+those objects to disk. 
+
+.. note::
+
+    A common misconception is that heavy queuing causes OOMs. While it's true that heavy 
+    object store use contributes to worker OOMs by leaving less memory for the heaps of 
+    tasks and actors, heavy queuing doesn't cause OOMs directly because Ray spills objects
+    to disk. If Ray Data queues too much data, you see out-of-disk errors instead.
+
+To limit heap memory use, Ray Data relies on memory hints from you to estimate how much
+heap memory each UDF needs. It passes those hints to Ray Core so the scheduler doesn't
+oversubscribe the cluster. These hints don't enforce any OS-level limit. They only guide
+scheduling.

doc/source/data/data.rst

@@ -0,0 +1,184 @@
+.. _data:
+
+===================================================
+Ray Data: Scalable Data Processing for AI Workloads
+===================================================
+
+.. toctree::
+    :hidden:
+
+    quickstart
+    key-concepts
+    user-guide
+    examples
+    contributing/contributing
+    comparisons
+    benchmark
+    data-internals
+
+Ray Data is a scalable data processing library for AI workloads built on Ray.
+Ray Data provides flexible and performant APIs for common operations such as :ref:`batch inference <batch_inference_home>`, data preprocessing, and data loading for ML training. Unlike other distributed data systems, Ray Data features a :ref:`streaming execution engine <streaming-execution>` to efficiently process large datasets and maintain high utilization across both CPU and GPU workloads.
+
+Quick start
+-----------
+
+First, install Ray Data. To learn more about installing Ray and its libraries, see
+:ref:`Installing Ray <installation>`:
+
+.. code-block:: console
+
+    $ pip install -U 'ray[data]'
+
+Here is an example of how to do perform a simple batch text classification task with Ray Data:
+
+.. testcode::
+
+    import ray
+    import pandas as pd
+
+    class ClassificationModel:
+        def __init__(self):
+            from transformers import pipeline
+            self.pipe = pipeline("text-classification")
+
+        def __call__(self, batch: pd.DataFrame):
+            results = self.pipe(list(batch["text"]))
+            result_df = pd.DataFrame(results)
+            return pd.concat([batch, result_df], axis=1)
+
+    ds = ray.data.read_text("s3://anonymous@ray-example-data/sms_spam_collection_subset.txt")
+    ds = ds.map_batches(
+        ClassificationModel,
+        compute=ray.data.ActorPoolStrategy(size=2),
+        batch_size=64,
+        batch_format="pandas"
+        # num_gpus=1  # this will set 1 GPU per worker
+    )
+    ds.show(limit=1)
+
+.. testoutput::
+    :options: +MOCK
+
+    {'text': 'ham\tGo until jurong point, crazy.. Available only in bugis n great world la e buffet... Cine there got amore wat...', 'label': 'NEGATIVE', 'score': 0.9935141801834106}
+
+
+Why choose Ray Data?
+--------------------
+
+Modern AI workloads revolve around the usage of deep learning models, which are computationally intensive and often require specialized hardware such as GPUs.
+Unlike CPUs, GPUs often come with less memory, have different semantics for scheduling, and are much more expensive to run.
+Systems built to support traditional data processing pipelines often don't utilize such resources well.
+
+Ray Data supports AI workloads as a first-class citizen and offers several key advantages:
+
+- **Faster and cheaper for deep learning**: Ray Data streams data between CPU preprocessing and GPU inference/training tasks, maximizing resource utilization and reducing costs by keeping GPUs active.
+
+- **Framework friendly**: Ray Data provides performant, first-class integration with common AI frameworks (vLLM, PyTorch, HuggingFace, TensorFlow) and common cloud providers (AWS, GCP, Azure)
+
+- **Support for multi-modal data**: Ray Data leverages Apache Arrow and Pandas and provides support for many data formats used in ML workloads such as Parquet, Lance, images, JSON, CSV, audio, video, and more.
+
+- **Scalable by default**: Built on Ray for automatic scaling across heterogeneous clusters with different CPU and GPU machines. Code runs unchanged from one machine to hundreds of nodes processing hundreds of TB of data.
+
+..
+  https://docs.google.com/drawings/d/16AwJeBNR46_TsrkOmMbGaBK7u-OPsf_V8fHjU-d2PPQ/edit
+
+
+Learn more
+----------
+
+.. grid:: 1 2 2 2
+    :gutter: 1
+    :class-container: container pb-5
+
+    .. grid-item-card::
+
+        **Quickstart**
+        ^^^
+
+        Get started with Ray Data with a simple example.
+
+        +++
+        .. button-ref:: data_quickstart
+            :color: primary
+            :outline:
+            :expand:
+
+            Quickstart
+
+    .. grid-item-card::
+
+        **Key Concepts**
+        ^^^
+
+        Learn the key concepts behind Ray Data. Learn what
+        Datasets are and how they're used.
+
+        +++
+        .. button-ref:: data_key_concepts
+            :color: primary
+            :outline:
+            :expand:
+
+            Key Concepts
+
+    .. grid-item-card::
+
+        **User Guides**
+        ^^^
+
+        Learn how to use Ray Data, from basic usage to end-to-end guides.
+
+        +++
+        .. button-ref:: data_user_guide
+            :color: primary
+            :outline:
+            :expand:
+
+            Learn how to use Ray Data
+
+    .. grid-item-card::
+
+        **Examples**
+        ^^^
+
+        Find both simple and scaling-out examples of using Ray Data.
+
+        +++
+        .. button-ref:: examples
+            :color: primary
+            :outline:
+            :expand:
+
+            Ray Data Examples
+
+    .. grid-item-card::
+
+        **API**
+        ^^^
+
+        Get more in-depth information about the Ray Data API.
+
+        +++
+        .. button-ref:: data-api
+            :color: primary
+            :outline:
+            :expand:
+
+            Read the API Reference
+
+
+Case studies for Ray Data
+-------------------------
+
+**Training ingest using Ray Data**
+
+- `Pinterest uses Ray Data to do last mile data processing for model training <https://medium.com/pinterest-engineering/last-mile-data-processing-with-ray-629affbf34ff>`_
+- `DoorDash elevates model training with Ray Data <https://www.youtube.com/watch?v=pzemMnpctVY>`_
+- `Instacart builds distributed machine learning model training on Ray Data <https://tech.instacart.com/distributed-machine-learning-at-instacart-4b11d7569423>`_
+- `Predibase speeds up image augmentation for model training using Ray Data <https://predibase.com/blog/ludwig-v0-7-fine-tuning-pretrained-image-and-text-models-50x-faster-and>`_
+
+**Batch inference using Ray Data**
+
+- `ByteDance scales offline inference with multi-modal LLMs to 200 TB on Ray Data <https://www.anyscale.com/blog/how-bytedance-scales-offline-inference-with-multi-modal-llms-to-200TB-data>`_
+- `Spotify's new ML platform built on Ray Data for batch inference <https://engineering.atspotify.com/2023/02/unleashing-ml-innovation-at-spotify-with-ray/>`_
+- `Sewer AI speeds up object detection on videos 3x using Ray Data <https://www.anyscale.com/blog/inspecting-sewer-line-safety-using-thousands-of-hours-of-video>`_

doc/source/data/doc_code/custom_datasource_example.py

@@ -0,0 +1,79 @@
+# flake8: noqa
+# fmt: off
+
+from typing import Iterator, Union, List
+
+import pyarrow
+
+from ray.data.block import Block
+
+# __datasource_constructor_start__
+from ray.data.datasource import FileBasedDatasource
+
+class ImageDatasource(FileBasedDatasource):
+    def __init__(self, paths: Union[str, List[str]], *, mode: str):
+        super().__init__(
+            paths,
+            file_extensions=["png", "jpg", "jpeg", "bmp", "gif", "tiff"],
+        )
+
+        self.mode = mode  # Specify read options in the constructor
+# __datasource_constructor_end__
+
+# __read_stream_start__
+    def _read_stream(self, f: "pyarrow.NativeFile", path: str) -> Iterator[Block]:
+        import io
+        import numpy as np
+        from PIL import Image
+        from ray.data._internal.delegating_block_builder import DelegatingBlockBuilder
+
+        data = f.readall()
+        image = Image.open(io.BytesIO(data))
+        image = image.convert(self.mode)
+
+        # Each block contains one row
+        builder = DelegatingBlockBuilder()
+        array = np.asarray(image)
+        item = {"image": array}
+        builder.add(item)
+        yield builder.build()
+# __read_stream_end__
+
+# __read_datasource_start__
+import ray
+
+ds = ray.data.read_datasource(
+    ImageDatasource("s3://anonymous@ray-example-data/batoidea", mode="RGB")
+)
+# __read_datasource_end__
+
+
+from typing import Any, Dict
+import pyarrow
+
+# __datasink_constructor_start__
+from ray.data.datasource import RowBasedFileDatasink
+
+class ImageDatasink(RowBasedFileDatasink):
+    def __init__(self, path: str, column: str, file_format: str):
+        super().__init__(path, file_format=file_format)
+
+        self.column = column
+        self.file_format = file_format  # Specify write options in the constructor
+# __datasink_constructor_end__
+
+# __write_row_to_file_start__
+    def write_row_to_file(self, row: Dict[str, Any], file: pyarrow.NativeFile):
+        import io
+        from PIL import Image
+
+        # PIL can't write to a NativeFile, so we have to write to a buffer first.
+        image = Image.fromarray(row[self.column])
+        buffer = io.BytesIO()
+        image.save(buffer, format=self.file_format)
+        file.write(buffer.getvalue())
+# __write_row_to_file_end__
+
+# __write_datasink_start__
+ds.write_datasink(ImageDatasink("/tmp/results", column="image", file_format="png"))
+# __write_datasink_end__

doc/source/data/doc_code/key_concepts.py

@@ -0,0 +1,27 @@
+# flake8: noqa
+
+# fmt: off
+# __resource_allocation_1_begin__
+import ray
+from ray import tune
+
+# This workload will use spare cluster resources for execution.
+def objective(*args):
+    ray.data.range(10).show()
+
+# Create a cluster with 4 CPU slots available.
+ray.init(num_cpus=4)
+
+# By setting `max_concurrent_trials=3`, this ensures the cluster will always
+# have a sparse CPU for Dataset. Try setting `max_concurrent_trials=4` here,
+# and notice that the experiment will appear to hang.
+tuner = tune.Tuner(
+    tune.with_resources(objective, {"cpu": 1}),
+    tune_config=tune.TuneConfig(
+        num_samples=1,
+        max_concurrent_trials=3
+    )
+)
+tuner.fit()
+# __resource_allocation_1_end__
+# fmt: on

doc/source/data/doc_code/working-with-llms/basic_llm_example.py

@@ -0,0 +1,462 @@
+"""
+This file serves as a documentation example and CI test for basic LLM batch inference.
+
+"""
+
+# __basic_llm_example_start__
+import os
+import shutil
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# __basic_config_example_start__
+# Basic vLLM configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,  # Reduce if CUDA OOM occurs
+        "max_model_len": 4096,  # Constrain to fit test GPU memory
+    },
+    concurrency=1,
+    batch_size=64,
+)
+# __basic_config_example_end__
+
+processor = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        messages=[
+            {"role": "system", "content": "You are a bot that responds with haikus."},
+            {"role": "user", "content": row["item"]},
+        ],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=250,
+        ),
+    ),
+    postprocess=lambda row: dict(
+        answer=row["generated_text"],
+        **row,  # This will return all the original columns in the dataset.
+    ),
+)
+
+ds = ray.data.from_items(["Start of the haiku is: Complete this for me..."])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            ds = processor(ds)
+            ds.show(limit=1)
+        else:
+            print("Skipping basic LLM run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping basic LLM run due to environment error: {e}")
+
+# __hf_token_config_example_start__
+# Configuration with Hugging Face token
+config_with_token = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    runtime_env={"env_vars": {"HF_TOKEN": "your_huggingface_token"}},
+    concurrency=1,
+    batch_size=64,
+)
+# __hf_token_config_example_end__
+
+# __parallel_config_example_start__
+# Model parallelism configuration for larger models
+# tensor_parallel_size=2: Split model across 2 GPUs for tensor parallelism
+# pipeline_parallel_size=2: Use 2 pipeline stages (total 4 GPUs needed)
+# Total GPUs required = tensor_parallel_size * pipeline_parallel_size = 4
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_model_len": 16384,
+        "tensor_parallel_size": 2,
+        "pipeline_parallel_size": 2,
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 2048,
+    },
+    concurrency=1,
+    batch_size=32,
+    accelerator_type="L4",
+)
+# __parallel_config_example_end__
+
+# __runai_config_example_start__
+# RunAI streamer configuration for optimized model loading
+# Note: Install vLLM with runai dependencies: pip install -U "vllm[runai]>=0.10.1"
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "load_format": "runai_streamer",
+        "max_model_len": 16384,
+    },
+    concurrency=1,
+    batch_size=64,
+)
+# __runai_config_example_end__
+
+# __lora_config_example_start__
+# Multi-LoRA configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_lora": True,
+        "max_lora_rank": 32,
+        "max_loras": 1,
+        "max_model_len": 16384,
+    },
+    concurrency=1,
+    batch_size=32,
+)
+# __lora_config_example_end__
+
+# __s3_config_example_start__
+# S3 hosted model configuration
+s3_config = vLLMEngineProcessorConfig(
+    model_source="s3://your-bucket/your-model-path/",
+    engine_kwargs={
+        "load_format": "runai_streamer",
+        "max_model_len": 16384,
+    },
+    concurrency=1,
+    batch_size=64,
+)
+# __s3_config_example_end__
+
+base_dir = "/tmp/llm_checkpoint_demo"
+input_path = os.path.join(base_dir, "input")
+output_path = os.path.join(base_dir, "output")
+checkpoint_path = os.path.join(base_dir, "checkpoint")
+
+# Reset directories
+for path in (input_path, output_path, checkpoint_path):
+    shutil.rmtree(path, ignore_errors=True)
+    os.makedirs(path)
+
+# __row_level_fault_tolerance_config_example_start__
+# Row-level fault tolerance configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    concurrency=1,
+    batch_size=64,
+    should_continue_on_error=True,
+)
+# __row_level_fault_tolerance_config_example_end__
+
+# Seed the input directory because Ray Data V2's `read_parquet` errors on empty dirs
+ray.data.from_items(
+    [{"id": i, "message": f"Question {i}: What is 2 + 2?"} for i in range(4)]
+).write_parquet(input_path)
+
+# __checkpoint_config_setup_example_start__
+from ray.data.checkpoint import CheckpointConfig
+
+ctx = ray.data.DataContext.get_current()
+ctx.checkpoint_config = CheckpointConfig(
+    id_column="id",
+    checkpoint_path=checkpoint_path,
+    delete_checkpoint_on_success=False,
+)
+# __checkpoint_config_setup_example_end__
+
+# __checkpoint_usage_example_start__
+processor_config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 4096,
+    },
+    concurrency=1,
+    batch_size=16,
+)
+
+processor = build_processor(
+    processor_config,
+    preprocess=lambda row: dict(
+        id=row["id"], # Preserve the ID column for checkpointing
+        messages=[{"role": "user", "content": row["message"]}],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=10,
+        ),
+    ),
+    postprocess=lambda row: {
+        "id": row["id"], # Preserve the ID column for checkpointing
+        "answer": row.get("generated_text"),
+    },
+)
+
+ds = ray.data.read_parquet(input_path)
+ds = processor(ds)
+ds.write_parquet(output_path)
+# __checkpoint_usage_example_end__
+
+
+# __gpu_memory_config_example_start__
+# GPU memory management configuration
+# If you encounter CUDA out of memory errors, try these optimizations:
+config_memory_optimized = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_model_len": 8192,
+        "max_num_batched_tokens": 2048,
+        "enable_chunked_prefill": True,
+        "gpu_memory_utilization": 0.85,
+        "block_size": 16,
+    },
+    concurrency=1,
+    batch_size=16,
+)
+
+# For very large models or limited GPU memory:
+config_minimal_memory = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "max_model_len": 4096,
+        "max_num_batched_tokens": 1024,
+        "enable_chunked_prefill": True,
+        "gpu_memory_utilization": 0.75,
+    },
+    concurrency=1,
+    batch_size=8,
+)
+# __gpu_memory_config_example_end__
+
+# __embedding_config_example_start__
+# Embedding model configuration
+embedding_config = vLLMEngineProcessorConfig(
+    model_source="sentence-transformers/all-MiniLM-L6-v2",
+    task_type="embed",
+    engine_kwargs=dict(
+        enable_prefix_caching=False,
+        enable_chunked_prefill=False,
+        max_model_len=256,
+        enforce_eager=True,
+    ),
+    batch_size=32,
+    concurrency=1,
+    chat_template_stage=False,  # Skip chat templating for embeddings
+    detokenize_stage=False,     # Skip detokenization for embeddings
+)
+
+# Example usage for embeddings
+def create_embedding_processor():
+    return build_processor(
+        embedding_config,
+        preprocess=lambda row: dict(prompt=row["text"]),
+        postprocess=lambda row: {
+            "text": row["prompt"],
+            "embedding": row["embeddings"],
+        },
+    )
+
+
+# __embedding_config_example_end__
+
+# __classification_config_example_start__
+# Sequence classification model configuration
+# Use task_type="classify" for classification models (e.g., sentiment, quality scoring)
+# Use task_type="score" for cross-encoder scoring models
+classification_config = vLLMEngineProcessorConfig(
+    model_source="nvidia/nemocurator-fineweb-nemotron-4-edu-classifier",
+    task_type="classify",
+    engine_kwargs=dict(
+        max_model_len=512,
+        enforce_eager=True,
+    ),
+    batch_size=8,
+    concurrency=1,
+    chat_template_stage=False,
+    detokenize_stage=False,
+)
+
+
+# Example usage for classification
+def create_classification_processor():
+    return build_processor(
+        classification_config,
+        preprocess=lambda row: dict(prompt=row["text"]),
+        postprocess=lambda row: {
+            "text": row["prompt"],
+            # Classification models return logits in the 'embeddings' field
+            "score": float(row["embeddings"][0])
+            if row.get("embeddings") is not None and len(row["embeddings"]) > 0
+            else None,
+        },
+    )
+
+
+# __classification_config_example_end__
+
+# __shared_vllm_engine_config_example_start__
+import ray
+from ray import serve
+from ray.data.llm import ServeDeploymentProcessorConfig, build_processor
+from ray.serve.llm import (
+    LLMConfig,
+    ModelLoadingConfig,
+    build_llm_deployment,
+)
+from ray.serve.llm.openai_api_models import CompletionRequest
+
+llm_config = LLMConfig(
+    model_loading_config=ModelLoadingConfig(
+        model_id="facebook/opt-1.3b",
+        model_source="facebook/opt-1.3b",
+    ),
+    deployment_config=dict(
+        name="demo_deployment_config",
+        autoscaling_config=dict(
+            min_replicas=1,
+            max_replicas=1,
+        ),
+    ),
+    engine_kwargs=dict(
+        enable_prefix_caching=True,
+        enable_chunked_prefill=True,
+        max_num_batched_tokens=4096,
+    ),
+)
+
+APP_NAME = "demo_app"
+DEPLOYMENT_NAME = "demo_deployment"
+override_serve_options = dict(name=DEPLOYMENT_NAME)
+
+llm_app = build_llm_deployment(
+    llm_config, override_serve_options=override_serve_options
+)
+app = serve.run(llm_app, name=APP_NAME)
+config = ServeDeploymentProcessorConfig(
+    deployment_name=DEPLOYMENT_NAME,
+    app_name=APP_NAME,
+    dtype_mapping={
+        "CompletionRequest": CompletionRequest,
+    },
+    concurrency=1,
+    batch_size=64,
+)
+
+processor1 = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        method="completions",
+        dtype="CompletionRequest",
+        request_kwargs=dict(
+            model="facebook/opt-1.3b",
+            prompt=f"This is a prompt for {row['id']}",
+            stream=False,
+        ),
+    ),
+    postprocess=lambda row: dict(
+        prompt=row["choices"][0]["text"],
+    ),
+)
+
+processor2 = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        method="completions",
+        dtype="CompletionRequest",
+        request_kwargs=dict(
+            model="facebook/opt-1.3b",
+            prompt=row["prompt"],
+            stream=False,
+        ),
+    ),
+    postprocess=lambda row: row,
+)
+
+ds = ray.data.range(10)
+ds = processor2(processor1(ds))
+print(ds.take_all())
+# __shared_vllm_engine_config_example_end__
+
+# __cross_node_parallelism_config_example_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+        "pipeline_parallel_size": 4,
+        "tensor_parallel_size": 4,
+        "distributed_executor_backend": "ray",
+    },
+    batch_size=32,
+    concurrency=1,
+)
+# __cross_node_parallelism_config_example_end__
+
+# __custom_placement_group_strategy_config_example_start__
+# Simple: specify resources per worker, auto-replicated by TP*PP (4 workers here)
+# Alternative: use "bundles": [{"GPU": 1}] * 4 for explicit bundle control
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+        "pipeline_parallel_size": 2,
+        "tensor_parallel_size": 2,
+        "distributed_executor_backend": "ray",
+    },
+    batch_size=32,
+    concurrency=1,
+    placement_group_config={
+        "bundle_per_worker": {"GPU": 1},
+        "strategy": "STRICT_PACK",
+    },
+)
+# __custom_placement_group_strategy_config_example_end__
+
+# __concurrent_config_example_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+    },
+    concurrency=10,
+    batch_size=64,
+)
+# __concurrent_config_example_end__
+
+
+# __concurrent_config_fixed_pool_example_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+        "max_model_len": 16384,
+    },
+    concurrency=(10, 10),
+    batch_size=64,
+)
+# __concurrent_config_fixed_pool_example_end__
+
+# __concurrent_batches_tuning_example_start__
+# Tuning concurrent batch processing
+# Configure both parameters together for optimal throughput
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={
+        "enable_chunked_prefill": True,
+        "max_num_batched_tokens": 4096,
+    },
+    batch_size=64,
+    # Dataset-level concurrency (number of actor replicas)
+    concurrency=1,
+    # Number of batches that can run concurrently per actor (default: 8)
+    max_concurrent_batches=8,
+    # Number of tasks Ray Data queues per actor (default: 16)
+    # Increase to keep actor task queue saturated
+    experimental={"max_tasks_in_flight_per_actor": 16},
+)
+# __concurrent_batches_tuning_example_end__
+# __basic_llm_example_end__

doc/source/data/doc_code/working-with-llms/classification_example.py

@@ -0,0 +1,69 @@
+"""
+Classification batch inference with Ray Data LLM.
+
+Uses sequence classification models for content classifiers and sentiment analyzers.
+"""
+
+# Dependency setup
+import subprocess
+import sys
+
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "ray[llm]"])
+subprocess.check_call(
+    [sys.executable, "-m", "pip", "install", "--upgrade", "transformers"]
+)
+subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy==1.26.4"])
+
+
+# __classification_example_start__
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# Configure vLLM for a sequence classification model
+classification_config = vLLMEngineProcessorConfig(
+    model_source="nvidia/nemocurator-fineweb-nemotron-4-edu-classifier",
+    task_type="classify",  # Use 'classify' for sequence classification models
+    engine_kwargs=dict(
+        max_model_len=512,
+        enforce_eager=True,
+    ),
+    batch_size=8,
+    concurrency=1,
+    chat_template_stage=False,
+    detokenize_stage=False,
+)
+
+classification_processor = build_processor(
+    classification_config,
+    preprocess=lambda row: dict(prompt=row["text"]),
+    postprocess=lambda row: {
+        "text": row["prompt"],
+        # Classification models return logits in the 'embeddings' field
+        "edu_score": float(row["embeddings"][0])
+        if row.get("embeddings") is not None and len(row["embeddings"]) > 0
+        else None,
+    },
+)
+
+# Sample texts with varying educational quality
+texts = [
+    "lol that was so funny haha",
+    "Photosynthesis converts light energy into chemical energy.",
+    "Newton's laws describe the relationship between forces and motion.",
+]
+ds = ray.data.from_items([{"text": text} for text in texts])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            classified_ds = classification_processor(ds)
+            classified_ds.show(limit=3)
+        else:
+            print("Skipping classification run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping classification run due to environment error: {e}")
+# __classification_example_end__
+
+

doc/source/data/doc_code/working-with-llms/custom_tokenizer_example.py

@@ -0,0 +1,195 @@
+"""
+Documentation example and test for custom tokenizer batch inference.
+
+Demonstrates how to use vLLM's tokenizer infrastructure for models whose
+tokenizers are not natively supported by HuggingFace (e.g. Mistral Tekken,
+DeepSeek-V3.2, Grok-2 tiktoken).
+
+This example uses a standard model to demonstrate the pattern. For models
+that truly require vLLM's custom tokenizer (e.g. deepseek-ai/DeepSeek-V3-0324),
+replace the model ID and adjust tokenizer_mode accordingly.
+"""
+
+# __custom_chat_template_start__
+from typing import Any, Dict, List
+from vllm.tokenizers import get_tokenizer
+
+
+class VLLMChatTemplate:
+    """Apply a chat template using vLLM's tokenizer."""
+
+    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
+        self.tokenizer = get_tokenizer(
+            model_id,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=True,
+        )
+
+    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        prompts: List[str] = []
+        all_messages: List[List[Dict[str, Any]]] = []
+
+        for messages in batch["messages"]:
+            if hasattr(messages, "tolist"):
+                messages = messages.tolist()
+            all_messages.append(messages)
+
+            add_generation_prompt = messages[-1]["role"] == "user"
+            prompt = self.tokenizer.apply_chat_template(
+                messages,
+                tokenize=False,
+                add_generation_prompt=add_generation_prompt,
+                continue_final_message=not add_generation_prompt,
+            )
+            prompts.append(prompt)
+
+        return {
+            "prompt": prompts,
+            "messages": all_messages,
+            "sampling_params": batch["sampling_params"],
+        }
+
+
+# __custom_chat_template_end__
+
+
+# __custom_tokenize_start__
+class VLLMTokenize:
+    """Tokenize text prompts using vLLM's tokenizer."""
+
+    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
+        self.tokenizer = get_tokenizer(
+            model_id,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=True,
+        )
+
+    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        all_tokenized: List[List[int]] = [
+            self.tokenizer.encode(prompt) for prompt in batch["prompt"]
+        ]
+
+        return {
+            "tokenized_prompt": all_tokenized,
+            "messages": batch["messages"],
+            "sampling_params": batch["sampling_params"],
+        }
+
+
+# __custom_tokenize_end__
+
+
+# __custom_detokenize_start__
+class VLLMDetokenize:
+    """Detokenize generated token IDs using vLLM's tokenizer."""
+
+    def __init__(self, model_id: str, tokenizer_mode: str = "auto"):
+        self.tokenizer = get_tokenizer(
+            model_id,
+            tokenizer_mode=tokenizer_mode,
+            trust_remote_code=True,
+        )
+
+    async def __call__(self, batch: Dict[str, Any]) -> Dict[str, Any]:
+        decoded: List[str] = []
+        for tokens in batch["generated_tokens"]:
+            if hasattr(tokens, "tolist"):
+                tokens = tokens.tolist()
+            decoded.append(self.tokenizer.decode(tokens, skip_special_tokens=True))
+
+        return {
+            **batch,
+            "generated_text_custom": decoded,
+        }
+
+
+# __custom_detokenize_end__
+
+
+def run_custom_tokenizer_example():
+    import ray
+    from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+    # Input dataset with sampling_params per row.
+    ds = ray.data.from_items(
+        [
+            {
+                "messages": [
+                    {"role": "user", "content": "What is the capital of France?"}
+                ],
+                "sampling_params": {"max_tokens": 256, "temperature": 0.7},
+            },
+            {
+                "messages": [
+                    {"role": "user", "content": "Write a haiku about computing."}
+                ],
+                "sampling_params": {"max_tokens": 256, "temperature": 0.7},
+            },
+        ]
+    )
+
+    # __custom_tokenizer_pipeline_start__
+    MODEL_ID = "unsloth/Llama-3.1-8B-Instruct"
+
+    config = vLLMEngineProcessorConfig(
+        model_source=MODEL_ID,
+        engine_kwargs=dict(
+            max_model_len=4096,
+            trust_remote_code=True,
+            tokenizer_mode="auto",
+        ),
+        batch_size=4,
+        concurrency=1,
+        # Disable built-in stages -- we handle them via map_batches.
+        chat_template_stage=False,
+        tokenize_stage=False,
+        detokenize_stage=False,
+    )
+
+    processor = build_processor(
+        config,
+        postprocess=lambda row: {
+            "generated_text": row.get("generated_text", ""),
+            "generated_tokens": row.get("generated_tokens", []),
+            "num_input_tokens": row.get("num_input_tokens", 0),
+            "num_generated_tokens": row.get("num_generated_tokens", 0),
+        },
+    )
+
+    ds = ds.map_batches(
+        VLLMChatTemplate,
+        fn_constructor_kwargs={"model_id": MODEL_ID},
+        concurrency=1,
+        batch_size=4,
+    )
+
+    ds = ds.map_batches(
+        VLLMTokenize,
+        fn_constructor_kwargs={"model_id": MODEL_ID},
+        concurrency=1,
+        batch_size=4,
+    )
+
+    ds = processor(ds)
+
+    ds = ds.map_batches(
+        VLLMDetokenize,
+        fn_constructor_kwargs={"model_id": MODEL_ID},
+        concurrency=1,
+        batch_size=4,
+    )
+
+    # __custom_tokenizer_pipeline_end__
+    ds.show(limit=2)
+
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_custom_tokenizer_example()
+        else:
+            print("Skipping custom tokenizer example (no GPU available)")
+    except Exception as e:
+        print(f"Skipping custom tokenizer example: {e}")

doc/source/data/doc_code/working-with-llms/embedding_example.py

@@ -0,0 +1,63 @@
+"""
+Documentation example and test for embedding model batch inference.
+
+"""
+
+import subprocess
+import sys
+
+subprocess.check_call([sys.executable, "-m", "pip", "install", "--upgrade", "ray[llm]"])
+subprocess.check_call([sys.executable, "-m", "pip", "install", "numpy==1.26.4"])
+
+
+def run_embedding_example():
+    # __embedding_example_start__
+    import ray
+    from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+    embedding_config = vLLMEngineProcessorConfig(
+        model_source="sentence-transformers/all-MiniLM-L6-v2",
+        task_type="embed",
+        engine_kwargs=dict(
+            enable_prefix_caching=False,
+            enable_chunked_prefill=False,
+            max_model_len=256,
+            enforce_eager=True,
+        ),
+        batch_size=32,
+        concurrency=1,
+        chat_template_stage=False,  # Skip chat templating for embeddings
+        detokenize_stage=False,     # Skip detokenization for embeddings
+    )
+
+    embedding_processor = build_processor(
+        embedding_config,
+        preprocess=lambda row: dict(prompt=row["text"]),
+        postprocess=lambda row: {
+            "text": row["prompt"],
+            "embedding": row["embeddings"],
+        },
+    )
+
+    texts = [
+        "Hello world",
+        "This is a test sentence",
+        "Embedding models convert text to vectors",
+    ]
+    ds = ray.data.from_items([{"text": text} for text in texts])
+
+    embedded_ds = embedding_processor(ds)
+    embedded_ds.show(limit=1)
+    # __embedding_example_end__
+
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_embedding_example()
+        else:
+            print("Skipping embedding example (no GPU available)")
+    except Exception as e:
+        print(f"Skipping embedding example: {e}")

doc/source/data/doc_code/working-with-llms/minimal_quickstart.py

@@ -0,0 +1,62 @@
+"""
+Quickstart: vLLM + Ray Data batch inference.
+
+1. Installation
+2. Dataset creation
+3. Processor configuration
+4. Running inference
+5. Getting results
+"""
+
+# __minimal_vllm_quickstart_start__
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# Initialize Ray
+ray.init()
+
+# simple dataset
+ds = ray.data.from_items([
+    {"prompt": "What is machine learning?"},
+    {"prompt": "Explain neural networks in one sentence."},
+])
+
+# Minimal vLLM configuration
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    concurrency=1,  # 1 vLLM engine replica
+    batch_size=32,  # 32 samples per batch
+    engine_kwargs={
+        "max_model_len": 4096, # Fit into test GPU memory
+    }
+)
+
+# Build processor
+# preprocess: converts input row to format expected by vLLM (OpenAI chat format)
+# postprocess: extracts generated text from vLLM output
+processor = build_processor(
+    config,
+    preprocess=lambda row: {
+        "messages": [{"role": "user", "content": row["prompt"]}],
+        "sampling_params": {"temperature": 0.7, "max_tokens": 100},
+    },
+    postprocess=lambda row: {
+        "prompt": row["prompt"],
+        "response": row["generated_text"],
+    },
+)
+
+# inference
+ds = processor(ds)
+
+# iterate through the results
+for result in ds.iter_rows():
+    print(f"Q: {result['prompt']}")
+    print(f"A: {result['response']}\n")
+
+# Alternative ways to get results:
+# results = ds.take(10)  # Get first 10 results
+# ds.show(limit=5)       # Print first 5 results
+# ds.write_parquet("output.parquet")  # Save to file
+# __minimal_vllm_quickstart_end__
+

doc/source/data/doc_code/working-with-llms/omni_audio_example.py

@@ -0,0 +1,212 @@
+"""
+This file serves as a documentation example and CI test for VLM batch inference with audio.
+
+Structure:
+1. Infrastructure setup: Dataset compatibility patches, dependency handling
+2. Docs example (between __vlm_audio_example_start/end__): Embedded in Sphinx docs via literalinclude
+3. Test validation and cleanup
+"""
+
+
+'''
+# __audio_message_format_example_start__
+"""Supported audio input formats: audio URL, audio binary data"""
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": "Provide a detailed description of the audio."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe what happens in this audio."},
+                # Option 1: Provide audio URL
+                {"type": "audio_url", "audio_url": {"url": "https://example.com/audio.wav"}},
+                # Option 2: Provide audio binary data
+                {"type": "input_audio", "input_audio": {"data": audio_base64, "format": "wav"}},
+            ]
+        },
+    ]
+}
+# __audio_message_format_example_end__
+'''
+
+# __omni_audio_example_start__
+import ray
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+)
+
+# __omni_audio_config_example_start__
+audio_processor_config = vLLMEngineProcessorConfig(
+    model_source="Qwen/Qwen2.5-Omni-3B",
+    task_type="generate",
+    engine_kwargs=dict(
+        limit_mm_per_prompt={"audio": 1},
+    ),
+    batch_size=16,
+    accelerator_type="L4",
+    concurrency=1,
+    prepare_multimodal_stage={
+        "enabled": True,
+        "chat_template_content_format": "openai",
+    },
+    chat_template_stage=True,
+    tokenize_stage=True,
+    detokenize_stage=True,
+)
+# __omni_audio_config_example_end__
+
+
+# __omni_audio_preprocess_example_start__
+def audio_preprocess(row: dict) -> dict:
+    """
+    Preprocessing function for audio-language model inputs.
+
+    Converts dataset rows into the format expected by the Omni model:
+    - System prompt for analysis instructions
+    - User message with text and audio content
+    - Sampling parameters
+    """
+    return {
+        "messages": [
+            {
+                "role": "system",
+                "content": "You are a helpful assistant that analyzes audio. "
+                "Listen to the audio carefully and provide detailed descriptions.",
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": row["text"],
+                    },
+                    {
+                        "type": "input_audio",
+                        "input_audio": {
+                            "data": row["audio_data"],
+                            "format": "wav",
+                        },
+                    },
+                ],
+            },
+        ],
+        "sampling_params": {
+            "temperature": 0.3,
+            "max_tokens": 150,
+            "detokenize": False,
+        },
+    }
+
+
+def audio_postprocess(row: dict) -> dict:
+    return {
+        "resp": row["generated_text"],
+    }
+
+
+# __omni_audio_preprocess_example_end__
+
+
+def load_audio_dataset():
+# __omni_audio_load_dataset_example_start__
+    """
+    Load audio dataset from MRSAudio Hugging Face dataset.
+    """
+    try:
+        from datasets import load_dataset
+        from huggingface_hub import hf_hub_download
+        import base64
+
+        dataset_name = "MRSAudio/MRSAudio"
+
+        dataset = load_dataset(dataset_name, split="train")
+
+        audio_items = []
+
+        # Limit to first 10 samples for the example
+        num_samples = min(10, len(dataset))
+        for i in range(num_samples):
+            item = dataset[i]
+
+            audio_path = hf_hub_download(
+                repo_id=dataset_name, filename=item["path"], repo_type="dataset"
+            )
+
+            with open(audio_path, "rb") as f:
+                audio_bytes = f.read()
+
+            audio_base64 = base64.b64encode(audio_bytes).decode("utf-8")
+            audio_items.append(
+                {
+                    "audio_data": audio_base64,
+                    "text": item.get("text", "Describe this audio."),
+                }
+            )
+
+        audio_dataset = ray.data.from_items(audio_items)
+        return audio_dataset
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return None
+# __omni_audio_load_dataset_example_end__
+
+
+def create_omni_audio_config():
+    """Create Omni audio configuration."""
+    return vLLMEngineProcessorConfig(
+        model_source="Qwen/Qwen2.5-Omni-3B",
+        task_type="generate",
+        engine_kwargs=dict(
+            enforce_eager=True,
+            limit_mm_per_prompt={"audio": 1},
+        ),
+        batch_size=16,
+        accelerator_type="L4",
+        concurrency=1,
+        prepare_multimodal_stage={
+            "enabled": True,
+            "chat_template_content_format": "openai",
+        },
+        chat_template_stage=True,
+        tokenize_stage=True,
+        detokenize_stage=True,
+    )
+
+def run_omni_audio_example():
+# __omni_audio_run_example_start__
+    """Run the complete Omni audio example workflow."""
+    config = create_omni_audio_config()
+    audio_dataset = load_audio_dataset()
+
+    if audio_dataset:
+        # Build processor with preprocessing and postprocessing
+        processor = build_processor(
+            config, preprocess=audio_preprocess, postprocess=audio_postprocess
+        )
+
+        print("Omni audio processor configured successfully")
+        print(f"Model: {config.model_source}")
+        print(f"Has multimodal support: {config.prepare_multimodal_stage.get('enabled', False)}")
+        result = processor(audio_dataset).take_all()
+        return config, processor, result
+# __omni_audio_run_example_end__
+    return None, None, None
+
+
+# __omni_audio_example_end__
+
+if __name__ == "__main__":
+    # Run the example Omni audio workflow only if GPU is available
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_omni_audio_example()
+        else:
+            print("Skipping Omni audio example run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping Omni audio example run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/openai_api_example.py

@@ -0,0 +1,99 @@
+"""
+This file serves as a documentation example and CI test for OpenAI API batch inference.
+
+"""
+
+import os
+from ray.data.llm import HttpRequestProcessorConfig, build_processor
+
+
+def run_openai_example():
+    # __openai_example_start__
+    import ray
+
+    OPENAI_KEY = os.environ["OPENAI_API_KEY"]
+    ds = ray.data.from_items(["Hand me a haiku."])
+
+    config = HttpRequestProcessorConfig(
+        url="https://api.openai.com/v1/chat/completions",
+        headers={"Authorization": f"Bearer {OPENAI_KEY}"},
+        qps=1,
+    )
+
+    processor = build_processor(
+        config,
+        preprocess=lambda row: dict(
+            payload=dict(
+                model="gpt-4o-mini",
+                messages=[
+                    {
+                        "role": "system",
+                        "content": "You are a bot that responds with haikus.",
+                    },
+                    {"role": "user", "content": row["item"]},
+                ],
+                temperature=0.0,
+                max_tokens=150,
+            ),
+        ),
+        postprocess=lambda row: dict(
+            response=row["http_response"]["choices"][0]["message"]["content"]
+        ),
+    )
+
+    ds = processor(ds)
+    print(ds.take_all())
+    # __openai_example_end__
+
+
+def run_openai_demo():
+    """Run the OpenAI API configuration demo."""
+    print("OpenAI API Configuration Demo")
+    print("=" * 30)
+    print("\nExample configuration:")
+    print("config = HttpRequestProcessorConfig(")
+    print("    url='https://api.openai.com/v1/chat/completions',")
+    print("    headers={'Authorization': f'Bearer {OPENAI_KEY}'},")
+    print("    qps=1,")
+    print(")")
+    print("\nThe processor handles:")
+    print("- Preprocessing: Convert text to OpenAI API format")
+    print("- HTTP requests: Send batched requests to OpenAI")
+    print("- Postprocessing: Extract response content")
+
+
+def preprocess_for_openai(row):
+    """Preprocess function for OpenAI API requests."""
+    return dict(
+        payload=dict(
+            model="gpt-4o-mini",
+            messages=[
+                {"role": "system", "content": "You are a helpful assistant."},
+                {"role": "user", "content": row["item"]},
+            ],
+            temperature=0.0,
+            max_tokens=150,
+        )
+    )
+
+
+def postprocess_openai_response(row):
+    """Postprocess function for OpenAI API responses."""
+    return dict(response=row["http_response"]["choices"][0]["message"]["content"])
+
+
+if __name__ == "__main__":
+    # Run live call if API key is set; otherwise show demo with mock output
+    if "OPENAI_API_KEY" in os.environ:
+        run_openai_example()
+    else:
+        # Mock response without API key
+        print(
+            [
+                {
+                    "response": (
+                        "Autumn leaves whisper\nSoft code flows in quiet lines\nBugs fall one by one"
+                    )
+                }
+            ]
+        )

doc/source/data/doc_code/working-with-llms/tokenization_aggregation_example.py

@@ -0,0 +1,53 @@
+"""
+This file serves as a documentation example for aggregated tokenization.
+
+"""
+
+import ray
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+    TokenizerStageConfig,
+    DetokenizeStageConfig,
+)
+
+# __aggregated_tokenization_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={"max_model_len": 4096},
+    concurrency=1,
+    batch_size=64,
+    tokenize_stage=TokenizerStageConfig(enabled=False),
+    detokenize_stage=DetokenizeStageConfig(enabled=False),
+)
+
+processor = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        messages=[
+            {"role": "user", "content": row["item"]},
+        ],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=250,
+            # Let the vLLM engine handle detokenization
+            detokenize=True,
+        ),
+    ),
+    postprocess=lambda row: dict(resp=row["generated_text"]),
+)
+# __aggregated_tokenization_end__
+
+ds = ray.data.from_items(["Hello world!"])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            ds = processor(ds)
+            ds.show(limit=1)
+        else:
+            print("Skipping aggregated run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping aggregated run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/tokenization_disaggregation_example.py

@@ -0,0 +1,48 @@
+"""
+This file serves as a documentation example for disaggregated tokenization.
+
+"""
+
+import ray
+from ray.data.llm import vLLMEngineProcessorConfig, build_processor
+
+# __disaggregated_tokenization_start__
+config = vLLMEngineProcessorConfig(
+    model_source="unsloth/Llama-3.1-8B-Instruct",
+    engine_kwargs={"max_model_len": 4096},
+    concurrency=1,
+    batch_size=64,
+    tokenize_stage=True,
+    detokenize_stage=True,
+)
+
+processor = build_processor(
+    config,
+    preprocess=lambda row: dict(
+        messages=[
+            {"role": "user", "content": row["item"]},
+        ],
+        sampling_params=dict(
+            temperature=0.3,
+            max_tokens=250,
+            # Let the vLLMEngineProcessor's CPU detokenize stage handle detokenization
+            detokenize=False,
+        ),
+    ),
+    postprocess=lambda row: dict(resp=row["generated_text"]),
+)
+# __disaggregated_tokenization_end__
+
+ds = ray.data.from_items(["Hello world!"])
+
+if __name__ == "__main__":
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            ds = processor(ds)
+            ds.show(limit=1)
+        else:
+            print("Skipping disaggregated run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping disaggregated run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/vlm_image_example.py

@@ -0,0 +1,217 @@
+"""
+This file serves as a documentation example and CI test for VLM batch inference with images.
+
+Structure:
+1. Infrastructure setup: Dataset compatibility patches, dependency handling
+2. Docs example (between __vlm_image_example_start/end__): Embedded in Sphinx docs via literalinclude
+3. Test validation and cleanup
+"""
+
+
+'''
+# __image_message_format_example_start__
+"""Supported image input formats: image URL, PIL Image object"""
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": "Provide a detailed description of the image."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe what happens in this image."},
+                # Option 1: Provide image URL
+                {"type": "image_url", "image_url": {"url": "https://example.com/image.jpg"}},
+                # Option 2: Provide PIL Image object
+                {"type": "image_pil", "image_pil": PIL.Image.open("path/to/image.jpg")}
+            ]
+        },
+    ]
+}
+# __image_message_format_example_end__
+'''
+
+
+# __vlm_image_example_start__
+import ray
+from PIL import Image
+from io import BytesIO
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+)
+from huggingface_hub import HfFileSystem
+
+# Load "LMMs-Eval-Lite" dataset from Hugging Face using HfFileSystem
+path = "hf://datasets/lmms-lab/LMMs-Eval-Lite/coco2017_cap_val/"
+fs = HfFileSystem()
+vision_dataset = ray.data.read_parquet(path, filesystem=fs)
+
+# __vlm_config_example_start__
+vision_processor_config = vLLMEngineProcessorConfig(
+    model_source="Qwen/Qwen2.5-VL-3B-Instruct",
+    engine_kwargs=dict(
+        tensor_parallel_size=1,
+        pipeline_parallel_size=1,
+        max_model_len=4096,
+        trust_remote_code=True,
+        limit_mm_per_prompt={"image": 1},
+    ),
+    batch_size=16,
+    concurrency=1,
+    prepare_multimodal_stage=True,
+)
+# __vlm_config_example_end__
+
+
+# __vlm_preprocess_example_start__
+def vision_preprocess(row: dict) -> dict:
+    """
+    Preprocessing function for vision-language model inputs.
+
+    Converts dataset rows into the format expected by the VLM:
+    - System prompt for analysis instructions
+    - User message with text and image content
+    - Multiple choice formatting
+    - Sampling parameters
+    """
+    choice_indices = ["A", "B", "C", "D", "E", "F", "G", "H"]
+
+    return {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "Analyze the image and question carefully, using step-by-step reasoning. "
+                    "First, describe any image provided in detail. Then, present your reasoning. "
+                    "And finally your final answer in this format: Final Answer: <answer> "
+                    "where <answer> is: The single correct letter choice A, B, C, D, E, F, etc. when options are provided. "
+                    "Only include the letter. Your direct answer if no options are given, as a single phrase or number. "
+                    "IMPORTANT: Remember, to end your answer with Final Answer: <answer>."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {"type": "text", "text": row["question"] + "\n\n"},
+                    {
+                        "type": "image_pil",
+                        "image_pil": Image.open(BytesIO(row["image"]["bytes"])),
+                    },
+                    {
+                        "type": "text",
+                        "text": "\n\nChoices:\n"
+                        + "\n".join(
+                            [
+                                f"{choice_indices[i]}. {choice}"
+                                for i, choice in enumerate(row["answer"])
+                            ]
+                        ),
+                    },
+                ],
+            },
+        ],
+        "sampling_params": {
+            "temperature": 0.3,
+            "max_tokens": 150,
+            "detokenize": False,
+        },
+        # Include original data for reference
+        "original_data": {
+            "question": row["question"],
+            "answer_choices": row["answer"],
+            "image_size": row["image"].get("width", 0) if row["image"] else 0,
+        },
+    }
+
+
+def vision_postprocess(row: dict) -> dict:
+    return {
+        "resp": row["generated_text"],
+    }
+
+
+# __vlm_preprocess_example_end__
+
+
+def load_vision_dataset():
+# __vlm_image_load_dataset_example_start__
+    """
+    Load vision dataset from Hugging Face.
+
+    This function loads the LMMs-Eval-Lite dataset which contains:
+    - Images with associated questions
+    - Multiple choice answers
+    - Various visual reasoning tasks
+    """
+    try:
+        from huggingface_hub import HfFileSystem
+
+        # Load "LMMs-Eval-Lite" dataset from Hugging Face using HfFileSystem
+        path = "hf://datasets/lmms-lab/LMMs-Eval-Lite/coco2017_cap_val/"
+        fs = HfFileSystem()
+        vision_dataset = ray.data.read_parquet(path, filesystem=fs)
+
+        return vision_dataset
+    except ImportError:
+        print(
+            "huggingface_hub package not available. Install with: pip install huggingface_hub"
+        )
+        return None
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return None
+# __vlm_image_load_dataset_example_end__
+
+
+def create_vlm_config():
+    """Create VLM configuration."""
+    return vLLMEngineProcessorConfig(
+        model_source="Qwen/Qwen2.5-VL-3B-Instruct",
+        engine_kwargs=dict(
+            tensor_parallel_size=1,
+            pipeline_parallel_size=1,
+            max_model_len=4096,
+            trust_remote_code=True,
+            limit_mm_per_prompt={"image": 1},
+        ),
+        batch_size=1,
+        concurrency=1,
+        prepare_multimodal_stage=True,
+    )
+
+
+def run_vlm_example():
+# __vlm_run_example_start__
+    """Run the complete VLM example workflow."""
+    config = create_vlm_config()
+    vision_dataset = load_vision_dataset()
+
+    if vision_dataset:
+        # Build processor with preprocessing and postprocessing
+        processor = build_processor(
+            config, preprocess=vision_preprocess, postprocess=vision_postprocess
+        )
+
+        print("VLM processor configured successfully")
+        print(f"Model: {config.model_source}")
+        result = processor(vision_dataset).take_all()
+        return config, processor, result
+# __vlm_run_example_end__
+    return None, None, None
+
+
+# __vlm_image_example_end__
+
+if __name__ == "__main__":
+    # Run the example VLM workflow only if GPU is available
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_vlm_example()
+        else:
+            print("Skipping VLM example run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping VLM example run due to environment error: {e}")

doc/source/data/doc_code/working-with-llms/vlm_video_example.py

@@ -0,0 +1,241 @@
+"""
+This file serves as a documentation example and CI test for VLM batch inference with videos.
+
+Structure:
+1. Infrastructure setup: Dataset compatibility patches, dependency handling
+2. Docs example (between __vlm_video_example_start/end__): Embedded in Sphinx docs via literalinclude
+3. Test validation and cleanup
+"""
+
+import os
+
+'''
+# __video_message_format_example_start__
+"""Supported video input formats: video URL"""
+{
+    "messages": [
+        {
+            "role": "system",
+            "content": "Provide a detailed description of the video."
+        },
+        {
+            "role": "user",
+            "content": [
+                {"type": "text", "text": "Describe what happens in this video."},
+                # Provide video URL
+                {"type": "video_url", "video_url": {"url": "https://example.com/video.mp4"}},
+            ]
+        },
+    ]
+}
+# __video_message_format_example_end__
+'''
+
+# __vlm_video_example_start__
+import ray
+from ray.data.llm import (
+    vLLMEngineProcessorConfig,
+    build_processor,
+)
+
+
+# __vlm_video_config_example_start__
+video_processor_config = vLLMEngineProcessorConfig(
+    model_source="Qwen/Qwen3-VL-4B-Instruct",
+    engine_kwargs=dict(
+        tensor_parallel_size=4,
+        pipeline_parallel_size=1,
+        trust_remote_code=True,
+        limit_mm_per_prompt={"video": 1},
+        mm_processor_kwargs={
+            "size": {
+                "shortest_edge": 65536,
+                "longest_edge": 20 * 1088 * 1920,
+            },
+            "do_sample_frames": False,
+        },
+    ),
+    batch_size=1,
+    accelerator_type="L4",
+    concurrency=1,
+    prepare_multimodal_stage={
+        "enabled": True,
+        "model_config_kwargs": dict(
+            # See available model config kwargs at https://docs.vllm.ai/en/latest/api/vllm/config/#vllm.config.ModelConfig
+            allowed_local_media_path="/tmp",
+            media_io_kwargs={"video": {"num_frames": 20, "fps": 2}},
+        ),
+    },
+    chat_template_stage=True,
+    tokenize_stage=True,
+    detokenize_stage=True,
+)
+# __vlm_video_config_example_end__
+
+
+# __vlm_video_preprocess_example_start__
+def video_preprocess(row: dict) -> dict:
+    """
+    Preprocessing function for video-language model inputs.
+
+    Converts dataset rows into the format expected by the VLM:
+    - System prompt for analysis instructions
+    - User message with text and video content
+    - Sampling parameters
+    - Multimodal processor kwargs for video processing
+    """
+    return {
+        "messages": [
+            {
+                "role": "system",
+                "content": (
+                    "You are a helpful assistant that analyzes videos. "
+                    "Watch the video carefully and provide detailed descriptions."
+                ),
+            },
+            {
+                "role": "user",
+                "content": [
+                    {
+                        "type": "text",
+                        "text": row["text"],
+                    },
+                    {
+                        "type": "video_url",
+                        "video_url": {"url": row["video_url"]},
+                    },
+                ],
+            },
+        ],
+        "sampling_params": {
+            "temperature": 0.3,
+            "max_tokens": 150,
+            "detokenize": False,
+        },
+    }
+
+
+def video_postprocess(row: dict) -> dict:
+    return {
+        "resp": row["generated_text"],
+    }
+
+
+# __vlm_video_preprocess_example_end__
+
+
+def load_video_dataset():
+# __vlm_video_load_dataset_example_start__
+    """
+    Load video dataset from ShareGPTVideo Hugging Face dataset.
+    """
+    try:
+        from huggingface_hub import hf_hub_download
+        import tarfile
+        from pathlib import Path
+
+        dataset_name = "ShareGPTVideo/train_raw_video"
+
+        tar_path = hf_hub_download(
+            repo_id=dataset_name,
+            filename="activitynet/chunk_0.tar.gz",
+            repo_type="dataset",
+        )
+
+        extract_dir = "/tmp/sharegpt_videos"
+        os.makedirs(extract_dir, exist_ok=True)
+
+        if not any(Path(extract_dir).glob("*.mp4")):
+            with tarfile.open(tar_path, "r:gz") as tar:
+                tar.extractall(extract_dir)
+
+        video_files = list(Path(extract_dir).rglob("*.mp4"))
+
+        # Limit to first 10 videos for the example
+        video_files = video_files[:10]
+
+        video_dataset = ray.data.from_items(
+            [
+                {
+                    "video_path": str(video_file),
+                    "video_url": f"file://{video_file}",
+                    "text": "Describe what happens in this video.",
+                }
+                for video_file in video_files
+            ]
+        )
+
+        return video_dataset
+    except Exception as e:
+        print(f"Error loading dataset: {e}")
+        return None
+# __vlm_video_load_dataset_example_end__
+
+def create_vlm_video_config():
+    """Create VLM video configuration."""
+    return vLLMEngineProcessorConfig(
+        model_source="Qwen/Qwen3-VL-4B-Instruct",
+        engine_kwargs=dict(
+            tensor_parallel_size=4,
+            pipeline_parallel_size=1,
+            trust_remote_code=True,
+            limit_mm_per_prompt={"video": 1},
+            mm_processor_kwargs={
+                "size": {
+                    "shortest_edge": 65536,
+                    "longest_edge": 20 * 1088 * 1920,
+                },
+                "do_sample_frames": False,
+            },
+        ),
+        batch_size=1,
+        accelerator_type="L4",
+        concurrency=1,
+        prepare_multimodal_stage={
+            "enabled": True,
+            "model_config_kwargs": dict(
+                # See available model config kwargs at https://docs.vllm.ai/en/latest/api/vllm/config/#vllm.config.ModelConfig
+                allowed_local_media_path="/tmp",
+                media_io_kwargs={"video": {"num_frames": 20, "fps": 2}},
+            ),
+        },
+        chat_template_stage=True,
+        tokenize_stage=True,
+        detokenize_stage=True,
+    )
+
+
+def run_vlm_video_example():
+# __vlm_video_run_example_start__
+    """Run the complete VLM video example workflow."""
+    config = create_vlm_video_config()
+    video_dataset = load_video_dataset()
+
+    if video_dataset:
+        # Build processor with preprocessing and postprocessing
+        processor = build_processor(
+            config, preprocess=video_preprocess, postprocess=video_postprocess
+        )
+
+        print("VLM video processor configured successfully")
+        print(f"Model: {config.model_source}")
+        print(f"Has multimodal support: {config.prepare_multimodal_stage.get('enabled', False)}")
+        result = processor(video_dataset).take_all()
+        return config, processor, result
+# __vlm_video_run_example_end__
+    return None, None, None
+
+
+# __vlm_video_example_end__
+
+if __name__ == "__main__":
+    # Run the example VLM video workflow only if GPU is available
+    try:
+        import torch
+
+        if torch.cuda.is_available():
+            run_vlm_video_example()
+        else:
+            print("Skipping VLM video example run (no GPU available)")
+    except Exception as e:
+        print(f"Skipping VLM video example run due to environment error: {e}")

doc/source/data/examples.yml

@@ -0,0 +1,79 @@
+# This file is used to auto-generate the Examples Gallery page.
+# Do not edit the generated examples.rst page directly.
+# To request formatting changes to the generated page, file an issue with the Ray docs team.
+# To reference the generated page, use examples.html.
+# When adding a new example, include the skill level and framework, if applicable.
+
+text: Below are examples for using Ray Data for batch inference workloads or large-scale data processing with a variety of frameworks and use cases.
+columns_to_show:
+  - frameworks
+groupby: skill_level
+examples:
+  - title: Image Classification Batch Inference with PyTorch ResNet152
+    skill_level: beginner
+    frameworks:
+      - PyTorch
+    use_cases:
+      - computer vision
+    link: examples/pytorch_resnet_batch_prediction
+  - title: Object Detection Batch Inference with PyTorch FasterRCNN_ResNet50
+    skill_level: beginner
+    frameworks:
+      - PyTorch
+    use_cases:
+      - computer vision
+    link: examples/batch_inference_object_detection
+  - title: Image Classification Batch Inference with Hugging Face Vision Transformer
+    skill_level: beginner
+    frameworks:
+      - Transformers
+    use_cases:
+      - computer vision
+    link: examples/huggingface_vit_batch_prediction
+  - title: Tabular Data Training and Batch Inference with XGBoost
+    skill_level: beginner
+    frameworks:
+      - xgboost
+    link: ../_collections/ray-overview/examples/e2e-xgboost/README
+  - title: LLM Batch Inference
+    skill_level: beginner
+    frameworks:
+      - vLLM
+    use_cases:
+      - large language models
+      - generative ai
+    link: ../_collections/data/examples/llm_batch_inference_text/README
+  - title: Batch Inference with Structural Output
+    skill_level: beginner
+    frameworks:
+      - vLLM
+    use_cases:
+      - large language models
+      - generative ai
+    link: ../llm/examples/batch/vllm-with-structural-output
+  - title: Batch Inference with LoRA Adapter
+    skill_level: beginner
+    frameworks:
+      - vLLM
+    use_cases:
+      - large language models
+      - generative ai
+    link: ../llm/examples/batch/vllm-with-lora
+  - title: Multimodal LLM Batch Inference
+    skill_level: beginner
+    frameworks:
+      - vLLM
+    use_cases:
+      - large language models
+      - generative ai
+      - computer vision
+    link: ../_collections/data/examples/llm_batch_inference_vision/README
+  - title: Unstructured Data Ingestion and Processing
+    skill_level: intermediate
+    frameworks:
+      - Transformers
+      - Unstructured
+    use_cases:
+      - document processing
+      - data ingestion
+    link: ../_collections/data/examples/unstructured_data_ingestion/README

doc/source/data/examples/BUILD.bazel

@@ -0,0 +1,28 @@
+load("//bazel:python.bzl", "py_test_run_all_notebooks")
+
+filegroup(
+    name = "data_examples",
+    srcs = glob(["*.ipynb"]),
+    visibility = ["//doc:__subpackages__"]
+)
+
+# --------------------------------------------------------------------
+# Test all doc/source/data/examples notebooks.
+# --------------------------------------------------------------------
+
+py_test_run_all_notebooks(
+    size = "large",
+    include = ["*.ipynb"],
+        exclude = [],
+        data = ["//doc/source/data/examples:data_examples"],
+    tags = ["exclusive", "team:data", "gpu"],
+)
+
+filegroup(
+    name = "data_examples_ci_configs",
+    srcs = glob([
+        "**/ci/aws.yaml",
+        "**/ci/gce.yaml"
+    ]),
+    visibility = ["//doc:__pkg__"],
+)

doc/source/data/execution-configurations.rst

@@ -0,0 +1,94 @@
+.. _execution_configurations:
+
+========================
+Execution Configurations
+========================
+
+Ray Data provides a number of configuration options that control various aspects
+of execution of Ray Data's :class:`~ray.data.Dataset` on top of configuration of the Ray Core cluster itself.
+
+Ray Data's configuration is primarily controlled through either of :class:`~ray.data.ExecutionOptions`
+or :class:`~ray.data.DataContext`.
+
+This guide describes the most important of these configurations and when to use them.
+
+Configuring :class:`~ray.data.ExecutionOptions`
+===============================================
+
+The :class:`~ray.data.ExecutionOptions` class is used to configure options during Ray Dataset execution.
+To use it, modify the attributes in the current :class:`~ray.data.DataContext` object's `execution_options`. For example:
+
+.. testcode::
+   :hide:
+
+   import ray
+
+.. testcode::
+
+   ctx = ray.data.DataContext.get_current()
+   ctx.execution_options.verbose_progress = True
+
+* `resource_limits`: Set a soft limit on the resource usage during execution. For example, if there are other parts of the code which require some minimum amount of resources, you may want to limit the amount of resources that Ray Data uses. Auto-detected by default.
+* `exclude_resources`: Amount of resources to exclude from Ray Data. Set this if you have other workloads running on the same cluster. Note:
+
+  * If you're using Ray Data with Ray Train, training resources are automatically reserved and you don't need to set ``exclude_resources`` for them. Otherwise, off by default.
+  * For each resource type, you can't set both ``resource_limits`` and ``exclude_resources``.
+
+* `preserve_order`: Set this to preserve the ordering between blocks processed by operators under the streaming executor. Off by default.
+* `actor_locality_enabled`: Whether to enable locality-aware task dispatch to actors. This parameter applies to stateful :meth:`~ray.data.Dataset.map` operations. This setting is useful if you know you are consuming the output data directly on the consumer node (such as for ML batch inference). However, other use cases can incur a performance penalty with this setting. Off by default.
+* `verbose_progress`: Whether to report progress individually per operator. By default, only AllToAll operators and global progress is reported. This option is useful for performance debugging. On by default.
+
+For more details on each of the preceding options, see :class:`~ray.data.ExecutionOptions`.
+
+Configuring :class:`~ray.data.DataContext`
+==========================================
+
+The :class:`~ray.data.DataContext` class is used to configure more general options for Ray Data usage, such as observability/logging options,
+error handling/retry behavior, and internal data formats. To use it, modify the attributes in the current :class:`~ray.data.DataContext` object. For example:
+
+.. testcode::
+	   :hide:
+
+	   import ray
+
+.. testcode::
+
+   ctx = ray.data.DataContext.get_current()
+   ctx.verbose_stats_logs = True
+
+Many of the options in :class:`~ray.data.DataContext` are intended for advanced use cases or debugging,
+and most users shouldn't need to modify them. However, some of the most important options are:
+
+* `max_errored_blocks`: Max number of blocks that are allowed to have errors, unlimited if negative. This option allows application-level exceptions in block processing tasks. These exceptions may be caused by UDFs (for example, due to corrupted data samples) or IO errors. Data in the failed blocks are dropped. This option can be useful to prevent a long-running job from failing due to a small number of bad blocks. By default, no retries are allowed.
+* `write_file_retry_on_errors`: A list of sub-strings of error messages that should trigger a retry when writing files. This is useful for handling transient errors when writing to remote storage systems. By default, retries on common transient AWS S3 errors.
+* `verbose_stats_logs`: Whether stats logs should be verbose. This includes fields such as ``extra_metrics`` in the stats output, which are excluded by default. Off by default.
+* `log_internal_stack_trace_to_stdout`: Whether to include internal Ray Data/Ray Core code stack frames when logging to ``stdout``. The full stack trace is always written to the Ray Data log file. Off by default.
+* `raise_original_map_exception`: Whether to raise the original exception encountered in map UDF instead of wrapping it in a `UserCodeException`.
+
+For more details on each of the preceding options, see :class:`~ray.data.DataContext`.
+
+Job-level Checkpointing
+-----------------------
+
+Ray Data supports job-level checkpointing to improve fault tolerance for
+long-running batch pipelines. When enabled, Ray Data can resume a failed job
+by skipping rows that were successfully processed in a previous run, instead
+of restarting from the beginning.
+
+To configure job-level checkpointing, specify a
+:class:`~ray.data.checkpoint.CheckpointConfig` on the current
+:class:`~ray.data.DataContext`.
+
+**Example configuration:**
+
+.. code-block:: python
+
+    import ray
+    from ray.data.checkpoint import CheckpointConfig
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.checkpoint_config = CheckpointConfig(
+        id_column="id",
+        checkpoint_path="s3://my-bucket/ray-data-checkpoints",  # Must be accessible by all nodes
+        delete_checkpoint_on_success=False,  # Preserves checkpoints after successful runs
+    ) 

doc/source/data/how-to-avoid-ooms.md

@@ -0,0 +1,222 @@
+# How to avoid out-of-memory errors (OOMs)
+
+Out-of-memory errors (OOMs) are one of the most common issues Ray Data users encounter.
+
+This guide describes what OOMs look like and provides practical guidance for mitigating them.
+
+For a lower-level explanation of how Ray Data treats memory, read 
+{ref}`Ray Data Memory Model <data_memory_management>` and  
+{doc}`Ray Core Resource Isolation </ray-core/resource-isolation-with-cgroupv2>`
+
+
+## What OOMs look like
+
+OOMs show up in several ways. If you see one or more of these error messages, your job might be using too much memory.
+
+### Ray OOM kills
+
+When the Ray OOM killer proactively kills a task or actor, you might see an error like this:
+
+```
+  Task hungry_hippo failed due to oom. There are infinite oom retries remaining, so the task will be retried. Error: 2 worker(s) were killed due to the node running low on memory. Memory on the node (IP: <ip address>, ID: 92edc4e97e4dac3cee61126133ee7ab6d0a2ee73803623d24a02979d) was 110.69GB / 124.35GB (0.890161)
+  OOM kill reason: user cgroup memory upper bound was met or exceeded
+  Object store memory usage: [- objects spillable: 0
+  - bytes spillable: 0
+  - objects unsealed: 0
+  - bytes unsealed: 0
+  - objects in use: 0
+  - bytes in use: 0
+  - objects evictable: 0
+  - bytes evictable: 0
+  
+  - objects created by worker: 0
+  - bytes created by worker: 0
+  - objects restored: 0
+  - bytes restored: 0
+  - objects received: 0
+  - bytes received: 0
+  - objects errored: 0
+  - bytes errored: 0
+  
+  Eviction Stats:
+  (global lru) capacity: 35098657996
+  (global lru) used: 0%
+  (global lru) num objects: 0
+  (global lru) num evictions: 0
+  (global lru) bytes evicted: 0]
+  Ray killed 2 worker(s) based on the killing policy
+  Considered workers: [
+  Selected to kill: (Task: job ID=01000000, lease ID=0600000001000000ffffffffffffffffffffffffffffffffffffffffffffffff, task name=hungry_hippo, required resources={CPU: 1}, pid=3310152, actual memory used=0.67GB, worker ID=3e3d8f80b70d48b643d79ed2292b5d4f779820a964e55ad65413687d)
+  Selected to kill: (Task: job ID=01000000, lease ID=0400000001000000ffffffffffffffffffffffffffffffffffffffffffffffff, task name=hungry_hippo, required resources={CPU: 1}, pid=3310153, actual memory used=21.64GB, worker ID=0e5649d39c15609c0db6a5cf95de94befded2ee7da2facbf64b52e6f)
+  (Task: job ID=01000000, lease ID=0500000001000000ffffffffffffffffffffffffffffffffffffffffffffffff, task name=hungry_hippo, required resources={CPU: 1}, pid=3310151, actual memory used=21.95GB, worker ID=34241048bfb59ac29bd5e32d706c9bd41eafc6972c9bcbada99464e7)
+  (Task: job ID=01000000, lease ID=0000000001000000ffffffffffffffffffffffffffffffffffffffffffffffff, task name=hungry_hippo, required resources={CPU: 1}, pid=3310149, actual memory used=21.87GB, worker ID=2cc6dbeef4ebc06789de65fb43e04fbe1feebf1e699902ece89a8328)
+  (Task: job ID=01000000, lease ID=0200000001000000ffffffffffffffffffffffffffffffffffffffffffffffff, task name=hungry_hippo, required resources={CPU: 1}, pid=3310155, actual memory used=21.85GB, worker ID=14d4cc84e2f21ba3edbc9948b780d67013afbffda99dd33e829d56d3)
+  (Task: job ID=01000000, lease ID=0100000001000000ffffffffffffffffffffffffffffffffffffffffffffffff, task name=hungry_hippo, required resources={CPU: 1}, pid=3310147, actual memory used=21.53GB, worker ID=c90db9af23d78530d1f848c1301bc4b877925fe9122bc70948ad9489)]
+  Total non-selected idle workers: 25
+  Total non-selected idle workers USS bytes: 1.00GB
+  To see more information about memory usage on this node, use `ray logs raylet.out -ip <ip address>`
+  Top 10 memory users: PID  MEM(GB) COMMAND
+  3310151   21.95   ray::hungry_hippo
+  3310149   21.87   ray::hungry_hippo
+  3310155   21.85   ray::hungry_hippo
+  3310153   21.64   ray::hungry_hippo
+  3310147   21.53   ray::hungry_hippo
+  3108574   1.95    bazel
+  3180337   1.61    ray::foo_actor
+  3310152   0.67    ray::hungry_hippo
+  2924839   0.53    ray::idle_worker
+  3149737   0.47    ray::idle_worker
+  Refer to the documentation on how to address the out of memory issue: https://docs.ray.io/en/latest/ray-core/scheduling/ray-oom-prevention.html. Consider provisioning more memory on this node or reducing task parallelism by requesting more CPUs per task. To adjust the kill threshold, set the environment variable `RAY_memory_usage_threshold` when starting Ray. To disable worker killing, set the environment variable `RAY_memory_monitor_refresh_ms` to zero. Since 2.56, Ray updated the oom killing policy to enabling killing multiple workers and selecting workers based on the time since the task start executing. To revert to the legacy policy of determining worker to oom kill based on owner group size or only selecting a single worker to kill at a time, set the environment variable `RAY_worker_killing_policy_by_group` to true before starting Ray. If the idle workers have a non-trivial memory footprint at the time of OOM (check OOM log for non-selected idle workers), consider setting the environment variable `RAY_idle_worker_killing_memory_threshold_bytes` to a lower value to consider idle workers with lower memory footprint for killing.
+```
+
+You can see the number of Ray OOM kills in the "Ray OOM Kills (Tasks and Actors)" chart in the Ray Core dashboard:
+
+![Ray OOM kills chart](ray-oom-kills-chart.png)
+
+### Kernel OOM kills
+
+When the kernel OOM killer kills a Ray process before the Ray OOM killer, you might see an error like this:
+
+```
+(raylet) Task _map_task failed. There are infinite retries remaining, so the task will be retried. Error:                                                                                      
+(raylet) A worker died or was killed while executing a task by an unexpected system error. To troubleshoot the problem, check the logs for the dead worker. Lease ID: 2100000005000000ffffffffffffffffffffffffffffffffffffffffffffffff Worker ID: 863d8a6a594d60f8d143462b96cd3bf4270eafe617aaf2d2ca7266cb Node ID: 4cb25bc084aeb5a31ca5402ad589ae042a71165a8e2dd8418fecee26 Worker IP address: 10.0.50.112 Worker port: 10015 Worker PID: 2938 Worker exit type: SYSTEM_ERROR Worker exit detail: Worker unexpectedly exits with a connection error code 2. End of file. Some common causes include: (1) the process was killed by the OOM killer due to high memory usage, (2) ray stop --force was called, or (3) the worker crashed unexpectedly due to SIGSEGV or another unexpected error.
+```
+
+You can see the number of unexpected worker deaths in the "Unexpected System Level Worker Failures" chart in the Ray Core  dashboard. Kernel OOM kills often cause unexpected worker death.
+
+![Unexpected system-level worker deaths chart](unexpected-system-level-chart.png)
+
+### Node death
+
+If you're using older versions of Ray without resource isolation, nodes can die under memory pressure, and you might see an error like this:
+
+```
+{"asctime":"2026-03-23 18:24:28,943","levelname":"E","message":":info_message: Attempting to recover 41 lost objects by resubmitting their tasks or setting a new primary location from existing copies. To disable object reconstruction, set @ray.remote(max_retries=0).","filename":"core_worker.cc","lineno":475}
+```
+
+You can see node death in the "Node Count" chart in the Ray Core dashboard:
+
+![Unexpected system-level worker deaths chart](node-count.png)
+
+Nodes can die for reasons unrelated to memory pressure. If you see node death along with other memory-related errors, memory pressure might have caused the death.
+
+## Best practices
+
+### Use ``batch_size="auto"`` or small batch sizes
+
+Choose the smallest batch size that achieves good performance, or if your UDF doesn't use GPUs, use ``batch_size="auto"``.
+
+:::{versionadded} 2.56
+``batch_size="auto"``
+:::
+
+<!-- We're recommending 16 MiB because we found that it's the smallest batch size
+that doesn't degrade throughput on a variety of UDFs. See https://docs.google.com/document/d/1sw9CVm9cKp1b6voLc5gWIJLJM57NSGQjJ-HxvD_92jQ/edit?tab=t.0 -->
+
+If your UDF runs on CPU and isn't vectorized, use `map` instead. If it's vectorized, a good rule of thumb is a batch size of about 16 MiB. Unlike GPUs, CPUs have limited ability to parallelize work, so they don't benefit from large batches the way a GPU does.
+
+<!-- The rule of thumb to use 1/4 of GRAM comes from @stephanie-wang -->
+
+If your UDF runs on GPU, a good rule of thumb is to use about 1/4 of the GPU memory. Keep in mind that large GPU batches increase the risk of not only GPU OOMs, but also of regular heap OOMs because Ray Data builds the batch in heap memory first.
+
+### Configure ``memory`` for reads and high-memory UDFs
+
+If a task or actor uses more than a few GiB of memory, set ``memory``. This tells Ray Data how much memory each task or actor needs so it doesn't launch too many at once.
+
+To pick a value for ``memory``, read the Ray Data log file and look for the `max_uss_bytes` field. Ray typically writes the log file to `/tmp/ray/session-latest/ray-data/ray-data.log`.
+
+```
+ReadRange->MapBatches(uses_lots_of_memory): {'average_num_outputs_per_task': 1.0, ..., 'max_uss_bytes': {'num_samples': 20, 'mean': 4393336422.4, 'variance': 26855731156.89417, 'min': 4393119744, 'max': 4393529344, 'p50': 4393418752.0, 'p90': 4393500672.0, 'p95': 4393529344.0, 'p99': 4393529344.0}, ...}
+```
+
+Ray Data also emits the information to stdout:
+
+```
+Operator 'ReadRange->MapBatches(uses_lots_of_memory)' uses 4.1GiB of
+memory per task on average, but Ray only requests 0.0B per task at the
+start of the pipeline.
+
+To avoid out-of-memory errors, consider setting `memory=4.1GiB` in the
+appropriate function or method call. (This might be unnecessary if the
+number of concurrent tasks is low.)
+
+To change the frequency of this warning, set
+`DataContext.get_current().issue_detectors_config.high_memory_detector_config.detection_time_interval_s`,
+or disable the warning by setting value to -1. (current value: 30)
+```
+
+### Enable default map memory
+
+Unless you specify a value, Ray Data assumes a UDF needs 0 ``memory``. So even if you've set ``memory`` correctly for some APIs, Ray Data can still oversubscribe tasks and actors for the ones you haven't.
+
+To avoid this, set ``DataContext.get_current().default_map_logical_memory = True``.
+
+:::{versionadded} 2.56
+``DataContext.default_map_logical_memory``
+:::
+
+### Start Ray with resource isolation
+
+If you encounter kernel OOM kills or memory pressure related node deaths, enable *resource isolation* to provide enhanced protection for critical system components and eliminate kernel OOMs and node deaths.
+
+To enable *resource isolation*, follow the guide in {doc}`Ray Core Resource Isolation </ray-core/resource-isolation-with-cgroupv2>`.
+
+:::{versionadded} 2.56
+The full implementation of resource isolation.
+:::
+
+### Configure system memory to cover the raylet and anything outside the container
+
+By default, Ray reserves 10% of physical memory for system use. "System" covers Ray processes that aren't worker tasks or actors, the OS itself, and anything else on the node that isn't Ray, including processes outside the container.
+
+If you run large non-Ray processes like Vector or still experience kernel OOM even with *resource isolation* enabled, your "system" processes are likely using more memory than the default reserved memory for system processes.
+
+Ray logs something similar to the following example if it detects the "system" processes using more memory than the reserved amount.
+
+```
+System slice memory usage 10869600256 bytes has exceeded the reserved system memory of 10737418240 bytes. This can prevent Ray from being able to provide the proper protection to critical system processes and can lead to node deaths and significant loss of progress. Please consider passing a system reserved memory value that is higher than the current system slice memory usage via the --system-reserved-memory flag when starting the raylet.
+```
+
+In this case, allocate more memory by passing in a custom byte value to the ray start flag `--system-reserved-memory`. Try to allocate at least a GiB (depending on host size) of buffer space between the reported/expected system slice memory usage and the reserved system memory.
+
+The default is usually fine unless you're on tiny nodes, like an m5.xlarge.
+
+### Isolate reads for large files
+
+Ray Data uses PyArrow to implement APIs like `read_parquet`, and PyArrow can allocate lots of memory that isn't reclaimed when the read tasks finish. Because Ray reuses workers across operators, a downstream operator can schedule tasks onto a worker that's still holding that allocation. As a result, downstream operators can appear to be consuming far more memory than they actually are.
+
+If you encounter this, try ``DataContext.get_current().isolate_read_workers = True``. The flag prevents Ray Data from scheduling downstream operators on the same workers as reads. It can improve memory safety at the cost of some performance.
+
+:::{versionadded} 2.56
+``DataContext.isolate_read_workers`` was added in Ray 2.56.
+:::
+
+### Don't increase RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION
+
+Older versions of Ray Data emit a warning that suggests you increase `RAY_DEFAULT_OBJECT_STORE_MEMORY_PROPORTION`. While this can improve performance for some workloads like shuffle, it can also increase the risk of OOMs because it decreases the amount of memory available for your UDFs.
+
+To improve memory safety, don't configure the knob.
+
+## What to expect after tuning
+
+If you do all of the following:
+
+- Start Ray with resource isolation enabled.
+- Set system memory large enough to cover everything used outside of Ray worker tasks and actors
+- Set logical memory to physical memory minus system memory minus object store memory.
+- Set ``memory`` for each API to at least the heap memory that API needs to run.
+
+Then you shouldn't see OOMs or node deaths.
+
+The main limitation of these configurations is performance. When you set ``memory`` based on worst-case heap memory use, the system might launch fewer tasks or actors than it might be able to, and that can decrease throughput.
+
+If you want to experiment with oversubscription at the risk of potential OOMs, decrease `memory`.
+
+## Further reading
+
+For a deeper understanding of how Ray handles memory, read the following guides:
+
+- {ref}`Ray Data Memory Model <data_memory_management>`
+- {doc}`Ray Core Resource Isolation </ray-core/resource-isolation-with-cgroupv2>`
+- {ref}`Ray Core Out-Of-Memory Prevention <ray-oom-prevention>`
+- {doc}`Debugging Ray Core Memory Issues </ray-observability/user-guides/debug-apps/debug-memory>`

doc/source/data/inspecting-data.rst

@@ -0,0 +1,230 @@
+.. _inspecting-data:
+
+===============
+Inspecting Data
+===============
+
+Inspect :class:`Datasets <ray.data.Dataset>` to better understand your data.
+
+This guide shows you how to:
+
+* `Describe datasets <#describing-datasets>`_
+* `Inspect rows <#inspecting-rows>`_
+* `Inspect batches <#inspecting-batches>`_
+* `Inspect execution statistics <#inspecting-execution-statistics>`_
+
+.. _describing-datasets:
+
+Describing datasets
+===================
+
+:class:`Datasets <ray.data.Dataset>` are tabular. To view a dataset's column names and
+types, call :meth:`Dataset.schema() <ray.data.Dataset.schema>`.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+    print(ds.schema())
+
+.. testoutput::
+
+    Column             Type
+    ------             ----
+    sepal length (cm)  double
+    sepal width (cm)   double
+    petal length (cm)  double
+    petal width (cm)   double
+    target             int64
+
+For more information like the number of rows, print the Dataset.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+    print(ds)
+
+.. testoutput::
+
+    Dataset(num_rows=..., schema=...)
+
+.. _inspecting-rows:
+
+Inspecting rows
+===============
+
+To get a list of rows, call :meth:`Dataset.take() <ray.data.Dataset.take>` or
+:meth:`Dataset.take_all() <ray.data.Dataset.take_all>`. Ray Data represents each row as
+a dictionary.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+    rows = ds.take(1)
+    print(rows)
+
+.. testoutput::
+
+    [{'sepal length (cm)': 5.1, 'sepal width (cm)': 3.5, 'petal length (cm)': 1.4, 'petal width (cm)': 0.2, 'target': 0}]
+
+
+For more information on working with rows, see
+:ref:`Transforming rows <transforming_rows>` and
+:ref:`Iterating over rows <iterating-over-rows>`.
+
+.. _inspecting-batches:
+
+Inspecting batches
+==================
+
+A batch contains data from multiple rows. To inspect batches, call
+`Dataset.take_batch() <ray.data.Dataset.take_batch>`.
+
+By default, Ray Data represents batches as dicts of NumPy ndarrays. To change the type
+of the returned batch, set ``batch_format``. The batch format is independent from how
+Ray Data stores the underlying blocks, so you can use any batch format regardless of
+the internal block representation.
+
+.. tab-set::
+
+    .. tab-item:: NumPy
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+
+            batch = ds.take_batch(batch_size=2, batch_format="numpy")
+            print("Batch:", batch)
+            print("Image shape", batch["image"].shape)
+
+        .. testoutput::
+            :options: +MOCK
+
+            Batch: {'image': array([[[[...]]]], dtype=uint8)}
+            Image shape: (2, 32, 32, 3)
+
+    .. tab-item:: pandas
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+            batch = ds.take_batch(batch_size=2, batch_format="pandas")
+            print(batch)
+
+        .. testoutput::
+            :options: +MOCK
+
+               sepal length (cm)  sepal width (cm)  ...  petal width (cm)  target
+            0                5.1               3.5  ...               0.2       0
+            1                4.9               3.0  ...               0.2       0
+    .. tab-item:: pyarrow
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+            batch = ds.take_batch(batch_size=2, batch_format="pyarrow")
+            print(batch)
+
+        .. testoutput::
+
+            pyarrow.Table
+            sepal length (cm): double
+            sepal width (cm): double
+            petal length (cm): double
+            petal width (cm): double
+            target: int64
+            ----
+            sepal length (cm): [[5.1,4.9]]
+            sepal width (cm): [[3.5,3]]
+            petal length (cm): [[1.4,1.4]]
+            petal width (cm): [[0.2,0.2]]
+            target: [[0,0]]
+
+For more information on working with batches, see
+:ref:`Transforming batches <transforming_batches>` and
+:ref:`Iterating over batches <iterating-over-batches>`.
+
+
+Inspecting execution statistics
+===============================
+
+Ray Data calculates statistics during execution for each operator, such as wall clock time and memory usage.
+
+To view stats about your :class:`Datasets <ray.data.Dataset>`, call :meth:`Dataset.stats() <ray.data.Dataset.stats>` on an executed dataset. The stats are also persisted under `/tmp/ray/session_*/logs/ray-data/ray-data.log`.
+For more on how to read this output, see :ref:`Monitoring Your Workload with the Ray Data Dashboard <monitoring-your-workload>`.
+
+.. This snippet below is skipped because of  https://github.com/ray-project/ray/issues/54101.
+
+.. testcode::
+    :skipif: True
+
+    import ray
+    from huggingface_hub import HfFileSystem
+
+    def f(batch):
+        return batch
+
+    def g(row):
+        return True
+
+    path = "hf://datasets/ylecun/mnist/mnist/"
+
+    fs = HfFileSystem()
+    train_files = [f["name"] for f in fs.ls(path) if "train" in f["name"] and f["name"].endswith(".parquet")]
+    ds = (
+        ray.data.read_parquet(train_files, filesystem=fs)
+        .map_batches(f)
+        .filter(g)
+        .materialize()
+    )
+
+    print(ds.stats())
+
+.. testoutput::
+    :options: +MOCK
+
+    Operator 1 ReadParquet->SplitBlocks(32): 1 tasks executed, 32 blocks produced in 2.92s
+    * Remote wall time: 103.38us min, 1.34s max, 42.14ms mean, 1.35s total
+    * Remote cpu time: 102.0us min, 164.66ms max, 5.37ms mean, 171.72ms total
+    * UDF time: 0us min, 0us max, 0.0us mean, 0us total
+    * Peak heap memory usage (MiB): 266375.0 min, 281875.0 max, 274491 mean
+    * Output num rows per block: 1875 min, 1875 max, 1875 mean, 60000 total
+    * Output size bytes per block: 537986 min, 555360 max, 545963 mean, 17470820 total
+    * Output rows per task: 60000 min, 60000 max, 60000 mean, 1 tasks used
+    * Tasks per node: 1 min, 1 max, 1 mean; 1 nodes used
+    * Operator throughput:
+        * Ray Data throughput: 20579.80984833993 rows/s
+        * Estimated single node throughput: 44492.67361278733 rows/s
+
+    Operator 2 MapBatches(f)->Filter(g): 32 tasks executed, 32 blocks produced in 3.63s
+    * Remote wall time: 675.48ms min, 1.0s max, 797.07ms mean, 25.51s total
+    * Remote cpu time: 673.41ms min, 897.32ms max, 768.09ms mean, 24.58s total
+    * UDF time: 661.65ms min, 978.04ms max, 778.13ms mean, 24.9s total
+    * Peak heap memory usage (MiB): 152281.25 min, 286796.88 max, 164231 mean
+    * Output num rows per block: 1875 min, 1875 max, 1875 mean, 60000 total
+    * Output size bytes per block: 530251 min, 547625 max, 538228 mean, 17223300 total
+    * Output rows per task: 1875 min, 1875 max, 1875 mean, 32 tasks used
+    * Tasks per node: 32 min, 32 max, 32 mean; 1 nodes used
+    * Operator throughput:
+        * Ray Data throughput: 16512.364546087643 rows/s
+        * Estimated single node throughput: 2352.3683708977856 rows/s
+
+    Dataset throughput:
+        * Ray Data throughput: 11463.372316361854 rows/s
+        * Estimated single node throughput: 25580.963670075285 rows/s

doc/source/data/iterating-over-data.rst

@@ -0,0 +1,293 @@
+.. _iterating-over-data:
+
+===================
+Iterating over Data
+===================
+
+Ray Data lets you iterate over rows or batches of data.
+
+This guide shows you how to:
+
+* `Iterate over rows <#iterating-over-rows>`_
+* `Iterate over batches <#iterating-over-batches>`_
+* `Iterate over batches with shuffling <#iterating-over-batches-with-shuffling>`_
+* `Split datasets for distributed parallel training <#splitting-datasets-for-distributed-parallel-training>`_
+
+.. _iterating-over-rows:
+
+Iterating over rows
+===================
+
+To iterate over the rows of your dataset, call
+:meth:`Dataset.iter_rows() <ray.data.Dataset.iter_rows>`. Ray Data represents each row
+as a dictionary.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+    for row in ds.iter_rows():
+        print(row)
+
+.. testoutput::
+
+    {'sepal length (cm)': 5.1, 'sepal width (cm)': 3.5, 'petal length (cm)': 1.4, 'petal width (cm)': 0.2, 'target': 0}
+    {'sepal length (cm)': 4.9, 'sepal width (cm)': 3.0, 'petal length (cm)': 1.4, 'petal width (cm)': 0.2, 'target': 0}
+    ...
+    {'sepal length (cm)': 5.9, 'sepal width (cm)': 3.0, 'petal length (cm)': 5.1, 'petal width (cm)': 1.8, 'target': 2}
+
+
+For more information on working with rows, see
+:ref:`Transforming rows <transforming_rows>` and
+:ref:`Inspecting rows <inspecting-rows>`.
+
+.. _iterating-over-batches:
+
+Iterating over batches
+======================
+
+A batch contains data from multiple rows. Iterate over batches of dataset in different
+formats by calling one of the following methods:
+
+* `Dataset.iter_batches() <ray.data.Dataset.iter_batches>`
+* `Dataset.iter_torch_batches() <ray.data.Dataset.iter_torch_batches>`
+* `Dataset.to_tf() <ray.data.Dataset.to_tf>`
+
+.. tab-set::
+
+    .. tab-item:: NumPy
+        :sync: NumPy
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+
+            for batch in ds.iter_batches(batch_size=2, batch_format="numpy"):
+                print(batch)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'image': array([[[[...]]]], dtype=uint8)}
+            ...
+            {'image': array([[[[...]]]], dtype=uint8)}
+
+    .. tab-item:: pandas
+        :sync: pandas
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+            for batch in ds.iter_batches(batch_size=2, batch_format="pandas"):
+                print(batch)
+
+        .. testoutput::
+            :options: +MOCK
+
+               sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)  target
+            0                5.1               3.5                1.4               0.2       0
+            1                4.9               3.0                1.4               0.2       0
+            ...
+               sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)  target
+            0                6.2               3.4                5.4               2.3       2
+            1                5.9               3.0                5.1               1.8       2
+
+    .. tab-item:: Torch
+        :sync: Torch
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+
+            for batch in ds.iter_torch_batches(batch_size=2):
+                print(batch)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'image': tensor([[[[...]]]], dtype=torch.uint8)}
+            ...
+            {'image': tensor([[[[...]]]], dtype=torch.uint8)}
+
+    .. tab-item:: TensorFlow
+        :sync: TensorFlow
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+            tf_dataset = ds.to_tf(
+                feature_columns="sepal length (cm)",
+                label_columns="target",
+                batch_size=2
+            )
+            for features, labels in tf_dataset:
+                print(features, labels)
+
+        .. testoutput::
+
+            tf.Tensor([5.1 4.9], shape=(2,), dtype=float64) tf.Tensor([0 0], shape=(2,), dtype=int64)
+            ...
+            tf.Tensor([6.2 5.9], shape=(2,), dtype=float64) tf.Tensor([2 2], shape=(2,), dtype=int64)
+
+For more information on working with batches, see
+:ref:`Transforming batches <transforming_batches>` and
+:ref:`Inspecting batches <inspecting-batches>`.
+
+.. _iterating-over-batches-with-shuffling:
+
+Iterating over batches with shuffling
+=====================================
+
+:class:`Dataset.random_shuffle <ray.data.Dataset.random_shuffle>` is slow because it
+shuffles all rows. If a full global shuffle isn't required, you can shuffle a subset of
+rows up to a provided buffer size during iteration by specifying
+``local_shuffle_buffer_size``. While this isn't a true global shuffle like
+``random_shuffle``, it's more performant because it doesn't require excessive data
+movement. For more details about these options, see :doc:`Shuffling Data <shuffling-data>`.
+
+.. tip::
+
+    To configure ``local_shuffle_buffer_size``, choose the smallest value that achieves
+    sufficient randomness. Higher values result in more randomness at the cost of slower
+    iteration. See :ref:`Local shuffle when iterating over batches <local_shuffle_buffer>`
+    on how to diagnose slowdowns.
+
+.. tab-set::
+
+    .. tab-item:: NumPy
+        :sync: NumPy
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+
+            for batch in ds.iter_batches(
+                batch_size=2,
+                batch_format="numpy",
+                local_shuffle_buffer_size=250,
+            ):
+                print(batch)
+
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'image': array([[[[...]]]], dtype=uint8)}
+            ...
+            {'image': array([[[[...]]]], dtype=uint8)}
+
+    .. tab-item:: pandas
+        :sync: pandas
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+            for batch in ds.iter_batches(
+                batch_size=2,
+                batch_format="pandas",
+                local_shuffle_buffer_size=250,
+            ):
+                print(batch)
+
+        .. testoutput::
+            :options: +MOCK
+
+               sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)  target
+            0                6.3               2.9                5.6               1.8       2
+            1                5.7               4.4                1.5               0.4       0
+            ...
+               sepal length (cm)  sepal width (cm)  petal length (cm)  petal width (cm)  target
+            0                5.6               2.7                4.2               1.3       1
+            1                4.8               3.0                1.4               0.1       0
+
+    .. tab-item:: Torch
+        :sync: Torch
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+            for batch in ds.iter_torch_batches(
+                batch_size=2,
+                local_shuffle_buffer_size=250,
+            ):
+                print(batch)
+
+        .. testoutput::
+            :options: +MOCK
+
+            {'image': tensor([[[[...]]]], dtype=torch.uint8)}
+            ...
+            {'image': tensor([[[[...]]]], dtype=torch.uint8)}
+
+    .. tab-item:: TensorFlow
+        :sync: TensorFlow
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+
+            tf_dataset = ds.to_tf(
+                feature_columns="sepal length (cm)",
+                label_columns="target",
+                batch_size=2,
+                local_shuffle_buffer_size=250,
+            )
+            for features, labels in tf_dataset:
+                print(features, labels)
+
+        .. testoutput::
+            :options: +MOCK
+
+            tf.Tensor([5.2 6.3], shape=(2,), dtype=float64) tf.Tensor([1 2], shape=(2,), dtype=int64)
+            ...
+            tf.Tensor([5.  5.8], shape=(2,), dtype=float64) tf.Tensor([0 0], shape=(2,), dtype=int64)
+
+Splitting datasets for distributed parallel training
+====================================================
+
+If you're performing distributed data parallel training, call
+:meth:`Dataset.streaming_split <ray.data.Dataset.streaming_split>` to split your dataset
+into disjoint shards.
+
+.. note::
+
+  If you're using :ref:`Ray Train <train-docs>`, you don't need to split the dataset.
+  Ray Train automatically splits your dataset for you. To learn more, see
+  :ref:`Data Loading for ML Training guide <data-ingest-torch>`.
+
+.. testcode::
+
+    import ray
+
+    @ray.remote
+    class Worker:
+
+        def train(self, data_iterator):
+            for batch in data_iterator.iter_batches(batch_size=8):
+                pass
+
+    ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+    workers = [Worker.remote() for _ in range(4)]
+    shards = ds.streaming_split(n=4, equal=True)
+    ray.get([w.train.remote(s) for w, s in zip(workers, shards)])

doc/source/data/joining-data.rst

@@ -0,0 +1,83 @@
+.. _joining-data:
+
+============
+Joining Data
+============
+
+.. note:: This is a new feature released in Ray 2.46. Note that this is an experimental feature and some things might not work as expected.
+
+Ray Data allows multiple :class:`~ray.data.dataset.Dataset` instances to be joined using different join types based on the provided key columns as follows:
+
+.. testcode::
+
+    import ray
+
+    doubles_ds = ray.data.range(4).map(
+        lambda row: {"id": row["id"], "double": int(row["id"]) * 2}
+    )
+
+    squares_ds = ray.data.range(4).map(
+        lambda row: {"id": row["id"], "square": int(row["id"]) ** 2}
+    )
+
+    doubles_and_squares_ds = doubles_ds.join(
+        squares_ds,
+        join_type="inner",
+        num_partitions=2,
+        on=("id",),
+    )
+
+Ray Data supports the following join types (check out `Dataset.join` docs for up-to-date list):
+
+**Inner/Outer Joins:**
+- Inner, Left Outer, Right Outer, Full Outer
+
+**Semi Joins:**
+- Left Semi, Right Semi (returns all rows that have at least one matching row in the other table,
+only returning columns from the requested side)
+
+**Anti Joins:**
+- Left Anti, Right Anti (return rows that have no matching rows in the other table, only returning
+columns from the requested side)
+
+Internally joins are currently powered by the :ref:`hash-shuffle backend <hash-shuffle>`.
+
+Configuring Joins
+----------------------------------
+
+Joins are generally memory-intensive operations that require accurate memory accounting and projection and hence are sensitive to skews and imbalances in the dataset.
+
+Ray Data provides the following levers to allow tuning the performance of joins for your workload:
+
+-   `num_partitions`: (required) specifies number of partitions both incoming datasets will be hash-partitioned into. Check out :ref:`configuring number of partitions <joins_configuring_num_partitions>` section for guidance on how to tune this up.
+-   `partition_size_hint`: (optional) Hint to joining operator about the estimated avg expected size of the individual partition (in bytes). If not specified, defaults to DataContext.target_max_block_size (128Mb by default).
+    -   Note that, `num_partitions * partition_size_hint` should ideally be approximating actual dataset size, ie `partition_size_hint` could be estimated as dataset size divided by `num_partitions` (assuming relatively evenly sized partitions)
+    -   However, in cases when dataset partitioning is expected to be heavily skewed `partition_size_hint` should approximate largest partition to prevent Out-of-Memory (OOM) errors
+
+.. _joins_configuring_num_partitions:
+
+Configuring number of partitions
+--------------------------------------------
+
+Number of partitions (also referred to as blocks) provide an important trade-off between the size of individual batch of rows handled by individual tasks against memory requirements of the operation performed on them
+
+**Rule of thumb**: *keep partitions large, but not too large to cause Out-of-Memory (OOM) errors*
+
+1.  It’s important to not “oversize” partitions for joins as that could lead to OOM errors (if joined partitions might be too large to fit in memory)
+2.  It’s also important to not create too many small partitions as this creates an overhead of passing large amount of smaller objects
+
+Configuring number of Aggregators
+----------------------------------------------
+
+“Aggregators” are worker actors that perform actual joins/aggregations/shuffling, they receive individual partition chunks from the incoming blocks and subsequently "aggregate" them in the way that's required to perform given operation.
+
+Following are important considerations for successfully configuring number of aggregators in your pool:
+
+- Defaults to 64 or `num_partitions` (in cases when there are less than 64 partitions)
+- Individual Aggregators might be assigned to handle more than one partition (partitions are evenly split in round-robin fashion among the aggregators)
+- Aggregators are stateful components that hold the state (partitions) during shuffling **in memory**
+
+.. note:: The rule of thumb is to avoid setting `num_partitions` >> number of aggregators as it might create bottlenecks
+
+1.  Setting `DataContext.max_hash_shuffle_aggregators` caps the number of aggregators
+2.  Setting it to large enough value has an effect of allocating 1 partition to 1 aggregator (when `max_hash_shuffle_aggregators >= num_partitions`)

doc/source/data/key-concepts.rst

@@ -0,0 +1,134 @@
+.. _data_key_concepts:
+
+Key Concepts
+============
+
+
+Datasets and blocks
+-------------------
+
+There are two main concepts in Ray Data: Datasets and Blocks.
+
+A :class:`Dataset <ray.data.Dataset>` represents a distributed data collection and defines data loading and processing operations and is the primary user-facing API for Ray Data.
+Users typically use the API by creating a :class:`Dataset <ray.data.Dataset>` from external storage or in-memory data, applying transformations to the data, and writing the outputs to external storage or feeding the outputs to training workers.
+
+The Dataset API is lazy, meaning that operations aren't executed until you materialize or consume the dataset, with methods like :meth:`~ray.data.Dataset.show`. This allows Ray Data to optimize the execution plan and execute operations in a pipelined, streaming fashion.
+
+A *block* is a set of rows representing single partition of the dataset. Blocks, as a collection of rows represented by columnar formats (like Arrow) are the basic unit of data processing in Ray Data:
+
+ 1. Every dataset is partitioned into a number of blocks, then
+ 2. Processing of the whole dataset is distributed and parallelized at the block level (blocks are processed in parallel and for the most part independently)
+
+The following figure visualizes a dataset with three blocks, each holding 1000 rows.
+Ray Data holds the :class:`~ray.data.Dataset` on the process that triggers execution
+(which is usually the entrypoint of the program, referred to as the :term:`driver`)
+and stores the blocks as objects in Ray's shared-memory :ref:`object store <objects-in-ray>`. Internally, Ray Data can natively handle blocks either
+as Pandas ``DataFrame`` or PyArrow ``Table``.
+
+.. image:: images/dataset-arch-with-blocks.svg
+..
+  https://docs.google.com/drawings/d/1kOYQqHdMrBp2XorDIn0u0G_MvFj-uSA4qm6xf9tsFLM/edit
+
+Operators and Plans
+-------------------
+
+Ray Data uses a two-phase planning process to execute operations efficiently. When you write a program using the Dataset API, Ray Data first builds a *logical plan* - a high-level description of what operations to perform. When execution begins, it converts this into a *physical plan* that specifies exactly how to execute those operations.
+
+This diagram illustrates the complete planning process:
+
+.. https://docs.google.com/drawings/d/1WrVAg3LwjPo44vjLsn17WLgc3ta2LeQGgRfE8UHrDA0/edit
+
+.. image:: images/get_execution_plan.svg
+   :width: 600
+   :align: center
+
+The building blocks of these plans are operators:
+
+* Logical plans consist of *logical operators* that describe *what* operation to perform. For example, when you write ``dataset = ray.data.read_parquet(...)``, Ray Data creates a ``ReadOp`` logical operator to specify what data to read.
+* Physical plans consist of *physical operators* that describe *how* to execute the operation. For example, Ray Data converts the ``ReadOp`` logical operator into a ``TaskPoolMapOperator`` physical operator that launches Ray tasks to read the data.
+
+Here is a simple example of how Ray Data builds a logical plan. As you chain operations together, Ray Data constructs the logical plan behind the scenes:
+
+.. testcode::
+    import ray
+
+    dataset = ray.data.range(100)
+    dataset = dataset.add_column("test", lambda x: x["id"] + 1)
+    dataset = dataset.select_columns("test")
+
+You can inspect the resulting logical plan by printing the dataset:
+
+.. code-block::
+
+    Project
+    +- MapBatches(add_column)
+       +- Dataset(schema={...})
+
+When execution begins, Ray Data optimizes the logical plan, then translates it into a physical plan - a series of operators that implement the actual data transformations. During this translation:
+
+1. A single logical operator may become multiple physical operators. For example, ``ReadOp`` becomes both ``InputDataBuffer`` and ``TaskPoolMapOperator``.
+2. Both logical and physical plans go through optimization passes. For example, ``OperatorFusionRule`` combines map operators to reduce serialization overhead.
+
+Physical operators work by:
+
+* Taking in a stream of block references
+* Performing their operation (either transforming data with Ray Tasks/Actors or manipulating references)
+* Outputting another stream of block references
+
+For more details on Ray Tasks and Actors, see :ref:`Ray Core Concepts <core-key-concepts>`.
+
+.. note:: A dataset's execution plan only runs when you materialize or consume the dataset through operations like :meth:`~ray.data.Dataset.show`.
+
+.. _streaming-execution:
+
+Streaming execution model
+-------------------------
+
+Ray Data can stream data through a pipeline of operators to efficiently process large datasets.
+
+This means that different operators in an execution can be scaled independently while running concurrently, allowing for more flexible and fine-grained resource allocation. For example, if two map operators require different amounts or types of resources, the streaming execution model can allow them to run concurrently and independently while still maintaining high performance.
+
+Note that this is primarily useful for non-shuffle operations. Shuffle operations like :meth:`ds.sort() <ray.data.Dataset.sort>` and :meth:`ds.groupby() <ray.data.Dataset.groupby>` require materializing data, which stops streaming until the shuffle is complete.
+
+Here is an example of how the streaming execution works in Ray Data.
+
+.. code-block:: python
+
+    import ray
+
+    # Create a dataset with 1K rows
+    ds = ray.data.read_parquet(...)
+
+    # Define a pipeline of operations
+    ds = ds.map(cpu_function, num_cpus=2)
+    ds = ds.map(GPUClass, num_gpus=1)
+    ds = ds.map(cpu_function2, num_cpus=4)
+    ds = ds.filter(filter_func)
+
+    # Data starts flowing when you call a method like show()
+    ds.show(5)
+
+This creates a logical plan like the following:
+
+.. code-block::
+
+    Filter(filter_func)
+    +- Map(cpu_function2)
+       +- Map(GPUClass)
+          +- Map(cpu_function)
+                +- Dataset(schema={...})
+
+
+The streaming topology looks like the following:
+
+.. https://docs.google.com/drawings/d/10myFIVtpI_ZNdvTSxsaHlOhA_gHRdUde_aHRC9zlfOw/edit
+
+.. image:: images/streaming-topology.svg
+   :width: 1000
+   :align: center
+
+In the streaming execution model, operators are connected in a pipeline, with each operator's output queue feeding directly into the input queue of the next downstream operator. This creates an efficient flow of data through the execution plan.
+
+This enables multiple stages to execute concurrently, improving overall performance and resource utilization. For example, if the map operator requires GPU resources, the streaming execution model can execute the map operator concurrently with the filter operator (which may run on CPUs), effectively utilizing the GPU through the entire duration of the pipeline.
+
+You can read more about the streaming execution model in this `blog post <https://www.anyscale.com/blog/streaming-distributed-execution-across-cpus-and-gpus>`__.

doc/source/data/mixing-data.rst

@@ -0,0 +1,158 @@
+.. _mixing_data:
+
+Weighted Dataset Mixing
+=======================
+
+Ray Data allows you to combine multiple datasets into a single streaming dataset with control over how often rows from each source appear. This is useful for:
+
+- **Class / scenario balancing**: upsample rare scenarios or harder tasks so that training batches see them more often.
+- **Multi-task pretraining**: combine code and web text datasets at fixed ratios.
+- **Catastrophic forgetting prevention**: keep a small fraction of an older dataset in the mix while training on a newer one.
+
+Quickstart
+----------
+
+.. testcode::
+
+    import ray.data
+    from ray.train.torch import TorchTrainer
+    from ray.train import ScalingConfig
+
+    # Read and preprocess each source independently.
+    # NOTE: These are mocked datasets for demonstration purposes.
+    def preprocess(row):
+        return row
+
+    ds1 = ray.data.from_items([{"x": 1} for _ in range(750)]).map(preprocess)
+    ds2 = ray.data.from_items([{"x": 2} for _ in range(250)]).map(preprocess)
+
+    # Output batches will contain 75% rows from ds1, 25% from ds2 (in expectation).
+    mixed = ds1.mix(ds2, weights=[0.75, 0.25])
+
+    def train_fn_per_worker(config):
+        shard = ray.train.get_dataset_shard("train")
+        for batch in shard.iter_torch_batches(batch_size=128):
+            print(batch)
+
+    trainer = TorchTrainer(
+        train_loop_per_worker=train_fn_per_worker,
+        scaling_config=ScalingConfig(num_workers=4),
+        datasets={"train": mixed},
+    )
+
+Mixing strategies
+-----------------
+
+You can compose :meth:`~ray.data.Dataset.mix` with other Ray Data operations to implement different mixing strategies, depending on how granular you want the mixing ratio to be. The sections below cover **per-block mixing** (:meth:`~ray.data.Dataset.mix` on its own) and **random mixing** (:meth:`~ray.data.Dataset.mix` followed by a shuffle).
+
+Per-block mixing
+~~~~~~~~~~~~~~~~
+
+By default, each output block comes from exactly one input dataset. :meth:`~ray.data.Dataset.mix` keeps a running row count per source and, on every step, pulls the next block from whichever dataset is furthest behind its target ratio. Over time, the cumulative row counts converge to the requested weights.
+
+Suppose you mix two datasets ``ds1`` and ``ds2`` with ``weights=[0.75, 0.25]``, and both sources produce blocks of equal size. This data pipeline then splits across 4 training workers, and data parallel training constructs a global batch across all workers.
+
+.. image:: /data/images/dataset_mixing/per_block_mix.png
+    :alt: Per-block mixing: blocks from ds1 and ds2 are interleaved in a 3:1 pattern, then split across 4 training workers to form a global batch.
+
+With uniform block sizes, the ratio is exact within any window of ``1 / min(weights)`` blocks. With ``weights=[0.9, 0.1]``, you're guaranteed a block from the second dataset at least once in every 10-block window.
+
+.. note::
+
+    :ref:`Blocks <dataset_concept>` are the unit of data transfer in Ray Data, and they don't map 1:1 to training batches. Workers construct each batch by pulling rows from one or more blocks. With per-block mixing, this means each local batch may contain data from one or more of the input datasets, depending on how block sizes compare to batch sizes. The next section covers how to align them with a streaming repartition.
+
+Advanced: Standardize input block sizes
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+If your input datasets produce blocks of very different sizes, a single large block can temporarily push that source ahead of its target ratio. :meth:`~ray.data.Dataset.mix` self-corrects on subsequent pulls, so the ratio is still correct in expectation---but a global batch built from a small number of those blocks can look skewed.
+
+To tighten the per-batch window, standardize input block sizes upstream with :meth:`ds.repartition(target_num_rows_per_block) <ray.data.Dataset.repartition>`:
+
+.. testcode::
+
+    LOCAL_BATCH_SIZE = 128
+
+    ds1 = ray.data.from_items([{"x": 1} for _ in range(750)]).map(preprocess)
+    ds2 = ray.data.from_items([{"x": 2} for _ in range(250)]).map(preprocess)
+
+    # Standardize block sizes so the ratio holds within tighter windows.
+    ds1 = ds1.repartition(target_num_rows_per_block=LOCAL_BATCH_SIZE)
+    ds2 = ds2.repartition(target_num_rows_per_block=LOCAL_BATCH_SIZE)
+
+    mixed = ds1.mix(ds2, weights=[0.75, 0.25])
+
+
+.. note::
+
+    You may want to repartition to some multiple of batch size (for example, ``N * LOCAL_BATCH_SIZE``) if your rows are small in terms of bytes. This prevents splitting blocks into extremely small pieces that increase overhead.
+
+Random mixing
+~~~~~~~~~~~~~
+
+The per-batch ratio quality of per-block mixing depends on two things: the sizes of the input blocks (covered in the preceding section) and the number of training workers contributing to each global batch. A global batch aggregates ``num_workers * grad_accum_steps`` local batches, each drawn from a single dataset, so the more local batches you have per global batch, the closer the ratio holds to the target.
+
+The extreme case: training on a single worker with no gradient accumulation means every global batch is a local batch, so every batch comes from a single dataset.
+
+Adding a streaming shuffle after :meth:`~ray.data.Dataset.mix` switches you to **random mixing**: the shuffle redistributes rows across block boundaries so each batch directly contains rows from multiple datasets in roughly the requested proportion, regardless of how many workers you're training on. :meth:`~ray.data.Dataset.mix` still governs the ratio; the shuffle just spreads it within each batch.
+
+.. image:: /data/images/dataset_mixing/random_mix.png
+    :alt: Random mixing: after mix(), a shuffle redistributes rows so that each worker batch contains rows from multiple datasets in the target proportion.
+
+Two streaming-friendly shuffle options in Ray Data:
+
+- :ref:`Local buffer shuffle <local_shuffle_buffer>` (:meth:`~ray.data.DataIterator.iter_batches` with ``local_shuffle_buffer_size``)
+- :ref:`map_batches shuffle <map_batches_shuffle>`
+
+.. testcode::
+
+    import numpy as np
+    import pyarrow as pa
+
+    LOCAL_BATCH_SIZE = 128
+
+    ds1 = ray.data.from_items([{"x": 1} for _ in range(750)]).map(preprocess)
+    ds2 = ray.data.from_items([{"x": 2} for _ in range(250)]).map(preprocess)
+
+    ds1 = ds1.repartition(target_num_rows_per_block=LOCAL_BATCH_SIZE)
+    ds2 = ds2.repartition(target_num_rows_per_block=LOCAL_BATCH_SIZE)
+
+    mixed = ds1.mix(ds2, weights=[0.75, 0.25])
+
+    # Add a shuffle after mix() to get random mixing.
+    def random_shuffle(batch: pa.Table) -> pa.Table:
+        indices = np.random.permutation(len(batch))
+        return batch.take(indices)
+
+    # Set the shuffle buffer size to be large enough for good mixing quality across datasets.
+    SHUFFLE_BUFFER_SIZE = 64 * LOCAL_BATCH_SIZE
+    mixed = mixed.map_batches(random_shuffle, batch_size=SHUFFLE_BUFFER_SIZE, batch_format="pyarrow")
+
+
+Stopping conditions
+-------------------
+
+.. list-table::
+   :header-rows: 1
+
+   * - Condition
+     - Behavior
+   * - ``STOP_ON_LONGEST_DROP`` (default)
+     - Pipeline ends when the longest dataset is exhausted. Shorter datasets drop out once exhausted; remaining batches come from the still-active datasets.
+   * - ``STOP_ON_SHORTEST``
+     - Pipeline ends when the shortest dataset is exhausted. Other datasets are truncated.
+
+See :class:`~ray.data.MixStoppingCondition` for more details.
+
+Limitations
+-----------
+
+- **Avoid** :meth:`~ray.data.Dataset.map` / :meth:`~ray.data.Dataset.filter` **after** :meth:`~ray.data.Dataset.mix`. Downstream transformations can combine or split blocks before they reach the trainer, which breaks the row-ratio guarantees :meth:`~ray.data.Dataset.mix` provides. Apply per-dataset transforms upstream of :meth:`~ray.data.Dataset.mix`.
+- **Schemas must match.** :meth:`~ray.data.Dataset.mix` does not unify schemas for you. Apply :meth:`~ray.data.Dataset.map` or :meth:`~ray.data.Dataset.select_columns` upstream to make all inputs structurally identical.
+- **Heavily skewed weights (current limitation).** All input datasets currently execute concurrently with some portion of cluster resources equally divided between them. With heavily skewed weights (for example, ``[0.95, 0.05]``), the high-weight dataset may bottleneck while the low-weight dataset idles. For now, keep weights within roughly 5x of each other (for example, ``[0.4, 0.3, 0.2, 0.1]``).
+
+See also
+--------
+
+- :ref:`Using Ray Data with Ray Train for distributed training and data ingest <data-ingest-torch>`
+- :ref:`Ray Data shuffling solutions <shuffling_data>`
+- :meth:`ray.data.Dataset.repartition`

doc/source/data/monitoring-your-workload.rst

@@ -0,0 +1,539 @@
+.. _monitoring-your-workload:
+
+Monitoring Your Workload
+========================
+
+This section helps you debug and monitor the execution of your :class:`~ray.data.Dataset` by viewing the:
+
+* :ref:`Ray Data progress bars <ray-data-progress-bars>`
+* :ref:`Ray Data dashboard <ray-data-dashboard>`
+* :ref:`Ray Data logs <ray-data-logs>`
+* :ref:`Ray Data stats <ray-data-stats>`
+
+.. _ray-data-progress-bars:
+
+Ray Data progress bars
+----------------------
+
+When you execute a :class:`~ray.data.Dataset`, Ray Data displays a set of progress bars in the console. These progress bars show various execution and progress-related metrics, including the number of rows completed/remaining, resource usage, and task/actor status. See the annotated image for a breakdown of how to interpret the progress bar outputs: 
+
+.. image:: images/dataset-progress-bar.png
+   :align: center
+
+
+Some additional notes on progress bars:
+
+* The progress bars are updated every second; resource usage, metrics, and task/actor status may take up to 5 seconds to update.
+* When the tasks section contains the label `[backpressure]`, it indicates that the operator is *backpressured*, meaning that the operator won't submit more tasks until the downstream operator is ready to accept more data.
+* The global resource usage is the sum of resources used by all operators, active and requested (includes pending scheduling and pending node assignment).
+
+Configuring the progress bar
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+Depending on your use case, you may not be interested in the full progress bar output, or wish to turn them off altogether. Ray Data provides several ways to accomplish this:
+
+* Disabling operator-level progress bars: Set `DataContext.get_current().enable_operator_progress_bars = False`. This only shows the global progress bar, and omits operator-level progress bars.
+* Disabling all progress bars: Set `DataContext.get_current().enable_progress_bars = False`. This disables all progress bars from Ray Data related to dataset execution.
+* Disabling `ray_tqdm`: Set `DataContext.get_current().use_ray_tqdm = False`. This configures Ray Data to use the base `tqdm` library instead of the custom distributed `tqdm` implementation, which could be useful when debugging logging issues in a distributed setting.
+
+For operator names longer than a threshold of 100 characters, Ray Data truncates the names by default, to prevent the case when the operator names are long and the progress bar is too wide to fit on the screen.
+
+* To turn off this behavior and show the full operator name, set `DataContext.get_current().enable_progress_bar_name_truncation = False`.
+* To change the threshold of truncating the name, update the constant `ray.data._internal.progress_bar.ProgressBar.MAX_NAME_LENGTH = 42`.
+
+.. tip::
+   There is a new experimental console UI to show progress bars. Set `DataContext.get_current().enable_rich_progress_bars = True` or set the `RAY_DATA_ENABLE_RICH_PROGRESS_BARS=1` environment variable to enable.
+
+.. _ray-data-dashboard:
+
+Ray Data dashboard
+------------------
+
+Ray Data emits Prometheus metrics in real-time while a Dataset is executing. These metrics are tagged by both dataset and operator, and are displayed in multiple views across the Ray dashboard.
+
+.. note::
+   Most metrics are only available for physical operators that use the map operation. For example, physical operators created by :meth:`~ray.data.Dataset.map_batches`, :meth:`~ray.data.Dataset.map`, and :meth:`~ray.data.Dataset.flat_map`.
+
+Jobs: Ray Data overview
+~~~~~~~~~~~~~~~~~~~~~~~
+
+For an overview of all datasets that have been running on your cluster, see the Ray Data Overview in the :ref:`jobs view <dash-jobs-view>`. This table appears once the first dataset starts executing on the cluster, and shows dataset details such as:
+
+* execution progress (measured in blocks)
+* execution state (running, failed, or finished)
+* dataset start/end time
+* dataset-level metrics (for example, sum of rows processed over all operators)
+
+.. image:: images/data-overview-table.png
+   :align: center
+
+For a more fine-grained overview, each dataset row in the table can also be expanded to display the same details for individual operators.
+
+.. image:: images/data-overview-table-expanded.png
+   :align: center
+
+.. tip::
+
+    To evaluate a dataset-level metric where it's not appropriate to sum the values of all the individual operators, it may be more useful to look at the operator-level metrics of the last operator. For example, to calculate a dataset's throughput, use the "Rows Outputted" of the dataset's last operator, because the dataset-level metric contains the sum of rows outputted over all operators.
+
+Ray dashboard metrics
+~~~~~~~~~~~~~~~~~~~~~
+
+For a time-series view of these metrics, see the Ray Data section in the :ref:`Metrics view <dash-metrics-view>`. This section contains time-series graphs of all metrics emitted by Ray Data. Execution metrics are grouped by dataset and operator, and iteration metrics are grouped by dataset.
+
+The metrics recorded include:
+
+* Bytes spilled by objects from object store to disk
+* Bytes of objects allocated in object store
+* Bytes of objects freed in object store
+* Current total bytes of objects in object store
+* Logical CPUs allocated to dataset operators
+* Logical GPUs allocated to dataset operators
+* Bytes outputted by dataset operators
+* Rows outputted by dataset operators
+* Input blocks received by data operators
+* Input blocks/bytes processed in tasks by data operators
+* Input bytes submitted to tasks by data operators
+* Output blocks/bytes/rows generated in tasks by data operators
+* Output blocks/bytes taken by downstream operators
+* Output blocks/bytes from finished tasks
+* Submitted tasks
+* Running tasks
+* Tasks with at least one output block
+* Finished tasks
+* Failed tasks
+* Operator internal inqueue size (in blocks/bytes)
+* Operator internal outqueue size (in blocks/bytes)
+* Size of blocks used in pending tasks
+* Freed memory in object store
+* Spilled memory in object store
+* Time spent generating blocks
+* Time spent in task submission backpressure
+* Time spent to initialize iteration.
+* Time user code is blocked during iteration.
+* Time spent in user code during iteration.
+
+.. image:: images/data-dashboard.png
+   :align: center
+
+
+To learn more about the Ray dashboard, including detailed setup instructions, see :ref:`Ray Dashboard <observability-getting-started>`.
+
+Prometheus metrics
+~~~~~~~~~~~~~~~~~~
+
+Ray Data emits Prometheus metrics that you can use to monitor dataset execution. The metrics are tagged with `dataset` and `operator` labels to help you identify which dataset and operator the metrics are coming from.
+
+To access these metrics, you can query the Prometheus server running on the Ray head node. The default Prometheus server URL is `http://<head-node-ip>:8080`.
+
+The following tables list all available Ray Data metrics grouped by category.
+
+Overview metrics
+^^^^^^^^^^^^^^^^
+
+These metrics provide high-level information about dataset execution and resource usage.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Metric name
+     - Description
+   * - `data_spilled_bytes`
+     - Bytes spilled by dataset operators. Set `DataContext.enable_get_object_locations_for_metrics` to `True` to report this metric.
+   * - `data_freed_bytes`
+     - Bytes freed by dataset operators
+   * - `data_current_bytes`
+     - Bytes of object store memory used by dataset operators
+   * - `data_cpu_usage_cores`
+     - CPUs allocated to dataset operators
+   * - `data_gpu_usage_cores`
+     - GPUs allocated to dataset operators
+   * - `data_output_bytes`
+     - Bytes outputted by dataset operators
+   * - `data_output_rows`
+     - Rows outputted by dataset operators
+
+Input metrics
+^^^^^^^^^^^^^
+
+These metrics track input data flowing into operators.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Metric name
+     - Description
+   * - `num_inputs_received`
+     - Number of input blocks received by operator
+   * - `num_row_inputs_received`
+     - Number of input rows received by operator
+   * - `bytes_inputs_received`
+     - Byte size of input blocks received by operator
+   * - `num_task_inputs_processed`
+     - Number of input blocks that the operator's tasks finished processing
+   * - `bytes_task_inputs_processed`
+     - Byte size of input blocks that the operator's tasks finished processing
+   * - `bytes_inputs_of_submitted_tasks`
+     - Byte size of input blocks passed to submitted tasks
+   * - `rows_inputs_of_submitted_tasks`
+     - Number of rows in the input blocks passed to submitted tasks
+   * - `average_num_inputs_per_task`
+     - Average number of input blocks per task, or `None` if no task finished
+   * - `average_bytes_inputs_per_task`
+     - Average size in bytes of ref bundles passed to tasks, or `None` if no tasks submitted
+   * - `average_rows_inputs_per_task`
+     - Average number of rows in input blocks per task, or `None` if no task submitted
+
+Output metrics
+^^^^^^^^^^^^^^
+
+These metrics track output data generated by operators.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Metric name
+     - Description
+   * - `num_task_outputs_generated`
+     - Number of output blocks generated by tasks
+   * - `bytes_task_outputs_generated`
+     - Byte size of output blocks generated by tasks
+   * - `rows_task_outputs_generated`
+     - Number of output rows generated by tasks
+   * - `row_outputs_taken`
+     - Number of rows that are already taken by downstream operators
+   * - `block_outputs_taken`
+     - Number of blocks that are already taken by downstream operators
+   * - `num_outputs_taken`
+     - Number of output blocks that are already taken by downstream operators
+   * - `bytes_outputs_taken`
+     - Byte size of output blocks that are already taken by downstream operators
+   * - `num_outputs_of_finished_tasks`
+     - Number of generated output blocks that are from finished tasks
+   * - `bytes_outputs_of_finished_tasks`
+     - Total byte size of generated output blocks produced by finished tasks
+   * - `rows_outputs_of_finished_tasks`
+     - Number of rows generated by finished tasks
+   * - `num_external_inqueue_blocks`
+     - Number of blocks in the external inqueue
+   * - `num_external_inqueue_bytes`
+     - Byte size of blocks in the external inqueue
+   * - `num_external_outqueue_blocks`
+     - Number of blocks in the external outqueue
+   * - `num_external_outqueue_bytes`
+     - Byte size of blocks in the external outqueue
+   * - `average_num_outputs_per_task`
+     - Average number of output blocks per task, or `None` if no task finished
+   * - `average_bytes_per_output`
+     - Average size in bytes of output blocks
+   * - `average_bytes_outputs_per_task`
+     - Average total output size of task in bytes, or `None` if no task finished
+   * - `average_rows_outputs_per_task`
+     - Average number of rows produced per task, or `None` if no task finished
+   * - `num_output_blocks_per_task_s`
+     - Average number of output blocks per task per second
+
+Task metrics
+^^^^^^^^^^^^
+
+These metrics track task execution and scheduling.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Metric name
+     - Description
+   * - `num_tasks_submitted`
+     - Number of submitted tasks
+   * - `num_tasks_running`
+     - Number of running tasks
+   * - `num_tasks_have_outputs`
+     - Number of tasks with at least one output
+   * - `num_tasks_finished`
+     - Number of finished tasks
+   * - `num_tasks_failed`
+     - Number of failed tasks
+   * - `block_generation_time`
+     - Time spent generating blocks in tasks
+   * - `task_submission_backpressure_time`
+     - Time spent in task submission backpressure
+   * - `task_output_backpressure_time`
+     - Time spent in task output backpressure
+   * - `task_completion_time`
+     - Histogram of time spent running tasks to completion
+   * - `block_completion_time`
+     - Histogram of time spent running a single block to completion. If multiple blocks are generated per task, Ray Data approximates this by assuming each block took an equal amount of time to process.
+   * - `task_completion_time_s`
+     - Time spent running tasks to completion
+   * - `task_completion_time_excl_backpressure_s`
+     - Time spent running tasks to completion without backpressure
+   * - `block_size_bytes`
+     - Histogram of block sizes in bytes generated by tasks
+   * - `block_size_rows`
+     - Histogram of number of rows in blocks generated by tasks
+   * - `average_total_task_completion_time_s`
+     - Average task completion time in seconds including throttling. This includes Ray Core and Ray Data backpressure.
+   * - `average_task_completion_excl_backpressure_time_s`
+     - Average task completion time in seconds excluding throttling
+   * - `average_max_uss_per_task`
+     - Average Unique Set Size (USS) memory usage of tasks. USS is the amount of memory unique to a process that would be freed if the process was terminated.
+
+Actor metrics
+^^^^^^^^^^^^^
+
+These metrics track actor lifecycle for operations that use actors.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Metric name
+     - Description
+   * - `num_alive_actors`
+     - Number of alive actors
+   * - `num_restarting_actors`
+     - Number of restarting actors
+   * - `num_pending_actors`
+     - Number of pending actors
+
+Object store memory metrics
+^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+These metrics track memory usage in the Ray object store.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Metric name
+     - Description
+   * - `obj_store_mem_internal_inqueue_blocks`
+     - Number of blocks in the operator's internal input queue
+   * - `obj_store_mem_internal_outqueue_blocks`
+     - Number of blocks in the operator's internal output queue
+   * - `obj_store_mem_freed`
+     - Byte size of freed memory in object store
+   * - `obj_store_mem_spilled`
+     - Byte size of spilled memory in object store
+   * - `obj_store_mem_used`
+     - Byte size of used memory in object store
+   * - `obj_store_mem_internal_inqueue`
+     - Byte size of input blocks in the operator's internal input queue
+   * - `obj_store_mem_internal_outqueue`
+     - Byte size of output blocks in the operator's internal output queue
+   * - `obj_store_mem_pending_task_inputs`
+     - Byte size of input blocks used by pending tasks
+
+Scheduling and resource metrics
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+These metrics track resource allocation and scheduling behavior in the streaming executor.
+
+.. list-table::
+   :header-rows: 1
+   :widths: 30 70
+
+   * - Metric name
+     - Description
+   * - `data_sched_loop_duration_s`
+     - Duration of the scheduling loop in seconds
+   * - `data_cpu_budget`
+     - CPU budget allocated per operator
+   * - `data_gpu_budget`
+     - GPU budget allocated per operator
+   * - `data_memory_budget`
+     - Memory budget allocated per operator
+   * - `data_object_store_memory_budget`
+     - Object store memory budget allocated per operator
+   * - `data_max_bytes_to_read`
+     - Maximum bytes to read from streaming generator buffer per operator
+
+.. _ray-data-logs:
+
+Ray Data logs
+-------------
+During execution, Ray Data periodically logs updates to `ray-data.log`.
+
+Every five seconds, Ray Data logs the execution progress of every operator in the dataset. For more frequent updates, set `RAY_DATA_TRACE_SCHEDULING=1` so that the progress is logged after each task is dispatched.
+
+.. code-block:: text
+
+   Execution Progress:
+   0: - Input: 0 active, 0 queued, 0.0 MiB objects, Blocks Outputted: 200/200
+   1: - ReadRange->MapBatches(<lambda>): 10 active, 190 queued, 381.47 MiB objects, Blocks Outputted: 100/200
+
+When an operator completes, the metrics for that operator are also logged.
+
+.. code-block:: text
+
+   Operator InputDataBuffer[Input] -> TaskPoolMapOperator[ReadRange->MapBatches(<lambda>)] completed. Operator Metrics:
+   {'num_inputs_received': 20, 'bytes_inputs_received': 46440, 'num_task_inputs_processed': 20, 'bytes_task_inputs_processed': 46440, 'num_task_outputs_generated': 20, 'bytes_task_outputs_generated': 800, 'rows_task_outputs_generated': 100, 'num_outputs_taken': 20, 'bytes_outputs_taken': 800, 'num_outputs_of_finished_tasks': 20, 'bytes_outputs_of_finished_tasks': 800, 'num_tasks_submitted': 20, 'num_tasks_running': 0, 'num_tasks_have_outputs': 20, 'num_tasks_finished': 20, 'obj_store_mem_freed': 46440, 'obj_store_mem_spilled': 0, 'block_generation_time': 1.191296085, 'cpu_usage': 0, 'gpu_usage': 0, 'ray_remote_args': {'num_cpus': 1, 'scheduling_strategy': 'SPREAD'}}
+
+This log file can be found locally at `/tmp/ray/{SESSION_NAME}/logs/ray-data/ray-data.log`. It can also be found on the Ray Dashboard under the head node's logs in the :ref:`Logs view <dash-logs-view>`.
+
+.. _ray-data-stats:
+
+Ray Data stats
+--------------
+To see detailed stats on the execution of a dataset you can use the :meth:`~ray.data.Dataset.stats` method.
+
+Operator stats
+~~~~~~~~~~~~~~
+The stats output includes a summary on the individual operator's execution stats for each operator. Ray Data calculates this
+summary across many different blocks, so some stats show the min, max, mean, and sum of the stats aggregated over all the blocks.
+The following are descriptions of the various stats included at the operator level:
+
+* **Remote wall time**: The wall time is the start to finish time for an operator. It includes the time where the operator
+  isn't processing data, sleeping, waiting for I/O, etc.
+* **Remote CPU time**: The CPU time is the process time for an operator which excludes time slept. This time includes both
+  user and system CPU time.
+* **UDF time**: The UDF time is time spent in functions defined by the user. This time includes functions you pass into Ray
+  Data methods, including :meth:`~ray.data.Dataset.map`, :meth:`~ray.data.Dataset.map_batches`, :meth:`~ray.data.Dataset.filter`,
+  etc. You can use this stat to track the time spent in functions you define and how much time optimizing those functions could save.
+* **Memory usage**: The output displays memory usage per block in MiB.
+* **Output stats**: The output includes stats on the number of rows output and size of output in bytes per block. The number of
+  output rows per task is also included. All of this together gives you insight into how much data Ray Data is outputting at a per
+  block and per task level.
+* **Task Stats**: The output shows the scheduling of tasks to nodes, which allows you to see if you are utilizing all of your nodes
+  as expected.
+* **Throughput**: The summary calculates the throughput for the operator, and for a point of comparison, it also computes an estimate of
+  the throughput of the same task on a single node. This estimate assumes the total time of the work remains the same, but with no
+  concurrency. The overall summary also calculates the throughput at the dataset level, including a single node estimate.
+
+Iterator stats
+~~~~~~~~~~~~~~
+If you iterate over the data, Ray Data also generates iteration stats. Even if you aren't directly iterating over the data, you
+might see iteration stats, for example, if you call :meth:`~ray.data.Dataset.take_all`. Some of the stats that Ray Data includes
+at the iterator level are:
+
+* **Iterator initialization**: The time Ray Data spent initializing the iterator. This time is internal to Ray Data.
+* **Time user thread is blocked**: The time Ray Data spent producing data in the iterator. This time is often the primary execution of a
+  dataset if you haven't previously materialized it.
+* **Time in user thread**: The time spent in the user thread that's iterating over the dataset outside of the Ray Data code.
+  If this time is high, consider optimizing the body of the loop that's iterating over the dataset.
+* **Batch iteration stats**: Ray Data also includes stats about the prefetching of batches. These times are internal to Ray
+  Data code, but you can further optimize this time by tuning the prefetching process.
+
+Verbose stats
+~~~~~~~~~~~~~~
+By default, Ray Data only logs the most important high-level stats. To enable verbose stats outputs, include
+the following snippet in your Ray Data code:
+
+.. testcode::
+
+   from ray.data import DataContext
+
+   context = DataContext.get_current()
+   context.verbose_stats_logs = True
+
+
+By enabling verbosity Ray Data adds a few more outputs:
+
+* **Extra metrics**: Operators, executors, etc. can add to this dictionary of various metrics. There is
+  some duplication of stats between the default output and this dictionary, but for advanced users this stat provides more
+  insight into the dataset's execution.
+* **Runtime metrics**: These metrics are a high-level breakdown of the runtime of the dataset execution. These stats are a per
+  operator summary of the time each operator took to complete and the fraction of the total execution time that the operator took
+  to complete. As there are potentially multiple concurrent operators, these percentages don't necessarily sum to 100%. Instead,
+  they show how long running each of the operators is in the context of the full dataset execution.
+
+Example stats
+~~~~~~~~~~~~~
+As a concrete example, below is a stats output from :doc:`Image Classification Batch Inference with PyTorch ResNet18 </data/examples/pytorch_resnet_batch_prediction>`:
+
+.. code-block:: text
+
+   Operator 1 ReadImage->Map(preprocess_image): 384 tasks executed, 386 blocks produced in 9.21s
+   * Remote wall time: 33.55ms min, 2.22s max, 1.03s mean, 395.65s total
+   * Remote cpu time: 34.93ms min, 3.36s max, 1.64s mean, 632.26s total
+   * UDF time: 535.1ms min, 2.16s max, 975.7ms mean, 376.62s total
+   * Peak heap memory usage (MiB): 556.32 min, 1126.95 max, 655 mean
+   * Output num rows per block: 4 min, 25 max, 24 mean, 9469 total
+   * Output size bytes per block: 6060399 min, 105223020 max, 31525416 mean, 12168810909 total
+   * Output rows per task: 24 min, 25 max, 24 mean, 384 tasks used
+   * Tasks per node: 32 min, 64 max, 48 mean; 8 nodes used
+   * Operator throughput:
+         * Ray Data throughput: 1028.5218637702708 rows/s
+         * Estimated single node throughput: 23.932674100499128 rows/s
+
+   Operator 2 MapBatches(ResnetModel): 14 tasks executed, 48 blocks produced in 27.43s
+   * Remote wall time: 523.93us min, 7.01s max, 1.82s mean, 87.18s total
+   * Remote cpu time: 523.23us min, 6.23s max, 1.76s mean, 84.61s total
+   * UDF time: 4.49s min, 17.81s max, 10.52s mean, 505.08s total
+   * Peak heap memory usage (MiB): 4025.42 min, 7920.44 max, 5803 mean
+   * Output num rows per block: 84 min, 334 max, 197 mean, 9469 total
+   * Output size bytes per block: 72317976 min, 215806447 max, 134739694 mean, 6467505318 total
+   * Output rows per task: 319 min, 720 max, 676 mean, 14 tasks used
+   * Tasks per node: 3 min, 4 max, 3 mean; 4 nodes used
+   * Operator throughput:
+         * Ray Data throughput: 345.1533728632648 rows/s
+         * Estimated single node throughput: 108.62003864820711 rows/s
+
+   Dataset iterator time breakdown:
+   * Total time overall: 38.53s
+      * Total time in Ray Data iterator initialization code: 16.86s
+      * Total time user thread is blocked by Ray Data iter_batches: 19.76s
+      * Total execution time for user thread: 1.9s
+   * Batch iteration time breakdown (summed across prefetch threads):
+      * In ray.get(): 70.49ms min, 2.16s max, 272.8ms avg, 13.09s total
+      * In batch creation: 3.6us min, 5.95us max, 4.26us avg, 204.41us total
+      * In batch formatting: 4.81us min, 7.88us max, 5.5us avg, 263.94us total
+
+   Dataset throughput:
+         * Ray Data throughput: 1026.5318925757008 rows/s
+         * Estimated single node throughput: 19.611578909587674 rows/s
+
+For the same example with verbosity enabled, the stats output is:
+
+.. code-block:: text
+
+   Operator 1 ReadImage->Map(preprocess_image): 384 tasks executed, 387 blocks produced in 9.49s
+   * Remote wall time: 22.81ms min, 2.5s max, 999.95ms mean, 386.98s total
+   * Remote cpu time: 24.06ms min, 3.36s max, 1.63s mean, 629.93s total
+   * UDF time: 552.79ms min, 2.41s max, 956.84ms mean, 370.3s total
+   * Peak heap memory usage (MiB): 550.95 min, 1186.28 max, 651 mean
+   * Output num rows per block: 4 min, 25 max, 24 mean, 9469 total
+   * Output size bytes per block: 4444092 min, 105223020 max, 31443955 mean, 12168810909 total
+   * Output rows per task: 24 min, 25 max, 24 mean, 384 tasks used
+   * Tasks per node: 39 min, 60 max, 48 mean; 8 nodes used
+   * Operator throughput:
+         * Ray Data throughput: 997.9207015895857 rows/s
+         * Estimated single node throughput: 24.46899945870273 rows/s
+   * Extra metrics: {'num_inputs_received': 384, 'bytes_inputs_received': 1104723940, 'num_task_inputs_processed': 384, 'bytes_task_inputs_processed': 1104723940, 'bytes_inputs_of_submitted_tasks': 1104723940, 'num_task_outputs_generated': 387, 'bytes_task_outputs_generated': 12168810909, 'rows_task_outputs_generated': 9469, 'num_outputs_taken': 387, 'bytes_outputs_taken': 12168810909, 'num_outputs_of_finished_tasks': 387, 'bytes_outputs_of_finished_tasks': 12168810909, 'num_tasks_submitted': 384, 'num_tasks_running': 0, 'num_tasks_have_outputs': 384, 'num_tasks_finished': 384, 'num_tasks_failed': 0, 'block_generation_time': 386.97945193799995, 'task_submission_backpressure_time': 7.263684450000142, 'obj_store_mem_internal_inqueue_blocks': 0, 'obj_store_mem_internal_inqueue': 0, 'obj_store_mem_internal_outqueue_blocks': 0, 'obj_store_mem_internal_outqueue': 0, 'obj_store_mem_pending_task_inputs': 0, 'obj_store_mem_freed': 1104723940, 'obj_store_mem_spilled': 0, 'obj_store_mem_used': 12582535566, 'cpu_usage': 0, 'gpu_usage': 0, 'ray_remote_args': {'num_cpus': 1, 'scheduling_strategy': 'SPREAD'}}
+
+   Operator 2 MapBatches(ResnetModel): 14 tasks executed, 48 blocks produced in 28.81s
+   * Remote wall time: 134.84us min, 7.23s max, 1.82s mean, 87.16s total
+   * Remote cpu time: 133.78us min, 6.28s max, 1.75s mean, 83.98s total
+   * UDF time: 4.56s min, 17.78s max, 10.28s mean, 493.48s total
+   * Peak heap memory usage (MiB): 3925.88 min, 7713.01 max, 5688 mean
+   * Output num rows per block: 125 min, 259 max, 197 mean, 9469 total
+   * Output size bytes per block: 75531617 min, 187889580 max, 134739694 mean, 6467505318 total
+   * Output rows per task: 325 min, 719 max, 676 mean, 14 tasks used
+   * Tasks per node: 3 min, 4 max, 3 mean; 4 nodes used
+   * Operator throughput:
+         * Ray Data throughput: 328.71474145609153 rows/s
+         * Estimated single node throughput: 108.6352856660782 rows/s
+   * Extra metrics: {'num_inputs_received': 387, 'bytes_inputs_received': 12168810909, 'num_task_inputs_processed': 0, 'bytes_task_inputs_processed': 0, 'bytes_inputs_of_submitted_tasks': 12168810909, 'num_task_outputs_generated': 1, 'bytes_task_outputs_generated': 135681874, 'rows_task_outputs_generated': 252, 'num_outputs_taken': 1, 'bytes_outputs_taken': 135681874, 'num_outputs_of_finished_tasks': 0, 'bytes_outputs_of_finished_tasks': 0, 'num_tasks_submitted': 14, 'num_tasks_running': 14, 'num_tasks_have_outputs': 1, 'num_tasks_finished': 0, 'num_tasks_failed': 0, 'block_generation_time': 7.229860895999991, 'task_submission_backpressure_time': 0, 'obj_store_mem_internal_inqueue_blocks': 13, 'obj_store_mem_internal_inqueue': 413724657, 'obj_store_mem_internal_outqueue_blocks': 0, 'obj_store_mem_internal_outqueue': 0, 'obj_store_mem_pending_task_inputs': 12168810909, 'obj_store_mem_freed': 0, 'obj_store_mem_spilled': 0, 'obj_store_mem_used': 1221136866.0, 'cpu_usage': 0, 'gpu_usage': 4}
+
+   Dataset iterator time breakdown:
+   * Total time overall: 42.29s
+      * Total time in Ray Data iterator initialization code: 20.24s
+      * Total time user thread is blocked by Ray Data iter_batches: 19.96s
+      * Total execution time for user thread: 2.08s
+   * Batch iteration time breakdown (summed across prefetch threads):
+      * In ray.get(): 73.0ms min, 2.15s max, 246.3ms avg, 11.82s total
+      * In batch creation: 3.62us min, 6.6us max, 4.39us avg, 210.7us total
+      * In batch formatting: 4.75us min, 8.67us max, 5.52us avg, 264.98us total
+
+   Dataset throughput:
+         * Ray Data throughput: 468.11051989434594 rows/s
+         * Estimated single node throughput: 972.8197093015862 rows/s
+
+   Runtime Metrics:
+   * ReadImage->Map(preprocess_image): 9.49s (46.909%)
+   * MapBatches(ResnetModel): 28.81s (142.406%)
+   * Scheduling: 6.16s (30.448%)
+   * Total: 20.23s (100.000%)

doc/source/data/performance-tips.rst

@@ -0,0 +1,363 @@
+.. _data_performance_tips:
+
+Advanced: Performance Tips and Tuning
+=====================================
+
+Optimizing transforms
+---------------------
+
+Batching transforms
+~~~~~~~~~~~~~~~~~~~
+
+If your transformation is vectorized like most NumPy or pandas operations, use
+:meth:`~ray.data.Dataset.map_batches` rather than :meth:`~ray.data.Dataset.map`. It's
+faster.
+
+If your transformation isn't vectorized, there's no performance benefit.
+
+Enabling Polars for sort operations
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+You can speed up :meth:`~ray.data.Dataset.sort` and operations that sort internally,
+such as :meth:`~ray.data.grouped_data.GroupedData.map_groups`, by enabling Polars:
+
+.. testcode::
+
+    import ray
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.use_polars_sort = True
+
+When you enable this flag, Ray Data uses Polars instead of PyArrow for the internal
+sorting step, which can improve performance for large tabular datasets.
+This flag doesn't affect other operations such as :meth:`~ray.data.Dataset.map_batches`.
+
+Optimizing reads
+----------------
+
+.. _read_output_blocks:
+
+Tuning output blocks for read
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+By default, Ray Data automatically selects the number of output blocks for read according to the following procedure:
+
+- The ``override_num_blocks`` parameter passed to Ray Data's :ref:`read APIs <loading-data-api>` specifies the number of output blocks, which is equivalent to the number of read tasks to create.
+- Usually, if the read is followed by a :func:`~ray.data.Dataset.map` or :func:`~ray.data.Dataset.map_batches`, the map is fused with the read; therefore ``override_num_blocks`` also determines the number of map tasks.
+
+Ray Data decides the default value for number of output blocks based on the following heuristics, applied in order:
+
+1. Start with the default value of 200. You can overwrite this by setting :class:`DataContext.read_op_min_num_blocks <ray.data.context.DataContext>`.
+2. Min block size (default=1 MiB). If number of blocks would make blocks smaller than this threshold, reduce number of blocks to avoid the overhead of tiny blocks. You can override by setting :class:`DataContext.target_min_block_size <ray.data.context.DataContext>` (bytes).
+3. Max block size (default=128 MiB). If number of blocks would make blocks larger than this threshold, increase number of blocks to avoid out-of-memory errors during processing. You can override by setting :class:`DataContext.target_max_block_size <ray.data.context.DataContext>` (bytes).
+4. Available CPUs. Increase number of blocks to utilize all of the available CPUs in the cluster. Ray Data chooses the number of read tasks to be at least 2x the number of available CPUs.
+
+Occasionally, it's advantageous to manually tune the number of blocks to optimize the application.
+For example, the following code batches multiple files into the same read task to avoid creating blocks that are too large.
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+    # Pretend there are two CPUs.
+    ray.init(num_cpus=2)
+
+    # Repeat the iris.csv file 16 times.
+    ds = ray.data.read_csv(["s3://anonymous@ray-example-data/iris.csv"] * 16)
+    print(ds.materialize())
+
+.. testoutput::
+    :options: +MOCK
+
+    MaterializedDataset(
+       num_blocks=4,
+       num_rows=2400,
+       ...
+    )
+
+But suppose that you knew that you wanted to read all 16 files in parallel.
+This could be, for example, because you know that additional CPUs should get added to the cluster by the autoscaler or because you want the downstream operator to transform each file's contents in parallel.
+You can get this behavior by setting the ``override_num_blocks`` parameter.
+Notice how the number of output blocks is equal to ``override_num_blocks`` in the following code:
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+    # Pretend there are two CPUs.
+    ray.init(num_cpus=2)
+
+    # Repeat the iris.csv file 16 times.
+    ds = ray.data.read_csv(["s3://anonymous@ray-example-data/iris.csv"] * 16, override_num_blocks=16)
+    print(ds.materialize())
+
+.. testoutput::
+    :options: +MOCK
+
+    MaterializedDataset(
+       num_blocks=16,
+       num_rows=2400,
+       ...
+    )
+
+
+When using the default auto-detected number of blocks, Ray Data attempts to cap each task's output to :class:`DataContext.target_max_block_size <ray.data.context.DataContext>` many bytes.
+Note however that Ray Data can't perfectly predict the size of each task's output, so it's possible that each task produces one or more output blocks.
+Thus, the total blocks in the final :class:`~ray.data.Dataset` may differ from the specified ``override_num_blocks``.
+Here's an example where we manually specify ``override_num_blocks=1``, but the one task still produces multiple blocks in the materialized Dataset:
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+    # Pretend there are two CPUs.
+    ray.init(num_cpus=2)
+
+    # Generate ~400MB of data.
+    ds = ray.data.range_tensor(5_000, shape=(10_000, ), override_num_blocks=1)
+    print(ds.materialize())
+
+.. testoutput::
+    :options: +MOCK
+
+    MaterializedDataset(
+       num_blocks=3,
+       num_rows=5000,
+       schema={data: ArrowTensorTypeV2(shape=(10000,), dtype=int64)}
+    )
+
+
+Currently, Ray Data can assign at most one read task per input file.
+Thus, if the number of input files is smaller than ``override_num_blocks``, the number of read tasks is capped to the number of input files.
+To ensure that downstream transforms can still execute with the desired number of blocks, Ray Data splits the read tasks' outputs into a total of ``override_num_blocks`` blocks and prevents fusion with the downstream transform.
+In other words, each read task's output blocks are materialized to Ray's object store before the consuming map task executes.
+For example, the following code executes :func:`~ray.data.read_csv` with only one task, but its output is split into 4 blocks before executing the :func:`~ray.data.Dataset.map`:
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+    # Pretend there are two CPUs.
+    ray.init(num_cpus=2)
+
+    ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv").map(lambda row: row)
+    print(ds.materialize().stats())
+
+.. testoutput::
+    :options: +MOCK
+
+    ...
+    Operator 1 ReadCSV->SplitBlocks(4): 1 tasks executed, 4 blocks produced in 0.01s
+    ...
+
+    Operator 2 Map(<lambda>): 4 tasks executed, 4 blocks produced in 0.3s
+    ...
+
+To turn off this behavior and allow the read and map operators to be fused, set ``override_num_blocks`` manually.
+For example, this code sets the number of files equal to ``override_num_blocks``:
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+    # Pretend there are two CPUs.
+    ray.init(num_cpus=2)
+
+    ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv", override_num_blocks=1).map(lambda row: row)
+    print(ds.materialize().stats())
+
+.. testoutput::
+    :options: +MOCK
+
+    ...
+    Operator 1 ReadCSV->Map(<lambda>): 1 tasks executed, 1 blocks produced in 0.01s
+    ...
+
+
+.. _tuning_read_resources:
+
+Tuning read resources
+~~~~~~~~~~~~~~~~~~~~~
+
+By default, Ray requests 1 CPU per read task, which means one read task per CPU can execute concurrently.
+For datasources that benefit from more IO parallelism, you can specify a lower ``num_cpus`` value for the read function with the ``ray_remote_args`` parameter.
+For example, use ``ray.data.read_parquet(path, ray_remote_args={"num_cpus": 0.25})`` to allow up to four read tasks per CPU.
+
+.. _parquet_column_pruning:
+
+Parquet column pruning (projection pushdown)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+By default, :func:`ray.data.read_parquet` reads all columns in the Parquet files into memory.
+If you only need a subset of the columns, make sure to specify the list of columns
+explicitly when calling :func:`ray.data.read_parquet` to
+avoid loading unnecessary data (projection pushdown). Note that this is more efficient than
+calling :func:`~ray.data.Dataset.select_columns`, since column selection is pushed down to the file scan.
+
+.. testcode::
+
+    import ray
+
+    # Read just two of the five columns of the Iris dataset.
+    ds = ray.data.read_parquet(
+        "s3://anonymous@ray-example-data/iris.parquet",
+    ).select_columns(["sepal.length", "variety"])
+    
+    print(ds.schema())
+
+.. testoutput::
+
+    Column        Type
+    ------        ----
+    sepal.length  double
+    variety       string
+
+
+.. _data_memory:
+
+Reducing memory usage
+---------------------
+
+Avoiding object spilling
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+A Dataset's intermediate and output blocks are stored in Ray's object store.
+Although Ray Data attempts to minimize object store usage with :ref:`streaming execution <streaming_execution>`, it's still possible that the working set exceeds the object store capacity.
+In this case, Ray begins spilling blocks to disk, which can slow down execution significantly or even cause out-of-disk errors.
+
+There are some cases where spilling is expected. In particular, if the total Dataset's size is larger than object store capacity, and one of the following is true:
+
+1. An :ref:`all-to-all shuffle operation <optimizing_shuffles>` is used. Or,
+2. There is a call to :meth:`ds.materialize() <ray.data.Dataset.materialize>`.
+
+Otherwise, it's best to tune your application to avoid spilling.
+The recommended strategy is to manually increase the :ref:`read output blocks <read_output_blocks>` or modify your application code to ensure that each task reads a smaller amount of data.
+
+.. note:: This is an active area of development. If your Dataset is causing spilling and you don't know why, `file a Ray Data issue on GitHub`_.
+
+Handling too-small blocks
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When different operators of your Dataset produce different-sized outputs, you may end up with very small blocks, which can hurt performance and even cause crashes from excessive metadata.
+Use :meth:`ds.stats() <ray.data.Dataset.stats>` to check that each operator's output blocks are each at least 1 MB and ideally >100 MB.
+
+If your blocks are smaller than this, consider repartitioning into larger blocks.
+There are two ways to do this:
+
+1. If you need control over the exact number of output blocks, use :meth:`ds.repartition(num_partitions) <ray.data.Dataset.repartition>`. Note that this is an :ref:`all-to-all operation <optimizing_shuffles>` and it materializes all blocks into memory before performing the repartition.
+2. If you don't need control over the exact number of output blocks and just want to produce larger blocks, use :meth:`ds.map_batches(lambda batch: batch, batch_size=batch_size) <ray.data.Dataset.map_batches>` and set ``batch_size`` to the desired number of rows per block. This is executed in a streaming fashion and avoids materialization.
+
+When :meth:`ds.map_batches() <ray.data.Dataset.map_batches>` is used, Ray Data coalesces blocks so that each map task can process at least this many rows.
+Note that the chosen ``batch_size`` is a lower bound on the task's input block size but it doesn't necessarily determine the task's final *output* block size.
+
+To illustrate these, the following code uses both strategies to coalesce the 10 tiny blocks with 1 row each into 1 larger block with 10 rows:
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+    # Pretend there are two CPUs.
+    ray.init(num_cpus=2)
+
+    # 1. Use ds.repartition().
+    ds = ray.data.range(10, override_num_blocks=10).repartition(1)
+    print(ds.materialize().stats())
+
+    # 2. Use ds.map_batches().
+    ds = ray.data.range(10, override_num_blocks=10).map_batches(lambda batch: batch, batch_size=10)
+    print(ds.materialize().stats())
+
+.. testoutput::
+    :options: +MOCK
+
+    # 1. ds.repartition() output.
+    Operator 1 ReadRange: 10 tasks executed, 10 blocks produced in 0.33s
+    ...
+    * Output num rows: 1 min, 1 max, 1 mean, 10 total
+    ...
+    Operator 2 Repartition: executed in 0.36s
+
+            Suboperator 0 RepartitionSplit: 10 tasks executed, 10 blocks produced
+            ...
+
+            Suboperator 1 RepartitionReduce: 1 tasks executed, 1 blocks produced
+            ...
+            * Output num rows: 10 min, 10 max, 10 mean, 10 total
+            ...
+
+
+    # 2. ds.map_batches() output.
+    Operator 1 ReadRange->MapBatches(<lambda>): 1 tasks executed, 1 blocks produced in 0s
+    ...
+    * Output num rows: 10 min, 10 max, 10 mean, 10 total
+
+Configuring execution
+---------------------
+
+Configuring resources and locality
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+By default, the CPU and GPU limits are set to the cluster size, and the object store memory limit conservatively to 1/4 of the total object store size to avoid the possibility of disk spilling.
+
+You may want to customize these limits in the following scenarios:
+
+- If running multiple concurrent jobs on the cluster, setting lower limits can avoid resource contention between the jobs.
+- If you want to fine-tune the memory limit to maximize performance.
+- For data loading into training jobs, you may want to set the object store memory to a low value (for example, 2 GB) to limit resource usage.
+
+You can configure execution options with the global DataContext. The options are applied for future jobs launched in the process:
+
+.. code-block::
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.execution_options.resource_limits = ctx.execution_options.resource_limits.copy(
+        cpu=10,
+        gpu=5,
+        object_store_memory=10e9,
+    )
+
+Reproducibility
+---------------
+
+Deterministic execution
+~~~~~~~~~~~~~~~~~~~~~~~
+
+.. code-block::
+
+   # By default, this is set to False.
+   ctx.execution_options.preserve_order = True
+
+To enable deterministic execution, set the preceding to True. This setting may decrease performance, but ensures block ordering is preserved through execution. This flag defaults to False.
+
+
+.. _`file a Ray Data issue on GitHub`: https://github.com/ray-project/ray/issues/new?assignees=&labels=bug%2Ctriage%2Cdata&projects=&template=bug-report.yml&title=[data]+

doc/source/data/quickstart.rst

@@ -0,0 +1,157 @@
+.. _data_quickstart:
+
+Ray Data Quickstart
+===================
+
+Get started with Ray Data's :class:`Dataset <ray.data.Dataset>` abstraction for distributed data processing.
+
+This guide introduces you to the core capabilities of Ray Data:
+
+* :ref:`Loading data <loading_key_concept>`
+* :ref:`Transforming data <transforming_key_concept>`
+* :ref:`Consuming data <consuming_key_concept>`
+* :ref:`Saving data <saving_key_concept>`
+
+Datasets
+--------
+
+Ray Data's main abstraction is a :class:`Dataset <ray.data.Dataset>`, which
+represents a distributed collection of data. Datasets are specifically designed for machine learning workloads
+and can efficiently handle data collections that exceed a single machine's memory.
+
+.. _loading_key_concept:
+
+Loading data
+------------
+
+Create datasets from various sources including local files, Python objects, and cloud storage services like S3 or GCS.
+Ray Data seamlessly integrates with any `filesystem supported by Arrow
+<http://arrow.apache.org/docs/python/generated/pyarrow.fs.FileSystem.html>`__.
+
+.. testcode::
+
+    import ray
+
+    # Load a CSV dataset directly from S3
+    ds = ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+    
+    # Preview the first record
+    ds.show(limit=1)
+
+.. testoutput::
+
+    {'sepal length (cm)': 5.1, 'sepal width (cm)': 3.5, 'petal length (cm)': 1.4, 'petal width (cm)': 0.2, 'target': 0}
+
+To learn more about creating datasets from different sources, read :ref:`Loading data <loading_data>`.
+
+.. _transforming_key_concept:
+
+Transforming data
+-----------------
+
+Apply user-defined functions (UDFs) to transform datasets. Ray automatically parallelizes these transformations
+across your cluster for better performance.
+
+.. testcode::
+
+    from typing import Dict
+    import numpy as np
+
+    # Define a transformation to compute a "petal area" attribute
+    def transform_batch(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        vec_a = batch["petal length (cm)"]
+        vec_b = batch["petal width (cm)"]
+        batch["petal area (cm^2)"] = np.round(vec_a * vec_b, 2)
+        return batch
+
+    # Apply the transformation to our dataset
+    transformed_ds = ds.map_batches(transform_batch, batch_size="auto")
+    
+    # View the updated schema with the new column
+    # .materialize() will execute all the lazy transformations and
+    # materialize the dataset into object store memory
+    print(transformed_ds.materialize())
+
+.. testoutput::
+
+    shape: (150, 6)
+    ╭───────────────────┬──────────────────┬───────────────────┬──────────────────┬────────┬───────────────────╮
+    │ sepal length (cm) ┆ sepal width (cm) ┆ petal length (cm) ┆ petal width (cm) ┆ target ┆ petal area (cm^2) │
+    │ ---               ┆ ---              ┆ ---               ┆ ---              ┆ ---    ┆ ---               │
+    │ double            ┆ double           ┆ double            ┆ double           ┆ int64  ┆ double            │
+    ╞═══════════════════╪══════════════════╪═══════════════════╪══════════════════╪════════╪═══════════════════╡
+    │ 5.1               ┆ 3.5              ┆ 1.4               ┆ 0.2              ┆ 0      ┆ 0.28              │
+    │ 4.9               ┆ 3.0              ┆ 1.4               ┆ 0.2              ┆ 0      ┆ 0.28              │
+    │ 4.7               ┆ 3.2              ┆ 1.3               ┆ 0.2              ┆ 0      ┆ 0.26              │
+    │ 4.6               ┆ 3.1              ┆ 1.5               ┆ 0.2              ┆ 0      ┆ 0.3               │
+    │ 5.0               ┆ 3.6              ┆ 1.4               ┆ 0.2              ┆ 0      ┆ 0.28              │
+    │ …                 ┆ …                ┆ …                 ┆ …                ┆ …      ┆ …                 │
+    │ 6.7               ┆ 3.0              ┆ 5.2               ┆ 2.3              ┆ 2      ┆ 11.96             │
+    │ 6.3               ┆ 2.5              ┆ 5.0               ┆ 1.9              ┆ 2      ┆ 9.5               │
+    │ 6.5               ┆ 3.0              ┆ 5.2               ┆ 2.0              ┆ 2      ┆ 10.4              │
+    │ 6.2               ┆ 3.4              ┆ 5.4               ┆ 2.3              ┆ 2      ┆ 12.42             │
+    │ 5.9               ┆ 3.0              ┆ 5.1               ┆ 1.8              ┆ 2      ┆ 9.18              │
+    ╰───────────────────┴──────────────────┴───────────────────┴──────────────────┴────────┴───────────────────╯
+    (Showing 10 of 150 rows)
+
+To explore more transformation capabilities, read :ref:`Transforming data <transforming_data>`.
+
+.. _consuming_key_concept:
+
+Consuming data
+--------------
+
+Access dataset contents through convenient methods like :meth:`~ray.data.Dataset.take_batch` and 
+:meth:`~ray.data.Dataset.iter_batches`. You can also pass datasets directly to Ray Tasks or Actors
+for distributed processing.
+
+.. testcode::
+
+    # Extract the first 3 rows as a batch for processing
+    print(transformed_ds.take_batch(batch_size=3))
+
+.. testoutput::
+    :options: +NORMALIZE_WHITESPACE
+
+    {'sepal length (cm)': array([5.1, 4.9, 4.7]),
+        'sepal width (cm)': array([3.5, 3. , 3.2]),
+        'petal length (cm)': array([1.4, 1.4, 1.3]),
+        'petal width (cm)': array([0.2, 0.2, 0.2]),
+        'target': array([0, 0, 0]),
+        'petal area (cm^2)': array([0.28, 0.28, 0.26])}
+
+For more details on working with dataset contents, see
+:ref:`Iterating over Data <iterating-over-data>` and :ref:`Saving Data <saving-data>`.
+
+.. _saving_key_concept:
+
+Saving data
+-----------
+
+Export processed datasets to a variety of formats and storage locations using methods
+like :meth:`~ray.data.Dataset.write_parquet`, :meth:`~ray.data.Dataset.write_csv`, and more.
+
+.. testcode::
+    :hide:
+
+    # The number of blocks can be non-deterministic. Repartition the dataset beforehand
+    # so that the number of written files is consistent.
+    transformed_ds = transformed_ds.repartition(2)
+
+.. testcode::
+
+    import os
+
+    # Save the transformed dataset as Parquet files
+    transformed_ds.write_parquet("/tmp/iris")
+
+    # Verify the files were created
+    print(os.listdir("/tmp/iris"))
+
+.. testoutput::
+    :options: +MOCK
+
+    ['..._000000.parquet', '..._000001.parquet']
+
+
+For more information on saving datasets, see :ref:`Saving data <saving-data>`.

doc/source/data/saving-data.rst

@@ -0,0 +1,406 @@
+.. _saving-data:
+
+===========
+Saving Data
+===========
+
+Ray Data lets you save data in files or other Python objects.
+
+This guide shows you how to:
+
+* `Write data to files <#writing-data-to-files>`_
+* `Convert Datasets to other Python libraries <#converting-datasets-to-other-python-libraries>`_
+
+Writing data to files
+=====================
+
+Ray Data writes to local disk and cloud storage.
+
+Writing data to local disk
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To save your :class:`~ray.data.dataset.Dataset` to local disk, call a method
+like :meth:`Dataset.write_parquet <ray.data.Dataset.write_parquet>`  and specify a local
+directory with the `local://` scheme.
+
+.. warning::
+
+    If your cluster contains multiple nodes and you don't use `local://`, Ray Data
+    writes different partitions of data to different nodes.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+    ds.write_parquet("local:///tmp/iris/")
+
+To write data to formats other than Parquet, see the
+:ref:`Saving Data API <saving-data-api>`.
+
+Writing data to cloud storage
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To save your :class:`~ray.data.dataset.Dataset` to cloud storage, authenticate all nodes
+with your cloud service provider. Then, call a method like
+:meth:`Dataset.write_parquet <ray.data.Dataset.write_parquet>` and specify a URI with
+the appropriate scheme. URI can point to buckets or folders.
+
+To write data to formats other than Parquet, see the :ref:`Saving Data API <saving-data-api>`.
+
+.. tab-set::
+
+    .. tab-item:: S3
+
+        To save data to Amazon S3, specify a URI with the ``s3://`` scheme.
+
+        .. testcode::
+            :skipif: True
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            ds.write_parquet("s3://my-bucket/my-folder")
+
+        Ray Data relies on PyArrow to authenticate with Amazon S3. For more on how to configure
+        your credentials to be compatible with PyArrow, see their
+        `S3 Filesystem docs <https://arrow.apache.org/docs/python/filesystems.html#s3>`_.
+
+    .. tab-item:: GCS
+
+        To save data to Google Cloud Storage, install the
+        `Filesystem interface to Google Cloud Storage <https://gcsfs.readthedocs.io/en/latest/>`_
+
+        .. code-block:: console
+
+            pip install gcsfs
+
+        Then, create a ``GCSFileSystem`` and specify a URI with the ``gcs://`` scheme.
+
+        .. testcode::
+            :skipif: True
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            filesystem = gcsfs.GCSFileSystem(project="my-google-project")
+            ds.write_parquet("gcs://my-bucket/my-folder", filesystem=filesystem)
+
+        Ray Data relies on PyArrow for authentication with Google Cloud Storage. For more on how
+        to configure your credentials to be compatible with PyArrow, see their
+        `GCS Filesystem docs <https://arrow.apache.org/docs/python/filesystems.html#google-cloud-storage-file-system>`_.
+
+    .. tab-item:: ABS
+
+        To save data to Azure Blob Storage, install the
+        `Filesystem interface to Azure-Datalake Gen1 and Gen2 Storage <https://pypi.org/project/adlfs/>`_
+
+        .. code-block:: console
+
+            pip install adlfs
+
+        Then, create a ``AzureBlobFileSystem`` and specify a URI with the ``az://`` scheme.
+
+        .. testcode::
+            :skipif: True
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            filesystem = adlfs.AzureBlobFileSystem(account_name="azureopendatastorage")
+            ds.write_parquet("az://my-bucket/my-folder", filesystem=filesystem)
+
+        Ray Data relies on PyArrow for authentication with Azure Blob Storage. For more on how
+        to configure your credentials to be compatible with PyArrow, see their
+        `fsspec-compatible filesystems docs <https://arrow.apache.org/docs/python/filesystems.html#using-fsspec-compatible-filesystems-with-arrow>`_.
+
+Writing data to NFS
+~~~~~~~~~~~~~~~~~~~
+
+To save your :class:`~ray.data.dataset.Dataset` to NFS file systems, call a method
+like :meth:`Dataset.write_parquet <ray.data.Dataset.write_parquet>` and specify a
+mounted directory.
+
+.. testcode::
+    :skipif: True
+
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+    ds.write_parquet("/mnt/cluster_storage/iris")
+
+To write data to formats other than Parquet, see the
+:ref:`Saving Data API <saving-data-api>`.
+
+.. _changing-number-output-files:
+
+Changing the number of output files
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When you call a write method, Ray Data writes your data to several files. To control the
+number of output files, configure ``min_rows_per_file``.
+
+.. note::
+
+    ``min_rows_per_file`` is a hint, not a strict limit. Ray Data might write more or
+    fewer rows to each file. Under the hood, if the number of rows per block is
+    larger than the specified value, Ray Data writes
+    the number of rows per block to each file.
+
+
+.. testcode::
+
+    import os
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+    ds.write_csv("/tmp/few_files/", min_rows_per_file=75)
+
+    print(os.listdir("/tmp/few_files/"))
+
+.. testoutput::
+    :options: +MOCK
+
+    ['0_000001_000000.csv', '0_000000_000000.csv', '0_000002_000000.csv']
+
+
+Writing into Partitioned Dataset
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+When writing partitioned dataset (using Hive-style, folder-based partitioning) it's recommended to repartition the dataset by the partition columns prior to writing into it. 
+This allows you to *have control over the file sizes and their number*. When the dataset is repartitioned by the partition columns every block should contain all of the rows corresponding to particular partition, 
+meaning that the number of files created should be controlled based on the configuration provided to, 
+for example, `write_parquet` method (such as `min_rows_per_file`, `max_rows_per_file`). 
+Since every block is written out independently, when writing the dataset without prior 
+repartitioning you could potentially get an N number of files per partition 
+(where N is the number of blocks in your dataset) with very limited ability to control the 
+number of files & their sizes (since every block could potentially carry the rows corresponding to any partition).
+
+.. testcode::
+    import ray
+    import pandas as pd
+    from ray.data import DataContext
+    from ray.data.context import ShuffleStrategy
+
+    def print_directory_tree(start_path: str) -> None:
+        """
+        Prints the directory tree structure starting from the given path.
+        """
+        for root, dirs, files in os.walk(start_path):
+            level = root.replace(start_path, '').count(os.sep)
+            indent = ' ' * 4 * (level)
+            print(f'{indent}{os.path.basename(root)}/')
+            subindent = ' ' * 4 * (level + 1)
+            for f in files:
+                print(f'{subindent}{f}')
+
+    # Sample dataset that we’ll partition by ``city`` and ``year``.
+    df = pd.DataFrame(
+        {
+            "city": ["SF", "SF", "NYC", "NYC", "SF", "NYC", "SF", "NYC"],
+            "year": [2023, 2024, 2023, 2024, 2023, 2023, 2024, 2024],
+            "sales": [100, 120, 90, 115, 105, 95, 130, 110],
+        }
+    )
+
+    ds = ray.data.from_pandas(df)
+    DataContext.shuffle_strategy=ShuffleStrategy.HASH_SHUFFLE
+
+    # ── Partitioned write ──────────────────────────────────────────────────────
+    # 1. Repartition so all rows with the same (city, year) land in the same
+    #    block – this minimises shuffling during the write.
+    # 2. Pass the same columns to ``partition_cols`` so Ray creates a
+    #    Hive-style directory layout:  city=<value>/year=<value>/....
+    # 3. Use ``min_rows_per_file`` / ``max_rows_per_file`` to control how many
+    #    rows Ray puts in each Parquet file.
+    ds.repartition(keys=["city", "year"], num_blocks=4).write_parquet(
+        "/tmp/sales_partitioned",
+        partition_cols=["city", "year"],
+        min_rows_per_file=2,     # At least 2 rows in each file …
+        max_rows_per_file=3,     # … but never more than 3.
+    )
+
+    print_directory_tree("/tmp/sales_partitioned")
+
+.. testoutput::
+    :options: +MOCK
+
+    sales_partitioned/
+        city=NYC/
+            year=2024/
+                1_a2b8b82cd2904a368ec39f42ae3cf830_000000_000000-0.parquet
+            year=2023/
+                1_a2b8b82cd2904a368ec39f42ae3cf830_000001_000000-0.parquet
+        city=SF/
+            year=2024/
+                1_a2b8b82cd2904a368ec39f42ae3cf830_000000_000000-0.parquet
+            year=2023/
+                1_a2b8b82cd2904a368ec39f42ae3cf830_000001_000000-0.parquet
+
+
+Converting Datasets to other Python libraries
+=============================================
+
+Converting Datasets to pandas
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To convert a :class:`~ray.data.dataset.Dataset` to a pandas DataFrame, call
+:meth:`Dataset.to_pandas() <ray.data.Dataset.to_pandas>`. Your data must fit in memory
+on the head node.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+    df = ds.to_pandas()
+    print(df)
+
+.. testoutput::
+    :options: +NORMALIZE_WHITESPACE
+
+         sepal length (cm)  sepal width (cm)  ...  petal width (cm)  target
+    0                  5.1               3.5  ...               0.2       0
+    1                  4.9               3.0  ...               0.2       0
+    2                  4.7               3.2  ...               0.2       0
+    3                  4.6               3.1  ...               0.2       0
+    4                  5.0               3.6  ...               0.2       0
+    ..                 ...               ...  ...               ...     ...
+    145                6.7               3.0  ...               2.3       2
+    146                6.3               2.5  ...               1.9       2
+    147                6.5               3.0  ...               2.0       2
+    148                6.2               3.4  ...               2.3       2
+    149                5.9               3.0  ...               1.8       2
+    <BLANKLINE>
+    [150 rows x 5 columns]
+
+Converting Datasets to distributed DataFrames
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Ray Data interoperates with distributed data processing frameworks like `Daft <https://www.daft.ai>`_,
+:ref:`Dask <dask-on-ray>`, :ref:`Spark <spark-on-ray>`, :ref:`Modin <modin-on-ray>`, and
+:ref:`Mars <mars-on-ray>`.
+
+.. tab-set::
+
+    .. tab-item:: Daft
+
+        To convert a :class:`~ray.data.dataset.Dataset` to a `Daft Dataframe <https://docs.daft.ai/en/stable/api/dataframe/>`_, call
+        :meth:`Dataset.to_daft() <ray.data.Dataset.to_daft>`.
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            df = ds.to_daft()
+            print(df)
+
+        .. testoutput::
+            :options: +MOCK
+
+            ╭───────────────────┬──────────────────┬───────────────────┬──────────────────┬────────╮
+            │ sepal length (cm) ┆ sepal width (cm) ┆ petal length (cm) ┆ petal width (cm) ┆ target │
+            │ ---               ┆ ---              ┆ ---               ┆ ---              ┆ ---    │
+            │ Float64           ┆ Float64          ┆ Float64           ┆ Float64          ┆ Int64  │
+            ╞═══════════════════╪══════════════════╪═══════════════════╪══════════════════╪════════╡
+            │ 5.1               ┆ 3.5              ┆ 1.4               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.9               ┆ 3                ┆ 1.4               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.7               ┆ 3.2              ┆ 1.3               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.6               ┆ 3.1              ┆ 1.5               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 5                 ┆ 3.6              ┆ 1.4               ┆ 0.2              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 5.4               ┆ 3.9              ┆ 1.7               ┆ 0.4              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 4.6               ┆ 3.4              ┆ 1.4               ┆ 0.3              ┆ 0      │
+            ├╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌╌┼╌╌╌╌╌╌╌╌┤
+            │ 5                 ┆ 3.4              ┆ 1.5               ┆ 0.2              ┆ 0      │
+            ╰───────────────────┴──────────────────┴───────────────────┴──────────────────┴────────╯
+
+            (Showing first 8 of 150 rows)
+
+
+    .. tab-item:: Dask
+
+        To convert a :class:`~ray.data.dataset.Dataset` to a
+        `Dask DataFrame <https://docs.dask.org/en/stable/dataframe.html>`__, call
+        :meth:`Dataset.to_dask() <ray.data.Dataset.to_dask>`.
+
+        ..
+          We skip the code snippet below because `to_dask` doesn't work with PyArrow 
+          14 and later. For more information, see https://github.com/ray-project/ray/issues/54837
+            
+        .. testcode::
+            :skipif: True
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            df = ds.to_dask()
+
+    .. tab-item:: Spark
+
+        To convert a :class:`~ray.data.dataset.Dataset` to a `Spark DataFrame
+        <https://spark.apache.org/docs/latest/api/python/reference/pyspark.sql/dataframe.html>`__,
+        call :meth:`Dataset.to_spark() <ray.data.Dataset.to_spark>`.
+
+        .. testcode::
+            :skipif: True
+
+            import ray
+            import raydp
+
+            spark = raydp.init_spark(
+                app_name = "example",
+                num_executors = 1,
+                executor_cores = 4,
+                executor_memory = "512M"
+            )
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+            df = ds.to_spark(spark)
+
+        .. testcode::
+            :skipif: True
+            :hide:
+
+            raydp.stop_spark()
+
+    .. tab-item:: Modin
+
+        To convert a :class:`~ray.data.dataset.Dataset` to a Modin DataFrame, call
+        :meth:`Dataset.to_modin() <ray.data.Dataset.to_modin>`.
+
+        .. testcode::
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            mdf = ds.to_modin()
+
+    .. tab-item:: Mars
+
+        To convert a :class:`~ray.data.dataset.Dataset` from a Mars DataFrame, call
+        :meth:`Dataset.to_mars() <ray.data.Dataset.to_mars>`.
+
+        .. testcode::
+            :skipif: True
+
+            import ray
+
+            ds = ray.data.read_csv("s3://anonymous@ray-example-data/iris.csv")
+
+            mdf = ds.to_mars()

doc/source/data/scaling-collation-functions.rst

@@ -0,0 +1,377 @@
+.. _scaling_collation_functions:
+
+Advanced: Scaling out expensive collate functions
+=================================================
+
+By default, the collate function executes on the training worker when you call :meth:`ray.data.DataIterator.iter_torch_batches`. This approach has two main drawbacks:
+
+- **Low scalability**: The collate function runs sequentially on each training worker, limiting parallelism.
+- **Resource competition**: The collate function consumes CPU and memory resources from the training worker, potentially slowing down model training.
+
+Scaling out the collate function to Ray Data allows you to scale collation across multiple CPU nodes independently of training workers, improving better overall pipeline throughput, especially with heavy collate functions.
+
+This optimization is particularly effective when the collate function is computationally expensive (such as tokenization, image augmentation, or complex feature engineering) and you have additional CPU resources available for data preprocessing.
+
+Moving the collate function to Ray Data
+---------------------------------------
+
+The following example shows a typical collate function that runs on the training worker:
+
+.. code-block:: python
+
+    train_dataset = read_parquet().map(...)
+
+    def train_func():
+        for batch in ray.train.get_dataset_shard("train").iter_torch_batches(
+            collate_fn=collate_fn,
+            batch_size=BATCH_SIZE
+        ):
+            # Training logic here
+            pass
+
+    trainer = TorchTrainer(
+        train_func,
+        datasets={"train": train_dataset},
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+    )
+
+    result = trainer.fit()
+
+If the collate function is time/compute intensive and you'd like to scale it out,you should:
+
+* Create a custom collate function that runs in Ray Data and use :meth:`ray.data.Dataset.map_batches` to scale it out.
+* Use :meth:`ray.data.Dataset.repartition` to ensure the batch size alignment.
+
+
+Creating a custom collate function that runs in Ray Data
+--------------------------------------------------------
+
+To scale out, move the ``collate_fn`` into a Ray Data ``map_batches`` operation:
+
+.. code-block:: python
+
+    def collate_fn(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        return batch
+
+    train_dataset = train_dataset.map_batches(collate_fn, batch_size=BATCH_SIZE)
+
+    def train_func():
+        for batch in ray.train.get_dataset_shard("train").iter_torch_batches(
+            collate_fn=None,
+            batch_size=BATCH_SIZE,
+        ):
+            # Training logic here
+            pass
+
+    trainer = TorchTrainer(
+        train_func,
+        datasets={"train": train_dataset},
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+    )
+
+    result = trainer.fit()
+
+A couple of things to note:
+
+- The ``collate_fn`` returns a dictionary of NumPy arrays, which is a standard Ray Data batch format.
+- The ``iter_torch_batches`` method uses ``collate_fn=None``, which reduces the amount of work is done on the training worker process.
+
+Ensuring batch size alignment
+-----------------------------
+
+Typically, collate functions are used to create complete batches of data with a target batch size.
+However, if you move the collate function to Ray Data using :meth:`ray.data.Dataset.map_batches`, it doesn't guarantee the batch size for each function call by default.
+
+There are two common problems that you may encounter.
+
+1. The collate function requires a certain number of rows provided as an input to work properly.
+2. You want to avoid any reformatting / rebatching of the data on the training worker process.
+
+To solve these problems, you can use :meth:`ray.data.Dataset.repartition` with ``target_num_rows_per_block`` to ensure the batch size alignment.
+
+By calling ``repartition`` before ``map_batches``, you ensure that the input blocks contain the desired number of rows.
+
+.. code-block:: python
+
+    # Note: If you only use map_batches(batch_size=BATCH_SIZE), you are not guaranteed to get the desired number of rows as an input.
+    dataset = dataset.repartition(target_num_rows_per_block=BATCH_SIZE).map_batches(collate_fn, batch_size=BATCH_SIZE)
+
+By calling ``repartition`` after ``map_batches``, you ensure that the output blocks contain the desired number of rows. This avoids any reformatting / rebatching of the data on the training worker process.
+
+.. code-block:: python
+
+    dataset = dataset.map_batches(collate_fn, batch_size=BATCH_SIZE).repartition(target_num_rows_per_block=BATCH_SIZE)
+
+    def train_func():
+        for batch in ray.train.get_dataset_shard("train").iter_torch_batches(
+            collate_fn=None,
+            batch_size=BATCH_SIZE,
+        ):
+            # Training logic here
+            pass
+
+    trainer = TorchTrainer(
+        train_func,
+        datasets={"train": train_dataset},
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+    )
+
+    result = trainer.fit()
+
+Putting things together
+-----------------------
+
+This guide uses a mock text dataset to demonstrate the optimization. You can find the implementation of the mock dataset in :ref:`random-text-generator`.
+
+.. tab-set::
+    .. tab-item:: Baseline implementation
+
+        The following example shows a typical collate function that runs on the training worker:
+
+        .. testcode::
+            :skipif: True
+
+            from transformers import AutoTokenizer
+            import torch
+            import numpy as np
+            from typing import Dict
+            from ray.train.torch import TorchTrainer
+            from ray.train import ScalingConfig
+            from mock_dataset import create_mock_ray_text_dataset
+
+            BATCH_SIZE = 10000
+
+            def vanilla_collate_fn(tokenizer: AutoTokenizer, batch: Dict[str, np.ndarray]) -> Dict[str, torch.Tensor]:
+                outputs = tokenizer(
+                    list(batch["text"]),
+                    truncation=True,
+                    padding="longest",
+                    return_tensors="pt",
+                )
+                outputs["labels"] = torch.LongTensor(batch["label"])
+                return outputs
+
+            def train_func():
+                tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+                collate_fn = lambda x: vanilla_collate_fn(tokenizer, x)
+
+                # Collate function runs on the training worker
+                for batch in ray.train.get_dataset_shard("train").iter_torch_batches(
+                    collate_fn=collate_fn,
+                    batch_size=BATCH_SIZE
+                ):
+                    # Training logic here
+                    pass
+
+            train_dataset = create_mock_ray_text_dataset(
+                dataset_size=1000000,
+                min_len=1000,
+                max_len=3000
+            )
+
+            trainer = TorchTrainer(
+                train_func,
+                datasets={"train": train_dataset},
+                scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+            )
+
+            result = trainer.fit()
+
+    .. tab-item:: Optimized implementation
+
+        The following example moves the collate function to Ray Data preprocessing:
+
+        .. testcode::
+            :skipif: True
+
+            from transformers import AutoTokenizer
+            import numpy as np
+            from typing import Dict
+            from ray.train.torch import TorchTrainer
+            from ray.train import ScalingConfig
+            from mock_dataset import create_mock_ray_text_dataset
+            import pyarrow as pa
+
+            BATCH_SIZE = 10000
+
+            class CollateFnRayData:
+                def __init__(self):
+                    self.tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+                def __call__(self, batch: pa.Table) -> Dict[str, np.ndarray]:
+                    results = self.tokenizer(
+                        batch["text"].to_pylist(),
+                        truncation=True,
+                        padding="longest",
+                        return_tensors="np",
+                    )
+                    results["labels"] = np.array(batch["label"])
+                    return results
+
+            def train_func():
+                # Collate function already ran in Ray Data
+                for batch in ray.train.get_dataset_shard("train").iter_torch_batches(
+                    collate_fn=None,
+                    batch_size=BATCH_SIZE,
+                ):
+                    # Training logic here
+                    pass
+
+            # Apply preprocessing in Ray Data
+            train_dataset = (
+                create_mock_ray_text_dataset(
+                    dataset_size=1000000,
+                    min_len=1000,
+                    max_len=3000
+                )
+                .map_batches(
+                    CollateFnRayData,
+                    batch_size=BATCH_SIZE,
+                    batch_format="pyarrow",
+                )
+                .repartition(target_num_rows_per_block=BATCH_SIZE)  # Ensure batch size alignment
+            )
+
+            trainer = TorchTrainer(
+                train_func,
+                datasets={"train": train_dataset},
+                scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+            )
+
+            result = trainer.fit()
+
+The optimized implementation makes these changes:
+
+- **Preprocessing in Ray Data**: The tokenization logic moves from ``train_func`` to ``CollateFnRayData``, which runs in ``map_batches``.
+- **NumPy output**: The collate function returns ``Dict[str, np.ndarray]`` instead of PyTorch tensors, which Ray Data natively supports.
+- **Batch alignment**: ``repartition(target_num_rows_per_block=BATCH_SIZE)`` after ``map_batches`` ensures the collate function receives exact batch sizes and output blocks align with the batch size.
+- **No ``collate_fn`` in iterator**: ``iter_torch_batches`` uses ``collate_fn=None`` because preprocessing already happened in Ray Data.
+
+Benchmark results
+~~~~~~~~~~~~~~~~~
+
+The following benchmarks demonstrate the performance improvement from scaling out the collate function. The test uses text tokenization with a batch size of 10,000 on a dataset of 1 million rows with text lengths between 1,000 and 3,000 characters.
+
+**Single node (g4dn.12xlarge: 48 vCPU, 4 NVIDIA T4 GPUs, 192 GiB memory)**
+
+.. list-table::
+   :header-rows: 1
+
+   * - Configuration
+     - Throughput
+   * - Collate in iterator (baseline)
+     - 1,588 rows/s
+   * - Collate in Ray Data
+     - 3,437 rows/s
+
+**With 2 additional CPU nodes (m5.8xlarge: 32 vCPU, 128 GiB memory each)**
+
+.. list-table::
+   :header-rows: 1
+
+   * - Configuration
+     - Throughput
+   * - Collate in iterator (baseline)
+     - 1,659 rows/s
+   * - Collate in Ray Data
+     - 10,717 rows/s
+
+The results show that scaling out the collate function to Ray Data provides a 2x speedup on a single node and a 6x speedup when adding CPU-only nodes for preprocessing.
+
+Advanced: Handling custom data types
+------------------------------------
+
+The preceding optimized implementation returns ``Dict[str, np.ndarray]``, which Ray Data natively supports. However, if your collate function needs to return PyTorch tensors or other custom data types that :meth:`ray.data.Dataset.map_batches` doesn't directly support, you need to serialize them.
+
+.. _train-tensor-serialization-utility:
+
+Tensor serialization utility
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The following utility serializes PyTorch tensors into PyArrow format. It flattens all tensors in a batch into a single binary buffer, stores metadata about tensor shapes and dtypes, and packs everything into a single-row PyArrow table. On the training side, it deserializes the table back into the original tensor structure.
+
+The serialization and deserialization operations are typically lightweight compared to the actual collate function work (such as tokenization or image processing), so the overhead is minimal relative to the performance gains from scaling the collate function.
+
+You can use :ref:`train-collate-utils` as a reference implementation and adapt it to your needs.
+
+Example with tensor serialization
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The following example demonstrates using tensor serialization when your collate function must return PyTorch tensors. This approach requires ``repartition`` before ``map_batches`` because the collate function changes the number of output rows (each batch becomes a single serialized row).
+
+.. testcode::
+    :skipif: True
+
+    from transformers import AutoTokenizer
+    import torch
+    from typing import Dict
+    from ray.data.collate_fn import ArrowBatchCollateFn
+    import pyarrow as pa
+    from collate_utils import serialize_tensors_to_table, deserialize_table_to_tensors
+    from ray.train.torch import TorchTrainer
+    from ray.train import ScalingConfig
+    from mock_dataset import create_mock_ray_text_dataset
+
+    BATCH_SIZE = 10000
+
+    class TextTokenizerCollateFn:
+        """Collate function that runs in Ray Data preprocessing."""
+        def __init__(self):
+            self.tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
+
+        def __call__(self, batch: pa.Table) -> pa.Table:
+            # Tokenize the batch
+            outputs = self.tokenizer(
+                batch["text"].to_pylist(),
+                truncation=True,
+                padding="longest",
+                return_tensors="pt",
+            )
+            outputs["labels"] = torch.LongTensor(batch["label"].to_numpy())
+            
+            # Serialize to single-row table using the utility
+            return serialize_tensors_to_table(outputs)
+
+    class IteratorCollateFn(ArrowBatchCollateFn):
+        """Collate function for iter_torch_batches that deserializes the batch."""
+        def __init__(self, pin_memory=False):
+            self._pin_memory = pin_memory
+
+        def __call__(self, batch: pa.Table) -> Dict[str, torch.Tensor]:
+            # Deserialize from single-row table using the utility
+            return deserialize_table_to_tensors(batch, pin_memory=self._pin_memory)
+
+    def train_func():
+        collate_fn = IteratorCollateFn()
+        
+        # Collate function only deserializes on the training worker
+        for batch in ray.train.get_dataset_shard("train").iter_torch_batches(
+            collate_fn=collate_fn,
+            batch_size=1  # Each "row" is actually a full batch
+        ):
+            # Training logic here
+            pass
+
+    # Apply preprocessing in Ray Data
+    # Use repartition BEFORE map_batches because output row count changes
+    train_dataset = (
+        create_mock_ray_text_dataset(
+            dataset_size=1000000,
+            min_len=1000,
+            max_len=3000
+        )
+        .repartition(target_num_rows_per_block=BATCH_SIZE)
+        .map_batches(
+            TextTokenizerCollateFn,
+            batch_size=BATCH_SIZE,
+            batch_format="pyarrow",
+        )
+    )
+
+    trainer = TorchTrainer(
+        train_func,
+        datasets={"train": train_dataset},
+        scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
+    )
+
+    result = trainer.fit()

doc/source/data/shuffling-data.rst

@@ -0,0 +1,345 @@
+.. _shuffling_data:
+
+==============
+Shuffling Data
+==============
+
+When consuming or iterating over Ray :class:`Datasets <ray.data.dataset.Dataset>`, it can be useful to
+shuffle or randomize the order of data (for example, randomizing data ingest order during ML training).
+This guide shows several different methods of shuffling data with Ray Data and their respective trade-offs.
+
+Types of shuffling
+==================
+
+Ray Data provides several different options for shuffling data, trading off the granularity of shuffle
+control with memory consumption and runtime. The list below presents options in increasing order of
+resource consumption and runtime. Choose the most appropriate method for your use case.
+
+.. _shuffling_file_order:
+
+Shuffle the ordering of files
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To randomly shuffle the ordering of input files before reading, call a :ref:`read function <loading-data-api>` function that supports shuffling, such as
+:func:`~ray.data.read_images`, and use the ``shuffle="files"`` parameter. This randomly assigns
+input files to workers for reading.
+
+This is the fastest "shuffle" option: it's purely a metadata operation---the system random-shuffles the list of files constituting the dataset before
+fetching them with reading tasks. This option, however, doesn't shuffle the rows inside files, so the randomness might not be
+sufficient for your needs in case of files with the large number of rows.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_images(
+        "s3://anonymous@ray-example-data/image-datasets/simple",
+        shuffle="files",
+    )
+
+.. _local_shuffle_buffer:
+
+Local buffer shuffle
+~~~~~~~~~~~~~~~~~~~~
+
+To locally shuffle a subset of rows using iteration methods, such as :meth:`~ray.data.Dataset.iter_batches`,
+:meth:`~ray.data.Dataset.iter_torch_batches`, and :meth:`~ray.data.Dataset.iter_tf_batches`,
+specify `local_shuffle_buffer_size`.
+
+This shuffles up to a `local_shuffle_buffer_size` number of rows buffered during iteration. See more details in
+:ref:`Iterating over batches with shuffling <iterating-over-batches-with-shuffling>`.
+
+This is slower than files shuffling, and shuffles rows locally without
+network transfer. You can use this local shuffle buffer together with shuffling
+ordering of files. See :ref:`Shuffle the ordering of files <shuffling_file_order>`.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+
+    for batch in ds.iter_batches(
+        batch_size=2,
+        batch_format="numpy",
+        local_shuffle_buffer_size=250,
+    ):
+        print(batch)
+
+.. tip::
+
+    If you observe reduced throughput when using ``local_shuffle_buffer_size``,
+    check the total time spent in batch creation by
+    examining the ``ds.stats()`` output (``In batch formatting``, under
+    ``Batch iteration time breakdown``). If this time is significantly larger than the
+    time spent in other steps, decrease ``local_shuffle_buffer_size`` or turn off the local
+    shuffle buffer altogether and only :ref:`shuffle the ordering of files <shuffling_file_order>`.
+
+.. _map_batches_shuffle:
+
+``map_batches`` shuffle
+~~~~~~~~~~~~~~~~~~~~~~~
+
+To shuffle data as a separate data stage, use :meth:`~ray.data.Dataset.map_batches`
+with a shuffle function that randomly permutes rows within each batch. Compared to local
+buffer shuffle, this approach has several advantages:
+
+- It **decouples shuffling from the iterator**, running as a separate Ray Data operator
+  that doesn't block downstream CPU/GPU processing.
+- Ray Data's resource management automatically schedules shuffle tasks based on available
+  cluster resources (CPU, memory), avoiding resource contention.
+- The shuffle work can happen in parallel across multiple machines, making it more scalable
+  for large datasets.
+
+The ``batch_size`` parameter controls the shuffle window---a larger value shuffles more rows
+together for better randomness but requires more memory.
+
+.. important::
+
+    Always set the ``memory`` parameter when using large batch sizes to avoid out-of-memory
+    errors. Estimate it as ``batch_size * row_bytes``:
+
+.. testcode::
+
+    import numpy as np
+    import pyarrow as pa
+    import ray
+
+    def random_shuffle(batch: pa.Table) -> pa.Table:
+        indices = np.random.permutation(len(batch))
+        return batch.take(indices)
+
+    row_bytes = 4096
+    shuffle_memory = int(2**30)  # 1 GB shuffle window
+    batch_size = int(shuffle_memory / row_bytes)
+
+    ds = ray.data.range(1000)
+    ds = ds.map_batches(
+        random_shuffle,
+        batch_size=batch_size,
+        batch_format="pyarrow",
+        memory=shuffle_memory,
+    )
+    ds.take(10)
+
+.. tip::
+
+    Combine ``map_batches`` shuffle with :ref:`file order shuffling <shuffling_file_order>` for
+    additional randomness. File order shuffling randomizes which files are read first, while
+    ``map_batches`` shuffle randomizes rows within each shuffle window.
+
+.. _map_batches_vs_local_shuffle:
+
+Comparing local buffer shuffle and ``map_batches`` shuffle
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+The following benchmark compares steady-state training throughput between local buffer shuffle
+and ``map_batches`` shuffle on a synthetic workload (``ray.data.range_tensor``, ~4 KB/row, 4 GPU
+workers, batch size 4096, 200 steps with 100 warmup):
+
+.. list-table:: Local buffer shuffle vs. ``map_batches`` shuffle
+   :header-rows: 1
+   :widths: 30 20 15
+
+   * - Method
+     - Throughput (rows/s)
+     - % of baseline
+   * - No shuffle (baseline)
+     - 1,759,282
+     - 100%
+   * - Local buffer shuffle 1 GB
+     - 225,181
+     - 13%
+   * - Local buffer shuffle 2 GB
+     - 220,644
+     - 13%
+   * - Local buffer shuffle 3 GB
+     - 153,256
+     - 9%
+   * - ``map_batches`` shuffle 1 GB
+     - 1,400,734
+     - 80%
+   * - ``map_batches`` shuffle 2 GB
+     - 1,460,037
+     - 83%
+   * - ``map_batches`` shuffle 3 GB
+     - 1,588,428
+     - 90%
+
+
+Randomizing block order
+~~~~~~~~~~~~~~~~~~~~~~~
+
+This option randomizes the order of :ref:`blocks <data_key_concepts>` in a dataset. While applying this operation alone doesn't involve heavy computation
+and communication, it requires Ray Data to materialize all blocks in memory before actually randomizing their ordering in the queue for subsequent operation.
+
+.. note:: Ray Data doesn't guarantee any particular ordering of the blocks when reading blocks from different files in parallel by default, unless you set `DataContext.execution_options.preserve_order` to true. Henceforth, this particular option
+    is primarily relevant in cases when the system yields blocks from relatively small set of very large files.
+
+.. note:: Only use this option when your dataset is small enough to fit into the object store memory.
+
+To perform block order shuffling, use :meth:`randomize_block_order <ray.data.Dataset.randomize_block_order>`.
+
+.. testcode::
+    import ray
+
+    ds = ray.data.read_text(
+        "s3://anonymous@ray-example-data/sms_spam_collection_subset.txt"
+    )
+
+    # Randomize the block order of this dataset.
+    ds = ds.randomize_block_order()
+
+Global shuffle
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+To shuffle all rows globally, across the whole dataset, multiple options are available
+
+1. *Random shuffling*: invoking :meth:`~ray.data.Dataset.random_shuffle` essentially permutes and shuffles individual rows
+from existing blocks into the new ones using an optionally provided seed.
+2. (**New in 2.46**) *Key-based repartitioning*: invoking :meth:`~ray.data.Dataset.repartition` with `keys` parameter triggers
+:ref:`hash-shuffle <hash-shuffle>` operation, shuffling the rows based on the hash of the values in the provided key columns, providing
+deterministic way of co-locating rows based on the hash of the column values.
+
+Note that shuffle is an expensive operation requiring materializing of the whole dataset in memory as well as serving as a synchronization
+barrier---subsequent operators won't be able to start executing until shuffle completion.
+
+Example of random shuffling with seed:
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+
+    # Random shuffle with seed
+    random_shuffled_ds = ds.random_shuffle(seed=123)
+
+
+Example of hash shuffling based on column `id`:
+
+.. testcode::
+
+    import ray
+    from ray.data.context import DataContext, ShuffleStrategy
+
+    # First enable hash-shuffle as shuffling strategy
+    DataContext.get_current().shuffle_strategy = ShuffleStrategy.HASH_SHUFFLE
+
+    # Hash-shuffle
+    hash_shuffled_ds = ds.repartition(keys="id", num_blocks=200)
+
+.. _optimizing_shuffles:
+
+Advanced: Optimizing shuffles
+=============================
+.. note:: This is an active area of development. If your Dataset uses a shuffle operation and you are having trouble configuring shuffle,
+    `file a Ray Data issue on GitHub <https://github.com/ray-project/ray/issues/new?assignees=&labels=bug%2Ctriage%2Cdata&projects=&template=bug-report.yml&title=[data]+>`_.
+
+When should you use global per-epoch shuffling?
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Use global per-epoch shuffling only if your model is sensitive to the
+randomness of the training data. Based on a
+`theoretical foundation <https://arxiv.org/abs/1709.10432>`__, all
+gradient-descent-based model trainers benefit from improved global shuffle quality.
+In practice, the benefit's particularly pronounced for tabular data/models.
+However, the more global the shuffle is, the more expensive the shuffling operation.
+The increase compounds with distributed data-parallel training on a multi-node cluster due
+to data transfer costs. This cost can be prohibitive when using very large datasets.
+
+The best route for determining the best tradeoff between preprocessing time and cost and
+per-epoch shuffle quality is to measure the precision gain per training step for your
+particular model under different shuffling policies such as no shuffling, local shuffling, or global shuffling.
+
+As long as your data loading and shuffling throughput is higher than your training throughput, your GPU should
+saturate. If you have shuffle-sensitive models, push the
+shuffle quality higher until you reach this threshold.
+
+.. _shuffle_performance_tips:
+
+Enabling push-based shuffle
+~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Some Dataset operations require a *shuffle* operation, meaning that the system shuffles data from all of the input partitions to all of the output partitions.
+These operations include :meth:`Dataset.random_shuffle <ray.data.Dataset.random_shuffle>`,
+:meth:`Dataset.sort <ray.data.Dataset.sort>` and :meth:`Dataset.groupby <ray.data.Dataset.groupby>`.
+For example, during a sort operation, the system reorders data between blocks and therefore requires shuffling across partitions.
+Shuffling can be challenging to scale to large data sizes and clusters, especially when the total dataset size can't fit into memory.
+
+Ray Data provides an alternative shuffle implementation known as push-based shuffle for improving large-scale performance.
+Try this out if your dataset has more than 1000 blocks or is larger than 1 TB in size.
+
+To try this out locally or on a cluster, you can start with the `nightly release test <https://github.com/ray-project/ray/blob/master/release/nightly_tests/dataset/sort_benchmark.py>`_ that Ray runs for :meth:`Dataset.random_shuffle <ray.data.Dataset.random_shuffle>` and :meth:`Dataset.sort <ray.data.Dataset.sort>`.
+To get an idea of the performance you can expect, here are some run time results for :meth:`Dataset.random_shuffle <ray.data.Dataset.random_shuffle>` on 1-10 TB of data on 20 machines - m5.4xlarge instances on AWS EC2, each with 16 vCPUs, 64 GB RAM.
+
+.. image:: https://docs.google.com/spreadsheets/d/e/2PACX-1vQvBWpdxHsW0-loasJsBpdarAixb7rjoo-lTgikghfCeKPQtjQDDo2fY51Yc1B6k_S4bnYEoChmFrH2/pubchart?oid=598567373&format=image
+   :align: center
+
+To try out push-based shuffle, set the environment variable ``RAY_DATA_PUSH_BASED_SHUFFLE=1`` when running your application:
+
+.. code-block:: bash
+
+    $ wget https://raw.githubusercontent.com/ray-project/ray/master/release/nightly_tests/dataset/sort_benchmark.py
+    $ RAY_DATA_PUSH_BASED_SHUFFLE=1 python sort_benchmark.py --num-partitions=10 --partition-size=1e7
+
+    # Dataset size: 10 partitions, 0.01GB partition size, 0.1GB total
+    # [dataset]: Run `pip install tqdm` to enable progress reporting.
+    # 2022-05-04 17:30:28,806	INFO push_based_shuffle.py:118 -- Using experimental push-based shuffle.
+    # Finished in 9.571171760559082
+    # ...
+
+You can also specify the shuffle implementation during program execution by
+setting the ``DataContext.use_push_based_shuffle`` flag:
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.use_push_based_shuffle = True
+
+    ds = (
+        ray.data.range(1000)
+        .random_shuffle()
+    )
+
+Large-scale shuffles can take a while to finish.
+For debugging purposes, shuffle operations support executing only part of the shuffle, so that you can collect an execution profile more quickly.
+Here is an example that shows how to limit a random shuffle operation to two output blocks:
+
+.. testcode::
+    :hide:
+
+    import ray
+    ray.shutdown()
+
+.. testcode::
+
+    import ray
+
+    ctx = ray.data.DataContext.get_current()
+    ctx.set_config(
+        "debug_limit_shuffle_execution_to_num_blocks", 2
+    )
+
+    ds = (
+        ray.data.range(1000, override_num_blocks=10)
+        .random_shuffle()
+        .materialize()
+    )
+    print(ds.stats())
+
+.. testoutput::
+    :options: +MOCK
+
+    Operator 1 ReadRange->RandomShuffle: executed in 0.08s
+
+        Suboperator 0 ReadRange->RandomShuffleMap: 2/2 blocks executed
+        ...

doc/source/data/transforming-data.rst

@@ -0,0 +1,597 @@
+.. _transforming_data:
+
+=================
+Transforming Data
+=================
+
+Transformations let you process and modify your dataset. You can compose transformations
+to express a chain of computations.
+
+.. note::
+    Transformations are lazy by default. They aren't executed until you trigger consumption of the data by :ref:`iterating over the Dataset <iterating-over-data>`, :ref:`saving the Dataset <saving-data>`, or :ref:`inspecting properties of the Dataset <inspecting-data>`.
+
+This guide shows you how to scale transformations (or user-defined functions (UDFs)) on your Ray Data dataset.
+
+.. _transforming_rows:
+
+Transforming rows
+=================
+
+.. tip::
+
+    If your transformation is vectorized, call :meth:`~ray.data.Dataset.map_batches` for
+    better performance. To learn more, see :ref:`Transforming batches <transforming_batches>`.
+
+Transforming rows with map
+~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If your transformation returns exactly one row for each input row, call
+:meth:`~ray.data.Dataset.map`. This transformation is automatically parallelized across your Ray cluster.
+
+.. testcode::
+
+    import os
+    from typing import Any, Dict
+    import ray
+
+    def parse_filename(row: Dict[str, Any]) -> Dict[str, Any]:
+        row["filename"] = os.path.basename(row["path"])
+        return row
+
+    ds = (
+        ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple", include_paths=True)
+        .map(parse_filename)
+    )
+
+The user defined function passed to :meth:`~ray.data.Dataset.map` should be of type
+`Callable[[Dict[str, Any]], Dict[str, Any]]`. In other words, your function should
+input and output a dictionary with keys of strings and values of any type. For example:
+
+.. testcode::
+
+    from typing import Any, Dict
+
+    def fn(row: Dict[str, Any]) -> Dict[str, Any]:
+        # access row data
+        value = row["col1"]
+
+        # add data to row
+        row["col2"] = ...
+
+        # return row
+        return row
+
+Transforming rows with flat map
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+If your transformation returns multiple rows for each input row, call
+:meth:`~ray.data.Dataset.flat_map`. This transformation is automatically parallelized across your Ray cluster.
+
+.. 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}]
+
+The user defined function passed to :meth:`~ray.data.Dataset.flat_map` should be of type
+`Callable[[Dict[str, Any]], List[Dict[str, Any]]]`. In other words your function should
+input a dictionary with keys of strings and values of any type and output a list of
+dictionaries that have the same type as the input, for example:
+
+.. testcode::
+
+    from typing import Any, Dict, List
+
+    def fn(row: Dict[str, Any]) -> List[Dict[str, Any]]:
+        # access row data
+        value = row["col1"]
+
+        # add data to row
+        row["col2"] = ...
+
+        # construct output list
+        output = [row, row]
+
+        # return list of output rows
+        return output
+
+.. _transforming_batches:
+
+Transforming batches
+====================
+
+If your transformation can be vectorized using NumPy, PyArrow or Pandas operations, transforming
+batches is considerably more performant than transforming individual rows.
+
+This transformation is automatically parallelized across your Ray cluster.
+
+.. testcode::
+
+    from typing import Dict
+    import numpy as np
+    import ray
+
+    def increase_brightness(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        batch["image"] = np.clip(batch["image"] + 4, 0, 255)
+        return batch
+
+    ds = (
+        ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+        .map_batches(increase_brightness)
+    )
+
+.. _configure_batch_format:
+
+Configuring batch format
+~~~~~~~~~~~~~~~~~~~~~~~~
+
+Ray Data represents batches as dicts of NumPy ndarrays, pandas DataFrames or Arrow Tables. By
+default, Ray Data represents batches as dicts of NumPy ndarrays. To configure the batch type,
+specify ``batch_format`` in :meth:`~ray.data.Dataset.map_batches`. You can return either
+format from your function, but ``batch_format`` should match the input of your function.
+
+When applying transformations to batches of rows, Ray Data could represent these batches as either NumPy's ``ndarrays``,
+Pandas ``DataFrame`` or PyArrow ``Table``.
+
+When using
+    * ``batch_format=numpy``, the input to the function is a dictionary where keys correspond to column names and values to column values represented as ``ndarrays``.
+    * ``batch_format=pyarrow``, the input to the function is a Pyarrow ``Table``.
+    * ``batch_format=pandas``, the input to the function is a Pandas ``DataFrame``.
+
+.. tab-set::
+
+    .. tab-item:: NumPy
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+            import ray
+
+            def increase_brightness(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                batch["image"] = np.clip(batch["image"] + 4, 0, 255)
+                return batch
+
+            ds = (
+                ray.data.read_images("s3://anonymous@ray-example-data/image-datasets/simple")
+                .map_batches(increase_brightness, batch_format="numpy")
+            )
+
+    .. tab-item:: pandas
+
+        .. testcode::
+
+            import pandas as pd
+            import ray
+
+            def drop_nas(batch: pd.DataFrame) -> pd.DataFrame:
+                return batch.dropna()
+
+            ds = (
+                ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+                .map_batches(drop_nas, batch_format="pandas")
+            )
+    .. tab-item:: pyarrow
+
+        .. testcode::
+
+            import pyarrow as pa
+            import pyarrow.compute as pc
+            import ray
+
+            def drop_nas(batch: pa.Table) -> pa.Table:
+                return pc.drop_null(batch)
+
+            ds = (
+                ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+                .map_batches(drop_nas, batch_format="pyarrow")
+            )
+
+The user defined function can also be a Python generator that yields batches, so the function can also
+be of type ``Callable[DataBatch, Iterator[[DataBatch]]``, where ``DataBatch = Union[pd.DataFrame, Dict[str, np.ndarray], pyarrow.Table]``.
+In this case, your function would look like:
+
+.. testcode::
+
+    from typing import Dict, Iterator
+    import numpy as np
+
+    def fn(batch: Dict[str, np.ndarray]) -> Iterator[Dict[str, np.ndarray]]:
+        # yield the same batch multiple times
+        for _ in range(10):
+            yield batch
+
+Choosing the right batch format
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+When choosing appropriate batch format for your ``map_batches`` primary consideration is a trade-off of convenience vs performance:
+
+1. Batches are a sliding window into the underlying block: the UDF is invoked with a subset of rows of the underlying block that make up the current batch of specified ``batch_size``. Specifying ``batch_size=None`` makes batch include all rows of the block in a single batch.
+2. Depending on the batch format, such view can either be a *zero-copy* (when batch format matches the block type of either ``pandas`` or ``pyarrow``) or copying one (when the batch format differs from the block type).
+
+For example, if the underlying block type is Arrow, specifying ``batch_format="numpy"`` or ``batch_format="pandas"``  might invoke a copy on the underlying data when converting it from the underlying block type.
+
+Ray Data also strives to minimize the amount of data conversions: for example, if your ``map_batches`` operation returns Pandas batches, then these batches are combined into blocks *without* conversion and propagated further as Pandas blocks. Most Ray Data datasources produce Arrow blocks, so using batch format ``pyarrow`` can avoid unnecessary data conversions.
+
+If you'd like to use a more ergonomic API for transformations but avoid performance overheads, you can consider using Polars inside your ``map_batches`` operation with ``batch_format="pyarrow"`` as follows:
+
+.. testcode::
+
+    import pyarrow as pa
+
+    def udf(table: pa.Table):
+        import polars as pl
+        df = polars.from_pyarrow(table)
+        df.summary()
+        return df.to_arrow()
+
+    ds.map_batches(udf, batch_format="pyarrow")
+
+
+Configuring batch size
+~~~~~~~~~~~~~~~~~~~~~~
+
+Increasing ``batch_size`` improves the performance of vectorized transformations as well
+as performance of model inference. However, if your batch size is too large, your
+program might run into out-of-memory (OOM) errors.
+
+Use ``batch_size="auto"`` to let Ray Data automatically determine an appropriate batch
+size based on the size of your data. For GPU workloads, you must specify an explicit
+integer batch size. If you encounter OOM errors with an explicit batch size, try decreasing it.
+
+.. _stateful_transforms:
+
+Stateful/Class-based Transforms
+===============================
+
+If your transform requires expensive setup such as downloading
+model weights, use a callable Python class instead of a function to make the transform stateful. When a Python class
+is used, the ``__init__`` method is called to perform setup exactly once on each worker.
+In contrast, functions are stateless, so any setup must be performed for each data item.
+
+Internally, Ray Data uses tasks to execute functions, and uses actors to execute classes.
+To learn more about tasks and actors, read the
+:ref:`Ray Core Key Concepts <core-key-concepts>`.
+
+To transform data with a Python class, complete these steps:
+
+1. Implement a class. Perform setup in ``__init__`` and transform data in ``__call__``.
+
+2. Call :meth:`~ray.data.Dataset.map_batches`, :meth:`~ray.data.Dataset.map`, or
+   :meth:`~ray.data.Dataset.flat_map`. Pass a ``ray.data.ActorPoolStrategy(...)`` object to
+   the  ``compute`` argument to control how many workers Ray uses. Each worker transforms a partition
+   of data in parallel.
+
+.. tab-set::
+
+    .. tab-item:: CPU
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+            import torch
+            import ray
+
+            class TorchPredictor:
+
+                def __init__(self):
+                    self.model = torch.nn.Identity()
+                    self.model.eval()
+
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                    inputs = torch.as_tensor(batch["data"], dtype=torch.float32)
+                    with torch.inference_mode():
+                        batch["output"] = self.model(inputs).detach().numpy()
+                    return batch
+
+            ds = (
+                ray.data.from_numpy(np.ones((32, 100)))
+                .map_batches(
+                    TorchPredictor,
+                    compute=ray.data.ActorPoolStrategy(size=2),
+                )
+            )
+
+        .. testcode::
+            :hide:
+
+            ds.materialize()
+
+    .. tab-item:: GPU
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+            import torch
+            import ray
+
+            class TorchPredictor:
+                def __init__(self):
+                    self.model = torch.nn.Identity().cuda()
+                    self.model.eval()
+
+                def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                    inputs = torch.as_tensor(batch["data"], dtype=torch.float32).cuda()
+                    with torch.inference_mode():
+                        batch["output"] = self.model(inputs).detach().cpu().numpy()
+                    return batch
+
+            ds = (
+                ray.data.from_numpy(np.ones((32, 100)))
+                .map_batches(
+                    TorchPredictor,
+                    # Two workers with one GPU each
+                    compute=ray.data.ActorPoolStrategy(size=2),
+                    # Batch size is required if you're using GPUs.
+                    batch_size=4,
+                    num_gpus=1
+                )
+            )
+
+        .. testcode::
+            :hide:
+
+            ds.materialize()
+
+Specifying CPUs, GPUs, and Memory
+=================================
+
+You can optionally specify logical resources per transformation by using one of the following parameters: ``num_cpus``, ``num_gpus``, ``memory``, ``resources``.
+
+* ``num_cpus``: The number of CPUs to use for the transformation.
+* ``num_gpus``: The number of GPUs to use for the transformation. Ray automatically configures the proper CUDA_VISIBLE_DEVICES environment variable so that GPUs are isolated from other tasks/actors.
+* ``memory``: The amount of memory to use for the transformation. This is useful for avoiding out-of-memory errors by telling Ray how much memory your function uses, and preventing Ray from scheduling too many tasks on a node.
+* ``resources``: A dictionary of resources to use for the transformation. This is useful for specifying custom resources.
+
+Note that these are logical resources and don't impose limits on actual physical resource usage.
+
+Also, both ``num_cpus`` and ``num_gpus`` support fractional values less than 1. For example, specifying ``num_cpus=0.5`` on a cluster with 4 CPUs allows 8 concurrent tasks/actors to run.
+You can read more about resources in Ray here: :ref:`resource-requirements`.
+
+.. testcode::
+    :hide:
+
+    import ray
+
+    ds = ray.data.range(1)
+
+.. testcode::
+
+    def uses_lots_of_memory(batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+        ...
+
+    # Tell Ray that the function uses 1 GiB of memory
+    ds.map_batches(uses_lots_of_memory, memory=1 * 1024 * 1024)
+
+Specifying Concurrency
+======================
+
+You can specify the concurrency of the transformation by using the ``compute`` parameter.
+
+For functions, use ``compute=ray.data.TaskPoolStrategy(size=n)`` to cap the number of concurrent tasks. By default, Ray Data automatically determines the number of concurrent tasks.
+For classes, use ``compute=ray.data.ActorPoolStrategy(size=n)`` to use a fixed size actor pool of ``n`` workers. If ``compute`` isn't specified, an autoscaling actor pool is used by default.
+
+.. testcode::
+
+    import ray
+
+    ds = ray.data.range(10).map_batches(lambda batch: {"id": batch["id"] * 2}, compute=ray.data.TaskPoolStrategy(size=2))
+    ds.take_all()
+
+.. testoutput::
+    :options: +MOCK
+
+    [{'id': 0}, {'id': 2}, {'id': 4}, {'id': 6}, {'id': 8}, {'id': 10}, {'id': 12}, {'id': 14}, {'id': 16}, {'id': 18}]
+
+.. _ordering_of_rows:
+
+Ordering of rows
+================
+
+When transforming data, the order of :ref:`blocks <data_key_concepts>` isn't preserved by default.
+
+If the order of blocks needs to be preserved/deterministic,
+you can use :meth:`~ray.data.Dataset.sort` method, or set :attr:`ray.data.ExecutionOptions.preserve_order` to `True`.
+Note that setting this flag may negatively impact performance on larger cluster setups where stragglers are more likely.
+
+.. testcode::
+
+   import ray
+
+   ctx = ray.data.DataContext().get_current()
+
+   # By default, this is set to False.
+   ctx.execution_options.preserve_order = True
+
+.. _transforming_groupby:
+
+Group-by and transforming groups
+================================
+
+To transform groups, call :meth:`~ray.data.Dataset.groupby` to group rows based on provided ``key`` column values.
+Then, call :meth:`~ray.data.grouped_data.GroupedData.map_groups` to execute a transformation on each group.
+
+.. tab-set::
+
+    .. tab-item:: NumPy
+
+        .. testcode::
+
+            from typing import Dict
+            import numpy as np
+            import ray
+
+            items = [
+                {"image": np.zeros((32, 32, 3)), "label": label}
+                for _ in range(10) for label in range(100)
+            ]
+
+            def normalize_images(group: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+                group["image"] = (group["image"] - group["image"].mean()) / group["image"].std()
+                return group
+
+            ds = (
+                ray.data.from_items(items)
+                .groupby("label")
+                .map_groups(normalize_images)
+            )
+
+    .. tab-item:: pandas
+
+        .. testcode::
+
+            import pandas as pd
+            import ray
+
+            def normalize_features(group: pd.DataFrame) -> pd.DataFrame:
+                target = group.drop("target")
+                group = (group - group.min()) / group.std()
+                group["target"] = target
+                return group
+
+            ds = (
+                ray.data.read_csv("s3://anonymous@air-example-data/iris.csv")
+                .groupby("target")
+                .map_groups(normalize_features)
+            )
+
+Advanced: Distributed UDFs with Placement Groups
+================================================
+
+While all transformations are automatically parallelized across your Ray cluster, often times these transformations can be distributed themselves. For example, if you're using
+a large model, you may want to distribute the model across multiple nodes.
+You can do this by using :ref:`placement groups <ray-placement-group-doc-ref>` and ``ray_remote_args_fn``, which can dynamically create placement groups for each model replica.
+
+.. testcode::
+
+    import ray
+    from typing import Dict
+    import numpy as np
+    import torch
+
+    NUM_SHARDS = 2
+    @ray.remote
+    class ModelShard:
+        def __init__(self):
+            self.model = torch.nn.Linear(10, 10)
+
+        def f(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+            return batch
+
+    class DistributedModel:
+        def __init__(self):
+            self.shards = [ModelShard.remote() for _ in range(NUM_SHARDS)]
+
+        def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+            return {"out": np.array(ray.get([shard.f.remote(batch) for shard in self.shards]))}
+
+    def ray_remote_args_fn():
+        from ray.util.scheduling_strategies import PlacementGroupSchedulingStrategy
+        pg = ray.util.placement_group([{"CPU": 1}] * NUM_SHARDS)
+        scheduling_strategy = PlacementGroupSchedulingStrategy(
+            placement_group=pg,
+            placement_group_capture_child_tasks=True,
+        )
+        return {"scheduling_strategy": scheduling_strategy}
+
+    ds = ray.data.range(10).map_batches(DistributedModel, ray_remote_args_fn=ray_remote_args_fn)
+    ds.take_all()
+
+Advanced: Asynchronous Transforms
+=================================
+
+Ray Data supports asynchronous functions by using the ``async`` keyword. This is useful for performing asynchronous operations such as fetching data from a database or making HTTP requests.
+Note that this only works when using a class-based transform function and currently requires ``uvloop==0.21.0``.
+
+.. testcode::
+
+    import ray
+    from typing import Dict
+    import numpy as np
+
+    class AsyncTransform:
+        async def __call__(self, batch: Dict[str, np.ndarray]) -> Dict[str, np.ndarray]:
+            return batch
+
+    ds = ray.data.range(10).map_batches(AsyncTransform)
+    ds.take_all()
+
+.. testoutput::
+    :options: +MOCK
+
+    [{'id': 0},
+    {'id': 1},
+    {'id': 2},
+    {'id': 3},
+    {'id': 4},
+    {'id': 5},
+    {'id': 6},
+    {'id': 7},
+    {'id': 8},
+    {'id': 9}]
+
+
+Expressions (Alpha)
+===================
+
+Ray Data expressions provide a way to specify column-based operations on datasets.
+Use :func:`~ray.data.expressions.col` to reference columns and :func:`~ray.data.expressions.lit` to create literal values.
+You can combine these with operators to create complex expressions for filtering,
+transformations, and computations.
+
+Expressions have to be used with :meth:`~ray.data.Dataset.with_column`. The core advantage of expressions
+is that because they operate on specific columns, Ray Data's optimizer can optimize the execution plan by reordering the operations.
+
+See :ref:`expressions-api` for more details.
+
+.. testcode::
+
+    import ray
+    from ray.data.expressions import col
+
+    ds = ray.data.range(10).with_column("id_2", col("id") * 2)
+    ds.show()
+
+To use a custom function with an expression, you can use :func:`~ray.data.expressions.udf`.
+
+.. testcode::
+
+    from ray.data.expressions import col, udf
+    from ray.data.datatype import DataType
+    import pyarrow as pa
+    import pyarrow.compute as pc
+    import ray
+
+    # UDF that operates on a batch of values (PyArrow Array)
+    @udf(return_dtype=DataType.int32())
+    def add_one(x: pa.Array) -> pa.Array:
+        return pc.add(x, 1)  # Vectorized operation on the entire Array
+
+    # UDF that combines multiple columns (each as a PyArrow Array)
+    @udf(return_dtype=DataType.string())
+    def format_name(first: pa.Array, last: pa.Array) -> pa.Array:
+        return pc.binary_join_element_wise(first, last, " ")  # Vectorized string concatenation
+
+    # Use in dataset operations
+    ds = ray.data.from_items([
+        {"value": 5, "first": "John", "last": "Doe"},
+        {"value": 10, "first": "Jane", "last": "Smith"}
+    ])
+    ds = ds.with_column("value_plus_one", add_one(col("value")))
+    ds = ds.with_column("full_name", format_name(col("first"), col("last")))
+    ds = ds.with_column("doubled_plus_one", add_one(col("value")) * 2)
+    ds.show()