(args, collections)
| 75 | |
| 76 | |
| 77 | def _finalize_args_collections(args, collections): |
| 78 | old_keys = [c.__dask_keys__()[0] for c in collections] |
| 79 | from dask._task_spec import cull |
| 80 | |
| 81 | collections = _ExprSequence(*collections).optimize() |
| 82 | new_keys = collections.__dask_keys__() |
| 83 | dsk = convert_legacy_graph(collections.__dask_graph__()) |
| 84 | annots = collections.__dask_annotations__() |
| 85 | outcollections = [] |
| 86 | for k in new_keys: |
| 87 | # Annotations are defined per HLG Layer but after this transformation |
| 88 | # these no longer properly exist which is why __dask_annotations__ |
| 89 | # returns a fully materialized dictionary {annot: {key: value}} |
| 90 | # Introducing a tombstone with a callable is the only way I found how we |
| 91 | # could revert this transformation (not necessarily efficient but |
| 92 | # well...) |
| 93 | layer_annotations = { |
| 94 | annot: partial( |
| 95 | _get_partial, dct=key_val, default=collections._annotations_tombstone() |
| 96 | ) |
| 97 | for annot, key_val in annots.items() |
| 98 | } |
| 99 | hlg = HighLevelGraph( |
| 100 | { |
| 101 | k[0]: MaterializedLayer( |
| 102 | cull(dsk, [k[0]]), |
| 103 | annotations=layer_annotations, |
| 104 | ) |
| 105 | }, |
| 106 | dependencies={k[0]: set()}, |
| 107 | ) |
| 108 | outcollections.append(Delayed(k[0], hlg)) |
| 109 | collections = tuple(outcollections) |
| 110 | subs = {old: new[0] for old, new in zip(old_keys, new_keys) if old != new} |
| 111 | args = args.substitute(subs) |
| 112 | return args, collections |
| 113 | |
| 114 | |
| 115 | def unpack_collections(expr, _return_collections=True): |
no test coverage detected