| 1089 | |
| 1090 | |
| 1091 | class _HLGExprSequence(Expr): |
| 1092 | |
| 1093 | def __getitem__(self, other): |
| 1094 | return self.operands[other] |
| 1095 | |
| 1096 | def _operands_for_repr(self): |
| 1097 | return [ |
| 1098 | f"name={self.operand('name')!r}", |
| 1099 | f"dsk={self.operand('dsk')!r}", |
| 1100 | ] |
| 1101 | |
| 1102 | def _tree_repr_lines(self, indent=0, recursive=True): |
| 1103 | return self._operands_for_repr() |
| 1104 | |
| 1105 | def finalize_compute(self): |
| 1106 | return _HLGExprSequence(*[op.finalize_compute() for op in self.operands]) |
| 1107 | |
| 1108 | def _tune_down(self): |
| 1109 | if len(self.operands) == 1: |
| 1110 | return None |
| 1111 | from dask.highlevelgraph import HighLevelGraph |
| 1112 | |
| 1113 | groups = toolz.groupby( |
| 1114 | lambda x: x.low_level_optimizer if isinstance(x, HLGExpr) else None, |
| 1115 | self.operands, |
| 1116 | ) |
| 1117 | exprs = [] |
| 1118 | changed = False |
| 1119 | for optimizer, group in groups.items(): |
| 1120 | if len(group) > 1: |
| 1121 | graphs = [expr.hlg for expr in group] |
| 1122 | |
| 1123 | changed = True |
| 1124 | dsk = HighLevelGraph.merge(*graphs) |
| 1125 | hlg_group = _HLGExprGroup( |
| 1126 | dsk=dsk, |
| 1127 | low_level_optimizer=optimizer, |
| 1128 | output_keys=[v.__dask_keys__() for v in group], |
| 1129 | postcompute=[g.postcompute for g in group], |
| 1130 | ) |
| 1131 | exprs.append(hlg_group) |
| 1132 | else: |
| 1133 | exprs.append(group[0]) |
| 1134 | if not changed: |
| 1135 | return None |
| 1136 | return _HLGExprSequence(*exprs) |
| 1137 | |
| 1138 | @functools.cached_property |
| 1139 | def _optimized_dsk(self) -> HighLevelGraph: |
| 1140 | from dask.highlevelgraph import HighLevelGraph |
| 1141 | |
| 1142 | hlgexpr: HLGExpr |
| 1143 | graphs = [] |
| 1144 | # simplify_down ensure there are only one HLGExpr per optimizer/finalizer |
| 1145 | for hlgexpr in self.operands: |
| 1146 | keys = hlgexpr.__dask_keys__() |
| 1147 | dsk = hlgexpr.hlg |
| 1148 | if (optimizer := hlgexpr.low_level_optimizer) is not None: |
no outgoing calls
no test coverage detected