Repository: dask/dask
File: dask/_expr.py (Python)

Code before:
                "This often means that you are attempting to use an unsupported "
                f"API function.."
            )


class SingletonExpr(Expr):
    """A singleton Expr class

    This is used to treat the subclassed expression as a singleton. Singletons
    are deduplicated by expr._name which is typically based on the dask.tokenize
    output.

    This is a crucial performance optimization for expressions that walk through
    an optimizer and are recreated repeatedly but isn't safe for objects that
    cannot be reliably or quickly tokenized.
    """

    _instances: weakref.WeakValueDictionary[str, Expr] = weakref.WeakValueDictionary()

    def __new__(cls, *args, _determ_token=None, **kwargs):
        inst = super().__new__(cls, *args, _determ_token=_determ_token, **kwargs)
        _name = inst._name
        if _name in SingletonExpr._instances:
            return SingletonExpr._instances[_name]

        SingletonExpr._instances[_name] = inst
        return inst


def collect_dependents(expr) -> defaultdict:
    dependents = defaultdict(list)
    stack = [expr]
    seen = set()
    while stack:
        node = stack.pop()
        if node._name in seen:
            continue
        seen.add(node._name)

        for dep in node.dependencies():
            stack.append(dep)
            dependents[dep._name].append(weakref.ref(node))
    return dependents


def optimize(expr: Expr, fuse: bool = True) -> Expr:
    """High level query optimization

    This leverages three optimization passes:

    1.  Class based simplification using the ``_simplify`` function and methods

Review comment (performance):
Consider using cls._instances instead of SingletonExpr._instances in the SingletonExpr __new__ method so that each subclass can have its own separate cache and avoid potential cross-type collisions.
```suggestion
        if not hasattr(cls, '_instances'):
            cls._instances = weakref.WeakValueDictionary()
        inst = super().__new__(cls, *args, _determ_token=_determ_token, **kwargs)
        _name = inst._name
        if _name in cls._instances:
            return cls._instances[_name]

        cls._instances[_name] = inst
```

Code after:
                f"{err}\n\n"
                "This often means that you are attempting to use an unsupported "
                f"API function.."
            )


class SingletonExpr(Expr):
    """A singleton Expr class

    This is used to treat the subclassed expression as a singleton. Singletons
    are deduplicated by expr._name which is typically based on the dask.tokenize
    output.

    This is a crucial performance optimization for expressions that walk through
    an optimizer and are recreated repeatedly but isn't safe for objects that
    cannot be reliably or quickly tokenized.
    """

    _instances: weakref.WeakValueDictionary[str, SingletonExpr]

    def __new__(cls, *args, _determ_token=None, **kwargs):
        if not hasattr(cls, "_instances"):
            cls._instances = weakref.WeakValueDictionary()
        inst = super().__new__(cls, *args, _determ_token=_determ_token, **kwargs)
        _name = inst._name
        if _name in cls._instances:
            return cls._instances[_name]

        cls._instances[_name] = inst
        return inst


def collect_dependents(expr) -> defaultdict:
    dependents = defaultdict(list)
    stack = [expr]
    seen = set()
    while stack:
        node = stack.pop()
        if node._name in seen:
            continue
        seen.add(node._name)

        for dep in node.dependencies():
            stack.append(dep)
            dependents[dep._name].append(weakref.ref(node))
    return dependents


def optimize(expr: Expr, fuse: bool = True) -> Expr:
    """High level query optimization