#NEWLINE# def set_partition( df , index , divisions , compute = False , drop = True , **kwargs ) : #NEWLINE# #INDENT# ' Group DataFrame by index\n\n    Sets a new index and partitions data along that index according to\n    divisions.  Divisions are often found by computing approximate quantiles.\n    The function ``set_index`` will do both of these steps.\n\n    Parameters\n    ----------\n    df: DataFrame/Series\n        Data that we want to re-partition\n    index: string or Series\n        Column to become the new index\n    divisions: list\n        Values to form new divisions between partitions\n    drop: bool, default True\n        Whether to delete columns to be used as the new index\n\n    See Also\n    --------\n    set_index\n    shuffle\n    partd\n    ' #NEWLINE# if isinstance ( index , Series ) : #NEWLINE# #INDENT# assert ( df . divisions == index . divisions ) #NEWLINE# metadata = df . _pd . set_index ( index . _pd , drop = drop ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# if np . isscalar ( index ) : #NEWLINE# #INDENT# metadata = df . _pd . set_index ( index , drop = drop ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# raise ValueError ( 'index must be Series or scalar, {0} given' . format ( type ( index ) ) ) #NEWLINE# #UNINDENT# #UNINDENT# token = tokenize ( df , index , divisions ) #NEWLINE# always_new_token = uuid . uuid1 ( ) . hex #NEWLINE# import partd #NEWLINE# p = ( ( 'zpartd-' + always_new_token ) , ) #NEWLINE# catname = ( 'set-partition--get-categories-old-' + always_new_token ) #NEWLINE# catname2 = ( 'set-partition--get-categories-new-' + always_new_token ) #NEWLINE# dsk1 = { #NEWLINE# #INDENT# catname : ( get_categories , df . _keys ( ) [ 0 ] ) , #NEWLINE# p : ( partd . PandasBlocks , ( partd . Buffer , ( partd . Dict , ) , ( partd . File , ) ) ) , #NEWLINE# catname2 : ( new_categories , catname , ( index . name if isinstance ( index , Series ) else index ) ) , #NEWLINE# #UNINDENT# } #NEWLINE# name = ( 'set-partition--partition-' + always_new_token ) #NEWLINE# if isinstance ( index , _Frame ) : #NEWLINE# #INDENT# dsk2 = dict ( ( ( ( name , i ) , ( _set_partition , part , ind , divisions , p , drop ) ) for ( i , ( part , ind ) ) in enumerate ( zip ( df . _keys ( ) , index . _keys ( ) ) ) ) ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# dsk2 = dict ( ( ( ( name , i ) , ( _set_partition , part , index , divisions , p , drop ) ) for ( i , part ) in enumerate ( df . _keys ( ) ) ) ) #NEWLINE# #UNINDENT# barrier_token = ( 'barrier-' + always_new_token ) #NEWLINE# dsk3 = { #NEWLINE# #INDENT# barrier_token : ( barrier , list ( dsk2 ) ) , #NEWLINE# #UNINDENT# } #NEWLINE# if compute : #NEWLINE# #INDENT# dsk = merge ( df . dask , dsk1 , dsk2 , dsk3 ) #NEWLINE# if isinstance ( index , _Frame ) : #NEWLINE# #INDENT# dsk . update ( index . dask ) #NEWLINE# #UNINDENT# ( p , barrier_token , categories ) = df . _get ( dsk , [ p , barrier_token , catname ] , ** kwargs ) #NEWLINE# dsk4 = { #NEWLINE# #INDENT# catname2 : categories , #NEWLINE# #UNINDENT# } #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# dsk4 = { #NEWLINE# #INDENT# #NEWLINE# #UNINDENT# } #NEWLINE# #UNINDENT# name = ( 'set-partition--collect-' + token ) #NEWLINE# if ( compute and ( not categories ) ) : #NEWLINE# #INDENT# dsk4 . update ( dict ( ( ( ( name , i ) , ( _set_collect , i , p , barrier_token , df . columns ) ) for i in range ( ( len ( divisions ) - 1 ) ) ) ) ) #NEWLINE# #UNINDENT# else : #NEWLINE# #INDENT# dsk4 . update ( dict ( ( ( ( name , i ) , ( _categorize , catname2 , ( _set_collect , i , p , barrier_token , df . columns ) ) ) for i in range ( ( len ( divisions ) - 1 ) ) ) ) ) #NEWLINE# #UNINDENT# dsk = merge ( df . dask , dsk1 , dsk2 , dsk3 , dsk4 ) #NEWLINE# if isinstance ( index , Series ) : #NEWLINE# #INDENT# dsk . update ( index . dask ) #NEWLINE# #UNINDENT# if compute : #NEWLINE# #INDENT# dsk = cull ( dsk , list ( dsk4 . keys ( ) ) ) #NEWLINE# #UNINDENT# return DataFrame ( dsk , name , metadata , divisions )