keys are the keys from the graph that are requested by a computation. We can't fuse those together.
(dsk, keys)
| 1085 | |
| 1086 | |
| 1087 | def fuse_linear_task_spec(dsk, keys): |
| 1088 | """ |
| 1089 | keys are the keys from the graph that are requested by a computation. We |
| 1090 | can't fuse those together. |
| 1091 | """ |
| 1092 | from dask.core import reverse_dict |
| 1093 | from dask.optimization import default_fused_keys_renamer |
| 1094 | |
| 1095 | keys = set(keys) |
| 1096 | dependencies = DependenciesMapping(dsk) |
| 1097 | dependents = reverse_dict(dependencies) |
| 1098 | |
| 1099 | seen = set() |
| 1100 | result = {} |
| 1101 | |
| 1102 | for key in dsk: |
| 1103 | if key in seen: |
| 1104 | continue |
| 1105 | |
| 1106 | seen.add(key) |
| 1107 | |
| 1108 | deps = dependencies[key] |
| 1109 | dependents_key = dependents[key] |
| 1110 | |
| 1111 | if len(deps) != 1 and len(dependents_key) != 1 or dsk[key].block_fusion: |
| 1112 | result[key] = dsk[key] |
| 1113 | continue |
| 1114 | |
| 1115 | linear_chain = [dsk[key]] |
| 1116 | top_key = key |
| 1117 | |
| 1118 | # Walk towards the leafs as long as the nodes have a single dependency |
| 1119 | # and a single dependent, we can't fuse two nodes of an intermediate node |
| 1120 | # is the source for 2 dependents |
| 1121 | while len(deps) == 1: |
| 1122 | (new_key,) = deps |
| 1123 | if new_key in seen: |
| 1124 | break |
| 1125 | seen.add(new_key) |
| 1126 | if new_key not in dsk: |
| 1127 | # This can happen if a future is in the graph, the dependency mapping |
| 1128 | # adds the key that is referenced by the future as a dependency |
| 1129 | # see test_futures_to_delayed_array |
| 1130 | break |
| 1131 | if ( |
| 1132 | len(dependents[new_key]) != 1 |
| 1133 | or dsk[new_key].block_fusion |
| 1134 | or new_key in keys |
| 1135 | ): |
| 1136 | result[new_key] = dsk[new_key] |
| 1137 | break |
| 1138 | # backwards comp for new names, temporary until is_rootish is removed |
| 1139 | linear_chain.insert(0, dsk[new_key]) |
| 1140 | deps = dependencies[new_key] |
| 1141 | |
| 1142 | # Walk the tree towards the root as long as the nodes have a single dependent |
| 1143 | # and a single dependency, we can't fuse two nodes if node has multiple |
| 1144 | # dependencies |