TensorDictBase.from_parquet#

Creates a TensorDict from a Parquet file.

Requires either pyarrow or pandas to be installed. Prefers pyarrow when available for better performance.

Parameters:

path (str or Path) – Path to the Parquet file.

Keyword Arguments:

auto_batch_size (bool, optional) – If True, the batch size will be computed automatically. Defaults to False.
batch_dims (int, optional) – If auto_batch_size is True, defines how many dimensions the output tensordict should have. Defaults to None.
device (device, optional) – The device for tensor data. Defaults to None.
batch_size (Size, optional) – The batch size. Defaults to [num_rows].
separator (str, optional) – If provided, column names are split on this separator to create nested TensorDicts. Defaults to None.
dtype (dtype, optional) – If provided, all numeric columns are cast to this dtype. Defaults to None.
columns (list of str, optional) – If provided, only read these columns from the file. Defaults to None (all columns).
**kwargs – Additional keyword arguments forwarded to the Parquet reader.

Returns:

A TensorDict representation of the Parquet data.

Examples

>>> import os
>>> import tempfile
>>> import torch
>>> from tensordict import TensorDict
>>> data = TensorDict(obs=torch.arange(3.0), reward=torch.ones(3), batch_size=[3])
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     path = os.path.join(tmpdir, "data.parquet")
...     data.to_parquet(path)
...     td = TensorDict.from_parquet(path)
...     td_obs = TensorDict.from_parquet(path, columns=["obs"])
>>> sorted(td.keys()), sorted(td_obs.keys())
(['obs', 'reward'], ['obs'])

Every tensordict class and every tensorclass has this method. For the other methods, see Methods by task.

TensorDictBase.from_parquet#

Docs

Tutorials

Resources