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