maxframe 0.1.0b5__cp39-cp39-macosx_11_0_arm64.whl

maxframe/dataframe/groupby/aggregation.py

@@ -0,0 +1,350 @@
+# Copyright 1999-2024 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(),
+}
+_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.
+    """
+
+    # 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)

maxframe/dataframe/groupby/apply.py

@@ -0,0 +1,251 @@
+# Copyright 1999-2024 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 numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import OutputType
+from ...core.operator import OperatorLogicKeyGeneratorMixin
+from ...serialization.serializables import (
+    AnyField,
+    BoolField,
+    DictField,
+    FunctionField,
+    StringField,
+    TupleField,
+)
+from ...utils import get_func_token, quiet_stdio, tokenize
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import make_dtype, make_dtypes, parse_index, validate_output_types
+
+
+class GroupByApplyLogicKeyGeneratorMixin(OperatorLogicKeyGeneratorMixin):
+    def _get_logic_key_token_values(self):
+        token_values = super()._get_logic_key_token_values()
+        if self.func:
+            return token_values + [get_func_token(self.func)]
+        else:  # pragma: no cover
+            return token_values
+
+
+class GroupByApply(
+    DataFrameOperator, DataFrameOperatorMixin, GroupByApplyLogicKeyGeneratorMixin
+):
+    _op_type_ = opcodes.APPLY
+    _op_module_ = "dataframe.groupby"
+
+    func = FunctionField("func")
+    args = TupleField("args", default_factory=tuple)
+    kwds = DictField("kwds", default_factory=dict)
+    maybe_agg = BoolField("maybe_agg", default=None)
+    logic_key = StringField("logic_key", default=None)
+    func_key = AnyField("func_key", default=None)
+    need_clean_up_func = BoolField("need_clean_up_func", default=False)
+
+    def __init__(self, output_types=None, **kw):
+        super().__init__(_output_types=output_types, **kw)
+
+    def _update_key(self):
+        values = [v for v in self._values_ if v is not self.func] + [
+            get_func_token(self.func)
+        ]
+        self._obj_set("_key", tokenize(type(self).__name__, *values))
+        return self
+
+    def _infer_df_func_returns(
+        self, in_groupby, in_df, dtypes=None, dtype=None, name=None, index=None
+    ):
+        index_value, output_type, new_dtypes = None, None, None
+
+        if self.output_types is not None and (dtypes is not None or dtype is not None):
+            ret_dtypes = dtypes if dtypes is not None else (dtype, name)
+            ret_index_value = parse_index(index) if index is not None else None
+            return ret_dtypes, ret_index_value
+
+        try:
+            infer_df = in_groupby.op.build_mock_groupby().apply(
+                self.func, *self.args, **self.kwds
+            )
+
+            if len(infer_df) <= 2:
+                # we create mock df with 4 rows, 2 groups
+                # if return df has 2 rows, we assume that
+                # it's an aggregation operation
+                self.maybe_agg = True
+
+            # todo return proper index when sort=True is implemented
+            index_value = parse_index(infer_df.index[:0], in_df.key, self.func)
+
+            # for backward compatibility
+            dtype = dtype if dtype is not None else dtypes
+            if isinstance(infer_df, pd.DataFrame):
+                output_type = output_type or OutputType.dataframe
+                new_dtypes = new_dtypes or infer_df.dtypes
+            elif isinstance(infer_df, pd.Series):
+                output_type = output_type or OutputType.series
+                new_dtypes = new_dtypes or (
+                    name or infer_df.name,
+                    dtype or infer_df.dtype,
+                )
+            else:
+                output_type = OutputType.series
+                new_dtypes = (name, dtype or pd.Series(infer_df).dtype)
+        except:  # noqa: E722  # nosec
+            pass
+
+        self.output_types = (
+            [output_type]
+            if not self.output_types and output_type
+            else self.output_types
+        )
+        dtypes = new_dtypes if dtypes is None else dtypes
+        index_value = index_value if index is None else parse_index(index)
+        return dtypes, index_value
+
+    def __call__(self, groupby, dtypes=None, dtype=None, name=None, index=None):
+        in_df = groupby
+        if self.output_types and self.output_types[0] == OutputType.df_or_series:
+            return self.new_df_or_series([groupby])
+        while in_df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
+            in_df = in_df.inputs[0]
+
+        with quiet_stdio():
+            dtypes, index_value = self._infer_df_func_returns(
+                groupby, in_df, dtypes, dtype=dtype, name=name, index=index
+            )
+        if index_value is None:
+            index_value = parse_index(None, (in_df.key, in_df.index_value.key))
+        for arg, desc in zip((self.output_types, dtypes), ("output_types", "dtypes")):
+            if arg is None:
+                raise TypeError(
+                    f"Cannot determine {desc} by calculating with enumerate data, "
+                    "please specify it as arguments"
+                )
+
+        if self.output_types[0] == OutputType.dataframe:
+            new_shape = (np.nan, len(dtypes))
+            return self.new_dataframe(
+                [groupby],
+                shape=new_shape,
+                dtypes=dtypes,
+                index_value=index_value,
+                columns_value=parse_index(dtypes.index, store_data=True),
+            )
+        else:
+            name = name or dtypes[0]
+            dtype = dtype or dtypes[1]
+            new_shape = (np.nan,)
+            return self.new_series(
+                [groupby],
+                name=name,
+                shape=new_shape,
+                dtype=dtype,
+                index_value=index_value,
+            )
+
+
+def groupby_apply(
+    groupby,
+    func,
+    *args,
+    output_type=None,
+    dtypes=None,
+    dtype=None,
+    name=None,
+    index=None,
+    skip_infer=None,
+    **kwargs,
+):
+    """
+    Apply function `func` group-wise and combine the results together.
+
+    The function passed to `apply` must take a dataframe as its first
+    argument and return a DataFrame, Series or scalar. `apply` will
+    then take care of combining the results back together into a single
+    dataframe or series. `apply` is therefore a highly flexible
+    grouping method.
+
+    While `apply` is a very flexible method, its downside is that
+    using it can be quite a bit slower than using more specific methods
+    like `agg` or `transform`. Pandas offers a wide range of method that will
+    be much faster than using `apply` for their specific purposes, so try to
+    use them before reaching for `apply`.
+
+    Parameters
+    ----------
+    func : callable
+        A callable that takes a dataframe as its first argument, and
+        returns a dataframe, a series or a scalar. In addition the
+        callable may take positional and keyword arguments.
+
+    output_type : {'dataframe', 'series'}, default None
+        Specify type of returned object. See `Notes` for more details.
+
+    dtypes : Series, default None
+        Specify dtypes of returned DataFrames. See `Notes` for more details.
+
+    dtype : numpy.dtype, default None
+        Specify dtype of returned Series. See `Notes` for more details.
+
+    name : str, default None
+        Specify name of returned Series. See `Notes` for more details.
+
+    index : Index, default None
+        Specify index of returned object. See `Notes` for more details.
+
+    skip_infer: bool, default False
+        Whether infer dtypes when dtypes or output_type is not specified.
+
+    args, kwargs : tuple and dict
+        Optional positional and keyword arguments to pass to `func`.
+
+    Returns
+    -------
+    applied : Series or DataFrame
+
+    See Also
+    --------
+    pipe : Apply function to the full GroupBy object instead of to each
+        group.
+    aggregate : Apply aggregate function to the GroupBy object.
+    transform : Apply function column-by-column to the GroupBy object.
+    Series.apply : Apply a function to a Series.
+    DataFrame.apply : Apply a function to each row or column of a DataFrame.
+
+    Notes
+    -----
+    When deciding output dtypes and shape of the return value, MaxFrame will
+    try applying ``func`` onto a mock grouped object, and the apply call
+    may fail. When this happens, you need to specify the type of apply
+    call (DataFrame or Series) in output_type.
+
+    * For DataFrame output, you need to specify a list or a pandas Series
+      as ``dtypes`` of output DataFrame. ``index`` of output can also be
+      specified.
+    * For Series output, you need to specify ``dtype`` and ``name`` of
+      output Series.
+    """
+    output_types = kwargs.pop("output_types", None)
+    object_type = kwargs.pop("object_type", None)
+    output_types = validate_output_types(
+        output_types=output_types, output_type=output_type, object_type=object_type
+    )
+    if output_types is None and skip_infer:
+        output_types = [OutputType.df_or_series]
+
+    dtypes = make_dtypes(dtypes)
+    dtype = make_dtype(dtype)
+    op = GroupByApply(func=func, args=args, kwds=kwargs, output_types=output_types)
+    return op(groupby, dtypes=dtypes, dtype=dtype, name=name, index=index)