original_docstring:  Splits dataframe into groups
  
      The group is determined by their final partition, and which stage we are in
      in the shuffle
  
      Parameters
      ----------
      df: DataFrame
      col: str
          Column name on which to split the dataframe
      stage: int
          We shuffle dataframes with many partitions we in a few stages to avoid
          a quadratic number of tasks.  This number corresponds to which stage
          we're in, starting from zero up to some small integer
      k: int
          Desired number of splits from this dataframe
      npartition: int
          Total number of output partitions for the full dataframe
  
      Returns
      -------
      out: Dict[int, DataFrame]
          A dictionary mapping integers in {0..k} to dataframes such that the
          hash values of ``df[col]`` are well partitioned.
      
docstring: Splits dataframe into groups
  The group is determined by their final partition, and which stage we are in
  in the shuffle
  
  Parameters
  
  DataFrame
  col: str
  Column name on which to split the dataframe
  stage: int
  We shuffle dataframes with many partitions we in a few stages to avoid
  a quadratic number of tasks.  This number corresponds to which stage
  we're in, starting from zero up to some small integer
  k: int
  Desired number of splits from this dataframe
  npartition: int
  Total number of output partitions for the full dataframe
  
  Returns
  
  Dict[int, DataFrame]
  A dictionary mapping integers in {0..k} to dataframes such that the
  hash values of ``df[col]`` are well partitioned.
docstring_tokens: Splits dataframe into groups The group is determined by their final partition and which stage we are in in the shuffle Parameters DataFrame col : str Column name on which to split the dataframe stage : int We shuffle dataframes with many partitions we in a few stages to avoid a quadratic number of tasks . This number corresponds to which stage we ' re in starting from zero up to some small integer k : int Desired number of splits from this dataframe npartition : int Total number of output partitions for the full dataframe Returns Dict [ int DataFrame ] A dictionary mapping integers in { 0 .. k } to dataframes such that the hash values of ` ` df [ col ] ` ` are well partitioned .
code: def shuffle_group(df, col, stage, k, npartitions):
      """ Splits dataframe into groups
  
      The group is determined by their final partition, and which stage we are in
      in the shuffle
  
      Parameters
      ----------
      df: DataFrame
      col: str
          Column name on which to split the dataframe
      stage: int
          We shuffle dataframes with many partitions we in a few stages to avoid
          a quadratic number of tasks.  This number corresponds to which stage
          we're in, starting from zero up to some small integer
      k: int
          Desired number of splits from this dataframe
      npartition: int
          Total number of output partitions for the full dataframe
  
      Returns
      -------
      out: Dict[int, DataFrame]
          A dictionary mapping integers in {0..k} to dataframes such that the
          hash values of ``df[col]`` are well partitioned.
      """
      if col == "_partitions":
          ind = df[col]
      else:
          ind = hash_pandas_object(df[col], index=False)
  
      c = ind._values
      typ = np.min_scalar_type(npartitions * 2)
  
      c = np.mod(c, npartitions).astype(typ, copy=False)
      np.floor_divide(c, k ** stage, out=c)
      np.mod(c, k, out=c)
  
      indexer, locations = groupsort_indexer(c.astype(np.int64), k)
      df2 = df.take(indexer)
      locations = locations.cumsum()
      parts = [df2.iloc[a:b] for a, b in zip(locations[:-1], locations[1:])]
  
      return dict(zip(range(k), parts))
code_tokens: def shuffle_group ( df , col , stage , k , npartitions ) : if col == "_partitions" : ind = df [ col ] else : ind = hash_pandas_object ( df [ col ] , index = False ) c = ind . _values typ = np . min_scalar_type ( npartitions * 2 ) c = np . mod ( c , npartitions ) . astype ( typ , copy = False ) np . floor_divide ( c , k ** stage , out = c ) np . mod ( c , k , out = c ) indexer , locations = groupsort_indexer ( c . astype ( np . int64 ) , k ) df2 = df . take ( indexer ) locations = locations . cumsum ( ) parts = [ df2 . iloc [ a : b ] for a , b in zip ( locations [ : - 1 ] , locations [ 1 : ] ) ] return dict ( zip ( range ( k ) , parts ) )
short_docstring: Splits dataframe into groups
  The group is determined by their final partition, and which stage we are in
  in the shuffle
short_docstring_tokens: Splits dataframe into groups The group is determined by their final partition and which stage we are in in the shuffle
comment: """ Splits dataframe into groups
  
      The group is determined by their final partition, and which stage we are in
      in the shuffle
  
      Parameters
      ----------
      df: DataFrame
      col: str
          Column name on which to split the dataframe
      stage: int
          We shuffle dataframes with many partitions we in a few stages to avoid
          a quadratic number of tasks.  This number corresponds to which stage
          we're in, starting from zero up to some small integer
      k: int
          Desired number of splits from this dataframe
      npartition: int
          Total number of output partitions for the full dataframe
  
      Returns
      -------
      out: Dict[int, DataFrame]
          A dictionary mapping integers in {0..k} to dataframes such that the
          hash values of ``df[col]`` are well partitioned.
      """