Skip to content

Commit 974bdfd

Browse files
committed
allow empty sub-containers in reductions
1 parent 9704e43 commit 974bdfd

File tree

1 file changed

+29
-7
lines changed

1 file changed

+29
-7
lines changed

arraycontext/container/traversal.py

Lines changed: 29 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -480,7 +480,7 @@ def rec_map_reduce_array_container(
480480
481481
or any other such traversal.
482482
"""
483-
def rec(_ary: ArrayOrContainerT) -> ArrayOrContainerT:
483+
def rec(_ary: ArrayOrContainerT) -> Optional[ArrayOrContainerT]:
484484
if type(_ary) is leaf_class:
485485
return map_func(_ary)
486486
else:
@@ -489,11 +489,22 @@ def rec(_ary: ArrayOrContainerT) -> ArrayOrContainerT:
489489
except NotAnArrayContainerError:
490490
return map_func(_ary)
491491
else:
492-
return reduce_func([
493-
rec(subary) for _, subary in iterable
494-
])
492+
subary_results = [
493+
rec(subary) for _, subary in iterable]
494+
filtered_subary_results = [
495+
result for result in subary_results
496+
if result is not None]
497+
if len(filtered_subary_results) > 0:
498+
return reduce_func(filtered_subary_results)
499+
else:
500+
return None
495501

496-
return rec(ary)
502+
result = rec(ary)
503+
504+
if result is None:
505+
raise ValueError("cannot reduce empty array container")
506+
507+
return result
497508

498509

499510
def rec_multimap_reduce_array_container(
@@ -519,12 +530,23 @@ def rec_multimap_reduce_array_container(
519530
# NOTE: this wrapper matches the signature of `deserialize_container`
520531
# to make plugging into `_multimap_array_container_impl` easier
521532
def _reduce_wrapper(ary: ContainerT, iterable: Iterable[Tuple[Any, Any]]) -> Any:
522-
return reduce_func([subary for _, subary in iterable])
533+
filtered_subary_results = [
534+
result for _, result in iterable
535+
if result is not None]
536+
if len(filtered_subary_results) > 0:
537+
return reduce_func(filtered_subary_results)
538+
else:
539+
return None
523540

524-
return _multimap_array_container_impl(
541+
result = _multimap_array_container_impl(
525542
map_func, *args,
526543
reduce_func=_reduce_wrapper, leaf_cls=leaf_class, recursive=True)
527544

545+
if result is None:
546+
raise ValueError("cannot reduce empty array container")
547+
548+
return result
549+
528550
# }}}
529551

530552

0 commit comments

Comments
 (0)