def shuffle_group_2(df, cols, ignore_index, nparts):
if not len(df):
return {}, df
if isinstance(cols, str):
cols = [cols]
if cols and cols[0] == "_partitions":
ind = df[cols[0]].astype(np.int32)
else:
ind = (hash_object_dispatch(df[cols] if cols else df, index=False) %
int(nparts)).astype(np.int32)
n = ind.max() + 1
result2 = group_split_dispatch(df, ind, n, ignore_index=ignore_index)
return result2, df.iloc[:0]
def shuffle_group(df, cols, stage, k, npartitions, ignore_index, nfinal):
"""Splits dataframe into groups
The group is determined by their final partition, and which stage we are in
in the shuffle
Parameters
----------
df: DataFrame
cols: str or list
Column name(s) on which to split the dataframe. If ``cols`` is not
"_partitions", hashing will be used to determine target partition
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
nfinal: int
Total number of output partitions after repartitioning
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 isinstance(cols, str):
cols = [cols]
if cols and cols[0] == "_partitions":
ind = df[cols[0]]
else:
ind = hash_object_dispatch(df[cols] if cols else df, index=False)
if nfinal and nfinal != npartitions:
ind = ind % int(nfinal)
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)
return group_split_dispatch(df, c, k, ignore_index=ignore_index)
def partitioning_index(df, npartitions):
"""
Computes a deterministic index mapping each record to a partition.
Identical rows are mapped to the same partition.
Parameters
----------
df : DataFrame/Series/Index
npartitions : int
The number of partitions to group into.
Returns
-------
partitions : ndarray
An array of int64 values mapping each record to a partition.
"""
return hash_object_dispatch(df, index=False) % int(npartitions)