Function optimize_until

dask/_expr.py:923–954

(expr: Expr, stage: OptimizerStage)

def optimize_until(expr: Expr, stage: OptimizerStage) -> Expr:
    result = expr
    if stage == "logical":
        return result

    # Simplify
    expr = result.simplify()
    if stage == "simplified-logical":
        return expr

    # Manipulate Expression to make it more efficient
    if dask.config.get("optimization.tune.active", True):
        expr = expr.rewrite(kind="tune", rewritten={})
    if stage == "tuned-logical":
        return expr

    # Lower
    expr = expr.lower_completely()
    if stage == "physical":
        return expr

    # Simplify again
    expr = expr.simplify()
    if stage == "simplified-physical":
        return expr

    # Final graph-specific optimizations
    expr = expr.fuse()
    if stage == "fused":
        return expr

    raise ValueError(f"Stage {stage!r} not supported.")


class LLGExpr(Expr):