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