Function are_co_aligned

dask/dataframe/dask_expr/_expr.py:3205–3242

Do inputs come from the same parents, modulo blockwise?

(*exprs)

def are_co_aligned(*exprs):
    """Do inputs come from the same parents, modulo blockwise?"""

    from dask.dataframe.dask_expr._cumulative import CumulativeAggregations
    from dask.dataframe.dask_expr._reductions import Reduction

    seen = set()
    # Scalars can always be broadcasted
    stack = [e for e in exprs if e.ndim > 0]
    ancestors = []
    while stack:
        e = stack.pop()
        if e._name in seen:
            continue
        seen.add(e._name)

        if isinstance(e, IO):
            ancestors.append(e)
        elif e.ndim == 0:
            # Scalars are valid ancestors that are always broadcastable,
            # so don't walk through them
            continue
        elif isinstance(e, (_DelayedExpr, Isin)):
            continue
        elif isinstance(e, (Blockwise, CumulativeAggregations, Reduction)):
            # TODO: Capture this in inheritance logic
            dependencies = e.dependencies()
            stack.extend(dependencies)
        else:
            ancestors.append(e)

    unique_ancestors = {
        # Account for column projection within IO expressions
        _tokenize_partial(item, ["columns", "_series", "_dataset_info_cache"])
        for item in ancestors
    }
    # Don't check divisions or npartitions at all
    return len(unique_ancestors) <= 1


## Utilities for Expr fusion