Skip to content

Commit 4a6c3c5

Browse files
authored
CLN: Simplify aggregation.aggregate (#37033)
1 parent 03709d4 commit 4a6c3c5

File tree

1 file changed

+18
-69
lines changed

1 file changed

+18
-69
lines changed

pandas/core/aggregation.py

+18-69
Original file line numberDiff line numberDiff line change
@@ -608,92 +608,41 @@ def aggregate(obj, arg: AggFuncType, *args, **kwargs):
608608

609609
from pandas.core.reshape.concat import concat
610610

611-
def _agg_1dim(name, how, subset=None):
612-
"""
613-
aggregate a 1-dim with how
614-
"""
615-
colg = obj._gotitem(name, ndim=1, subset=subset)
616-
if colg.ndim != 1:
617-
raise SpecificationError(
618-
"nested dictionary is ambiguous in aggregation"
619-
)
620-
return colg.aggregate(how)
621-
622-
def _agg_2dim(how):
623-
"""
624-
aggregate a 2-dim with how
625-
"""
626-
colg = obj._gotitem(obj._selection, ndim=2, subset=selected_obj)
627-
return colg.aggregate(how)
628-
629-
def _agg(arg, func):
630-
"""
631-
run the aggregations over the arg with func
632-
return a dict
633-
"""
634-
result = {}
635-
for fname, agg_how in arg.items():
636-
result[fname] = func(fname, agg_how)
637-
return result
611+
if selected_obj.ndim == 1:
612+
# key only used for output
613+
colg = obj._gotitem(obj._selection, ndim=1)
614+
results = {key: colg.agg(how) for key, how in arg.items()}
615+
else:
616+
# key used for column selection and output
617+
results = {
618+
key: obj._gotitem(key, ndim=1).agg(how) for key, how in arg.items()
619+
}
638620

639621
# set the final keys
640622
keys = list(arg.keys())
641623

642-
if obj._selection is not None:
643-
644-
sl = set(obj._selection_list)
645-
646-
# we are a Series like object,
647-
# but may have multiple aggregations
648-
if len(sl) == 1:
649-
650-
result = _agg(
651-
arg, lambda fname, agg_how: _agg_1dim(obj._selection, agg_how)
652-
)
653-
654-
# we are selecting the same set as we are aggregating
655-
elif not len(sl - set(keys)):
656-
657-
result = _agg(arg, _agg_1dim)
658-
659-
# we are a DataFrame, with possibly multiple aggregations
660-
else:
661-
662-
result = _agg(arg, _agg_2dim)
663-
664-
# no selection
665-
else:
666-
667-
try:
668-
result = _agg(arg, _agg_1dim)
669-
except SpecificationError:
670-
671-
# we are aggregating expecting all 1d-returns
672-
# but we have 2d
673-
result = _agg(arg, _agg_2dim)
674-
675624
# combine results
676625

677626
def is_any_series() -> bool:
678627
# return a boolean if we have *any* nested series
679-
return any(isinstance(r, ABCSeries) for r in result.values())
628+
return any(isinstance(r, ABCSeries) for r in results.values())
680629

681630
def is_any_frame() -> bool:
682631
# return a boolean if we have *any* nested series
683-
return any(isinstance(r, ABCDataFrame) for r in result.values())
632+
return any(isinstance(r, ABCDataFrame) for r in results.values())
684633

685-
if isinstance(result, list):
686-
return concat(result, keys=keys, axis=1, sort=True), True
634+
if isinstance(results, list):
635+
return concat(results, keys=keys, axis=1, sort=True), True
687636

688637
elif is_any_frame():
689638
# we have a dict of DataFrames
690639
# return a MI DataFrame
691640

692-
keys_to_use = [k for k in keys if not result[k].empty]
641+
keys_to_use = [k for k in keys if not results[k].empty]
693642
# Have to check, if at least one DataFrame is not empty.
694643
keys_to_use = keys_to_use if keys_to_use != [] else keys
695644
return (
696-
concat([result[k] for k in keys_to_use], keys=keys_to_use, axis=1),
645+
concat([results[k] for k in keys_to_use], keys=keys_to_use, axis=1),
697646
True,
698647
)
699648

@@ -702,7 +651,7 @@ def is_any_frame() -> bool:
702651
# we have a dict of Series
703652
# return a MI Series
704653
try:
705-
result = concat(result)
654+
result = concat(results)
706655
except TypeError as err:
707656
# we want to give a nice error here if
708657
# we have non-same sized objects, so
@@ -720,7 +669,7 @@ def is_any_frame() -> bool:
720669
from pandas import DataFrame, Series
721670

722671
try:
723-
result = DataFrame(result)
672+
result = DataFrame(results)
724673
except ValueError:
725674
# we have a dict of scalars
726675

@@ -731,7 +680,7 @@ def is_any_frame() -> bool:
731680
else:
732681
name = None
733682

734-
result = Series(result, name=name)
683+
result = Series(results, name=name)
735684

736685
return result, True
737686
elif is_list_like(arg):

0 commit comments

Comments
 (0)