wehub-resource-sync
449 lines
16 KiB
Python
449 lines
16 KiB
Python
import warnings
|
|
from functools import partial
|
|
from typing import TYPE_CHECKING, Callable, Dict, List, Optional, Tuple, Type, TypeVar
|
|
|
|
if TYPE_CHECKING:
|
|
import pyarrow.fs
|
|
|
|
from ray.data._internal.execution.execution_callback import ExecutionCallback
|
|
from ray.data._internal.execution.interfaces import PhysicalOperator
|
|
from ray.data._internal.execution.operators.aggregate_num_rows import (
|
|
AggregateNumRows,
|
|
)
|
|
from ray.data._internal.execution.operators.hash_shuffle_v2 import (
|
|
_SHUFFLE_MAP_RUNTIME_ENV,
|
|
_make_hash_partition_fn,
|
|
)
|
|
from ray.data._internal.execution.operators.input_data_buffer import (
|
|
InputDataBuffer,
|
|
)
|
|
from ray.data._internal.execution.operators.join import (
|
|
JoinOperator,
|
|
_make_join_reduce_fn,
|
|
)
|
|
from ray.data._internal.execution.operators.limit_operator import LimitOperator
|
|
from ray.data._internal.execution.operators.mix_operator import MixOperator
|
|
from ray.data._internal.execution.operators.output_splitter import OutputSplitter
|
|
from ray.data._internal.execution.operators.shuffle_operators.shuffle_map_operator import (
|
|
ShuffleMapOp,
|
|
)
|
|
from ray.data._internal.execution.operators.shuffle_operators.shuffle_reduce_operator import (
|
|
ShuffleReduceOp,
|
|
)
|
|
from ray.data._internal.execution.operators.union_operator import UnionOperator
|
|
from ray.data._internal.execution.operators.zip_operator import ZipOperator
|
|
from ray.data._internal.logical.interfaces import (
|
|
LogicalOperator,
|
|
LogicalPlan,
|
|
PhysicalPlan,
|
|
)
|
|
from ray.data._internal.logical.operators import (
|
|
AbstractAllToAll,
|
|
AbstractFrom,
|
|
AbstractUDFMap,
|
|
Count,
|
|
Download,
|
|
Filter,
|
|
InputData,
|
|
Join,
|
|
JoinType,
|
|
Limit,
|
|
ListFiles,
|
|
Mix,
|
|
Project,
|
|
Read,
|
|
ReadFiles,
|
|
StreamingRepartition,
|
|
StreamingSplit,
|
|
Union,
|
|
Write,
|
|
Zip,
|
|
)
|
|
from ray.data._internal.planner.checkpoint import (
|
|
plan_read_files_op_with_checkpoint_filter,
|
|
plan_read_op_with_checkpoint_filter,
|
|
plan_write_op_with_checkpoint_writer,
|
|
)
|
|
from ray.data._internal.planner.plan_all_to_all_op import plan_all_to_all_op
|
|
from ray.data._internal.planner.plan_download_op import plan_download_op
|
|
from ray.data._internal.planner.plan_list_files_op import plan_list_files_op
|
|
from ray.data._internal.planner.plan_read_files_op import plan_read_files_op
|
|
from ray.data._internal.planner.plan_read_op import plan_read_op
|
|
from ray.data._internal.planner.plan_udf_map_op import (
|
|
plan_filter_op,
|
|
plan_project_op,
|
|
plan_streaming_repartition_op,
|
|
plan_udf_map_op,
|
|
)
|
|
from ray.data._internal.planner.plan_write_op import plan_write_op
|
|
from ray.data._internal.usage import create_usage_callback
|
|
from ray.data.checkpoint.load_checkpoint_callback import LoadCheckpointCallback
|
|
from ray.data.context import DataContext
|
|
from ray.data.datasource.file_datasink import _FileDatasink
|
|
|
|
LogicalOperatorType = TypeVar("LogicalOperatorType", bound=LogicalOperator)
|
|
PlanLogicalOpFn = Callable[
|
|
[LogicalOperatorType, List[PhysicalOperator], DataContext], PhysicalOperator
|
|
]
|
|
|
|
|
|
def plan_input_data_op(
|
|
logical_op: InputData,
|
|
physical_children: List[PhysicalOperator],
|
|
data_context: DataContext,
|
|
) -> PhysicalOperator:
|
|
"""Get the corresponding DAG of physical operators for InputData."""
|
|
assert len(physical_children) == 0
|
|
|
|
return InputDataBuffer(
|
|
data_context,
|
|
input_data=logical_op.input_data,
|
|
)
|
|
|
|
|
|
def plan_from_op(
|
|
op: AbstractFrom,
|
|
physical_children: List[PhysicalOperator],
|
|
data_context: DataContext,
|
|
) -> PhysicalOperator:
|
|
assert len(physical_children) == 0
|
|
return InputDataBuffer(data_context, op.input_data)
|
|
|
|
|
|
def plan_zip_op(_, physical_children, data_context):
|
|
assert len(physical_children) >= 2
|
|
return ZipOperator(data_context, *physical_children)
|
|
|
|
|
|
def plan_mix_op(logical_op, physical_children, data_context):
|
|
assert len(physical_children) >= 1
|
|
return MixOperator(
|
|
data_context,
|
|
*physical_children,
|
|
weights=logical_op.weights,
|
|
stopping_condition=logical_op.stopping_condition,
|
|
)
|
|
|
|
|
|
def plan_union_op(_, physical_children, data_context):
|
|
assert len(physical_children) >= 2
|
|
return UnionOperator(data_context, *physical_children)
|
|
|
|
|
|
def plan_limit_op(logical_op, physical_children, data_context):
|
|
assert len(physical_children) == 1
|
|
return LimitOperator(logical_op.limit, physical_children[0], data_context)
|
|
|
|
|
|
def plan_count_op(logical_op, physical_children, data_context):
|
|
assert len(physical_children) == 1
|
|
return AggregateNumRows(
|
|
[physical_children[0]], data_context, column_name=Count.COLUMN_NAME
|
|
)
|
|
|
|
|
|
def _plan_join_shuffle_v2(
|
|
logical_op: Join,
|
|
physical_children: List[PhysicalOperator],
|
|
data_context: DataContext,
|
|
) -> PhysicalOperator:
|
|
left_keys = list(logical_op.left_key_columns)
|
|
right_keys = list(logical_op.right_key_columns)
|
|
num_partitions = logical_op.num_partitions
|
|
join_type = JoinType(logical_op.join_type)
|
|
|
|
left_map = ShuffleMapOp(
|
|
physical_children[0],
|
|
data_context,
|
|
num_partitions=num_partitions,
|
|
partition_fn=_make_hash_partition_fn(left_keys, num_partitions),
|
|
map_runtime_env=_SHUFFLE_MAP_RUNTIME_ENV,
|
|
name=f"JoinShuffleMapLeft(keys={tuple(left_keys)}, parts={num_partitions})",
|
|
)
|
|
right_map = ShuffleMapOp(
|
|
physical_children[1],
|
|
data_context,
|
|
num_partitions=num_partitions,
|
|
partition_fn=_make_hash_partition_fn(right_keys, num_partitions),
|
|
map_runtime_env=_SHUFFLE_MAP_RUNTIME_ENV,
|
|
name=f"JoinShuffleMapRight(keys={tuple(right_keys)}, parts={num_partitions})",
|
|
)
|
|
|
|
reduce_fn = _make_join_reduce_fn(
|
|
join_type=join_type,
|
|
left_key_col_names=tuple(left_keys),
|
|
right_key_col_names=tuple(right_keys),
|
|
left_columns_suffix=logical_op.left_columns_suffix,
|
|
right_columns_suffix=logical_op.right_columns_suffix,
|
|
left_schema=logical_op.input_dependencies[0].infer_schema(),
|
|
right_schema=logical_op.input_dependencies[1].infer_schema(),
|
|
)
|
|
return ShuffleReduceOp(
|
|
[left_map, right_map],
|
|
data_context,
|
|
num_partitions=num_partitions,
|
|
reduce_fn=reduce_fn,
|
|
disallow_block_splitting=False,
|
|
reduce_ray_remote_args=logical_op.aggregator_ray_remote_args,
|
|
name=f"JoinShuffleReduce(num_partitions={num_partitions})",
|
|
)
|
|
|
|
|
|
def plan_join_op(
|
|
logical_op: Join,
|
|
physical_children: List[PhysicalOperator],
|
|
data_context: DataContext,
|
|
) -> PhysicalOperator:
|
|
assert len(physical_children) == 2
|
|
if data_context.use_hash_shuffle_v2:
|
|
return _plan_join_shuffle_v2(logical_op, physical_children, data_context)
|
|
return JoinOperator(
|
|
data_context=data_context,
|
|
left_input_op=physical_children[0],
|
|
right_input_op=physical_children[1],
|
|
join_type=logical_op.join_type,
|
|
left_key_columns=logical_op.left_key_columns,
|
|
right_key_columns=logical_op.right_key_columns,
|
|
left_columns_suffix=logical_op.left_columns_suffix,
|
|
right_columns_suffix=logical_op.right_columns_suffix,
|
|
num_partitions=logical_op.num_outputs,
|
|
partition_size_hint=logical_op.partition_size_hint,
|
|
aggregator_ray_remote_args_override=logical_op.aggregator_ray_remote_args,
|
|
)
|
|
|
|
|
|
def plan_streaming_split_op(
|
|
logical_op: StreamingSplit,
|
|
physical_children: List[PhysicalOperator],
|
|
data_context: DataContext,
|
|
):
|
|
assert len(physical_children) == 1
|
|
return OutputSplitter(
|
|
physical_children[0],
|
|
n=logical_op.num_splits,
|
|
equal=logical_op.equal,
|
|
data_context=data_context,
|
|
locality_hints=logical_op.locality_hints,
|
|
)
|
|
|
|
|
|
class Planner:
|
|
"""The planner to convert optimized logical to physical operators.
|
|
|
|
Note that planner is only doing operators conversion. Physical optimization work is
|
|
done by physical optimizer.
|
|
"""
|
|
|
|
_DEFAULT_PLAN_FNS = {
|
|
Read: plan_read_op,
|
|
ReadFiles: plan_read_files_op,
|
|
ListFiles: plan_list_files_op,
|
|
InputData: plan_input_data_op,
|
|
Write: plan_write_op,
|
|
AbstractFrom: plan_from_op,
|
|
Filter: plan_filter_op,
|
|
AbstractUDFMap: plan_udf_map_op,
|
|
AbstractAllToAll: plan_all_to_all_op,
|
|
Mix: plan_mix_op,
|
|
Union: plan_union_op,
|
|
Zip: plan_zip_op,
|
|
Limit: plan_limit_op,
|
|
Count: plan_count_op,
|
|
Project: plan_project_op,
|
|
StreamingRepartition: plan_streaming_repartition_op,
|
|
Join: plan_join_op,
|
|
StreamingSplit: plan_streaming_split_op,
|
|
Download: plan_download_op,
|
|
}
|
|
# Operators that support checkpoint filtering. Subclasses can override.
|
|
_CHECKPOINT_FILTER_OPS = (Read, ReadFiles)
|
|
|
|
def __init__(self):
|
|
self._supports_checkpointing = False
|
|
self._plan_fns_for_checkpointing = {}
|
|
|
|
def plan(
|
|
self, logical_plan: LogicalPlan
|
|
) -> Tuple[PhysicalPlan, List["ExecutionCallback"]]:
|
|
"""Convert logical to physical operators recursively in post-order."""
|
|
checkpoint_config = logical_plan.context.checkpoint_config
|
|
|
|
callbacks = [cls() for cls in logical_plan.context.execution_callback_classes]
|
|
callbacks.append(create_usage_callback(logical_plan))
|
|
|
|
if checkpoint_config is not None and self._check_supports_checkpointing(
|
|
logical_plan
|
|
):
|
|
self._supports_checkpointing = True
|
|
data_file_dir, data_file_fs = self._get_data_file_info(logical_plan)
|
|
|
|
checkpoint_callback = self._create_checkpoint_callback(
|
|
checkpoint_config,
|
|
)
|
|
|
|
callbacks.append(checkpoint_callback)
|
|
|
|
# Dynamically set the plan functions for checkpointing because they
|
|
# need to a reference to the checkpoint ref.
|
|
self._plan_fns_for_checkpointing = self._get_plan_fns_for_checkpointing(
|
|
data_file_dir, data_file_fs
|
|
)
|
|
|
|
elif checkpoint_config is not None:
|
|
assert not self._check_supports_checkpointing(logical_plan)
|
|
warnings.warn(
|
|
"You've enabled checkpointing, but the logical plan doesn't support "
|
|
"checkpointing. Checkpointing will be disabled."
|
|
)
|
|
physical_dag, op_map = self._plan_recursively(
|
|
logical_plan.dag, logical_plan.context
|
|
)
|
|
physical_plan = PhysicalPlan(physical_dag, op_map, logical_plan.context)
|
|
return physical_plan, callbacks
|
|
|
|
def get_plan_fn(self, logical_op: LogicalOperator) -> PlanLogicalOpFn:
|
|
if self._supports_checkpointing:
|
|
assert self._plan_fns_for_checkpointing
|
|
plan_fn = find_plan_fn(logical_op, self._plan_fns_for_checkpointing)
|
|
if plan_fn is not None:
|
|
return plan_fn
|
|
|
|
plan_fn = find_plan_fn(logical_op, self._DEFAULT_PLAN_FNS)
|
|
if plan_fn is not None:
|
|
return plan_fn
|
|
|
|
raise ValueError(
|
|
f"Found unknown logical operator during planning: {logical_op}"
|
|
)
|
|
|
|
def _plan_recursively(
|
|
self, logical_op: LogicalOperator, data_context: DataContext
|
|
) -> Tuple[PhysicalOperator, Dict[LogicalOperator, PhysicalOperator]]:
|
|
"""Plan a logical operator and its input dependencies recursively.
|
|
|
|
Args:
|
|
logical_op: The logical operator to plan.
|
|
data_context: The data context.
|
|
|
|
Returns:
|
|
A tuple of the physical operator corresponding to the logical operator, and
|
|
a mapping from physical to logical operators.
|
|
"""
|
|
op_map: Dict[PhysicalOperator, LogicalOperator] = {}
|
|
|
|
# Plan the input dependencies first.
|
|
physical_children = []
|
|
for child in logical_op.input_dependencies:
|
|
physical_child, child_op_map = self._plan_recursively(child, data_context)
|
|
physical_children.append(physical_child)
|
|
op_map.update(child_op_map)
|
|
|
|
plan_fn = self.get_plan_fn(logical_op)
|
|
# We will call `set_logical_operators()` in the following for-loop,
|
|
# no need to do it here.
|
|
physical_op = plan_fn(logical_op, physical_children, data_context)
|
|
|
|
# Traverse up the DAG, and set the mapping from physical to logical operators.
|
|
# At this point, all physical operators without logical operators set
|
|
# must have been created by the current logical operator.
|
|
queue = [physical_op]
|
|
while queue:
|
|
curr_physical_op = queue.pop()
|
|
if curr_physical_op._logical_operators:
|
|
continue
|
|
|
|
curr_physical_op.set_logical_operators(logical_op)
|
|
# Add this operator to the op_map so optimizer can find it
|
|
op_map[curr_physical_op] = logical_op
|
|
queue.extend(curr_physical_op.input_dependencies)
|
|
|
|
# Also add the final operator (in case the loop didn't catch it)
|
|
op_map[physical_op] = logical_op
|
|
return physical_op, op_map
|
|
|
|
def _create_checkpoint_callback(
|
|
self,
|
|
checkpoint_config,
|
|
) -> LoadCheckpointCallback:
|
|
"""Factory method to create the LoadCheckpointCallback.
|
|
|
|
Subclasses can override this to use a different callback implementation.
|
|
"""
|
|
return LoadCheckpointCallback(
|
|
checkpoint_config,
|
|
)
|
|
|
|
@staticmethod
|
|
def _get_data_file_info(logical_plan: LogicalPlan):
|
|
"""Extract the data file directory and filesystem from the Write op's datasink.
|
|
|
|
Returns (path, filesystem) for file-based datasinks, or (None, None)
|
|
for non-file datasinks.
|
|
"""
|
|
last_op = logical_plan.dag
|
|
if isinstance(last_op, Write):
|
|
datasink = last_op.datasink_or_legacy_datasource
|
|
if isinstance(datasink, _FileDatasink):
|
|
return datasink.unresolved_path, datasink.filesystem
|
|
return None, None
|
|
|
|
def _get_plan_fns_for_checkpointing(
|
|
self,
|
|
data_file_dir: Optional[str] = None,
|
|
data_file_filesystem: Optional["pyarrow.fs.FileSystem"] = None,
|
|
) -> Dict[Type[LogicalOperator], PlanLogicalOpFn]:
|
|
plan_fns = {
|
|
Read: partial(
|
|
plan_read_op_with_checkpoint_filter, data_file_dir, data_file_filesystem
|
|
),
|
|
ReadFiles: partial(
|
|
plan_read_files_op_with_checkpoint_filter,
|
|
data_file_dir,
|
|
data_file_filesystem,
|
|
),
|
|
Write: plan_write_op_with_checkpoint_writer,
|
|
}
|
|
return plan_fns
|
|
|
|
def _check_supports_checkpointing(self, logical_plan: LogicalPlan) -> bool:
|
|
"""Check if the logical plan supports checkpointing.
|
|
|
|
Subclasses can override _CHECKPOINT_FILTER_OPS to support more operators.
|
|
"""
|
|
if not isinstance(logical_plan.dag, (Write, StreamingSplit)):
|
|
return False
|
|
|
|
def _all_paths_contain_checkpoint_filter(op: LogicalOperator) -> bool:
|
|
if isinstance(op, self._CHECKPOINT_FILTER_OPS):
|
|
return True
|
|
return all(
|
|
_all_paths_contain_checkpoint_filter(input_dep)
|
|
for input_dep in op.input_dependencies
|
|
)
|
|
|
|
return _all_paths_contain_checkpoint_filter(logical_plan.dag)
|
|
|
|
|
|
def find_plan_fn(
|
|
logical_op: LogicalOperator, plan_fns: Dict[Type[LogicalOperator], PlanLogicalOpFn]
|
|
) -> Optional[PlanLogicalOpFn]:
|
|
"""Find the plan function for a logical operator.
|
|
|
|
This function goes through the plan functions in order and returns the first one
|
|
that is an instance of the logical operator type.
|
|
|
|
Args:
|
|
logical_op: The logical operator to find the plan function for.
|
|
plan_fns: The dictionary of plan functions.
|
|
|
|
Returns:
|
|
The plan function for the logical operator, or None if no plan function is
|
|
found.
|
|
"""
|
|
# TODO: This implementation doesn't account for type hierarchies conflicts or
|
|
# multiple inheritance.
|
|
for op_type, plan_fn in plan_fns.items():
|
|
if isinstance(logical_op, op_type):
|
|
return plan_fn
|
|
return None
|