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

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

@@ -623,6 +623,46 @@
   settings take effect: `primary_key=True` adds the column to the primary key,
   which can turn a model's key composite, and `nullable=True` makes the column
   nullable. Review such fields before upgrading.
+- **core:** a `date` field is inferred as a `Date` column and a `time` field
+  as a `Time` column (`loom.core.model.Date`, `loom.core.model.Time`, new),
+  where both used to be inferred as an unbounded `String`; an enum mixed with
+  `date` or `time` follows to `Date` or `Time`. The SQLAlchemy repository
+  binds `date` and `time` values natively, so on Postgres with asyncpg writing
+  either field failed (`expected str, got date`, and the same for `time`). The
+  SQLAlchemy repository now refuses a `time` that carries a UTC offset when it
+  is written to, or compared with, a column that would drop the offset, with a
+  `BoundaryValidationError` (a `ValueError`) naming the column and not the
+  value: `boundary_validation` for a write (REST answers `422`),
+  `query_validation` for a filter, `get_by`, `exists_by` or `list_paginated`
+  value (`400`). Without the check, Postgres and SQLite would store
+  `08:30+02:00` in a `Time` column as `08:30`, and an aware filter value would
+  match the row holding `08:30`. A JSON column and a `Time` column declared
+  with `ColumnType("Time", kwargs={"timezone": True})` accept it; the JSON
+  column keeps the written offset, and so does the `Time` column on Postgres
+  (`timetz`), which takes only a `time` with an offset: one without (naive, or
+  with a zone but no fixed offset) is refused the same way, under
+  `time must carry a UTC offset`, since asyncpg cannot encode it for `timetz`;
+  every refused column of a write is reported, in name order. On any other
+  dialect startup refuses a `Time(timezone=True)` column with a `ConfigError`
+  naming `table.column`, since the dialect would not keep the offset. A
+  `datetime` field (`DateTime(tz=True)`) keeps the instant on Postgres
+  (`timestamptz`) and MongoDB, read back in UTC, not the written offset.
+  `get_by`, `exists_by` and the `filters` of `list_paginated` now convert
+  their value as a write converts it, as `QuerySpec` filters already did, so a
+  plain `Enum` member matches the value the write stored instead of failing to
+  bind. A `list_paginated` filter, `get_by` or `exists_by` on a name that is
+  not a mapped column of the model (a missing field, a relation, an attribute
+  such as `metadata`) raises `ValueError` naming it. An existing Postgres
+  table holding these columns as `varchar` needs a migration: Alembic
+  autogenerate detects the type change, but Postgres refuses it without a
+  cast, so the revision passes `postgresql_using="c::date"` (`"c::time"` for a
+  `time` field) to `op.alter_column`, which is
+  `ALTER TABLE t ALTER COLUMN c TYPE date USING c::date`. `c::time` drops any
+  offset the stored text holds (`'08:30+02'::time` is `08:30`): check first
+  with `SELECT c FROM t WHERE c ~ '[+-]\d{2}(:?\d{2})?$'`, and if offsets are
+  present convert those rows deliberately (for example to UTC with
+  `(c::timetz AT TIME ZONE 'UTC')::time`) before changing the type. MongoDB is
+  unaffected: it stores `date` and `time` values as ISO 8601 strings.
 
 ### ai
 
@@ -909,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
@@ -1367,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-6.0.2 → loom_kernel-8.0.0}/PKG-INFO

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

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/backend/sqlalchemy.py

@@ -15,6 +15,7 @@ from sqlalchemy import (
     CheckConstraint,
     Column,
     Constraint,
+    Date,
     DateTime,
     Float,
     ForeignKey,
@@ -27,6 +28,7 @@ from sqlalchemy import (
     String,
     Table,
     Text,
+    Time,
     UniqueConstraint,
     func,
     text,
@@ -80,7 +82,9 @@ _SA_TYPE_MAP: dict[str, type] = {
     "Text": Text,
     "JSON": JSON,
     "Bytes": LargeBinary,
+    "Date": Date,
     "DateTime": DateTime,
+    "Time": Time,
     "Numeric": Numeric,
     "Postgres.JSONB": PG_JSONB,
     "Postgres.UUID": PG_UUID,

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/model/__init__.py

@@ -44,6 +44,7 @@ from loom.core.model.types import (
     BigInteger,
     Boolean,
     Bytes,
+    Date,
     DateTime,
     Float,
     Integer,
@@ -51,6 +52,7 @@ from loom.core.model.types import (
     Numeric,
     String,
     Text,
+    Time,
 )
 from loom.core.model.types_postgres import Postgres
 
@@ -85,6 +87,7 @@ __all__ = [
     "ColumnField",
     "ScopedField",
     "ColumnType",
+    "Date",
     "DateTime",
     "Field",
     "Float",
@@ -105,6 +108,7 @@ __all__ = [
     "ServerOnUpdate",
     "String",
     "Text",
+    "Time",
     "UnsupportedBoundaryType",
     "get_column_fields",
     "get_id_attribute",

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/model/_loom_type.py

@@ -107,7 +107,9 @@ class BoundaryValidationError(LoomError, ValueError):
     reasons drop the detail pydantic derives from the input; the msgspec
     message is forwarded as one line, without the value msgspec echoes for an
     invalid enum value or tag. A message raised by the type's own validator
-    keeps its text. It is also a
+    keeps its text. A repository raises it for a value its storage cannot
+    hold, such as the SQLAlchemy repository for a ``time`` whose UTC offset,
+    present or missing, its column cannot take. It is also a
     :class:`ValueError`, so an ``except ValueError`` that used to catch
     ``msgspec.ValidationError`` or ``pydantic.ValidationError`` keeps
     catching it.

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/model/introspection.py

@@ -31,11 +31,13 @@ from loom.core.model.types import (
     JSON,
     Boolean,
     Bytes,
+    Date,
     DateTime,
     Float,
     Integer,
     Numeric,
     String,
+    Time,
 )
 from loom.core.schema_names import sql_identifier
 
@@ -405,8 +407,8 @@ _SCALAR_TYPE_MAP: dict[type, ColumnType] = {
     bool: Boolean,
     bytes: Bytes,
     datetime: DateTime(tz=True),
-    date: String(None),
-    time: String(None),
+    date: Date,
+    time: Time,
     Decimal: Numeric(),
 }
 

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/model/types.py

@@ -20,6 +20,8 @@ Boolean = ColumnType("Boolean")
 Text = ColumnType("Text")
 JSON = ColumnType("JSON")
 Bytes = ColumnType("Bytes")
+Date = ColumnType("Date")
+Time = ColumnType("Time")
 
 
 def DateTime(*, tz: bool = True) -> ColumnType:

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/repository/sqlalchemy/backend.py

@@ -9,7 +9,7 @@ from contextlib import asynccontextmanager
 from typing import ClassVar, Final, Literal
 
 import msgspec
-from sqlalchemy import inspect, make_url, text
+from sqlalchemy import MetaData, Time, inspect, make_url, text
 from sqlalchemy.ext.asyncio import AsyncConnection
 
 from loom.core.authz.elevation import ElevationSink
@@ -205,6 +205,31 @@ def startup_checks(
     return unprotected
 
 
+def check_time_zone_columns(dialect: str, metadata: MetaData) -> None:
+    """Refuse a ``Time(timezone=True)`` column on a dialect that drops its offset.
+
+    Postgres stores such a column as ``time with time zone`` and keeps the
+    offset; SQLite ignores the zone and stores the time without it.
+
+    Raises:
+        ConfigError: When *metadata* holds such a column and *dialect* is not
+            Postgres, naming each ``table.column``.
+    """
+    if dialect == POSTGRES_DIALECT:
+        return
+    columns = sorted(
+        f"{table.name}.{column.name}"
+        for table in metadata.tables.values()
+        for column in table.columns
+        if isinstance(column.type, Time) and column.type.timezone
+    )
+    if columns:
+        raise ConfigError(
+            f"dialect {dialect!r} would drop the UTC offset of the Time(timezone=True) "
+            f"columns {', '.join(columns)}; declare a JSON column to keep the offset"
+        )
+
+
 @asynccontextmanager
 async def _lifespan(
     session_manager: SessionManager,
@@ -223,6 +248,7 @@ async def _lifespan(
         scoped,
         has_session_settings=session_manager.has_session_settings,
     )
+    check_time_zone_columns(session_manager.engine.dialect.name, get_metadata())
     if unprotected:
         _logger.warning(
             "scoped tables left unprotected on this dialect: %s", ", ".join(unprotected)

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/repository/sqlalchemy/mixins.py

@@ -11,6 +11,7 @@ from sqlalchemy import update as sa_update
 from sqlalchemy.ext.asyncio import AsyncSession
 
 from loom.core.backend.sqlalchemy import get_compiled_core
+from loom.core.errors.codes import ErrorCode
 from loom.core.model.convert import to_struct
 from loom.core.model.introspection import (
     get_column_fields,
@@ -33,7 +34,12 @@ from loom.core.repository.mutation import MutationEvent
 from loom.core.repository.sqlalchemy.integrity import handle_integrity_errors
 from loom.core.repository.sqlalchemy.query_compiler.compiler import QuerySpecCompiler
 from loom.core.repository.sqlalchemy.transactional import record_mutation
-from loom.core.repository.sqlalchemy.values import is_json_column, to_column_value
+from loom.core.repository.sqlalchemy.values import (
+    check_time_offsets,
+    is_json_column,
+    to_column_value,
+    to_comparison_value,
+)
 
 _SENTINEL = object()
 _TOTAL_COUNT_ALIAS = "__loom_total_count"
@@ -161,7 +167,14 @@ class SQLAlchemyContextMixin(Generic[OutputT, IdT]):
             builtins = to_column_value(data)
             if not isinstance(builtins, dict):
                 raise TypeError("Struct payload must serialize to dict")
-            return self._encode_json_columns(_to_internal_field_names(type(data), builtins))
+            values = self._encode_json_columns(_to_internal_field_names(type(data), builtins))
+            sa_model = self._effective_sa_model
+            names = (self._column_field_names or frozenset()).intersection(values)
+            check_time_offsets(
+                ((getattr(sa_model, name), values[name]) for name in names),
+                code=ErrorCode.BOUNDARY_VALIDATION,
+            )
+            return values
         return data
 
     def _encode_json_columns(self, values: dict[str, Any]) -> dict[str, Any]:
@@ -174,10 +187,11 @@ class SQLAlchemyContextMixin(Generic[OutputT, IdT]):
         """Resolve a model column attribute by field name.
 
         Raises:
-            ValueError: If ``field`` does not exist on the SQLAlchemy model.
+            ValueError: If ``field`` is not a mapped column of the SQLAlchemy
+                model.
         """
         sa_model = self._effective_sa_model
-        if not hasattr(sa_model, field):
+        if field not in inspect(sa_model).column_attrs:
             raise ValueError(f"Unknown field '{field}' for model '{self.model.__name__}'")
         return getattr(sa_model, field)
 
@@ -189,7 +203,12 @@ class SQLAlchemyContextMixin(Generic[OutputT, IdT]):
         """Apply page-mode filters to an offset items statement."""
         if not filter_params or not filter_params.filters:
             return items_stmt
-        return items_stmt.filter_by(**filter_params.filters)
+        return items_stmt.filter_by(
+            **{
+                name: to_comparison_value(self._column_for_field(name), value)
+                for name, value in filter_params.filters.items()
+            }
+        )
 
     def _to_output(
         self,
@@ -441,7 +460,9 @@ class SQLAlchemyReadMixin(SQLAlchemyContextMixin[OutputT, IdT], Generic[OutputT,
         """Fetch one entity by equality over an arbitrary column field."""
         async with self._session_scope() as scoped_session:
             core_model = self._effective_core_model
-            stmt = core_model.select(profile).where(self._column_for_field(field) == value).limit(1)
+            column = self._column_for_field(field)
+            stmt = core_model.select(profile).where(column == to_comparison_value(column, value))
+            stmt = stmt.limit(1)
             loaded = await core_model.fetch_one(scoped_session, stmt, profile=profile)
             return cast(OutputT | None, loaded)
 
@@ -575,7 +596,8 @@ class SQLAlchemyReadMixin(SQLAlchemyContextMixin[OutputT, IdT], Generic[OutputT,
 
     async def exists_by(self, field: str, value: Any) -> bool:
         """Check whether an entity exists by an arbitrary column field."""
-        return await self._exists_where(self._column_for_field(field) == value)
+        column = self._column_for_field(field)
+        return await self._exists_where(column == to_comparison_value(column, value))
 
     async def _exists_where(self, clause: Any) -> bool:
         async with self._session_scope() as scoped_session:

{loom_kernel-6.0.2 → loom_kernel-8.0.0}/src/loom/core/repository/sqlalchemy/query_compiler/filters.py

@@ -15,7 +15,7 @@ from loom.core.repository.sqlalchemy.query_compiler.subquery import (
     compile_exists,
     compile_not_exists,
 )
-from loom.core.repository.sqlalchemy.values import is_json_column, to_column_value
+from loom.core.repository.sqlalchemy.values import to_comparison_value
 
 _RELATION_OPS = frozenset([FilterOp.EXISTS, FilterOp.NOT_EXISTS])
 _SIMPLE_OPS: dict[FilterOp, Callable[[Any, Any], Any]] = {
@@ -50,6 +50,8 @@ def compile_filter_group(
     Raises:
         UnsafeFilterError: If a field is not in ``allowed_fields``.
         FilterPathError: If a field path cannot be resolved.
+        BoundaryValidationError: ``query_validation``, for a ``time`` whose
+            offset the compared column cannot take.
     """
     clauses = [_compile_spec(sa_model, spec, allowed_fields) for spec in group.filters]
     if group.op == "OR":
@@ -78,17 +80,22 @@ def _compile_spec(
 def _apply_op(col: Any, op: FilterOp, value: Any) -> Any:
     """Apply *op*, converting *value* as the repository converts a value it writes.
 
-    A SQLAlchemy expression, and an ``IN`` value that is a string or not
-    iterable, reach SQLAlchemy as they are.
+    A ``time`` whose offset the column cannot take is refused, as
+    :func:`to_comparison_value` refuses it. A SQLAlchemy expression, and an
+    ``IN`` value that is a string or not iterable, reach SQLAlchemy as they
+    are.
+
+    Raises:
+        ValueError: If *op* is not a supported operator.
+        BoundaryValidationError: ``query_validation``, for a refused ``time``.
     """
     handler = _SIMPLE_OPS.get(op)
     if handler is None:
         raise ValueError(f"Unsupported FilterOp: {op!r}")
     if isinstance(value, ClauseElement):
         return handler(col, value)
-    json = is_json_column(col)
     if op is not FilterOp.IN:
-        return handler(col, to_column_value(value, json=json))
+        return handler(col, to_comparison_value(col, value))
     if isinstance(value, str | bytes) or not isinstance(value, Iterable):
         return handler(col, value)
-    return handler(col, [to_column_value(item, json=json) for item in value])
+    return handler(col, [to_comparison_value(col, item) for item in value])

loom_kernel-8.0.0/src/loom/core/repository/sqlalchemy/values.py

@@ -0,0 +1,130 @@
+"""Column form of model values in SQLAlchemy rows.
+
+SQLAlchemy binds ``bytes``, ``datetime``, ``date``, ``time``, ``Decimal`` and
+``UUID`` natively; every other value is reduced to msgspec builtins (an
+``Enum`` to its value, a struct to a dict). A JSON column holds builtins only,
+so its values are reduced completely.
+
+:func:`to_column_value` is the conversion the repository applies to a value it
+writes, and :func:`to_comparison_value` the one it applies to a value compared
+with a column (a filter, ``get_by``, ``exists_by``), so a comparison matches
+exactly what the row holds.
+
+A ``time`` with a UTC offset is refused for a column that cannot hold the
+offset: SQL ``TIME`` without time zone (Postgres ``time``, SQLite text) stores
+``08:30+02:00`` as ``08:30``, and an aware value compared with it matches the
+row holding ``08:30``. A JSON column, which keeps the ISO text, takes any
+``time``. A ``Time`` column with ``timezone=True`` (Postgres ``timetz``) takes
+only a ``time`` with an offset: asyncpg cannot encode a naive one for it. The
+SQLAlchemy backend refuses such a column at startup on any dialect other than
+Postgres.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable
+from datetime import date, datetime, time
+from decimal import Decimal
+from typing import Any
+from uuid import UUID
+
+import msgspec
+from sqlalchemy import JSON, Time
+
+from loom.core.errors.codes import ErrorCode
+from loom.core.model import BoundaryValidationError
+
+# Values SQLAlchemy binds natively; encoding them to strings breaks typed columns.
+_NATIVE_COLUMN_TYPES = (bytes, datetime, date, time, Decimal, UUID)
+
+
+def to_column_value(value: object, *, json: bool = False) -> Any:
+    """Return the form ``value`` takes in a column.
+
+    Args:
+        value: A model value or a filter value.
+        json: ``True`` when the column is a JSON column.
+
+    Returns:
+        ``value`` reduced to msgspec builtins, keeping the natively bound
+        types unless the column is a JSON column.
+    """
+    if json:
+        return msgspec.to_builtins(value)
+    return msgspec.to_builtins(value, builtin_types=_NATIVE_COLUMN_TYPES)
+
+
+_OFFSET_REFUSED = "time must not carry a UTC offset"
+_OFFSET_REQUIRED = "time must carry a UTC offset"
+
+
+def check_time_offsets(bindings: Iterable[tuple[Any, object]], *, code: str) -> None:
+    """Raise when a time's offset, present or missing, does not fit the column it is bound to.
+
+    A JSON column takes any ``time``. A ``Time(timezone=True)`` column keeps the
+    offset and takes only a ``time`` that has one (``utcoffset()`` is not
+    ``None``); every other column drops the offset and takes only a ``time``
+    without one. A model that needs the offset declares a JSON column or
+    ``ColumnType("Time", kwargs={"timezone": True})`` (Postgres only); on
+    Postgres a ``datetime`` field keeps the instant instead. The error names
+    the columns, not the values.
+
+    Args:
+        bindings: ``(column, value)`` pairs, each a SQLAlchemy column or mapped
+            attribute and a value about to be written to or compared with it.
+        code: ``boundary_validation`` for written values, ``query_validation``
+            for compared values.
+
+    Raises:
+        BoundaryValidationError: With one violation per refused column, in
+            column name order.
+    """
+    violations = sorted(
+        (column.key, reason)
+        for column, value in bindings
+        if (reason := _time_offset_refusal(column, value)) is not None
+    )
+    if violations:
+        message = "; ".join(f"{name}: {reason}" for name, reason in violations)
+        raise BoundaryValidationError(message, tuple(violations), code=code)
+
+
+def _time_offset_refusal(column: Any, value: object) -> str | None:
+    if not isinstance(value, time) or is_json_column(column):
+        return None
+    column_type = getattr(column, "type", None)
+    keeps_offset = isinstance(column_type, Time) and bool(column_type.timezone)
+    if (value.utcoffset() is not None) == keeps_offset:
+        return None
+    return _OFFSET_REQUIRED if keeps_offset else _OFFSET_REFUSED
+
+
+def to_comparison_value(column: Any, value: object) -> Any:
+    """Return *value* in the form it is compared with *column*.
+
+    Refused as :func:`check_time_offsets` refuses it, and converted as
+    :func:`to_column_value` converts a written value. A refusal is raised under
+    ``query_validation`` whatever the value's source, so ``get_by`` and
+    ``exists_by`` called with a value from a command body raise it too.
+
+    Args:
+        column: A SQLAlchemy column or mapped attribute.
+        value: A filter or lookup value.
+
+    Raises:
+        BoundaryValidationError: ``query_validation``, for a ``time`` whose
+            offset *column* cannot take, as :func:`check_time_offsets` decides.
+    """
+    check_time_offsets(((column, value),), code=ErrorCode.QUERY_VALIDATION)
+    return to_column_value(value, json=is_json_column(column))
+
+
+def is_json_column(column: object) -> bool:
+    """Whether *column*, a SQLAlchemy column or mapped attribute, holds JSON.
+
+    ``JSON`` and the dialect types derived from it (Postgres ``JSONB``) count.
+    """
+    return isinstance(getattr(column, "type", None), JSON)
+
+
+__all__ = ["check_time_offsets", "is_json_column", "to_column_value", "to_comparison_value"]

{loom_kernel-6.0.2 → 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-6.0.2 → 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(