Repository: ray-project/ray
File: python/ray/data/_internal/plan.py (Python)

Code before:

    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._logical_plan.dag.estimated_num_outputs()

    def schema(
        self, fetch_if_missing: bool = False
    ) -> Union[type, "pyarrow.lib.Schema"]:
        """Get the schema after applying all execution plan optimizations,
        but prior to fully executing the dataset
        (unless `fetch_if_missing` is set to True).

        Args:
            fetch_if_missing: Whether to execute the plan to fetch the schema.

        Returns:
            The schema of the output dataset.
        """
        schema = self._cache.get_schema(self._logical_plan.dag)
        if schema is None:
            schema = self._logical_plan.dag.infer_schema()
        if schema is None and fetch_if_missing:
            # For consistency with the previous implementation, we fetch the schema if
            # the plan is read-only even if `fetch_if_missing` is False.

            iter_ref_bundles, _, executor = self.execute_to_iterator()
            if executor is not None:
                # Make sure executor is fully shutdown upon exiting
                with executor:
                    schema = _take_first_non_empty_schema(
                        bundle.schema for bundle in iter_ref_bundles
                    )
        if schema is not None:
            self.cache_schema(schema)
        return schema

    def cache_schema(self, schema: Union[type, "pyarrow.lib.Schema"]):
        self._cache.set_schema(self._logical_plan.dag, schema)

    def input_files(self) -> Optional[List[str]]:
        """Get the input files of the dataset, if available."""
        return self._logical_plan.dag.infer_metadata().input_files

    def meta_count(self) -> Optional[int]:
        """Get the number of rows after applying all plan optimizations, if possible.

        This method will never trigger any computation.

        Returns:

Review comment (suggestion):
> we fetch the schema if the plan is read-only even if `fetch_if_missing` is False.

Is this still true

Code after:
        return plan_copy

    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._logical_plan.dag.estimated_num_outputs()

    def schema(
        self, fetch_if_missing: bool = False
    ) -> Union[type, "pyarrow.lib.Schema"]:
        """Get the schema after applying all execution plan optimizations,
        but prior to fully executing the dataset
        (unless `fetch_if_missing` is set to True).

        Args:
            fetch_if_missing: Whether to execute the plan to fetch the schema.

        Returns:
            The schema of the output dataset.
        """

        def _build_limited_plan(plan: "ExecutionPlan") -> "ExecutionPlan":
            limited_dag = Limit(plan._logical_plan.dag, limit=1)
            limited_plan = plan.copy()
            limited_plan.link_logical_plan(LogicalPlan(limited_dag, plan._context))
            return limited_plan

        schema = self._cache.get_schema(self._logical_plan.dag)
        if schema is None:
            schema = self._logical_plan.dag.infer_schema()
        if schema is None and fetch_if_missing:
            # Lazily execute only the first block to minimize computation.
            # We achieve this by appending a Limit[1] operation to a copy of
            # this plan, which we then execute to get its schema.
            limited_plan = _build_limited_plan(self)
            iter_ref_bundles, _, executor = limited_plan.execute_to_iterator()
            if executor is not None:
                # Make sure executor is fully shutdown upon exiting
                with executor:
                    schema = _take_first_non_empty_schema(
                        bundle.schema for bundle in iter_ref_bundles
                    )
        if schema is not None:
            self._cache.set_schema(self._logical_plan.dag, schema)
        return schema

    def input_files(self) -> Optional[List[str]]:
        """Get the input files of the dataset, if available."""
        return self._logical_plan.dag.infer_metadata().input_files