chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,32 @@
|
||||
from .logical_operator import (
|
||||
LogicalOperator,
|
||||
LogicalOperatorPreservesSchema,
|
||||
LogicalOperatorSupportsPredicatePassThrough,
|
||||
LogicalOperatorSupportsPredicatePushdown,
|
||||
LogicalOperatorSupportsProjectionPushdown,
|
||||
LogicalOperatorUnifiesInputSchemas,
|
||||
PredicatePassThroughBehavior,
|
||||
)
|
||||
from .logical_plan import LogicalPlan
|
||||
from .operator import Operator
|
||||
from .optimizer import Optimizer, Rule
|
||||
from .physical_plan import PhysicalPlan
|
||||
from .plan import Plan
|
||||
from .source_operator import SourceOperator
|
||||
|
||||
__all__ = [
|
||||
"LogicalOperator",
|
||||
"LogicalPlan",
|
||||
"Operator",
|
||||
"Optimizer",
|
||||
"PhysicalPlan",
|
||||
"Plan",
|
||||
"Rule",
|
||||
"SourceOperator",
|
||||
"LogicalOperatorPreservesSchema",
|
||||
"LogicalOperatorSupportsProjectionPushdown",
|
||||
"LogicalOperatorSupportsPredicatePushdown",
|
||||
"LogicalOperatorSupportsPredicatePassThrough",
|
||||
"LogicalOperatorUnifiesInputSchemas",
|
||||
"PredicatePassThroughBehavior",
|
||||
]
|
||||
@@ -0,0 +1,247 @@
|
||||
import copy
|
||||
from abc import ABC, abstractmethod
|
||||
from dataclasses import dataclass, field, fields, replace
|
||||
from enum import Enum
|
||||
from typing import TYPE_CHECKING, Any, Callable, Dict, Iterator, List, Optional
|
||||
|
||||
from .operator import Operator
|
||||
from ray.data.block import BlockMetadata
|
||||
from ray.data.expressions import Expr
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.block import Schema
|
||||
|
||||
|
||||
@dataclass(frozen=True, repr=False, eq=False)
|
||||
class LogicalOperator(Operator, ABC):
|
||||
"""Abstract class for logical operators.
|
||||
|
||||
A logical operator describes transformation, and later is converted into
|
||||
physical operator.
|
||||
"""
|
||||
|
||||
_name: Optional[str] = field(init=False, default=None, repr=False)
|
||||
_input_dependencies: List["LogicalOperator"] = field(
|
||||
init=False, default_factory=list, repr=False
|
||||
)
|
||||
|
||||
@property
|
||||
def name(self) -> str:
|
||||
return self._name or self.__class__.__name__
|
||||
|
||||
@property
|
||||
def num_outputs(self) -> Optional[int]:
|
||||
"""Expected number of output blocks, if known."""
|
||||
return None
|
||||
|
||||
def estimated_num_outputs(self) -> Optional[int]:
|
||||
"""Returns the estimated number of blocks that
|
||||
would be outputted by this logical operator.
|
||||
|
||||
This method does not execute the plan, so it does not take into consideration
|
||||
block splitting. This method only considers high-level block constraints like
|
||||
`Dataset.repartition(num_blocks=X)`. A more accurate estimation can be given by
|
||||
`PhysicalOperator.num_outputs_total()` during execution.
|
||||
"""
|
||||
if self.num_outputs is not None:
|
||||
return self.num_outputs
|
||||
elif len(self.input_dependencies) == 1:
|
||||
return self.input_dependencies[0].estimated_num_outputs()
|
||||
return None
|
||||
|
||||
# Override the following 3 methods to correct type hints.
|
||||
|
||||
@property
|
||||
def input_dependencies(self) -> List["LogicalOperator"]:
|
||||
value = self._input_dependencies
|
||||
for x in value:
|
||||
assert isinstance(x, LogicalOperator), x
|
||||
return value
|
||||
|
||||
@input_dependencies.setter
|
||||
def input_dependencies(self, value: List["LogicalOperator"]) -> None:
|
||||
for x in value:
|
||||
assert isinstance(x, LogicalOperator), x
|
||||
object.__setattr__(self, "_input_dependencies", value)
|
||||
|
||||
def post_order_iter(self) -> Iterator["LogicalOperator"]:
|
||||
return super().post_order_iter() # type: ignore
|
||||
|
||||
def _apply_transform(
|
||||
self, transform: Callable[["LogicalOperator"], "LogicalOperator"]
|
||||
) -> "LogicalOperator":
|
||||
input_dependencies = self.input_dependencies
|
||||
transformed_inputs = [
|
||||
input_op._apply_transform(transform) for input_op in input_dependencies
|
||||
]
|
||||
if all(
|
||||
transformed_input is input_op
|
||||
for transformed_input, input_op in zip(
|
||||
transformed_inputs, input_dependencies
|
||||
)
|
||||
):
|
||||
target = self
|
||||
else:
|
||||
target = self._with_new_input_dependencies(transformed_inputs)
|
||||
return transform(target)
|
||||
|
||||
def _with_new_input_dependencies(
|
||||
self, input_dependencies: List["LogicalOperator"]
|
||||
) -> "LogicalOperator":
|
||||
if "input_dependencies" in {field.name for field in fields(self)}:
|
||||
return replace(self, input_dependencies=input_dependencies)
|
||||
|
||||
target = copy.copy(self)
|
||||
object.__setattr__(target, "_input_dependencies", input_dependencies)
|
||||
return target
|
||||
|
||||
def _get_args(self) -> Dict[str, Any]:
|
||||
"""This Dict must be serializable"""
|
||||
args: Dict[str, Any] = {}
|
||||
for dataclass_field in fields(self):
|
||||
key = dataclass_field.name
|
||||
value = getattr(self, key)
|
||||
# Keep underscore-prefixed keys to preserve legacy export schema.
|
||||
args[key if key.startswith("_") else f"_{key}"] = value
|
||||
args["_name"] = self.name
|
||||
# Preserve legacy export shape even though output deps are no longer tracked.
|
||||
args["_output_dependencies"] = []
|
||||
# Do not include input dependencies, since we only want to export this
|
||||
# operator-specific args. Adding input_dependencies isn't wrong, but can
|
||||
# lead to slow recursive calls with `sanitize_for_struct`, since logical
|
||||
# operators are dataclasses.
|
||||
args["_input_dependencies"] = []
|
||||
return args
|
||||
|
||||
def infer_schema(self) -> Optional["Schema"]:
|
||||
"""Returns the inferred schema of the output blocks."""
|
||||
return None
|
||||
|
||||
def infer_metadata(self) -> "BlockMetadata":
|
||||
"""A ``BlockMetadata`` that represents the aggregate metadata of the outputs.
|
||||
|
||||
This method is used by methods like :meth:`~ray.data.Dataset.schema` to
|
||||
efficiently return metadata.
|
||||
"""
|
||||
return BlockMetadata(None, None, None, None)
|
||||
|
||||
def is_lineage_serializable(self) -> bool:
|
||||
"""Returns whether the lineage of this operator can be serialized.
|
||||
|
||||
An operator is lineage serializable if you can serialize it on one machine and
|
||||
deserialize it on another without losing information. Operators that store
|
||||
object references (e.g., ``InputData``) aren't lineage serializable because the
|
||||
objects aren't available on the deserialized machine.
|
||||
"""
|
||||
return True
|
||||
|
||||
|
||||
class LogicalOperatorSupportsProjectionPushdown(LogicalOperator):
|
||||
"""Mixin for reading operators supporting projection pushdown"""
|
||||
|
||||
def supports_projection_pushdown(self) -> bool:
|
||||
return False
|
||||
|
||||
def get_projection_map(self) -> Optional[Dict[str, str]]:
|
||||
return None
|
||||
|
||||
def apply_projection(
|
||||
self,
|
||||
projection_map: Optional[Dict[str, str]],
|
||||
) -> LogicalOperator:
|
||||
return self
|
||||
|
||||
|
||||
class LogicalOperatorPreservesSchema(LogicalOperator):
|
||||
"""Mixin for operators whose output column layout is identical to their
|
||||
single input's. Provides a default ``infer_schema()`` that delegates to
|
||||
the input. Use for ops like ``Filter``, ``Sort``, ``Limit``, etc., that
|
||||
only re-order or filter rows.
|
||||
|
||||
List this mixin last in the bases of subclasses so the concrete operator
|
||||
base (e.g., ``AbstractMap``, ``AbstractAllToAll``) drives ``__init__`` /
|
||||
``super()`` chains.
|
||||
"""
|
||||
|
||||
def infer_schema(self) -> Optional["Schema"]:
|
||||
assert len(self.input_dependencies) == 1, len(self.input_dependencies)
|
||||
return self.input_dependencies[0].infer_schema()
|
||||
|
||||
|
||||
class LogicalOperatorUnifiesInputSchemas(LogicalOperator):
|
||||
"""Mixin for n-ary operators whose output schema is the unification of
|
||||
all inputs' schemas (e.g., ``Union``, ``Mix``). Provides a default
|
||||
``infer_schema()`` that returns the result of
|
||||
``unify_schemas_with_validation`` over each input's schema, or
|
||||
``None`` if any input's schema is unresolvable.
|
||||
|
||||
List this mixin last in the bases of subclasses so the concrete operator
|
||||
base (e.g., ``NAry``) drives ``__init__`` / ``super()`` chains.
|
||||
"""
|
||||
|
||||
def infer_schema(self) -> Optional["Schema"]:
|
||||
import pyarrow as pa
|
||||
|
||||
from ray.data._internal.util import unify_schemas_with_validation
|
||||
|
||||
input_schemas = [op.infer_schema() for op in self.input_dependencies]
|
||||
if not all(isinstance(s, pa.Schema) for s in input_schemas):
|
||||
return None
|
||||
try:
|
||||
return unify_schemas_with_validation(input_schemas)
|
||||
except (pa.ArrowTypeError, pa.ArrowInvalid):
|
||||
return None
|
||||
|
||||
|
||||
class LogicalOperatorSupportsPredicatePushdown(LogicalOperator):
|
||||
"""Mixin for reading operators supporting predicate pushdown"""
|
||||
|
||||
def supports_predicate_pushdown(self) -> bool:
|
||||
return False
|
||||
|
||||
def get_current_predicate(self) -> Optional[Expr]:
|
||||
return None
|
||||
|
||||
def apply_predicate(
|
||||
self,
|
||||
predicate_expr: Expr,
|
||||
) -> LogicalOperator:
|
||||
return self
|
||||
|
||||
|
||||
class PredicatePassThroughBehavior(Enum):
|
||||
"""Defines how predicates can be passed through an operator."""
|
||||
|
||||
# Predicate can be pushed through as-is (e.g., Sort, Repartition, RandomShuffle, Limit)
|
||||
PASSTHROUGH = "passthrough"
|
||||
|
||||
# Predicate can be pushed through but needs column rebinding (e.g., Project)
|
||||
PASSTHROUGH_WITH_SUBSTITUTION = "passthrough_with_substitution"
|
||||
|
||||
# Predicate can be pushed into each branch (e.g., Union)
|
||||
PUSH_INTO_BRANCHES = "push_into_branches"
|
||||
|
||||
# Predicate can be conditionally pushed based on columns (e.g., Join)
|
||||
CONDITIONAL = "conditional"
|
||||
|
||||
|
||||
class LogicalOperatorSupportsPredicatePassThrough(ABC):
|
||||
"""Mixin for operators that allow predicates to be pushed through them.
|
||||
|
||||
This is distinct from LogicalOperatorSupportsPredicatePushdown, which is for
|
||||
operators that can *accept* predicates (like Read). This trait is for operators
|
||||
that allow predicates to *pass through* them.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def predicate_passthrough_behavior(self) -> PredicatePassThroughBehavior:
|
||||
"""Returns the predicate passthrough behavior for this operator."""
|
||||
pass
|
||||
|
||||
def get_column_substitutions(self) -> Optional[Dict[str, str]]:
|
||||
"""Returns column renames needed when pushing through (for PASSTHROUGH_WITH_SUBSTITUTION).
|
||||
|
||||
Returns:
|
||||
Dict mapping from old_name -> new_name, or None if no rebinding needed
|
||||
"""
|
||||
return None
|
||||
@@ -0,0 +1,56 @@
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
from .logical_operator import LogicalOperator
|
||||
from .plan import Plan
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.context import DataContext
|
||||
|
||||
|
||||
class LogicalPlan(Plan):
|
||||
"""The plan with a DAG of logical operators."""
|
||||
|
||||
def __init__(self, dag: LogicalOperator, context: "DataContext"):
|
||||
super().__init__(context)
|
||||
self._dag = dag
|
||||
|
||||
@property
|
||||
def dag(self) -> LogicalOperator:
|
||||
"""Get the DAG of logical operators."""
|
||||
return self._dag
|
||||
|
||||
def sources(self) -> List[LogicalOperator]:
|
||||
"""List of operators that are sources for this plan's DAG."""
|
||||
# If an operator has no input dependencies, it's a source.
|
||||
if not any(self._dag.input_dependencies):
|
||||
return [self._dag]
|
||||
|
||||
sources = []
|
||||
for op in self._dag.input_dependencies:
|
||||
sources.extend(LogicalPlan(op, self.context).sources())
|
||||
return sources
|
||||
|
||||
def has_lazy_input(self) -> bool:
|
||||
"""Return whether this plan has lazy input blocks."""
|
||||
from ray.data._internal.logical.operators import Read
|
||||
|
||||
return all(isinstance(op, Read) for op in self.sources())
|
||||
|
||||
def require_preserve_order(self) -> bool:
|
||||
"""Whether this plan requires to preserve order."""
|
||||
from ray.data._internal.logical.operators import Zip
|
||||
|
||||
return any(isinstance(op, Zip) for op in self.dag.post_order_iter())
|
||||
|
||||
def input_files(self) -> Optional[List[str]]:
|
||||
"""Get the input files of the dataset, if available."""
|
||||
input_files = self.dag.infer_metadata().input_files
|
||||
if input_files is None:
|
||||
return None
|
||||
return list(set(input_files))
|
||||
|
||||
def initial_num_blocks(self) -> Optional[int]:
|
||||
"""Get the estimated number of blocks from the logical plan
|
||||
after applying execution plan optimizations, but prior to
|
||||
fully executing the dataset."""
|
||||
return self.dag.estimated_num_outputs()
|
||||
@@ -0,0 +1,56 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import Callable, Iterator, List
|
||||
|
||||
|
||||
class Operator(ABC):
|
||||
"""Abstract class for operators.
|
||||
|
||||
Operators live on the driver side of the Dataset only.
|
||||
"""
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def name(self) -> str:
|
||||
"""Name for this operator."""
|
||||
...
|
||||
|
||||
@property
|
||||
def dag_str(self) -> str:
|
||||
"""String representation of the whole DAG."""
|
||||
if self.input_dependencies:
|
||||
out_str = ", ".join([x.dag_str for x in self.input_dependencies])
|
||||
out_str += " -> "
|
||||
else:
|
||||
out_str = ""
|
||||
out_str += f"{self.__class__.__name__}[{self.name}]"
|
||||
return out_str
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def input_dependencies(self) -> List["Operator"]:
|
||||
"""List of operators that provide inputs for this operator."""
|
||||
...
|
||||
|
||||
def post_order_iter(self) -> Iterator["Operator"]:
|
||||
"""Depth-first traversal of this operator and its input dependencies."""
|
||||
for op in self.input_dependencies:
|
||||
yield from op.post_order_iter()
|
||||
yield self
|
||||
|
||||
@abstractmethod
|
||||
def _apply_transform(
|
||||
self, transform: Callable[["Operator"], "Operator"]
|
||||
) -> "Operator":
|
||||
"""Recursively applies transformation (in post-order) to the operators DAG
|
||||
|
||||
NOTE: This operation should be opting in to avoid in-place modifications,
|
||||
instead creating new operations whenever any operator needs to be
|
||||
updated.
|
||||
"""
|
||||
...
|
||||
|
||||
def __repr__(self) -> str:
|
||||
return f"{self.__class__.__name__}[{self.name}]"
|
||||
|
||||
def __str__(self) -> str:
|
||||
return repr(self)
|
||||
@@ -0,0 +1,52 @@
|
||||
from typing import List, Type
|
||||
|
||||
from .plan import Plan
|
||||
|
||||
|
||||
class Rule:
|
||||
"""Abstract class for optimization rule."""
|
||||
|
||||
def apply(self, plan: Plan) -> Plan:
|
||||
"""Apply the optimization rule to the execution plan."""
|
||||
raise NotImplementedError
|
||||
|
||||
@classmethod
|
||||
def dependencies(cls) -> List[Type["Rule"]]:
|
||||
"""List of rules that must be applied before this rule."""
|
||||
return []
|
||||
|
||||
@classmethod
|
||||
def dependents(cls) -> List[Type["Rule"]]:
|
||||
"""List of rules that must be applied after this rule."""
|
||||
return []
|
||||
|
||||
|
||||
class Optimizer:
|
||||
"""Abstract class for optimizers.
|
||||
|
||||
An optimizers transforms a DAG of operators with a list of predefined rules.
|
||||
"""
|
||||
|
||||
@property
|
||||
def rules(self) -> List[Rule]:
|
||||
"""List of predefined rules for this optimizer."""
|
||||
raise NotImplementedError
|
||||
|
||||
def optimize(self, plan: Plan) -> Plan:
|
||||
"""Optimize operators with a list of rules."""
|
||||
# Apply rules until the plan is not changed
|
||||
previous_plan = plan
|
||||
while True:
|
||||
for rule in self.rules:
|
||||
plan = rule.apply(plan)
|
||||
# TODO: Eventually we should implement proper equality.
|
||||
# Using str to check equality seems brittle
|
||||
if plan.dag.dag_str == previous_plan.dag.dag_str:
|
||||
break
|
||||
previous_plan = plan
|
||||
return self._post_optimize(plan)
|
||||
|
||||
def _post_optimize(self, plan: Plan) -> Plan:
|
||||
"""Post optimize is used for rules or other post-processing
|
||||
that needs to be executed only once after the `optimize` loop."""
|
||||
return plan
|
||||
@@ -0,0 +1,34 @@
|
||||
from typing import TYPE_CHECKING, Dict
|
||||
|
||||
from .logical_operator import LogicalOperator
|
||||
from .plan import Plan
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data._internal.execution.interfaces import PhysicalOperator
|
||||
from ray.data.context import DataContext
|
||||
|
||||
|
||||
class PhysicalPlan(Plan):
|
||||
"""The plan with a DAG of physical operators."""
|
||||
|
||||
def __init__(
|
||||
self,
|
||||
dag: "PhysicalOperator",
|
||||
op_map: Dict["PhysicalOperator", LogicalOperator],
|
||||
context: "DataContext",
|
||||
):
|
||||
super().__init__(context)
|
||||
self._dag = dag
|
||||
self._op_map = op_map
|
||||
|
||||
@property
|
||||
def dag(self) -> "PhysicalOperator":
|
||||
"""Get the DAG of physical operators."""
|
||||
return self._dag
|
||||
|
||||
@property
|
||||
def op_map(self) -> Dict["PhysicalOperator", LogicalOperator]:
|
||||
"""
|
||||
Get a mapping from physical operators to their corresponding logical operator.
|
||||
"""
|
||||
return self._op_map
|
||||
@@ -0,0 +1,29 @@
|
||||
from typing import TYPE_CHECKING
|
||||
|
||||
from .operator import Operator
|
||||
|
||||
if TYPE_CHECKING:
|
||||
from ray.data.context import DataContext
|
||||
|
||||
|
||||
class Plan:
|
||||
"""Abstract class for logical/physical execution plans.
|
||||
|
||||
This plan should hold an operator representing the plan DAG and any auxiliary data
|
||||
that's useful for plan optimization or execution.
|
||||
"""
|
||||
|
||||
def __init__(self, context: "DataContext"):
|
||||
self._context = context
|
||||
|
||||
@property
|
||||
def dag(self) -> Operator:
|
||||
raise NotImplementedError
|
||||
|
||||
@property
|
||||
def context(self) -> "DataContext":
|
||||
return self._context
|
||||
|
||||
@context.setter
|
||||
def context(self, context: "DataContext") -> None:
|
||||
self._context = context
|
||||
@@ -0,0 +1,17 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from typing import TYPE_CHECKING, List, Optional
|
||||
|
||||
if TYPE_CHECKING:
|
||||
|
||||
from ray.data.dataset import RefBundle
|
||||
|
||||
|
||||
class SourceOperator(ABC):
|
||||
"""Mixin for Logical operators that can be logical source nodes.
|
||||
Subclasses: Read, InputData, FromAbstract.
|
||||
"""
|
||||
|
||||
@abstractmethod
|
||||
def output_data(self) -> Optional[List["RefBundle"]]:
|
||||
"""The output data of this operator if already known, or ``None``."""
|
||||
pass
|
||||
Reference in New Issue
Block a user