mirror of
https://github.com/Hestia-Homes/Model.git
synced 2026-07-12 13:29:04 +00:00
cc4bf4394e
Review feedback (#1481): the address batcher and the Modelling Run batcher implemented the same greedy packing; the core moves to utilities/grouped_batching.py and both become thin wrappers. Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
48 lines
1.4 KiB
Python
48 lines
1.4 KiB
Python
"""Greedy group-preserving batching.
|
|
|
|
Packs items into batches of at most ``max_batch_size`` without ever splitting
|
|
a group (items sharing a key) across batches; a single group larger than the
|
|
cap becomes its own oversized batch. Shared by the address postcode batcher
|
|
(``domain/addresses/postcode_batching.py``) and the Modelling Run distributor
|
|
(``backend/app/modelling/batching.py``).
|
|
"""
|
|
|
|
from collections.abc import Callable, Hashable, Iterable, Iterator
|
|
from typing import TypeVar
|
|
|
|
T = TypeVar("T")
|
|
|
|
|
|
def iter_grouped_batches(
|
|
items: Iterable[T],
|
|
*,
|
|
key: Callable[[T], Hashable],
|
|
max_batch_size: int,
|
|
) -> Iterator[list[T]]:
|
|
if max_batch_size < 1:
|
|
raise ValueError("max_batch_size must be >= 1")
|
|
|
|
groups: dict[Hashable, list[T]] = {}
|
|
for item in items:
|
|
groups.setdefault(key(item), []).append(item)
|
|
|
|
buffer: list[T] = []
|
|
for group in groups.values():
|
|
# Oversize single-key group: flush the buffer first, then dispatch
|
|
# the group as its own batch.
|
|
if len(group) >= max_batch_size:
|
|
if buffer:
|
|
yield buffer
|
|
buffer = []
|
|
yield group
|
|
continue
|
|
|
|
# Adding this group would overflow: flush the buffer before appending.
|
|
if len(buffer) + len(group) > max_batch_size:
|
|
yield buffer
|
|
buffer = []
|
|
|
|
buffer.extend(group)
|
|
|
|
if buffer:
|
|
yield buffer
|