original_docstring:  Order divisions of DataFrame so that all values within column align
  
      This enacts a task-based shuffle.  It contains most of the tricky logic
      around the complex network of tasks.  Typically before this function is
      called a new column, ``"_partitions"`` has been added to the dataframe,
      containing the output partition number of every row.  This function
      produces a new dataframe where every row is in the proper partition.  It
      accomplishes this by splitting each input partition into several pieces,
      and then concatenating pieces from different input partitions into output
      partitions.  If there are enough partitions then it does this work in
      stages to avoid scheduling overhead.
  
      Parameters
      ----------
      df: dask.dataframe.DataFrame
      column: str
          A column name on which we want to split, commonly ``"_partitions"``
          which is assigned by functions upstream
      max_branch: int
          The maximum number of splits per input partition.  Defaults to 32.
          If there are more partitions than this then the shuffling will occur in
          stages in order to avoid creating npartitions**2 tasks
          Increasing this number increases scheduling overhead but decreases the
          number of full-dataset transfers that we have to make.
      npartitions: Optional[int]
          The desired number of output partitions
  
      Returns
      -------
      df3: dask.dataframe.DataFrame
  
      See also
      --------
      rearrange_by_column_disk: same operation, but uses partd
      rearrange_by_column: parent function that calls this or rearrange_by_column_disk
      shuffle_group: does the actual splitting per-partition
      
docstring: Order divisions of DataFrame so that all values within column align
  This enacts a task-based shuffle.  It contains most of the tricky logic
  around the complex network of tasks.  Typically before this function is
  called a new column, ``"_partitions"`` has been added to the dataframe,
  containing the output partition number of every row.  This function
  produces a new dataframe where every row is in the proper partition.  It
  accomplishes this by splitting each input partition into several pieces,
  and then concatenating pieces from different input partitions into output
  partitions.  If there are enough partitions then it does this work in
  stages to avoid scheduling overhead.
  
  Parameters
  
  dask.dataframe.DataFrame
  column: str
  A column name on which we want to split, commonly ``"_partitions"``
  which is assigned by functions upstream
  max_branch: int
  The maximum number of splits per input partition.  Defaults to 32.
  If there are more partitions than this then the shuffling will occur in
  stages in order to avoid creating npartitions**2 tasks
  Increasing this number increases scheduling overhead but decreases the
  number of full-dataset transfers that we have to make.
  npartitions: Optional[int]
  The desired number of output partitions
  
  Returns
  
  
  
  See also
  
  same operation, but uses partd
  rearrange_by_column: parent function that calls this or rearrange_by_column_disk
  shuffle_group: does the actual splitting per-partition
docstring_tokens: Order divisions of DataFrame so that all values within column align This enacts a task - based shuffle . It contains most of the tricky logic around the complex network of tasks . Typically before this function is called a new column ` ` " _partitions " ` ` has been added to the dataframe containing the output partition number of every row . This function produces a new dataframe where every row is in the proper partition . It accomplishes this by splitting each input partition into several pieces and then concatenating pieces from different input partitions into output partitions . If there are enough partitions then it does this work in stages to avoid scheduling overhead . Parameters dask . dataframe . DataFrame column : str A column name on which we want to split commonly ` ` " _partitions " ` ` which is assigned by functions upstream max_branch : int The maximum number of splits per input partition . Defaults to 32 . If there are more partitions than this then the shuffling will occur in stages in order to avoid creating npartitions ** 2 tasks Increasing this number increases scheduling overhead but decreases the number of full - dataset transfers that we have to make . npartitions : Optional [ int ] The desired number of output partitions Returns See also same operation but uses partd rearrange_by_column : parent function that calls this or rearrange_by_column_disk shuffle_group : does the actual splitting per - partition
code: def rearrange_by_column_tasks(df, column, max_branch=32, npartitions=None):
      """ Order divisions of DataFrame so that all values within column align
  
      This enacts a task-based shuffle.  It contains most of the tricky logic
      around the complex network of tasks.  Typically before this function is
      called a new column, ``"_partitions"`` has been added to the dataframe,
      containing the output partition number of every row.  This function
      produces a new dataframe where every row is in the proper partition.  It
      accomplishes this by splitting each input partition into several pieces,
      and then concatenating pieces from different input partitions into output
      partitions.  If there are enough partitions then it does this work in
      stages to avoid scheduling overhead.
  
      Parameters
      ----------
      df: dask.dataframe.DataFrame
      column: str
          A column name on which we want to split, commonly ``"_partitions"``
          which is assigned by functions upstream
      max_branch: int
          The maximum number of splits per input partition.  Defaults to 32.
          If there are more partitions than this then the shuffling will occur in
          stages in order to avoid creating npartitions**2 tasks
          Increasing this number increases scheduling overhead but decreases the
          number of full-dataset transfers that we have to make.
      npartitions: Optional[int]
          The desired number of output partitions
  
      Returns
      -------
      df3: dask.dataframe.DataFrame
  
      See also
      --------
      rearrange_by_column_disk: same operation, but uses partd
      rearrange_by_column: parent function that calls this or rearrange_by_column_disk
      shuffle_group: does the actual splitting per-partition
      """
      max_branch = max_branch or 32
      n = df.npartitions
  
      stages = int(math.ceil(math.log(n) / math.log(max_branch)))
      if stages > 1:
          k = int(math.ceil(n ** (1 / stages)))
      else:
          k = n
  
      groups = []
      splits = []
      joins = []
  
      inputs = [tuple(digit(i, j, k) for j in range(stages)) for i in range(k ** stages)]
  
      token = tokenize(df, column, max_branch)
  
      start = {
          ("shuffle-join-" + token, 0, inp): (df._name, i)
          if i < df.npartitions
          else df._meta
          for i, inp in enumerate(inputs)
      }
  
      for stage in range(1, stages + 1):
          group = {  # Convert partition into dict of dataframe pieces
              ("shuffle-group-" + token, stage, inp): (
                  shuffle_group,
                  ("shuffle-join-" + token, stage - 1, inp),
                  column,
                  stage - 1,
                  k,
                  n,
              )
              for inp in inputs
          }
  
          split = {  # Get out each individual dataframe piece from the dicts
              ("shuffle-split-" + token, stage, i, inp): (
                  getitem,
                  ("shuffle-group-" + token, stage, inp),
                  i,
              )
              for i in range(k)
              for inp in inputs
          }
  
          join = {  # concatenate those pieces together, with their friends
              ("shuffle-join-" + token, stage, inp): (
                  _concat,
                  [
                      (
                          "shuffle-split-" + token,
                          stage,
                          inp[stage - 1],
                          insert(inp, stage - 1, j),
                      )
                      for j in range(k)
                  ],
              )
              for inp in inputs
          }
          groups.append(group)
          splits.append(split)
          joins.append(join)
  
      end = {
          ("shuffle-" + token, i): ("shuffle-join-" + token, stages, inp)
          for i, inp in enumerate(inputs)
      }
  
      dsk = toolz.merge(start, end, *(groups + splits + joins))
      graph = HighLevelGraph.from_collections("shuffle-" + token, dsk, dependencies=[df])
      df2 = DataFrame(graph, "shuffle-" + token, df, df.divisions)
  
      if npartitions is not None and npartitions != df.npartitions:
          parts = [i % df.npartitions for i in range(npartitions)]
          token = tokenize(df2, npartitions)
  
          dsk = {
              ("repartition-group-" + token, i): (shuffle_group_2, k, column)
              for i, k in enumerate(df2.__dask_keys__())
          }
          for p in range(npartitions):
              dsk[("repartition-get-" + token, p)] = (
                  shuffle_group_get,
                  ("repartition-group-" + token, parts[p]),
                  p,
              )
  
          graph2 = HighLevelGraph.from_collections(
              "repartition-get-" + token, dsk, dependencies=[df2]
          )
          df3 = DataFrame(
              graph2, "repartition-get-" + token, df2, [None] * (npartitions + 1)
          )
      else:
          df3 = df2
          df3.divisions = (None,) * (df.npartitions + 1)
  
      return df3
code_tokens: def rearrange_by_column_tasks ( df , column , max_branch = 32 , npartitions = None ) : max_branch = max_branch or 32 n = df . npartitions stages = int ( math . ceil ( math . log ( n ) / math . log ( max_branch ) ) ) if stages > 1 : k = int ( math . ceil ( n ** ( 1 / stages ) ) ) else : k = n groups = [ ] splits = [ ] joins = [ ] inputs = [ tuple ( digit ( i , j , k ) for j in range ( stages ) ) for i in range ( k ** stages ) ] token = tokenize ( df , column , max_branch ) start = { ( "shuffle-join-" + token , 0 , inp ) : ( df . _name , i ) if i < df . npartitions else df . _meta for i , inp in enumerate ( inputs ) } for stage in range ( 1 , stages + 1 ) : group = { ( "shuffle-group-" + token , stage , inp ) : ( shuffle_group , ( "shuffle-join-" + token , stage - 1 , inp ) , column , stage - 1 , k , n , ) for inp in inputs } split = { ( "shuffle-split-" + token , stage , i , inp ) : ( getitem , ( "shuffle-group-" + token , stage , inp ) , i , ) for i in range ( k ) for inp in inputs } join = { ( "shuffle-join-" + token , stage , inp ) : ( _concat , [ ( "shuffle-split-" + token , stage , inp [ stage - 1 ] , insert ( inp , stage - 1 , j ) , ) for j in range ( k ) ] , ) for inp in inputs } groups . append ( group ) splits . append ( split ) joins . append ( join ) end = { ( "shuffle-" + token , i ) : ( "shuffle-join-" + token , stages , inp ) for i , inp in enumerate ( inputs ) } dsk = toolz . merge ( start , end , * ( groups + splits + joins ) ) graph = HighLevelGraph . from_collections ( "shuffle-" + token , dsk , dependencies = [ df ] ) df2 = DataFrame ( graph , "shuffle-" + token , df , df . divisions ) if npartitions is not None and npartitions != df . npartitions : parts = [ i % df . npartitions for i in range ( npartitions ) ] token = tokenize ( df2 , npartitions ) dsk = { ( "repartition-group-" + token , i ) : ( shuffle_group_2 , k , column ) for i , k in enumerate ( df2 . __dask_keys__ ( ) ) } for p in range ( npartitions ) : dsk [ ( "repartition-get-" + token , p ) ] = ( shuffle_group_get , ( "repartition-group-" + token , parts [ p ] ) , p , ) graph2 = HighLevelGraph . from_collections ( "repartition-get-" + token , dsk , dependencies = [ df2 ] ) df3 = DataFrame ( graph2 , "repartition-get-" + token , df2 , [ None ] * ( npartitions + 1 ) ) else : df3 = df2 df3 . divisions = ( None , ) * ( df . npartitions + 1 ) return df3
short_docstring: Order divisions of DataFrame so that all values within column align
  This enacts a task-based shuffle.
short_docstring_tokens: Order divisions of DataFrame so that all values within column align This enacts a task - based shuffle .
comment: """ Order divisions of DataFrame so that all values within column align
  
      This enacts a task-based shuffle.  It contains most of the tricky logic
      around the complex network of tasks.  Typically before this function is
      called a new column, ``"_partitions"`` has been added to the dataframe,
      containing the output partition number of every row.  This function
      produces a new dataframe where every row is in the proper partition.  It
      accomplishes this by splitting each input partition into several pieces,
      and then concatenating pieces from different input partitions into output
      partitions.  If there are enough partitions then it does this work in
      stages to avoid scheduling overhead.
  
      Parameters
      ----------
      df: dask.dataframe.DataFrame
      column: str
          A column name on which we want to split, commonly ``"_partitions"``
          which is assigned by functions upstream
      max_branch: int
          The maximum number of splits per input partition.  Defaults to 32.
          If there are more partitions than this then the shuffling will occur in
          stages in order to avoid creating npartitions**2 tasks
          Increasing this number increases scheduling overhead but decreases the
          number of full-dataset transfers that we have to make.
      npartitions: Optional[int]
          The desired number of output partitions
  
      Returns
      -------
      df3: dask.dataframe.DataFrame
  
      See also
      --------
      rearrange_by_column_disk: same operation, but uses partd
      rearrange_by_column: parent function that calls this or rearrange_by_column_disk
      shuffle_group: does the actual splitting per-partition
      """$SEP$# Convert partition into dict of dataframe pieces$SEP$# Get out each individual dataframe piece from the dicts$SEP$# concatenate those pieces together, with their friends