Creating Datasets#

Ray Datasets can be created from:

generated synthetic data,
local and distributed in-memory data, and
local and external storage systems (local disk, cloud storage, HDFS, etc.).

This guide surveys the many ways to create a Dataset. If none of these meet your needs, please reach out to us on Discourse or open a feature request on the Ray GitHub repo, and check out our guide for implementing a custom Datasets datasource if you’re interested in rolling your own integration!

Generating Synthetic Data#

Int Range

Create a Dataset from a range of integers.

# Create a Dataset of Python objects.
ds = ray.data.range(10000)
# -> Dataset(num_blocks=200, num_rows=10000, schema=<class 'int'>)

ds.take(5)
# -> [0, 1, 2, 3, 4]

Tabular Range

Create an Arrow (tabular) Dataset from a range of integers, with a single column containing this integer range.

# Create a Dataset of Arrow records.
ds = ray.data.range_table(10000)
# -> Dataset(num_blocks=200, num_rows=10000, schema={value: int64})

ds.take(5)
# -> [{'value': 0}, {'value': 1}, {'value': 2}, {'value': 3}, {'value': 4}]

Tensor Range

Create a tensor dataset from a range of integers, packing this integer range into tensors of the provided shape.

# Create a Dataset of tensors.
ds = ray.data.range_tensor(100 * 64 * 64, shape=(64, 64))
# -> Dataset(
#       num_blocks=200,
#       num_rows=409600,
#       schema={value: <ArrowTensorType: shape=(64, 64), dtype=int64>}
#    )

ds.take(2)
# -> [array([[0, 0, 0, ..., 0, 0, 0],
#         [0, 0, 0, ..., 0, 0, 0],
#         [0, 0, 0, ..., 0, 0, 0],
#         ...,
#         [0, 0, 0, ..., 0, 0, 0],
#         [0, 0, 0, ..., 0, 0, 0],
#         [0, 0, 0, ..., 0, 0, 0]]),
#  array([[1, 1, 1, ..., 1, 1, 1],
#         [1, 1, 1, ..., 1, 1, 1],
#         [1, 1, 1, ..., 1, 1, 1],
#         ...,
#         [1, 1, 1, ..., 1, 1, 1],
#         [1, 1, 1, ..., 1, 1, 1],
#         [1, 1, 1, ..., 1, 1, 1]])]

From In-Memory Data#

Datasets can be constructed from existing in-memory data. In addition to being able to construct a Dataset from plain Python objects, Datasets also interoperates with popular single-node libraries (Pandas, NumPy, Arrow) as well as distributed frameworks (Dask, Spark, Modin, Mars).

From Single-Node Data Libraries#

In this section, we demonstrate creating a Dataset from single-node in-memory data.

Pandas

Create a Dataset from a Pandas DataFrame. This constructs a Dataset backed by a single Pandas DataFrame block.

import pandas as pd

# Create a tabular Dataset from a Pandas DataFrame.
df = pd.DataFrame({"col1": list(range(10000)), "col2": list(map(str, range(10000)))})
ds = ray.data.from_pandas(df)
# -> Dataset(num_blocks=1, num_rows=10000, schema={col1: int64, col2: object})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

We can also build a Dataset from more than one Pandas DataFrame, where each said DataFrame will become a block in the Dataset.

import pandas as pd

data = list(range(10000))
num_chunks = 10
chunk_size = len(data) // num_chunks
chunks = [data[i : i + chunk_size] for i in range(0, len(data), chunk_size)]
dfs = [
    pd.DataFrame({"col1": list(chunk), "col2": list(map(str, chunk))})
    for chunk in chunks
]
# Create a tabular Dataset from multiple Pandas DataFrames.
ds = ray.data.from_pandas(dfs)
# -> Dataset(num_blocks=10, num_rows=10000, schema={col1: int64, col2: object})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

NumPy

Create a Dataset from a NumPy ndarray. This constructs a Dataset backed by a single-column Arrow table block; the outer dimension of the ndarray will be treated as the row dimension, and the column will have name "__value__".

import numpy as np

# Create a tensor Dataset from a 3D NumPy ndarray.
arr = np.ones((3, 4, 4))
# The outer dimension is treated as the row dimension.
ds = ray.data.from_numpy(arr)
# -> Dataset(
#        num_blocks=1,
#        num_rows=3,
#        schema={value: <ArrowTensorType: shape=(4, 4), dtype=int64>},
#    )

ds.show(2)
# -> {'value': array([[1., 1., 1., 1.],
#        [1., 1., 1., 1.],
#        [1., 1., 1., 1.],
#        [1., 1., 1., 1.]])}
# -> {'value': array([[1., 1., 1., 1.],
#        [1., 1., 1., 1.],
#        [1., 1., 1., 1.],
#        [1., 1., 1., 1.]])}

We can also build a Dataset from more than one NumPy ndarray, where each said ndarray will become a single-column Arrow table block in the Dataset.

import numpy as np

# Create a tensor Dataset from multiple 3D NumPy ndarray.
arrs = [np.random.rand(2, 4, 4) for _ in range(4)]
# The outer dimension is treated as the row dimension.
ds = ray.data.from_numpy(arrs)
# -> Dataset(
#        num_blocks=4,
#        num_rows=8,
#        schema={value: <ArrowTensorType: shape=(4, 4), dtype=int64>},
#    )

ds.show(2)
# -> {'value': array([[0.06587483, 0.67808656, 0.76461924, 0.83428549],
#        [0.04932103, 0.25112165, 0.26476714, 0.24599738],
#        [0.67624391, 0.58689537, 0.12594709, 0.94663371],
#        [0.32435665, 0.97719096, 0.03234169, 0.71563231]])}
# -> {'value': array([[0.98570318, 0.65956399, 0.82168898, 0.09798336],
#        [0.22426704, 0.34209978, 0.02605247, 0.48200137],
#        [0.17312096, 0.38789983, 0.42663678, 0.92652456],
#        [0.80787394, 0.92437162, 0.11185822, 0.3319638 ]])}

Arrow

Create a Dataset from an Arrow Table. This constructs a Dataset backed by a single Arrow Table block.

import pyarrow as pa

# Create a tabular Dataset from an Arrow Table.
t = pa.table({"col1": list(range(10000)), "col2": list(map(str, range(10000)))})
ds = ray.data.from_arrow(t)
# -> Dataset(num_blocks=1, num_rows=10000, schema={col1: int64, col2: string})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

We can also build a Dataset from more than one Arrow Table, where each said Table will become a block in the Dataset.

import pyarrow as pa

data = list(range(10000))
num_chunks = 10
chunk_size = len(data) // num_chunks
chunks = [data[i : i + chunk_size] for i in range(0, len(data), chunk_size)]
ts = [
    pa.table({"col1": list(chunk), "col2": list(map(str, chunk))})
    for chunk in chunks
]
# Create a tabular Dataset from multiple Arrow Tables.
ds = ray.data.from_arrow(ts)
# -> Dataset(num_blocks=10, num_rows=10000, schema={col1: int64, col2: string})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

Python Objects

Create a Dataset from a list of Python objects; since each object in this particular list is a dictionary, Datasets will treat this list as a list of tabular records, and will construct an Arrow Dataset.

# Create a Dataset of tabular (Arrow) records.
ds = ray.data.from_items([{"col1": i, "col2": str(i)} for i in range(10000)])
# -> Dataset(num_blocks=200, num_rows=10000, schema={col1: int64, col2: string})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

From Distributed Data Processing Frameworks#

In addition to working with single-node in-memory data, Datasets can be constructed from distributed (multi-node) in-memory data, interoperating with popular distributed data processing frameworks such as Dask, Spark, Modin, and Mars.

These conversions work by running Ray tasks converting each Dask/Spark/Modin/Mars data partition to a block format supported by Datasets (copying data if needed), and using the futures representing the return value of those conversion tasks as the Dataset block futures.

Note

These data processing frameworks must be running on Ray in order for these Datasets integrations to work. See how these frameworks can be run on Ray in our data processing integrations docs.

Dask

Create a Dataset from a Dask DataFrame. This constructs a Dataset backed by the distributed Pandas DataFrame partitions that underly the Dask DataFrame.

This conversion has near-zero overhead, since Datasets simply reinterprets existing Dask-in-Ray partition objects as Dataset blocks.

import pandas as pd
import dask.dataframe as dd

df = pd.DataFrame({"col1": list(range(10000)), "col2": list(map(str, range(10000)))})
ddf = dd.from_pandas(df, npartitions=4)
# Create a tabular Dataset from a Dask DataFrame.
ds = ray.data.from_dask(ddf)
# -> Dataset(num_blocks=10, num_rows=10000, schema={col1: int64, col2: object})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

Spark

Create a Dataset from a Spark DataFrame. This constructs a Dataset backed by the distributed Spark DataFrame partitions that underly the Spark DataFrame. When this conversion happens, Spark-on-Ray (RayDP) will save the Spark DataFrame partitions to Ray’s object store in the Arrow format, which Datasets will then interpret as its blocks.

import raydp

spark = raydp.init_spark(app_name="Spark -> Datasets Example",
                         num_executors=2,
                         executor_cores=2,
                         executor_memory="500MB")
df = spark.createDataFrame([(i, str(i)) for i in range(10000)], ["col1", "col2"])
# Create a tabular Dataset from a Spark DataFrame.
ds = ray.data.from_dask(df)
# -> Dataset(num_blocks=10, num_rows=10000, schema={col1: int64, col2: string})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

Modin

Create a Dataset from a Modin DataFrame. This constructs a Dataset backed by the distributed Pandas DataFrame partitions that underly the Modin DataFrame.

This conversion has near-zero overhead, since Datasets simply reinterprets existing Modin partition objects as Dataset blocks.

import modin.pandas as md

df = pd.DataFrame({"col1": list(range(10000)), "col2": list(map(str, range(10000)))})
mdf = md.DataFrame(df)
# Create a tabular Dataset from a Modin DataFrame.
ds = ray.data.from_modin(mdf)
# -> Dataset(num_blocks=8, num_rows=10000, schema={col1: int64, col2: object})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

Mars

Create a Dataset from a Mars DataFrame. This constructs a Dataset backed by the distributed Pandas DataFrame partitions that underly the Mars DataFrame.

This conversion has near-zero overhead, since Datasets simply reinterprets existing Mars partition objects as Dataset blocks.

import mars
import mars.dataframe as md
import pandas as pd

cluster = mars.new_cluster_in_ray(worker_num=2, worker_cpu=1)

df = pd.DataFrame({"col1": list(range(10000)), "col2": list(map(str, range(10000)))})
mdf = md.DataFrame(df, num_partitions=8)
# Create a tabular Dataset from a Mars DataFrame.
ds = ray.data.from_mars(mdf)
# -> Dataset(num_blocks=8, num_rows=10000, schema={col1: int64, col2: object})

ds.show(3)
# -> {'col1': 0, 'col2': '0'}
# -> {'col1': 1, 'col2': '1'}
# -> {'col1': 2, 'col2': '2'}

From Torch and TensorFlow#

PyTorch

If you already have a Torch dataset available, you can create a Ray Dataset using from_torch.

Warning

from_torch doesn’t support parallel reads. You should only use this datasource for small datasets like MNIST or CIFAR.

import ray
import torchvision

dataset = torchvision.datasets.MNIST("data", download=True)
dataset = ray.data.from_torch(dataset)
dataset.take(1)
# (<PIL.Image.Image image mode=L size=28x28 at 0x1142CCA60>, 5)

TensorFlow

If you already have a TensorFlow dataset available, you can create a Ray Dataset using from_tf.

Warning

from_tf doesn’t support parallel reads. You should only use this function with small datasets like MNIST or CIFAR.

import ray
import tensorflow_datasets as tfds

dataset, _ = tfds.load("cifar10", split=["train", "test"])
dataset = ray.data.from_tf(dataset)

dataset
# -> Dataset(num_blocks=200, num_rows=50000, schema={id: binary, image: ArrowTensorType(shape=(32, 32, 3), dtype=uint8), label: int64})

From 🤗 (Hugging Face) Datasets#

You can convert 🤗 Datasets into Ray Datasets by using from_huggingface. This function accesses the underlying Arrow table and converts it into a Ray Dataset directly.

Warning

from_huggingface doesn’t support parallel reads. This will not usually be an issue with in-memory 🤗 Datasets, but may fail with large memory-mapped 🤗 Datasets. 🤗 IterableDataset objects are not supported.

import ray.data
from datasets import load_dataset

hf_datasets = load_dataset("wikitext", "wikitext-2-raw-v1")
ray_datasets = ray.data.from_huggingface(hf_datasets)
ray_datasets["train"].take(2)
# [{'text': ''}, {'text': ' = Valkyria Chronicles III = \n'}]

From MongoDB#

A Dataset can also be created from MongoDB with read_mongo. This interacts with MongoDB similar to external filesystems, except here you will need to specify the MongoDB source by its uri, database and collection, and specify a pipeline to run against the collection. The execution results are then used to create a Dataset.

Note

This example is not runnable as-is; you’ll need to point it at your MongoDB instance.

import ray

# Read a local MongoDB.
ds = ray.data.read_mongo(
    uri="mongodb://localhost:27017",
    database="my_db",
    collection="my_collection",
    pipeline=[{"$match": {"col": {"$gte": 0, "$lt": 10}}}, {"$sort": "sort_col"}],
)

# Reading a remote MongoDB is the same.
ds = ray.data.read_mongo(
    uri="mongodb://username:password@mongodb0.example.com:27017/?authSource=admin",
    database="my_db",
    collection="my_collection",
    pipeline=[{"$match": {"col": {"$gte": 0, "$lt": 10}}}, {"$sort": "sort_col"}],
)

# Write back to MongoDB.
ds.write_mongo(
    MongoDatasource(),
    uri="mongodb://username:password@mongodb0.example.com:27017/?authSource=admin",
    database="my_db",
    collection="my_collection",
)

Reading From SQL Databases#

Call read_sql() to read data from a database that provides a Python DB API2-compliant connector.

MySQL

To read from MySQL, install MySQL Connector/Python. It’s the first-party MySQL database connector.

pip install mysql-connector-python

Then, define your connection login and query the database.

import mysql.connector

import ray

def create_connection():
    return mysql.connector.connect(
        user="admin",
        password=...,
        host="example-mysql-database.c2c2k1yfll7o.us-west-2.rds.amazonaws.com",
        connection_timeout=30,
        database="example",
    )

# Get all movies
dataset = ray.data.read_sql("SELECT * FROM movie", create_connection)
# Get movies after the year 1980
dataset = ray.data.read_sql(
    "SELECT title, score FROM movie WHERE year >= 1980", create_connection
)
# Get the number of movies per year
dataset = ray.data.read_sql(
    "SELECT year, COUNT(*) FROM movie GROUP BY year", create_connection
)

PostgreSQL

To read from PostgreSQL, install Psycopg 2. It’s the most popular PostgreSQL database connector.

pip install psycopg2-binary

Then, define your connection login and query the database.

import psycopg2

import ray

def create_connection():
    return psycopg2.connect(
        user="postgres",
        password=...,
        host="example-postgres-database.c2c2k1yfll7o.us-west-2.rds.amazonaws.com",
        dbname="example",
    )

# Get all movies
dataset = ray.data.read_sql("SELECT * FROM movie", create_connection)
# Get movies after the year 1980
dataset = ray.data.read_sql(
    "SELECT title, score FROM movie WHERE year >= 1980", create_connection
)
# Get the number of movies per year
dataset = ray.data.read_sql(
    "SELECT year, COUNT(*) FROM movie GROUP BY year", create_connection
)

Snowflake

To read from Snowflake, install the Snowflake Connector for Python.

pip install snowflake-connector-python

Then, define your connection login and query the database.

import snowflake.connector

import ray

def create_connection():
    return snowflake.connector.connect(
        user=...,
        password=...
        account="ZZKXUVH-IPB52023",
        database="example",
    )

# Get all movies
dataset = ray.data.read_sql("SELECT * FROM movie", create_connection)
# Get movies after the year 1980
dataset = ray.data.read_sql(
    "SELECT title, score FROM movie WHERE year >= 1980", create_connection
)
# Get the number of movies per year
dataset = ray.data.read_sql(
    "SELECT year, COUNT(*) FROM movie GROUP BY year", create_connection
)

Databricks

To read from Databricks, install the Databricks SQL Connector for Python.

pip install databricks-sql-connector

Then, define your connection logic and read from the Databricks SQL warehouse.

from databricks import sql

import ray

def create_connection():
    return sql.connect(
        server_hostname="dbc-1016e3a4-d292.cloud.databricks.com",
        http_path="/sql/1.0/warehouses/a918da1fc0b7fed0",
        access_token=...,


# Get all movies
dataset = ray.data.read_sql("SELECT * FROM movie", create_connection)
# Get movies after the year 1980
dataset = ray.data.read_sql(
    "SELECT title, score FROM movie WHERE year >= 1980", create_connection
)
# Get the number of movies per year
dataset = ray.data.read_sql(
    "SELECT year, COUNT(*) FROM movie GROUP BY year", create_connection
)

BigQuery

To read from BigQuery, install the Python Client for Google BigQuery. This package includes a DB API2-compliant database connector.

pip install google-cloud-bigquery

Then, define your connection login and query the dataset.

from google.cloud import bigquery
from google.cloud.bigquery import dbapi

import ray

def create_connection():
    client = bigquery.Client(...)
    return dbapi.Connection(client)

# Get all movies
dataset = ray.data.read_sql("SELECT * FROM movie", create_connection)
# Get movies after the year 1980
dataset = ray.data.read_sql(
    "SELECT title, score FROM movie WHERE year >= 1980", create_connection
)
# Get the number of movies per year
dataset = ray.data.read_sql(
    "SELECT year, COUNT(*) FROM movie GROUP BY year", create_connection
)

Custom Datasources#

Datasets can read and write in parallel to custom datasources defined in Python. Once you have implemented YourCustomDataSource, you can use it like any other source in Ray Data:

# Read from a custom datasource.
ds = ray.data.read_datasource(YourCustomDatasource(), **read_args)

# Write to a custom datasource.
ds.write_datasource(YourCustomDatasource(), **write_args)

For more details, check out guide for implementing a custom Datasets datasource.

Performance Considerations#

Read Parallelism#

Datasets automatically selects the read parallelism according to the following procedure:

The number of available CPUs is estimated. If in a placement group, the number of CPUs in the cluster is scaled by the size of the placement group compared to the cluster size. If not in a placement group, this is the number of CPUs in the cluster.
The parallelism is set to the estimated number of CPUs multiplied by 2. If the parallelism is less than 8, it is set to 8.
The in-memory data size is estimated. If the parallelism would create in-memory blocks that are larger on average than the target block size (512MiB), the parallelism is increased until the blocks are < 512MiB in size.
The parallelism is truncated to min(num_files, parallelism).

To perform the read, parallelism parallel read tasks will be launched, each reading one or more files and each creating a single block of data. When reading from remote datasources, these parallel read tasks will be spread across the nodes in your Ray cluster, creating the distributed collection of blocks that makes up a distributed Ray Dataset.

This default parallelism can be overridden via the parallelism argument; see the performance guide for tips on how to tune this read parallelism.

Deferred Read Task Execution#

Datasets created via the ray.data.read_*() APIs are lazy: no read tasks are executed until a downstream consumption operation triggers execution. Metadata inspection functions like ds.schema() and ds.show() will trigger execution of only one or some tasks, instead of all tasks. This allows metadata to be inspected right away. Execution of all read tasks can be triggered manually using the ds.materialize() API.

Ray 2.4.0

Creating Datasets

Contents

Creating Datasets#

Generating Synthetic Data#

Reading Files From Storage#

Supported File Formats#

Reading from Remote Storage#

Reading from Local Storage#

Reading Compressed Files#

From In-Memory Data#

From Single-Node Data Libraries#

From Distributed Data Processing Frameworks#

From Torch and TensorFlow#

From 🤗 (Hugging Face) Datasets#

From MongoDB#

Reading From SQL Databases#

Custom Datasources#

Performance Considerations#

Read Parallelism#

Deferred Read Task Execution#