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

Code before:

        for name, type in zip(self.names, self.types):
            output += name
            output += " " * ((column_width + padding) - len(name))
            output += f"{type}\n"

        output = output.rstrip()
        return output


def _block_to_df(block: Block) -> "pandas.DataFrame":
    block = BlockAccessor.for_block(block)
    return block.to_pandas()


def _block_to_ndarray(block: Block, column: Optional[str]):
    block = BlockAccessor.for_block(block)
    return block.to_numpy(column)


def _block_to_arrow(block: Block):
    block = BlockAccessor.for_block(block)
    return block.to_arrow()


class _ExecutionCache:
    """Consolidated cache for Dataset execution results.

    Caches the output bundle, execution stats, and lightweight metadata
    (schema, num_rows, size_bytes), and tracks which operator produced
    the cached data.

    There are two "layers" of cache:
      1. Bundle layer: the full RefBundle from eager execution (execute()).
         Valid only when _operator matches the current DAG.
      2. Metadata layer: schema, num_rows, size_bytes cached as scalars.
         Populated when a streaming iterator is fully exhausted
         (CacheMetadataIterator in legacy_compat.py).

    Getters check the bundle layer first, then the metadata layer.
    """

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

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

Review comment (performance):
I think we can simplify and narrow the `_ExecutionCache` interface by removing `has_cached_metadata` and the `_metadata_cached` flag entirely, and replacing `set_metadata` with `set_num_rows` and `set_size_bytes`:
* The setters and getters are asymmetric. You call `set_metadata` to set the count and byte size, but call `get_num_rows` and `get_size_bytes` to get them.
* `has_cached_metadata` returns different values depending on *how* you set cache the metadata. In the code snippet below, `has_cached_metadata` returns `False` even though we've cached the schema, which is a type of metadata
```python
>>> cache1.set_schema(dag, schema) 
>>> cache1.has_cached_metadata(dag)
False  
```

I think `has_cached_metadata` is called in only one place in `get_plan_as_string`. Would it make sense to just replace that call with `self._cache.get_schema(dag) is not None and self._cache.get_num_rows(dag) is not None`? IIUC that's what we actually care about for that branch, not whether the cache has been populated by a particular method?

Code after:

        for name, type in zip(self.names, self.types):
            output += name
            output += " " * ((column_width + padding) - len(name))
            output += f"{type}\n"

        output = output.rstrip()
        return output


def _block_to_df(block: Block) -> "pandas.DataFrame":
    block = BlockAccessor.for_block(block)
    return block.to_pandas()


def _block_to_ndarray(block: Block, column: Optional[str]):
    block = BlockAccessor.for_block(block)
    return block.to_numpy(column)


def _block_to_arrow(block: Block):
    block = BlockAccessor.for_block(block)
    return block.to_arrow()


class _ExecutionCache:
    """Consolidated cache for Dataset execution results.

    Caches the output bundle, execution stats, and lightweight metadata
    (schema, num_rows, size_bytes), and tracks which operator produced
    the cached data.

    There are two "layers" of cache:
      1. Bundle layer: the full RefBundle from eager execution (execute()).
         Valid only when _operator matches the current DAG.
      2. Metadata layer: schema, num_rows, size_bytes cached as scalars.
         Populated when a streaming iterator is fully exhausted
         (CacheMetadataIterator in legacy_compat.py).

    Getters check the bundle layer first, then the metadata layer.
    """

    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.