Class ProhibitReuse

dask/_expr.py:1353–1420

An expression that guarantees that all keys are suffixes with a unique id. This can be used to break a common subexpression apart.

class ProhibitReuse(Expr):
    """
    An expression that guarantees that all keys are suffixes with a unique id.
    This can be used to break a common subexpression apart.
    """

    _parameters = ["expr"]
    _ALLOWED_TYPES = [HLGExpr, LLGExpr, HLGFinalizeCompute, _HLGExprSequence]

    def __dask_keys__(self):
        return self._modify_keys(self.expr.__dask_keys__())

    @staticmethod
    def _identity(obj):
        return obj

    @functools.cached_property
    def _suffix(self):
        return uuid.uuid4().hex

    def _modify_keys(self, k):
        if isinstance(k, list):
            return [self._modify_keys(kk) for kk in k]
        elif isinstance(k, tuple):
            return (self._modify_keys(k[0]),) + k[1:]
        elif isinstance(k, (int, float)):
            k = str(k)
        return f"{k}-{self._suffix}"

    def _simplify_down(self):
        # FIXME: Shuffling cannot be rewritten since the barrier key is
        # hardcoded. Skipping this here should do the trick most of the time
        if not isinstance(
            self.expr,
            tuple(self._ALLOWED_TYPES),
        ):
            return self.expr

    def __dask_graph__(self):
        try:
            from distributed.shuffle._core import P2PBarrierTask
        except ModuleNotFoundError:
            P2PBarrierTask = type(None)
        dsk = convert_legacy_graph(self.expr.__dask_graph__())

        subs = {old_key: self._modify_keys(old_key) for old_key in dsk}
        dsk2 = {}
        for old_key, new_key in subs.items():
            t = dsk[old_key]
            if isinstance(t, P2PBarrierTask):
                warnings.warn(
                    "Cannot block reusing for graphs including a "
                    "P2PBarrierTask. This may cause unexpected results. "
                    "This typically happens when converting a dask "
                    "DataFrame to delayed objects.",
                    UserWarning,
                )
                return dsk