code: def load_balance_func(ddf_, by, client=None):
      # Load balances the sorted dask_cudf DataFrame.
      # Input is a dask_cudf dataframe ddf_ which is sorted by
      # the column name passed as the 'by' argument.
  
      client = default_client() if client is None else client
  
      parts = persist_distributed_data(ddf_, client)
      wait(parts)
  
      who_has = client.who_has(parts)
      key_to_part = [(part.key, part) for part in parts]
      gpu_fututres = [
          (first(who_has[key]), part.key[1], part)
          for key, part in key_to_part
      ]
      worker_to_data = create_dict(gpu_fututres)
  
      # Calculate cumulative sum in each dataframe partition
      cumsum_parts = [
          client.submit(
              get_cumsum, wf[1][0][0], by, workers=[wf[0]]
          ).result()
          for idx, wf in enumerate(worker_to_data.items())
      ]
  
      num_rows = []
      for cumsum in cumsum_parts:
          num_rows.append(cumsum.iloc[-1])
  
      # Calculate current partition divisions.
      divisions = [
          sum(num_rows[0:x:1]) for x in range(0, len(num_rows) + 1)
      ]
      divisions[-1] = divisions[-1] - 1
      divisions = tuple(divisions)
  
      # Set global index from 0 to len(dask_cudf_dataframe) so that global
      # indexing of divisions can be used for repartitioning.
      futures = [
          client.submit(
              set_global_index,
              wf[1][0][0],
              divisions[wf[1][0][1]],
              workers=[wf[0]],
          )
          for idx, wf in enumerate(worker_to_data.items())
      ]
      wait(futures)
  
      ddf = dask_cudf.from_delayed(futures)
      ddf.divisions = divisions
  
      # Repartition the data
      ddf = repartition(ddf, cumsum_parts)
  
      return ddf
  
imports: import numpy
  from dask.distributed import futures_of, default_client, wait
  from toolz import first
  import collections
  import dask_cudf
  from dask.array.core import Array
  from dask_cudf.core import DataFrame
  from dask_cudf.core import Series
  from functools import reduce
  import cugraph.dask.comms.comms
  from dask.delayed import delayed
  import cudf
label_desc: Use image node_id 14 for calling a global function with example usage: load_balance_func(ddf_, by, client) and returns: ddf