Repository: ray-project/ray
File: python/ray/data/_internal/execution/bundle_queue/base.py (Python)

Code before:
import abc
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
)

from ray.data._internal.execution.util import memory_string

if TYPE_CHECKING:
    from ray.data._internal.execution.interfaces import RefBundle


class _QueueMetricRecorder:
    """Mixin for recording stats about a queue. Subclasses
    may choose to use the _on_dequeue and _on_enqueue methods to
    track num_blocks, nbytes, etc... If not, they should override
    those methods."""

    def __init__(self):
        self._nbytes: int = 0
        self._num_blocks: int = 0
        self._num_bundles: int = 0
        self._num_rows: int = 0

    def _on_enqueue(self, bundle: RefBundle):
        self._nbytes += bundle.size_bytes()
        self._num_blocks += len(bundle.block_refs)
        self._num_bundles += 1
        self._num_rows += bundle.num_rows() or 0

    def _on_dequeue(self, bundle: RefBundle):
        self._nbytes -= bundle.size_bytes()
        self._num_blocks -= len(bundle.block_refs)
        self._num_bundles -= 1
        self._num_rows -= bundle.num_rows() or 0

    def estimate_size_bytes(self) -> int:
        """Return the estimated size in bytes of all bundles."""
        return self._nbytes

    def num_blocks(self) -> int:
        """Return the total # of blocks across all bundles."""
        return self._num_blocks

    def num_rows(self) -> int:
        """Return the total # of rows across all bundles."""
        return self._num_rows

    def _reset_metrics(self):
        self._num_rows = 0

Review comment (suggestion):
```suggestion
    def _on_enqueued_bundle(self, bundle: RefBundle):
```

Code after:
import abc
from typing import (
    TYPE_CHECKING,
    Any,
    Optional,
)

from ray.data._internal.execution.util import memory_string

if TYPE_CHECKING:
    from ray.data._internal.execution.interfaces import RefBundle


class BaseBundleQueue:
    """Base class for storing bundles. Here and subclasses should adhere to the mental
    model that "first", "front", or "head" is the next bundle to be dequeued. Consequently,
    "last", "back", or "tail" is the last bundle to be dequeued.

    Subclasses may choose to use the _on_dequeue_bundle and _on_enqueue_bundle methods to
    track num_blocks, nbytes, etc... If not, they should override those methods.
    """

    def __init__(self):
        self._nbytes: int = 0
        self._num_blocks: int = 0
        self._num_bundles: int = 0
        self._num_rows: int = 0

    def _on_enqueue_bundle(self, bundle: RefBundle):
        self._nbytes += bundle.size_bytes()
        self._num_blocks += len(bundle.block_refs)
        self._num_bundles += 1
        self._num_rows += bundle.num_rows() or 0

    def _on_dequeue_bundle(self, bundle: RefBundle):
        self._nbytes -= bundle.size_bytes()
        self._num_blocks -= len(bundle.block_refs)
        self._num_bundles -= 1
        self._num_rows -= bundle.num_rows() or 0

    def estimate_size_bytes(self) -> int:
        """Return the estimated size in bytes of all bundles."""
        return self._nbytes

    def num_blocks(self) -> int:
        """Return the total # of blocks across all bundles."""
        return self._num_blocks

    def num_rows(self) -> int:
        """Return the total # of rows across all bundles."""
        return self._num_rows