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

Code before:

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
        self._num_blocks = 0
        self._num_bundles = 0
        self._nbytes = 0

    def __len__(self) -> int:
        """Return the total # bundles."""

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

Code after:

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

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