Ray Data - Scalable ML Data Processing

Distributed data processing library for ML and AI workloads.

When to use Ray Data

Use Ray Data when:

• Processing large datasets (>100GB) for ML training
• Need distributed data preprocessing across cluster
• Building batch inference pipelines
• Loading multi-modal data (images, audio, video)
• Scaling data processing from laptop to cluster

Key features:

• Streaming execution: Process data larger than memory
• GPU support: Accelerate transforms with GPUs
• Framework integration: PyTorch, TensorFlow, HuggingFace
• Multi-modal: Images, Parquet, CSV, JSON, audio, video

Use alternatives instead:

• Pandas: Small data (<1GB) on single machine
• Dask: Tabular data, SQL-like operations
• Spark: Enterprise ETL, SQL queries

Quick start

Installation

pip install -U 'ray[data]'

Load and transform data

import ray

# Read Parquet files
ds = ray.data.read_parquet("s3://bucket/data/*.parquet")

# Transform data (lazy execution)
ds = ds.map_batches(lambda batch: {"processed": batch["text"].str.lower()})

# Consume data
for batch in ds.iter_batches(batch_size=100):
    print(batch)

Integration with Ray Train

import ray
from ray.train import ScalingConfig
from ray.train.torch import TorchTrainer

# Create dataset
train_ds = ray.data.read_parquet("s3://bucket/train/*.parquet")

def train_func(config):
    # Access dataset in training
    train_ds = ray.train.get_dataset_shard("train")

    for epoch in range(10):
        for batch in train_ds.iter_batches(batch_size=32):
            # Train on batch
            pass

# Train with Ray
trainer = TorchTrainer(
    train_func,
    datasets={"train": train_ds},
    scaling_config=ScalingConfig(num_workers=4, use_gpu=True)
)
trainer.fit()

Reading data

From cloud storage

import ray

# Parquet (recommended for ML)
ds = ray.data.read_parquet("s3://bucket/data/*.parquet")