Function determine_column_projection

dask/dataframe/dask_expr/_expr.py:3857–3894

(
    expr: Expr,
    parent: Expr,
    dependents: dict[str, Collection[weakref.ref[BaseExpr]]],
    additional_columns: list | None = None,
)

def determine_column_projection(
    expr: Expr,
    parent: Expr,
    dependents: dict[str, Collection[weakref.ref[BaseExpr]]],
    additional_columns: list | None = None,
) -> object:
    if isinstance(parent, Index):
        column_union = []
    else:
        column_union = parent.columns.copy()
    parents: list[Expr]
    parents = [inst for x in dependents[expr._name] if isinstance((inst := x()), Expr)]

    seen = set()
    for p in parents:
        if p._name in seen:
            continue
        seen.add(p._name)
        column_union.extend(p._projection_columns)

    if additional_columns is not None:
        column_union.extend(flatten(additional_columns, container=list))

    # We can end up with MultiIndex columns from groupby ops, needs to be
    # accounted for in the sort
    flattened_columns = set(column_union)
    try:
        column_union = sorted(flattened_columns)
    except TypeError:
        # mixed type columns
        column_union = _sort_mixed(pd.Index(list(flattened_columns))).tolist()
    if (
        len(column_union) == 1
        and parent.ndim == 1
        and all(p.ndim == 1 for p in parents)
    ):
        return column_union[0]
    return column_union


def _sort_mixed(values):