Method _dask_postpersist

xarray/core/dataset.py:721–770

(
        self, dsk: Mapping, *, rename: Mapping[str, str] | None = None
    )

        )

    def _dask_postpersist(
        self, dsk: Mapping, *, rename: Mapping[str, str] | None = None
    ) -> Self:
        from dask import is_dask_collection
        from dask.highlevelgraph import HighLevelGraph
        from dask.optimization import cull

        variables = {}

        for k, v in self._variables.items():
            if not is_dask_collection(v):
                variables[k] = v
                continue

            if isinstance(dsk, HighLevelGraph):
                # dask >= 2021.3
                # __dask_postpersist__() was called by dask.highlevelgraph.
                # Don't use dsk.cull(), as we need to prevent partial layers:
                # https://github.com/dask/dask/issues/7137
                layers = v.__dask_layers__()
                if rename:
                    layers = [rename.get(k, k) for k in layers]
                dsk2 = dsk.cull_layers(layers)
            elif rename:  # pragma: nocover
                # At the moment of writing, this is only for forward compatibility.
                # replace_name_in_key requires dask >= 2021.3.
                from dask.base import flatten, replace_name_in_key

                keys = [
                    replace_name_in_key(k, rename) for k in flatten(v.__dask_keys__())
                ]
                dsk2, _ = cull(dsk, keys)
            else:
                # __dask_postpersist__() was called by dask.optimize or dask.persist
                dsk2, _ = cull(dsk, v.__dask_keys__())

            rebuild, args = v.__dask_postpersist__()
            # rename was added in dask 2021.3
            kwargs = {"rename": rename} if rename else {}
            variables[k] = rebuild(dsk2, *args, **kwargs)

        return type(self)._construct_direct(
            variables,
            self._coord_names,
            self._dims,
            self._attrs,
            self._indexes,
            self._encoding,
            self._close,
        )

    def compute(self, **kwargs) -> Self:
        """Trigger loading data into memory and return a new dataset.