Optimize several dask collections at once. Returns equivalent dask collections that all share the same merged and optimized underlying graph. This can be useful if converting multiple collections to delayed objects, or to manually apply the optimizations at strategic points. No
(*args, traverse=True, **kwargs)
| 597 | |
| 598 | |
| 599 | def optimize(*args, traverse=True, **kwargs): |
| 600 | """Optimize several dask collections at once. |
| 601 | |
| 602 | Returns equivalent dask collections that all share the same merged and |
| 603 | optimized underlying graph. This can be useful if converting multiple |
| 604 | collections to delayed objects, or to manually apply the optimizations at |
| 605 | strategic points. |
| 606 | |
| 607 | Note that in most cases you shouldn't need to call this function directly. |
| 608 | |
| 609 | Warning:: |
| 610 | |
| 611 | This function triggers a materialization of the collections and looses |
| 612 | any annotations attached to HLG layers. |
| 613 | |
| 614 | Parameters |
| 615 | ---------- |
| 616 | *args : objects |
| 617 | Any number of objects. If a dask object, its graph is optimized and |
| 618 | merged with all those of all other dask objects before returning an |
| 619 | equivalent dask collection. Non-dask arguments are passed through |
| 620 | unchanged. |
| 621 | traverse : bool, optional |
| 622 | By default dask traverses builtin python collections looking for dask |
| 623 | objects passed to ``optimize``. For large collections this can be |
| 624 | expensive. If none of the arguments contain any dask objects, set |
| 625 | ``traverse=False`` to avoid doing this traversal. |
| 626 | optimizations : list of callables, optional |
| 627 | Additional optimization passes to perform. |
| 628 | **kwargs |
| 629 | Extra keyword arguments to forward to the optimization passes. |
| 630 | |
| 631 | Examples |
| 632 | -------- |
| 633 | >>> import dask |
| 634 | >>> import dask.array as da |
| 635 | >>> a = da.arange(10, chunks=2).sum() |
| 636 | >>> b = da.arange(10, chunks=2).mean() |
| 637 | >>> a2, b2 = dask.optimize(a, b) |
| 638 | |
| 639 | >>> a2.compute() == a.compute() |
| 640 | np.True_ |
| 641 | >>> b2.compute() == b.compute() |
| 642 | np.True_ |
| 643 | """ |
| 644 | # TODO: This API is problematic. The approach to using postpersist forces us |
| 645 | # to materialize the graph. Most low level optimizations will materialize as |
| 646 | # well |
| 647 | collections, repack = unpack_collections(*args, traverse=traverse) |
| 648 | if not collections: |
| 649 | return args |
| 650 | |
| 651 | from dask._expr import CompositeExpr, _ExprSequence |
| 652 | |
| 653 | dsk = collections_to_expr(collections) |
| 654 | collection_exprs = list(dsk.operands) if isinstance(dsk, _ExprSequence) else [dsk] |
| 655 | if len(collection_exprs) != len(collections): |
| 656 | raise RuntimeError( |