| 972 | |
| 973 | |
| 974 | class HLGExpr(Expr): |
| 975 | _parameters = [ |
| 976 | "dsk", |
| 977 | "low_level_optimizer", |
| 978 | "output_keys", |
| 979 | "postcompute", |
| 980 | "_cached_optimized", |
| 981 | ] |
| 982 | _defaults = { |
| 983 | "low_level_optimizer": None, |
| 984 | "output_keys": None, |
| 985 | "postcompute": None, |
| 986 | "_cached_optimized": None, |
| 987 | } |
| 988 | |
| 989 | @property |
| 990 | def hlg(self): |
| 991 | return self.operand("dsk") |
| 992 | |
| 993 | @staticmethod |
| 994 | def from_collection(collection, optimize_graph=True): |
| 995 | from dask.highlevelgraph import HighLevelGraph |
| 996 | |
| 997 | if hasattr(collection, "dask"): |
| 998 | dsk = collection.dask.copy() |
| 999 | else: |
| 1000 | dsk = collection.__dask_graph__() |
| 1001 | |
| 1002 | # Delayed objects still ship with low level graphs as `dask` when going |
| 1003 | # through optimize / persist |
| 1004 | if not isinstance(dsk, HighLevelGraph): |
| 1005 | |
| 1006 | dsk = HighLevelGraph.from_collections( |
| 1007 | str(id(collection)), dsk, dependencies=() |
| 1008 | ) |
| 1009 | if optimize_graph and not hasattr(collection, "__dask_optimize__"): |
| 1010 | warnings.warn( |
| 1011 | f"Collection {type(collection)} does not define a " |
| 1012 | "`__dask_optimize__` method. In the future this will raise. " |
| 1013 | "If no optimization is desired, please set this to `None`.", |
| 1014 | PendingDeprecationWarning, |
| 1015 | ) |
| 1016 | low_level_optimizer = None |
| 1017 | else: |
| 1018 | low_level_optimizer = ( |
| 1019 | collection.__dask_optimize__ if optimize_graph else None |
| 1020 | ) |
| 1021 | return HLGExpr( |
| 1022 | dsk=dsk, |
| 1023 | low_level_optimizer=low_level_optimizer, |
| 1024 | output_keys=collection.__dask_keys__(), |
| 1025 | postcompute=collection.__dask_postcompute__(), |
| 1026 | ) |
| 1027 | |
| 1028 | def finalize_compute(self): |
| 1029 | return HLGFinalizeCompute( |
| 1030 | self, |
| 1031 | low_level_optimizer=self.low_level_optimizer, |
no outgoing calls