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

Code before:

    def __init__(self):
        # --- Bundle layer (from eager execution) ---
        self._operator: Optional["LogicalOperator"] = None
        self._bundle: Optional[RefBundle] = None
        self._stats: Optional[DatasetStats] = None

        # --- Metadata layer (from streaming iteration) ---
        self._metadata_cached = False
        self._schema: Optional["Schema"] = None
        self._num_rows: Optional[int] = None
        self._size_bytes: Optional[int] = None

    # --- Consolidated Getters ---

    def cache_is_fresh(self, dag: "LogicalOperator") -> bool:
        # This _ExecutionCache is only fresh if the current logical
        # plan dag ends with the same operator. Otherwise, the plan
        # has changed, so there may have been a change in schema,
        # count, etc.
        return self._operator == dag

    def has_computed_output(self, dag: "LogicalOperator") -> bool:
        return self.get_bundle(dag) is not None

    def has_iterator_cache(self, dag: "LogicalOperator") -> bool:
        return self.cache_is_fresh(dag) and self._metadata_cached

    def get_bundle(self, dag: "LogicalOperator") -> Optional[RefBundle]:
        if self.cache_is_fresh(dag):
            return self._bundle
        return None

    def get_stats(self) -> Optional[DatasetStats]:
        # We don't check for cache freshness just for stats for behaviorial
        # backwards compatibility.
        return self._stats

    def get_schema(self, dag: "LogicalOperator") -> Optional["Schema"]:
        if self.cache_is_fresh(dag):
            if self._bundle is not None and self._bundle.schema is not None:
                return self._bundle.schema
            return self._schema
        return None

    def get_num_rows(self, dag: "LogicalOperator") -> Optional[int]:
        if self.cache_is_fresh(dag):
            if self._bundle is not None and self._bundle.num_rows() is not None:
                return self._bundle.num_rows()
            return self._num_rows
        return None

Review comment (performance):
```suggestion
    def has_cached_metadata(self, dag: "LogicalOperator") -> bool:
```
I don't think callers of `ExecutionCache` care where the metadata from, and I think it might confuse readers when the look at the `ExecutionCache` interface in isolation. Like, this only makes sense if you know the implementation detail that we cache metadata when we iterate over a full dataset.

I think what callers actually care about is -- has the metadata been cached?

Code after:

    def __init__(self):
        # --- Bundle layer (from eager execution) ---
        self._operator: Optional["LogicalOperator"] = None
        self._bundle: Optional[RefBundle] = None

        # --- Metadata layer (from streaming iteration) ---
        self._num_rows: Optional[int] = None
        self._size_bytes: Optional[int] = None
        # Note schema can be cached via other means as well.
        self._schema: Optional["Schema"] = None

        # --- Other ---
        self._stats: Optional[DatasetStats] = None

    # --- Consolidated Getters ---

    def _cache_is_fresh(self, dag: "LogicalOperator") -> bool:
        # This _ExecutionCache is only fresh if the current logical
        # plan dag ends with the same operator. Otherwise, the plan
        # has changed, so there may have been a change in schema,
        # count, etc.
        return self._operator == dag

    def get_bundle(self, dag: "LogicalOperator") -> Optional[RefBundle]:
        if self._cache_is_fresh(dag):
            return self._bundle
        return None

    def get_stats(self) -> Optional[DatasetStats]:
        # We don't check for cache freshness just for stats for behaviorial
        # backwards compatibility.
        return self._stats

    def get_schema(self, dag: "LogicalOperator") -> Optional["Schema"]:
        if self._cache_is_fresh(dag):
            if self._bundle is not None and self._bundle.schema is not None:
                return self._bundle.schema
            return self._schema
        return None

    def get_num_rows(self, dag: "LogicalOperator") -> Optional[int]:
        if self._cache_is_fresh(dag):
            if self._bundle is not None and self._bundle.num_rows() is not None:
                return self._bundle.num_rows()
            return self._num_rows
        return None

    def get_size_bytes(self, dag: "LogicalOperator") -> Optional[int]:
        if self._cache_is_fresh(dag):
            if self._bundle is not None: