text: # Copyright (c) 2020-2021, NVIDIA CORPORATION.
  
  import collections
  import pickle
  import warnings
  
  import pandas as pd
  from nvtx import annotate
  
  import cudf
  from cudf._lib import groupby as libgroupby
  from cudf._typing import DataFrameOrSeries
  from cudf.api.types import is_list_like
  from cudf.core.abc import Serializable
  from cudf.core.column.column import arange
  from cudf.utils.utils import GetAttrGetItemMixin, cached_property
  
  
  # The three functions below return the quantiles [25%, 50%, 75%]
  # respectively, which are called in the describe() method to ouput
  # the summary stats of a GroupBy object
  def _quantile_25(x):
      return x.quantile(0.25)
  
  
  def _quantile_50(x):
      return x.quantile(0.50)
  
  
  def _quantile_75(x):
      return x.quantile(0.75)
  
  
  # Note that all valid aggregation methods (e.g. GroupBy.min) are bound to the
  # class after its definition (see below).
  class GroupBy(Serializable):
  
      _MAX_GROUPS_BEFORE_WARN = 100
  
      def __init__(
          self, obj, by=None, level=None, sort=False, as_index=True, dropna=True
      ):
          """
          Group a DataFrame or Series by a set of columns.
  
          Parameters
          ----------
          by : optional
              Specifies the grouping columns. Can be any of the following:
              - A Python function called on each value of the object's index
              - A dict or Series that maps index labels to group names
              - A cudf.Index object
              - A str indicating a column name
              - An array of the same length as the object
              - A Grouper object
              - A list of the above
          level : int, level_name or list, optional
              For objects with a MultiIndex, `level` can be used to specify
              grouping by one or more levels of the MultiIndex.
          sort : bool, default False
              Sort the result by group keys. Differ from Pandas, cudf defaults
              to False for better performance.
          as_index : bool, optional
              If as_index=True (default), the group names appear
              as the keys of the resulting DataFrame.
              If as_index=False, the groups are returned as ordinary
              columns of the resulting DataFrame, *if they are named columns*.
          dropna : bool, optional
              If True (default), do not include the "null" group.
          """
          self.obj = obj
          self._as_index = as_index
          self._sort = sort
          self._dropna = dropna
  
          if isinstance(by, _Grouping):
              self.grouping = by
          else:
              self.grouping = _Grouping(obj, by, level)
  
      def __iter__(self):
          group_names, offsets, _, grouped_values = self._grouped()
          if isinstance(group_names, cudf.BaseIndex):
              group_names = group_names.to_pandas()
          for i, name in enumerate(group_names):
              yield name, grouped_values[offsets[i] : offsets[i + 1]]
  
      @cached_property
      def groups(self):
          """
          Returns a dictionary mapping group keys to row labels.
          """
          group_names, offsets, _, grouped_values = self._grouped()
          grouped_index = grouped_values.index
  
          if len(group_names) > self._MAX_GROUPS_BEFORE_WARN:
              warnings.warn(
                  f"GroupBy.groups() performance scales poorly with "
                  f"number of groups. Got {len(group_names)} groups."
              )
  
          return dict(
              zip(group_names.to_pandas(), grouped_index._split(offsets[1:-1]))
          )
  
      def get_group(self, name, obj=None):
          """
          Construct DataFrame from group with provided name.
  
          Parameters
          ----------
          name : object
              The name of the group to get as a DataFrame.
          obj : DataFrame, default None
              The DataFrame to take the DataFrame out of.  If
              it is None, the object groupby was called on will
              be used.
  
          Returns
          -------
          group : same type as obj
  
          Examples
          --------
          >>> import cudf
          >>> df = cudf.DataFrame({"X": ["A", "B", "A", "B"], "Y": [1, 4, 3, 2]})
          >>> df
             X  Y
          0  A  1
          1  B  4
          2  A  3
          3  B  2
          >>> df.groupby("X").get_group("A")
             X  Y
          0  A  1
          2  A  3
          """
          if obj is None:
              obj = self.obj
  
          return obj.loc[self.groups[name]]
  
      def size(self):
          """
          Return the size of each group.
          """
          return (
              cudf.Series(
                  cudf.core.column.column_empty(
                      len(self.obj), "int8", masked=False
                  )
              )
              .groupby(self.grouping, sort=self._sort)
              .agg("size")
          )
  
      def cumcount(self):
          """
          Return the cumulative count of keys in each group.
          """
          return (
              cudf.Series(
                  cudf.core.column.column_empty(
                      len(self.obj), "int8", masked=False
                  )
              )
              .groupby(self.grouping, sort=self._sort)
              .agg("cumcount")
              .reset_index(drop=True)
          )
  
      @cached_property
      def _groupby(self):
          return libgroupby.GroupBy(self.grouping.keys, dropna=self._dropna)
  
      @annotate("GROUPBY_AGG", domain="cudf_python")
      def agg(self, func):
          """
          Apply aggregation(s) to the groups.
  
          Parameters
          ----------
          func : str, callable, list or dict
  
          Returns
          -------
          A Series or DataFrame containing the combined results of the
          aggregation.
  
          Examples
          --------
          >>> import cudf
          >>> a = cudf.DataFrame(
              {'a': [1, 1, 2], 'b': [1, 2, 3], 'c': [2, 2, 1]})
          >>> a.groupby('a').agg('sum')
             b  c
          a
          2  3  1
          1  3  4
  
          Specifying a list of aggregations to perform on each column.
  
          >>> a.groupby('a').agg(['sum', 'min'])
              b       c
            sum min sum min
          a
          2   3   3   1   1
          1   3   1   4   2
  
          Using a dict to specify aggregations to perform per column.
  
          >>> a.groupby('a').agg({'a': 'max', 'b': ['min', 'mean']})
              a   b
            max min mean
          a
          2   2   3  3.0
          1   1   1  1.5
  
          Using lambdas/callables to specify aggregations taking parameters.
  
          >>> f1 = lambda x: x.quantile(0.5); f1.__name__ = "q0.5"
          >>> f2 = lambda x: x.quantile(0.75); f2.__name__ = "q0.75"
          >>> a.groupby('a').agg([f1, f2])
               b          c
            q0.5 q0.75 q0.5 q0.75
          a
          1  1.5  1.75  2.0   2.0
          2  3.0  3.00  1.0   1.0
          """
          normalized_aggs = self._normalize_aggs(func)
  
          # Note: When there are no key columns, the below produces
          # a Float64Index, while Pandas returns an Int64Index
          # (GH: 6945)
          result = cudf.DataFrame._from_data(
              *self._groupby.aggregate(self.obj, normalized_aggs)
          )
  
          if self._sort:
              result = result.sort_index()
  
          if not _is_multi_agg(func):
              if result.columns.nlevels == 1:
                  # make sure it's a flat index:
                  result.columns = result.columns.get_level_values(0)
  
              if result.columns.nlevels > 1:
                  try:
                      # drop the last level
                      result.columns = result.columns.droplevel(-1)
                  except IndexError:
                      # Pandas raises an IndexError if we are left
                      # with an all-nan MultiIndex when dropping
                      # the last level
                      if result.shape[1] == 1:
                          result.columns = [None]
                      else:
                          raise
  
          if libgroupby._is_all_scan_aggregate(normalized_aggs):
              # Scan aggregations return rows in original index order
              return self._mimic_pandas_order(result)
  
          # set index names to be group key names
          if len(result):
              result.index.names = self.grouping.names
  
          # copy categorical information from keys to the result index:
          result.index._copy_type_metadata(self.grouping.keys)
          result._index = cudf.Index(result._index)
  
          if not self._as_index:
              for col_name in reversed(self.grouping._named_columns):
                  result.insert(
                      0,
                      col_name,
                      result.index.get_level_values(col_name)._values,
                  )
              result.index = cudf.core.index.RangeIndex(len(result))
  
          return result
  
      aggregate = agg
  
      def nth(self, n):
          """
          Return the nth row from each group.
          """
          result = self.agg(lambda x: x.nth(n)).sort_index()
          sizes = self.size().sort_index()
  
          return result[sizes > n]
  
      def serialize(self):
          header = {}
          frames = []
  
          header["kwargs"] = {
              "sort": self._sort,
              "dropna": self._dropna,
              "as_index": self._as_index,
          }
  
          obj_header, obj_frames = self.obj.serialize()
          header["obj"] = obj_header
          header["obj_type"] = pickle.dumps(type(self.obj))
          header["num_obj_frames"] = len(obj_frames)
          frames.extend(obj_frames)
  
          grouping_header, grouping_frames = self.grouping.serialize()
          header["grouping"] = grouping_header
          header["num_grouping_frames"] = len(grouping_frames)
          frames.extend(grouping_frames)
  
          return header, frames
  
      @classmethod
      def deserialize(cls, header, frames):
          kwargs = header["kwargs"]
  
          obj_type = pickle.loads(header["obj_type"])
          obj = obj_type.deserialize(
              header["obj"], frames[: header["num_obj_frames"]]
          )
          grouping = _Grouping.deserialize(
              header["grouping"], frames[header["num_obj_frames"] :]
          )
          return cls(obj, grouping, **kwargs)
  
      def _grouped(self):
          grouped_keys, grouped_values, offsets = self._groupby.groups(self.obj)
          grouped_values = self.obj.__class__._from_data(*grouped_values)
          grouped_values._copy_type_metadata(self.obj)
          group_names = grouped_keys.unique()
          return (group_names, offsets, grouped_keys, grouped_values)
  
      def _normalize_aggs(self, aggs):
          """
          Normalize aggs to a dict mapping column names
          to a list of aggregations.
          """
          if not isinstance(aggs, collections.abc.Mapping):
              # Make col_name->aggs mapping from aggs.
              # Do not include named key columns
  
              # Can't do set arithmetic here as sets are
              # not ordered
              if isinstance(self, SeriesGroupBy):
                  columns = [self.obj.name]
              else:
                  columns = [
                      col_name
                      for col_name in self.obj._data
                      if col_name not in self.grouping._named_columns
                  ]
              out = dict.fromkeys(columns, aggs)
          else:
              out = aggs.copy()
  
          # Convert all values to list-like:
          for col, agg in out.items():
              if not is_list_like(agg):
                  out[col] = [agg]
              else:
                  out[col] = list(agg)
  
          return out
  
      def pipe(self, func, *args, **kwargs):
          """
          Apply a function `func` with arguments to this GroupBy
          object and return the function’s result.
  
          Parameters
          ----------
          func : function
              Function to apply to this GroupBy object or,
              alternatively, a ``(callable, data_keyword)`` tuple where
              ``data_keyword`` is a string indicating the keyword of
              ``callable`` that expects the GroupBy object.
          args : iterable, optional
              Positional arguments passed into ``func``.
          kwargs : mapping, optional
              A dictionary of keyword arguments passed into ``func``.
  
          Returns
          -------
          object : the return type of ``func``.
  
          See also
          --------
          cudf.Series.pipe
              Apply a function with arguments to a series.
  
          cudf.DataFrame.pipe
              Apply a function with arguments to a dataframe.
  
          apply
              Apply function to each group instead of to the full GroupBy object.
  
          Examples
          --------
          >>> import cudf
          >>> df = cudf.DataFrame({'A': ['a', 'b', 'a', 'b'], 'B': [1, 2, 3, 4]})
          >>> df
             A  B
          0  a  1
          1  b  2
          2  a  3
          3  b  4
  
          To get the difference between each groups maximum and minimum value
          in one pass, you can do
  
          >>> df.groupby('A').pipe(lambda x: x.max() - x.min())
             B
          A
          a  2
          b  2
          """
          return cudf.core.common.pipe(self, func, *args, **kwargs)
  
      def apply(self, function):
          """Apply a python transformation function over the grouped chunk.
  
          Parameters
          ----------
          func : function
            The python transformation function that will be applied
            on the grouped chunk.
  
          Examples
          --------
          .. code-block:: python
  
            from cudf import DataFrame
            df = DataFrame()
            df['key'] = [0, 0, 1, 1, 2, 2, 2]
            df['val'] = [0, 1, 2, 3, 4, 5, 6]
            groups = df.groupby(['key'])
  
            # Define a function to apply to each row in a group
            def mult(df):
              df['out'] = df['key'] * df['val']
              return df
  
            result = groups.apply(mult)
            print(result)
  
          Output:
  
          .. code-block:: python
  
               key  val  out
            0    0    0    0
            1    0    1    0
            2    1    2    2
            3    1    3    3
            4    2    4    8
            5    2    5   10
            6    2    6   12
  
          .. pandas-compat::
              **groupby.apply**
  
              cuDF's ``groupby.apply`` is limited compared to pandas.
              In some situations, Pandas returns the grouped keys as part of
              the index while cudf does not due to redundancy. For example:
  
              .. code-block::
  
                  >>> df = pd.DataFrame({
                      'a': [1, 1, 2, 2],
                      'b': [1, 2, 1, 2],
                      'c': [1, 2, 3, 4]})
                  >>> gdf = cudf.from_pandas(df)
                  >>> df.groupby('a').apply(lambda x: x.iloc[[0]])
                          a  b  c
                      a
                      1 0  1  1  1
                      2 2  2  1  3
                  >>> gdf.groupby('a').apply(lambda x: x.iloc[[0]])
                          a  b  c
                      0  1  1  1
                      2  2  1  3
          """
          if not callable(function):
              raise TypeError(f"type {type(function)} is not callable")
          group_names, offsets, _, grouped_values = self._grouped()
  
          ngroups = len(offsets) - 1
          if ngroups > self._MAX_GROUPS_BEFORE_WARN:
              warnings.warn(
                  f"GroupBy.apply() performance scales poorly with "
                  f"number of groups. Got {ngroups} groups."
              )
  
          chunks = [
              grouped_values[s:e] for s, e in zip(offsets[:-1], offsets[1:])
          ]
          chunk_results = [function(chk) for chk in chunks]
  
          if not len(chunk_results):
              return self.obj.head(0)
  
          if cudf.api.types.is_scalar(chunk_results[0]):
              result = cudf.Series(chunk_results, index=group_names)
              result.index.names = self.grouping.names
          elif isinstance(chunk_results[0], cudf.Series):
              result = cudf.concat(chunk_results, axis=1).T
              result.index.names = self.grouping.names
          else:
              result = cudf.concat(chunk_results)
  
          if self._sort:
              result = result.sort_index()
          return result
  
      def apply_grouped(self, function, **kwargs):
          """Apply a transformation function over the grouped chunk.
  
          This uses numba's CUDA JIT compiler to convert the Python
          transformation function into a CUDA kernel, thus will have a
          compilation overhead during the first run.
  
          Parameters
          ----------
          func : function
            The transformation function that will be executed on the CUDA GPU.
          incols: list
            A list of names of input columns.
          outcols: list
            A dictionary of output column names and their dtype.
          kwargs : dict
            name-value of extra arguments. These values are passed directly into
            the function.
  
          Examples
          --------
          .. code-block:: python
  
              from cudf import DataFrame
              from numba import cuda
              import numpy as np
  
              df = DataFrame()
              df['key'] = [0, 0, 1, 1, 2, 2, 2]
              df['val'] = [0, 1, 2, 3, 4, 5, 6]
              groups = df.groupby(['key'])
  
              # Define a function to apply to each group
              def mult_add(key, val, out1, out2):
                  for i in range(cuda.threadIdx.x, len(key), cuda.blockDim.x):
                      out1[i] = key[i] * val[i]
                      out2[i] = key[i] + val[i]
  
              result = groups.apply_grouped(mult_add,
                                            incols=['key', 'val'],
                                            outcols={'out1': np.int32,
                                                     'out2': np.int32},
                                            # threads per block
                                            tpb=8)
  
              print(result)
  
          Output:
  
          .. code-block:: python
  
                 key  val out1 out2
              0    0    0    0    0
              1    0    1    0    1
              2    1    2    2    3
              3    1    3    3    4
              4    2    4    8    6
              5    2    5   10    7
              6    2    6   12    8
  
  
  
          .. code-block:: python
  
              import cudf
              import numpy as np
              from numba import cuda
              import pandas as pd
              from random import randint
  
  
              # Create a random 15 row dataframe with one categorical
              # feature and one random integer valued feature
              df = cudf.DataFrame(
                      {
                          "cat": [1] * 5 + [2] * 5 + [3] * 5,
                          "val": [randint(0, 100) for _ in range(15)],
                      }
                   )
  
              # Group the dataframe by its categorical feature
              groups = df.groupby("cat")
  
              # Define a kernel which takes the moving average of a
              # sliding window
              def rolling_avg(val, avg):
                  win_size = 3
                  for i in range(cuda.threadIdx.x, len(val), cuda.blockDim.x):
                      if i < win_size - 1:
                          # If there is not enough data to fill the window,
                          # take the average to be NaN
                          avg[i] = np.nan
                      else:
                          total = 0
                          for j in range(i - win_size + 1, i + 1):
                              total += val[j]
                          avg[i] = total / win_size
  
              # Compute moving averages on all groups
              results = groups.apply_grouped(rolling_avg,
                                             incols=['val'],
                                             outcols=dict(avg=np.float64))
              print("Results:", results)
  
              # Note this gives the same result as its pandas equivalent
              pdf = df.to_pandas()
              pd_results = pdf.groupby('cat')['val'].rolling(3).mean()
  
  
          Output:
  
          .. code-block:: python
  
              Results:
                 cat  val                 avg
              0    1   16
              1    1   45
              2    1   62                41.0
              3    1   45  50.666666666666664
              4    1   26  44.333333333333336
              5    2    5
              6    2   51
              7    2   77  44.333333333333336
              8    2    1                43.0
              9    2   46  41.333333333333336
              [5 more rows]
  
          This is functionally equivalent to `pandas.DataFrame.Rolling
          <https://pandas.pydata.org/pandas-docs/stable/reference/api/pandas.DataFrame.rolling.html>`_
  
          """
          if not callable(function):
              raise TypeError(f"type {type(function)} is not callable")
  
          _, offsets, _, grouped_values = self._grouped()
          kwargs.update({"chunks": offsets})
          return grouped_values.apply_chunks(function, **kwargs)
  
      def rolling(self, *args, **kwargs):
          """
          Returns a `RollingGroupby` object that enables rolling window
          calculations on the groups.
  
          See also
          --------
          cudf.core.window.Rolling
          """
          return cudf.core.window.rolling.RollingGroupby(self, *args, **kwargs)
  
      def count(self, dropna=True):
          """Compute the number of values in each column.
  
          Parameters
          ----------
          dropna : bool
              If ``True``, don't include null values in the count.
          """
  
          def func(x):
              return getattr(x, "count")(dropna=dropna)
  
          return self.agg(func)
  
      def describe(self, include=None, exclude=None):
          """
          Generate descriptive statistics that summarizes the central tendency,
          dispersion and shape of a dataset’s distribution, excluding NaN values.
  
          Analyzes numeric DataFrames only
  
          Parameters
          ----------
          include: ‘all’, list-like of dtypes or None (default), optional
              list of data types to include in the result.
              Ignored for Series.
  
          exclude: list-like of dtypes or None (default), optional,
              list of data types to omit from the result.
              Ignored for Series.
  
          Returns
          -------
          Series or DataFrame
              Summary statistics of the Dataframe provided.
  
          Examples
          --------
          >>> import cudf
          >>> gdf = cudf.DataFrame({"Speed": [380.0, 370.0, 24.0, 26.0],
                                    "Score": [50, 30, 90, 80]})
          >>> gdf
          Speed  Score
          0  380.0     50
          1  370.0     30
          2   24.0     90
          3   26.0     80
          >>> gdf.groupby('Score').describe()
               Speed
               count   mean   std    min    25%    50%    75%     max
          Score
          30        1  370.0  <NA>  370.0  370.0  370.0  370.0  370.0
          50        1  380.0  <NA>  380.0  380.0  380.0  380.0  380.0
          80        1   26.0  <NA>   26.0   26.0   26.0   26.0   26.0
          90        1   24.0  <NA>   24.0   24.0   24.0   24.0   24.0
  
          """
          if exclude is not None and include is not None:
              raise NotImplementedError
  
          res = self.agg(
              [
                  "count",
                  "mean",
                  "std",
                  "min",
                  _quantile_25,
                  _quantile_50,
                  _quantile_75,
                  "max",
              ]
          )
          res.rename(
              columns={
                  "_quantile_25": "25%",
                  "_quantile_50": "50%",
                  "_quantile_75": "75%",
              },
              level=1,
              inplace=True,
          )
          return res
  
      def sum(self):
          """Compute the column-wise sum of the values in each group."""
          return self.agg("sum")
  
      def prod(self):
          """Compute the column-wise product of the values in each group."""
          return self.agg("prod")
  
      def idxmin(self):
          """Get the column-wise index of the minimum value in each group."""
          return self.agg("idxmin")
  
      def idxmax(self):
          """Get the column-wise index of the maximum value in each group."""
          return self.agg("idxmax")
  
      def min(self):
          """Get the column-wise minimum value in each group."""
          return self.agg("min")
  
      def max(self):
          """Get the column-wise maximum value in each group."""
          return self.agg("max")
  
      def mean(self):
          """Compute the column-wise mean of the values in each group."""
          return self.agg("mean")
  
      def median(self):
          """Get the column-wise median of the values in each group."""
          return self.agg("median")
  
      def var(self, ddof=1):
          """Compute the column-wise variance of the values in each group.
  
          Parameters
          ----------
          ddof : int
              The delta degrees of freedom. N - ddof is the divisor used to
              normalize the variance.
          """
  
          def func(x):
              return getattr(x, "var")(ddof=ddof)
  
          return self.agg(func)
  
      def std(self, ddof=1):
          """Compute the column-wise std of the values in each group.
  
          Parameters
          ----------
          ddof : int
              The delta degrees of freedom. N - ddof is the divisor used to
              normalize the standard deviation.
          """
  
          def func(x):
              return getattr(x, "std")(ddof=ddof)
  
          return self.agg(func)
  
      def quantile(self, q=0.5, interpolation="linear"):
          """Compute the column-wise quantiles of the values in each group.
  
          Parameters
          ----------
          q : float or array-like
              The quantiles to compute.
          interpolation : {"linear", "lower", "higher", "midpoint", "nearest"}
              The interpolation method to use when the desired quantile lies
              between two data points. Defaults to "linear".
         """
  
          def func(x):
              return getattr(x, "quantile")(q=q, interpolation=interpolation)
  
          return self.agg(func)
  
      def nunique(self):
          """Compute the number of unique values in each column in each group."""
          return self.agg("nunique")
  
      def collect(self):
          """Get a list of all the values for each column in each group."""
          return self.agg("collect")
  
      def unique(self):
          """Get a list of the unique values for each column in each group."""
          return self.agg("unique")
  
      def cumsum(self):
          """Compute the column-wise cumulative sum of the values in
          each group."""
          return self.agg("cumsum")
  
      def cummin(self):
          """Get the column-wise cumulative minimum value in each group."""
          return self.agg("cummin")
  
      def cummax(self):
          """Get the column-wise cumulative maximum value in each group."""
          return self.agg("cummax")
  
      def first(self):
          """Get the first non-null value in each group."""
          return self.agg("first")
  
      def last(self):
          """Get the last non-null value in each group."""
          return self.agg("last")
  
      def _scan_fill(self, method: str, limit: int) -> DataFrameOrSeries:
          """Internal implementation for `ffill` and `bfill`
          """
          value_columns = self.grouping.values
          result = self.obj.__class__._from_data(
              self._groupby.replace_nulls(
                  cudf.core.frame.Frame(value_columns._data), method
              )
          )
          result = self._mimic_pandas_order(result)
          return result._copy_type_metadata(value_columns)
  
      def pad(self, limit=None):
          """Forward fill NA values.
  
          Parameters
          ----------
          limit : int, default None
              Unsupported
          """
  
          if limit is not None:
              raise NotImplementedError("Does not support limit param yet.")
  
          return self._scan_fill("ffill", limit)
  
      ffill = pad
  
      def backfill(self, limit=None):
          """Backward fill NA values.
  
          Parameters
          ----------
          limit : int, default None
              Unsupported
          """
          if limit is not None:
              raise NotImplementedError("Does not support limit param yet.")
  
          return self._scan_fill("bfill", limit)
  
      bfill = backfill
  
      def fillna(
          self,
          value=None,
          method=None,
          axis=0,
          inplace=False,
          limit=None,
          downcast=None,
      ):
          """Fill NA values using the specified method.
  
          Parameters
          ----------
          value : scalar, dict
              Value to use to fill the holes. Cannot be specified with method.
          method : {'backfill', 'bfill', 'pad', 'ffill', None}, default None
              Method to use for filling holes in reindexed Series
  
              - pad/ffill: propagate last valid observation forward to next valid
              - backfill/bfill: use next valid observation to fill gap
          axis : {0 or 'index', 1 or 'columns'}
              Unsupported
          inplace : bool, default False
              If `True`, fill inplace. Note: this will modify other views on this
              object.
          limit : int, default None
              Unsupported
          downcast : dict, default None
              Unsupported
  
          Returns
          -------
          DataFrame or Series
  
          .. pandas-compat::
              **groupby.fillna**
  
              This function may return result in different format to the method
              Pandas supports. For example:
  
              .. code-block::
  
                  >>> df = pd.DataFrame({'k': [1, 1, 2], 'v': [2, None, 4]})
                  >>> gdf = cudf.from_pandas(df)
                  >>> df.groupby('k').fillna({'v': 4}) # pandas
                         v
                  k
                  1 0  2.0
                    1  4.0
                  2 2  4.0
                  >>> gdf.groupby('k').fillna({'v': 4}) # cudf
                       v
                  0  2.0
                  1  4.0
                  2  4.0
          """
          if inplace:
              raise NotImplementedError("Does not support inplace yet.")
          if limit is not None:
              raise NotImplementedError("Does not support limit param yet.")
          if downcast is not None:
              raise NotImplementedError("Does not support downcast yet.")
          if not axis == 0:
              raise NotImplementedError("Only support axis == 0.")
  
          if value is None and method is None:
              raise ValueError("Must specify a fill 'value' or 'method'.")
          if value is not None and method is not None:
              raise ValueError("Cannot specify both 'value' and 'method'.")
  
          if method is not None:
              if method not in {"pad", "ffill", "backfill", "bfill"}:
                  raise ValueError(
                      "Method can only be of 'pad', 'ffill',"
                      "'backfill', 'bfill'."
                  )
              return getattr(self, method, limit)()
  
          value_columns = self.grouping.values
          _, (data, index), _ = self._groupby.groups(
              cudf.core.frame.Frame(value_columns._data)
          )
  
          grouped = self.obj.__class__._from_data(data, index)
          result = grouped.fillna(
              value=value, inplace=inplace, axis=axis, limit=limit
          )
          result = self._mimic_pandas_order(result)
          return result._copy_type_metadata(value_columns)
  
      def shift(self, periods=1, freq=None, axis=0, fill_value=None):
          """
          Shift each group by ``periods`` positions.
  
          Parameters
          ----------
          periods : int, default 1
              Number of periods to shift.
          freq : str, unsupported
          axis : 0, axis to shift
              Shift direction. Only row-wise shift is supported
          fill_value : scalar or list of scalars, optional
              The scalar value to use for newly introduced missing values. Can be
              specified with `None`, a single value or multiple values:
  
              - `None` (default): sets all indeterminable values to null.
              - Single value: fill all shifted columns with this value. Should
                match the data type of all columns.
              - List of values: fill shifted columns with corresponding value in
                the list. The length of the list should match the number of
                columns shifted. Each value should match the data type of the
                column to fill.
  
          Returns
          -------
          Series or DataFrame
              Object shifted within each group.
  
          Notes
          -----
          Parameter ``freq`` is unsupported.
          """
  
          if freq is not None:
              raise NotImplementedError("Parameter freq is unsupported.")
  
          if not axis == 0:
              raise NotImplementedError("Only axis=0 is supported.")
  
          value_columns = self.grouping.values
          if is_list_like(fill_value):
              if not len(fill_value) == len(value_columns._data):
                  raise ValueError(
                      "Mismatched number of columns and values to fill."
                  )
          else:
              fill_value = [fill_value] * len(value_columns._data)
  
          result = self.obj.__class__._from_data(
              *self._groupby.shift(
                  cudf.core.frame.Frame(value_columns), periods, fill_value
              )
          )
          result = self._mimic_pandas_order(result)
          return result._copy_type_metadata(value_columns)
  
      def _mimic_pandas_order(
          self, result: DataFrameOrSeries
      ) -> DataFrameOrSeries:
          """Given a groupby result from libcudf, reconstruct the row orders
          matching that of pandas. This also adds appropriate indices.
          """
          sorted_order_column = arange(0, result._data.nrows)
          _, (order, _), _ = self._groupby.groups(
              cudf.core.frame.Frame({"sorted_order_column": sorted_order_column})
          )
          gather_map = order["sorted_order_column"].argsort()
          result = result.take(gather_map)
          result.index = self.obj.index
          return result
  
  
  class DataFrameGroupBy(GroupBy, GetAttrGetItemMixin):
      """
      Group DataFrame using a mapper or by a Series of columns.
  
      A groupby operation involves some combination of splitting the object,
      applying a function, and combining the results. This can be used to
      group large amounts of data and compute operations on these groups.
  
      Parameters
      ----------
      by : mapping, function, label, or list of labels
          Used to determine the groups for the groupby. If by is a
          function, it’s called on each value of the object’s index.
          If a dict or Series is passed, the Series or dict VALUES will
          be used to determine the groups (the Series’ values are first
          aligned; see .align() method). If a cupy array is passed, the
          values are used as-is determine the groups. A label or list
          of labels may be passed to group by the columns in self.
          Notice that a tuple is interpreted as a (single) key.
      level : int, level name, or sequence of such, default None
          If the axis is a MultiIndex (hierarchical), group by a particular
          level or levels.
      as_index : bool, default True
          For aggregated output, return object with group labels as
          the index. Only relevant for DataFrame input.
          as_index=False is effectively “SQL-style” grouped output.
      sort : bool, default False
          Sort result by group key. Differ from Pandas, cudf defaults to
          ``False`` for better performance. Note this does not influence
          the order of observations within each group. Groupby preserves
          the order of rows within each group.
      dropna : bool, optional
          If True (default), do not include the "null" group.
  
      Returns
      -------
          DataFrameGroupBy
              Returns a groupby object that contains information
              about the groups.
  
      Examples
      --------
      >>> import cudf
      >>> import pandas as pd
      >>> df = cudf.DataFrame({'Animal': ['Falcon', 'Falcon',
      ...                               'Parrot', 'Parrot'],
      ...                    'Max Speed': [380., 370., 24., 26.]})
      >>> df
         Animal  Max Speed
      0  Falcon      380.0
      1  Falcon      370.0
      2  Parrot       24.0
      3  Parrot       26.0
      >>> df.groupby(['Animal']).mean()
              Max Speed
      Animal
      Falcon      375.0
      Parrot       25.0
  
      >>> arrays = [['Falcon', 'Falcon', 'Parrot', 'Parrot'],
      ... ['Captive', 'Wild', 'Captive', 'Wild']]
      >>> index = pd.MultiIndex.from_arrays(arrays, names=('Animal', 'Type'))
      >>> df = cudf.DataFrame({'Max Speed': [390., 350., 30., 20.]},
              index=index)
      >>> df
                      Max Speed
      Animal Type
      Falcon Captive      390.0
             Wild         350.0
      Parrot Captive       30.0
             Wild          20.0
      >>> df.groupby(level=0).mean()
              Max Speed
      Animal
      Falcon      370.0
      Parrot       25.0
      >>> df.groupby(level="Type").mean()
              Max Speed
      Type
      Wild         185.0
      Captive      210.0
      """
  
      _PROTECTED_KEYS = frozenset(("obj",))
  
      def __init__(
          self, obj, by=None, level=None, sort=False, as_index=True, dropna=True
      ):
          super().__init__(
              obj=obj,
              by=by,
              level=level,
              sort=sort,
              as_index=as_index,
              dropna=dropna,
          )
  
      def __getitem__(self, key):
          return self.obj[key].groupby(
              self.grouping, dropna=self._dropna, sort=self._sort
          )
  
      def nunique(self):
          """
          Return the number of unique values per group
          """
          return self.agg("nunique")
  
  
  class SeriesGroupBy(GroupBy):
      """
      Group Series using a mapper or by a Series of columns.
  
      A groupby operation involves some combination of splitting the object,
      applying a function, and combining the results. This can be used to
      group large amounts of data and compute operations on these groups.
  
      Parameters
      ----------
      by : mapping, function, label, or list of labels
          Used to determine the groups for the groupby. If by is a
          function, it’s called on each value of the object’s index.
          If a dict or Series is passed, the Series or dict VALUES will
          be used to determine the groups (the Series’ values are first
          aligned; see .align() method). If an cupy array is passed, the
          values are used as-is determine the groups. A label or list
          of labels may be passed to group by the columns in self.
          Notice that a tuple is interpreted as a (single) key.
      level : int, level name, or sequence of such, default None
          If the axis is a MultiIndex (hierarchical), group by a particular
          level or levels.
      as_index : bool, default True
          For aggregated output, return object with group labels as
          the index. Only relevant for DataFrame input.
          as_index=False is effectively “SQL-style” grouped output.
      sort : bool, default False
          Sort result by group key. Differ from Pandas, cudf defaults to
          ``False`` for better performance. Note this does not influence
          the order of observations within each group. Groupby preserves
          the order of rows within each group.
  
      Returns
      -------
          SeriesGroupBy
              Returns a groupby object that contains information
              about the groups.
  
      Examples
      --------
      >>> ser = cudf.Series([390., 350., 30., 20.],
      ...                 index=['Falcon', 'Falcon', 'Parrot', 'Parrot'],
      ...                 name="Max Speed")
      >>> ser
      Falcon    390.0
      Falcon    350.0
      Parrot     30.0
      Parrot     20.0
      Name: Max Speed, dtype: float64
      >>> ser.groupby(level=0).mean()
      Falcon    370.0
      Parrot     25.0
      Name: Max Speed, dtype: float64
      >>> ser.groupby(ser > 100).mean()
      Max Speed
      False     25.0
      True     370.0
      Name: Max Speed, dtype: float64
      """
  
      def __init__(
          self, obj, by=None, level=None, sort=False, as_index=True, dropna=True
      ):
          super().__init__(
              obj=obj,
              by=by,
              level=level,
              sort=sort,
              as_index=as_index,
              dropna=dropna,
          )
  
      def agg(self, func):
          result = super().agg(func)
  
          # downcast the result to a Series:
          if len(result._data):
              if result.shape[1] == 1 and not is_list_like(func):
                  return result.iloc[:, 0]
  
          # drop the first level if we have a multiindex
          if (
              isinstance(result.columns, pd.MultiIndex)
              and result.columns.nlevels > 1
          ):
              result.columns = result.columns.droplevel(0)
  
          return result
  
      def apply(self, func):
          result = super().apply(func)
  
          # apply Series name to result
          result.name = self.obj.name
  
          return result
  
  
  class Grouper(object):
      def __init__(self, key=None, level=None):
          if key is not None and level is not None:
              raise ValueError("Grouper cannot specify both key and level")
          if key is None and level is None:
              raise ValueError("Grouper must specify either key or level")
          self.key = key
          self.level = level
  
  
  class _Grouping(Serializable):
      def __init__(self, obj, by=None, level=None):
          self._obj = obj
          self._key_columns = []
          self.names = []
          # For transform operations, we want to filter out only the value
          # columns that will be used. When part of the key columns are composed
          # from columns in `obj`, these columns are not included in the value
          # columns. This only happens when `by` is specified with
          # column labels or `Grouper` object with `key` param. To avoid that
          # external objects overlaps in names to `obj`, these column
          # names are recorded separately in this list.
          self._key_column_names_from_obj = []
  
          # Need to keep track of named key columns
          # to support `as_index=False` correctly
          self._named_columns = []
          self._handle_by_or_level(by, level)
  
          if len(obj) and not len(self._key_columns):
              raise ValueError("No group keys passed")
  
      def _handle_by_or_level(self, by=None, level=None):
          if level is not None:
              if by is not None:
                  raise ValueError("Cannot specify both by and level")
              level_list = level if isinstance(level, list) else [level]
              for level in level_list:
                  self._handle_level(level)
          else:
              by_list = by if isinstance(by, list) else [by]
  
              for by in by_list:
                  if callable(by):
                      self._handle_callable(by)
                  elif isinstance(by, cudf.Series):
                      self._handle_series(by)
                  elif isinstance(by, cudf.BaseIndex):
                      self._handle_index(by)
                  elif isinstance(by, collections.abc.Mapping):
                      self._handle_mapping(by)
                  elif isinstance(by, Grouper):
                      self._handle_grouper(by)
                  else:
                      try:
                          self._handle_label(by)
                      except (KeyError, TypeError):
                          self._handle_misc(by)
  
      @property
      def keys(self):
          """Return grouping key columns as index
          """
          nkeys = len(self._key_columns)
  
          if nkeys == 0:
              return cudf.core.index.as_index([], name=None)
          elif nkeys > 1:
              return cudf.MultiIndex._from_data(
                  dict(zip(range(nkeys), self._key_columns))
              )._set_names(self.names)
          else:
              return cudf.core.index.as_index(
                  self._key_columns[0], name=self.names[0]
              )
  
      @property
      def values(self):
          """Return value columns as a frame.
  
          Note that in aggregation, value columns can be arbitrarily
          specified. While this method returns all non-key columns from `obj` as
          a frame.
  
          This is mainly used in transform-like operations.
          """
          # If the key columns are in `obj`, filter them out
          value_column_names = [
              x
              for x in self._obj._data.names
              if x not in self._key_column_names_from_obj
          ]
          value_columns = self._obj._data.select_by_label(value_column_names)
          return self._obj.__class__._from_data(value_columns)
  
      def _handle_callable(self, by):
          by = by(self._obj.index)
          self.__init__(self._obj, by)
  
      def _handle_series(self, by):
          by = by._align_to_index(self._obj.index, how="right")
          self._key_columns.append(by._column)
          self.names.append(by.name)
  
      def _handle_index(self, by):
          self._key_columns.extend(by._data.columns)
          self.names.extend(by._data.names)
  
      def _handle_mapping(self, by):
          by = cudf.Series(by.values(), index=by.keys())
          self._handle_series(by)
  
      def _handle_label(self, by):
          self._key_columns.append(self._obj._data[by])
          self.names.append(by)
          self._named_columns.append(by)
          self._key_column_names_from_obj.append(by)
  
      def _handle_grouper(self, by):
          if by.key:
              self._handle_label(by.key)
          else:
              self._handle_level(by.level)
  
      def _handle_level(self, by):
          level_values = self._obj.index.get_level_values(by)
          self._key_columns.append(level_values._values)
          self.names.append(level_values.name)
  
      def _handle_misc(self, by):
          by = cudf.core.column.as_column(by)
          if len(by) != len(self._obj):
              raise ValueError("Grouper and object must have same length")
          self._key_columns.append(by)
          self.names.append(None)
  
      def serialize(self):
          header = {}
          frames = []
          header["names"] = pickle.dumps(self.names)
          header["_named_columns"] = pickle.dumps(self._named_columns)
          column_header, column_frames = cudf.core.column.serialize_columns(
              self._key_columns
          )
          header["columns"] = column_header
          frames.extend(column_frames)
          return header, frames
  
      @classmethod
      def deserialize(cls, header, frames):
          names = pickle.loads(header["names"])
          _named_columns = pickle.loads(header["_named_columns"])
          key_columns = cudf.core.column.deserialize_columns(
              header["columns"], frames
          )
          out = _Grouping.__new__(_Grouping)
          out.names = names
          out._named_columns = _named_columns
          out._key_columns = key_columns
          return out
  
  
  def _is_multi_agg(aggs):
      """
      Returns True if more than one aggregation is performed
      on any of the columns as specified in `aggs`.
      """
      if isinstance(aggs, collections.abc.Mapping):
          return any(is_list_like(agg) for agg in aggs.values())
      if is_list_like(aggs):
          return True
      return False
  