| 324 | def wrapper(*args: expr.ColumnExpression | api.Value) -> ColumnExpression: |
| 325 | @stateful_many |
| 326 | def stateful_wrapper( |
| 327 | packed_state: tuple[api.Value, tuple, int] | None, |
| 328 | rows: list[tuple[list[api.Value], int]], |
| 329 | ) -> tuple[api.Value, tuple, int] | None: |
| 330 | if packed_state is not None: |
| 331 | serialized_state, _positive_updates_tuple, _cnt = packed_state |
| 332 | state = reducer_cls.deserialize(serialized_state) |
| 333 | _positive_updates = list(_positive_updates_tuple) |
| 334 | else: |
| 335 | state = None |
| 336 | positive_updates: list[tuple[api.Value, ...]] = [] |
| 337 | negative_updates = [] |
| 338 | for row, count in rows: |
| 339 | if count > 0: |
| 340 | positive_updates.extend([tuple(row)] * count) |
| 341 | else: |
| 342 | negative_updates.extend([tuple(row)] * (-count)) |
| 343 | |
| 344 | if not retract_available and len(negative_updates) > 0: |
| 345 | if state is not None: |
| 346 | assert _positive_updates is not None |
| 347 | positive_updates.extend(_positive_updates) |
| 348 | _positive_updates = [] |
| 349 | state = None |
| 350 | acc = Counter(positive_updates) |
| 351 | acc.subtract(negative_updates) |
| 352 | assert all(x >= 0 for x in acc.values()) |
| 353 | positive_updates = list(acc.elements()) |
| 354 | negative_updates = [] |
| 355 | |
| 356 | if state is None: |
| 357 | if neutral_available: |
| 358 | state = reducer_cls.neutral() |
| 359 | _positive_updates = [] |
| 360 | _cnt = 0 |
| 361 | elif len(positive_updates) == 0: |
| 362 | if len(negative_updates) == 0: |
| 363 | return None |
| 364 | else: |
| 365 | raise ValueError( |
| 366 | "Unable to process negative update with this stateful reducer." |
| 367 | ) |
| 368 | else: |
| 369 | if sort_by_available: |
| 370 | positive_updates.sort( |
| 371 | key=lambda x: reducer_cls.sort_by(list(x)) |
| 372 | ) |
| 373 | state = reducer_cls.from_row(list(positive_updates[0])) |
| 374 | if not retract_available: |
| 375 | _positive_updates = positive_updates[0:1] |
| 376 | _cnt = 0 |
| 377 | else: |
| 378 | _positive_updates = [] |
| 379 | _cnt = 1 |
| 380 | positive_updates = positive_updates[1:] |
| 381 | |
| 382 | updates = [(row_up, False) for row_up in positive_updates] + [ |
| 383 | (row_up, True) for row_up in negative_updates |