maxframe 2.2.0__cp312-cp312-macosx_10_9_universal2.whl

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,
+    )

maxframe/dataframe/groupby/sample.py

@@ -0,0 +1,214 @@
+# 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 copy
+from typing import List, Optional, Sequence, Union
+
+import numpy as np
+import pandas as pd
+
+from ... import opcodes
+from ...core import ENTITY_TYPE, EntityData, OutputType, get_output_types
+from ...serialization.serializables import (
+    BoolField,
+    DictField,
+    Float32Field,
+    Int32Field,
+    Int64Field,
+    KeyField,
+    NDArrayField,
+    StringField,
+)
+from ...tensor.random import RandomStateField
+from ..initializer import Series as asseries
+from ..operators import DataFrameOperator, DataFrameOperatorMixin
+from ..utils import parse_index
+
+
+class GroupBySample(DataFrameOperator, DataFrameOperatorMixin):
+    _op_type_ = opcodes.RAND_SAMPLE
+    _op_module_ = "dataframe.groupby"
+
+    groupby_params = DictField("groupby_params", default=None)
+    size = Int64Field("size", default=None)
+    frac = Float32Field("frac", default=None)
+    replace = BoolField("replace", default=None)
+    weights = KeyField("weights", default=None)
+    seed = Int32Field("seed", default=None)
+    _random_state = RandomStateField("random_state", default=None)
+    errors = StringField("errors", default=None)
+    # for chunks
+    # num of instances for chunks
+    input_nsplits = NDArrayField("input_nsplits", default=None)
+
+    def __init__(self, random_state=None, **kw):
+        super().__init__(_random_state=random_state, **kw)
+
+    @property
+    def random_state(self):
+        return self._random_state
+
+    @classmethod
+    def _set_inputs(cls, op: "GroupBySample", inputs: List[EntityData]):
+        super()._set_inputs(op, inputs)
+        input_iter = iter(inputs)
+        next(input_iter)
+        if isinstance(op.weights, ENTITY_TYPE):
+            op.weights = next(input_iter)
+
+    def __call__(self, groupby):
+        df = groupby
+        while df.op.output_types[0] not in (OutputType.dataframe, OutputType.series):
+            df = df.inputs[0]
+
+        selection = groupby.op.groupby_params.pop("selection", None)
+        if df.ndim > 1 and selection:
+            if isinstance(selection, tuple) and selection not in df.dtypes:
+                selection = list(selection)
+            result_df = df[selection]
+        else:
+            result_df = df
+
+        params = result_df.params
+        params["shape"] = (
+            (np.nan,) if result_df.ndim == 1 else (np.nan, result_df.shape[-1])
+        )
+        params["index_value"] = parse_index(result_df.index_value.to_pandas()[:0])
+
+        input_dfs = [df]
+        if isinstance(self.weights, ENTITY_TYPE):
+            input_dfs.append(self.weights)
+
+        self._output_types = get_output_types(result_df)
+        return self.new_tileable(input_dfs, **params)
+
+
+def groupby_sample(
+    groupby,
+    n: Optional[int] = None,
+    frac: Optional[float] = None,
+    replace: bool = False,
+    weights: Union[Sequence, pd.Series, None] = None,
+    random_state: Optional[np.random.RandomState] = None,
+    errors: str = "ignore",
+):
+    """
+    Return a random sample of items from each group.
+
+    You can use `random_state` for reproducibility.
+
+    Parameters
+    ----------
+    n : int, optional
+        Number of items to return for each group. Cannot be used with
+        `frac` and must be no larger than the smallest group unless
+        `replace` is True. Default is one if `frac` is None.
+    frac : float, optional
+        Fraction of items to return. Cannot be used with `n`.
+    replace : bool, default False
+        Allow or disallow sampling of the same row more than once.
+    weights : list-like, optional
+        Default None results in equal probability weighting.
+        If passed a list-like then values must have the same length as
+        the underlying DataFrame or Series object and will be used as
+        sampling probabilities after normalization within each group.
+        Values must be non-negative with at least one positive element
+        within each group.
+    random_state : int, array-like, BitGenerator, np.random.RandomState, optional
+        If int, array-like, or BitGenerator (NumPy>=1.17), seed for
+        random number generator
+        If np.random.RandomState, use as numpy RandomState object.
+    errors : {'ignore', 'raise'}, default 'ignore'
+        If ignore, errors will not be raised when `replace` is False
+        and size of some group is less than `n`.
+
+    Returns
+    -------
+    Series or DataFrame
+        A new object of same type as caller containing items randomly
+        sampled within each group from the caller object.
+
+    See Also
+    --------
+    DataFrame.sample: Generate random samples from a DataFrame object.
+    numpy.random.choice: Generate a random sample from a given 1-D numpy
+        array.
+
+    Examples
+    --------
+    >>> import maxframe.dataframe as md
+    >>> df = md.DataFrame(
+    ...     {"a": ["red"] * 2 + ["blue"] * 2 + ["black"] * 2, "b": range(6)}
+    ... )
+    >>> df.execute()
+           a  b
+    0    red  0
+    1    red  1
+    2   blue  2
+    3   blue  3
+    4  black  4
+    5  black  5
+
+    Select one row at random for each distinct value in column a. The
+    `random_state` argument can be used to guarantee reproducibility:
+
+    >>> df.groupby("a").sample(n=1, random_state=1).execute()
+           a  b
+    4  black  4
+    2   blue  2
+    1    red  1
+
+    Set `frac` to sample fixed proportions rather than counts:
+
+    >>> df.groupby("a")["b"].sample(frac=0.5, random_state=2).execute()
+    5    5
+    2    2
+    0    0
+    Name: b, dtype: int64
+
+    Control sample probabilities within groups by setting weights:
+
+    >>> df.groupby("a").sample(
+    ...     n=1,
+    ...     weights=[1, 1, 1, 0, 0, 1],
+    ...     random_state=1,
+    ... ).execute()
+           a  b
+    5  black  5
+    2   blue  2
+    0    red  0
+    """
+    groupby_params = groupby.op.groupby_params.copy()
+    groupby_params.pop("as_index", None)
+
+    if weights is not None and not isinstance(weights, ENTITY_TYPE):
+        weights = asseries(weights)
+
+    n = 1 if n is None and frac is None else n
+    rs = copy.deepcopy(
+        random_state.to_numpy() if hasattr(random_state, "to_numpy") else random_state
+    )
+    if not isinstance(rs, np.random.RandomState):  # pragma: no cover
+        rs = np.random.RandomState(rs)
+
+    op = GroupBySample(
+        size=n,
+        frac=frac,
+        replace=replace,
+        weights=weights,
+        random_state=rs,
+        groupby_params=groupby_params,
+        errors=errors,
+    )
+    return op(groupby)

maxframe/dataframe/groupby/shift.py

@@ -0,0 +1,114 @@
+# 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, Int64Field
+from .core import BaseGroupByWindowOp
+
+
+class GroupByShift(BaseGroupByWindowOp):
+    _op_type_ = opcodes.SHIFT
+
+    periods = Int64Field("periods", default=None)
+    freq = AnyField("freq", default=None)
+    fill_value = AnyField("fill_value", default=None)
+
+    def _calc_mock_result_df(self, mock_groupby):
+        return mock_groupby.shift(
+            self.periods, freq=self.freq, fill_value=self.fill_value
+        )
+
+
+def shift(
+    groupby, periods=1, freq=None, fill_value=None, order_cols=None, ascending=True
+):
+    """
+    Shift each group by periods observations.
+
+    If freq is passed, the index will be increased using the periods and the freq.
+
+    Parameters
+    ----------
+    periods : int | Sequence[int], default 1
+        Number of periods to shift. If a list of values, shift each group by
+        each period.
+    freq : str, optional
+        Frequency string.
+
+    fill_value : optional
+        The scalar value to use for newly introduced missing values.
+
+    Returns
+    -------
+    Series or DataFrame
+        Object shifted within each group.
+
+    See Also
+    --------
+    Index.shift : Shift values of Index.
+
+    Examples
+    --------
+
+    For SeriesGroupBy:
+
+    >>> import maxframe.dataframe as md
+    >>> lst = ['a', 'a', 'b', 'b']
+    >>> ser = md.Series([1, 2, 3, 4], index=lst)
+    >>> ser.execute()
+    a    1
+    a    2
+    b    3
+    b    4
+    dtype: int64
+    >>> ser.groupby(level=0).shift(1).execute()
+    a    NaN
+    a    1.0
+    b    NaN
+    b    3.0
+    dtype: float64
+
+    For DataFrameGroupBy:
+
+    >>> data = [[1, 2, 3], [1, 5, 6], [2, 5, 8], [2, 6, 9]]
+    >>> df = md.DataFrame(data, columns=["a", "b", "c"],
+    ...                   index=["tuna", "salmon", "catfish", "goldfish"])
+    >>> df.execute()
+               a  b  c
+        tuna   1  2  3
+      salmon   1  5  6
+     catfish   2  5  8
+    goldfish   2  6  9
+    >>> df.groupby("a").shift(1).execute()
+                  b    c
+        tuna    NaN  NaN
+      salmon    2.0  3.0
+     catfish    NaN  NaN
+    goldfish    5.0  8.0
+    """
+    if not isinstance(ascending, list):
+        ascending = [ascending]
+
+    window_params = dict(
+        order_cols=order_cols,
+        ascending=ascending,
+    )
+    op = GroupByShift(
+        periods=periods,
+        freq=freq,
+        fill_value=fill_value,
+        groupby_params=groupby.op.groupby_params,
+        window_params=window_params,
+    )
+    return op(groupby)

maxframe/dataframe/groupby/tests/__init__.py

@@ -0,0 +1,13 @@
+# 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.