Loading data#

Ray Data loads data from various sources. This guide shows you how to read files such as images, load in-memory data such as pandas DataFrames, and read databases such as MySQL.

Read databases#

Ray Data reads from databases such as MySQL, PostgreSQL, MongoDB, and BigQuery.

Read 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 logic 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.

pip install psycopg2-binary

Then, define your connection logic 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 logic 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, set the DATABRICKS_TOKEN environment variable to your Databricks warehouse access token.

export DATABRICKS_TOKEN=...

If you’re not running your program on the Databricks runtime, also set the DATABRICKS_HOST environment variable.

export DATABRICKS_HOST=adb-<workspace-id>.<random-number>.azuredatabricks.net

Then, call ray.data.read_databricks_tables() to read from the Databricks SQL warehouse.

import ray

dataset = ray.data.read_databricks_tables(
    warehouse_id='...',  # Databricks SQL warehouse ID
    catalog='catalog_1',  # Unity catalog name
    schema='db_1',  # Schema name
    query="SELECT title, score FROM movie WHERE year >= 1980",
)

BigQuery

To read from BigQuery, install the Python Client for Google BigQuery and the Python Client for Google BigQueryStorage.

pip install google-cloud-bigquery
pip install google-cloud-bigquery-storage

To read data from BigQuery, call read_bigquery() and specify the project ID and either a dataset or a query.

import ray

# Read an entire table, specified as "dataset_id.table_id". Do not specify query.
ds = ray.data.read_bigquery(
    project_id="my_gcloud_project_id",
    dataset="my_dataset.my_table",
)

# Read from a SQL query of the dataset. Do not specify dataset.
ds = ray.data.read_bigquery(
    project_id="my_gcloud_project_id",
    query = "SELECT * FROM `bigquery-public-data.ml_datasets.iris` LIMIT 50",
)

# Write back to BigQuery
ds.write_bigquery(
    project_id="my_gcloud_project_id",
    dataset="destination_dataset.destination_table",
    overwrite_table=True,
)

Read from MongoDB#

To read data from MongoDB, call read_mongo() and specify the source URI, database, and collection. You can also pass an aggregation pipeline to run against the collection. Without one, Ray Data reads the entire collection.

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(
    uri="mongodb://username:password@mongodb0.example.com:27017/?authSource=admin",
    database="my_db",
    collection="my_collection",
)

Read from Kafka#

Ray Data reads from message queues such as Kafka.

To read data from Kafka topics, call read_kafka() and specify the topic names and broker addresses. Ray Data performs bounded reads between a start and end offset. You can specify each offset as an integer, as a datetime object for a time-based range, or as a dictionary of per-partition offsets that maps {topic: {partition_id: offset}}. The start offset also accepts "earliest", and the end offset also accepts "latest". Partitions that a per-partition dictionary doesn’t list fall back to "earliest" for the start offset and "latest" for the end offset.

First, install the required dependencies:

pip install confluent-kafka

Then, specify your Kafka configuration and read from topics.

import ray

# Read from a single topic with offset range
ds = ray.data.read_kafka(
    topics="my-topic",
    bootstrap_servers="localhost:9092",
    start_offset=0,
    end_offset=1000,
)

# Read from multiple topics
ds = ray.data.read_kafka(
    topics=["topic1", "topic2"],
    bootstrap_servers="localhost:9092",
    start_offset="earliest",
    end_offset="latest",
)

# Read messages within a datetime range (datetimes with no timezone info are treated as UTC)
from datetime import datetime
ds = ray.data.read_kafka(
    topics="my-topic",
    bootstrap_servers="localhost:9092",
    start_offset=datetime(2025, 1, 1),
    end_offset=datetime(2025, 1, 2),
)

# Read with authentication (Confluent/librdkafka options)
ds = ray.data.read_kafka(
    topics="secure-topic",
    bootstrap_servers="localhost:9092",
    consumer_config={
        "security.protocol": "SASL_SSL",
        "sasl.mechanism": "PLAIN",
        "sasl.username": "your-username",
        "sasl.password": "your-password",
    },
)

print(ds.schema())

Column          Type
------          ----
offset          int64
key             binary
value           binary
topic           string
partition       int32
timestamp       int64
timestamp_type  int32
headers         map<string, binary>

Create synthetic data#

Synthetic datasets are useful for testing and benchmarking.

Integer range

To create a synthetic Dataset from a range of integers, call range(). Ray Data stores the integer range in a single column called id.

import ray

ds = ray.data.range(10000)

print(ds.schema())

Column  Type
------  ----
id      int64

Tensor range

To create a synthetic Dataset containing arrays, call range_tensor(). Ray Data packs an integer range into ndarrays of the provided shape. In the schema, the column name defaults to data.

import ray

ds = ray.data.range_tensor(10, shape=(64, 64))

print(ds.schema())

Column  Type
------  ----
data    ArrowTensorTypeV2(shape=(64, 64), dtype=int64)

Load other datasources#

If Ray Data can’t load your data, subclass Datasource. Then, construct an instance of your custom datasource and pass it to read_datasource(). To write results, you might also need to subclass ray.data.Datasink. Then, create an instance of your custom datasink and pass it to write_datasink(). For more details, see Advanced: Read and write custom file types.

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

# Write to a custom datasink.
ds.write_datasink(YourCustomDatasink())

Community-maintained connectors#

The community maintains the following connectors, which integrate Ray Data with additional data systems:

Apache Doris Ray Connector: Reads and writes data between Ray Data and Apache Doris.
Kinetica Ray Connector: Reads and writes data between Ray Data and Kinetica.

Performance considerations#

By default, Ray Data decides the number of output blocks from all read tasks dynamically, based on input data size and available resources. This default should work well in most cases. To override it, set the override_num_blocks argument. Ray Data decides internally how many read tasks to run concurrently to make the best use of the cluster, from 1 to override_num_blocks tasks. The higher the override_num_blocks value, the smaller the data blocks in the dataset, and the more opportunities for parallel execution.

To tune the number of output blocks and find other ways to optimize read performance, see Optimize reads.

Loading data#

Read files#

Read files from shared local storage#

Read files from cloud storage#

Read files from the Hadoop Distributed File System#

Handle compressed files#

Download files from URIs#

Load data from other libraries#

Load data from single-node data libraries#

Load data from distributed DataFrame libraries#

Load Hugging Face datasets#

Load data from ML libraries#

Read databases#

Read SQL databases#

Read from MongoDB#

Read from Kafka#

Create synthetic data#

Load other datasources#

Community-maintained connectors#

Performance considerations#