File: C:/Users/fred/anaconda3/Lib/site-packages/dask_expr/_reductions.py
import functools
from collections.abc import Callable
import numpy as np
import pandas as pd
import toolz
from dask.array import chunk
from dask.array.reductions import moment_agg, moment_chunk, moment_combine, nannumel
from dask.dataframe import hyperloglog, methods
from dask.dataframe.core import (
_concat,
_cov_corr_agg,
_cov_corr_chunk,
_cov_corr_combine,
_mode_aggregate,
idxmaxmin_agg,
idxmaxmin_chunk,
idxmaxmin_combine,
is_dataframe_like,
is_index_like,
is_series_like,
make_meta,
meta_nonempty,
total_mem_usage,
)
from dask.typing import no_default
from dask.utils import M, apply, funcname
from dask_expr._concat import Concat
from dask_expr._expr import (
Blockwise,
Expr,
Index,
Projection,
RenameFrame,
RenameSeries,
ResetIndex,
ToFrame,
determine_column_projection,
plain_column_projection,
)
from dask_expr._util import _tokenize_deterministic, is_scalar
class Chunk(Blockwise):
"""Partition-wise component of `ApplyConcatApply`
This class is used within `ApplyConcatApply._lower`.
See Also
--------
ApplyConcatApply
"""
_parameters = ["frame", "kind", "chunk", "chunk_kwargs"]
@property
def operation(self):
return self.chunk
@functools.cached_property
def _args(self) -> list:
return [self.frame]
@functools.cached_property
def _kwargs(self) -> dict:
return self.chunk_kwargs or {}
def _tree_repr_lines(self, indent=0, recursive=True):
header = f"{funcname(self.kind)}({funcname(type(self))}):"
lines = []
if recursive:
for dep in self.dependencies():
lines.extend(dep._tree_repr_lines(2))
for k, v in self._kwargs.items():
try:
if v != self.kind._defaults[k]:
header += f" {k}={v}"
except KeyError:
header += f" {k}={v}"
lines = [header] + lines
lines = [" " * indent + line for line in lines]
return lines
class Aggregate(Chunk):
"""Partition-wise aggregation component of `ApplyConcatApply`
This class is used within `ApplyConcatApply._lower`.
See Also
--------
ApplyConcatApply
"""
_parameters = ["frame", "kind", "aggregate", "aggregate_kwargs"]
@functools.cached_property
def aggregate_args(self):
return self.operands[len(self._parameters) :]
@staticmethod
def _call_with_list_arg(func, *args, **kwargs):
return func(list(args), **kwargs)
@property
def operation(self):
return functools.partial(self._call_with_list_arg, self.aggregate)
@functools.cached_property
def _args(self) -> list:
args = [self.frame]
if self.aggregate_args is not None:
args.extend(self.aggregate_args)
return args
@functools.cached_property
def _kwargs(self) -> dict:
return self.aggregate_kwargs or {}
class ShuffleReduce(Expr):
"""Shuffle-reduction component of `ApplyConcatApply`
when `split_out > 1`
This class is used within `ApplyConcatApply._lower`.
See Also
--------
ApplyConcatApply
"""
_parameters = [
"frame",
"kind",
"_meta",
"combine",
"aggregate",
"combine_kwargs",
"aggregate_args",
"aggregate_kwargs",
"split_by",
"split_every",
"split_out",
"sort",
"shuffle_by_index",
"shuffle_method",
"ignore_index",
]
_defaults = {
"split_every": 8,
"split_out": True,
"sort": None,
"shuffle_by_index": None,
"shuffle_method": None,
"ignore_index": True,
}
@property
def split_out(self):
if "split_out" in self._parameters:
split_out = self.operand("split_out")
if split_out is True:
return self.frame.npartitions
return split_out
else:
return 1
def _lower(self):
from dask_expr._repartition import Repartition
from dask_expr._shuffle import RearrangeByColumn, SetIndexBlockwise, SortValues
if is_index_like(self.frame._meta):
columns = [
self.frame._meta.name
if self.frame._meta.name is not None
else "__index__"
]
elif is_series_like(self.frame._meta):
columns = [
self.frame._meta.name
if self.frame._meta.name is not None
else "__series__"
]
else:
columns = self.frame.columns
# Find what columns we are shuffling by
split_by = self.split_by or columns
if not isinstance(split_by, (list, tuple)):
split_by = [split_by]
split_by_index = bool(set(split_by) - set(columns))
# Make sure we have dataframe-like data to shuffle
if split_by_index:
if self.frame.ndim == 1:
chunked = ResetIndex(self.frame, drop=False, name=self.frame.name)
else:
chunked = ResetIndex(self.frame, drop=False)
if split_by == [None]:
split_by = ["index"]
elif is_index_like(self.frame._meta) or is_series_like(self.frame._meta):
chunked = ToFrame(self.frame, name=columns[0])
else:
chunked = self.frame
# Map Tuple[str] column names to str before the shuffle
map_columns = {col: str(col) for col in chunked.columns if col != str(col)}
unmap_columns = {v: k for k, v in map_columns.items()}
if map_columns:
chunked = RenameFrame(chunked, map_columns)
split_by = [c if c not in map_columns else map_columns[c] for c in split_by]
# Sort or shuffle
split_every = getattr(self, "split_every", 0) or chunked.npartitions
ignore_index = getattr(self, "ignore_index", True)
if self.shuffle_by_index is not None:
ignore_index = not self.shuffle_by_index
shuffle_npartitions = max(
chunked.npartitions // split_every,
self.split_out,
)
if self.sort:
shuffled = SortValues(
chunked,
split_by,
npartitions=shuffle_npartitions,
ignore_index=ignore_index,
)
else:
shuffled = RearrangeByColumn(
chunked,
split_by,
shuffle_npartitions,
ignore_index=ignore_index,
index_shuffle=not split_by_index and self.shuffle_by_index,
method=self.shuffle_method,
)
# Unmap column names if necessary
if unmap_columns:
shuffled = RenameFrame(shuffled, unmap_columns)
# Reset the index if we we used it for shuffling
if split_by_index:
shuffled = SetIndexBlockwise(shuffled, split_by, True, None)
# Convert back to Series if necessary
if self.shuffle_by_index is not False:
if is_series_like(self._meta) and is_series_like(self.frame._meta):
shuffled = shuffled[shuffled.columns[0]]
if shuffled.name == "__series__":
shuffled = RenameSeries(shuffled, self.frame._meta.name)
elif is_index_like(self._meta):
column = shuffled.columns[0]
divs = None if shuffled.divisions[0] is None else shuffled.divisions
shuffled = Index(SetIndexBlockwise(shuffled, column, True, divs))
if column == "__index__":
shuffled = RenameSeries(shuffled, self.frame._meta.name)
# Blockwise aggregate
result = Aggregate(
shuffled,
self.kind,
self.aggregate,
self.aggregate_kwargs,
*self.aggregate_args,
)
# Repartition and return
if self.split_out < result.npartitions:
return Repartition(result, new_partitions=self.split_out)
return result
@property
def _meta(self):
return self.operand("_meta")
def _divisions(self):
return (None,) * (self.split_out + 1)
def __str__(self):
chunked = str(self.frame)
split_every = getattr(self, "split_every", 0)
return f"{type(self).__name__}({chunked}, kind={funcname(self.kind)}, split_every={split_every})"
class TreeReduce(Expr):
"""Tree-reduction component of `ApplyConcatApply`
This class is used within `ApplyConcatApply._lower`.
See Also
--------
ApplyConcatApply
"""
_parameters = [
"frame",
"kind",
"_meta",
"combine",
"aggregate",
"combine_kwargs",
"aggregate_kwargs",
"split_every",
]
_defaults = {"split_every": 8}
@functools.cached_property
def _name(self):
if funcname(self.combine) in ("combine", "aggregate"):
name = funcname(self.combine.__self__).lower() + "-tree"
else:
name = funcname(self.combine)
return name + "-" + _tokenize_deterministic(*self.operands)
def __dask_postcompute__(self):
return toolz.first, ()
@functools.cached_property
def split_every(self):
out = self.operand("split_every")
if out is None:
return 8
if out is False or isinstance(out, int) and out >= 2:
return out
raise ValueError("split_every must be greater than 1 or False")
def _layer(self):
# apply combine to batches of intermediate results
j = 1
d = {}
keys = self.frame.__dask_keys__()
split_every = self.split_every
while split_every is not False and len(keys) > split_every:
new_keys = []
for i, batch in enumerate(
toolz.partition_all(split_every or len(keys), keys)
):
batch = list(batch)
if self.combine_kwargs:
d[self._name, j, i] = (
apply,
self.combine,
[batch],
self.combine_kwargs,
)
else:
d[self._name, j, i] = (self.combine, batch)
new_keys.append((self._name, j, i))
j += 1
keys = new_keys
# apply aggregate to the final result
d[self._name, 0] = (apply, self.aggregate, [keys], self.aggregate_kwargs)
return d
@property
def _meta(self):
return self.operand("_meta")
def _divisions(self):
return (None, None)
def __str__(self):
chunked = str(self.frame)
split_every = getattr(self, "split_every", 0)
return f"{type(self).__name__}({chunked}, kind={funcname(self.kind)}, split_every={split_every})"
def _tree_repr_lines(self, indent=0, recursive=True):
header = f"{funcname(self.kind)}({funcname(type(self))}):"
lines = []
if recursive:
for dep in self.dependencies():
lines.extend(dep._tree_repr_lines(2))
split_every = getattr(self, "split_every", 0)
header += f" split_every={split_every}"
lines = [header] + lines
lines = [" " * indent + line for line in lines]
return lines
class ApplyConcatApply(Expr):
"""Perform reduction-like operation on dataframes
This pattern is commonly used for reductions, groupby-aggregations, and
more. It requires three methods to be implemented:
- `chunk`: applied to each input partition
- `combine`: applied to lists of intermediate partitions as they are
combined in batches
- `aggregate`: applied at the end to finalize the computation
These methods should be easy to serialize, and can take in keyword
arguments defined in `chunks/combine/aggregate_kwargs`.
In many cases people don't define all three functions. In these cases
combine takes from aggregate and aggregate takes from chunk.
"""
_parameters = ["frame"]
chunk = None
combine = None
aggregate = None
chunk_kwargs = {}
combine_kwargs = {}
aggregate_args = []
aggregate_kwargs = {}
_chunk_cls = Chunk
@property
def split_out(self):
if "split_out" in self._parameters:
split_out = self.operand("split_out")
if isinstance(split_out, Callable):
split_out = split_out(self.frame.npartitions)
if split_out is None:
raise ValueError("split_out can't be None")
return split_out
else:
return 1
def _layer(self):
# This is an abstract expression
raise NotImplementedError()
@functools.cached_property
def _meta_chunk(self):
meta = meta_nonempty(self.frame._meta)
return self.chunk(meta, **self.chunk_kwargs)
@functools.cached_property
def _meta(self):
meta = self._meta_chunk
aggregate = self.aggregate or (lambda x: x)
if self.combine:
combine = self.combine
combine_kwargs = self.combine_kwargs
else:
combine = aggregate
combine_kwargs = self.aggregate_kwargs
meta = combine([meta], **combine_kwargs)
meta = aggregate([meta], **self.aggregate_kwargs)
return make_meta(meta)
def _divisions(self):
if getattr(self, "sort", False):
return (None, None)
if self.split_out is True:
return (None,) * (self.frame.npartitions + 1)
return (None,) * (self.split_out + 1)
@property
def _chunk_cls_args(self):
return []
@property
def should_shuffle(self):
sort = getattr(self, "sort", False)
return not (
not isinstance(self.split_out, bool) and self.split_out == 1 or sort
)
@functools.cached_property
def need_to_shuffle(self):
split_by = self.split_by or self.frame.columns
if any(
set(split_by) >= (set(cols) if isinstance(cols, tuple) else {cols})
for cols in self.frame.unique_partition_mapping_columns_from_shuffle
):
return False
return True
@functools.cached_property
def unique_partition_mapping_columns_from_shuffle(self):
if self.should_shuffle and not self.need_to_shuffle:
return self.frame.unique_partition_mapping_columns_from_shuffle
elif self.should_shuffle:
return (
{self.split_by}
if not isinstance(self.split_by, list)
else {tuple(self.split_by)}
)
else:
return set()
def _lower(self):
# Normalize functions in case not all are defined
chunk = self.chunk
chunk_kwargs = self.chunk_kwargs
if self.aggregate:
aggregate = self.aggregate
aggregate_kwargs = self.aggregate_kwargs
else:
aggregate = chunk
aggregate_kwargs = chunk_kwargs
if self.combine:
combine = self.combine
combine_kwargs = self.combine_kwargs
else:
combine = aggregate
combine_kwargs = aggregate_kwargs
split_every = getattr(self, "split_every", None)
chunked = self._chunk_cls(
self.frame, type(self), chunk, chunk_kwargs, *self._chunk_cls_args
)
if not self.should_shuffle:
# Lower into TreeReduce(Chunk)
return TreeReduce(
chunked,
type(self),
self._meta,
combine,
aggregate,
combine_kwargs,
aggregate_kwargs,
split_every=split_every,
)
elif not self.need_to_shuffle:
# Repartition and return
result = Aggregate(
chunked,
type(self),
aggregate,
aggregate_kwargs,
*self.aggregate_args,
)
if self.split_out is not True and self.split_out < result.npartitions:
from dask_expr import Repartition
return Repartition(result, new_partitions=self.split_out)
if self.ndim < result.ndim:
result = result[result.columns[0]]
return result
# Lower into ShuffleReduce
return ShuffleReduce(
chunked,
type(self),
self._meta,
combine,
aggregate,
combine_kwargs,
self.aggregate_args,
aggregate_kwargs,
split_by=self.split_by,
split_out=self.split_out,
split_every=split_every,
sort=getattr(self, "sort", False),
shuffle_by_index=getattr(self, "shuffle_by_index", None),
shuffle_method=getattr(self, "shuffle_method", None),
ignore_index=getattr(self, "ignore_index", True),
)
class Unique(ApplyConcatApply):
_parameters = ["frame", "split_every", "split_out", "shuffle_method"]
_defaults = {"split_every": None, "split_out": True, "shuffle_method": "tasks"}
chunk = staticmethod(methods.unique)
aggregate_func = staticmethod(methods.unique)
@functools.cached_property
def _meta(self):
return self.chunk(
meta_nonempty(self.frame._meta), series_name=self.frame._meta.name
)
@property
def split_by(self):
return self.name
@property
def chunk_kwargs(self):
return {"series_name": self._meta.name}
@property
def aggregate_kwargs(self):
return self.chunk_kwargs
@classmethod
def combine(cls, inputs: list, **kwargs):
return _concat(inputs)
@classmethod
def aggregate(cls, inputs: list, **kwargs):
df = _concat(inputs)
return cls.aggregate_func(df, **kwargs)
class DropDuplicates(Unique):
_parameters = [
"frame",
"subset",
"ignore_index",
"split_every",
"split_out",
"shuffle_method",
"keep",
]
_defaults = {
"subset": None,
"ignore_index": False,
"split_every": None,
"split_out": True,
"shuffle_method": "tasks",
"keep": "first",
}
chunk = M.drop_duplicates
aggregate_func = M.drop_duplicates
@property
def split_by(self):
return self.subset
@functools.cached_property
def _meta(self):
return make_meta(
self.chunk(meta_nonempty(self.frame._meta), **self.chunk_kwargs)
)
@property
def chunk_kwargs(self):
out = {"keep": self.keep}
if is_dataframe_like(self.frame._meta):
out["subset"] = self.subset
if not is_index_like(self.frame._meta):
out["ignore_index"] = self.ignore_index
return out
def _simplify_up(self, parent, dependents):
if self.subset is not None and isinstance(parent, Projection):
columns = determine_column_projection(
self, parent, dependents, additional_columns=self.subset
)
if set(columns) == set(self.frame.columns):
# Don't add unnecessary Projections, protects against loops
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.operands[1:],
)
class PivotTable(ApplyConcatApply):
_parameters = ["frame", "columns", "index", "values", "aggfunc"]
_defaults = {"columns": None, "index": None, "values": None, "aggfunc": "mean"}
@functools.cached_property
def _meta(self):
df = self.frame._meta
columns = self.operand("columns")
values = self.operand("values")
index = self.operand("index")
columns_contents = pd.CategoricalIndex(df[columns].cat.categories, name=columns)
if is_scalar(values):
new_columns = columns_contents
else:
new_columns = pd.MultiIndex.from_product(
(sorted(values), columns_contents), names=[None, columns]
)
if self.operand("aggfunc") in ["first", "last"]:
# Infer datatype as non-numeric values are allowed
if is_scalar(values):
meta = pd.DataFrame(
columns=new_columns,
dtype=df[values].dtype,
index=pd.Index(df[index]),
)
else:
meta = pd.DataFrame(
columns=new_columns,
index=pd.Index(df[index]),
)
for value_col in values:
meta[value_col] = meta[value_col].astype(
df[values].dtypes[value_col]
)
else:
# Use float64 as other aggregate functions require numerical data
meta = pd.DataFrame(
columns=new_columns, dtype=np.float64, index=pd.Index(df[index])
)
return meta
def _lower(self):
args = [
self.frame,
self.operand("columns"),
self.operand("index"),
self.operand("values"),
]
if self.aggfunc == "sum":
return PivotTableSum(*args)
elif self.aggfunc == "mean":
return PivotTableSum(*args) / PivotTableCount(*args)
elif self.aggfunc == "count":
return PivotTableCount(*args)
elif self.aggfunc == "first":
return PivotTableFirst(*args)
elif self.aggfunc == "last":
return PivotTableLast(*args)
else:
raise NotImplementedError(f"{self.aggfunc=} is not implemented")
class PivotTableAbstract(ApplyConcatApply):
_parameters = ["frame", "columns", "index", "values", "aggfunc"]
_defaults = {"columns": None, "index": None, "values": None, "aggfunc": "mean"}
@property
def chunk_kwargs(self):
return {
"index": self.operand("index"),
"columns": self.operand("columns"),
"values": self.operand("values"),
}
@classmethod
def combine(cls, inputs: list, **kwargs):
return _concat(inputs)
@classmethod
def aggregate(cls, inputs: list, **kwargs):
df = _concat(inputs)
return cls.aggregate_func(df, **kwargs)
class PivotTableSum(PivotTableAbstract):
chunk = staticmethod(methods.pivot_sum)
aggregate_func = staticmethod(methods.pivot_agg)
class PivotTableCount(PivotTableAbstract):
chunk = staticmethod(methods.pivot_count)
aggregate_func = staticmethod(methods.pivot_agg)
class PivotTableFirst(PivotTableAbstract):
chunk = staticmethod(methods.pivot_first)
aggregate_func = staticmethod(methods.pivot_agg_first)
class PivotTableLast(PivotTableAbstract):
chunk = staticmethod(methods.pivot_last)
aggregate_func = staticmethod(methods.pivot_agg_last)
class Reduction(ApplyConcatApply):
"""A common pattern of apply concat apply
Common reductions like sum/min/max/count/... have some shared code around
`_concat` and so on. This class inherits from `ApplyConcatApply` in order
to leverage this shared structure.
I wouldn't be surprised if there was a way to merge them both into a single
abstraction in the future.
This class implements `{chunk,combine,aggregate}` methods of
`ApplyConcatApply` by depending on `reduction_{chunk,combine,aggregate}`
methods.
"""
_defaults = {
"skipna": True,
"numeric_only": False,
"min_count": 0,
"dropna": True,
}
reduction_chunk = None
reduction_combine = None
reduction_aggregate = None
@property
def _projection_columns(self):
return self.frame.columns
@classmethod
def chunk(cls, df, **kwargs):
out = cls.reduction_chunk(df, **kwargs)
# Return a dataframe so that the concatenated version is also a dataframe
return out.to_frame().T if is_series_like(out) else out
@classmethod
def combine(cls, inputs: list, **kwargs):
func = cls.reduction_combine or cls.reduction_aggregate or cls.reduction_chunk
df = _concat(inputs)
out = func(df, **kwargs)
# Return a dataframe so that the concatenated version is also a dataframe
return out.to_frame().T if is_series_like(out) else out
@classmethod
def aggregate(cls, inputs, **kwargs):
func = cls.reduction_aggregate or cls.reduction_chunk
df = _concat(inputs)
return func(df, **kwargs)
def __dask_postcompute__(self):
return toolz.first, ()
def _divisions(self):
if self.ndim == 0 or len(self.frame.columns) == 0:
return (None, None)
return (min(self.frame.columns), max(self.frame.columns))
def __str__(self):
params = {param: self.operand(param) for param in self._parameters[1:]}
s = ", ".join(
k + "=" + repr(v)
for k, v in params.items()
if v is not self._defaults.get(k)
)
base = str(self.frame)
if " " in base:
base = "(" + base + ")"
return f"{base}.{self.__class__.__name__.lower()}({s})"
def _simplify_up(self, parent, dependents):
if isinstance(parent, Projection):
return plain_column_projection(self, parent, dependents)
class CustomReduction(Reduction):
_parameters = [
"frame",
"meta",
"chunk_kwargs",
"aggregate_kwargs",
"combine_kwargs",
"split_every",
"token",
]
@functools.cached_property
def _name(self):
name = self.operand("token") or funcname(type(self)).lower()
return name + "-" + _tokenize_deterministic(*self.operands)
@classmethod
def chunk(cls, df, **kwargs):
func = kwargs.pop("func")
out = func(df, **kwargs)
# Return a dataframe so that the concatenated version is also a dataframe
return out.to_frame().T if is_series_like(out) else out
@classmethod
def combine(cls, inputs: list, **kwargs):
func = kwargs.pop("func")
df = _concat(inputs)
out = func(df, **kwargs)
# Return a dataframe so that the concatenated version is also a dataframe
return out.to_frame().T if is_series_like(out) else out
@classmethod
def aggregate(cls, inputs, **kwargs):
func = kwargs.pop("func")
df = _concat(inputs)
return func(df, **kwargs)
@functools.cached_property
def _meta(self):
if self.operand("meta") is not no_default:
return self.operand("meta")
return super()._meta
@property
def chunk_kwargs(self):
return self.operand("chunk_kwargs")
@property
def combine_kwargs(self):
return self.operand("combine_kwargs")
@property
def aggregate_kwargs(self):
return self.operand("aggregate_kwargs")
def _simplify_up(self, parent, dependents):
return
def _divisions(self):
return (None, None)
class Sum(Reduction):
_parameters = ["frame", "skipna", "numeric_only", "split_every", "axis"]
_defaults = {
"split_every": False,
"numeric_only": False,
"skipna": True,
"axis": 0,
}
reduction_chunk = M.sum
@property
def chunk_kwargs(self):
return dict(skipna=self.skipna, numeric_only=self.numeric_only, axis=self.axis)
@property
def combine_kwargs(self):
return dict(skipna=self.skipna, axis=self.axis)
@property
def aggregate_kwargs(self):
return dict(skipna=self.skipna, axis=self.axis)
class Prod(Sum):
reduction_chunk = M.prod
class Max(Reduction):
_parameters = ["frame", "skipna", "numeric_only", "split_every", "axis"]
_defaults = {
"split_every": False,
"numeric_only": False,
"skipna": True,
"axis": 0,
}
reduction_chunk = M.max
@property
def chunk_kwargs(self):
if self.frame._meta.ndim < 2:
return dict(skipna=self.skipna)
else:
return dict(
skipna=self.skipna, numeric_only=self.numeric_only, axis=self.axis
)
@property
def combine_kwargs(self):
return dict(skipna=self.skipna, axis=self.axis)
@property
def aggregate_kwargs(self):
return dict(skipna=self.skipna, axis=self.axis)
class Min(Max):
reduction_chunk = M.min
class Any(Reduction):
_parameters = ["frame", "skipna", "split_every"]
_defaults = {"skipna": True, "split_every": False}
reduction_chunk = M.any
@property
def chunk_kwargs(self):
return dict(
skipna=self.skipna,
)
class All(Reduction):
_parameters = ["frame", "skipna", "split_every"]
_defaults = {"split_every": False}
reduction_chunk = M.all
@property
def chunk_kwargs(self):
return dict(
skipna=self.skipna,
)
class IdxMin(Reduction):
_parameters = ["frame", "skipna", "numeric_only", "split_every"]
_defaults = {"skipna": True, "numeric_only": False, "split_every": False}
reduction_chunk = idxmaxmin_chunk
reduction_combine = idxmaxmin_combine
reduction_aggregate = idxmaxmin_agg
_reduction_attribute = "idxmin"
@property
def chunk_kwargs(self):
return dict(
skipna=self.skipna,
fn=self._reduction_attribute,
numeric_only=self.numeric_only,
)
@property
def combine_kwargs(self):
return dict(skipna=self.skipna, fn=self._reduction_attribute)
@property
def aggregate_kwargs(self):
return {**self.chunk_kwargs, "scalar": is_series_like(self.frame._meta)}
class IdxMax(IdxMin):
_reduction_attribute = "idxmax"
class Cov(Reduction):
_parameters = ["frame", "min_periods", "split_every", "scalar"]
_defaults = {"min_periods": 2, "split_every": False, "scalar": False}
reduction_chunk = staticmethod(_cov_corr_chunk)
reduction_combine = staticmethod(_cov_corr_combine)
reduction_aggregate = staticmethod(_cov_corr_agg)
corr = False
@property
def chunk_kwargs(self):
return {"corr": self.corr}
@property
def combine_kwargs(self):
return self.chunk_kwargs
@property
def aggregate_kwargs(self):
return {
**self.chunk_kwargs,
"scalar": self.scalar,
"like_df": self.frame._meta,
"cols": self.frame.columns,
}
class Corr(Cov):
corr = True
class Len(Reduction):
reduction_chunk = staticmethod(len)
reduction_aggregate = sum
def _simplify_down(self):
from dask_expr.io.io import IO
# We introduce Index nodes sometimes. We special case around them.
if isinstance(self.frame, Index) and self.frame.frame._is_length_preserving:
return Len(self.frame.frame)
# Pass through Elemwises, unless we just introduced an Index
if self.frame._is_length_preserving and not isinstance(self.frame, Index):
child = max(self.frame.dependencies(), key=lambda expr: expr.npartitions)
return Len(child)
# Let the child handle it. They often know best
if isinstance(self.frame, IO):
return self
if isinstance(self.frame, Concat) and self.frame.operand("axis") == 0:
return sum(Len(obj) for obj in self.frame.dependencies())
# Drop all of the columns, just pass through the index
if self.frame.ndim == 2 and len(self.frame.columns):
return Len(self.frame.index)
def _simplify_up(self, parent, dependents):
return
class Size(Reduction):
reduction_aggregate = sum
@staticmethod
def reduction_chunk(df):
return df.size
def _simplify_down(self):
if is_dataframe_like(self.frame._meta) and len(self.frame.columns) > 1:
return len(self.frame.columns) * Len(self.frame)
else:
return Len(self.frame)
def _simplify_up(self, parent, dependents):
return
class NBytes(Reduction):
# Only supported for Series objects
reduction_aggregate = sum
@staticmethod
def reduction_chunk(ser):
return ser.nbytes
class ArrayReduction(Reduction):
@classmethod
def chunk(cls, df, **kwargs):
return cls.reduction_chunk(df, **kwargs)
@classmethod
def combine(cls, inputs: list, **kwargs):
func = cls.reduction_combine or cls.reduction_aggregate or cls.reduction_chunk
return func(inputs, **kwargs)
@classmethod
def aggregate(cls, inputs, meta, index, **kwargs):
func = cls.reduction_aggregate or cls.reduction_chunk
result = func(inputs, **kwargs)
if is_series_like(meta):
return type(meta)(result, name=meta.name, index=index)
else:
return result
class Var(ArrayReduction):
# Uses the parallel version of Welford's online algorithm (Chan 79')
# (http://i.stanford.edu/pub/cstr/reports/cs/tr/79/773/CS-TR-79-773.pdf)
_parameters = ["frame", "skipna", "ddof", "numeric_only", "split_every"]
_defaults = {"skipna": True, "ddof": 1, "numeric_only": False, "split_every": False}
@functools.cached_property
def _meta(self):
return make_meta(
meta_nonempty(self.frame._meta).var(
skipna=self.skipna, numeric_only=self.numeric_only
)
)
@property
def chunk_kwargs(self):
return dict(skipna=self.skipna)
@property
def combine_kwargs(self):
return {"skipna": self.skipna}
@property
def aggregate_kwargs(self):
cols = self.frame.columns if self.frame.ndim == 1 else self.frame._meta.columns
return dict(
ddof=self.ddof,
skipna=self.skipna,
meta=self._meta,
index=cols,
)
@classmethod
def reduction_chunk(cls, x, skipna):
values = x.values.astype("f8")
if skipna:
return moment_chunk(
values, sum=chunk.nansum, numel=nannumel, keepdims=True, axis=(0,)
)
else:
return moment_chunk(values, keepdims=True, axis=(0,))
@classmethod
def reduction_combine(cls, parts, skipna):
if skipna:
return moment_combine(parts, sum=np.nansum, axis=(0,))
else:
return moment_combine(parts, axis=(0,))
@classmethod
def reduction_aggregate(cls, vals, ddof, skipna):
if skipna:
result = moment_agg(vals, sum=np.nansum, ddof=ddof, axis=(0,))
else:
result = moment_agg(vals, ddof=ddof, axis=(0,))
return result
class Moment(ArrayReduction):
_parameters = ["frame", "order"]
@functools.cached_property
def _meta(self):
# Use var as proxy for result dimension
return make_meta(meta_nonempty(self.frame._meta).var())
@property
def chunk_kwargs(self):
return dict(order=self.order)
@property
def combine_kwargs(self):
return self.chunk_kwargs
@property
def aggregate_kwargs(self):
return dict(
order=self.order,
meta=self._meta,
index=self.frame.columns,
)
@classmethod
def reduction_chunk(cls, x, order):
values = x.values.astype("f8")
return moment_chunk(values, order=order, axis=(0,), keepdims=True)
@classmethod
def reduction_combine(cls, parts, order):
return moment_combine(parts, order=order, axis=(0,))
@classmethod
def reduction_aggregate(cls, vals, order):
result = moment_agg(vals, order=order, axis=(0,))
return result
class Mean(Reduction):
_parameters = ["frame", "skipna", "numeric_only", "split_every", "axis"]
_defaults = {"skipna": True, "numeric_only": False, "split_every": False, "axis": 0}
@functools.cached_property
def _meta(self):
return make_meta(
meta_nonempty(self.frame._meta).mean(
skipna=self.skipna, numeric_only=self.numeric_only, axis=self.axis
)
)
def _lower(self):
s = self.frame.sum(
skipna=self.skipna,
numeric_only=self.numeric_only,
split_every=self.split_every,
)
c = self.frame.count(
split_every=self.split_every, numeric_only=self.numeric_only
)
if self.axis is None and s.ndim == 1:
return s.sum() / c.sum()
else:
return MeanAggregate(s, c)
class MeanAggregate(Blockwise):
_parameters = ["frame", "counter"]
operation = staticmethod(methods.mean_aggregate)
class Count(Reduction):
_parameters = ["frame", "numeric_only", "split_every"]
_defaults = {"split_every": False, "numeric_only": False}
reduction_chunk = M.count
@classmethod
def reduction_aggregate(cls, df):
return df.sum().astype("int64")
@property
def chunk_kwargs(self):
if self.frame._meta.ndim < 2:
return dict()
else:
return dict(numeric_only=self.numeric_only)
class IndexCount(Reduction):
_parameters = ["frame", "split_every"]
_defaults = {"split_every": False}
reduction_chunk = staticmethod(methods.index_count)
reduction_aggregate = staticmethod(np.sum)
# aggregate_chunk = staticmethod(np.sum)
class Mode(Reduction):
"""
Mode was a bit more complicated than class reductions, so we retreat back
to ApplyConcatApply
"""
_parameters = ["frame", "dropna", "split_every"]
_defaults = {"dropna": True, "split_every": False}
reduction_chunk = M.value_counts
reduction_combine = M.sum
reduction_aggregate = staticmethod(_mode_aggregate)
def _divisions(self):
return self.frame.divisions[0], self.frame.divisions[-1]
@property
def chunk_kwargs(self):
return {"dropna": self.dropna}
@property
def aggregate_kwargs(self):
return {"dropna": self.dropna}
class NuniqueApprox(Reduction):
_parameters = ["frame", "b", "split_every"]
_defaults = {"b": 16, "split_every": None}
reduction_chunk = hyperloglog.compute_hll_array
reduction_combine = hyperloglog.reduce_state
reduction_aggregate = hyperloglog.estimate_count
@functools.cached_property
def _meta(self):
return 1.0
@property
def chunk_kwargs(self):
return {"b": self.b}
@property
def combine_kwargs(self):
return self.chunk_kwargs
@property
def aggregate_kwargs(self):
return self.chunk_kwargs
class ReductionConstantDim(Reduction):
"""
Some reductions reduce the number of rows in your object but keep the original
dimension, e.g. a DataFrame stays a DataFrame instead of getting reduced to
a Series.
"""
@classmethod
def chunk(cls, df, **kwargs):
return cls.reduction_chunk(df, **kwargs)
@classmethod
def combine(cls, inputs: list, **kwargs):
func = cls.reduction_combine or cls.reduction_aggregate or cls.reduction_chunk
df = _concat(inputs)
return func(df, **kwargs)
def _divisions(self):
# TODO: We can do better in some cases
return (None, None)
class NLargest(ReductionConstantDim):
_parameters = ["frame", "n", "_columns", "split_every"]
_defaults = {"n": 5, "_columns": None, "split_every": None}
reduction_chunk = M.nlargest
reduction_aggregate = M.nlargest
def _columns_kwarg(self):
if self._columns is None:
return {}
return {"columns": self._columns}
@property
def chunk_kwargs(self):
return {"n": self.n, **self._columns_kwarg()}
@property
def combine_kwargs(self):
return self.chunk_kwargs
@property
def aggregate_kwargs(self):
return self.chunk_kwargs
def _nfirst(df, columns, n, ascending):
return df.sort_values(by=columns, ascending=ascending).head(n)
def _nlast(df, columns, n, ascending):
return df.sort_values(by=columns, ascending=ascending).tail(n)
class NFirst(NLargest):
_parameters = ["frame", "n", "_columns", "ascending", "split_every"]
_defaults = {"n": 5, "_columns": None, "ascending": None, "split_every": None}
reduction_chunk = staticmethod(_nfirst)
reduction_aggregate = staticmethod(_nfirst)
@property
def chunk_kwargs(self):
return {"ascending": self.ascending, **super().chunk_kwargs}
class NLast(NFirst):
reduction_chunk = staticmethod(_nlast)
reduction_aggregate = staticmethod(_nlast)
class NSmallest(NLargest):
reduction_chunk = M.nsmallest
reduction_aggregate = M.nsmallest
class ValueCounts(ReductionConstantDim):
_defaults = {
"sort": None,
"ascending": False,
"dropna": True,
"normalize": False,
"split_every": None,
"split_out": 1,
"total_length": None,
}
_parameters = [
"frame",
"sort",
"ascending",
"dropna",
"normalize",
"split_every",
"split_out",
"total_length",
]
reduction_chunk = M.value_counts
reduction_aggregate = methods.value_counts_aggregate
reduction_combine = methods.value_counts_combine
@functools.cached_property
def _meta(self):
return self.frame._meta.value_counts(normalize=self.normalize)
@classmethod
def aggregate(cls, inputs, **kwargs):
func = cls.reduction_aggregate or cls.reduction_chunk
if is_scalar(inputs[-1]):
return func(_concat(inputs[:-1]), inputs[-1], observed=True, **kwargs)
else:
return func(_concat(inputs), observed=True, **kwargs)
@property
def split_by(self):
return self.frame._meta.name
@property
def chunk_kwargs(self):
return {"sort": self.sort, "ascending": self.ascending, "dropna": self.dropna}
@property
def aggregate_args(self):
if self.normalize and (self.split_out > 1 or self.split_out is True):
return [self.total_length]
return []
@property
def aggregate_kwargs(self):
return {**self.chunk_kwargs, "normalize": self.normalize}
@property
def combine_kwargs(self):
return self.chunk_kwargs
def _simplify_up(self, parent, dependents):
# We are already a Series
return
@functools.cached_property
def split_by(self):
return self.frame._meta.name
def _divisions(self):
if self.sort:
return (None, None)
if self.split_out is True:
return (None,) * (self.frame.npartitions + 1)
return (None,) * (self.split_out + 1)
class MemoryUsage(Reduction):
reduction_chunk = M.memory_usage
reduction_aggregate = M.sum
split_every = False
def _divisions(self):
# TODO: We can do better, but not high priority
return (None, None)
class MemoryUsageIndex(MemoryUsage):
_parameters = ["frame", "deep"]
_defaults = {"deep": False}
@property
def chunk_kwargs(self):
return {"deep": self.deep}
class MemoryUsageFrame(MemoryUsage):
_parameters = ["frame", "deep", "_index"]
_defaults = {"deep": False, "_index": True}
@property
def chunk_kwargs(self):
return {"deep": self.deep, "index": self._index}
@property
def combine_kwargs(self):
return {"is_dataframe": is_dataframe_like(self.frame._meta)}
@staticmethod
def reduction_combine(x, is_dataframe):
if is_dataframe:
return x.groupby(x.index).sum()
return x.sum()
class TotalMemoryUsageFrame(MemoryUsageFrame):
reduction_chunk = total_mem_usage
@staticmethod
def reduction_combine(x, is_dataframe):
return x
@staticmethod
def reduction_aggregate(x):
return x
class IsMonotonicIncreasing(Reduction):
@functools.cached_property
def _meta(self):
return make_meta(bool)
reduction_chunk = methods.monotonic_increasing_chunk
reduction_combine = methods.monotonic_increasing_combine
reduction_aggregate = methods.monotonic_increasing_aggregate
class IsMonotonicDecreasing(IsMonotonicIncreasing):
reduction_chunk = methods.monotonic_decreasing_chunk
reduction_combine = methods.monotonic_decreasing_combine
reduction_aggregate = methods.monotonic_decreasing_aggregate