wehub-resource-sync
57 lines
1.9 KiB
Python
57 lines
1.9 KiB
Python
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()
|