def from_pandas_refs(
dfs: Union[ObjectRef["pandas.DataFrame"], List[ObjectRef["pandas.DataFrame"]]]
) -> Dataset[ArrowRow]:
"""Create a dataset from a list of Ray object references to Pandas
dataframes.
Args:
dfs: A Ray object references to pandas dataframe, or a list of
Ray object references to pandas dataframes.
Returns:
Dataset holding Arrow records read from the dataframes.
"""
if isinstance(dfs, ray.ObjectRef):
dfs = [dfs]
elif isinstance(dfs, list):
for df in dfs:
if not isinstance(df, ray.ObjectRef):
raise ValueError(
"Expected list of Ray object refs, "
f"got list containing {type(df)}"
)
else:
raise ValueError(
"Expected Ray object ref or list of Ray object refs, " f"got {type(df)}"
)
context = DatasetContext.get_current()
if context.enable_pandas_block:
get_metadata = cached_remote_fn(_get_metadata)
metadata = ray.get([get_metadata.remote(df) for df in dfs])
return Dataset(
ExecutionPlan(
BlockList(dfs, metadata),
DatasetStats(stages={"from_pandas_refs": metadata}, parent=None),
),
0,
False,
)
df_to_block = cached_remote_fn(_df_to_block, num_returns=2)
res = [df_to_block.remote(df) for df in dfs]
blocks, metadata = map(list, zip(*res))
metadata = ray.get(metadata)
return Dataset(
ExecutionPlan(
BlockList(blocks, metadata),
DatasetStats(stages={"from_pandas_refs": metadata}, parent=None),
),
0,
False,
)
def _test_equal_split_balanced(block_sizes, num_splits):
blocks = []
metadata = []
total_rows = 0
for block_size in block_sizes:
block = list(range(total_rows, total_rows + block_size))
blocks.append(ray.put(block))
metadata.append(
BlockAccessor.for_block(block).get_metadata(None, None))
total_rows += block_size
block_list = BlockList(blocks, metadata)
ds = Dataset(
ExecutionPlan(block_list, DatasetStats.TODO()),
0,
False,
)
splits = ds.split(num_splits, equal=True)
split_counts = [split.count() for split in splits]
assert len(split_counts) == num_splits
expected_block_size = total_rows // num_splits
# Check that all splits are the expected size.
assert all([count == expected_block_size for count in split_counts])
expected_total_rows = sum(split_counts)
# Check that the expected number of rows were dropped.
assert total_rows - expected_total_rows == total_rows % num_splits
# Check that all rows are unique (content check).
split_rows = [row for split in splits for row in split.take(total_rows)]
assert len(set(split_rows)) == len(split_rows)
def from_arrow_refs(
tables: Union[ObjectRef[Union["pyarrow.Table", bytes]],
List[ObjectRef[Union["pyarrow.Table", bytes]]], ]
) -> Dataset[ArrowRow]:
"""Create a dataset from a set of Arrow tables.
Args:
tables: A Ray object reference to Arrow table, or list of Ray object
references to Arrow tables, or its streaming format in bytes.
Returns:
Dataset holding Arrow records from the tables.
"""
if isinstance(tables, ray.ObjectRef):
tables = [tables]
get_metadata = cached_remote_fn(_get_metadata)
metadata = ray.get([get_metadata.remote(t) for t in tables])
return Dataset(
ExecutionPlan(
BlockList(tables, metadata),
DatasetStats(stages={"from_arrow_refs": metadata}, parent=None),
),
0,
False,
)
def _optimize_stages(self):
"""Optimize this pipeline, fusing stages together as possible."""
context = DatasetContext.get_current()
if not context.optimize_fuse_stages:
self._optimized_stages = self._stages
return
# This dummy dataset will be used to get a set of optimized stages.
dummy_ds = Dataset(
ExecutionPlan(BlockList([], []),
DatasetStats(stages={}, parent=None)),
0,
True,
)
# Apply all pipeline operations to the dummy dataset.
for stage in self._stages:
dummy_ds = stage(dummy_ds)
# Get the optimized stages.
_, _, stages = dummy_ds._plan._optimize()
# Apply these optimized stages to the datasets underlying the pipeline.
# These optimized stages will be executed by the PipelineExecutor.
optimized_stages = []
for stage in stages:
optimized_stages.append(lambda ds, stage=stage: Dataset(
ds._plan.with_stage(stage), ds._epoch, True))
self._optimized_stages = optimized_stages
def from_items(items: List[Any], *, parallelism: int = -1) -> Dataset[Any]:
"""Create a dataset from a list of local Python objects.
Examples:
>>> import ray
>>> ds = ray.data.from_items([1, 2, 3, 4, 5]) # doctest: +SKIP
>>> ds # doctest: +SKIP
Dataset(num_blocks=5, num_rows=5, schema=<class 'int'>)
>>> ds.take(2) # doctest: +SKIP
[1, 2]
Args:
items: List of local Python objects.
parallelism: The amount of parallelism to use for the dataset.
Parallelism may be limited by the number of items.
Returns:
Dataset holding the items.
"""
detected_parallelism, _ = _autodetect_parallelism(
parallelism,
ray.util.get_current_placement_group(),
DatasetContext.get_current(),
)
block_size = max(
1,
len(items) // detected_parallelism,
)
blocks: List[ObjectRef[Block]] = []
metadata: List[BlockMetadata] = []
i = 0
while i < len(items):
stats = BlockExecStats.builder()
builder = DelegatingBlockBuilder()
for item in items[i:i + block_size]:
builder.add(item)
block = builder.build()
blocks.append(ray.put(block))
metadata.append(
BlockAccessor.for_block(block).get_metadata(
input_files=None, exec_stats=stats.build()))
i += block_size
return Dataset(
ExecutionPlan(
BlockList(blocks, metadata),
DatasetStats(stages={"from_items": metadata}, parent=None),
),
0,
False,
)
def from_numpy_refs(
ndarrays: Union[ObjectRef[np.ndarray], List[ObjectRef[np.ndarray]]],
) -> Dataset[ArrowRow]:
"""Create a dataset from a list of NumPy ndarray futures.
Args:
ndarrays: A Ray object reference to a NumPy ndarray or a list of Ray object
references to NumPy ndarrays.
Returns:
Dataset holding the given ndarrays.
"""
if isinstance(ndarrays, ray.ObjectRef):
ndarrays = [ndarrays]
elif isinstance(ndarrays, list):
for ndarray in ndarrays:
if not isinstance(ndarray, ray.ObjectRef):
raise ValueError(
"Expected list of Ray object refs, "
f"got list containing {type(ndarray)}"
)
else:
raise ValueError(
f"Expected Ray object ref or list of Ray object refs, got {type(ndarray)}"
)
ndarray_to_block = cached_remote_fn(_ndarray_to_block, num_returns=2)
res = [ndarray_to_block.remote(ndarray) for ndarray in ndarrays]
blocks, metadata = map(list, zip(*res))
metadata = ray.get(metadata)
return Dataset(
ExecutionPlan(
BlockList(blocks, metadata),
DatasetStats(stages={"from_numpy_refs": metadata}, parent=None),
),
0,
False,
)
def read_datasource(
datasource: Datasource[T],
*,
parallelism: int = -1,
ray_remote_args: Dict[str, Any] = None,
**read_args,
) -> Dataset[T]:
"""Read a dataset from a custom data source.
Args:
datasource: The datasource to read data from.
parallelism: The requested parallelism of the read. Parallelism may be
limited by the available partitioning of the datasource. If set to -1,
parallelism will be automatically chosen based on the available cluster
resources and estimated in-memory data size.
read_args: Additional kwargs to pass to the datasource impl.
ray_remote_args: kwargs passed to ray.remote in the read tasks.
Returns:
Dataset holding the data read from the datasource.
"""
ctx = DatasetContext.get_current()
# TODO(ekl) remove this feature flag.
force_local = "RAY_DATASET_FORCE_LOCAL_METADATA" in os.environ
cur_pg = ray.util.get_current_placement_group()
pa_ds = _lazy_import_pyarrow_dataset()
if pa_ds:
partitioning = read_args.get("dataset_kwargs",
{}).get("partitioning", None)
if isinstance(partitioning, pa_ds.Partitioning):
logger.info(
"Forcing local metadata resolution since the provided partitioning "
f"{partitioning} is not serializable.")
force_local = True
if force_local:
requested_parallelism, min_safe_parallelism, read_tasks = _get_read_tasks(
datasource, ctx, cur_pg, parallelism, read_args)
else:
# Prepare read in a remote task so that in Ray client mode, we aren't
# attempting metadata resolution from the client machine.
get_read_tasks = cached_remote_fn(_get_read_tasks,
retry_exceptions=False,
num_cpus=0)
requested_parallelism, min_safe_parallelism, read_tasks = ray.get(
get_read_tasks.remote(
datasource,
ctx,
cur_pg,
parallelism,
_wrap_and_register_arrow_serialization_workaround(read_args),
))
if read_tasks and len(read_tasks) < min_safe_parallelism * 0.7:
perc = 1 + round(
(min_safe_parallelism - len(read_tasks)) / len(read_tasks), 1)
logger.warning(
f"{WARN_PREFIX} The blocks of this dataset are estimated to be {perc}x "
"larger than the target block size "
f"of {int(ctx.target_max_block_size / 1024 / 1024)} MiB. This may lead to "
"out-of-memory errors during processing. Consider reducing the size of "
"input files or using `.repartition(n)` to increase the number of "
"dataset blocks.")
elif len(read_tasks) < requested_parallelism and (
len(read_tasks) < ray.available_resources().get("CPU", 1) // 2):
logger.warning(
f"{WARN_PREFIX} The number of blocks in this dataset ({len(read_tasks)}) "
f"limits its parallelism to {len(read_tasks)} concurrent tasks. "
"This is much less than the number "
"of available CPU slots in the cluster. Use `.repartition(n)` to "
"increase the number of "
"dataset blocks.")
if ray_remote_args is None:
ray_remote_args = {}
if ("scheduling_strategy" not in ray_remote_args
and ctx.scheduling_strategy == DEFAULT_SCHEDULING_STRATEGY):
ray_remote_args["scheduling_strategy"] = "SPREAD"
block_list = LazyBlockList(read_tasks, ray_remote_args=ray_remote_args)
block_list.compute_first_block()
block_list.ensure_metadata_for_first_block()
return Dataset(
ExecutionPlan(block_list, block_list.stats()),
0,
False,
)
def fast_repartition(blocks, num_blocks):
from ray.data.dataset import Dataset
wrapped_ds = Dataset(
ExecutionPlan(blocks, DatasetStats(stages={}, parent=None)), 0, lazy=False
)
# Compute the (n-1) indices needed for an equal split of the data.
count = wrapped_ds.count()
dataset_format = wrapped_ds._dataset_format()
indices = []
cur_idx = 0
for _ in range(num_blocks - 1):
cur_idx += count / num_blocks
indices.append(int(cur_idx))
assert len(indices) < num_blocks, (indices, num_blocks)
if indices:
splits = wrapped_ds.split_at_indices(indices)
else:
splits = [wrapped_ds]
# TODO(ekl) include stats for the split tasks. We may also want to
# consider combining the split and coalesce tasks as an optimization.
# Coalesce each split into a single block.
reduce_task = cached_remote_fn(_ShufflePartitionOp.reduce).options(num_returns=2)
reduce_bar = ProgressBar("Repartition", position=0, total=len(splits))
reduce_out = [
reduce_task.remote(False, None, *s.get_internal_block_refs())
for s in splits
if s.num_blocks() > 0
]
# Early-release memory.
del splits, blocks, wrapped_ds
new_blocks, new_metadata = zip(*reduce_out)
new_blocks, new_metadata = list(new_blocks), list(new_metadata)
new_metadata = reduce_bar.fetch_until_complete(new_metadata)
reduce_bar.close()
# Handle empty blocks.
if len(new_blocks) < num_blocks:
from ray.data._internal.arrow_block import ArrowBlockBuilder
from ray.data._internal.pandas_block import PandasBlockBuilder
from ray.data._internal.simple_block import SimpleBlockBuilder
num_empties = num_blocks - len(new_blocks)
if dataset_format == "arrow":
builder = ArrowBlockBuilder()
elif dataset_format == "pandas":
builder = PandasBlockBuilder()
else:
builder = SimpleBlockBuilder()
empty_block = builder.build()
empty_meta = BlockAccessor.for_block(empty_block).get_metadata(
input_files=None, exec_stats=None
) # No stats for empty block.
empty_blocks, empty_metadata = zip(
*[(ray.put(empty_block), empty_meta) for _ in range(num_empties)]
)
new_blocks += empty_blocks
new_metadata += empty_metadata
return BlockList(new_blocks, new_metadata), {}
def read_datasource(
datasource: Datasource[T],
*,
parallelism: int = 200,
ray_remote_args: Dict[str, Any] = None,
**read_args,
) -> Dataset[T]:
"""Read a dataset from a custom data source.
Args:
datasource: The datasource to read data from.
parallelism: The requested parallelism of the read. Parallelism may be
limited by the available partitioning of the datasource.
read_args: Additional kwargs to pass to the datasource impl.
ray_remote_args: kwargs passed to ray.remote in the read tasks.
Returns:
Dataset holding the data read from the datasource.
"""
ctx = DatasetContext.get_current()
# TODO(ekl) remove this feature flag.
force_local = "RAY_DATASET_FORCE_LOCAL_METADATA" in os.environ
pa_ds = _lazy_import_pyarrow_dataset()
if pa_ds:
partitioning = read_args.get("dataset_kwargs", {}).get("partitioning", None)
if isinstance(partitioning, pa_ds.Partitioning):
logger.info(
"Forcing local metadata resolution since the provided partitioning "
f"{partitioning} is not serializable."
)
force_local = True
if force_local:
read_tasks = datasource.prepare_read(parallelism, **read_args)
else:
# Prepare read in a remote task so that in Ray client mode, we aren't
# attempting metadata resolution from the client machine.
prepare_read = cached_remote_fn(
_prepare_read, retry_exceptions=False, num_cpus=0
)
read_tasks = ray.get(
prepare_read.remote(
datasource,
ctx,
parallelism,
_wrap_and_register_arrow_serialization_workaround(read_args),
)
)
if len(read_tasks) < parallelism and (
len(read_tasks) < ray.available_resources().get("CPU", 1) // 2
):
logger.warning(
"The number of blocks in this dataset ({}) limits its parallelism to {} "
"concurrent tasks. This is much less than the number of available "
"CPU slots in the cluster. Use `.repartition(n)` to increase the number of "
"dataset blocks.".format(len(read_tasks), len(read_tasks))
)
if ray_remote_args is None:
ray_remote_args = {}
if (
"scheduling_strategy" not in ray_remote_args
and ctx.scheduling_strategy == DEFAULT_SCHEDULING_STRATEGY
):
ray_remote_args["scheduling_strategy"] = "SPREAD"
block_list = LazyBlockList(read_tasks, ray_remote_args=ray_remote_args)
block_list.compute_first_block()
block_list.ensure_metadata_for_first_block()
return Dataset(
ExecutionPlan(block_list, block_list.stats()),
0,
False,
)