ray-data

用途

Ray Data 是一个面向 ML 与 AI 工作负载的分布式数据处理库，用于为机器学习训练和推理准备、转换和处理数据。根据资料，适用场景包括：为 ML 训练处理大型数据集（>100GB）；需要在集群上做分布式数据预处理；构建批量推理流水线；加载图片、音频、视频等多模态数据；以及从单机扩展到集群的数据处理。

资料同时给出了替代选择：小于 1GB 的单机小数据可使用 Pandas；表格数据与类似 SQL 的操作可使用 Dask；企业级 ETL 与 SQL 查询可使用 Spark。

主要能力

流式执行：可处理比内存更大的数据，通过 iter_batches 分批消费而非一次性加载。
GPU 支持：通过 num_gpus 参数在 map_batches 中请求 GPU 加速转换，例如图片预处理。
多格式支持：资料中的表格列出 Parquet（读/写，推荐用于 ML）、CSV（读/写）、JSON（读/写）、图片（读，不支持写）、NumPy（读/写）、Pandas（读，不支持写）。
数据类型来源：可从云存储（如 s3://、gs://）读取 Parquet、CSV、JSON、图片；也可从 Python 对象构造，包括 from_items 列表、range 合成数据、from_pandas。
转换操作：map_batches（向量化批处理，较快）、map（逐行，较慢）、filter、groupby 及 count 聚合、map_groups 自定义聚合。
框架集成：可与 Ray Train 配合，在训练函数中通过 ray.train.get_dataset_shard 获取数据分片；可转换为 PyTorch（to_torch）与 TensorFlow（to_tf）数据集。
性能优化手段：repartition 控制并行度（如 100 个块对应 100 核集群）、调整 batch_size、流式执行。

使用方式

安装：

pip install -U 'ray[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)

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

写入数据：

# Write to Parquet
ds.write_parquet("s3://bucket/output/")

# Write to CSV
ds.write_csv("output/")

# Write to JSON
ds.write_json("output/")

GPU 加速转换示例：

# Use GPU for preprocessing
def preprocess_images_gpu(batch):
    import torch
    images = torch.tensor(batch["image"]).cuda()
    # GPU preprocessing
    processed = images * 255
    return {"processed": processed.cpu().numpy()}

ds = ds.map_batches(
    preprocess_images_gpu,
    batch_size=64,
    num_gpus=1  # Request GPU
)

适用场景与常见模式

资料列出两类常见模式：

批量推理：先用 BatchInference 类在 worker 内加载模型，再通过 ds.map_batches(BatchInference, batch_size=32, num_gpus=1) 做分布式推理，最后 write_parquet 输出。
数据预处理流水线：对读取的数据依次 map_batches 执行 clean_data、tokenize、augment，再写入输出路径。

资料提到的生产用例包括：Pinterest 用于模型训练的最后一步数据处理；ByteDance 用于扩展多模态大模型的离线推理；Spotify 用于批量推理的 ML 平台。

性能参考

资料给出的基准数据（处理 100GB 数据）：单节点 16 核约 30 分钟；4 节点 64 核约 8 分钟；16 节点 256 核约 2 分钟。图片预处理：纯 CPU 约 1,000 张/秒；1 块 GPU 约 5,000 张/秒；4 块 GPU 约 18,000 张/秒。以上为资料中提供的参考数值。

限制与说明

图片格式仅支持读取，不支持写入；Pandas 仅支持读取，不支持写入。
转换采用惰性执行（lazy execution）。
map 为逐行处理，资料标注其比 map_batches 慢。
使用时需注意按数据规模与任务类型选择工具，小数据、表格类与 SQL 类任务可使用资料建议的替代方案。
版本信息：资料标注为 Ray 2.40.0+。

参考资源

Docs: https://docs.ray.io/en/latest/data/data.html GitHub: https://github.com/ray-project/ray Examples: https://docs.ray.io/en/latest/data/examples/overview.html