Do inputs come from the same parents, modulo blockwise?
(*exprs)
| 3180 | |
| 3181 | |
| 3182 | def are_co_aligned(*exprs): |
| 3183 | """Do inputs come from the same parents, modulo blockwise?""" |
| 3184 | |
| 3185 | from dask.dataframe.dask_expr._cumulative import CumulativeAggregations |
| 3186 | from dask.dataframe.dask_expr._reductions import Reduction |
| 3187 | |
| 3188 | seen = set() |
| 3189 | # Scalars can always be broadcasted |
| 3190 | stack = [e for e in exprs if e.ndim > 0] |
| 3191 | ancestors = [] |
| 3192 | while stack: |
| 3193 | e = stack.pop() |
| 3194 | if e._name in seen: |
| 3195 | continue |
| 3196 | seen.add(e._name) |
| 3197 | |
| 3198 | if isinstance(e, IO): |
| 3199 | ancestors.append(e) |
| 3200 | elif e.ndim == 0: |
| 3201 | # Scalars are valid ancestors that are always broadcastable, |
| 3202 | # so don't walk through them |
| 3203 | continue |
| 3204 | elif isinstance(e, (_DelayedExpr, Isin)): |
| 3205 | continue |
| 3206 | elif isinstance(e, (Blockwise, CumulativeAggregations, Reduction)): |
| 3207 | # TODO: Capture this in inheritance logic |
| 3208 | dependencies = e.dependencies() |
| 3209 | stack.extend(dependencies) |
| 3210 | else: |
| 3211 | ancestors.append(e) |
| 3212 | |
| 3213 | unique_ancestors = { |
| 3214 | # Account for column projection within IO expressions |
| 3215 | _tokenize_partial(item, ["columns", "_series", "_dataset_info_cache"]) |
| 3216 | for item in ancestors |
| 3217 | } |
| 3218 | # Don't check divisions or npartitions at all |
| 3219 | return len(unique_ancestors) <= 1 |
| 3220 | |
| 3221 | |
| 3222 | ## Utilities for Expr fusion |