chore: import upstream snapshot with attribution
This commit is contained in:
wehub-resource-sync
@@ -0,0 +1,254 @@
|
||||
from __future__ import annotations
|
||||
|
||||
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 BundleQueue(abc.ABC):
|
||||
@abc.abstractmethod
|
||||
def estimate_size_bytes(self) -> int:
|
||||
"""Returns the estimated size in bytes of all bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def num_blocks(self) -> int:
|
||||
"""Returns the total # of blocks across all bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def num_bundles(self) -> int:
|
||||
"""Returns the total # of bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def num_rows(self) -> int:
|
||||
"""Return the total # of rows across all bundles."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _on_enqueue_bundle(self, bundle: RefBundle):
|
||||
"""Hook called before a bundle is added to the queue."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _on_dequeue_bundle(self, bundle: RefBundle):
|
||||
"""Hook called after a bundle is removed from the queue."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _add_inner(self, bundle: RefBundle, **kwargs: Any):
|
||||
"""Add a bundle to the internal data structure."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
"""Remove and return the next bundle from the internal data structure."""
|
||||
...
|
||||
|
||||
def add(self, bundle: RefBundle, **kwargs: Any):
|
||||
"""Add a bundle to the tail(end) of the queue. Base classes should override
|
||||
the `_add_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Args:
|
||||
bundle: The bundle to add.
|
||||
**kwargs: Additional queue-specific parameters (e.g., `key` for ordered queues).
|
||||
This is used for `finalize`.
|
||||
"""
|
||||
self._on_enqueue_bundle(bundle)
|
||||
self._add_inner(bundle, **kwargs)
|
||||
|
||||
def get_next(self) -> RefBundle:
|
||||
"""Remove and return the head of the queue. Base classes should override
|
||||
the `_get_next_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Raises:
|
||||
IndexError: If the queue is empty.
|
||||
|
||||
Returns:
|
||||
The `RefBundle` at the head of the queue.
|
||||
"""
|
||||
bundle = self._get_next_inner()
|
||||
self._on_dequeue_bundle(bundle)
|
||||
return bundle
|
||||
|
||||
@abc.abstractmethod
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
"""Return the head of the queue. The only invariant is
|
||||
that the # of blocks, rows, and bytes must remain unchanged
|
||||
before and after this method call.
|
||||
|
||||
If queue.has_next() == False, return `None`.
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def has_next(self) -> bool:
|
||||
"""Check if the queue has a valid bundle."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def clear(self):
|
||||
"""Remove all bundles from the queue."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def finalize(self, **kwargs: Any):
|
||||
"""Signal that no additional bundles will be added to the bundler so
|
||||
the bundler can be finalized. The keys of kwargs provided should be the same
|
||||
as the ones passed into the `add()` method. This is important for ordered
|
||||
queues."""
|
||||
...
|
||||
|
||||
def __len__(self) -> int:
|
||||
"""Return the total # bundles."""
|
||||
return self.num_bundles()
|
||||
|
||||
def __repr__(self) -> str:
|
||||
"""Return a string representation showing queue metrics."""
|
||||
nbytes = memory_string(self.estimate_size_bytes())
|
||||
return (
|
||||
f"{self.__class__.__name__}("
|
||||
f"num_bundles={len(self)}, "
|
||||
f"num_blocks={self.num_blocks()}, "
|
||||
f"num_rows={self.num_rows()}, "
|
||||
f"nbytes={nbytes})"
|
||||
)
|
||||
|
||||
|
||||
class BaseBundleQueue(BundleQueue):
|
||||
"""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_bundles(self) -> int:
|
||||
return self._num_bundles
|
||||
|
||||
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 add(self, bundle: RefBundle, **kwargs: Any):
|
||||
"""Add a bundle to the tail(end) of the queue. Base classes should override
|
||||
the `_add_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Args:
|
||||
bundle: The bundle to add.
|
||||
**kwargs: Additional queue-specific parameters (e.g., `key` for ordered queues).
|
||||
This is used for `finalize`.
|
||||
"""
|
||||
self._on_enqueue_bundle(bundle)
|
||||
self._add_inner(bundle, **kwargs)
|
||||
|
||||
def _add_inner(self, bundle: RefBundle, **kwargs: Any) -> None:
|
||||
raise NotImplementedError
|
||||
|
||||
def get_next(self) -> RefBundle:
|
||||
"""Remove and return the head of the queue. Base classes should override
|
||||
the `_get_next_inner` method for simple use cases. For more complex metrics tracking,
|
||||
they can override this method.
|
||||
|
||||
Raises:
|
||||
IndexError: If the queue is empty.
|
||||
|
||||
Returns:
|
||||
The `RefBundle` at the head of the queue.
|
||||
"""
|
||||
bundle = self._get_next_inner()
|
||||
self._on_dequeue_bundle(bundle)
|
||||
return bundle
|
||||
|
||||
def _get_next_inner(self) -> RefBundle:
|
||||
raise NotImplementedError
|
||||
|
||||
@abc.abstractmethod
|
||||
def peek_next(self) -> Optional[RefBundle]:
|
||||
"""Return the head of the queue. The only invariant is
|
||||
that the # of blocks, rows, and bytes must remain unchanged
|
||||
before and after this method call.
|
||||
|
||||
If queue.has_next() == False, return `None`.
|
||||
"""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def has_next(self) -> bool:
|
||||
"""Check if the queue has a valid bundle."""
|
||||
...
|
||||
|
||||
@abc.abstractmethod
|
||||
def clear(self):
|
||||
"""Remove all bundles from the queue."""
|
||||
...
|
||||
|
||||
def finalize(self, **kwargs: Any):
|
||||
"""Signal that no additional bundles will be added to the bundler so
|
||||
the bundler can be finalized. The keys of kwargs provided should be the same
|
||||
as the ones passed into the `add()` method. This is important for ordered
|
||||
queues."""
|
||||
return None
|
||||
|
||||
|
||||
class QueueWithRemoval(BaseBundleQueue):
|
||||
"""Base class for storing bundles AND supporting remove(bundle)
|
||||
and contains(bundle) operations."""
|
||||
|
||||
def __contains__(self, bundle: RefBundle) -> bool:
|
||||
"""Return whether the key is in the queue."""
|
||||
...
|
||||
|
||||
def remove(self, bundle: RefBundle) -> RefBundle:
|
||||
"""Remove the specified bundle from the queue. If multiple instances exist, remove the first one."""
|
||||
bundle = self._remove_inner(bundle)
|
||||
self._on_dequeue_bundle(bundle)
|
||||
return bundle
|
||||
|
||||
def _remove_inner(self, bundle: RefBundle) -> RefBundle:
|
||||
raise NotImplementedError
|
||||
Reference in New Issue
Block a user