# Copyright 1999-2025 Alibaba Group Holding Ltd. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. import functools import logging from typing import Callable, Dict import numpy as np import pandas as pd from ... import opcodes from ...core import ENTITY_TYPE, OutputType from ...serialization.serializables import ( AnyField, DictField, Int32Field, Int64Field, ListField, StringField, ) from ...utils import lazy_import, pd_release_version from ..core import GROUPBY_TYPE from ..operators import DataFrameOperator, DataFrameOperatorMixin from ..reduction.aggregation import ( compile_reduction_funcs, is_funcs_aggregate, normalize_reduction_funcs, ) from ..utils import is_cudf, parse_index cp = lazy_import("cupy", rename="cp") cudf = lazy_import("cudf") logger = logging.getLogger(__name__) CV_THRESHOLD = 0.2 MEAN_RATIO_THRESHOLD = 2 / 3 _support_get_group_without_as_index = pd_release_version[:2] > (1, 0) class SizeRecorder: def __init__(self): self._raw_records = [] self._agg_records = [] def record(self, raw_record: int, agg_record: int): self._raw_records.append(raw_record) self._agg_records.append(agg_record) def get(self): return self._raw_records, self._agg_records _agg_functions = { "sum": lambda x: x.sum(), "prod": lambda x: x.prod(), "product": lambda x: x.product(), "min": lambda x: x.min(), "max": lambda x: x.max(), "all": lambda x: x.all(), "any": lambda x: x.any(), "count": lambda x: x.count(), "size": lambda x: x._reduction_size(), "mean": lambda x: x.mean(), "var": lambda x, ddof=1: x.var(ddof=ddof), "std": lambda x, ddof=1: x.std(ddof=ddof), "sem": lambda x, ddof=1: x.sem(ddof=ddof), "skew": lambda x, bias=False: x.skew(bias=bias), "kurt": lambda x, bias=False: x.kurt(bias=bias), "kurtosis": lambda x, bias=False: x.kurtosis(bias=bias), "nunique": lambda x: x.nunique(), "median": lambda x: x.median(), } _series_col_name = "col_name" def _patch_groupby_kurt(): try: from pandas.core.groupby import DataFrameGroupBy, SeriesGroupBy if not hasattr(DataFrameGroupBy, "kurt"): # pragma: no branch def _kurt_by_frame(a, *args, **kwargs): data = a.to_frame().kurt(*args, **kwargs).iloc[0] if is_cudf(data): # pragma: no cover data = data.copy() return data def _group_kurt(x, *args, **kwargs): if kwargs.get("numeric_only") is not None: return x.agg(functools.partial(_kurt_by_frame, *args, **kwargs)) else: return x.agg(functools.partial(pd.Series.kurt, *args, **kwargs)) DataFrameGroupBy.kurt = DataFrameGroupBy.kurtosis = _group_kurt SeriesGroupBy.kurt = SeriesGroupBy.kurtosis = _group_kurt except (AttributeError, ImportError): # pragma: no cover pass _patch_groupby_kurt() del _patch_groupby_kurt def build_mock_agg_result( groupby: GROUPBY_TYPE, groupby_params: Dict, raw_func: Callable, **raw_func_kw, ): try: agg_result = groupby.op.build_mock_groupby().aggregate(raw_func, **raw_func_kw) except ValueError: if ( groupby_params.get("as_index") or _support_get_group_without_as_index ): # pragma: no cover raise agg_result = ( groupby.op.build_mock_groupby(as_index=True) .aggregate(raw_func, **raw_func_kw) .to_frame() ) agg_result.index.names = [None] * agg_result.index.nlevels return agg_result class DataFrameGroupByAgg(DataFrameOperator, DataFrameOperatorMixin): _op_type_ = opcodes.GROUPBY_AGG raw_func = AnyField("raw_func") raw_func_kw = DictField("raw_func_kw") func = AnyField("func") func_rename = ListField("func_rename", default=None) raw_groupby_params = DictField("raw_groupby_params") groupby_params = DictField("groupby_params") method = StringField("method") # for chunk chunk_store_limit = Int64Field("chunk_store_limit") pre_funcs = ListField("pre_funcs") agg_funcs = ListField("agg_funcs") post_funcs = ListField("post_funcs") index_levels = Int32Field("index_levels") size_recorder_name = StringField("size_recorder_name") def _set_inputs(self, inputs): super()._set_inputs(inputs) inputs_iter = iter(self._inputs[1:]) if len(self._inputs) > 1: by = [] for v in self.groupby_params["by"]: if isinstance(v, ENTITY_TYPE): by.append(next(inputs_iter)) else: by.append(v) self.groupby_params["by"] = by def _get_inputs(self, inputs): if isinstance(self.groupby_params["by"], list): for v in self.groupby_params["by"]: if isinstance(v, ENTITY_TYPE): inputs.append(v) return inputs def _get_index_levels(self, groupby, mock_index): if not self.groupby_params["as_index"]: try: as_index_agg_df = groupby.op.build_mock_groupby( as_index=True ).aggregate(self.raw_func, **self.raw_func_kw) except: # noqa: E722 # nosec # pylint: disable=bare-except # handling cases like mdf.groupby("b", as_index=False).b.agg({"c": "count"}) if isinstance(self.groupby_params["by"], list): return len(self.groupby_params["by"]) raise # pragma: no cover pd_index = as_index_agg_df.index else: pd_index = mock_index return 1 if not isinstance(pd_index, pd.MultiIndex) else len(pd_index.levels) def _fix_as_index(self, result_index: pd.Index): # make sure if as_index=False takes effect if isinstance(result_index, pd.MultiIndex): # if MultiIndex, as_index=False definitely takes no effect self.groupby_params["as_index"] = True elif result_index.name is not None: # if not MultiIndex and agg_df.index has a name # means as_index=False takes no effect self.groupby_params["as_index"] = True def _call_dataframe(self, groupby, input_df): compile_reduction_funcs(self, input_df) agg_df = build_mock_agg_result( groupby, self.groupby_params, self.raw_func, **self.raw_func_kw ) shape = (np.nan, agg_df.shape[1]) if isinstance(agg_df.index, pd.RangeIndex): index_value = parse_index( pd.RangeIndex(-1), groupby.key, groupby.index_value.key ) else: index_value = parse_index( agg_df.index, groupby.key, groupby.index_value.key ) # make sure if as_index=False takes effect self._fix_as_index(agg_df.index) # determine num of indices to group in intermediate steps self.index_levels = self._get_index_levels(groupby, agg_df.index) inputs = self._get_inputs([input_df]) return self.new_dataframe( inputs, shape=shape, dtypes=agg_df.dtypes, index_value=index_value, columns_value=parse_index(agg_df.columns, store_data=True), ) def _call_series(self, groupby, in_series): compile_reduction_funcs(self, in_series) agg_result = build_mock_agg_result( groupby, self.groupby_params, self.raw_func, **self.raw_func_kw ) # make sure if as_index=False takes effect self._fix_as_index(agg_result.index) index_value = parse_index( agg_result.index, groupby.key, groupby.index_value.key ) inputs = self._get_inputs([in_series]) # determine num of indices to group in intermediate steps self.index_levels = self._get_index_levels(groupby, agg_result.index) # update value type if isinstance(agg_result, pd.DataFrame): return self.new_dataframe( inputs, shape=(np.nan, len(agg_result.columns)), dtypes=agg_result.dtypes, index_value=index_value, columns_value=parse_index(agg_result.columns, store_data=True), ) else: return self.new_series( inputs, shape=(np.nan,), dtype=agg_result.dtype, name=agg_result.name, index_value=index_value, ) def __call__(self, groupby): normalize_reduction_funcs(self, ndim=groupby.ndim) df = groupby while df.op.output_types[0] not in (OutputType.dataframe, OutputType.series): df = df.inputs[0] if self.raw_func == "size": self.output_types = [OutputType.series] else: self.output_types = ( [OutputType.dataframe] if groupby.op.output_types[0] == OutputType.dataframe_groupby else [OutputType.series] ) if self.output_types[0] == OutputType.dataframe: return self._call_dataframe(groupby, df) else: return self._call_series(groupby, df) def agg(groupby, func=None, method="auto", *args, **kwargs): """ Aggregate using one or more operations on grouped data. Parameters ---------- groupby : MaxFrame Groupby Groupby data. func : str or list-like Aggregation functions. method : {'auto', 'shuffle', 'tree'}, default 'auto' 'tree' method provide a better performance, 'shuffle' is recommended if aggregated result is very large, 'auto' will use 'shuffle' method in distributed mode and use 'tree' in local mode. Returns ------- Series or DataFrame Aggregated result. Examples -------- >>> import maxframe.dataframe as md >>> df = md.DataFrame( ... { ... "A": [1, 1, 2, 2], ... "B": [1, 2, 3, 4], ... "C": [0.362838, 0.227877, 1.267767, -0.562860], ... } ... ).execute() A B C 0 1 1 0.362838 1 1 2 0.227877 2 2 3 1.267767 3 2 4 -0.562860 The aggregation is for each column. >>> df.groupby('A').agg('min').execute() B C A 1 1 0.227877 2 3 -0.562860 Multiple aggregations. >>> df.groupby('A').agg(['min', 'max']).execute() B C min max min max A 1 1 2 0.227877 0.362838 2 3 4 -0.562860 1.267767 Different aggregations per column >>> df.groupby('A').agg({'B': ['min', 'max'], 'C': 'sum'}).execute() B C min max sum A 1 1 2 0.590715 2 3 4 0.704907 To control the output names with different aggregations per column, pandas supports “named aggregation” >>> from maxframe.dataframe.groupby import NamedAgg >>> df.groupby("A").agg( ... b_min=NamedAgg(column="B", aggfunc="min"), ... c_sum=NamedAgg(column="C", aggfunc="sum")).execute() b_min c_sum A 1 1 0.590715 2 3 0.704907 """ # When perform a computation on the grouped data, we won't shuffle # the data in the stage of groupby and do shuffle after aggregation. if not isinstance(groupby, GROUPBY_TYPE): raise TypeError(f"Input should be type of groupby, not {type(groupby)}") if method is None: method = "auto" if method not in ["shuffle", "tree", "auto"]: raise ValueError( f"Method {method} is not available, please specify 'tree' or 'shuffle" ) if not is_funcs_aggregate(func, ndim=groupby.ndim): # pass index to transform, otherwise it will lose name info for index agg_result = build_mock_agg_result( groupby, groupby.op.groupby_params, func, **kwargs ) if isinstance(agg_result.index, pd.RangeIndex): # set -1 to represent unknown size for RangeIndex index_value = parse_index( pd.RangeIndex(-1), groupby.key, groupby.index_value.key ) else: index_value = parse_index( agg_result.index, groupby.key, groupby.index_value.key ) return groupby.transform( func, *args, _call_agg=True, index=index_value, **kwargs ) agg_op = DataFrameGroupByAgg( raw_func=func, raw_func_kw=kwargs, method=method, raw_groupby_params=groupby.op.groupby_params, groupby_params=groupby.op.groupby_params, ) return agg_op(groupby)