ray.data.Dataset#

class ray.data.Dataset(plan: ExecutionPlan, logical_plan: LogicalPlan)[source]#

A Dataset is a distributed data collection for data loading and processing.

Datasets are distributed pipelines that produce ObjectRef[Block] outputs, where each block holds data in Arrow format, representing a shard of the overall data collection. The block also determines the unit of parallelism. For more details, see Ray Data Internals.

Datasets can be created in multiple ways: from synthetic data via range_*() APIs, from existing memory data via from_*() APIs (this creates a subclass of Dataset called MaterializedDataset), or from external storage systems such as local disk, S3, HDFS etc. via the read_*() APIs. The (potentially processed) Dataset can be saved back to external storage systems via the write_*() APIs.

Examples

import ray
# Create dataset from synthetic data.
ds = ray.data.range(1000)
# Create dataset from in-memory data.
ds = ray.data.from_items(
    [{"col1": i, "col2": i * 2} for i in range(1000)]
)
# Create dataset from external storage system.
ds = ray.data.read_parquet("s3://bucket/path")
# Save dataset back to external storage system.
ds.write_csv("s3://bucket/output")

Dataset has two kinds of operations: transformation, which takes in Dataset and outputs a new Dataset (e.g. map_batches()); and consumption, which produces values (not a data stream) as output (e.g. iter_batches()).

Dataset transformations are lazy, with execution of the transformations being triggered by downstream consumption.

Dataset supports parallel processing at scale: transformations such as map_batches(), aggregations such as min()/max()/mean(), grouping via groupby(), shuffling operations such as sort(), random_shuffle(), and repartition().

Examples

>>> import ray
>>> ds = ray.data.range(1000)
>>> # Transform batches (Dict[str, np.ndarray]) with map_batches().
>>> ds.map_batches(lambda batch: {"id": batch["id"] * 2})  
MapBatches(<lambda>)
+- Dataset(num_rows=1000, schema={id: int64})
>>> # Compute the maximum.
>>> ds.max("id")
999
>>> # Shuffle this dataset randomly.
>>> ds.random_shuffle()  
RandomShuffle
+- Dataset(num_rows=1000, schema={id: int64})
>>> # Sort it back in order.
>>> ds.sort("id")  
Sort
+- Dataset(num_rows=1000, schema={id: int64})

Both unexecuted and materialized Datasets can be passed between Ray tasks and actors without incurring a copy. Dataset supports conversion to/from several more featureful dataframe libraries (e.g., Spark, Dask, Modin, MARS), and are also compatible with distributed TensorFlow / PyTorch.

Methods

__init__

Construct a Dataset (internal API).

add_column

Add the given column to the dataset.

aggregate

Aggregate values using one or more functions.

columns

Returns the columns of this Dataset.

count

Count the number of records in the dataset.

deserialize_lineage

Deserialize the provided lineage-serialized Dataset.

drop_columns

Drop one or more columns from the dataset.

filter

Filter out rows that don't satisfy the given predicate.

flat_map

Apply the given function to each row and then flatten results.

get_internal_block_refs

Get a list of references to the underlying blocks of this dataset.

groupby

Group rows of a Dataset according to a column.

has_serializable_lineage

Whether this dataset's lineage is able to be serialized for storage and later deserialized, possibly on a different cluster.

input_files

Return the list of input files for the dataset.

iter_batches

Return an iterable over batches of data.

iter_rows

Return an iterable over the rows in this dataset.

iter_tf_batches

Return an iterable over batches of data represented as TensorFlow tensors.

iter_torch_batches

Return an iterable over batches of data represented as Torch tensors.

iterator

Return a DataIterator over this dataset.

limit

Truncate the dataset to the first limit rows.

map

Apply the given function to each row of this dataset.

map_batches

Apply the given function to batches of data.

materialize

Execute and materialize this dataset into object store memory.

max

Return the maximum of one or more columns.

mean

Compute the mean of one or more columns.

min

Return the minimum of one or more columns.

num_blocks

Return the number of blocks of this Dataset.

random_sample

Returns a new Dataset containing a random fraction of the rows.

random_shuffle

Randomly shuffle the rows of this Dataset.

randomize_block_order

Randomly shuffle the blocks of this Dataset.

repartition

Repartition the Dataset into exactly this number of blocks.

schema

Return the schema of the dataset.

select_columns

Select one or more columns from the dataset.

serialize_lineage

Serialize this dataset's lineage, not the actual data or the existing data futures, to bytes that can be stored and later deserialized, possibly on a different cluster.

show

Print up to the given number of rows from the Dataset.

size_bytes

Return the in-memory size of the dataset.

sort

Sort the dataset by the specified key column or key function.

split

Materialize and split the dataset into n disjoint pieces.

split_at_indices

Materialize and split the dataset at the given indices (like np.split).

split_proportionately

Materialize and split the dataset using proportions.

stats

Returns a string containing execution timing information.

std

Compute the standard deviation of one or more columns.

streaming_split

Returns n DataIterators that can be used to read disjoint subsets of the dataset in parallel.

sum

Compute the sum of one or more columns.

take

Return up to limit rows from the Dataset.

take_all

Return all of the rows in this Dataset.

take_batch

Return up to batch_size rows from the Dataset in a batch.

to_arrow_refs

Convert this Dataset into a distributed set of PyArrow tables.

to_dask

Convert this Dataset into a Dask DataFrame.

to_mars

Convert this Dataset into a Mars DataFrame.

to_modin

Convert this Dataset into a Modin DataFrame.

to_numpy_refs

Converts this Dataset into a distributed set of NumPy ndarrays or dictionary of NumPy ndarrays.

to_pandas

Convert this Dataset to a single pandas DataFrame.

to_pandas_refs

Converts this Dataset into a distributed set of Pandas dataframes.

to_random_access_dataset

Convert this dataset into a distributed RandomAccessDataset (EXPERIMENTAL).

to_spark

Convert this Dataset into a Spark DataFrame.

to_tf

Return a TensorFlow Dataset over this Dataset.

to_torch

Return a Torch IterableDataset over this Dataset.

train_test_split

Materialize and split the dataset into train and test subsets.

union

Concatenate Datasets across rows.

unique

List the unique elements in a given column.

write_bigquery

Write the dataset to a BigQuery dataset table.

write_csv

Writes the Dataset to CSV files.

write_datasink

Writes the dataset to a custom Datasink.

write_datasource

Writes the dataset to a custom Datasource.

write_images

Writes the Dataset to images.

write_json

Writes the Dataset to JSON and JSONL files.

write_mongo

Writes the Dataset to a MongoDB database.

write_numpy

Writes a column of the Dataset to .npy files.

write_parquet

Writes the Dataset to parquet files under the provided path.

write_sql

Write to a database that provides a Python DB API2-compliant connector.

write_tfrecords

Write the Dataset to TFRecord files.

write_webdataset

Writes the dataset to WebDataset files.

zip

Materialize and zip the columns of this dataset with the columns of another.

Attributes

context

Return the DataContext used to create this Dataset.