from __future__ import annotations import functools import math import operator import uuid import numpy as np import pandas as pd import tlz as toolz from pandas import CategoricalDtype from dask import compute from dask._task_spec import Task, TaskRef from dask.dataframe.core import _concat from dask.dataframe.dask_expr._expr import ( Assign, Blockwise, Expr, Filter, PartitionsFiltered, Projection, ToSeriesIndex, determine_column_projection, is_filter_pushdown_available, ) from dask.dataframe.dask_expr._reductions import ( All, Any, Count, DropDuplicates, Len, Max, Mean, MemoryUsage, Min, Mode, NBytes, NFirst, NLargest, NLast, NSmallest, Prod, Size, Sum, Unique, ValueCounts, ) from dask.dataframe.dask_expr._repartition import Repartition, RepartitionToFewer from dask.dataframe.dask_expr._util import LRU, _convert_to_list from dask.dataframe.dispatch import is_categorical_dtype, make_meta from dask.dataframe.shuffle import ( barrier, collect, ensure_cleanup_on_exception, maybe_buffered_partd, partitioning_index, set_partitions_pre, shuffle_group, shuffle_group_2, shuffle_group_get, ) from dask.utils import ( M, digit, get_default_shuffle_method, insert, is_index_like, is_series_like, ) class ShuffleBase(Expr): _parameters = [ "frame", "partitioning_index", "npartitions_out", "ignore_index", "method", "options", "index_shuffle", "original_partitioning_index", ] _defaults = { "ignore_index": False, "method": None, "options": None, "index_shuffle": None, "original_partitioning_index": None, } _is_length_preserving = True _filter_passthrough = True def __str__(self): return f"Shuffle({self._name[-7:]})" def _node_label_args(self): return [self.frame, self.partitioning_index] @functools.cached_property def _partitioning_index(self): partitioning_index = self.partitioning_index if isinstance(partitioning_index, (str, int)): partitioning_index = [partitioning_index] return partitioning_index @functools.cached_property def unique_partition_mapping_columns_from_shuffle(self): idx = self.original_partitioning_index or self._partitioning_index return {tuple(idx)} if isinstance(idx, list) else set() def _simplify_up(self, parent, dependents): if isinstance(parent, Filter) and self._filter_passthrough_available( parent, dependents ): return self._filter_simplification(parent) if isinstance(parent, Projection): # Move the column projection to come # before the abstract Shuffle projection = _convert_to_list( determine_column_projection(self, parent, dependents) ) partitioning_index = self._partitioning_index target = self.frame new_projection = [ col for col in target.columns if (col in partitioning_index or col in projection) ] if set(new_projection) < set(target.columns): return type(self)(target[new_projection], *self.operands[1:])[ parent.operand("columns") ] if isinstance( parent, ( Unique, DropDuplicates, Sum, Prod, Max, Any, All, Min, Len, Size, NBytes, Mean, Count, Mode, NLargest, NSmallest, ValueCounts, MemoryUsage, ), ): return type(parent)(self.frame, *parent.operands[1:]) def _layer(self): raise NotImplementedError( f"{self} is abstract! Please call `simplify`" f"before generating a task graph." ) @functools.cached_property def _meta(self): meta = self.frame._meta if self.ignore_index and self.method == "tasks": meta = meta.reset_index(drop=True) return meta def _divisions(self): return (None,) * (self.npartitions_out + 1) class Shuffle(ShuffleBase): """Abstract shuffle class Parameters ---------- frame: Expr The DataFrame-like expression to shuffle. partitioning_index: str, list Column and/or index names to hash and partition by. npartitions: int Number of output partitions. ignore_index: bool Whether to ignore the index during this shuffle operation. method: str or Callable Label or callback function to convert a shuffle operation to its necessary components. options: dict Algorithm-specific options. index_shuffle : bool Whether to perform the shuffle on the index. """ def _lower(self): # Use `method` to decide how to compose a # shuffle operation from concerete expressions # Reduce partition count if necessary frame = self.frame npartitions_out = self.npartitions_out method = self.method or get_default_shuffle_method() if npartitions_out < frame.npartitions and method != "p2p": frame = Repartition(frame, new_partitions=npartitions_out) ops = [ self.partitioning_index, self.npartitions_out, self.ignore_index, self.options, self.original_partitioning_index, ] if method == "p2p": return P2PShuffle(frame, *ops) elif method == "disk": return DiskShuffle(frame, *ops) elif method == "simple": return SimpleShuffle(frame, *ops) elif method == "tasks": return TaskShuffle(frame, *ops) else: raise ValueError(f"{method} not supported") def _is_numeric_cast_type(dtype): return ( pd.api.types.is_numeric_dtype(dtype) or isinstance(dtype, CategoricalDtype) and pd.api.types.is_numeric_dtype(dtype.categories) ) class RearrangeByColumn(ShuffleBase): def _lower(self): frame = self.frame partitioning_index = self.partitioning_index npartitions_out = self.npartitions_out ignore_index = self.ignore_index options = self.options index_shuffle = self.index_shuffle # Normalize partitioning_index if isinstance(partitioning_index, str): partitioning_index = [partitioning_index] if index_shuffle: pass elif not isinstance(partitioning_index, (list, Expr)): raise ValueError( f"{type(partitioning_index)} not a supported type for partitioning_index" ) if not isinstance(partitioning_index, Expr) and not index_shuffle: cs = [col for col in partitioning_index if col not in frame.columns] if len(cs) == 1: frame = Assign(frame, "_partitions_0", frame.index) partitioning_index = partitioning_index.copy() partitioning_index[partitioning_index.index(cs[0])] = "_partitions_0" # Assign new "_partitions" column index_added = AssignPartitioningIndex( frame, partitioning_index, "_partitions", npartitions_out, frame._meta, index_shuffle, ) # Apply shuffle shuffled = Shuffle( index_added, "_partitions", npartitions_out, ignore_index, self.method, options, original_partitioning_index=self._partitioning_index, ) if frame.ndim == 1: # Reduce back to series return shuffled[index_added.columns[0]] # Drop "_partitions" column and return return shuffled[ [c for c in shuffled.columns if c not in ["_partitions", "_partitions_0"]] ] class SimpleShuffle(PartitionsFiltered, Shuffle): _parameters = [ "frame", "partitioning_index", "npartitions_out", "ignore_index", "options", "original_partitioning_index", "_partitions", ] _defaults = { "_partitions": None, "original_partitioning_index": None, "partitioning_index": "_partitions", "ignore_index": False, "options": None, } @functools.cached_property def _meta(self): return self.frame._meta @staticmethod def _shuffle_group(df, _filter, *args): """Filter the output of `shuffle_group`""" if _filter is None: return shuffle_group(df, *args) return {k: v for k, v in shuffle_group(df, *args).items() if k in _filter} def _layer(self): """Construct graph for a simple shuffle operation.""" shuffle_group_name = "group-" + self._name split_name = "split-" + self._name npartitions = self.npartitions_out dsk = {} _filter = self._partitions if self._filtered else None for global_part, part_out in enumerate(self._partitions): _concat_list = [ (split_name, part_out, part_in) for part_in in range(self.frame.npartitions) ] dsk[(self._name, global_part)] = ( _concat, _concat_list, self.ignore_index, ) for _, _part_out, _part_in in _concat_list: dsk[(split_name, _part_out, _part_in)] = ( operator.getitem, (shuffle_group_name, _part_in), _part_out, ) if (shuffle_group_name, _part_in) not in dsk: dsk[(shuffle_group_name, _part_in)] = ( self._shuffle_group, (self.frame._name, _part_in), _filter, self.partitioning_index, 0, npartitions, npartitions, self.ignore_index, npartitions, ) return dsk def _lower(self): return None class TaskShuffle(SimpleShuffle): """Staged task-based shuffle implementation""" @functools.cached_property def _meta(self): meta = self.frame._meta if self.ignore_index: meta = meta.reset_index(drop=True) return meta def _layer(self): max_branch = (self.options or {}).get("max_branch", None) or 32 npartitions_input = self.frame.npartitions if len(self._partitions) <= max_branch or npartitions_input <= max_branch: # We are creating a small number of output partitions, # or starting with a small number of input partitions. # No need for staged shuffling. Staged shuffling will # sometimes require extra work/communication in this case. return super()._layer() # Calculate number of stages and splits per stage npartitions = self.npartitions_out stages = int(math.ceil(math.log(npartitions_input) / math.log(max_branch))) if stages > 1: nsplits = int(math.ceil(npartitions_input ** (1 / stages))) else: nsplits = npartitions_input # Construct global data-movement plan inputs = [ tuple(digit(i, j, nsplits) for j in range(stages)) for i in range(nsplits**stages) ] inp_part_map = {inp: i for i, inp in enumerate(inputs)} parts_out = range(len(inputs)) # Build graph dsk = {} name = self.frame._name meta_input = make_meta(self.frame._meta) for stage in range(stages): # Define names name_input = name if stage == (stages - 1) and npartitions == npartitions_input: name = self._name parts_out = self._partitions _filter = parts_out if self._filtered else None else: name = f"stage-{stage}-{self._name}" _filter = None shuffle_group_name = "group-" + name split_name = "split-" + name for global_part, part in enumerate(parts_out): out = inputs[part] _concat_list = [] # get_item tasks to concat for this output partition for i in range(nsplits): # Get out each individual dataframe piece from the dicts _inp = insert(out, stage, i) _idx = out[stage] _concat_list.append((split_name, _idx, _inp)) # concatenate those pieces together, with their friends dsk[(name, global_part)] = ( _concat, _concat_list, self.ignore_index, ) for _, _idx, _inp in _concat_list: dsk[(split_name, _idx, _inp)] = ( operator.getitem, (shuffle_group_name, _inp), _idx, ) if (shuffle_group_name, _inp) not in dsk: # Initial partitions (output of previous stage) _part = inp_part_map[_inp] if stage == 0: if _part < npartitions_input: input_key = (name_input, _part) else: # In order to make sure that to_serialize() serialize the # empty dataframe input, we add it as a key. input_key = (shuffle_group_name, _inp, "empty") dsk[input_key] = meta_input else: input_key = (name_input, _part) # Convert partition into dict of dataframe pieces dsk[(shuffle_group_name, _inp)] = ( self._shuffle_group, input_key, _filter, self.partitioning_index, stage, nsplits, npartitions_input, self.ignore_index, npartitions, ) if npartitions != npartitions_input: repartition_group_name = "repartition-group-" + name dsk2 = { (repartition_group_name, i): ( shuffle_group_2, (name, i), self.partitioning_index, self.ignore_index, npartitions, ) for i in range(npartitions_input) } for i, p in enumerate(self._partitions): dsk2[(self._name, i)] = ( shuffle_group_get, (repartition_group_name, p % npartitions_input), p, ) dsk.update(dsk2) return dsk class DiskShuffle(SimpleShuffle): """Disk-based shuffle implementation""" @staticmethod def _shuffle_group(df, col, _filter, p): with ensure_cleanup_on_exception(p): g = df.groupby(col) d = {i: g.get_group(i) for i in g.groups if i in _filter} p.append(d, fsync=True) def _layer(self): from dask.dataframe.dispatch import partd_encode_dispatch column = self.partitioning_index df = self.frame 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): (self._shuffle_group, key, column, self._partitions, p) for i, key in enumerate(df.__dask_keys__()) } # Barrier barrier_token = ("barrier-" + always_new_token,) dsk3 = {barrier_token: (barrier, list(dsk2))} # Collect groups dsk4 = { (self._name, j): (collect, p, k, df._meta, barrier_token) for j, k in enumerate(self._partitions) } return toolz.merge(dsk1, dsk2, dsk3, dsk4) def _shuffle_transfer( input: pd.DataFrame, id, input_partition: int, ) -> int: from distributed.shuffle._shuffle import shuffle_transfer return shuffle_transfer( input, id, input_partition, ) class P2PShuffle(SimpleShuffle): """P2P worker-based shuffle implementation""" @functools.cached_property def _meta(self): return self.frame._meta.drop(columns=self.partitioning_index) def _layer(self): from distributed.shuffle._core import ( P2PBarrierTask, ShuffleId, barrier_key, p2p_barrier, ) from distributed.shuffle._shuffle import DataFrameShuffleSpec, shuffle_unpack dsk = {} token = self._name.split("-")[-1] shuffle_id = ShuffleId(token) _barrier_key = barrier_key(shuffle_id) name = "shuffle-transfer-" + token parts_out = ( self._partitions if self._filtered else list(range(self.npartitions_out)) ) # Avoid embedding a materialized list unless necessary parts_out_arg = ( tuple(self._partitions) if self._filtered else self.npartitions_out ) transfer_keys = list() for i in range(self.frame.npartitions): t = Task( (name, i), _shuffle_transfer, TaskRef((self.frame._name, i)), token, i, ) dsk[t.key] = t transfer_keys.append(t.ref()) barrier = P2PBarrierTask( _barrier_key, p2p_barrier, token, *transfer_keys, spec=DataFrameShuffleSpec( id=shuffle_id, npartitions=self.npartitions_out, column=self.partitioning_index, meta=self.frame._meta, parts_out=parts_out_arg, disk=True, drop_column=True, ), ) dsk[barrier.key] = barrier # TODO: Decompose p2p Into transfer/barrier + unpack name = self._name for i, part_out in enumerate(parts_out): t = Task( (name, i), shuffle_unpack, token, part_out, barrier.ref(), ) dsk[t.key] = t return dsk # # Helper logic # def _select_columns_or_index(df, columns_or_index): """ Make a column selection that may include the index Parameters ---------- columns_or_index Column or index name, or a list of these """ # Ensure columns_or_index is a list columns_or_index = ( columns_or_index if isinstance(columns_or_index, list) else [columns_or_index] ) column_names = [n for n in columns_or_index if _is_column_label_reference(df, n)] selected_df = df[column_names] if _contains_index_name(df, columns_or_index): # Index name was included selected_df = selected_df.assign(_index=df.index) return selected_df def _is_column_label_reference(df, key): """ Test whether a key is a column label reference To be considered a column label reference, `key` must match the name of at least one column. """ return ( not isinstance(key, Expr) and (np.isscalar(key) or pd.api.types.is_scalar(key) or isinstance(key, tuple)) and key in df.columns ) def _contains_index_name(df, columns_or_index): """ Test whether the input contains a reference to the index of the df """ if isinstance(columns_or_index, list): return any(_is_index_level_reference(df, n) for n in columns_or_index) else: return _is_index_level_reference(df, columns_or_index) def _is_index_level_reference(df, key): """ Test whether a key is an index level reference To be considered an index level reference, `key` must match the index name and must NOT match the name of any column. """ index_name = df.index._meta.name if isinstance(df, Expr) else df.index.name return ( index_name is not None and not isinstance(key, Expr) and (np.isscalar(key) or pd.api.types.is_scalar(key) or isinstance(key, tuple)) and key == index_name and key not in getattr(df, "columns", ()) ) class AssignPartitioningIndex(Blockwise): """Assign a partitioning index This class is used to construct a hash-based partitioning index for shuffling. Parameters ---------- frame: Expr Frame-like expression being partitioned. partitioning_index: Expr or list Index-like expression or list of columns to construct the partitioning-index from. index_name: str New column name to assign. npartitions_out: int Number of partitions after repartitioning is finished. index_shuffle : bool, default False Whether we are using solely the index for the shuffle """ _parameters = [ "frame", "partitioning_index", "index_name", "npartitions_out", "meta", "index_shuffle", ] _defaults = { "cast_dtype": None, "index_shuffle": False, "partitioning_index": "_partitions", "index_name": "_partitions", } _preserves_partitioning_information = True @staticmethod def operation(df, index, name: str, npartitions: int, meta, index_shuffle: bool): # type: ignore """Construct a hash-based partitioning index""" def _get_index(idx, obj): if hasattr(idx, "ndim"): if idx.ndim == 1: idx = idx.to_frame() elif index_shuffle: # set default index name, otherwise we will end up with 0 name = {"name": "_index"} if idx == ["_index"] else {} idx = obj.index.to_frame(**name) else: idx = _select_columns_or_index(obj, idx) return idx # meta and df dtypes can deviate, this is why we do the cast here meta_index = _get_index(index, meta) index = _get_index(index, df) dtypes = {} for col, dtype in meta_index.dtypes.items(): if _is_numeric_cast_type(dtype): dtypes[col] = np.float64 if dtypes: index = index.astype(dtypes, errors="ignore") index = partitioning_index(index, npartitions) if df.ndim == 1: df = df.to_frame() return df.assign(**{name: index}) class BaseSetIndexSortValues(Expr): _is_length_preserving = True def _divisions(self): if "user_divisions" in self._parameters and self.user_divisions is not None: return self.user_divisions if self._npartitions_input == 1: return (None, None) if ( is_index_like(self._divisions_column._meta) and self._divisions_column.known_divisions and self._divisions_column.npartitions == self.frame.npartitions ): return self.other.divisions divisions, mins, maxes, presorted = _get_divisions( self.frame, self._divisions_column, self._npartitions_input, self.ascending, upsample=self.upsample, ) if presorted and len(mins) == self._npartitions_input: divisions = mins.copy() + [maxes[-1]] return divisions @property def _npartitions_input(self): return self.operand("npartitions") or self.frame.npartitions @property def npartitions(self): return self.operand("npartitions") or len(self._divisions()) - 1 class SetIndex(BaseSetIndexSortValues): """Abstract ``set_index`` class. Simplifies (later lowers) either to Blockwise ops if we are already sorted or to ``SetPartition`` which handles shuffling. Parameters ---------- frame: Expr Frame-like expression where the index is set. _other: Expr | Scalar Either a Series-like expression to use as Index or a scalar defining the column. drop: bool Whether we drop the old column. sorted: str No need for shuffling if we are already sorted. user_divisions: int Divisions as passed by the user. upsample: float Used to increase the number of samples for quantiles. """ _parameters = [ "frame", "_other", "drop", "user_divisions", "partition_size", "ascending", "npartitions", "upsample", "shuffle_method", "append", "options", # Options for the chosen shuffle method ] _defaults = { "drop": True, "user_divisions": None, "partition_size": 128e6, "ascending": True, "npartitions": None, "upsample": 1.0, "shuffle_method": None, "options": None, "append": False, } _filter_passthrough = True @property def _projection_columns(self): return self.columns + ( [self._other] if not isinstance(self._other, Expr) else [] ) @functools.cached_property def _meta(self): if isinstance(self._other, Expr): other = self._other._meta else: other = self._other return self.frame._meta.set_index(other, drop=self.drop) @property def _divisions_column(self): return self.other @property def other(self): if isinstance(self._other, Expr): return self._other return self.frame[self._other] def _lower(self): if ( self.operand("npartitions") == 1 or self.frame.npartitions == 1 and (self.user_divisions is None or len(self.user_divisions) == 2) and self.operand("npartitions") is None ): expr = self.frame if self.frame.npartitions > 1: expr = RepartitionToFewer(expr, 1) index_set = SetIndexBlockwise(expr, self._other, self.drop, None) return SortIndexBlockwise(index_set) if self.user_divisions is None: divisions = self._divisions() if ( is_index_like(self._divisions_column._meta) and self.other.divisions == divisions ): presorted = True else: presorted = _get_divisions( self.frame, self.other, self._npartitions_input, self.ascending, upsample=self.upsample, )[3] if presorted and self.npartitions == self.frame.npartitions: index_set = SetIndexBlockwise( self.frame, self._other, self.drop, divisions, self.append ) return SortIndexBlockwise(index_set) return SetPartition( self.frame, self._other, self.drop, self._npartitions_input if self.user_divisions is None else None, self.ascending, self.upsample, self.user_divisions, self.shuffle_method, self.options, ) def _simplify_up(self, parent, dependents): from dask.dataframe.dask_expr._expr import Filter, Head, Tail # TODO, handle setting index with other frame if ( isinstance(parent, Head) and isinstance(self._other, (int, str)) and self._other in self.frame.columns ): head = NFirst(self.frame, n=parent.n, _columns=self._other, ascending=True) return SetIndex(head, _other=self._other) if ( isinstance(parent, Tail) and isinstance(self._other, (int, str)) and self._other in self.frame.columns ): tail = NLast(self.frame, n=parent.n, _columns=self._other, ascending=True) return SetIndex(tail, _other=self._other) if isinstance(parent, Projection): addition_columns = ( [self._other] if not isinstance(self._other, Expr) else [] ) columns = determine_column_projection( self, parent, dependents, additional_columns=addition_columns ) columns = _convert_to_list(columns) columns = [c for c in self.frame.columns if c in columns] if self.frame.columns == columns: return return type(parent)( type(self)(self.frame[columns], *self.operands[1:]), parent.operand("columns"), ) if isinstance(parent, Filter) and self._filter_passthrough_available( parent, dependents ): return self._filter_simplification(parent) def _filter_passthrough_available(self, parent, dependents): if is_filter_pushdown_available(self, parent, dependents): from dask.dataframe.dask_expr._expr import Index p = parent.predicate return not any(isinstance(x, Index) for x in p.walk()) return False class SortValues(BaseSetIndexSortValues): _parameters = [ "frame", "by", "ascending", "na_position", "npartitions", "partition_size", "sort_function", "sort_function_kwargs", "upsample", "ignore_index", "shuffle_method", "options", # Options for the chosen shuffle method ] _defaults = { "partition_size": 128e6, "ascending": True, "npartitions": None, "na_position": "last", "sort_function": None, "sort_function_kwargs": None, "upsample": 1.0, "ignore_index": False, "shuffle_method": None, } _filter_passthrough = True def _divisions(self): if self.frame.npartitions == 1: # Protect against triggering calculations when we only have one division return (None, None) divisions, mins, maxes, presorted = _get_divisions( self.frame, self.frame[self.by[0]], self._npartitions_input, self._divisions_ascending, upsample=self.upsample, ) if presorted: return self.frame.divisions return (None,) * len(divisions) @property def _divisions_ascending(self) -> bool: divisions_ascending = self.ascending if not isinstance(divisions_ascending, bool): divisions_ascending = divisions_ascending[0] assert isinstance(divisions_ascending, bool) return divisions_ascending @property def sort_function(self): if self.operand("sort_function") is not None: return self.operand("sort_function") return M.sort_values @property def sort_function_kwargs(self): sort_kwargs = { "by": self.by, "ascending": self.ascending, "na_position": self.na_position, "ignore_index": self.ignore_index, } if self.operand("sort_function_kwargs") is not None: sort_kwargs.update(self.operand("sort_function_kwargs")) return sort_kwargs @functools.cached_property def _meta(self): return self.frame._meta @functools.cached_property def _meta_by_dtype(self): dtype = self._meta.dtypes[self.by] if is_series_like(dtype): dtype = dtype.iloc[0] return dtype def _lower(self): if self.frame.npartitions == 1: return SortValuesBlockwise( self.frame, self.sort_function, self.sort_function_kwargs ) _divisions_by = self.frame[self.by[0]] divisions, _, _, presorted = _get_divisions( self.frame, _divisions_by, self._npartitions_input, self._divisions_ascending, upsample=self.upsample, ) if presorted and self.npartitions == self.frame.npartitions: return SortValuesBlockwise( self.frame, self.sort_function, self.sort_function_kwargs ) partitions = _SetPartitionsPreSetIndex( _divisions_by, _divisions_by._meta._constructor(divisions).sort_values(), ascending=self._divisions_ascending, ) assigned = Assign(self.frame, "_partitions", partitions) shuffled = Shuffle( assigned, "_partitions", npartitions_out=len(divisions) - 1, ignore_index=self.ignore_index, method=self.shuffle_method, options=self.options, ) shuffled = Projection(shuffled, self.frame.columns) return SortValuesBlockwise( shuffled, self.sort_function, self.sort_function_kwargs ) def _simplify_up(self, parent, dependents): from dask.dataframe.dask_expr._expr import Filter, Head, Tail if isinstance(parent, Head): return NFirst( self.frame, n=parent.n, _columns=self.by, ascending=self.ascending ) if isinstance(parent, Tail): return NLast( self.frame, n=parent.n, _columns=self.by, ascending=self.ascending ) if isinstance(parent, Filter) and self._filter_passthrough_available( parent, dependents ): return self._filter_simplification(parent) if isinstance(parent, Projection): columns = determine_column_projection( self, parent, dependents, additional_columns=self.by ) columns = _convert_to_list(columns) columns = [col for col in self.frame.columns if col in columns] if self.frame.columns == columns: return return type(parent)( type(self)(self.frame[columns], *self.operands[1:]), parent.operand("columns"), ) if ( isinstance(parent, Repartition) and parent.operand("new_partitions") is not None ): return type(self)( type(parent)(self.frame, *parent.operands[1:]), *self.operands[1:] ) class SetPartition(SetIndex): """Shuffles the DataFrame according to its new divisions. Simplifies the Expression to blockwise pre-processing, shuffle and blockwise post-processing expressions. Parameters ---------- frame: Expr Frame-like expression where the index is set. _other: Expr | Scalar Either a Series-like expression to use as Index or a scalar defining the column. drop: bool Whether to drop the old column. new_divisions: int Divisions of the resulting expression. """ _parameters = [ "frame", "_other", "drop", "npartitions", "ascending", "upsample", "user_divisions", "shuffle_method", "options", # Shuffle method options ] def _lower(self): divisions = self.other._meta._constructor(self._divisions()) partitions = _SetPartitionsPreSetIndex(self.other, divisions) assigned = Assign(self.frame, "_partitions", partitions) if isinstance(self._other, Expr): assigned = Assign(assigned, "_index", self._other) shuffled = Shuffle( assigned, "_partitions", npartitions_out=len(self._divisions()) - 1, ignore_index=True, method=self.shuffle_method, options=self.options, ) shuffled = Projection( shuffled, [c for c in assigned.columns if c != "_partitions"] ) if isinstance(self._other, Expr): drop, set_name = True, "_index" else: drop, set_name = self.drop, self.other._meta.name lru_key = ( self.other._name, self._npartitions_input, self.ascending, 128e6, self.upsample, ) computed_divisions = divisions_lru.get(lru_key) index_set = _SetIndexPost( shuffled, self.other._meta.name, drop, set_name, self.frame._meta.columns.dtype, computed_divisions, self.user_divisions, ) return SortIndexBlockwise(index_set) class _SetPartitionsPreSetIndex(Blockwise): _parameters = ["frame", "new_divisions", "ascending", "na_position"] _defaults = {"ascending": True, "na_position": "last"} operation = staticmethod(set_partitions_pre) _is_length_preserving = True @functools.cached_property def _meta(self): return make_meta(self.frame._meta._constructor([0])) class _SetIndexPost(Blockwise): _parameters = [ "frame", "index_name", "drop", "set_name", "column_dtype", "computed_divisions", "user_divisions", ] _is_length_preserving = True @property def _args(self) -> list: return self.operands[:5] @staticmethod def operation(df, index_name, drop, set_name, column_dtype): df = df.set_index(set_name, drop=drop) df.index.name = index_name df.columns = df.columns.astype(column_dtype) return df def _get_culled_divisions(self, divisions): if self.frame.npartitions < len(divisions) - 1: part_filter = list(self.frame.find_operations(PartitionsFiltered)) if len(part_filter) > 0: return tuple( [ div for i, div in enumerate(divisions) if i in part_filter[0]._partitions ] + [divisions[-1]] ) else: return self.frame.divisions return divisions def _divisions(self): if self.operand("user_divisions") is not None: return self._get_culled_divisions(self.operand("user_divisions")) assert self.computed_divisions is not None return self._get_culled_divisions(self.computed_divisions[0]) class SortIndexBlockwise(Blockwise): _projection_passthrough = True _parameters = ["frame"] operation = M.sort_index _is_length_preserving = True class SortValuesBlockwise(Blockwise): _projection_passthrough = False _parameters = ["frame", "sort_function", "sort_kwargs"] _keyword_only = ["sort_function", "sort_kwargs"] _is_length_preserving = True @staticmethod def operation(*args, **kwargs): sort_func = kwargs.pop("sort_function") sort_kwargs = kwargs.pop("sort_kwargs") return sort_func(*args, **kwargs, **sort_kwargs) @functools.cached_property def _meta(self): return self.frame._meta class SetIndexBlockwise(Blockwise): _parameters = ["frame", "other", "drop", "new_divisions", "append"] _defaults = {"append": False, "new_divisions": None, "drop": True} _keyword_only = ["drop", "new_divisions", "append"] _is_length_preserving = True _preserves_partitioning_information = True @staticmethod def operation(df, *args, new_divisions, **kwargs): return df.set_index(*args, **kwargs) def _divisions(self): if self.new_divisions is None: return (None,) * (self.frame.npartitions + 1) return tuple(self.new_divisions) def _simplify_up(self, parent, dependents): if isinstance(parent, Projection): columns = determine_column_projection( self, parent, dependents, additional_columns=_convert_to_list(self.other), ) if self.frame.columns == columns: return columns = [col for col in self.frame.columns if col in columns] return type(parent)( type(self)(self.frame[columns], *self.operands[1:]), parent.operand("columns"), ) divisions_lru = LRU(10) # type: ignore def _get_divisions( frame, other, npartitions: int, ascending: bool = True, partition_size: float = 128e6, upsample: float = 1.0, ): key = (other._name, npartitions, ascending, partition_size, upsample) if key in divisions_lru: return divisions_lru[key] result = _calculate_divisions( frame, other, npartitions, ascending, partition_size, upsample ) divisions_lru[key] = result return result def _calculate_divisions( frame, other, npartitions: int, ascending: bool = True, partition_size: float = 128e6, upsample: float = 1.0, ): from dask.dataframe.dask_expr import RepartitionQuantiles, new_collection if is_index_like(other._meta): other = ToSeriesIndex(other) if is_categorical_dtype(other._meta.dtype): other = new_collection(other).cat.as_ordered()._expr try: divisions, mins, maxes = compute( new_collection(RepartitionQuantiles(other, npartitions, upsample=upsample)), new_collection(other).map_partitions(M.min), new_collection(other).map_partitions(M.max), ) 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 pd.api.types.is_numeric_dtype(other._meta.dtype): obj, suggested_method = ( ("column", f"`.dropna(subset=['{other.name}'])`") if any(other._name == frame[c]._name for c in frame.columns) else ("series", "`.loc[series[~series.isna()]]`") ) raise NotImplementedError( f"Divisions calculation failed for non-numeric {obj} '{other.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 sizes = [] # type: ignore empty_dataframe_detected = pd.isna(divisions).all() if empty_dataframe_detected: total = sum(sizes) npartitions = max(math.ceil(total / partition_size), 1) npartitions = min(npartitions, frame.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(other._meta.dtype, pd.CategoricalDtype): dtype = other._meta.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