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
        self._stats: Optional[DatasetStats] = None

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

Review comment (performance):
I felt confused trying to understand this interface. I think it's because the interface is wide and leaks implementation details (e.g., that we cache metadata at the end of iteration).

How can we make this interface easier to reason about?

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.