chore: import upstream snapshot with attribution

python/ray/data/_internal/logical/interfaces/logical_plan.py

@@ -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()