Skip to content

Commit 8bd3252

Browse files
committed
allow empty sub-containers in reductions
1 parent 4c6a39e commit 8bd3252

File tree

1 file changed

+29
-7
lines changed

1 file changed

+29
-7
lines changed

arraycontext/container/traversal.py

+29-7
Original file line numberDiff line numberDiff line change
@@ -464,7 +464,7 @@ def rec_map_reduce_array_container(
464464
465465
or any other such traversal.
466466
"""
467-
def rec(_ary: ArrayOrContainerT) -> ArrayOrContainerT:
467+
def rec(_ary: ArrayOrContainerT) -> Optional[ArrayOrContainerT]:
468468
if type(_ary) is leaf_class:
469469
return map_func(_ary)
470470
else:
@@ -473,11 +473,22 @@ def rec(_ary: ArrayOrContainerT) -> ArrayOrContainerT:
473473
except NotAnArrayContainerError:
474474
return map_func(_ary)
475475
else:
476-
return reduce_func([
477-
rec(subary) for _, subary in iterable
478-
])
476+
subary_results = [
477+
rec(subary) for _, subary in iterable]
478+
filtered_subary_results = [
479+
result for result in subary_results
480+
if result is not None]
481+
if len(filtered_subary_results) > 0:
482+
return reduce_func(filtered_subary_results)
483+
else:
484+
return None
479485

480-
return rec(ary)
486+
result = rec(ary)
487+
488+
if result is None:
489+
raise ValueError("cannot reduce empty array container")
490+
491+
return result
481492

482493

483494
def rec_multimap_reduce_array_container(
@@ -503,12 +514,23 @@ def rec_multimap_reduce_array_container(
503514
# NOTE: this wrapper matches the signature of `deserialize_container`
504515
# to make plugging into `_multimap_array_container_impl` easier
505516
def _reduce_wrapper(ary: ContainerT, iterable: Iterable[Tuple[Any, Any]]) -> Any:
506-
return reduce_func([subary for _, subary in iterable])
517+
filtered_subary_results = [
518+
result for _, result in iterable
519+
if result is not None]
520+
if len(filtered_subary_results) > 0:
521+
return reduce_func(filtered_subary_results)
522+
else:
523+
return None
507524

508-
return _multimap_array_container_impl(
525+
result = _multimap_array_container_impl(
509526
map_func, *args,
510527
reduce_func=_reduce_wrapper, leaf_cls=leaf_class, recursive=True)
511528

529+
if result is None:
530+
raise ValueError("cannot reduce empty array container")
531+
532+
return result
533+
512534
# }}}
513535

514536

0 commit comments

Comments
 (0)