loom-kernel 7.0.0__tar.gz → 8.0.0__tar.gz

{loom_kernel-7.0.0 → loom_kernel-8.0.0}/CHANGELOG.md

@@ -949,6 +949,24 @@
 
 ### streaming
 
+- **streaming:** `BatchExpandStep.execute` returns one list of payloads per
+  message of its batch, in batch order, instead of one flat list. List `i`
+  holds what `messages[i]` expands into; each output takes that message's
+  metadata, and an empty list filters the message out. A result whose length
+  differs from the batch, a flat list of payloads, a `Message` among the
+  outputs, or a `str`, `bytes`, `bytearray` or mapping in place of the result
+  or of one of its lists raises `TypeError`, and the batch becomes error
+  envelopes.
+
+  ```python
+  # before
+  def execute(self, messages: list[Message[Event]], **kwargs: object) -> list[Score]:
+      return [Score(value=m.payload.amount) for m in messages]
+
+  # after
+  def execute(self, messages: list[Message[Event]], **kwargs: object) -> list[list[Score]]:
+      return [[Score(value=m.payload.amount)] for m in messages]
+  ```
 - **streaming:** compile issues come from the code catalogue
   (`StreamingErrorCode.X.issue(...)`), and the compiler's private issue
   factories are removed. `loom.streaming.compiler.CompilationIssue` is a type
@@ -1407,6 +1425,14 @@
   `Drain` default failed with `ROUTER_BRANCH_SHAPE_DIVERGENCE`, and a step
   after a `Router` without a default whose branches all end their path failed
   with `SHAPE_MISMATCH`; both compile now.
+- **streaming:** a `BatchExpandStep` expands. Every output had to pair with
+  one input message, so a step returning more or fewer outputs than its batch
+  raised and each message of the batch became an error envelope. Under
+  at-least-once delivery each source offset now completes once per output,
+  after that output's write, and a message with no output completes its
+  offset at the step, also when the batch holds several messages of one
+  offset that an earlier expanding node produced. The new return type is under
+  Behaviour changes.
 
 ### prefect
 

{loom_kernel-7.0.0 → loom_kernel-8.0.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.5
 Name: loom-kernel
-Version: 7.0.0
+Version: 8.0.0
 Summary: Loom Python project
 License-File: LICENSE
 Classifier: Programming Language :: Python :: 3

{loom_kernel-7.0.0 → loom_kernel-8.0.0}/src/loom/streaming/bytewax/handlers/_shared.py

@@ -36,6 +36,7 @@ if TYPE_CHECKING:
 
 Stream: TypeAlias = Any
 AwaitT = TypeVar("AwaitT")
+ExpandT = TypeVar("ExpandT", bound=Iterable[Any])
 
 
 @runtime_checkable
@@ -91,11 +92,8 @@ class _ExecutableBatchExpandStep(Protocol):
         self,
         messages: list[Message[StreamPayload]],
         **kwargs: object,
-    ) -> (
-        Iterable[StreamPayload | Message[StreamPayload]]
-        | Awaitable[Iterable[StreamPayload | Message[StreamPayload]]]
-    ):
-        """Expand one batch into many payloads or replacement messages."""
+    ) -> Iterable[Iterable[StreamPayload]] | Awaitable[Iterable[Iterable[StreamPayload]]]:
+        """Expand each message of one batch into its own list of payloads."""
         ...
 
 
@@ -259,10 +257,9 @@ def _resolve_batch_result(
 
 
 def _resolve_expand_result(
-    result: Iterable[StreamPayload | Message[StreamPayload]]
-    | Awaitable[Iterable[StreamPayload | Message[StreamPayload]]],
+    result: ExpandT | Awaitable[ExpandT],
     node_type: str,
-) -> Iterable[StreamPayload | Message[StreamPayload]]:
+) -> ExpandT:
     """Resolve a synchronous expanding result and reject awaitables."""
     if isinstance(result, Awaitable):
         raise TypeError(f"{node_type} returned an awaitable outside WithAsync.")
@@ -353,9 +350,11 @@ def _reconcile_fanout(
 
     An expanding node turns one input record into N outputs that all carry the
     same source offset, and every one of them completes that offset when it
-    reaches a terminal. The tracker starts each record expecting a single
-    completion, so accounting has to be corrected here — while the outputs are
-    still a plain list and none of them can have reached a terminal yet.
+    reaches a terminal. Each input already owes one completion, and a batch can
+    hold several inputs of one offset when an earlier node expanded it, so per
+    offset the outputs are compared with the inputs: extra outputs fork the
+    difference, and missing ones complete it here — while the outputs are still
+    a plain list and none of them can have reached a terminal yet.
 
     Two failures this prevents, both silent:
 
@@ -372,20 +371,21 @@ def _reconcile_fanout(
     """
     if tracker is None:
         return
-    produced: Counter[tuple[str, int, int]] = Counter()
-    for item in results:
-        key = _commit_key(item)
-        if key is not None:
-            produced[key] += 1
-    for key, count in produced.items():
-        if count > 1:
-            tracker.fork(key[0], key[1], key[2], count - 1)
-    for original in originals:
-        key = _commit_key(original)
-        if key is not None and key not in produced:
+    expected = _commit_key_counts(originals)
+    produced = _commit_key_counts(results)
+    for key in expected | produced:
+        delta = produced[key] - expected[key]
+        if delta > 0:
+            tracker.fork(key[0], key[1], key[2], delta)
+        for _ in range(-delta):
             tracker.complete(key[0], key[1], key[2])
 
 
+def _commit_key_counts(items: Sequence[Any]) -> Counter[tuple[str, int, int]]:
+    """Count the items accountable for each source offset."""
+    return Counter(key for item in items if (key := _commit_key(item)) is not None)
+
+
 def _drop_and_commit(item: Any, tracker: CommitCompletionPort) -> tuple[()]:
     """Drop one item and mark it complete for commit tracking.
 

{loom_kernel-7.0.0 → loom_kernel-8.0.0}/src/loom/streaming/bytewax/handlers/steps.py

@@ -2,7 +2,7 @@
 
 from __future__ import annotations
 
-from collections.abc import Iterable
+from collections.abc import Iterable, Mapping
 from typing import Any
 
 from bytewax.operators import flat_map
@@ -134,8 +134,6 @@ def _apply_expand_step(
 
     sid = _step_id(f"expand_{idx}_{name}", ctx)
     mapped = bw_map(sid, stream, step_fn)
-    from bytewax.operators import flat_map
-
     flattened_stream = flat_map(_step_id(f"flatten_expand_{idx}", ctx), mapped, _identity)
     return _split_node_result(flattened_stream, sid, ctx, ErrorKind.TASK)
 
@@ -296,9 +294,44 @@ def _execute_batch_expand_step(
         correlation_id=messages[0].meta.correlation_id if messages else None,
     ):
         result = _resolve_expand_result(batch_expand_step.execute(messages), name)
-        if not isinstance(result, Iterable):
-            raise TypeError(f"{name} must return an iterable of payloads.")
-        return _replace_payloads(messages, list(result))
+        return _outputs_per_message(messages, _result_list(result, name), name)
+
+
+def _outputs_per_message(
+    messages: list[Message[StreamPayload]],
+    groups: list[Any],
+    name: str,
+) -> list[Message[StreamPayload]]:
+    """Give each output the metadata of the message whose list it is in.
+
+    ``groups[i]`` holds the outputs of ``messages[i]``; an empty list filters
+    that message out.
+    """
+    if len(groups) != len(messages):
+        raise TypeError(
+            f"{name} returned {len(groups)} output lists for {len(messages)} messages; "
+            "return one list of payloads per message, in batch order."
+        )
+    outputs: list[Message[StreamPayload]] = []
+    for message, payloads in zip(messages, groups, strict=True):
+        for payload in _result_list(payloads, name):
+            if isinstance(payload, Message):
+                raise TypeError(
+                    f"{name} must return payloads, not messages: each output "
+                    "takes the metadata of the message it came from."
+                )
+            outputs.append(_replace_payload(message, payload))
+    return outputs
+
+
+def _result_list(value: object, name: str) -> list[Any]:
+    """Materialize one list of a batch-expanding result, iterating it once.
+
+    ``str``, ``bytes`` and mappings are iterable but are not lists of results.
+    """
+    if not isinstance(value, Iterable) or isinstance(value, (str, bytes, bytearray, Mapping)):
+        raise TypeError(f"{name} must return one list of payloads per message.")
+    return list(value)
 
 
 def _execute_explode_step(

{loom_kernel-7.0.0 → loom_kernel-8.0.0}/src/loom/streaming/nodes/_step.py

@@ -130,7 +130,7 @@ class ExpandStep(Step[InT, OutT], ABC):
 
 
 class BatchExpandStep(Step[InT, OutT], ABC):
-    """Streaming step that expands one batch into many output messages.
+    """Streaming step that expands each message of a batch into its own outputs.
 
     Declare the subclass itself in a flow, not an instance. The compiler
     materializes the class during binding resolution.
@@ -138,8 +138,17 @@ class BatchExpandStep(Step[InT, OutT], ABC):
     Pattern:
         Batch fan-out step.
 
-    Subclasses define ``execute(self, messages, **kwargs) -> Iterable[OutT]``
-    with the explicit batch and dependency signature they need.
+    Subclasses define ``execute(self, messages, **kwargs) -> list[list[OutT]]``
+    with the explicit batch and dependency signature they need. The result
+    holds one list per message, in batch order: list ``i`` holds the payloads
+    ``messages[i]`` expands into, and each takes that message's metadata. An
+    empty list filters the message out. The result and each of its lists are
+    iterated once, so generators work.
+
+    These results raise ``TypeError`` and every message of the batch becomes
+    an error envelope: a result whose length differs from the batch, a flat
+    list of payloads, a ``Message`` among the outputs, and a ``str``,
+    ``bytes`` or mapping in place of the result or of one of its lists.
     """
 
     router_branch_safe: ClassVar[bool] = True