diff --git a/backend/app/modelling/batching.py b/backend/app/modelling/batching.py index 7cf02afef..785855c41 100644 --- a/backend/app/modelling/batching.py +++ b/backend/app/modelling/batching.py @@ -6,6 +6,7 @@ same packing the trigger script has been running. """ from backend.app.modelling.property_filters import FilteredProperty +from utilities.grouped_batching import iter_grouped_batches BATCH_SIZE = 50 @@ -16,18 +17,8 @@ def pack_postcode_batches( """Pack *properties* into batches of ~*batch_size*, never splitting a postcode across batches. A single postcode larger than *batch_size* becomes its own oversized batch.""" - groups: dict[str, list[FilteredProperty]] = {} - for prop in properties: - groups.setdefault(prop.postcode or "", []).append(prop) - - batches: list[list[FilteredProperty]] = [] - current: list[FilteredProperty] = [] - for group in groups.values(): - if current and len(current) + len(group) > batch_size: - batches.append(current) - current = list(group) - else: - current.extend(group) - if current: - batches.append(current) - return batches + return list( + iter_grouped_batches( + properties, key=lambda p: p.postcode or "", max_batch_size=batch_size + ) + ) diff --git a/domain/addresses/postcode_batching.py b/domain/addresses/postcode_batching.py index ca4cd7529..7de6a5b43 100644 --- a/domain/addresses/postcode_batching.py +++ b/domain/addresses/postcode_batching.py @@ -3,7 +3,7 @@ from __future__ import annotations from collections.abc import Iterable, Iterator from domain.addresses.unstandardised_address import AddressList, UnstandardisedAddress -from domain.postcode import Postcode +from utilities.grouped_batching import iter_grouped_batches def iter_postcode_grouped_batches( @@ -11,41 +11,7 @@ def iter_postcode_grouped_batches( *, max_batch_size: int = 500, ) -> Iterator[AddressList]: - if max_batch_size < 1: - raise ValueError("max_batch_size must be >= 1") - - groups = _group_by_postcode_in_order(addresses) - - buffer: AddressList = AddressList([]) - for group in groups.values(): - group_len = len(group) - - # Oversize single-Postcode group: flush buffer first, then dispatch - # the group as its own batch. Mirrors the legacy - # ``if group_len >= batch_size`` branch. - if group_len >= max_batch_size: - if buffer: - yield buffer - buffer = AddressList([]) - yield group - continue - - # Adding this group would overflow: flush buffer before appending. - if len(buffer) + group_len > max_batch_size: - yield buffer - buffer = AddressList([]) - - buffer.extend(group) - - # Final flush. - if buffer: - yield buffer - - -def _group_by_postcode_in_order( - addresses: Iterable[UnstandardisedAddress], -) -> dict[Postcode, AddressList]: - groups: dict[Postcode, AddressList] = {} - for address in addresses: - groups.setdefault(address.postcode, AddressList([])).append(address) - return groups + for batch in iter_grouped_batches( + addresses, key=lambda a: a.postcode, max_batch_size=max_batch_size + ): + yield AddressList(batch) diff --git a/utilities/grouped_batching.py b/utilities/grouped_batching.py new file mode 100644 index 000000000..b5e42b3f9 --- /dev/null +++ b/utilities/grouped_batching.py @@ -0,0 +1,48 @@ +"""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