File: C:/Users/fred/anaconda3/Lib/site-packages/dask/dataframe/shuffle.py
from __future__ import annotations
import contextlib
import logging
import math
import shutil
import tempfile
import uuid
import warnings
from collections.abc import Callable, Mapping, Sequence
from typing import Any, Literal
import numpy as np
import pandas as pd
import tlz as toolz
from pandas.api.types import is_numeric_dtype
from dask import config
from dask.base import compute, compute_as_if_collection, is_dask_collection, tokenize
from dask.dataframe import methods
from dask.dataframe._compat import PANDAS_GE_300
from dask.dataframe.core import (
DataFrame,
Series,
_deprecated_kwarg,
_Frame,
map_partitions,
new_dd_object,
)
from dask.dataframe.dispatch import (
group_split_dispatch,
hash_object_dispatch,
partd_encode_dispatch,
)
from dask.dataframe.utils import UNKNOWN_CATEGORIES
from dask.highlevelgraph import HighLevelGraph
from dask.layers import ShuffleLayer, SimpleShuffleLayer
from dask.sizeof import sizeof
from dask.utils import M, digit, get_default_shuffle_method
logger = logging.getLogger(__name__)
def _calculate_divisions(
df: DataFrame,
partition_col: Series,
repartition: bool,
npartitions: int,
upsample: float = 1.0,
partition_size: float = 128e6,
ascending: bool = True,
) -> tuple[list, list, list, bool]:
"""
Utility function to calculate divisions for calls to `map_partitions`
"""
sizes = df.map_partitions(sizeof) if repartition else []
divisions = partition_col._repartition_quantiles(npartitions, upsample=upsample)
mins = partition_col.map_partitions(M.min)
maxes = partition_col.map_partitions(M.max)
try:
divisions, sizes, mins, maxes = compute(divisions, sizes, mins, maxes)
except TypeError as e:
# When there are nulls and a column is non-numeric, a TypeError is sometimes raised as a result of
# 1) computing mins/maxes above, 2) every null being switched to NaN, and 3) NaN being a float.
# Also, Pandas ExtensionDtypes may cause TypeErrors when dealing with special nulls such as pd.NaT or pd.NA.
# If this happens, we hint the user about eliminating nulls beforehand.
if not is_numeric_dtype(partition_col.dtype):
obj, suggested_method = (
("column", f"`.dropna(subset=['{partition_col.name}'])`")
if any(partition_col._name == df[c]._name for c in df)
else ("series", "`.loc[series[~series.isna()]]`")
)
raise NotImplementedError(
f"Divisions calculation failed for non-numeric {obj} '{partition_col.name}'.\n"
f"This is probably due to the presence of nulls, which Dask does not entirely support in the index.\n"
f"We suggest you try with {suggested_method}."
) from e
# For numeric types there shouldn't be problems with nulls, so we raise as-it-is this particular TypeError
else:
raise e
empty_dataframe_detected = pd.isna(divisions).all()
if repartition or empty_dataframe_detected:
total = sum(sizes)
npartitions = max(math.ceil(total / partition_size), 1)
npartitions = min(npartitions, df.npartitions)
n = divisions.size
try:
divisions = np.interp(
x=np.linspace(0, n - 1, npartitions + 1),
xp=np.linspace(0, n - 1, n),
fp=divisions.tolist(),
).tolist()
except (TypeError, ValueError): # str type
indexes = np.linspace(0, n - 1, npartitions + 1).astype(int)
divisions = divisions.iloc[indexes].tolist()
else:
# Drop duplicate divisions returned by partition quantiles
n = divisions.size
divisions = (
list(divisions.iloc[: n - 1].unique()) + divisions.iloc[n - 1 :].tolist()
)
mins = mins.bfill()
maxes = maxes.bfill()
if isinstance(partition_col.dtype, pd.CategoricalDtype):
dtype = partition_col.dtype
mins = mins.astype(dtype)
maxes = maxes.astype(dtype)
if mins.isna().any() or maxes.isna().any():
presorted = False
else:
n = mins.size
maxes2 = (maxes.iloc[: n - 1] if ascending else maxes.iloc[1:]).reset_index(
drop=True
)
mins2 = (mins.iloc[1:] if ascending else mins.iloc[: n - 1]).reset_index(
drop=True
)
presorted = (
mins.tolist() == mins.sort_values(ascending=ascending).tolist()
and maxes.tolist() == maxes.sort_values(ascending=ascending).tolist()
and (maxes2 < mins2).all()
)
return divisions, mins.tolist(), maxes.tolist(), presorted
@_deprecated_kwarg("shuffle", "shuffle_method")
def sort_values(
df: DataFrame,
by: str | list[str],
npartitions: int | Literal["auto"] | None = None,
shuffle_method: str | None = None,
ascending: bool | list[bool] = True,
na_position: Literal["first"] | Literal["last"] = "last",
upsample: float = 1.0,
partition_size: float = 128e6,
sort_function: Callable[[pd.DataFrame], pd.DataFrame] | None = None,
sort_function_kwargs: Mapping[str, Any] | None = None,
) -> DataFrame:
"""See DataFrame.sort_values for docstring"""
if na_position not in ("first", "last"):
raise ValueError("na_position must be either 'first' or 'last'")
if not isinstance(by, list):
by = [by]
if any(not isinstance(b, str) for b in by):
raise NotImplementedError(
"Dataframes only support sorting by named columns which must be passed as a "
"string or a list of strings.\n"
"You passed %s" % str(by)
)
if (
ascending is not None
and not isinstance(ascending, bool)
and not len(ascending) == len(by)
):
raise ValueError(f"Length of {ascending=} != length of {by=}")
sort_kwargs = {
"by": by,
"ascending": ascending,
"na_position": na_position,
}
if sort_function is None:
sort_function = M.sort_values
if sort_function_kwargs is not None:
sort_kwargs.update(sort_function_kwargs)
if df.npartitions == 1:
return df.map_partitions(sort_function, **sort_kwargs)
if npartitions == "auto":
warnings.warn(
"`npartitions='auto'` is deprecated, either set it as an integer or leave as `None`.",
FutureWarning,
2,
)
repartition = True
npartitions = max(100, df.npartitions)
else:
if npartitions is None:
npartitions = df.npartitions
repartition = False
sort_by_col = df[by[0]]
divisions_ascending = ascending
if divisions_ascending and not isinstance(divisions_ascending, bool):
divisions_ascending = divisions_ascending[0]
assert divisions_ascending is None or isinstance(divisions_ascending, bool)
divisions, _, _, presorted = _calculate_divisions(
df,
sort_by_col,
repartition,
npartitions,
upsample,
partition_size,
divisions_ascending,
)
if len(divisions) == 2:
return df.repartition(npartitions=1).map_partitions(
sort_function, **sort_kwargs
)
if presorted and npartitions == df.npartitions:
# divisions are in the right place
return df.map_partitions(sort_function, **sort_kwargs)
df = rearrange_by_divisions(
df,
by[0],
divisions,
shuffle_method=shuffle_method,
ascending=divisions_ascending,
na_position=na_position,
duplicates=False,
)
df = df.map_partitions(sort_function, **sort_kwargs)
return df
@_deprecated_kwarg("compute")
@_deprecated_kwarg("shuffle", "shuffle_method")
def set_index(
df: DataFrame,
index: str | Series,
npartitions: int | Literal["auto"] | None = None,
shuffle_method: str | None = None,
compute: bool = False,
drop: bool = True,
upsample: float = 1.0,
divisions: Sequence | None = None,
partition_size: float = 128e6,
sort: bool = True,
**kwargs,
) -> DataFrame:
"""See _Frame.set_index for docstring"""
if not sort:
return df.map_partitions(
M.set_index, index, align_dataframes=False, drop=drop, **kwargs
).clear_divisions()
if npartitions == "auto":
warnings.warn(
"`npartitions='auto'` is deprecated, either set it as an integer or leave as `None`.",
FutureWarning,
2,
)
repartition = True
npartitions = max(100, df.npartitions)
else:
if npartitions is None:
npartitions = df.npartitions
repartition = False
if not isinstance(index, Series):
index2 = df[index]
else:
index2 = index
if divisions is None:
divisions, mins, maxes, presorted = _calculate_divisions(
df, index2, repartition, npartitions, upsample, partition_size
)
if presorted and npartitions == df.npartitions:
divisions = mins + [maxes[-1]]
result = set_sorted_index(df, index, drop=drop, divisions=divisions)
return result.map_partitions(M.sort_index)
return set_partition(
df,
index,
divisions,
shuffle_method=shuffle_method,
drop=drop,
compute=compute,
**kwargs,
)
@_deprecated_kwarg("shuffle", "shuffle_method")
def set_partition(
df: DataFrame,
index: str | Series,
divisions: Sequence,
max_branch: int = 32,
drop: bool = True,
shuffle_method: str | None = None,
compute: bool | None = None,
) -> DataFrame:
"""Group DataFrame by index
Sets a new index and partitions data along that index according to
divisions. Divisions are often found by computing approximate quantiles.
The function ``set_index`` will do both of these steps.
Parameters
----------
df: DataFrame/Series
Data that we want to re-partition
index: string or Series
Column to become the new index
divisions: list
Values to form new divisions between partitions
drop: bool, default True
Whether to delete columns to be used as the new index
shuffle_method: str (optional)
Either 'disk' for an on-disk shuffle or 'tasks' to use the task
scheduling framework. Use 'disk' if you are on a single machine
and 'tasks' if you are on a distributed cluster.
max_branch: int (optional)
If using the task-based shuffle, the amount of splitting each
partition undergoes. Increase this for fewer copies but more
scheduler overhead.
See Also
--------
set_index
shuffle
partd
"""
if isinstance(divisions, tuple):
# pd.isna considers tuples to be scalars. Convert to a list.
divisions = list(divisions)
if not isinstance(index, Series):
dtype = df[index].dtype
else:
dtype = index.dtype
if pd.isna(divisions).any() and pd.api.types.is_integer_dtype(dtype):
# Can't construct a Series[int64] when any / all of the divisions are NaN.
divisions = df._meta._constructor_sliced(divisions)
elif (
isinstance(dtype, pd.CategoricalDtype)
and UNKNOWN_CATEGORIES in dtype.categories
):
# If categories are unknown, leave as a string dtype instead.
divisions = df._meta._constructor_sliced(divisions)
else:
divisions = df._meta._constructor_sliced(divisions, dtype=dtype)
meta = df._meta._constructor_sliced([0])
# Ensure that we have the same index as before to avoid alignment
# when calculating meta dtypes later on
meta.index = df._meta_nonempty.index[:1]
if not isinstance(index, Series):
partitions = df[index].map_partitions(
set_partitions_pre, divisions=divisions, meta=meta
)
df2 = df.assign(_partitions=partitions)
else:
partitions = index.map_partitions(
set_partitions_pre, divisions=divisions, meta=meta
)
df2 = df.assign(_partitions=partitions, _index=index)
df3 = rearrange_by_column(
df2,
"_partitions",
max_branch=max_branch,
npartitions=len(divisions) - 1,
shuffle_method=shuffle_method,
compute=compute,
ignore_index=True,
)
if not isinstance(index, Series):
df4 = df3.map_partitions(
set_index_post_scalar,
index_name=index,
drop=drop,
column_dtype=df.columns.dtype,
)
else:
df4 = df3.map_partitions(
set_index_post_series,
index_name=index.name,
drop=drop,
column_dtype=df.columns.dtype,
)
divisions = methods.tolist(divisions)
# None and pd.NA values are not sortable
df4.divisions = tuple(i if not pd.isna(i) else np.nan for i in divisions)
return df4.map_partitions(M.sort_index)
@_deprecated_kwarg("shuffle", "shuffle_method")
def shuffle(
df,
index,
shuffle_method=None,
npartitions=None,
max_branch=32,
ignore_index=False,
compute=None,
):
"""Group DataFrame by index
Hash grouping of elements. After this operation all elements that have
the same index will be in the same partition. Note that this requires
full dataset read, serialization and shuffle. This is expensive. If
possible you should avoid shuffles.
This does not preserve a meaningful index/partitioning scheme. This is not
deterministic if done in parallel.
See Also
--------
set_index
set_partition
shuffle_disk
"""
list_like = pd.api.types.is_list_like(index) and not is_dask_collection(index)
shuffle_method = shuffle_method or get_default_shuffle_method()
if not isinstance(index, _Frame):
if list_like:
# Make sure we don't try to select with pd.Series/pd.Index
index = list(index)
index = df._select_columns_or_index(index)
elif hasattr(index, "to_frame"):
# If this is an index, we should still convert to a
# DataFrame. Otherwise, the hashed values of a column
# selection will not match (important when merging).
index = index.to_frame()
dtypes = {}
for col, dtype in index.dtypes.items():
if pd.api.types.is_numeric_dtype(dtype):
dtypes[col] = np.float64
if not dtypes:
dtypes = None
meta = df._meta._constructor_sliced([0])
# Ensure that we have the same index as before to avoid alignment
# when calculating meta dtypes later on
meta.index = df._meta_nonempty.index[:1]
partitions = index.map_partitions(
partitioning_index,
npartitions=npartitions or df.npartitions,
meta=meta,
transform_divisions=False,
cast_dtype=dtypes,
)
df2 = df.assign(_partitions=partitions)
df2._meta.index.name = df._meta.index.name
df3 = rearrange_by_column(
df2,
"_partitions",
npartitions=npartitions,
max_branch=max_branch,
shuffle_method=shuffle_method,
compute=compute,
ignore_index=ignore_index,
)
del df3["_partitions"]
return df3
@_deprecated_kwarg("shuffle", "shuffle_method")
def rearrange_by_divisions(
df,
column,
divisions,
max_branch=None,
shuffle_method=None,
ascending=True,
na_position="last",
duplicates=True,
):
"""Shuffle dataframe so that column separates along divisions"""
divisions = df._meta._constructor_sliced(divisions)
# duplicates need to be removed sometimes to properly sort null dataframes
if not duplicates:
divisions = divisions.drop_duplicates()
meta = df._meta._constructor_sliced([0])
# Ensure that we have the same index as before to avoid alignment
# when calculating meta dtypes later on
meta.index = df._meta_nonempty.index[:1]
# Assign target output partitions to every row
partitions = df[column].map_partitions(
set_partitions_pre,
divisions=divisions,
ascending=ascending,
na_position=na_position,
meta=meta,
)
df2 = df.assign(_partitions=partitions)
# Perform shuffle
df3 = rearrange_by_column(
df2,
"_partitions",
max_branch=max_branch,
npartitions=len(divisions) - 1,
shuffle_method=shuffle_method,
)
del df3["_partitions"]
return df3
@_deprecated_kwarg("shuffle", "shuffle_method")
def rearrange_by_column(
df,
col,
npartitions=None,
max_branch=None,
shuffle_method=None,
compute=None,
ignore_index=False,
):
shuffle_method = shuffle_method or get_default_shuffle_method()
# if the requested output partitions < input partitions
# we repartition first as shuffling overhead is
# proportionate to the number of input partitions
if (
shuffle_method != "p2p"
and npartitions is not None
and npartitions < df.npartitions
):
df = df.repartition(npartitions=npartitions)
if shuffle_method == "disk":
return rearrange_by_column_disk(df, col, npartitions, compute=compute)
elif shuffle_method == "tasks":
df2 = rearrange_by_column_tasks(
df, col, max_branch, npartitions, ignore_index=ignore_index
)
if ignore_index:
df2._meta = df2._meta.reset_index(drop=True)
return df2
elif shuffle_method == "p2p":
from distributed.shuffle import rearrange_by_column_p2p
return rearrange_by_column_p2p(df, col, npartitions)
else:
raise NotImplementedError("Unknown shuffle method %s" % shuffle_method)
class maybe_buffered_partd:
"""
If serialized, will return non-buffered partd. Otherwise returns a buffered partd
"""
def __init__(self, encode_cls=None, buffer=True, tempdir=None):
self.tempdir = tempdir or config.get("temporary_directory", None)
self.buffer = buffer
self.compression = config.get("dataframe.shuffle.compression", None)
self.encode_cls = encode_cls
if encode_cls is None:
import partd
self.encode_cls = partd.PandasBlocks
def __reduce__(self):
if self.tempdir:
return (maybe_buffered_partd, (self.encode_cls, False, self.tempdir))
else:
return (maybe_buffered_partd, (self.encode_cls, False))
def __call__(self, *args, **kwargs):
import partd
path = tempfile.mkdtemp(suffix=".partd", dir=self.tempdir)
try:
partd_compression = (
getattr(partd.compressed, self.compression)
if self.compression
else None
)
except AttributeError as e:
raise ImportError(
"Not able to import and load {} as compression algorithm."
"Please check if the library is installed and supported by Partd.".format(
self.compression
)
) from e
file = partd.File(path)
partd.file.cleanup_files.append(path)
# Envelope partd file with compression, if set and available
if partd_compression:
file = partd_compression(file)
if self.buffer:
return self.encode_cls(partd.Buffer(partd.Dict(), file))
else:
return self.encode_cls(file)
def rearrange_by_column_disk(df, column, npartitions=None, compute=False):
"""Shuffle using local disk
See Also
--------
rearrange_by_column_tasks:
Same function, but using tasks rather than partd
Has a more informative docstring
"""
if npartitions is None:
npartitions = df.npartitions
token = tokenize(df, column, npartitions)
always_new_token = uuid.uuid1().hex
p = ("zpartd-" + always_new_token,)
encode_cls = partd_encode_dispatch(df._meta)
dsk1 = {p: (maybe_buffered_partd(encode_cls=encode_cls),)}
# Partition data on disk
name = "shuffle-partition-" + always_new_token
dsk2 = {
(name, i): (shuffle_group_3, key, column, npartitions, p)
for i, key in enumerate(df.__dask_keys__())
}
dependencies = []
if compute:
graph = HighLevelGraph.merge(df.dask, dsk1, dsk2)
graph = HighLevelGraph.from_collections(name, graph, dependencies=[df])
keys = [p, sorted(dsk2)]
pp, values = compute_as_if_collection(DataFrame, graph, keys)
dsk1 = {p: pp}
dsk2 = dict(zip(sorted(dsk2), values))
else:
dependencies.append(df)
# Barrier
barrier_token = "barrier-" + always_new_token
dsk3 = {barrier_token: (barrier, list(dsk2))}
# Collect groups
name = "shuffle-collect-" + token
dsk4 = {
(name, i): (collect, p, i, df._meta, barrier_token) for i in range(npartitions)
}
divisions = (None,) * (npartitions + 1)
layer = toolz.merge(dsk1, dsk2, dsk3, dsk4)
graph = HighLevelGraph.from_collections(name, layer, dependencies=dependencies)
return new_dd_object(graph, name, df._meta, divisions)
def _noop(x, cleanup_token):
"""
A task that does nothing.
"""
return x
def rearrange_by_column_tasks(
df, column, max_branch=32, npartitions=None, ignore_index=False
):
"""Order divisions of DataFrame so that all values within column(s) 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.
Lets explain the motivation for this further. Imagine that we have 1000
input partitions and 1000 output partitions. In theory we could split each
input into 1000 pieces, and then move the 1 000 000 resulting pieces
around, and then concatenate them all into 1000 output groups. This would
be fine, but the central scheduling overhead of 1 000 000 tasks would
become a bottleneck. Instead we do this in stages so that we split each of
the 1000 inputs into 30 pieces (we now have 30 000 pieces) move those
around, concatenate back down to 1000, and then do the same process again.
This has the same result as the full transfer, but now we've moved data
twice (expensive) but done so with only 60 000 tasks (cheap).
Note that the `column` input may correspond to a list of columns (rather
than just a single column name). In this case, the `shuffle_group` and
`shuffle_group_2` functions will use hashing to map each row to an output
partition. This approach may require the same rows to be hased multiple
times, but avoids the need to assign a new "_partitions" column.
Parameters
----------
df: dask.dataframe.DataFrame
column: str or list
A column name on which we want to split, commonly ``"_partitions"``
which is assigned by functions upstream. This could also be a list of
columns (in which case shuffle_group will create a hash array/column).
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
if (npartitions or df.npartitions) <= max_branch:
# We are creating a small number of output partitions.
# No need for staged shuffling. Staged shuffling will
# sometimes require extra work/communication in this case.
token = tokenize(df, column, npartitions)
shuffle_name = f"simple-shuffle-{token}"
npartitions = npartitions or df.npartitions
shuffle_layer = SimpleShuffleLayer(
shuffle_name,
column,
npartitions,
df.npartitions,
ignore_index,
df._name,
df._meta,
)
graph = HighLevelGraph.from_collections(
shuffle_name, shuffle_layer, dependencies=[df]
)
return new_dd_object(graph, shuffle_name, df._meta, [None] * (npartitions + 1))
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
inputs = [tuple(digit(i, j, k) for j in range(stages)) for i in range(k**stages)]
npartitions_orig = df.npartitions
token = tokenize(df, stages, column, n, k)
for stage in range(stages):
stage_name = f"shuffle-{stage}-{token}"
stage_layer = ShuffleLayer(
stage_name,
column,
inputs,
stage,
npartitions,
n,
k,
ignore_index,
df._name,
df._meta,
)
graph = HighLevelGraph.from_collections(
stage_name, stage_layer, dependencies=[df]
)
df = new_dd_object(graph, stage_name, df._meta, df.divisions)
if npartitions is not None and npartitions != npartitions_orig:
token = tokenize(df, npartitions)
repartition_group_token = "repartition-group-" + token
dsk = {
(repartition_group_token, i): (
shuffle_group_2,
k,
column,
ignore_index,
npartitions,
)
for i, k in enumerate(df.__dask_keys__())
}
repartition_get_name = "repartition-get-" + token
for p in range(npartitions):
dsk[(repartition_get_name, p)] = (
shuffle_group_get,
(repartition_group_token, p % npartitions_orig),
p,
)
graph2 = HighLevelGraph.from_collections(
repartition_get_name, dsk, dependencies=[df]
)
df2 = new_dd_object(
graph2, repartition_get_name, df._meta, [None] * (npartitions + 1)
)
else:
df2 = df
df2.divisions = (None,) * (npartitions_orig + 1)
return df2
########################################################
# Various convenience functions to be run by the above #
########################################################
def partitioning_index(df, npartitions, cast_dtype=None):
"""
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.
cast_dtype : dtype, optional
The dtype to cast to to avoid nullability issues
Returns
-------
partitions : ndarray
An array of int64 values mapping each record to a partition.
"""
if cast_dtype is not None:
# Fixme: astype raises with strings in numeric columns, but raising
# here might be very noisy
df = df.astype(cast_dtype, errors="ignore")
res = hash_object_dispatch(df, index=False) % int(npartitions)
# Note: Use a signed integer since pandas is more efficient at handling
# this since there is not always a fastpath for uints
return res.astype(np.min_scalar_type(-(npartitions - 1)))
def barrier(args):
list(args)
return 0
def cleanup_partd_files(p, keys):
"""
Cleanup the files in a partd.File dataset.
Parameters
----------
p : partd.Interface
File or Encode wrapping a file should be OK.
keys: List
Just for scheduling purposes, not actually used.
"""
import partd
if isinstance(p, partd.Encode):
maybe_file = p.partd
else:
maybe_file = None
if isinstance(maybe_file, partd.File):
path = maybe_file.path
else:
path = None
if path:
shutil.rmtree(path, ignore_errors=True)
def collect(p, part, meta, barrier_token):
"""Collect partitions from partd, yield dataframes"""
with ensure_cleanup_on_exception(p):
res = p.get(part)
return res if len(res) > 0 else meta
def set_partitions_pre(s, divisions, ascending=True, na_position="last"):
try:
if ascending:
partitions = divisions.searchsorted(s, side="right") - 1
else:
partitions = len(divisions) - divisions.searchsorted(s, side="right") - 1
except (TypeError, ValueError):
# `searchsorted` fails if either `divisions` or `s` contains nulls and strings
partitions = np.empty(len(s), dtype="int32")
not_null = s.notna()
divisions_notna = divisions[divisions.notna()]
if ascending:
partitions[not_null] = (
divisions_notna.searchsorted(s[not_null], side="right") - 1
)
else:
partitions[not_null] = (
len(divisions)
- divisions_notna.searchsorted(s[not_null], side="right")
- 1
)
partitions[(partitions < 0) | (partitions >= len(divisions) - 1)] = (
len(divisions) - 2 if ascending else 0
)
nas = s.isna()
# We could be a ndarray already (datetime dtype)
nas = getattr(nas, "values", nas)
partitions[nas] = len(divisions) - 2 if na_position == "last" else 0
return partitions
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_get(g_head, i):
g, head = g_head
if i in g:
return g[i]
else:
return head
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)
typ = np.min_scalar_type(npartitions * 2)
# Here we convert the final output index `ind` into the output index
# for the current stage.
kwargs = {} if PANDAS_GE_300 else {"copy": False}
ind = (ind % npartitions).astype(typ, **kwargs) // k**stage % k
return group_split_dispatch(df, ind, k, ignore_index=ignore_index)
@contextlib.contextmanager
def ensure_cleanup_on_exception(p):
"""Ensure a partd.File is cleaned up.
We have several tasks referring to a `partd.File` instance. We want to
ensure that the file is cleaned up if and only if there's an exception
in the tasks using the `partd.File`.
"""
try:
yield
except Exception:
# the function (e.g. shuffle_group_3) had an internal exception.
# We'll cleanup our temporary files and re-raise.
try:
p.drop()
except Exception:
logger.exception("ignoring exception in ensure_cleanup_on_exception")
raise
def shuffle_group_3(df, col, npartitions, p):
with ensure_cleanup_on_exception(p):
g = df.groupby(col)
d = {i: g.get_group(i) for i in g.groups}
p.append(d, fsync=True)
def set_index_post_scalar(df, index_name, drop, column_dtype):
df2 = df.drop("_partitions", axis=1).set_index(index_name, drop=drop)
df2.columns = df2.columns.astype(column_dtype)
return df2
def set_index_post_series(df, index_name, drop, column_dtype):
df2 = df.drop("_partitions", axis=1).set_index("_index", drop=True)
df2.index.name = index_name
df2.columns = df2.columns.astype(column_dtype)
return df2
def drop_overlap(df, index):
return df.drop(index) if index in df.index else df
def get_overlap(df, index):
return df.loc[[index]] if index in df.index else df._constructor()
def fix_overlap(ddf, mins, maxes, lens):
"""Ensures that the upper bound on each partition of ddf (except the last) is exclusive
This is accomplished by first removing empty partitions, then altering existing
partitions as needed to include all the values for a particular index value in
one partition.
"""
name = "fix-overlap-" + tokenize(ddf, mins, maxes, lens)
non_empties = [i for i, length in enumerate(lens) if length != 0]
# If all empty, collapse into one partition
if len(non_empties) == 0:
divisions = (None, None)
dsk = {(name, 0): (ddf._name, 0)}
graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
return new_dd_object(graph, name, ddf._meta, divisions)
# drop empty partitions by mapping each partition in a new graph to a particular
# partition on the old graph.
dsk = {(name, i): (ddf._name, div) for i, div in enumerate(non_empties)}
ddf_keys = list(dsk.values())
divisions = tuple(mins) + (maxes[-1],)
overlap = [i for i in range(1, len(mins)) if mins[i] >= maxes[i - 1]]
frames = []
for i in overlap:
# `frames` is a list of data from previous partitions that we may want to
# move to partition i. Here, we add "overlap" from the previous partition
# (i-1) to this list.
frames.append((get_overlap, ddf_keys[i - 1], divisions[i]))
# Make sure that any data added from partition i-1 to `frames` is removed
# from partition i-1.
dsk[(name, i - 1)] = (drop_overlap, dsk[(name, i - 1)], divisions[i])
# We do not want to move "overlap" from the previous partition (i-1) into
# this partition (i) if the data from this partition will need to be moved
# to the next partition (i+1) anyway. If we concatenate data too early,
# we may lose rows (https://github.com/dask/dask/issues/6972).
if divisions[i] == divisions[i + 1] and i + 1 in overlap:
continue
frames.append(ddf_keys[i])
dsk[(name, i)] = (methods.concat, frames)
frames = []
graph = HighLevelGraph.from_collections(name, dsk, dependencies=[ddf])
return new_dd_object(graph, name, ddf._meta, divisions)
def _compute_partition_stats(
column: Series, allow_overlap: bool = False, **kwargs
) -> tuple[list, list, list[int]]:
"""For a given column, compute the min, max, and len of each partition.
And make sure that the partitions are sorted relative to each other.
NOTE: this does not guarantee that every partition is internally sorted.
"""
mins = column.map_partitions(M.min, meta=column)
maxes = column.map_partitions(M.max, meta=column)
lens = column.map_partitions(len, meta=column)
mins, maxes, lens = compute(mins, maxes, lens, **kwargs)
mins = mins.bfill().tolist()
maxes = maxes.bfill().tolist()
non_empty_mins = [m for m, length in zip(mins, lens) if length != 0]
non_empty_maxes = [m for m, length in zip(maxes, lens) if length != 0]
if (
sorted(non_empty_mins) != non_empty_mins
or sorted(non_empty_maxes) != non_empty_maxes
):
raise ValueError(
f"Partitions are not sorted ascending by {column.name or 'the index'}",
f"In your dataset the (min, max, len) values of {column.name or 'the index'} "
f"for each partition are : {list(zip(mins, maxes, lens))}",
)
if not allow_overlap and any(
a <= b for a, b in zip(non_empty_mins[1:], non_empty_maxes[:-1])
):
warnings.warn(
"Partitions have overlapping values, so divisions are non-unique."
"Use `set_index(sorted=True)` with no `divisions` to allow dask to fix the overlap. "
f"In your dataset the (min, max, len) values of {column.name or 'the index'} "
f"for each partition are : {list(zip(mins, maxes, lens))}",
UserWarning,
)
lens = methods.tolist(lens)
if not allow_overlap:
return (mins, maxes, lens)
else:
return (non_empty_mins, non_empty_maxes, lens)
def compute_divisions(df: DataFrame, col: Any | None = None, **kwargs) -> tuple:
column = df.index if col is None else df[col]
mins, maxes, _ = _compute_partition_stats(column, allow_overlap=False, **kwargs)
return tuple(mins) + (maxes[-1],)
def compute_and_set_divisions(df: DataFrame, **kwargs) -> DataFrame:
mins, maxes, lens = _compute_partition_stats(df.index, allow_overlap=True, **kwargs)
if len(mins) == len(df.divisions) - 1:
df._divisions = tuple(mins) + (maxes[-1],)
if not any(mins[i] >= maxes[i - 1] for i in range(1, len(mins))):
return df
return fix_overlap(df, mins, maxes, lens)
def set_sorted_index(
df: DataFrame,
index: str | Series,
drop: bool = True,
divisions: Sequence | None = None,
**kwargs,
) -> DataFrame:
if isinstance(index, Series):
meta = df._meta.set_index(index._meta, drop=drop)
else:
meta = df._meta.set_index(index, drop=drop)
result = map_partitions(
M.set_index,
df,
index,
drop=drop,
meta=meta,
align_dataframes=False,
transform_divisions=False,
)
if not divisions:
return compute_and_set_divisions(result, **kwargs)
elif len(divisions) != len(df.divisions):
msg = (
"When doing `df.set_index(col, sorted=True, divisions=...)`, "
"divisions indicates known splits in the index column. In this "
"case divisions must be the same length as the existing "
"divisions in `df`\n\n"
"If the intent is to repartition into new divisions after "
"setting the index, you probably want:\n\n"
"`df.set_index(col, sorted=True).repartition(divisions=divisions)`"
)
raise ValueError(msg)
result.divisions = tuple(divisions)
return result