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):
`self._metadata_cached` is true only if you call `set_from_iteration`, but I think that's confusing in this scenario:
* You cache the output bundles but not the explicit metadata
* You can get the cached metadata from the output bundles, but `_metadata_cached` returns False 



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:
                return self._bundle.size_bytes()