Ray Data quickstart

Get started with Ray Data’s Dataset abstraction for distributed data processing. This guide walks through a complete workflow. You load data into a dataset, transform it, and then consume the results or save them to storage.

What’s a dataset?

Ray Data’s main abstraction is a Dataset, which represents a distributed collection of data. Datasets are designed for machine learning workloads and can efficiently handle data collections larger than a single machine’s memory.

Load data

Create datasets from sources such as local files, Python objects, and cloud storage services, including S3 and GCS. Ray Data integrates with any filesystem that Arrow supports.

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)

{'sepal length (cm)': 5.1, 'sepal width (cm)': 3.5, 'petal length (cm)': 1.4, 'petal width (cm)': 0.2, 'target': 0}

For more on creating datasets from different sources, see Loading data.

Transform data

Apply user-defined functions (UDFs) to transform datasets. Ray automatically parallelizes these transformations across your cluster.

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

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)

For more transformation options, see Transforming data.

Consume data

Access dataset contents with methods such as take_batch() and iter_batches(). You can also pass datasets directly to Ray tasks or actors for distributed processing.

# Extract the first 3 rows as a batch for processing
print(transformed_ds.take_batch(batch_size=3))

{'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 on working with dataset contents, see Iterating over data and Saving data.

Save data

Export processed datasets to a variety of formats and storage locations with methods such as write_parquet() and write_csv().

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

['..._000000.parquet', '..._000001.parquet']

For more information on saving datasets, see Saving data.