| 1333 | |
| 1334 | |
| 1335 | class HLGFinalizeCompute(HLGExpr): |
| 1336 | |
| 1337 | def _simplify_down(self): |
| 1338 | if not self.postcompute: |
| 1339 | return self.dsk |
| 1340 | |
| 1341 | from dask.delayed import Delayed |
| 1342 | |
| 1343 | # Skip finalization for Delayed |
| 1344 | if self.dsk.postcompute == Delayed.__dask_postcompute__(self.dsk): |
| 1345 | return self.dsk |
| 1346 | return self |
| 1347 | |
| 1348 | @property |
| 1349 | def _name(self): |
| 1350 | return f"finalize-{super()._name}" |
| 1351 | |
| 1352 | def __dask_graph__(self): |
| 1353 | # The base class __dask_graph__ will not just materialize this layer but |
| 1354 | # also that of its dependencies, i.e. it will render the finalized and |
| 1355 | # the non-finalized graph and combine them. We only want the finalized |
| 1356 | # so we're overriding this. |
| 1357 | # This is an artifact generated because the wrapped expression is |
| 1358 | # identified automatically as a dependency but HLG expressions are not |
| 1359 | # working in this layered way. |
| 1360 | return self._layer() |
| 1361 | |
| 1362 | @property |
| 1363 | def hlg(self): |
| 1364 | expr = self.operand("dsk") |
| 1365 | layers = expr.dsk.layers.copy() |
| 1366 | deps = expr.dsk.dependencies.copy() |
| 1367 | keys = expr.__dask_keys__() |
| 1368 | if isinstance(expr.postcompute, list): |
| 1369 | postcomputes = expr.postcompute |
| 1370 | else: |
| 1371 | postcomputes = [expr.postcompute] |
| 1372 | tasks = [ |
| 1373 | Task(self._name, func, _convert_dask_keys(keys), *extra_args) |
| 1374 | for func, extra_args in postcomputes |
| 1375 | ] |
| 1376 | from dask.highlevelgraph import HighLevelGraph, MaterializedLayer |
| 1377 | |
| 1378 | leafs = set(deps) |
| 1379 | for val in deps.values(): |
| 1380 | leafs -= val |
| 1381 | for t in tasks: |
| 1382 | layers[t.key] = MaterializedLayer({t.key: t}) |
| 1383 | deps[t.key] = leafs |
| 1384 | return HighLevelGraph(layers, dependencies=deps) |
| 1385 | |
| 1386 | def __dask_keys__(self): |
| 1387 | return [self._name] |
| 1388 | |
| 1389 | |
| 1390 | class CompositeFinalizeCompute(CompositeExpr): |