67 lines
1.6 KiB
Python
67 lines
1.6 KiB
Python
from __future__ import annotations
|
|
|
|
from collections import deque
|
|
from typing import TYPE_CHECKING, Any, Deque, Iterator, List, Optional
|
|
|
|
from typing_extensions import override
|
|
|
|
from .base import BaseBundleQueue
|
|
|
|
if TYPE_CHECKING:
|
|
from ray.data._internal.execution.interfaces import RefBundle
|
|
|
|
|
|
class FIFOBundleQueue(BaseBundleQueue):
|
|
"""A bundle queue that follows fifo-policy. Conceptually
|
|
[ ] <- [ ] <- [ ] ...
|
|
^ where the leftmost is popped first
|
|
|
|
NOTE: Not thread-safe
|
|
"""
|
|
|
|
def __init__(self, bundles: Optional[List[RefBundle]] = None):
|
|
super().__init__()
|
|
self._inner: Deque[RefBundle] = deque([])
|
|
if bundles is not None:
|
|
for bundle in bundles:
|
|
self.add(bundle)
|
|
|
|
@override
|
|
def _add_inner(self, bundle: RefBundle, **kwargs: Any):
|
|
self._inner.append(bundle)
|
|
|
|
@override
|
|
def _get_next_inner(self) -> RefBundle:
|
|
if not self.has_next():
|
|
raise ValueError(
|
|
f"Popping from empty {self.__class__.__name__} is prohibited"
|
|
)
|
|
|
|
bundle = self._inner.popleft()
|
|
return bundle
|
|
|
|
@override
|
|
def peek_next(self) -> Optional[RefBundle]:
|
|
if not self.has_next():
|
|
return None
|
|
return self._inner[0]
|
|
|
|
@override
|
|
def has_next(self) -> bool:
|
|
return len(self) > 0
|
|
|
|
@override
|
|
def finalize(self, **kwargs: Any):
|
|
pass
|
|
|
|
def __iter__(self) -> Iterator[RefBundle]:
|
|
yield from self._inner
|
|
|
|
def to_list(self) -> List[RefBundle]:
|
|
return list(self._inner)
|
|
|
|
@override
|
|
def clear(self):
|
|
self._reset_metrics()
|
|
self._inner.clear()
|