maxframe 2.0.0b2__cp311-cp311-win_amd64.whl → 2.2.0__cp311-cp311-win_amd64.whl

maxframe/dataframe/groupby/expanding.py

@@ -0,0 +1,264 @@
+# 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.
+
+from typing import Any, Dict
+
+import pandas as pd
+
+from ... import opcodes
+from ...serialization.serializables import AnyField, DictField
+from ..window.expanding import Expanding
+from .core import BaseGroupByWindowOp
+
+_supported_funcs = {
+    "sum",
+    "prod",
+    "mean",
+    "std",
+    "var",
+    "median",
+    "min",
+    "max",
+    "count",
+}
+
+
+class GroupByExpandingAgg(BaseGroupByWindowOp):
+    _op_type_ = opcodes.EXPANDING_AGG
+
+    func = AnyField("func", default=None)
+    kwargs = DictField("kwargs", default_factory=dict)
+
+    def _calc_mock_result_df(self, mock_groupby):
+        expanding_args = self.window_params.copy()
+        # exclude MF-specific args
+        for key in Expanding._mf_specific_fields:
+            expanding_args.pop(key, None)
+
+        def apply_func(frame, **_):
+            func = self.func
+            if self.func == "prod":
+                func = lambda x: x.prod()
+            return frame.expanding(**expanding_args).agg(func)
+
+        is_cumcount = self.kwargs.get("cumcount", False)
+        # special handling for cumcount which returns a series
+        if is_cumcount:
+            return mock_groupby.cumcount(ascending=self.kwargs.get("ascending", True))
+        if self.func in ["sum", "max", "min", "prod"]:
+            return getattr(mock_groupby, f"cum{self.func}")()
+        return mock_groupby.apply(apply_func, include_groups=False)
+
+    def get_sort_cols_to_asc(self) -> Dict[Any, bool]:
+        res = super().get_sort_cols_to_asc()
+        if self.func == "count" and self.kwargs.get("ascending") is not None:
+            res = {k: not v for k, v in res.items()}
+        return res
+
+
+class ExpandingGroupby(Expanding):
+    def aggregate(self, func, **kwargs):
+        if func not in _supported_funcs:
+            raise NotImplementedError(f"func {func} is not supported")
+        op = GroupByExpandingAgg(
+            func=func,
+            groupby_params=self.input.op.groupby_params,
+            window_params=self.params,
+            kwargs=kwargs,
+        )
+        return op(self.input)
+
+    agg = aggregate
+
+
+def expanding(
+    groupby,
+    min_periods=1,
+    *,
+    shift=0,
+    reverse_range=False,
+    order_cols=None,
+    ascending=True,
+):
+    """
+    Return an expanding grouper, providing expanding
+    functionality per group.
+
+    Parameters
+    ----------
+    min_periods : int, default None
+        Minimum number of observations in window required to have a value;
+        otherwise, result is ``np.nan``.
+
+    shift : int, default 0
+        If specified, the window will be shifted by `shift` rows (or data will be
+        shifted by `-shift` rows) before computing window function.
+
+    reverse_range : bool, default False
+        If True, the window for current row is expanded from the last row to
+        the current instead of the first row.
+
+    Returns
+    -------
+    maxframe.dataframe.groupby.ExpandingGroupby
+
+    See Also
+    --------
+    Series.groupby : Apply a function groupby to a Series.
+    DataFrame.groupby : Apply a function groupby
+        to each row or column of a DataFrame.
+    """
+    return ExpandingGroupby(
+        input=groupby,
+        min_periods=min_periods,
+        shift=shift,
+        reverse_range=reverse_range,
+        order_cols=order_cols,
+        ascending=ascending,
+    )
+
+
+_cum_doc_template = """
+Cumulative %(func_name)s for each group.
+
+Returns
+-------
+Series or DataFrame
+
+See Also
+--------
+Series.groupby : Apply a function groupby to a Series.
+DataFrame.groupby : Apply a function groupby
+    to each row or column of a DataFrame.
+
+Examples
+--------
+For SeriesGroupBy:
+
+>>> import maxframe.dataframe as md
+>>> lst = ['a', 'a', 'b']
+>>> ser = md.Series([6, 2, 0], index=lst)
+>>> ser.execute()
+a    6
+a    2
+b    0
+dtype: int64
+>>> ser.groupby(level=0).cum%(func_name)s().execute()
+%(series_result)s
+
+For DataFrameGroupBy:
+
+>>> data = [[1, 8, 2], [1, 2, 5], [2, 6, 9]]
+>>> df = md.DataFrame(data, columns=["a", "b", "c"],
+...                   index=["fox", "gorilla", "lion"])
+>>> df.execute()
+          a   b   c
+fox       1   8   2
+gorilla   1   2   5
+lion      2   6   9
+>>> df.groupby("a").groups.execute()
+{1: ['fox', 'gorilla'], 2: ['lion']}
+>>> df.groupby("a").cum%(func_name)s().execute()
+%(df_result)s
+"""
+
+
+def _build_cum_docs(func_name):
+    ser = pd.Series([6, 2, 0], index=list("aab"))
+    data = [[1, 8, 2], [1, 2, 5], [2, 6, 9]]
+    df = pd.DataFrame(data, columns=["a", "b", "c"], index=["fox", "gorilla", "lion"])
+    cum_func_name = f"cum{func_name}"
+    return _cum_doc_template % dict(
+        func_name=func_name,
+        series_result=getattr(ser.groupby(level=0), cum_func_name)().to_string(),
+        df_result=getattr(df.groupby("a"), cum_func_name)().to_string(),
+    )
+
+
+def cumcount(groupby, ascending: bool = True):
+    """
+    Number each item in each group from 0 to the length of that group - 1.
+
+    Essentially this is equivalent to
+
+    .. code-block:: python
+
+        self.apply(lambda x: pd.Series(np.arange(len(x)), x.index))
+
+    Parameters
+    ----------
+    ascending : bool, default True
+        If False, number in reverse, from length of group - 1 to 0.
+
+    Returns
+    -------
+    Series
+        Sequence number of each element within each group.
+
+    See Also
+    --------
+    .ngroup : Number the groups themselves.
+
+    Examples
+    --------
+    >>> import maxframe.tensor as mt
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame([['a'], ['a'], ['a'], ['b'], ['b'], ['a']],
+    ...                   columns=['A'])
+    >>> df.execute()
+       A
+    0  a
+    1  a
+    2  a
+    3  b
+    4  b
+    5  a
+    >>> df.groupby('A').cumcount().execute()
+    0    0
+    1    1
+    2    2
+    3    0
+    4    1
+    5    3
+    dtype: int64
+    >>> df.groupby('A').cumcount(ascending=False).execute()
+    0    3
+    1    2
+    2    1
+    3    1
+    4    0
+    5    0
+    dtype: int64
+    """
+    return groupby.expanding().agg("count", ascending=ascending, cumcount=True)
+
+
+def cummin(groupby):
+    return groupby.expanding().min()
+
+
+def cummax(groupby):
+    return groupby.expanding().max()
+
+
+def cumprod(groupby):
+    return groupby.expanding().prod()
+
+
+def cumsum(groupby):
+    return groupby.expanding().sum()
+
+
+for _func in (cummin, cummax, cumprod, cumsum):
+    _func.__doc__ = _build_cum_docs(_func.__name__[3:])

maxframe/dataframe/groupby/fill.py

@@ -23,7 +23,7 @@ from ..utils import parse_index
 
 class GroupByFill(DataFrameOperator, DataFrameOperatorMixin):
     _op_module_ = "dataframe.groupby"
-    _legacy_name = "GroupByFillOperator"
+    _legacy_name = "GroupByFillOperator"  # since v2.0.0
 
     value = AnyField("value", default=None)
     method = StringField("method", default=None)

maxframe/dataframe/groupby/getitem.py

@@ -15,7 +15,7 @@
 from collections.abc import Iterable
 
 from ... import opcodes
-from ...core import OutputType
+from ...core import ENTITY_TYPE, OutputType
 from ...serialization.serializables import AnyField
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import parse_index
@@ -38,7 +38,8 @@ class GroupByIndex(DataFrameOperatorMixin, DataFrameOperator):
 
     def build_mock_groupby(self, **kwargs):
         groupby_op = self.inputs[0].op
-        return groupby_op.build_mock_groupby(**kwargs)[self.selection]
+        selection = kwargs.pop("selection", None) or self.selection
+        return groupby_op.build_mock_groupby(**kwargs)[selection]
 
     def __call__(self, groupby):
         indexed = groupby.op.build_mock_groupby()[self.selection]
@@ -55,8 +56,10 @@ class GroupByIndex(DataFrameOperatorMixin, DataFrameOperator):
         else:
             self.output_types = [OutputType.dataframe_groupby]
 
-            if isinstance(self.selection, Iterable) and not isinstance(
-                self.selection, str
+            if (
+                isinstance(self.selection, Iterable)
+                and not isinstance(self.selection, str)
+                and not isinstance(self.selection, ENTITY_TYPE)
             ):
                 item_list = list(self.selection)
             else:
@@ -80,7 +83,11 @@ def df_groupby_getitem(df_groupby, item):
 
     if hashable and item in df_groupby.dtypes:
         output_types = [OutputType.series_groupby]
-    elif isinstance(item, Iterable) and all(it in df_groupby.dtypes for it in item):
+    elif (
+        isinstance(item, Iterable)
+        and not isinstance(item, ENTITY_TYPE)
+        and all(it in df_groupby.dtypes for it in item)
+    ):
         output_types = [OutputType.dataframe_groupby]
     else:
         raise NameError(f"Cannot slice groupby with {item!r}")

maxframe/dataframe/groupby/head.py

@@ -17,8 +17,10 @@ import pandas as pd
 
 from ... import opcodes
 from ...core import OutputType
+from ...serialization import PickleContainer
 from ...serialization.serializables import BoolField, DictField, Int64Field
-from ...utils import pd_release_version
+from ...udf import BuiltinFunction
+from ...utils import find_objects, pd_release_version
 from ..core import IndexValue
 from ..operators import DataFrameOperator, DataFrameOperatorMixin
 from ..utils import parse_index
@@ -34,6 +36,14 @@ class GroupByHead(DataFrameOperator, DataFrameOperatorMixin):
     groupby_params = DictField("groupby_params", default=dict())
     enable_negative = BoolField("enable_negative", default=_pandas_enable_negative)
 
+    def has_custom_code(self) -> bool:
+        callable_bys = find_objects(
+            self.groupby_params.get("by"), types=PickleContainer, checker=callable
+        )
+        if not callable_bys:
+            return False
+        return any(not isinstance(fun, BuiltinFunction) for fun in callable_bys)
+
     def __call__(self, groupby):
         df = groupby
         while df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):

maxframe/dataframe/groupby/rank.py

@@ -0,0 +1,136 @@
+# 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.
+
+from typing import Any, Dict
+
+from ... import opcodes
+from ...serialization.serializables import BoolField, FieldTypes, ListField, StringField
+from ..utils import make_column_list
+from .core import BaseGroupByWindowOp
+
+
+class GroupByRank(BaseGroupByWindowOp):
+    _op_type_ = opcodes.RANK
+
+    method = StringField("method", default=None)
+    ascending = ListField("ascending", FieldTypes.bool, default=True)
+    na_option = StringField("na_option", default=None)
+    pct = BoolField("pct", default=False)
+
+    def _calc_mock_result_df(self, mock_groupby):
+        return mock_groupby.rank(
+            method=self.method,
+            ascending=self.ascending,
+            na_option=self.na_option,
+            pct=self.pct,
+        )
+
+    def get_sort_cols_to_asc(self) -> Dict[Any, bool]:
+        if self.inputs[0].ndim == 1:
+            return {self.inputs[0].name: self.ascending}
+        selections = make_column_list(
+            self.groupby_params.get("selection", []), self.inputs[0].dtypes
+        )
+        if not selections:
+            by_cols = set(
+                make_column_list(
+                    self.groupby_params.get("by") or [], self.inputs[0].dtypes
+                )
+            )
+            selections = [c for c in self.inputs[0].dtypes.index if c not in by_cols]
+
+        if len(self.ascending) < len(selections):
+            ascending = [self.ascending[0]] * len(selections)
+        else:
+            ascending = self.ascending
+        return dict(zip(selections, ascending))
+
+
+def rank(groupby, method="average", ascending=True, na_option="keep", pct=False):
+    """
+    Provide the rank of values within each group.
+
+    Parameters
+    ----------
+    method : {'average', 'min', 'max', 'first', 'dense'}, default 'average'
+        * average: average rank of group.
+        * min: lowest rank in group.
+        * max: highest rank in group.
+        * first: ranks assigned in order they appear in the array.
+        * dense: like 'min', but rank always increases by 1 between groups.
+    ascending : bool, default True
+        False for ranks by high (1) to low (N).
+    na_option : {'keep', 'top', 'bottom'}, default 'keep'
+        * keep: leave NA values where they are.
+        * top: smallest rank if ascending.
+        * bottom: smallest rank if descending.
+    pct : bool, default False
+        Compute percentage rank of data within each group.
+
+    Returns
+    -------
+    DataFrame with ranking of values within each group
+
+    See Also
+    --------
+    Series.groupby : Apply a function groupby to a Series.
+    DataFrame.groupby : Apply a function groupby
+        to each row or column of a DataFrame.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame(
+    ...     {
+    ...         "group": ["a", "a", "a", "a", "a", "b", "b", "b", "b", "b"],
+    ...         "value": [2, 4, 2, 3, 5, 1, 2, 4, 1, 5],
+    ...     }
+    ... )
+    >>> df.execute()
+      group  value
+    0     a      2
+    1     a      4
+    2     a      2
+    3     a      3
+    4     a      5
+    5     b      1
+    6     b      2
+    7     b      4
+    8     b      1
+    9     b      5
+    >>> for method in ['average', 'min', 'max', 'dense', 'first']:
+    ...     df[f'{method}_rank'] = df.groupby('group')['value'].rank(method)
+    >>> df.execute()
+      group  value  average_rank  min_rank  max_rank  dense_rank  first_rank
+    0     a      2           1.5       1.0       2.0         1.0         1.0
+    1     a      4           4.0       4.0       4.0         3.0         4.0
+    2     a      2           1.5       1.0       2.0         1.0         2.0
+    3     a      3           3.0       3.0       3.0         2.0         3.0
+    4     a      5           5.0       5.0       5.0         4.0         5.0
+    5     b      1           1.5       1.0       2.0         1.0         1.0
+    6     b      2           3.0       3.0       3.0         2.0         3.0
+    7     b      4           4.0       4.0       4.0         3.0         4.0
+    8     b      1           1.5       1.0       2.0         1.0         2.0
+    9     b      5           5.0       5.0       5.0         4.0         5.0
+    """
+    if not isinstance(ascending, list):
+        ascending = [ascending]
+    op = GroupByRank(
+        method=method,
+        ascending=ascending,
+        na_option=na_option,
+        pct=pct,
+        groupby_params=groupby.op.groupby_params,
+    )
+    return op(groupby)

maxframe/dataframe/groupby/rolling.py

@@ -0,0 +1,206 @@
+# 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.
+
+from ... import opcodes
+from ...serialization.serializables import AnyField, DictField
+from ..window.rolling import Rolling
+from .core import BaseGroupByWindowOp
+
+_supported_funcs = {"sum", "mean", "std", "var", "median", "min", "max", "count"}
+
+
+class GroupByRollingAgg(BaseGroupByWindowOp):
+    _op_type_ = opcodes.ROLLING_AGG
+
+    func = AnyField("func", default=None)
+    kwargs = DictField("kwargs", default_factory=dict)
+
+    def _calc_mock_result_df(self, mock_groupby):
+        rolling_args = self.window_params.copy()
+        # exclude MF-specific args
+        for key in Rolling._mf_specific_fields:
+            rolling_args.pop(key, None)
+
+        def apply_func(frame, **_):
+            return getattr(frame.rolling(**rolling_args), self.func)()
+
+        return mock_groupby.apply(apply_func, include_groups=False)
+
+
+class RollingGroupby(Rolling):
+    def aggregate(self, func, **kwargs):
+        if func not in _supported_funcs:
+            raise NotImplementedError(f"func {func} is not supported")
+        op = GroupByRollingAgg(
+            func=func,
+            groupby_params=self.input.op.groupby_params,
+            window_params=self.params,
+            kwargs=kwargs,
+        )
+        return op(self.input)
+
+
+def rolling(
+    groupby,
+    window,
+    min_periods=None,
+    *,
+    center=False,
+    win_type=None,
+    on=None,
+    axis=0,
+    closed=None,
+    shift=0,
+    order_cols=None,
+    ascending=True,
+) -> RollingGroupby:
+    """
+    Return a rolling grouper, providing rolling functionality per group.
+
+    Parameters
+    ----------
+    window : int, timedelta, str, offset, or BaseIndexer subclass
+        Size of the moving window.
+
+        If an integer, the fixed number of observations used for
+        each window.
+
+        If a timedelta, str, or offset, the time period of each window. Each
+        window will be a variable sized based on the observations included in
+        the time-period. This is only valid for datetimelike indexes.
+        To learn more about the offsets & frequency strings, please see `this link
+        <https://pandas.pydata.org/pandas-docs/stable/user_guide/timeseries.html#offset-aliases>`__.
+
+        If a BaseIndexer subclass, the window boundaries
+        based on the defined ``get_window_bounds`` method. Additional rolling
+        keyword arguments, namely ``min_periods``, ``center``, ``closed`` and
+        ``step`` will be passed to ``get_window_bounds``.
+
+    min_periods : int, default None
+        Minimum number of observations in window required to have a value;
+        otherwise, result is ``np.nan``.
+
+        For a window that is specified by an offset,
+        ``min_periods`` will default to 1.
+
+        For a window that is specified by an integer, ``min_periods`` will default
+        to the size of the window.
+
+    center : bool, default False
+        If False, set the window labels as the right edge of the window index.
+
+        If True, set the window labels as the center of the window index.
+
+    win_type : str, default None
+        If ``None``, all points are evenly weighted.
+
+        If a string, it must be a valid `scipy.signal window function
+        <https://docs.scipy.org/doc/scipy/reference/signal.windows.html#module-scipy.signal.windows>`__.
+
+        Certain Scipy window types require additional parameters to be passed
+        in the aggregation function. The additional parameters must match
+        the keywords specified in the Scipy window type method signature.
+
+    on : str, optional
+        For a DataFrame, a column label or Index level on which
+        to calculate the rolling window, rather than the DataFrame's index.
+
+        Provided integer column is ignored and excluded from result since
+        an integer index is not used to calculate the rolling window.
+
+    axis : int or str, default 0
+        If ``0`` or ``'index'``, roll across the rows.
+
+        If ``1`` or ``'columns'``, roll across the columns.
+
+        For `Series` this parameter is unused and defaults to 0.
+
+    closed : str, default None
+        If ``'right'``, the first point in the window is excluded from calculations.
+
+        If ``'left'``, the last point in the window is excluded from calculations.
+
+        If ``'both'``, no points in the window are excluded from calculations.
+
+        If ``'neither'``, the first and last points in the window are excluded
+        from calculations.
+
+        Default ``None`` (``'right'``).
+
+    shift : int, default 0
+        If specified, the window will be shifted by `shift` rows (or data will be
+        shifted by `-shift` rows) before computing window function.
+
+    Returns
+    -------
+    maxframe.dataframe.groupby.RollingGroupby
+        Return a new grouper with our rolling appended.
+
+    See Also
+    --------
+    Series.rolling : Calling object with Series data.
+    DataFrame.rolling : Calling object with DataFrames.
+    Series.groupby : Apply a function groupby to a Series.
+    DataFrame.groupby : Apply a function groupby.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame({'A': [1, 1, 2, 2],
+    ...                    'B': [1, 2, 3, 4],
+    ...                    'C': [0.362, 0.227, 1.267, -0.562]})
+    >>> df.execute()
+          A  B      C
+    0     1  1  0.362
+    1     1  2  0.227
+    2     2  3  1.267
+    3     2  4 -0.562
+
+    >>> df.groupby('A').rolling(2).sum().execute()
+        B      C
+    A
+    1 0  NaN    NaN
+      1  3.0  0.589
+    2 2  NaN    NaN
+      3  7.0  0.705
+
+    >>> df.groupby('A').rolling(2, min_periods=1).sum().execute()
+        B      C
+    A
+    1 0  1.0  0.362
+      1  3.0  0.589
+    2 2  3.0  1.267
+      3  7.0  0.705
+
+    >>> df.groupby('A').rolling(2, on='B').sum().execute()
+        B      C
+    A
+    1 0  1    NaN
+      1  2  0.589
+    2 2  3    NaN
+      3  4  0.705
+    """
+    return RollingGroupby(
+        input=groupby,
+        window=window,
+        min_periods=min_periods,
+        center=center,
+        win_type=win_type,
+        on=on,
+        axis=axis,
+        closed=closed,
+        shift=shift,
+        order_cols=order_cols,
+        ascending=ascending,
+    )