wehub-resource-sync
112 lines
3.4 KiB
Python
112 lines
3.4 KiB
Python
import collections
|
|
from dataclasses import dataclass, field
|
|
from typing import Generic, List, Optional, TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
@dataclass
|
|
class _WeightedBucket(Generic[T]):
|
|
"""A bucket of weighted items."""
|
|
|
|
items: List[T] = field(default_factory=list)
|
|
weight: float = 0
|
|
|
|
def add(
|
|
self,
|
|
item: T,
|
|
weight: float,
|
|
):
|
|
self.items.append(item)
|
|
self.weight += weight
|
|
|
|
def clear(self):
|
|
self.items.clear()
|
|
self.weight = 0
|
|
|
|
|
|
class WeightedRoundRobinPartitioner(Generic[T]):
|
|
"""Partitions weighted items into round-robin buckets.
|
|
|
|
Each item has an optional weight. If a weight is missing, the item is still
|
|
spread round-robin but doesn't contribute to bucket fullness.
|
|
"""
|
|
|
|
def __init__(
|
|
self,
|
|
*,
|
|
min_bucket_size: int,
|
|
max_bucket_size: int,
|
|
num_buckets: int,
|
|
emit_before_overflow: bool = False,
|
|
):
|
|
self._num_buckets = max(1, num_buckets)
|
|
self._min_bucket_size = min_bucket_size
|
|
self._max_bucket_size = max_bucket_size
|
|
self._emit_before_overflow = emit_before_overflow
|
|
|
|
self._buckets = [_WeightedBucket[T]() for _ in range(self._num_buckets)]
|
|
self._current_bucket_index = 0
|
|
self._output_queue: collections.deque[List[T]] = collections.deque()
|
|
|
|
def add_item(self, item: T, weight: Optional[float]) -> None:
|
|
current_bucket = self._current_bucket
|
|
|
|
# If a weight estimate isn't available, add the item to the current
|
|
# bucket and move on. This spreads unknown-size items evenly across
|
|
# buckets without pretending to know their size.
|
|
if weight is None:
|
|
current_bucket.add(item, 0)
|
|
self._advance_bucket()
|
|
return
|
|
|
|
# Do not truncate to int: in-memory size estimates are floating-point
|
|
# (e.g. file_size * encoding_ratio), and truncating each per-item
|
|
# estimate would shift bucket-fullness and emit/advance timing relative
|
|
# to accumulating the raw values.
|
|
weight = max(0.0, weight)
|
|
while (
|
|
self._emit_before_overflow
|
|
and current_bucket.items
|
|
and current_bucket.weight + weight > self._max_bucket_size
|
|
):
|
|
self._emit_current_bucket()
|
|
current_bucket = self._current_bucket
|
|
|
|
current_bucket.add(item, weight)
|
|
if current_bucket.weight >= self._max_bucket_size:
|
|
self._emit_current_bucket()
|
|
elif current_bucket.weight >= self._min_bucket_size:
|
|
self._advance_bucket()
|
|
|
|
def has_partition(self) -> bool:
|
|
return len(self._output_queue) > 0
|
|
|
|
def next_partition(self) -> List[T]:
|
|
return self._output_queue.popleft()
|
|
|
|
def finalize(self):
|
|
for bucket in self._buckets:
|
|
if bucket.items:
|
|
self._output_queue.append(list(bucket.items))
|
|
bucket.clear()
|
|
|
|
@property
|
|
def num_buckets(self) -> int:
|
|
return self._num_buckets
|
|
|
|
@property
|
|
def _current_bucket(self) -> _WeightedBucket[T]:
|
|
return self._buckets[self._current_bucket_index]
|
|
|
|
def _advance_bucket(self):
|
|
self._current_bucket_index = (
|
|
self._current_bucket_index + 1
|
|
) % self._num_buckets
|
|
|
|
def _emit_current_bucket(self):
|
|
current_bucket = self._current_bucket
|
|
self._output_queue.append(list(current_bucket.items))
|
|
current_bucket.clear()
|
|
self._advance_bucket()
|