core/maxframe/dataframe/groupby/sample.py

# 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 Optional, Sequence, Union

import numpy as np
import pandas as pd

from ... import opcodes
from ...core import ENTITY_TYPE, 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

    def _set_inputs(self, inputs):
        super()._set_inputs(inputs)
        input_iter = iter(inputs)
        next(input_iter)
        if isinstance(self.weights, ENTITY_TYPE):
            self.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)