_expr.py

from __future__ import annotations

import datetime
import functools
import numbers
import operator
import warnings
from collections import defaultdict
from collections.abc import Callable, Mapping

import dask
import numpy as np
import pandas as pd
from dask.array import Array
from dask.base import normalize_token
from dask.core import flatten
from dask.dataframe import methods
from dask.dataframe._pyarrow import to_pyarrow_string
from dask.dataframe.core import (
    _concat,
    _get_divisions_map_partitions,
    _rename,
    apply_and_enforce,
    has_parallel_type,
    is_dataframe_like,
    is_index_like,
    is_series_like,
    make_meta,
    pd_split,
    safe_head,
    total_mem_usage,
)
from dask.dataframe.dispatch import meta_nonempty
from dask.dataframe.rolling import CombinedOutput, _head_timedelta, overlap_chunk
from dask.dataframe.shuffle import drop_overlap, get_overlap
from dask.dataframe.utils import (
    clear_known_categories,
    drop_by_shallow_copy,
    raise_on_meta_error,
    valid_divisions,
)
from dask.typing import no_default
from dask.utils import (
    M,
    apply,
    funcname,
    get_meta_library,
    has_keyword,
    is_arraylike,
    partial_by_order,
    pseudorandom,
    random_state_data,
)
from pandas.errors import PerformanceWarning
from tlz import merge, merge_sorted, partition, unique

from dask_expr import _core as core
from dask_expr._util import (
    _calc_maybe_new_divisions,
    _convert_to_list,
    _tokenize_deterministic,
    _tokenize_partial,
    is_scalar,
)


class Expr(core.Expr):
    """Primary class for all Expressions

    This mostly includes Dask protocols and various Pandas-like method
    definitions to make us look more like a DataFrame.
    """

    _is_length_preserving = False
    _filter_passthrough = False

    def _filter_passthrough_available(self, parent, dependents):
        return self._filter_passthrough and is_filter_pushdown_available(
            self, parent, dependents
        )

    @functools.cached_property
    def ndim(self):
        meta = self._meta
        try:
            return meta.ndim
        except AttributeError:
            return 0

    @functools.cached_property
    def _resources(self):
        dep_resources = merge(dep._resources or {} for dep in self.dependencies())
        return dep_resources or None

    def __dask_keys__(self):
        return [(self._name, i) for i in range(self.npartitions)]

    def optimize(self, **kwargs):
        return optimize(self, **kwargs)

    def __hash__(self):
        return hash(self._name)

    @property
    def index(self):
        return Index(self)

    @property
    def size(self):
        return Size(self)

    @property
    def nbytes(self):
        return NBytes(self)

    def __getitem__(self, other):
        if isinstance(other, Expr):
            if not are_co_aligned(self, other):
                return FilterAlign(self, other)
            return Filter(self, other)
        else:
            return Projection(self, other)  # df[["a", "b", "c"]]

    def __bool__(self):
        raise ValueError(
            f"The truth value of a {self.__class__.__name__} is ambiguous. "
            "Use a.any() or a.all()."
        )

    def __add__(self, other):
        return Add(self, other)

    def __radd__(self, other):
        return Add(other, self)

    def __sub__(self, other):
        return Sub(self, other)

    def __rsub__(self, other):
        return Sub(other, self)

    def __mul__(self, other):
        return Mul(self, other)

    def __rmul__(self, other):
        return Mul(other, self)

    def __pow__(self, power):
        return Pow(self, power)

    def __rpow__(self, power):
        return Pow(power, self)

    def __truediv__(self, other):
        return Div(self, other)

    def __rtruediv__(self, other):
        return Div(other, self)

    def __lt__(self, other):
        return LT(self, other)

    def __rlt__(self, other):
        return LT(other, self)

    def __gt__(self, other):
        return GT(self, other)

    def __rgt__(self, other):
        return GT(other, self)

    def __le__(self, other):
        return LE(self, other)

    def __rle__(self, other):
        return LE(other, self)

    def __ge__(self, other):
        return GE(self, other)

    def __rge__(self, other):
        return GE(other, self)

    def __eq__(self, other):
        return EQ(self, other)

    def __ne__(self, other):
        return NE(self, other)

    def __and__(self, other):
        return And(self, other)

    def __rand__(self, other):
        return And(other, self)

    def __or__(self, other):
        return Or(self, other)

    def __ror__(self, other):
        return Or(other, self)

    def __xor__(self, other):
        return XOr(self, other)

    def __rxor__(self, other):
        return XOr(other, self)

    def __invert__(self):
        return Invert(self)

    def __neg__(self):
        return Neg(self)

    def __pos__(self):
        return Pos(self)

    def __mod__(self, other):
        return Mod(self, other)

    def __rmod__(self, other):
        return Mod(other, self)

    def __floordiv__(self, other):
        return FloorDiv(self, other)

    def __rfloordiv__(self, other):
        return FloorDiv(other, self)

    def __divmod__(self, other):
        res1 = self // other
        res2 = self % other
        return res1, res2

    def __rdivmod__(self, other):
        res1 = other // self
        res2 = other % self
        return res1, res2

    def sum(self, skipna=True, numeric_only=False, split_every=False, axis=0):
        return Sum(self, skipna, numeric_only, split_every, axis)

    def prod(self, skipna=True, numeric_only=False, split_every=False, axis=0):
        return Prod(self, skipna, numeric_only, split_every, axis)

    def var(self, axis=0, skipna=True, ddof=1, numeric_only=False, split_every=False):
        if axis == 0:
            return Var(self, skipna, ddof, numeric_only, split_every)
        elif axis == 1:
            return VarColumns(self, skipna, ddof, numeric_only)
        else:
            raise ValueError(f"axis={axis} not supported. Please specify 0 or 1")

    def std(self, axis=0, skipna=True, ddof=1, numeric_only=False, split_every=False):
        return Sqrt(self.var(axis, skipna, ddof, numeric_only, split_every=split_every))

    def mean(self, skipna=True, numeric_only=False, split_every=False, axis=0):
        return Mean(
            self,
            skipna=skipna,
            numeric_only=numeric_only,
            split_every=split_every,
            axis=axis,
        )

    def max(self, skipna=True, numeric_only=False, split_every=False, axis=0):
        return Max(self, skipna, numeric_only, split_every, axis)

    def any(self, skipna=True, split_every=False):
        return Any(self, skipna=skipna, split_every=split_every)

    def all(self, skipna=True, split_every=False):
        return All(self, skipna=skipna, split_every=split_every)

    def idxmin(self, skipna=True, numeric_only=False, split_every=False):
        return IdxMin(
            self, skipna=skipna, numeric_only=numeric_only, split_every=split_every
        )

    def idxmax(self, skipna=True, numeric_only=False, split_every=False):
        return IdxMax(
            self, skipna=skipna, numeric_only=numeric_only, split_every=split_every
        )

    def mode(self, dropna=True, split_every=False):
        return Mode(self, dropna=dropna, split_every=split_every)

    def min(self, skipna=True, numeric_only=False, split_every=False, axis=0):
        return Min(self, skipna, numeric_only, split_every=split_every, axis=axis)

    def count(self, numeric_only=False, split_every=False):
        return Count(self, numeric_only, split_every)

    def cumsum(self, skipna=True):
        from dask_expr._cumulative import CumSum

        return CumSum(self, skipna=skipna)

    def cumprod(self, skipna=True):
        from dask_expr._cumulative import CumProd

        return CumProd(self, skipna=skipna)

    def cummax(self, skipna=True):
        from dask_expr._cumulative import CumMax

        return CumMax(self, skipna=skipna)

    def cummin(self, skipna=True):
        from dask_expr._cumulative import CumMin

        return CumMin(self, skipna=skipna)

    def abs(self):
        return Abs(self)

    def astype(self, dtypes):
        return AsType(self, dtypes)

    def clip(self, lower=None, upper=None, axis=None):
        return Clip(self, lower=lower, upper=upper, axis=axis)

    def combine_first(self, other):
        if are_co_aligned(self, other):
            return CombineFirst(self, other=other)
        else:
            return CombineFirstAlign(self, other)

    def to_timestamp(self, freq=None, how="start"):
        return ToTimestamp(self, freq=freq, how=how)

    def isna(self):
        return IsNa(self)

    def isnull(self):
        # These are the same anyway
        return IsNa(self)

    def round(self, decimals=0):
        return Round(self, decimals=decimals)

    def where(self, cond, other=np.nan):
        if not are_co_aligned(self, *[c for c in [cond, other] if isinstance(c, Expr)]):
            return WhereAlign(self, cond=cond, other=other)
        return Where(self, cond=cond, other=other)

    def mask(self, cond, other=np.nan):
        if not are_co_aligned(self, *[c for c in [cond, other] if isinstance(c, Expr)]):
            return MaskAlign(self, cond=cond, other=other)
        return Mask(self, cond=cond, other=other)

    def apply(self, function, *args, meta=None, **kwargs):
        return Apply(self, function, args, meta, kwargs)

    def replace(self, to_replace=None, value=no_default, regex=False):
        return Replace(self, to_replace=to_replace, value=value, regex=regex)

    def fillna(self, value=None):
        if isinstance(value, Expr) and not are_co_aligned(self, value):
            return FillnaAlign(self, value=value)
        return Fillna(self, value=value)

    def rename_axis(
        self, mapper=no_default, index=no_default, columns=no_default, axis=0
    ):
        return RenameAxis(self, mapper=mapper, index=index, columns=columns, axis=axis)

    def align(self, other, join="outer", axis=None, fill_value=None):
        from dask_expr._collection import new_collection

        if not are_co_aligned(self, other):
            aligned = AlignAlignPartitions(self, other, join, axis, fill_value)
        else:
            aligned = _Align(self, other, join, axis=axis, fill_value=fill_value)

        return new_collection(AlignGetitem(aligned, position=0)), new_collection(
            AlignGetitem(aligned, position=1)
        )

    def nunique_approx(self, split_every=None):
        return NuniqueApprox(self, b=16, split_every=split_every)

    def memory_usage_per_partition(self, index=True, deep=False):
        return MemoryUsagePerPartition(self, index, deep)

    @functools.cached_property
    def divisions(self):
        return tuple(self._divisions())

    def _divisions(self):
        raise NotImplementedError()

    @property
    def known_divisions(self):
        """Whether divisions are already known"""
        return len(self.divisions) > 0 and self.divisions[0] is not None

    @property
    def npartitions(self):
        if "npartitions" in self._parameters:
            idx = self._parameters.index("npartitions")
            return self.operands[idx]
        else:
            return len(self.divisions) - 1

    @property
    def columns(self) -> list:
        try:
            return list(self._meta.columns)
        except AttributeError:
            if self.ndim == 1:
                return [self.name]
            return []
        except Exception:
            raise

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self) -> set:
        """Preserves the columns defining the partition mapping from shuffles.

        This property specifies if a column or a set of columns have a unique
        partition mapping that was defined by a shuffle operation. The mapping
        is created by hasing the values and the separating them onto partitions.
        It is important that this property is only propagated if the values
        in those columns did not change in this expression. The property is
        only populated if the mapping was created by the ``partitioning_index``
        function.

        Simply knowing that every value is in only one partition is not a
        satisfying condition, because we also use this property on merge
        operations, where we need these values to be in matching partitions.

        This is also the reason why set_index or sort_values can't set the
        property, they fullfil a weaker condition than what this property enforcey.

        Normally, this set contains one tuple of either one or multiple columns.
        It can contain 2, when the operation shuffles multiple columns of the
        result, i.e. a merge operation and the left and right join columns.


        Returns
        -------
            A set of column groups that have a unique partition mapping as
            defined by a shuffle.
        """
        return set()

    @property
    def _projection_columns(self):
        return self.columns

    @property
    def name(self):
        return self._meta.name

    @property
    def dtypes(self):
        return self._meta.dtypes

    def _filter_simplification(self, parent, predicate=None):
        if predicate is None:
            predicate = parent.predicate.substitute(self, self.frame)
        return type(self)(self.frame[predicate], *self.operands[1:])


class Literal(Expr):
    """Represent a literal (known) value as an `Expr`"""

    _parameters = ["value"]

    def _divisions(self):
        return (None, None)

    @functools.cached_property
    def _meta(self):
        return make_meta(self.value)

    def _task(self, index: int):
        assert index == 0
        return self.value


class Blockwise(Expr):
    """Super-class for block-wise operations

    This is fairly generic, and includes definitions for `_meta`, `divisions`,
    `_layer` that are often (but not always) correct.  Mostly this helps us
    avoid duplication in the future.

    Note that `Fused` expressions rely on every `Blockwise`
    expression defining a proper `_task` method.
    """

    operation = None
    _keyword_only = []
    _projection_passthrough = False
    _preserves_partitioning_information = False

    @functools.cached_property
    def _meta(self):
        args = [op._meta if isinstance(op, Expr) else op for op in self._args]
        return self.operation(*args, **self._kwargs)

    @functools.cached_property
    def _kwargs(self) -> dict:
        if self._keyword_only:
            return {
                p: self.operand(p)
                for p in self._parameters
                if p in self._keyword_only and self.operand(p) is not no_default
            }
        return {}

    @functools.cached_property
    def _args(self) -> list:
        if self._keyword_only:
            args = [
                self.operand(p) for p in self._parameters if p not in self._keyword_only
            ] + self.operands[len(self._parameters) :]
            return args
        return self.operands

    def _broadcast_dep(self, dep: Expr):
        # Checks if a dependency should be broadcasted to
        # all partitions of this `Blockwise` operation
        return dep.npartitions == 1 and dep.ndim < self.ndim

    def _divisions(self):
        # This is an issue.  In normal Dask we re-divide everything in a step
        # which combines divisions and graph.
        # We either have to create a new Align layer (ok) or combine divisions
        # and graph into a single operation.
        dependencies = self.dependencies()
        for arg in dependencies:
            if not self._broadcast_dep(arg):
                assert arg.divisions == dependencies[0].divisions
        return dependencies[0].divisions

    @functools.cached_property
    def _name(self):
        if self.operation:
            head = funcname(self.operation)
        else:
            head = funcname(type(self)).lower()
        return head + "-" + _tokenize_deterministic(*self.operands)

    def _blockwise_arg(self, arg, i):
        """Return a Blockwise-task argument"""
        if isinstance(arg, Expr):
            # Make key for Expr-based argument
            if self._broadcast_dep(arg):
                return (arg._name, 0)
            else:
                return (arg._name, i)

        else:
            return arg

    def _task(self, index: int):
        """Produce the task for a specific partition

        Parameters
        ----------
        index:
            Partition index for this task.

        Returns
        -------
        task: tuple
        """
        args = [self._blockwise_arg(op, index) for op in self._args]
        if self._kwargs:
            return apply, self.operation, args, self._kwargs
        else:
            return (self.operation,) + tuple(args)

    def _simplify_up(self, parent, dependents):
        if self._projection_passthrough and isinstance(parent, Projection):
            return plain_column_projection(self, parent, dependents)

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        if self._preserves_partitioning_information:
            return self.frame.unique_partition_mapping_columns_from_shuffle
        return set()


class MapPartitions(Blockwise):
    _parameters = [
        "frame",
        "func",
        "meta",
        "enforce_metadata",
        "transform_divisions",
        "clear_divisions",
        "align_dataframes",
        "parent_meta",
        "token",
        "kwargs",
    ]
    _defaults = {
        "kwargs": None,
        "align_dataframes": True,
        "parent_meta": None,
        "token": None,
    }

    @functools.cached_property
    def token(self):
        if "token" in self._parameters:
            return self.operand("token")
        return None

    def __str__(self):
        return f"MapPartitions({funcname(self.func)})"

    @functools.cached_property
    def _name(self):
        if self.token is not None:
            head = self.token
        else:
            head = funcname(self.func).lower()
        return head + "-" + _tokenize_deterministic(*self.operands)

    def _broadcast_dep(self, dep: Expr):
        # Always broadcast single-partition dependencies in MapPartitions
        return dep.npartitions == 1

    @functools.cached_property
    def args(self):
        return [self.frame] + self.operands[len(self._parameters) :]

    @functools.cached_property
    def _meta(self):
        meta = self.operand("meta")
        return _get_meta_map_partitions(
            self.args,
            [e for e in self.args if isinstance(e, Expr)],
            self.func,
            self.kwargs,
            meta,
            self.parent_meta,
        )

    def _divisions(self):
        # Unknown divisions
        dfs = [arg for arg in self.args if isinstance(arg, Expr)]

        if self.clear_divisions:
            max_partitions = max(df.npartitions for df in dfs)
            return (None,) * (max_partitions + 1)

        # (Possibly) known divisions
        return _get_divisions_map_partitions(
            self.align_dataframes,
            self.transform_divisions,
            dfs,
            self.func,
            self.args,
            self.kwargs,
        )

    @functools.cached_property
    def _has_partition_info(self):
        return has_keyword(self.func, "partition_info")

    def _task(self, index: int):
        args = [self._blockwise_arg(op, index) for op in self.args]
        kwargs = (self.kwargs if self.kwargs is not None else {}).copy()
        if self._has_partition_info:
            kwargs["partition_info"] = {
                "number": index,
                "division": self.divisions[index],
            }

        if self.enforce_metadata:
            kwargs.update(
                {
                    "_func": self.func,
                    "_meta": self._meta,
                }
            )
            return (apply, apply_and_enforce, args, kwargs)
        else:
            return (
                apply,
                self.func,
                args,
                kwargs,
            )


def _get_meta_ufunc(dfs, args, func):
    dasks = [arg for arg in args if isinstance(arg, (Expr, Array))]

    if len(dfs) >= 2 and not all(hasattr(d, "npartitions") for d in dasks):
        # should not occur in current funcs
        msg = "elemwise with 2 or more DataFrames and Scalar is not supported"
        raise NotImplementedError(msg)
    # For broadcastable series, use no rows.
    parts = [
        (
            d._meta
            if d.ndim == 0
            else (
                np.empty((), dtype=d.dtype)
                if isinstance(d, Array)
                else meta_nonempty(d._meta)
            )
        )
        for d in dasks
    ]

    other = [
        (i, arg) for i, arg in enumerate(args) if not isinstance(arg, (Expr, Array))
    ]

    return partial_by_order(*parts, function=func, other=other)


class UFuncElemwise(MapPartitions):
    _parameters = [
        "frame",
        "func",
        "meta",
        "transform_divisions",
        "kwargs",
    ]
    enforce_metadata = False

    def __str__(self):
        return f"UFunc({funcname(self.func)})"

    @functools.cached_property
    def args(self):
        return self.operands[len(self._parameters) :]

    @functools.cached_property
    def _dfs(self):
        return [df for df in self.args if isinstance(df, Expr) and df.ndim > 0]

    @functools.cached_property
    def _meta(self):
        if self.operand("meta") is not no_default:
            meta = self.operand("meta")
        else:
            meta = _get_meta_ufunc(self._dfs, self.args, self.func)
        return make_meta(meta)

    def _divisions(self):
        if (
            self.transform_divisions
            and isinstance(self._dfs[0], Index)
            and len(self._dfs) == 1
        ):
            try:
                divisions = self.func(
                    *[
                        pd.Index(arg.divisions) if arg is self._dfs[0] else arg
                        for arg in self.args
                    ],
                    **self.kwargs,
                )
                if isinstance(divisions, pd.Index):
                    divisions = methods.tolist(divisions)
            except Exception:
                pass
            else:
                if not valid_divisions(divisions):
                    divisions = [None] * (self._dfs[0].npartitions + 1)
                return divisions

        return self._dfs[0].divisions


class MapOverlapAlign(Expr):
    _parameters = [
        "frame",
        "func",
        "before",
        "after",
        "meta",
        "enforce_metadata",
        "transform_divisions",
        "clear_divisions",
        "align_dataframes",
        "token",
        "kwargs",
    ]
    _defaults = {
        "meta": None,
        "enfore_metadata": True,
        "transform_divisions": True,
        "kwargs": None,
        "clear_divisions": False,
        "align_dataframes": False,
        "token": None,
    }

    @functools.cached_property
    def _meta(self):
        meta = self.operand("meta")
        args = [self.frame._meta] + [
            arg._meta if isinstance(arg, Expr) else arg
            for arg in self.operands[len(self._parameters) :]
        ]
        return _get_meta_map_partitions(
            args,
            [self.dependencies()[0]],
            self.func,
            self.kwargs,
            meta,
            self.kwargs.pop("parent_meta", None),
        )

    def _divisions(self):
        args = [self.frame] + self.operands[len(self._parameters) :]
        return calc_divisions_for_align(*args, allow_shuffle=False)

    def _lower(self):
        args = [self.frame] + self.operands[len(self._parameters) :]
        args = maybe_align_partitions(*args, divisions=self._divisions())
        return MapOverlap(
            args[0],
            self.func,
            self.before,
            self.after,
            self._meta,
            self.enforce_metadata,
            self.transform_divisions,
            self.clear_divisions,
            self.align_dataframes,
            self.token,
            self.kwargs,
            *args[1:],
        )


class MapOverlap(MapPartitions):
    _parameters = [
        "frame",
        "func",
        "before",
        "after",
        "meta",
        "enforce_metadata",
        "transform_divisions",
        "clear_divisions",
        "align_dataframes",
        "token",
        "kwargs",
    ]
    _defaults = {
        "meta": None,
        "enfore_metadata": True,
        "transform_divisions": True,
        "kwargs": None,
        "clear_divisions": False,
        "align_dataframes": False,
        "token": None,
    }

    @functools.cached_property
    def _kwargs(self) -> dict:
        kwargs = self.kwargs
        if kwargs is None:
            kwargs = {}
        return kwargs

    @property
    def args(self):
        return (
            [self.frame]
            + [self.func, self.before, self.after]
            + self.operands[len(self._parameters) :]
        )

    @functools.cached_property
    def _meta(self):
        meta = self.operand("meta")
        args = [self.frame._meta] + [
            arg._meta if isinstance(arg, Expr) else arg
            for arg in self.operands[len(self._parameters) :]
        ]
        return _get_meta_map_partitions(
            args,
            [self.dependencies()[0]],
            self.func,
            self.kwargs,
            meta,
            self.kwargs.pop("parent_meta", None),
        )

    @functools.cached_property
    def before(self):
        before = self.operand("before")
        if isinstance(before, str):
            return pd.to_timedelta(before)
        return before

    @functools.cached_property
    def after(self):
        after = self.operand("after")
        if isinstance(after, str):
            return pd.to_timedelta(after)
        return after

    def _lower(self):
        overlapped = CreateOverlappingPartitions(self.frame, self.before, self.after)

        return MapPartitions(
            overlapped,
            _overlap_chunk,
            self._meta,
            self.enforce_metadata,
            self.transform_divisions,
            self.clear_divisions,
            self.align_dataframes,
            None,
            self.token,
            self._kwargs,
            *self.args[1:],
        )


class CreateOverlappingPartitions(Expr):
    _parameters = ["frame", "before", "after"]

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    def _divisions(self):
        # Keep divisions alive, MapPartitions will handle the actual division logic
        return self.frame.divisions

    def _layer(self) -> dict:
        dsk, prevs, nexts = {}, [], []

        name_prepend = "overlap-prepend" + self.frame._name
        if self.before:
            prevs.append(None)
            if isinstance(self.before, numbers.Integral):
                before = self.before
                for i in range(self.frame.npartitions - 1):
                    dsk[(name_prepend, i)] = (M.tail, (self.frame._name, i), before)
                    prevs.append((name_prepend, i))
            elif isinstance(self.before, datetime.timedelta):
                # Assumes monotonic (increasing?) index
                divs = pd.Series(self.frame.divisions)
                deltas = divs.diff().iloc[1:-1]

                # In the first case window-size is larger than at least one partition, thus it is
                # necessary to calculate how many partitions must be used for each rolling task.
                # Otherwise, these calculations can be skipped (faster)

                if (self.before > deltas).any():
                    pt_z = divs[0]
                    for i in range(self.frame.npartitions - 1):
                        # Select all indexes of relevant partitions between the current partition and
                        # the partition with the highest division outside the rolling window (before)
                        pt_i = divs[i + 1]

                        # lower-bound the search to the first division
                        lb = max(pt_i - self.before, pt_z)

                        first, j = divs[i], i
                        while first > lb and j > 0:
                            first = first - deltas[j]
                            j = j - 1

                        dsk[(name_prepend, i)] = (
                            _tail_timedelta,
                            (self.frame._name, i + 1),
                            [(self.frame._name, k) for k in range(j, i + 1)],
                            self.before,
                        )
                        prevs.append((name_prepend, i))
                else:
                    for i in range(self.frame.npartitions - 1):
                        dsk[(name_prepend, i)] = (
                            _tail_timedelta,
                            (self.frame._name, i + 1),
                            [(self.frame._name, i)],
                            self.before,
                        )
                        prevs.append((name_prepend, i))
        else:
            prevs.extend([None] * self.frame.npartitions)

        name_append = "overlap-append" + self.frame._name
        if self.after:
            if isinstance(self.after, numbers.Integral):
                after = self.after
                for i in range(1, self.frame.npartitions):
                    dsk[(name_append, i)] = (M.head, (self.frame._name, i), after)
                    nexts.append((name_append, i))
            else:
                # We don't want to look at the divisions, so take twice the step and
                # validate later.
                after = 2 * self.after
                for i in range(1, self.frame.npartitions):
                    dsk[(name_append, i)] = (
                        _head_timedelta,
                        (self.frame._name, i - 1),
                        (self.frame._name, i),
                        after,
                    )
                    nexts.append((name_append, i))

            nexts.append(None)

        else:
            nexts.extend([None] * self.frame.npartitions)

        for i, (prev, next) in enumerate(zip(prevs, nexts)):
            dsk[(self._name, i)] = (
                _combined_parts,
                prev,
                (self.frame._name, i),
                next,
                self.before,
                self.after,
            )
        return dsk


def _tail_timedelta(current, prev_, before):
    selected = methods.concat(
        [prev[prev.index > (current.index.min() - before)] for prev in prev_]
    )
    return selected


def _overlap_chunk(df, func, before, after, *args, **kwargs):
    return overlap_chunk(func, before, after, df, *args, **kwargs)


def _combined_parts(prev_part, current_part, next_part, before, after):
    msg = (
        "Partition size is less than overlapping "
        "window size. Try using ``df.repartition`` "
        "to increase the partition size."
    )

    if prev_part is not None:
        if isinstance(before, numbers.Integral):
            if prev_part.shape[0] != before:
                raise NotImplementedError(msg)
        else:
            prev_part_input = prev_part
            prev_part = _tail_timedelta(current_part, [prev_part], before)
            if (
                len(prev_part_input) == len(prev_part)
                and len(prev_part_input) > 0
                and not isinstance(before, datetime.timedelta)
            ):
                raise NotImplementedError(msg)

    if next_part is not None:
        if isinstance(after, numbers.Integral):
            if next_part.shape[0] != after:
                raise NotImplementedError(msg)
        else:
            next_part_input = next_part
            next_part = _head_timedelta(current_part, next_part, after)
            if len(next_part_input) == len(next_part) and len(next_part_input) > 0:
                raise NotImplementedError(msg)

    parts = [p for p in (prev_part, current_part, next_part) if p is not None]
    combined = methods.concat(parts)

    return CombinedOutput(
        (
            combined,
            len(prev_part) if prev_part is not None and len(prev_part) > 0 else None,
            len(next_part) if next_part is not None and len(next_part) > 0 else None,
        )
    )


class _Align(Blockwise):
    _parameters = ["frame", "other", "join", "axis", "fill_value"]
    _defaults = {"join": "outer", "fill_value": None, "axis": None}
    _keyword_only = ["join", "fill_value", "axis"]
    operation = M.align
    _preserves_partitioning_information = True

    def _divisions(self):
        # Aligning, so take first frames divisions
        return self.frame._divisions()


class AlignGetitem(Blockwise):
    _parameters = ["frame", "position"]
    operation = operator.getitem
    _preserves_partitioning_information = True

    @functools.cached_property
    def _meta(self):
        return self.frame._meta[self.position]

    def _divisions(self):
        return self.frame.divisions


class ScalarToSeries(Blockwise):
    _parameters = ["frame", "index"]
    _defaults = {"index": 0}

    @staticmethod
    def operation(value, index=0):
        return pd.Series(value, index=[index])


class DropnaSeries(Blockwise):
    _parameters = ["frame"]
    operation = M.dropna
    _preserves_partitioning_information = True


class DropnaFrame(Blockwise):
    _parameters = ["frame", "how", "subset", "thresh"]
    _defaults = {"how": no_default, "subset": None, "thresh": no_default}
    _keyword_only = ["how", "subset", "thresh"]
    operation = M.dropna
    _preserves_partitioning_information = True

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection) and self.subset is not None:
            columns = determine_column_projection(
                self, parent, dependents, additional_columns=self.subset
            )
            columns = [col for col in self.frame.columns if col in columns]

            if columns == self.frame.columns:
                # Don't add unnecessary Projections
                return

            return type(parent)(
                type(self)(self.frame[columns], *self.operands[1:]),
                *parent.operands[1:],
            )


class CombineFirst(Blockwise):
    _parameters = ["frame", "other"]
    operation = M.combine_first

    @functools.cached_property
    def _meta(self):
        return make_meta(
            self.operation(
                meta_nonempty(self.frame._meta),
                meta_nonempty(self.other._meta),
            ),
        )

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            columns = determine_column_projection(self, parent, dependents)
            frame_columns = [col for col in self.frame.columns if col in columns]
            other_columns = [col for col in self.other.columns if col in columns]
            if (
                self.frame.columns == frame_columns
                and self.other.columns == other_columns
            ):
                return

            return type(parent)(
                type(self)(self.frame[frame_columns], self.other[other_columns]),
                *parent.operands[1:],
            )


class Sample(Blockwise):
    _parameters = ["frame", "state_data", "frac", "replace"]
    operation = staticmethod(methods.sample)

    @functools.cached_property
    def _meta(self):
        args = [self.operands[0]._meta] + [self.operands[1][0]] + self.operands[2:]
        return self.operation(*args)

    def _task(self, index: int):
        args = [self._blockwise_arg(self.frame, index)] + [
            self.state_data[index],
            self.frac,
            self.operand("replace"),
        ]
        return (self.operation,) + tuple(args)


class Query(Blockwise):
    _parameters = ["frame", "_expr", "expr_kwargs"]
    _defaults = {"expr_kwargs": {}}
    _keyword_only = ["expr_kwargs"]
    operation = M.query

    @functools.cached_property
    def _kwargs(self) -> dict:
        return {**self.expr_kwargs}


class MemoryUsagePerPartition(Blockwise):
    _parameters = ["frame", "index", "deep"]
    _defaults = {"index": True, "deep": False}

    @staticmethod
    def operation(*args, **kwargs):
        if is_series_like(args[0]):
            return args[0]._constructor([total_mem_usage(*args, **kwargs)])
        return args[0]._constructor_sliced([total_mem_usage(*args, **kwargs)])

    def _divisions(self):
        return (None,) * (self.frame.npartitions + 1)


class DropDuplicatesBlockwise(Blockwise):
    _parameters = ["frame"]
    operation = M.drop_duplicates
    _preserves_partitioning_information = True


class Elemwise(Blockwise):
    """
    This doesn't really do anything, but we anticipate that future
    optimizations, like `len` will care about which operations preserve length
    """

    _is_length_preserving = True

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Filter) and self._filter_passthrough_available(
            parent, dependents
        ):
            predicate = None
            if self.frame.ndim == 1 and self.ndim == 2:
                name = self.frame._meta.name
                # Avoid Projection since we are already a Series
                subs = Projection(self, name)
                predicate = parent.predicate.substitute(subs, self.frame)
            return self._filter_simplification(parent, predicate)
        return super()._simplify_up(parent, dependents)


class RenameFrame(Elemwise):
    _parameters = ["frame", "columns"]

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        result = set()
        columns = self.operand("columns")
        for elem in self.frame.unique_partition_mapping_columns_from_shuffle:
            if isinstance(elem, tuple):
                subset = self.frame._meta[list(elem)].rename(columns=columns)
                result.add(tuple(list(subset.columns)))
            else:
                # scalar
                subset = self.frame._meta[[elem]]
                result.add(subset.columns[0])
        return result

    @staticmethod
    def operation(df, columns):
        return df.rename(columns=columns)

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection) and isinstance(
            self.operand("columns"), Mapping
        ):
            reverse_mapping = {val: key for key, val in self.operand("columns").items()}

            columns = determine_column_projection(self, parent, dependents)
            columns = _convert_to_list(columns)
            frame_columns = set(self.frame.columns)
            columns = [
                reverse_mapping[col]
                if col in reverse_mapping and reverse_mapping[col] in frame_columns
                else col
                for col in columns
            ]
            columns = [col for col in self.frame.columns if col in columns]
            if columns == self.frame.columns:
                return

            return type(parent)(
                type(self)(self.frame[columns], *self.operands[1:]),
                *parent.operands[1:],
            )


class ColumnsSetter(RenameFrame):
    _preserves_partitioning_information = True

    @staticmethod
    def operation(df, columns):
        return _rename(columns, df)


class _DeepCopy(Elemwise):
    _parameters = ["frame"]
    _projection_passthrough = True
    _filter_passthrough = True
    _preserves_partitioning_information = True

    @staticmethod
    def operation(df):
        return df.copy(deep=True)


class ResourceBarrier(Expr):
    @property
    def _resources(self):
        raise NotImplementedError()

    def __str__(self):
        return f"{type(self).__name__}({self._resources})"


class ElemwiseResourceBarrier(Elemwise, ResourceBarrier):
    _parameters = ["frame", "resource_spec"]
    _projection_passthrough = True
    _filter_passthrough = True
    _preserves_partitioning_information = True

    @property
    def _resources(self):
        return self.resource_spec

    def _task(self, index: int):
        return (self.frame._name, index)

    @property
    def _meta(self):
        return self.frame._meta

    def _divisions(self):
        return self.frame.divisions


class RenameSeries(Elemwise):
    _parameters = ["frame", "index", "sorted_index"]
    _defaults = {"sorted_index": False}
    _filter_passthrough = True
    _preserves_partitioning_information = True

    @functools.cached_property
    def _meta(self):
        args = [
            meta_nonempty(op._meta) if isinstance(op, Expr) else op for op in self._args
        ]
        return make_meta(self.operation(*args, **self._kwargs))

    @staticmethod
    def operation(df, index, sorted_index):
        if is_series_like(df):
            return df.rename(index=index)
        return df.rename(name=index)

    def _divisions(self):
        index = self.operand("index")
        if is_scalar(index) and not isinstance(index, Callable):
            return self.frame.divisions
        elif self.sorted_index and self.frame.known_divisions:
            old = pd.Series(1, index=self.frame.divisions)
            new_divisions = old.rename(index).index
            if not new_divisions.is_monotonic_increasing:
                raise ValueError(
                    "The renamer creates an Index with non-monotonic divisions. "
                    "This is not allowed. Please set sorted_index=False."
                )
            return tuple(new_divisions.tolist())
        else:
            return (None,) * (self.frame.npartitions + 1)


class Fillna(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "value"]
    _defaults = {"value": None}
    operation = M.fillna


class Replace(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "to_replace", "value", "regex"]
    _defaults = {"to_replace": None, "value": no_default, "regex": False}
    _keyword_only = ["value", "regex"]
    operation = M.replace


class Isin(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "values"]
    operation = M.isin

    @functools.cached_property
    def _meta(self):
        return make_meta(
            meta_nonempty(self.frame._meta).isin(
                meta_nonempty(self.frame._meta).iloc[[0]]
            )
        )

    def _broadcast_dep(self, dep: Expr):
        return dep.npartitions == 1


class Clip(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "lower", "upper", "axis"]
    _defaults = {"lower": None, "upper": None, "axis": None}
    _keyword_only = ["axis"]
    operation = M.clip

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            return plain_column_projection(self, parent, dependents)


class ArrowStringConversion(Elemwise):
    _projection_passthrough = True
    _filter_passthrough = True
    _parameters = ["frame"]
    operation = staticmethod(to_pyarrow_string)
    _preserves_partitioning_information = True


class Between(Elemwise):
    _parameters = ["frame", "left", "right", "inclusive"]
    _defaults = {"inclusive": "both"}
    operation = M.between


class ToTimestamp(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "freq", "how"]
    _defaults = {"freq": None, "how": "start"}
    operation = M.to_timestamp
    _filter_passthrough = True
    _preserves_partitioning_information = True

    def _divisions(self):
        return tuple(
            pd.Index(self.frame.divisions).to_timestamp(freq=self.freq, how=self.how)
        )


class CombineSeries(Elemwise):
    _parameters = ["frame", "other", "func", "fill_value"]
    _defaults = {"fill_value": None}
    operation = M.combine

    @functools.cached_property
    def _meta(self):
        return make_meta(
            meta_nonempty(self.frame._meta).combine(
                meta_nonempty(self.other._meta), func=self.func
            )
        )


class CombineFrame(CombineSeries):
    _parameters = CombineSeries._parameters + ["overwrite"]
    _defaults = {"fill_value": None, "overwrite": True}


class ToNumeric(Elemwise):
    _parameters = ["frame", "errors", "downcast", "meta"]
    _defaults = {"errors": "raise", "downcast": None, "meta": None}
    _keyword_only = ["meta"]
    operation = staticmethod(pd.to_numeric)

    @functools.cached_property
    def _kwargs(self):
        kwargs = super()._kwargs
        kwargs.pop("meta", None)
        return kwargs

    @functools.cached_property
    def _meta(self):
        if self.operand("meta") is not None:
            return self.operand("meta")
        return super()._meta


class ToDatetime(Elemwise):
    _parameters = ["frame", "kwargs", "meta"]
    _defaults = {"kwargs": None}
    _keyword_only = ["kwargs", "meta"]

    @staticmethod
    def operation(*args, **kwargs):
        return get_meta_library(args[0]).to_datetime(*args, **kwargs)

    @functools.cached_property
    def _kwargs(self):
        if (kwargs := self.operand("kwargs")) is None:
            return {}
        return kwargs


class ToTimedelta(Elemwise):
    _parameters = ["frame", "unit", "errors"]
    _defaults = {"unit": None, "errors": "raise"}
    operation = staticmethod(pd.to_timedelta)


class AsType(Elemwise):
    """A good example of writing a trivial blockwise operation"""

    _parameters = ["frame", "dtypes"]
    operation = M.astype
    _filter_passthrough = True

    @functools.cached_property
    def _meta(self):
        def _cat_dtype_without_categories(dtype):
            return (
                isinstance(pd.api.types.pandas_dtype(dtype), pd.CategoricalDtype)
                and getattr(dtype, "categories", None) is None
            )

        meta = super()._meta
        dtypes = self.operand("dtypes")
        if hasattr(dtypes, "items"):
            set_unknown = [
                k for k, v in dtypes.items() if _cat_dtype_without_categories(v)
            ]
            meta = clear_known_categories(meta, cols=set_unknown)

        elif _cat_dtype_without_categories(dtypes):
            meta = clear_known_categories(meta)
        return meta

    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):
            dtypes = self.operand("dtypes")
            columns = determine_column_projection(self, parent, dependents)
            if isinstance(dtypes, dict):
                dtypes = {key: val for key, val in dtypes.items() if key in columns}
                if not dtypes:
                    return type(parent)(self.frame, *parent.operands[1:])
            if isinstance(columns, list):
                columns = [col for col in self.frame.columns if col in columns]
            if self.frame.columns == columns:
                return
            result = type(self)(self.frame[columns], dtypes)
            if not isinstance(columns, list):
                return result
            return type(parent)(result, *parent.operands[1:])


class IsNa(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame"]
    operation = M.isna


class Mask(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "cond", "other"]
    _defaults = {"other": np.nan}
    operation = M.mask


class Round(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "decimals"]
    operation = M.round


class Where(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "cond", "other"]
    _defaults = {"other": np.nan}
    operation = M.where


def _check_divisions(df, i, division_min, division_max, last):
    # Check divisions
    real_min = df.index.min()
    real_max = df.index.max()
    # Upper division of the last partition is often set to
    # the max value. For all other partitions, the upper
    # division should be greater than the maximum value.
    valid_min = real_min >= division_min
    valid_max = (real_max <= division_max) if last else (real_max < division_max)
    if not (valid_min and valid_max):
        raise RuntimeError(
            f"`enforce_runtime_divisions` failed for partition {i}."
            f" Expected a range of [{division_min}, {division_max}), "
            f" but the real range was [{real_min}, {real_max}]."
        )
    return df


class EnforceRuntimeDivisions(Blockwise):
    _parameters = ["frame"]
    operation = staticmethod(_check_divisions)
    _preserves_partitioning_information = True

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    def _task(self, index: int):
        args = [self._blockwise_arg(op, index) for op in self._args]
        args = args + [
            index,
            self.divisions[index],
            self.divisions[index + 1],
            index == (self.npartitions - 1),
        ]
        return (self.operation,) + tuple(args)


class Abs(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame"]
    operation = M.abs


class RenameAxis(Elemwise):
    _projection_passthrough = True
    _filter_passthrough = True
    _parameters = ["frame", "mapper", "index", "columns", "axis"]
    _defaults = {
        "mapper": no_default,
        "index": no_default,
        "columns": no_default,
        "axis": 0,
    }
    _keyword_only = ["mapper", "index", "columns", "axis"]
    operation = M.rename_axis
    _preserves_partitioning_information = True


class NotNull(Elemwise):
    _parameters = ["frame"]
    operation = M.notnull
    _projection_passthrough = True


class ToFrame(Elemwise):
    _parameters = ["frame", "name"]
    _defaults = {"name": no_default}
    _keyword_only = ["name"]
    operation = M.to_frame
    _filter_passthrough = True

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        result = set()
        name_mapping = dict(zip(self.frame.columns, self.columns))
        for elem in self.frame.unique_partition_mapping_columns_from_shuffle:
            if isinstance(elem, tuple):
                result.add(tuple(name_mapping.get(v, v) for v in elem))
            else:
                result.add(name_mapping.get(elem, elem))
        return result


class ToFrameIndex(ToFrame):
    _parameters = ["frame", "index", "name"]
    _defaults = {"name": no_default, "index": True}
    _keyword_only = ["name", "index"]
    operation = M.to_frame
    _filter_passthrough = True


class ToSeriesIndex(ToFrameIndex):
    _defaults = {"name": no_default, "index": None}
    operation = M.to_series
    _preserves_partitioning_information = True


def pd_split(df, p, random_state=None, shuffle=False):
    """Split DataFrame into multiple pieces pseudorandomly

    >>> df = pd.DataFrame({'a': [1, 2, 3, 4, 5, 6],
    ...                    'b': [2, 3, 4, 5, 6, 7]})

    >>> a, b = pd_split(
    ...     df, [0.5, 0.5], random_state=123, shuffle=True
    ... )  # roughly 50/50 split
    >>> a
       a  b
    3  4  5
    0  1  2
    5  6  7
    >>> b
       a  b
    1  2  3
    4  5  6
    2  3  4
    """
    p = list(p)
    if shuffle:
        if not isinstance(random_state, np.random.RandomState):
            random_state = np.random.RandomState(random_state)
        df = df.sample(frac=1.0, random_state=random_state)
    index = pseudorandom(len(df), p, random_state)
    if df.ndim == 1:
        df = df.to_frame()
    return df.assign(_split=index)


class Split(Elemwise):
    _parameters = ["frame", "frac", "random_state", "shuffle"]
    _keyword_only = ["random_state", "shuffle"]
    operation = staticmethod(pd_split)

    @functools.cached_property
    def _kwargs(self) -> dict:
        return {"shuffle": self.shuffle}

    @functools.cached_property
    def random_state_data(self):
        return random_state_data(self.frame.npartitions, self.random_state)

    def _task(self, index: int):
        args = [self._blockwise_arg(op, index) for op in self._args]
        kwargs = self._kwargs.copy()
        kwargs["random_state"] = self.random_state_data[index]
        return apply, self.operation, args, kwargs


def _random_split_take(df, i, ndim):
    df = df[df["_split"] == i].drop(columns="_split")
    if ndim == 1:
        return df[df.columns[0]]
    return df


class SplitTake(Blockwise):
    _parameters = ["frame", "i", "ndim"]
    operation = staticmethod(_random_split_take)


class Apply(Elemwise):
    """A good example of writing a less-trivial blockwise operation"""

    _parameters = ["frame", "function", "args", "meta", "kwargs"]
    _defaults = {"args": (), "kwargs": {}}
    operation = M.apply

    @functools.cached_property
    def _meta(self):
        return make_meta(self.operand("meta"), parent_meta=self.frame._meta)

    def _task(self, index: int):
        return (
            apply,
            M.apply,
            [
                (self.frame._name, index),
                self.function,
            ]
            + list(self.args),
            self.kwargs,
        )


class Map(Elemwise):
    _projection_passthrough = True
    _parameters = ["frame", "arg", "na_action", "meta", "is_monotonic"]
    _defaults = {"na_action": None, "meta": None, "is_monotonic": True}
    _keyword_only = ["meta", "is_monotonic"]
    operation = M.map

    @functools.cached_property
    def _meta(self):
        if self.operand("meta") is None:
            args = [
                meta_nonempty(op._meta) if isinstance(op, Expr) else op
                for op in self._args
            ]
            return make_meta(self.operation(*args, **self._kwargs))
        return make_meta(
            self.operand("meta"),
            parent_meta=self.frame._meta,
            index=getattr(self.frame._meta, "index", None),  # could be an index
        )

    @functools.cached_property
    def _kwargs(self) -> dict:
        return {}

    def _divisions(self):
        if not self.is_monotonic:
            # Implement this consistently with dask.dataframe, e.g. add option to
            # control monotonic map func
            return (None,) * len(self.frame.divisions)
        elif is_index_like(self.frame._meta):
            return tuple(
                pd.Series(self.frame.divisions).map(self.arg, na_action=self.na_action)
            )
        elif isinstance(self.arg, Expr):
            if self.arg.divisions == self.frame.divisions:
                return self.frame.divisions
            else:
                # We only get here when we have only one partition
                return (None,) * (self.frame.npartitions + 1)
        return super()._divisions()


class VarColumns(Elemwise):
    _parameters = ["frame", "skipna", "ddof", "numeric_only"]
    _defaults = {"skipna": True, "ddof": 1, "numeric_only": False}
    _keyword_only = ["skipna", "ddof", "numeric_only"]
    operation = M.var
    _is_length_preserving = True

    @functools.cached_property
    def _kwargs(self) -> dict:
        return {"axis": 1, **super()._kwargs}


class NUniqueColumns(Elemwise):
    _parameters = ["frame", "axis", "dropna"]
    _defaults = {"axis": 1, "dropna": True}
    operation = M.nunique


class Sqrt(Elemwise):
    _parameters = ["frame"]
    operation = np.sqrt


class ExplodeSeries(Blockwise):
    _parameters = ["frame"]
    operation = M.explode


class ExplodeFrame(ExplodeSeries):
    _parameters = ["frame", "column"]

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            return plain_column_projection(self, parent, dependents, [self.column])


class Drop(Elemwise):
    _parameters = ["frame", "columns", "errors"]
    _defaults = {"errors": "raise"}
    operation = staticmethod(drop_by_shallow_copy)
    _preserves_partitioning_information = True

    def _simplify_down(self):
        col_op = self.operand("columns")
        if is_scalar(col_op):
            col_op = [col_op]
        columns = [col for col in self.frame.columns if col not in col_op]
        return Projection(self.frame, columns)


def assign(df, *pairs):
    pairs = dict(partition(2, pairs))
    df = df.copy(deep=False)
    with warnings.catch_warnings():
        warnings.filterwarnings(
            "ignore",
            message="DataFrame is highly fragmented *",
            category=PerformanceWarning,
        )
        for name, val in pairs.items():
            if isinstance(val, Callable):
                val = val(df)
            df[name] = val
    return df


class Assign(Elemwise):
    """Column Assignment"""

    _parameters = ["frame"]
    operation = staticmethod(assign)

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        keys = set(self.keys)
        return {
            col
            for col in self.frame.unique_partition_mapping_columns_from_shuffle
            if not isinstance(col, tuple)
            and col not in keys
            or not set(col).intersection(keys)
        }

    @functools.cached_property
    def keys(self):
        return self.operands[1::2]

    @functools.cached_property
    def vals(self):
        return self.operands[2::2]

    @functools.cached_property
    def _meta(self):
        args = [op._meta if isinstance(op, Expr) else op for op in self._args]
        return make_meta(self.operation(*args, **self._kwargs))

    def _tree_repr_argument_construction(self, i, op, header):
        if i == 0:
            return super()._tree_repr_argument_construction(i, op, header)
        if i % 2 == 1:
            sep = "" if i == 1 else ","
            header += f"{sep} {repr(op)[1:-1]}="
        else:
            header += f"{repr(op)}"
        return header

    def _node_label_args(self):
        return self.operands

    def _simplify_down(self):
        if isinstance(self.frame, Assign):
            if self._check_for_previously_created_column(self.frame):
                # don't squash if we are using a column that was previously created
                return
            return Assign(*self.frame.operands, *self.operands[1:])

    def _check_for_previously_created_column(self, child):
        input_columns = []
        for v in self.vals:
            if isinstance(v, Expr):
                input_columns.extend(v.columns)
        return bool(set(input_columns) & set(child.keys))

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            columns = determine_column_projection(self, parent, dependents)
            columns = _convert_to_list(columns)

            cols = set(columns) - set(self.keys)
            if cols == set(self.frame.columns):
                # Protect against pushing the same projection twice
                return

            diff = set(self.keys) - set(columns)
            if len(diff) == len(self.keys):
                return type(parent)(self.frame, *parent.operands[1:])
            elif len(diff) > 0:
                new_args = []
                for k, v in zip(self.keys, self.vals):
                    if k in columns:
                        new_args.extend([k, v])
            else:
                new_args = self.operands[1:]

            columns = [col for col in self.frame.columns if col in cols]
            return type(parent)(
                type(self)(self.frame[sorted(columns)], *new_args),
                *parent.operands[1:],
            )


class Eval(Elemwise):
    _parameters = ["frame", "_expr", "expr_kwargs"]
    _defaults = {"expr_kwargs": {}}
    _keyword_only = ["expr_kwargs"]
    operation = M.eval

    @functools.cached_property
    def _kwargs(self) -> dict:
        return {**self.expr_kwargs}


class CaseWhen(Elemwise):
    _parameters = ["frame"]

    @functools.cached_property
    def caselist(self):
        c = self.operands[1:]
        return [(c[i], c[i + 1]) for i in range(0, len(c), 2)]

    @functools.cached_property
    def _meta(self):
        c = self.operands[1:]
        caselist = [
            (
                meta_nonempty(c[i]._meta) if isinstance(c[i], Expr) else c[i],
                (
                    meta_nonempty(c[i + 1]._meta)
                    if isinstance(c[i + 1], Expr)
                    else c[i + 1]
                ),
            )
            for i in range(0, len(c), 2)
        ]
        return make_meta(meta_nonempty(self.frame._meta).case_when(caselist))

    @staticmethod
    def operation(ser, *caselist):
        caselist = [(caselist[i], caselist[i + 1]) for i in range(0, len(caselist), 2)]
        return ser.case_when(list(caselist))


class Filter(Blockwise):
    _projection_passthrough = True
    _filter_passthrough = True
    _parameters = ["frame", "predicate"]
    operation = operator.getitem
    _preserves_partitioning_information = True

    def _simplify_up(self, parent, dependents):
        if isinstance(self.predicate, Or):
            result = rewrite_filters(self.predicate)
            if result._name != self.predicate._name:
                return type(parent)(
                    type(self)(self.frame, result), *parent.operands[1:]
                )

        if isinstance(parent, (FilterAlign, Filter)) and not isinstance(
            self.frame, (FilterAlign, Filter)
        ):
            if not self.frame._filter_passthrough_available(self, dependents):
                # We want to collect filters again when we can't move them
                # anymore. Otherwise, a chain of Filters might block Projections
                # We start with the first Filter, e.g. if my child isn't a filter
                # anymore. Filters above that don't know if the first Filter can
                # still move further
                if is_filter_pushdown_available(
                    self, parent, dependents, allow_reduction=False
                ):
                    # We can only squash 2 filters together if the predicate of parent
                    # does not directly depend on self, e.g. if
                    # sum is in the predicate of parent, then removing self would
                    # alter the condition of parent because the sum changes, this is
                    # only relevant in broadcasting cases
                    return self.frame[
                        self.predicate & parent.predicate.substitute(self, self.frame)
                    ]
        if isinstance(parent, Projection):
            if self.frame._filter_passthrough_available(self, dependents):
                # We can't push Projections through filters if the preceding operation
                # allows us to push filters further down the graph because Projections
                # block filter pushdown
                if not isinstance(self.frame, (FilterAlign, Filter)):
                    return
                elif is_filter_pushdown_available(
                    self.frame, self, dependents, allow_reduction=False
                ):
                    return

            return plain_column_projection(self, parent, dependents)
        if isinstance(parent, Index):
            return self.frame.index[self.predicate]


class Projection(Elemwise):
    """Column Selection"""

    _parameters = ["frame", "columns"]
    operation = operator.getitem

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        col_op = self.operand("columns")
        columns = set(col_op) if isinstance(col_op, list) else {col_op}
        return {
            c
            for c in self.frame.unique_partition_mapping_columns_from_shuffle
            if c in columns or isinstance(c, tuple) and set(c).issubset(columns)
        }

    @property
    def columns(self):
        cols = self.operand("columns")
        if isinstance(cols, list):
            return cols
        elif isinstance(cols, pd.Index):
            return list(cols)
        else:
            return [cols]

    @functools.cached_property
    def _meta(self):
        if is_dataframe_like(self.frame._meta):
            return super()._meta
        # if we are not a DataFrame and have a scalar, we reduce to a scalar
        if not isinstance(self.operand("columns"), (list, slice)) and not hasattr(
            self.operand("columns"), "dtype"
        ):
            return meta_nonempty(self.frame._meta).iloc[0]
        # Avoid column selection for Series/Index
        return self.frame._meta

    def _node_label_args(self):
        return [self.frame, self.operand("columns")]

    def __str__(self):
        base = str(self.frame)
        if " " in base:
            base = "(" + base + ")"
        return f"{base}[{repr(self.operand('columns'))}]"

    def _divisions(self):
        if self.ndim == 0:
            return (None, None)
        return super()._divisions()

    def _simplify_down(self):
        if (
            str(self.frame.columns) == str(self.columns)
            and self._meta.ndim == self.frame._meta.ndim
        ):
            # TODO: we should get more precise around Expr.columns types
            return self.frame
        if isinstance(self.frame, Projection):
            # df[a][b]
            a = self.frame.operand("columns")
            b = self.operand("columns")

            if not isinstance(a, list):
                # df[scalar][b] -> First selection coerces to Series
                return
            elif isinstance(b, list):
                assert all(bb in a for bb in b)
            else:
                assert b in a

            return self.frame.frame[b]


class Index(Elemwise):
    """Column Selection"""

    _parameters = ["frame"]
    operation = getattr

    @functools.cached_property
    def _meta(self):
        meta = self.frame._meta
        # Handle scalar results
        if is_series_like(meta) or is_dataframe_like(meta):
            return self.frame._meta.index
        return meta

    @property
    def _projection_columns(self):
        return []

    def _task(self, index: int):
        return (
            getattr,
            (self.frame._name, index),
            "index",
        )

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        name = self.frame._meta.index.name
        if name in self.frame.unique_partition_mapping_columns_from_shuffle:
            return {name}
        elif (name,) in self.frame.unique_partition_mapping_columns_from_shuffle:
            return {(name,)}
        else:
            return set()


def _return_input(df, divisions=None):
    return df


class ClearDivisions(Elemwise):
    _parameters = ["frame"]
    operation = staticmethod(_return_input)
    _preserves_partitioning_information = True

    def _divisions(self):
        return (None,) * (self.frame.npartitions + 1)


class SetDivisions(Elemwise):
    _parameters = ["frame", "divisions"]
    operation = staticmethod(_return_input)
    _preserves_partitioning_information = True

    def _divisions(self):
        return self.operand("divisions")


class ResolveOverlappingDivisions(Expr):
    _parameters = ["frame", "mins", "maxes", "lens"]

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    def _divisions(self):
        non_empties = [i for i, length in enumerate(self.lens) if length != 0]
        if len(non_empties) == 0:
            return (None, None)

        return tuple(self.mins) + (self.maxes[-1],)

    def _layer(self):
        non_empties = [i for i, length in enumerate(self.lens) if length != 0]
        # If all empty, collapse into one partition
        if len(non_empties) == 0:
            return {(self._name, 0): (self.frame._name, 0)}

        # drop empty partitions by mapping each partition in a new graph to a particular
        # partition on the old graph.
        dsk = {
            (self._name, i): (self.frame._name, div)
            for i, div in enumerate(non_empties)
        }
        ddf_keys = list(dsk.values())

        overlap = [
            i for i in range(1, len(self.mins)) if self.mins[i] >= self.maxes[i - 1]
        ]
        divisions = self.divisions

        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[(self._name, i - 1)] = (
                drop_overlap,
                dsk[(self._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[(self._name, i)] = (methods.concat, frames)
            frames = []
        return dsk


class Lengths(Expr):
    """Returns a tuple of partition lengths"""

    _parameters = ["frame"]

    @functools.cached_property
    def _meta(self):
        return tuple()

    def _divisions(self):
        return (None, None)

    def _simplify_down(self):
        if isinstance(self.frame, Elemwise):
            child = max(self.frame.dependencies(), key=lambda expr: expr.npartitions)
            return Lengths(child)

    def _layer(self):
        name = "part-" + self._name
        dsk = {
            (name, i): (len, (self.frame._name, i))
            for i in range(self.frame.npartitions)
        }
        dsk[(self._name, 0)] = (tuple, list(dsk.keys()))
        return dsk


class ResetIndex(Elemwise):
    """Reset the index of a Series or DataFrame"""

    _parameters = ["frame", "drop", "name"]
    _defaults = {"drop": False, "name": no_default}
    _keyword_only = ["drop", "name"]
    operation = M.reset_index
    _filter_passthrough = True
    _preserves_partitioning_information = True

    @functools.cached_property
    def _kwargs(self) -> dict:
        kwargs = {"drop": self.drop}
        if self.operand("name") is not no_default:
            kwargs.update({"name": self.operand("name")})
        return kwargs

    def _divisions(self):
        return (None,) * (self.frame.npartitions + 1)

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Filter) and self._filter_passthrough_available(
            parent, dependents
        ):
            parents = [
                p().columns
                for p in dependents[self._name]
                if p() is not None and not isinstance(p(), Filter)
            ]
            predicate = None
            if not set(flatten(parents, list)).issubset(set(self.frame.columns)):
                # one of the filters is the Index
                name = self.operand("name") or self.frame._meta.index.name
                if name is no_default and self.frame._meta.index.name is None:
                    name = "index"
                elif self.frame._meta.index.name is not None:
                    name = self.frame._meta.index.name
                # replace the projection of the former index with the actual index
                subs = Projection(self, name)
                predicate = parent.predicate.substitute(subs, Index(self.frame))
            elif self.frame.ndim == 1 and not self.operand("drop"):
                name = self.frame._meta.name
                # Avoid Projection since we are already a Series
                subs = Projection(self, name)
                predicate = parent.predicate.substitute(subs, self.frame)
            return self._filter_simplification(parent, predicate)

        if isinstance(parent, Projection):
            if self.frame.ndim == 1 and not self.drop:
                if isinstance(parent.operand("columns"), list):
                    # Don't bother, dimensionality changes are tricky here and
                    # potential improvement is tiny
                    return
                col = parent.operand("columns")
                if col in (self.name, "index", self.frame._meta.index.name):
                    return
                if all(
                    isinstance(d(), Projection) and d().operand("columns") == col
                    for d in dependents[self._name]
                ):
                    return type(self)(self.frame, True, self.name)
                return
            result = plain_column_projection(self, parent, dependents)
            if result is not None and not set(result.columns) == set(
                result.frame.columns
            ):
                result = result.substitute_parameters({"drop": True})
            return result


class AddPrefixSeries(Elemwise):
    _parameters = ["frame", "prefix"]
    operation = M.add_prefix
    _filter_passthrough = True
    _preserves_partitioning_information = True

    def _divisions(self):
        return tuple(self.prefix + str(division) for division in self.frame.divisions)


class AddSuffixSeries(AddPrefixSeries):
    _parameters = ["frame", "suffix"]
    operation = M.add_suffix
    _preserves_partitioning_information = True

    def _divisions(self):
        return tuple(str(division) + self.suffix for division in self.frame.divisions)


class AddPrefix(Elemwise):
    _parameters = ["frame", "prefix"]
    operation = M.add_prefix

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        return {
            f"{self.prefix}{c}"
            if not isinstance(c, tuple)
            else tuple(self.prefix + t for t in c)
            for c in self.frame.unique_partition_mapping_columns_from_shuffle
        }

    def _convert_columns(self, columns):
        len_prefix = len(self.prefix)
        return [col[len_prefix:] for col in columns]

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            columns = determine_column_projection(self, parent, dependents)
            columns = self._convert_columns(_convert_to_list(columns))
            if set(columns) == set(self.frame.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"),
            )


class AddSuffix(AddPrefix):
    _parameters = ["frame", "suffix"]
    operation = M.add_suffix

    @functools.cached_property
    def unique_partition_mapping_columns_from_shuffle(self):
        return {
            f"{c}{self.suffix}"
            if not isinstance(c, tuple)
            else tuple(t + self.suffix for t in c)
            for c in self.frame.unique_partition_mapping_columns_from_shuffle
        }

    def _convert_columns(self, columns):
        len_suffix = len(self.suffix)
        return [col[:-len_suffix] for col in columns]


class AssignIndex(Elemwise):
    _parameters = ["frame", "value"]
    operation = staticmethod(methods.assign_index)
    _preserves_partitioning_information = True

    def _divisions(self):
        return self.value.divisions


class Head(Expr):
    """Take the first `n` rows of the first partition"""

    _parameters = ["frame", "n", "npartitions"]
    _defaults = {"n": 5, "npartitions": 1}

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    @functools.cached_property
    def npartitions(self):
        return 1

    def _divisions(self):
        if self.operand("npartitions") <= -1:
            return self.frame.divisions[0], self.frame.divisions[-1]
        return (
            self.frame.divisions[0],
            self.frame.divisions[self.operand("npartitions")],
        )

    def _task(self, index: int):
        raise NotImplementedError()

    def _simplify_down(self):
        if isinstance(self.frame, Elemwise):
            operands = [
                Head(op, self.n, self.operand("npartitions"))
                if isinstance(op, Expr) and not isinstance(op, _DelayedExpr)
                else op
                for op in self.frame.operands
            ]
            return type(self.frame)(*operands)
        if isinstance(self.frame, Head):
            return Head(
                self.frame.frame, min(self.n, self.frame.n), self.operand("npartitions")
            )

    def _simplify_up(self, parent, dependents):
        from dask_expr import Repartition

        if isinstance(parent, Repartition) and parent.new_partitions == 1:
            return self

    def _lower(self):
        if not isinstance(self, BlockwiseHead):
            # Lower to Blockwise
            npartitions = self.operand("npartitions")
            if self.operand("npartitions") > self.frame.npartitions:
                raise ValueError(
                    f"only {self.frame.npartitions} partitions, head received {npartitions}"
                )
            partitions = self._partitions
            if is_index_like(self._meta):
                return BlockwiseHeadIndex(
                    Partitions(self.frame, partitions), self.n, safe=False
                )

            safe = True if npartitions == 1 and self.frame.npartitions != 1 else False
            frame = BlockwiseHead(
                Partitions(self.frame, partitions), self.n, npartitions, safe
            )
            if npartitions != 1:
                from dask_expr import Repartition

                safe = npartitions != self.frame.npartitions and npartitions != -1
                frame = BlockwiseHead(
                    Repartition(frame, new_partitions=1), self.n, 1, safe
                )
            return frame

    @property
    def _partitions(self):
        if isinstance(self, PartitionsFiltered):
            partitions = self.frame._partitions
        else:
            partitions = list(range(self.frame.npartitions))
        if self.operand("npartitions") > -1:
            partitions = partitions[: self.operand("npartitions")]
        return partitions


class BlockwiseHead(Head, Blockwise):
    """Take the first `n` rows of every partition

    Typically used after `Partitions(..., [0])` to take
    the first `n` rows of an entire collection.
    """

    _parameters = ["frame", "n", "npartitions", "safe"]
    _preserves_partitioning_information = True

    def _simplify_down(self):
        return

    def _simplify_up(self, parent, dependents):
        return

    @functools.cached_property
    def npartitions(self):
        return len(self._divisions()) - 1

    def _divisions(self):
        return self.frame.divisions[: len(self._partitions) + 1]

    def _task(self, index: int):
        if self.safe:
            op = safe_head
        else:
            op = M.head
        return (op, (self.frame._name, index), self.n)


class BlockwiseHeadIndex(BlockwiseHead):
    def _task(self, index: int):
        return (operator.getitem, (self.frame._name, index), slice(0, self.n))


class Tail(Expr):
    """Take the last `n` rows of the last partition"""

    _parameters = ["frame", "n"]
    _defaults = {"n": 5}

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    def _divisions(self):
        return self.frame.divisions[-2:]

    def _task(self, index: int):
        raise NotImplementedError()

    def _simplify_down(self):
        if isinstance(self.frame, Elemwise):
            operands = [
                Tail(op, self.n) if isinstance(op, Expr) else op
                for op in self.frame.operands
            ]
            return type(self.frame)(*operands)
        if isinstance(self.frame, Tail):
            return Tail(self.frame.frame, min(self.n, self.frame.n))

    def _simplify_up(self, parent, dependents):
        from dask_expr import Repartition

        if isinstance(parent, Repartition) and parent.new_partitions == 1:
            return self

    def _lower(self):
        if not isinstance(self, BlockwiseTail):
            # Lower to Blockwise
            if is_index_like(self._meta):
                return BlockwiseTailIndex(
                    Partitions(self.frame, [self.frame.npartitions - 1]), self.n
                )
            return BlockwiseTail(
                Partitions(self.frame, [self.frame.npartitions - 1]), self.n
            )


class BlockwiseTail(Tail, Blockwise):
    """Take the last `n` rows of every partition

    Typically used after `Partitions(..., [-1])` to take
    the last `n` rows of an entire collection.
    """

    _preserves_partitioning_information = True

    def _divisions(self):
        return self.frame.divisions

    def _task(self, index: int):
        return (M.tail, (self.frame._name, index), self.n)


class BlockwiseTailIndex(BlockwiseTail):
    def _task(self, index: int):
        return (operator.getitem, (self.frame._name, index), slice(-self.n, None))


class Binop(Elemwise):
    _parameters = ["left", "right"]

    def __str__(self):
        return f"{self.left} {self._operator_repr} {self.right}"

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            changed = False
            columns = determine_column_projection(self, parent, dependents)
            columns = _convert_to_list(columns)
            columns = [col for col in self.columns if col in columns]
            if (
                isinstance(self.left, Expr)
                and self.left.ndim > 1
                and self.left.columns != columns
            ):
                left = self.left[columns]  # TODO: filter just the correct columns
                changed = True
            else:
                left = self.left
            if (
                isinstance(self.right, Expr)
                and self.right.ndim > 1
                and self.right.columns != columns
            ):
                right = self.right[columns]  # TODO: filter just the correct columns
                changed = True
            else:
                right = self.right
            if not changed:
                return

            return type(parent)(type(self)(left, right), *parent.operands[1:])

    def _node_label_args(self):
        return [self.left, self.right]

    def _divisions(self):
        if is_index_like(self._meta):
            left_divisions = (
                pd.Series(self.left.divisions)
                if isinstance(self.left, Expr)
                else self.left
            )
            right_divisions = (
                pd.Series(self.right.divisions)
                if isinstance(self.right, Expr)
                else self.right
            )

            return tuple(self.operation(left_divisions, right_divisions))
        else:
            return super()._divisions()


class Add(Binop):
    operation = operator.add
    _operator_repr = "+"


class MethodOperator(Binop):
    _parameters = ["name", "left", "right", "axis", "level", "fill_value"]
    _defaults = {"axis": "columns", "level": None, "fill_value": None}
    _keyword_only = ["axis", "level", "fill_value"]

    @property
    def _operator_repr(self):
        return self.name

    @staticmethod
    def operation(name, left, right, **kwargs):
        return getattr(left, name)(right, **kwargs)


class Sub(Binop):
    operation = operator.sub
    _operator_repr = "-"


class Mul(Binop):
    operation = operator.mul
    _operator_repr = "*"

    def _simplify_down(self):
        if (
            isinstance(self.right, Mul)
            and isinstance(self.left, numbers.Number)
            and isinstance(self.right.left, numbers.Number)
        ):
            return (self.left * self.right.left) * self.right.right


class Pow(Binop):
    operation = operator.pow
    _operator_repr = "**"


class Div(Binop):
    operation = operator.truediv
    _operator_repr = "/"


class LT(Binop):
    operation = operator.lt
    _operator_repr = "<"


class BinOpSeries(Binop):
    _parameters = ["left", "right", "level", "fill_value"]
    _defaults = {"fill_value": None, "level": None}


class BinOpFrame(Binop):
    _parameters = ["left", "right", "axis"]
    _defaults = {"axis": 1}


class LTSeries(BinOpSeries):
    operation = M.lt
    _operator_repr = "<"


class LTFrame(BinOpFrame):
    operation = M.lt
    _operator_repr = "<"


class LESeries(BinOpSeries):
    operation = M.le
    _operator_repr = "<="


class LEFrame(BinOpFrame):
    operation = M.le
    _operator_repr = "<="


class GTSeries(BinOpSeries):
    operation = M.gt
    _operator_repr = ">"


class GTFrame(BinOpFrame):
    operation = M.gt
    _operator_repr = ">"


class GESeries(BinOpSeries):
    operation = M.ge
    _operator_repr = ">="


class GEFrame(BinOpFrame):
    operation = M.ge
    _operator_repr = ">="


class NESeries(BinOpSeries):
    operation = M.ne
    _operator_repr = "!="


class NEFrame(BinOpFrame):
    operation = M.ne
    _operator_repr = "!="


class EQSeries(BinOpSeries):
    operation = M.eq
    _operator_repr = "=="


class EQFrame(BinOpFrame):
    operation = M.eq
    _operator_repr = "=="


class LE(Binop):
    operation = operator.le
    _operator_repr = "<="


class GT(Binop):
    operation = operator.gt
    _operator_repr = ">"


class GE(Binop):
    operation = operator.ge
    _operator_repr = ">="


class EQ(Binop):
    operation = operator.eq
    _operator_repr = "=="


class NE(Binop):
    operation = operator.ne
    _operator_repr = "!="


class And(Binop):
    operation = operator.and_
    _operator_repr = "&"


class Or(Binop):
    operation = operator.or_
    _operator_repr = "|"


class XOr(Binop):
    operation = operator.xor
    _operator_repr = "^"


class Mod(Binop):
    operation = operator.mod
    _operator_repr = "%"


class FloorDiv(Binop):
    operation = operator.floordiv
    _operator_repr = "//"


class Unaryop(Elemwise):
    _parameters = ["frame"]

    def __str__(self):
        return f"{self._operator_repr} {self.frame}"

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            if isinstance(self.frame, Expr):
                return plain_column_projection(self, parent, dependents)
            else:
                frame = self.frame
            return type(self)(frame)

    def _node_label_args(self):
        return [self.frame]

    def _divisions(self):
        if is_index_like(self._meta):
            return (None,) * (self.frame.npartitions + 1)
        else:
            return super()._divisions()


class Invert(Unaryop):
    operation = operator.inv
    _operator_repr = "~"


class Neg(Unaryop):
    operation = operator.neg
    _operator_repr = "-"


class Pos(Unaryop):
    operation = operator.pos
    _operator_repr = "+"


class Partitions(Expr):
    """Select one or more partitions"""

    _parameters = ["frame", "partitions"]

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    def _divisions(self):
        divisions = []
        for part in self.partitions:
            divisions.append(self.frame.divisions[part])
        divisions.append(self.frame.divisions[part + 1])
        return tuple(divisions)

    def _task(self, index: int):
        return (self.frame._name, self.partitions[index])

    def _simplify_down(self):
        from dask_expr import SetIndexBlockwise
        from dask_expr._resample import ResampleAggregation

        if isinstance(self.frame, Blockwise) and not isinstance(
            self.frame, (BlockwiseIO, Fused, SetIndexBlockwise, ResampleAggregation)
        ):
            operands = [
                (
                    Partitions(op, self.partitions)
                    if (isinstance(op, Expr) and not self.frame._broadcast_dep(op))
                    else op
                )
                for op in self.frame.operands
            ]
            return type(self.frame)(*operands)
        elif isinstance(self.frame, PartitionsFiltered):
            if self.frame._partitions:
                partitions = [self.frame._partitions[p] for p in self.partitions]
            else:
                partitions = self.partitions
            # We assume that expressions defining a special "_partitions"
            # parameter can internally capture the same logic as `Partitions`
            return self.frame.substitute_parameters({"_partitions": partitions})

    def _node_label_args(self):
        return [self.frame, self.partitions]


class PartitionsFiltered(Expr):
    """Mixin class for partition filtering

    A ``PartitionsFiltered`` subclass must define a ``_partitions`` parameter. When
    ``_partitions`` is defined, the following expressions must produce the same output
    for :cls:`PartitionsFiltered`:

    - ``cls(expr: Expr, ..., _partitions)``
    - ``Partitions(cls(expr: Expr, ...), _partitions)``

    In order to leverage the default ``Expr._layer`` method, subclasses should define
    ``_filtered_task`` instead of ``_task``.
    """

    @property
    def _filtered(self) -> bool:
        """Whether output partitions have been filtered"""
        return self.operand("_partitions") is not None

    @property
    def _partitions(self) -> list | tuple | range:
        """Selected partition indices"""
        if self._filtered:
            return self.operand("_partitions")
        else:
            return range(self.npartitions)

    @functools.cached_property
    def divisions(self):
        # Common case: Use self._divisions()
        full_divisions = super().divisions
        if not self._filtered:
            return full_divisions

        # Specific case: Specific partitions were selected
        new_divisions = []
        for part in self._partitions:
            new_divisions.append(full_divisions[part])
        new_divisions.append(full_divisions[part + 1])
        return tuple(new_divisions)

    @property
    def npartitions(self):
        if self._filtered:
            return len(self._partitions)
        return super().npartitions

    def _task(self, index: int):
        return self._filtered_task(self._partitions[index])

    def _filtered_task(self, index: int):
        raise NotImplementedError()


class _DelayedExpr(Expr):
    # Wraps a Delayed object to make it an Expr for now. This is hacky and we should
    # integrate this properly...
    # TODO
    _parameters = ["obj"]

    def __init__(self, obj):
        self.obj = obj
        self.operands = [obj]

    def __str__(self):
        return f"{type(self).__name__}({str(self.obj)})"

    @property
    def _name(self):
        return self.obj.key

    def _layer(self) -> dict:
        dc = self.obj.__dask_optimize__(self.obj.dask, self.obj.key).to_dict().copy()
        dc[(self.obj.key, 0)] = dc[self.obj.key]
        dc.pop(self.obj.key)
        return dc

    def _divisions(self):
        return (None, None)

    @property
    def ndim(self):
        return 0


class DelayedsExpr(Expr):
    _parameters = []

    def __init__(self, *delayed_objects):
        self.operands = delayed_objects

    def __str__(self):
        return f"{type(self).__name__}({str(self.operands[0])})"

    @property
    def _name(self):
        return "delayed-container-" + _tokenize_deterministic(*self.operands)

    def _layer(self) -> dict:
        dask = {}
        for i, obj in enumerate(self.operands):
            dc = obj.__dask_optimize__(obj.dask, obj.key).to_dict().copy()
            dc[(self._name, i)] = dc[obj.key]
            dc.pop(obj.key)
            dask.update(dc)
        return dask

    def _divisions(self):
        return (None,) * (len(self.operands) + 1)

    @property
    def ndim(self):
        return 0


@normalize_token.register(Expr)
def normalize_expression(expr):
    return expr._name


def optimize_until(expr: Expr, stage: core.OptimizerStage) -> Expr:
    result = expr
    if stage == "logical":
        return result

    # Simplify
    expr = result.simplify()
    if stage == "simplified-logical":
        return expr

    # Manipulate Expression to make it more efficient
    expr = expr.rewrite(kind="tune")
    if stage == "tuned-logical":
        return expr

    # Lower
    expr = expr.lower_completely()
    if stage == "physical":
        return expr

    # Simplify again
    expr = expr.simplify()
    if stage == "simplified-physical":
        return expr

    # Final graph-specific optimizations
    expr = optimize_blockwise_fusion(expr)
    if stage == "fused":
        return expr

    raise ValueError(f"Stage {stage!r} not supported.")


def optimize(expr: Expr, fuse: bool = True) -> Expr:
    """High level query optimization

    This leverages three optimization passes:

    1.  Class based simplification using the ``_simplify`` function and methods
    2.  Blockwise fusion

    Parameters
    ----------
    expr:
        Input expression to optimize
    fuse:
        whether or not to turn on blockwise fusion

    See Also
    --------
    simplify
    optimize_blockwise_fusion
    """
    stage: core.OptimizerStage = "fused" if fuse else "simplified-physical"

    return optimize_until(expr, stage)


def is_broadcastable(dfs, s):
    """
    This Series is broadcastable against another dataframe in the sequence
    """

    def compare(s, df):
        try:
            return s.divisions == (min(df.columns), max(df.columns))
        except (TypeError, ValueError):
            return False

    return (
        s.ndim == 1
        and s.npartitions == 1
        and s.known_divisions
        and any(compare(s, df) for df in dfs if df.ndim == 2)
        or s.ndim == 0
    )


def are_co_aligned(*exprs):
    """Do inputs come from the same parents, modulo blockwise?"""

    from dask_expr._cumulative import CumulativeAggregations
    from dask_expr._reductions import Reduction

    seen = set()
    # Scalars can always be broadcasted
    stack = [e for e in exprs if e.ndim > 0]
    ancestors = []
    while stack:
        e = stack.pop()
        if e._name in seen:
            continue
        seen.add(e._name)

        if isinstance(e, IO):
            ancestors.append(e)
        elif e.ndim == 0:
            # Scalars are valid ancestors that are always broadcastable,
            # so don't walk through them
            continue
        elif isinstance(e, (_DelayedExpr, Isin)):
            continue
        elif isinstance(e, (Blockwise, CumulativeAggregations, Reduction)):
            # TODO: Capture this in inheritance logic
            dependencies = e.dependencies()
            stack.extend(dependencies)
        else:
            ancestors.append(e)

    unique_ancestors = {
        # Account for column projection within IO expressions
        _tokenize_partial(item, ["columns", "_series", "_dataset_info_cache"])
        for item in ancestors
    }
    # Don't check divisions or npartitions at all
    return len(unique_ancestors) <= 1


## Utilites for Expr fusion


def is_valid_blockwise_op(expr):
    return isinstance(expr, Blockwise) and not isinstance(
        expr, (FromPandas, FromArray, FromDelayed, ResourceBarrier)
    )


def optimize_blockwise_fusion(expr):
    """Traverse the expression graph and apply fusion"""

    def _fusion_pass(expr):
        # Full pass to find global dependencies
        seen = set()
        stack = [expr]
        dependents = defaultdict(set)
        dependencies = {}
        expr_mapping = {}

        while stack:
            next = stack.pop()

            if next._name in seen:
                continue
            seen.add(next._name)

            if is_valid_blockwise_op(next):
                dependencies[next._name] = set()
                if next._name not in dependents:
                    dependents[next._name] = set()
                    expr_mapping[next._name] = next

            for operand in next.operands:
                if isinstance(operand, Expr):
                    stack.append(operand)
                    if is_valid_blockwise_op(operand):
                        if next._name in dependencies:
                            dependencies[next._name].add(operand._name)
                        dependents[operand._name].add(next._name)
                        expr_mapping[operand._name] = operand
                        expr_mapping[next._name] = next

        # Traverse each "root" until we find a fusable sub-group.
        # Here we use root to refer to a Blockwise Expr node that
        # has no Blockwise dependents
        roots = [
            expr_mapping[k]
            for k, v in dependents.items()
            if v == set()
            or all(not is_valid_blockwise_op(expr_mapping[_expr]) for _expr in v)
        ]
        while roots:
            root = roots.pop()
            seen = set()
            stack = [root]
            group = []
            while stack:
                next = stack.pop()

                if next._name in seen:
                    continue
                seen.add(next._name)

                group.append(next)
                for dep_name in dependencies[next._name]:
                    dep = expr_mapping[dep_name]

                    stack_names = {s._name for s in stack}
                    group_names = {g._name for g in group}
                    if (
                        dep.npartitions == root.npartitions or next._broadcast_dep(dep)
                    ) and not (dependents[dep._name] - stack_names - group_names):
                        # All of deps dependents are contained
                        # in the local group (or the local stack
                        # of expr nodes that we know we will be
                        # adding to the local group).
                        # All nodes must also have the same number
                        # of partitions, since broadcasting within
                        # a group is not allowed.
                        stack.append(dep)
                    elif dependencies[dep._name] and dep._name not in [
                        r._name for r in roots
                    ]:
                        # Couldn't fuse dep, but we may be able to
                        # use it as a new root on the next pass
                        roots.append(dep)

            # Replace fusable sub-group
            if len(group) > 1:
                group_deps = []
                local_names = [_expr._name for _expr in group]
                for _expr in group:
                    group_deps += [
                        operand
                        for operand in _expr.dependencies()
                        if operand._name not in local_names
                    ]
                _ret = expr.substitute(group[0], Fused(group, *group_deps))
                return _ret, not roots

        # Return original expr if no fusable sub-groups were found
        return expr, True

    while True:
        original_name = expr._name
        expr, done = _fusion_pass(expr)
        if done or expr._name == original_name:
            break

    return expr


class Diff(MapOverlap):
    _parameters = ["frame", "periods"]
    _defaults = {"periods": 1}
    func = M.diff
    enforce_metadata = True
    transform_divisions = False
    clear_divisions = False
    align_dataframes = True

    def _divisions(self):
        return self.frame.divisions

    @functools.cached_property
    def _meta(self):
        return make_meta(meta_nonempty(self.frame._meta).diff(**self.kwargs))

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            return plain_column_projection(self, parent, dependents)

    @functools.cached_property
    def kwargs(self):
        return dict(periods=self.periods)

    @property
    def before(self):
        return self.periods if self.periods > 0 else 0

    @property
    def after(self):
        return 0 if self.periods > 0 else -self.periods


class FillnaCheck(Blockwise):
    _parameters = ["frame", "method", "skip_check"]
    operation = staticmethod(methods.fillna_check)
    _projection_passthrough = True

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    def _task(self, index: int):
        args = [self._blockwise_arg(op, index) for op in self._args]
        args[-1] = index != self.skip_check(self.frame)
        return (self.operation,) + tuple(args)


class FFill(MapOverlap):
    _parameters = ["frame", "limit"]
    _defaults = {"limit": None}
    func = M.ffill
    enforce_metadata = True
    transform_divisions = False
    clear_divisions = False
    align_dataframes = True

    def _divisions(self):
        return self.frame.divisions

    @functools.cached_property
    def _meta(self):
        return self.frame._meta

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            return plain_column_projection(self, parent, dependents)

    @functools.cached_property
    def kwargs(self):
        return dict(limit=self.limit)

    @property
    def before(self):
        return 1 if self.limit is None else self.limit

    @property
    def after(self):
        return 0


class BFill(FFill):
    func = M.bfill

    @property
    def before(self):
        # bfill is the opposite direction of ffill, so
        # we swap before with after of ffill.
        return super().after

    @property
    def after(self):
        # bfill is the opposite direction of ffill, so
        # we swap after with before of ffill.
        return super().before


class Shift(MapOverlap):
    _parameters = ["frame", "periods", "freq"]
    _defaults = {"periods": 1, "freq": None}

    func = M.shift
    enforce_metadata = True
    transform_divisions = False
    align_dataframes = True

    @functools.cached_property
    def clear_divisions(self):
        # TODO We can do better if freq is given, but this needs adjustments in
        #  map_partitions
        return True if self._divisions()[0] is None or self.freq is not None else False

    def _divisions(self):
        if self.freq is None:
            return self.frame.divisions
        divisions = _calc_maybe_new_divisions(self.frame, self.periods, self.freq)
        if divisions is None:
            divisions = (None,) * (self.frame.npartitions + 1)
        return divisions

    @functools.cached_property
    def _meta(self):
        return make_meta(meta_nonempty(self.frame._meta).shift(**self.kwargs))

    @functools.cached_property
    def kwargs(self):
        return dict(periods=self.periods, freq=self.freq)

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection):
            return plain_column_projection(self, parent, dependents)

    @property
    def before(self):
        return self.periods if self.periods > 0 else 0

    @property
    def after(self):
        return 0 if self.periods > 0 else -self.periods


class ShiftIndex(Blockwise):
    _parameters = ["frame", "periods", "freq"]
    _defaults = {"periods": 1, "freq": None}
    _keyword_only = ["freq"]
    operation = M.shift

    def _divisions(self):
        freq = self.freq
        if freq is None:
            freq = self._meta.freq
        divisions = _calc_maybe_new_divisions(self.frame, self.periods, freq)
        if divisions is None:
            divisions = (None,) * (self.frame.npartitions + 1)
        return divisions

    @functools.cached_property
    def _kwargs(self) -> dict:
        return {"freq": self.freq} if self.freq is not None else {}


class MaybeAlignPartitions(Expr):
    _projection_passthrough = False
    _expr_cls = None

    def _divisions(self):
        if {df.npartitions for df in self.args} == {1}:
            divs = []
            for df in self.args:
                divs.extend(list(df.divisions))
            try:
                return min(divs), max(divs)
            except TypeError:
                # either unknown divisions or int-str mix
                return None, None
        return calc_divisions_for_align(*self.args)

    def _simplify_up(self, parent, dependents):
        if isinstance(parent, Projection) and self._projection_passthrough:
            return plain_column_projection(self, parent, dependents)

    @functools.cached_property
    def args(self):
        dfs = [op for op in self.operands if isinstance(op, Expr)]
        return [op for op in dfs if not is_broadcastable(dfs, op)]

    def _lower(self):
        # This can be expensive when something that has expensive division
        # calculation is in the Expression
        dfs = self.args
        if (
            len(dfs) == 1
            or all(
                dfs[0].divisions == df.divisions and df.known_divisions for df in dfs
            )
            or len(self.divisions) == 2
        ):
            return self._expr_cls(*self.operands)
        elif self.divisions[0] is None:
            # We have to shuffle
            npartitions = max(df.npartitions for df in dfs)
            dtypes = {df._meta.index.dtype for df in dfs}
            if not _are_dtypes_shuffle_compatible(dtypes):
                raise TypeError(
                    "DataFrames are not aligned. We need to shuffle to align partitions "
                    "with each other. This is not possible because the indexes of the "
                    f"DataFrames have differing dtypes={dtypes}. Please ensure that "
                    "all Indexes have the same dtype or align manually for this to "
                    "work."
                )

            from dask_expr._shuffle import RearrangeByColumn

            args = [
                (
                    RearrangeByColumn(df, None, npartitions, index_shuffle=True)
                    if isinstance(df, Expr)
                    else df
                )
                for df in self.operands
            ]
            return self._expr_cls(*args)

        args = maybe_align_partitions(*self.operands, divisions=self.divisions)
        return self._expr_cls(*args)

    @functools.cached_property
    def _meta(self):
        return self._expr_cls(*self.operands)._meta


def _are_dtypes_shuffle_compatible(dtypes):
    if len(dtypes) == 1:
        return True
    if all(pd.api.types.is_numeric_dtype(d) for d in dtypes):
        return True
    return False


class CombineFirstAlign(MaybeAlignPartitions):
    _parameters = ["frame", "other"]
    _expr_cls = CombineFirst

    def _simplify_up(self, parent, dependents):
        # TODO: de-duplicate
        if isinstance(parent, Projection):
            columns = determine_column_projection(self, parent, dependents)
            frame_columns = [col for col in self.frame.columns if col in columns]
            other_columns = [col for col in self.other.columns if col in columns]
            if (
                self.frame.columns == frame_columns
                and self.other.columns == other_columns
            ):
                return

            return type(parent)(
                type(self)(self.frame[frame_columns], self.other[other_columns]),
                *parent.operands[1:],
            )


class FillnaAlign(MaybeAlignPartitions):
    _projection_passthrough = True
    _parameters = ["frame", "value"]
    _expr_cls = Fillna


class AlignAlignPartitions(MaybeAlignPartitions):
    _parameters = ["frame", "other", "join", "axis", "fill_value"]
    _expr_cls = _Align


class CombineSeriesAlign(MaybeAlignPartitions):
    _parameters = ["frame", "other", "func", "fill_value"]
    _expr_cls = CombineSeries


class CombineFrameAlign(MaybeAlignPartitions):
    _parameters = ["frame", "other", "func", "fill_value", "overwrite"]
    _expr_cls = CombineSeries


class FilterAlign(MaybeAlignPartitions):
    _projection_passthrough = True
    _filter_passthrough = True
    _parameters = ["frame", "predicate"]
    _expr_cls = Filter

    def _simplify_up(self, parent, dependents):
        return Filter._simplify_up(self, parent, dependents)


class AssignAlign(MaybeAlignPartitions):
    _parameters = ["frame", "column", "value"]
    _expr_cls = Assign

    def _simplify_up(self, parent, dependents):
        # TODO: de-duplicate
        if isinstance(parent, Projection):
            columns = determine_column_projection(self, parent, dependents)
            if not isinstance(columns, list):
                columns = [columns]

            cols = set(columns) - {self.column}
            if cols == set(self.frame.columns):
                # Protect against pushing the same projection twice
                return

            diff = {self.column} - set(columns)
            if len(diff) == 1:
                return type(parent)(self.frame, *parent.operands[1:])
            else:
                new_args = self.operands[1:]

            columns = [col for col in self.frame.columns if col in cols]
            return type(parent)(
                type(self)(self.frame[sorted(columns)], *new_args),
                *parent.operands[1:],
            )


class MaskAlign(MaybeAlignPartitions):
    _parameters = ["frame", "cond", "other"]
    _expr_cls = Mask


class WhereAlign(MaskAlign):
    _expr_cls = Where


class MapAlign(MaybeAlignPartitions):
    _parameters = ["frame", "other", "op", "na_action", "meta"]
    _projection_passthrough = False
    _expr_cls = Map


class MapIndexAlign(MapAlign):
    _parameters = MaskAlign._parameters + ["is_monotonic"]


class OpAlignPartitions(MaybeAlignPartitions):
    _parameters = ["frame", "other", "op"]
    _projection_passthrough = True

    @functools.cached_property
    def _meta(self):
        return getattr(self.frame._meta, self.op)(self.other._meta)

    def _lower(self):
        # This can be expensive when something that has expensive division
        # calculation is in the Expression
        dfs = self.args
        if (
            len(dfs) == 1
            or all(dfs[0].divisions == df.divisions for df in dfs)
            or len(self.divisions) == 2
        ):
            return self._op(self.frame, self.op, self.other, *self.operands[3:])

        from dask_expr._repartition import RepartitionDivisions

        frame = RepartitionDivisions(
            self.frame, new_divisions=self.divisions, force=True
        )
        other = RepartitionDivisions(
            self.other, new_divisions=self.divisions, force=True
        )
        return self._op(frame, self.op, other, *self.operands[3:])

    @staticmethod
    def _op(frame, op, other, *args, **kwargs):
        return getattr(frame, op)(other)


class MethodOperatorAlign(OpAlignPartitions):
    _parameters = ["frame", "other", "op", "axis", "level", "fill_value"]

    @staticmethod
    def _op(frame, op, other, *args, **kwargs):
        return MethodOperator(op, frame, other, *args, **kwargs)


class UFuncAlign(MaybeAlignPartitions):
    _parameters = ["frame", "func", "meta", "kwargs"]
    enforce_metadata = False

    def __str__(self):
        return f"UFunc({funcname(self.func)})"

    @functools.cached_property
    def args(self):
        return self.operands[len(self._parameters) :]

    @functools.cached_property
    def _dfs(self):
        return [df for df in self.args if isinstance(df, Expr)]

    @functools.cached_property
    def _meta(self):
        if self.operand("meta") is not no_default:
            return self.operand("meta")
        return _get_meta_ufunc(self._dfs, self.args, self.func)

    def _lower(self):
        args = maybe_align_partitions(*self.args, divisions=self._divisions())
        dfs = [x for x in args if isinstance(x, Expr) and x.ndim > 0]
        return UFuncElemwise(dfs[0], self.func, self._meta, False, self.kwargs, *args)


class Fused(Blockwise):
    """Fused ``Blockwise`` expression

    A ``Fused`` corresponds to the fusion of multiple
    ``Blockwise`` expressions into a single ``Expr`` object.
    Before graph-materialization time, the behavior of this
    object should be identical to that of the first element
    of ``Fused.exprs`` (i.e. the top-most expression in
    the fused group).

    Parameters
    ----------
    exprs : List[Expr]
        Group of original ``Expr`` objects being fused together.
    *dependencies:
        List of external ``Expr`` dependencies. External-``Expr``
        dependencies correspond to any ``Expr`` operand that is
        not already included in ``exprs``. Note that these
        dependencies should be defined in the order of the ``Expr``
        objects that require them (in ``exprs``). These
        dependencies do not include literal operands, because those
        arguments should already be captured in the fused subgraphs.
    """

    _parameters = ["exprs"]

    @functools.cached_property
    def _meta(self):
        return self.exprs[0]._meta

    def _tree_repr_lines(self, indent=0, recursive=True):
        header = f"Fused({self._name[-5:]}):"
        if not recursive:
            return [header]

        seen = set()
        lines = []
        stack = [(self.exprs[0], 2)]
        fused_group = [_expr._name for _expr in self.exprs]
        dependencies = {dep._name: dep for dep in self.dependencies()}
        while stack:
            expr, _indent = stack.pop()

            if expr._name in seen:
                continue
            seen.add(expr._name)

            line = expr._tree_repr_lines(_indent, recursive=False)[0]
            lines.append(line.replace(" ", "|", 1))
            for dep in expr.dependencies():
                if dep._name in fused_group:
                    stack.append((dep, _indent + 2))
                elif dep._name in dependencies:
                    dependencies.pop(dep._name)
                    lines.extend(dep._tree_repr_lines(_indent + 2))

        for dep in dependencies.values():
            lines.extend(dep._tree_repr_lines(2))

        lines = [header] + lines
        lines = [" " * indent + line for line in lines]

        return lines

    def __str__(self):
        exprs = sorted(self.exprs, key=M._depth)
        names = [expr._name.split("-")[0] for expr in exprs]
        if len(names) > 4:
            return names[0] + "-fused-" + names[-1]
        else:
            return "-".join(names)

    @functools.cached_property
    def _name(self):
        return f"{str(self)}-{_tokenize_deterministic(*self.operands)}"

    def _divisions(self):
        return self.exprs[0]._divisions()

    def _broadcast_dep(self, dep: Expr):
        # Always broadcast single-partition dependencies in Fused
        return dep.npartitions == 1

    def _task(self, index):
        graph = {self._name: (self.exprs[0]._name, index)}
        for _expr in self.exprs:
            if isinstance(_expr, Fused):
                subgraph, name = _expr._task(index)[1:3]
                graph.update(subgraph)
                graph[(name, index)] = name
            elif self._broadcast_dep(_expr):
                # When _expr is being broadcasted, we only
                # want to define a fused task for index 0
                graph[(_expr._name, 0)] = _expr._task(0)
            else:
                graph[(_expr._name, index)] = _expr._task(index)

        for i, dep in enumerate(self.dependencies()):
            graph[self._blockwise_arg(dep, index)] = "_" + str(i)

        return (
            Fused._execute_task,
            graph,
            self._name,
        ) + tuple(self._blockwise_arg(dep, index) for dep in self.dependencies())

    @staticmethod
    def _execute_task(graph, name, *deps):
        for i, dep in enumerate(deps):
            graph["_" + str(i)] = dep
        return dask.core.get(graph, name)


# Used for sorting with None
@functools.total_ordering
class MinType:
    def __le__(self, other):
        return True


def determine_column_projection(expr, parent, dependents, additional_columns=None):
    if isinstance(parent, Index):
        column_union = []
    else:
        column_union = parent.columns.copy()
    parents = [x() for x in dependents[expr._name] if x() is not None]

    seen = set()
    for p in parents:
        if p._name in seen:
            continue
        seen.add(p._name)

        column_union.extend(p._projection_columns)

    if additional_columns is not None:
        column_union.extend(flatten(additional_columns, container=list))

    # We can end up with MultiIndex columns from groupby ops, needs to be
    # accounted for in the sort
    flattened_columns = set(column_union)
    try:
        column_union = sorted(flattened_columns)
    except TypeError:
        # mixed type columns
        column_union = _sort_mixed(pd.Index(list(flattened_columns))).tolist()
    if (
        len(column_union) == 1
        and parent.ndim == 1
        and all(p.ndim == 1 for p in parents)
    ):
        return column_union[0]
    return column_union


def _sort_mixed(values):
    """order ints before strings before nulls in 1d arrays"""
    str_pos = np.array([isinstance(x, str) for x in values], dtype=bool)
    tuple_pos = np.array([isinstance(x, tuple) for x in values], dtype=bool)
    null_pos = np.array([pd.isna(x) for x in values], dtype=bool)
    num_pos = ~str_pos & ~null_pos & ~tuple_pos
    str_argsort = np.argsort(values[str_pos])
    tuple_argsort = np.argsort(values[tuple_pos])
    num_argsort = np.argsort(values[num_pos])
    # convert boolean arrays to positional indices, then order by underlying values
    str_locs = str_pos.nonzero()[0].take(str_argsort)
    tuple_locs = tuple_pos.nonzero()[0].take(tuple_argsort)
    num_locs = num_pos.nonzero()[0].take(num_argsort)
    null_locs = null_pos.nonzero()[0]
    locs = np.concatenate([num_locs, str_locs, tuple_locs, null_locs])
    return values.take(locs)


def plain_column_projection(expr, parent, dependents, additional_columns=None):
    column_union = determine_column_projection(
        expr, parent, dependents, additional_columns=additional_columns
    )
    if isinstance(column_union, list):
        column_union = [col for col in expr.frame.columns if col in column_union]
    elif column_union not in expr.frame.columns:
        # we are accesing the index
        column_union = []

    if column_union == expr.frame.columns:
        return
    result = type(expr)(expr.frame[column_union], *expr.operands[1:])
    if column_union == parent.operand("columns"):
        return result
    return type(parent)(result, parent.operand("columns"))


def is_filter_pushdown_available(expr, parent, dependents, allow_reduction=True):
    parents = [x() for x in dependents[expr._name] if x() is not None]
    filters = {e._name for e in parents if isinstance(e, Filter)}
    if len(filters) != 1:
        # Don't push down if not exactly one Filter
        return False
    if len(parents) == 1:
        return True

    # We have to see if the non-filter ops are all exclusively part of the predicates
    others = {e._name for e in parents if not isinstance(e, Filter)}
    return _check_dependents_are_predicates(
        expr, others, parent, dependents, allow_reduction
    )


def rewrite_filters(predicate):
    """Rewriting a filter to decompose OR clauses. If a predicate part is part of
    all OR clauses, we can move it to the front so that we can push it down.
    """
    or_components = _get_predicate_components(predicate, [])
    if len(or_components) == 1:
        return predicate
    result = _replace_common_or_components(or_components[0], or_components[1:])
    if result is None:
        return predicate
    return result


def _get_predicate_components(predicate, components, type_=Or):
    if not isinstance(predicate, type_):
        components.append(predicate)
        return components
    if isinstance(predicate.left, type_):
        components = _get_predicate_components(predicate.left, components, type_)
    else:
        components.append(predicate.left)
    if isinstance(predicate.right, type_):
        components = _get_predicate_components(predicate.right, components, type_)
    else:
        components.append(predicate.right)
    return components


def _convert_mapping(components):
    return dict(zip([e._name for e in components], components))


def _replace_common_or_components(expr, or_components):
    and_component = _get_predicate_components(expr, [], type_=And)
    mapping = _convert_mapping(and_component)
    and_components = [
        _get_predicate_components(c, [], type_=And) for c in or_components
    ]
    and_components = list(map(_convert_mapping, and_components))

    replacements = []
    for c in mapping.keys():
        if all(c in comp for comp in and_components):
            # We can pull this component out if it's part of all or components
            replacements.append(c)
    if len(replacements) == 0:
        # exit if we can't replace anything
        return

    outer_component = mapping[replacements[0]]
    for r in replacements[1:]:
        # construct the outer component
        outer_component = outer_component & mapping[r]

    #
    result_components = []
    for comp in [mapping] + and_components:
        keep_components = [c for c in comp if c not in replacements]
        if len(keep_components) == 0:
            # Just return outer_component if we can replace a whole OR component
            return outer_component
        result_component = comp[keep_components[0]]
        for c in keep_components[1:]:
            result_component = result_component & comp[c]
        result_components.append(result_component)

    or_component = result_components[0]
    for c in result_components[1:]:
        or_component = or_component | c
    return outer_component & or_component


def _check_dependents_are_predicates(
    expr, other_names, parent: Expr, dependents, allow_reduction=True
):
    # singleton approach should make this easier

    # Walk down the predicate side from the filter to see if we can arrive at
    # other_names without hitting an expression that has other dependents that
    # are not part of the predicate, see test_filter_pushdown_unavailable
    allowed_expressions = {parent._name}
    stack = parent.dependencies()
    seen = set()
    all_dependents = set()

    while stack:
        e = stack.pop()
        if expr._name == e._name:
            continue

        if e._name in seen:
            continue
        seen.add(e._name)

        if isinstance(e, _DelayedExpr):
            continue

        all_dependents.update(
            {x()._name for x in dependents[e._name] if x() is not None}
        )

        if not allow_reduction:
            if isinstance(e, (ApplyConcatApply, TreeReduce, ShuffleReduce)):
                return False

        allowed_expressions.add(e._name)
        stack.extend(e.dependencies())

    return all_dependents.issubset(allowed_expressions) and other_names.issubset(
        allowed_expressions
    )


def calc_divisions_for_align(*exprs, allow_shuffle=True):
    dfs = [df for df in exprs if isinstance(df, Expr) and df.ndim > 0]
    if not all(df.known_divisions for df in dfs):
        return (None,) * (max(df.npartitions for df in dfs) + 1)
    divisions = list(unique(merge_sorted(*[df.divisions for df in dfs])))
    if len(divisions) == 1:  # single value for index
        divisions = (divisions[0], divisions[0])
    return divisions


def maybe_align_partitions(*exprs, divisions):
    from dask_expr._repartition import Repartition

    return [
        (
            Repartition(df, new_divisions=divisions, force=True)
            if isinstance(df, Expr) and df.ndim > 0
            else df
        )
        for df in exprs
    ]


def _extract_meta(x, nonempty=False):
    """
    Extract internal cache data (``_meta``) from dd.DataFrame / dd.Series
    """
    if isinstance(x, Expr):
        return meta_nonempty(x._meta) if nonempty else x._meta
    elif isinstance(x, list):
        return [_extract_meta(_x, nonempty) for _x in x]
    elif isinstance(x, tuple):
        return tuple(_extract_meta(_x, nonempty) for _x in x)
    elif isinstance(x, dict):
        res = {}
        for k in x:
            res[k] = _extract_meta(x[k], nonempty)
        return res
    elif hasattr(x, "expr"):
        return _extract_meta(x.expr, nonempty)
    else:
        return x


def _emulate(func, *args, udf=False, **kwargs):
    """
    Apply a function using args / kwargs. If arguments contain dd.DataFrame /
    dd.Series, using internal cache (``_meta``) for calculation
    """
    with raise_on_meta_error(funcname(func), udf=udf):
        return func(*_extract_meta(args, True), **_extract_meta(kwargs, True))


def _get_meta_map_partitions(args, dfs, func, kwargs, meta, parent_meta):
    """
    Helper to generate metadata for map_partitions and map_overlap output.
    """
    meta_index = getattr(make_meta(dfs[0]), "index", None) if dfs else None
    if parent_meta is None and dfs:
        parent_meta = dfs[0]._meta
    if meta is no_default:
        # Use non-normalized kwargs here, as we want the real values (not
        # delayed values)
        a = [meta_nonempty(arg._meta) if isinstance(arg, Expr) else arg for arg in args]
        meta = _emulate(func, *a, udf=True, **kwargs)
        meta_is_emulated = True
    else:
        meta = make_meta(meta, index=meta_index, parent_meta=parent_meta)
        meta_is_emulated = False

    if not (has_parallel_type(meta) or is_arraylike(meta) and meta.shape) and not all(
        isinstance(arg, Expr) and arg.ndim == 0 for arg in args
    ):
        if not meta_is_emulated:
            warnings.warn(
                "Meta is not valid, `map_partitions` and `map_overlap` expects output to be a pandas object. "
                "Try passing a pandas object as meta or a dict or tuple representing the "
                "(name, dtype) of the columns. In the future the meta you passed will not work.",
                FutureWarning,
            )
        # If `meta` is not a pandas object, the concatenated results will be a
        # different type
        meta = make_meta(_concat([meta]), index=meta_index)

    # Ensure meta is empty series
    meta = make_meta(meta, parent_meta=parent_meta)

    return meta


from dask_expr._reductions import (
    All,
    Any,
    ApplyConcatApply,
    Count,
    IdxMax,
    IdxMin,
    Max,
    Mean,
    Min,
    Mode,
    NBytes,
    NuniqueApprox,
    Prod,
    ShuffleReduce,
    Size,
    Sum,
    TreeReduce,
    Var,
)
from dask_expr.io import IO, BlockwiseIO, FromArray, FromPandas
from dask_expr.io._delayed import FromDelayed