ronantakizawa/codereview-bench huggingface dataset:
Repository: dask/dask
File: dask/dataframe/utils.py (Python)

Code before:
    False
    """
    if not isinstance(divisions, (tuple, list)):
        return False

    for i, x in enumerate(divisions[:-2]):
        if x >= divisions[i + 1]:
            return False
        if isinstance(x, numbers.Number) and math.isnan(x):
            return False

    for x in divisions[-2:]:
        if isinstance(x, numbers.Number) and math.isnan(x):
            return False

    if divisions[-2] > divisions[-1]:
        return False

    return True


def drop_by_shallow_copy(df, columns):
    """ Use shallow copy to drop columns in place
    """
    df2 = df.copy(deep=False)
    if isinstance(columns, str):
        columns = [columns]
    for column in columns:
        del df2[column]
    return df2


Review comment (suggestion):
Pandas columns can also be tuples or numbers I think, so something like the following might be wiser:

```suggestion
    if not isinstance(columns, (list, pd.Index)):
```

Code after:
    False
    """
    if not isinstance(divisions, (tuple, list)):
        return False

    for i, x in enumerate(divisions[:-2]):
        if x >= divisions[i + 1]:
            return False
        if isinstance(x, numbers.Number) and math.isnan(x):
            return False

    for x in divisions[-2:]:
        if isinstance(x, numbers.Number) and math.isnan(x):
            return False

    if divisions[-2] > divisions[-1]:
        return False

    return True


def drop_by_shallow_copy(df, columns, errors="raise"):
    """ Use shallow copy to drop columns in place
    """
    df2 = df.copy(deep=False)
    if not pd.api.types.is_list_like(columns):
        columns = [columns]
    df2.drop(columns=columns, inplace=True, errors=errors)
    return df2
Repository: dask/dask
File: dask/dataframe/methods.py (Python)

Code before:
    if is_series_like(x) or is_dataframe_like(x):
        return x.where((x > y) | x.isnull(), y, axis=x.ndim - 1)
    else:  # scalar
        return x if x > y else y


def set_column(df, name, value):
    df = df.copy(deep=False)
    df[name] = value
    return df


def assign(df, *pairs):
    for name, val in dict(partition(2, pairs)).items():
        df = set_column(df, name, val)
    return df


def unique(x, series_name=None):
    out = x.unique()
    # out can be either an np.ndarray or may already be a series
    # like object.  When out is an np.ndarray, it must be wrapped.
    if not is_series_like(out):
        out = pd.Series(out, name=series_name)
    return out


def value_counts_combine(x):
    return x.groupby(level=0).sum()


def value_counts_aggregate(x):
    return x.groupby(level=0).sum().sort_values(ascending=False)


def nbytes(x):
    return x.nbytes


def size(x):
    return x.size


def values(df):
    return df.values


def sample(df, state, frac, replace):
    rs = np.random.RandomState(state)
    return df.sample(random_state=rs, frac=frac, replace=replace) if len(df) > 0 else df


Review comment (suggestion):
We probably don't need the set_column function if it's only used here.  I recommend that we just implement this directly within assign

```python
def assign(...):
    df = df.copy(deep=False)
    for ...
        ...
    return df
```

I personally find that reducing indirection like this is easier on future maintenance

Code after:
        return aggregate(x, y)


def cummin_aggregate(x, y):
    if is_series_like(x) or is_dataframe_like(x):
        return x.where((x < y) | x.isnull(), y, axis=x.ndim - 1)
    else:  # scalar
        return x if x < y else y


def cummax_aggregate(x, y):
    if is_series_like(x) or is_dataframe_like(x):
        return x.where((x > y) | x.isnull(), y, axis=x.ndim - 1)
    else:  # scalar
        return x if x > y else y


def assign(df, *pairs):
    # Only deep copy when updating an element
    # (to avoid modifying the original)
    pairs = dict(partition(2, pairs))
    deep = bool(set(pairs) & set(df.columns))
    df = df.copy(deep=bool(deep))
    for name, val in pairs.items():
        df[name] = val
    return df


def unique(x, series_name=None):
    out = x.unique()
    # out can be either an np.ndarray or may already be a series
    # like object.  When out is an np.ndarray, it must be wrapped.
    if not is_series_like(out):
        out = pd.Series(out, name=series_name)
    return out


def value_counts_combine(x):
    return x.groupby(level=0).sum()


def value_counts_aggregate(x):
    return x.groupby(level=0).sum().sort_values(ascending=False)


def nbytes(x):
    return x.nbytes


def size(x):
    return x.size
Repository: dask/dask
File: dask/array/tests/test_array_core.py (Python)

Code before:
    assert_eq(x[ind], dx[dind])
    assert_eq(x[:, ind], dx[:, dind])


def test_setitem_1d():
    x = np.arange(10)
    dx = da.from_array(x.copy(), chunks=(5,))

    x[x > 6] = -1
    x[x % 2 == 0] = -2

    dx[dx > 6] = -1
    dx[dx % 2 == 0] = -2

    assert_eq(x, dx)

    index = da.arange(3)
    with pytest.raises(ValueError, match="Boolean index assignment in Dask"):
        dx[index] = 1


def test_setitem_masked():
    # Test np.ma.masked assignment to object-type arrays
    x = np.ma.array(["a", 1, 3.14], dtype=object)
    dx = da.from_array(x.copy(), chunks=2)

    x[...] = np.ma.masked
    dx[...] = np.ma.masked

    assert_eq(x.mask, da.ma.getmaskarray(dx))


def test_setitem_hardmask():
    x = np.ma.array([1, 2, 3, 4], dtype=int)
    x.harden_mask()

    y = x.copy()
    assert y.hardmask

    x[0] = np.ma.masked
    x[0:2] = np.ma.masked

    dx = da.from_array(y)
    dx[0] = np.ma.masked
    dx[0:2] = np.ma.masked

    assert_eq(x, dx)


def test_setitem_2d():
    x = np.arange(24).reshape((4, 6))

Review comment (bug):
I notice that if I try to do `assert_eq(x, dx)` it fails, but I would have expected not to. 

```python
>>> x
masked_array(data=[--, --, --],
             mask=[ True,  True,  True],
       fill_value='?',
            dtype=object)
>>> dx 
dask.array<setitem, shape=(3,), dtype=object, chunksize=(2,), chunktype=numpy.MaskedArray>

>>>dx.compute() 
masked_array(data=[--, --, --],
             mask=[ True,  True,  True],
       fill_value='?',
            dtype=object)           
```
not sure why the `assert_eq(x, dx)` does not work. 
But with `assert_eq(x.mask, da.ma.getmaskarray(dx))`
we are asserting the mask is the same, but not the resulting array like in all the other tests. 

@jrbourbeau Do you have any thoughts?

Code after:
def test_setitem_1d():
    x = np.arange(10)
    dx = da.from_array(x.copy(), chunks=(5,))

    x[x > 6] = -1
    x[x % 2 == 0] = -2

    dx[dx > 6] = -1
    dx[dx % 2 == 0] = -2

    assert_eq(x, dx)

    index = da.arange(3)
    with pytest.raises(ValueError, match="Boolean index assignment in Dask"):
        dx[index] = 1


def test_setitem_masked():
    # Test np.ma.masked assignment to object-type arrays
    x = np.ma.array(["a", 1, 3.14], dtype=object)
    dx = da.from_array(x.copy(), chunks=2)

    x[...] = np.ma.masked
    dx[...] = np.ma.masked

    assert_eq(x.mask, da.ma.getmaskarray(dx))


def test_setitem_hardmask():
    x = np.ma.array([1, 2, 3, 4], dtype=int)
    x.harden_mask()

    y = x.copy()
    assert y.hardmask

    x[0] = np.ma.masked
    x[0:2] = np.ma.masked

    dx = da.from_array(y)
    dx[0] = np.ma.masked
    dx[0:2] = np.ma.masked
    assert_eq(x, dx)


def test_setitem_slice_twice():
    x = np.array([1, 2, 3, 4, 5, 6], dtype=int)
    val = np.array([0, 0], dtype=int)
    y = x.copy()

    x[0:2] = val
    x[4:6] = val
Repository: dask/dask
File: dask/tests/test_order.py (Python)

Code before:
    connected_roots, max_dependents = _connecting_to_roots(dependencies, dependents)
    assert connected_roots == {
        "a0": set(),
        "a1": set(),
        "a2": set(),
        "a3": set(),
        "a4": set(),
        "b1": {"a0", "a1"},
        "c1": {"a0", "a1"},
        "c2": {"a0", "a1"},
        "d1": {"a0", "a1"},
        "b2": {"a2"},
        "b3": {"a3"},
        "c3": {"a2", "a3"},
        "d2": {"a2", "a3", "a4"},
    }
    # Max dependents is just pre-computed for performance but it is itself just
    # a derived property of connected_roots
    assert max_dependents == {
        k: max((len(dependents[r]) for r in v), default=len(dependents[k]))
        for k, v in connected_roots.items()
    }


def test_donot_mutate_input():
    # Internally we may modify the graph but we don't want to mutatet the
    # external dsk
    np = pytest.importorskip("numpy")
    dsk = {
        "a": "data",
        "b": (f, 1),
        "c": np.array([[1, 2], [3, 4]]),
        "d": ["a", "b", "c"],
    }
    dsk_copy = dsk.copy()
    o = order(dsk)
    assert o["d"] == len(dsk) - 1
    assert len(dsk) == len(dsk_copy)


Review comment (suggestion):
```suggestion
def test_do_not_mutate_input():
    # Internally we may modify the graph but we don't want to mutate the
```

Code after:
        "a0": set(),
        "a1": set(),
        "a2": set(),
        "a3": set(),
        "a4": set(),
        "b1": {"a0", "a1"},
        "c1": {"a0", "a1"},
        "c2": {"a0", "a1"},
        "d1": {"a0", "a1"},
        "b2": {"a2"},
        "b3": {"a3"},
        "c3": {"a2", "a3"},
        "d2": {"a2", "a3", "a4"},
    }
    # Max dependents is just pre-computed for performance but it is itself just
    # a derived property of connected_roots
    assert max_dependents == {
        k: max((len(dependents[r]) for r in v), default=len(dependents[k]))
        for k, v in connected_roots.items()
    }


def test_do_not_mutate_input():
    # Internally we may modify the graph but we don't want to mutate the
    # external dsk
    np = pytest.importorskip("numpy")
    dsk = {
        "a": "data",
        "b": (f, 1),
        "c": np.array([[1, 2], [3, 4]]),
        "d": ["a", "b", "c"],
    }
    dsk_copy = dsk.copy()
    o = order(dsk)
    assert o["d"] == len(dsk) - 1
    assert len(dsk) == len(dsk_copy)
Repository: dask/dask
File: dask/dataframe/tests/test_dataframe.py (Python)

Code before:

    completions = ddf._ipython_key_completions_()
    assert "a" in completions
    assert "b" in completions
    assert "c" not in completions


def test_diff():
    df = pd.DataFrame(np.random.randn(100, 5), columns=list("abcde"))
    ddf = dd.from_pandas(df, 5)

    assert_eq(ddf.diff(), df.diff())
    assert_eq(ddf.diff(0), df.diff(0))
    assert_eq(ddf.diff(2), df.diff(2))
    assert_eq(ddf.diff(-2), df.diff(-2))

    assert_eq(ddf.diff(2, axis=1), df.diff(2, axis=1))

    assert_eq(ddf.a.diff(), df.a.diff())
    assert_eq(ddf.a.diff(0), df.a.diff(0))
    assert_eq(ddf.a.diff(2), df.a.diff(2))
    assert_eq(ddf.a.diff(-2), df.a.diff(-2))

    assert ddf.diff(2)._name == ddf.diff(2)._name
    assert ddf.diff(2)._name != ddf.diff(3)._name
    pytest.raises((TypeError, ValueError), lambda: ddf.diff(1.5))


def test_shift():
    df = _compat.makeTimeDataFrame()
    ddf = dd.from_pandas(df, npartitions=4)

    # DataFrame
    assert_eq(ddf.shift(), df.shift())
    assert_eq(ddf.shift(0), df.shift(0))
    assert_eq(ddf.shift(2), df.shift(2))
    assert_eq(ddf.shift(-2), df.shift(-2))

    assert_eq(ddf.shift(2, axis=1), df.shift(2, axis=1))

    # Series
    assert_eq(ddf.A.shift(), df.A.shift())
    assert_eq(ddf.A.shift(0), df.A.shift(0))
    assert_eq(ddf.A.shift(2), df.A.shift(2))
    assert_eq(ddf.A.shift(-2), df.A.shift(-2))

    with pytest.raises(TypeError):
        ddf.shift(1.5)


@pytest.mark.parametrize("data_freq,divs1", [("B", False), ("D", True), ("h", True)])

Review comment (suggestion):
```suggestion
    with pytest.raises((TypeError, ValueError)):
        ddf.diff(1.5)
```

Code after:

    completions = ddf._ipython_key_completions_()
    assert "a" in completions
    assert "b" in completions
    assert "c" not in completions


def test_diff():
    df = pd.DataFrame(np.random.randn(100, 5), columns=list("abcde"))
    ddf = dd.from_pandas(df, 5)

    assert_eq(ddf.diff(), df.diff())
    assert_eq(ddf.diff(0), df.diff(0))
    assert_eq(ddf.diff(2), df.diff(2))
    assert_eq(ddf.diff(-2), df.diff(-2))

    assert_eq(ddf.diff(2, axis=1), df.diff(2, axis=1))

    assert_eq(ddf.a.diff(), df.a.diff())
    assert_eq(ddf.a.diff(0), df.a.diff(0))
    assert_eq(ddf.a.diff(2), df.a.diff(2))
    assert_eq(ddf.a.diff(-2), df.a.diff(-2))

    assert ddf.diff(2)._name == ddf.diff(2)._name
    assert ddf.diff(2)._name != ddf.diff(3)._name
    with pytest.raises((TypeError, ValueError)):
        ddf.diff(1.5)


def test_shift():
    df = _compat.makeTimeDataFrame()
    ddf = dd.from_pandas(df, npartitions=4)

    # DataFrame
    assert_eq(ddf.shift(), df.shift())
    assert_eq(ddf.shift(0), df.shift(0))
    assert_eq(ddf.shift(2), df.shift(2))
    assert_eq(ddf.shift(-2), df.shift(-2))

    assert_eq(ddf.shift(2, axis=1), df.shift(2, axis=1))

    # Series
    assert_eq(ddf.A.shift(), df.A.shift())
    assert_eq(ddf.A.shift(0), df.A.shift(0))
    assert_eq(ddf.A.shift(2), df.A.shift(2))
    assert_eq(ddf.A.shift(-2), df.A.shift(-2))

    with pytest.raises(TypeError):
        ddf.shift(1.5)
Repository: dask/dask
File: dask/dataframe/tests/test_groupby.py (Python)

Code before:
import dask
import dask.dataframe as dd
from dask.dataframe import _compat
from dask.dataframe._compat import (
    PANDAS_GE_140,
    PANDAS_GE_150,
    PANDAS_GE_200,
    PANDAS_GE_210,
    check_nuisance_columns_warning,
    check_numeric_only_deprecation,
    check_observed_deprecation,
    tm,
)
from dask.dataframe._pyarrow import to_pyarrow_string
from dask.dataframe.backends import grouper_dispatch
from dask.dataframe.groupby import NUMERIC_ONLY_NOT_IMPLEMENTED
from dask.dataframe.utils import (
    assert_dask_graph,
    assert_eq,
    assert_max_deps,
    pyarrow_strings_enabled,
)
from dask.utils import M
from dask.utils_test import _check_warning, hlg_layer

dask_expr_enabled = dd._dask_expr_enabled()
AGG_FUNCS = [
    "sum",
    pytest.param(
        "mean",
        marks=pytest.mark.xfail(
            condition=PANDAS_GE_200,
            reason="numeric_only=False not implemented",
            strict=False,
        ),
    ),
    "median",
    "min",
    "max",
    "count",
    "size",
    pytest.param(
        "std",
        marks=pytest.mark.xfail(
            condition=PANDAS_GE_200,
            reason="numeric_only=False not implemented",
            strict=False,
        ),
    ),
    pytest.param(
        "var",

Review comment (suggestion):
```suggestion
DASK_EXPR_ENABLED = dd._dask_expr_enabled()
```

Code after:
import dask
import dask.dataframe as dd
from dask.dataframe import _compat
from dask.dataframe._compat import (
    PANDAS_GE_140,
    PANDAS_GE_150,
    PANDAS_GE_200,
    PANDAS_GE_210,
    check_nuisance_columns_warning,
    check_numeric_only_deprecation,
    check_observed_deprecation,
    tm,
)
from dask.dataframe._pyarrow import to_pyarrow_string
from dask.dataframe.backends import grouper_dispatch
from dask.dataframe.groupby import NUMERIC_ONLY_NOT_IMPLEMENTED
from dask.dataframe.utils import (
    assert_dask_graph,
    assert_eq,
    assert_max_deps,
    pyarrow_strings_enabled,
)
from dask.utils import M
from dask.utils_test import _check_warning, hlg_layer

DASK_EXPR_ENABLED = dd._dask_expr_enabled()
AGG_FUNCS = [
    "sum",
    pytest.param(
        "mean",
        marks=pytest.mark.xfail(
            condition=PANDAS_GE_200,
            reason="numeric_only=False not implemented",
            strict=False,
        ),
    ),
    "median",
    "min",
    "max",
    "count",
    "size",
    pytest.param(
        "std",
        marks=pytest.mark.xfail(
            condition=PANDAS_GE_200,
            reason="numeric_only=False not implemented",
            strict=False,
        ),
    ),
    pytest.param(
        "var",
Repository: dask/dask
File: dask/dataframe/tests/test_dataframe.py (Python)

Code before:
    with pytest.raises(ValueError, match="can be aligned|shape"):
        ddf.where(cond=cond, other=5)

    # but works when DataFrame is provided, with matching index
    dd_cond = pd.DataFrame(cond, index=df.index, columns=df.columns)
    expected = df.mask(cond=cond, other=5)
    result = ddf.mask(cond=dd_cond, other=5)
    assert_eq(expected, result)

    expected = df.where(cond=cond, other=5)
    result = ddf.where(cond=dd_cond, other=5)
    assert_eq(expected, result)


@pytest.mark.parametrize(
    "func, kwargs",
    [
        ("select_dtypes", {"include": "integer"}),
        ("describe", {"include": "integer"}),
        ("nunique", {}),
        ("quantile", {}),
    ],
)
def test_duplicate_columns(func, kwargs):
    if DASK_EXPR_ENABLED:
        pytest.skip("duplicated columns not supported")
    df = pd.DataFrame(
        {
            "int": [1, 2, 3],
            "float": [1.0, 2.0, 3.0],
            "d": 1,
        }
    )
    df.columns = ["a", "a", "d"]
    ddf = dd.from_pandas(df, npartitions=1)

    assert_eq(
        getattr(df, func)(**kwargs),
        getattr(ddf, func)(**kwargs),
    )


def test_mask_where_callable():
    """https://github.com/dask/dask/issues/10282"""
    pdf = pd.DataFrame({"x": [1, None]})
    ddf = dd.from_pandas(pdf, npartitions=1)

    # dataframe
    assert_eq(pdf.where(lambda d: d.isna(), 3), ddf.where(lambda d: d.isna(), 3))

    #  series

Review comment (bug):
```suggestion
@pytest.mark.xfail(DASK_EXPR_ENABLED, reason="duplicated columns not supported")
def test_duplicate_columns(func, kwargs):
```

Code after:
        ddf.where(cond=cond, other=5)

    # but works when DataFrame is provided, with matching index
    dd_cond = pd.DataFrame(cond, index=df.index, columns=df.columns)
    expected = df.mask(cond=cond, other=5)
    result = ddf.mask(cond=dd_cond, other=5)
    assert_eq(expected, result)

    expected = df.where(cond=cond, other=5)
    result = ddf.where(cond=dd_cond, other=5)
    assert_eq(expected, result)


@pytest.mark.parametrize(
    "func, kwargs",
    [
        ("select_dtypes", {"include": "integer"}),
        ("describe", {"include": "integer"}),
        ("nunique", {}),
        ("quantile", {}),
    ],
)
def test_duplicate_columns(func, kwargs):
    if DASK_EXPR_ENABLED:
        pytest.xfail("duplicated columns not supported")
    df = pd.DataFrame(
        {
            "int": [1, 2, 3],
            "float": [1.0, 2.0, 3.0],
            "d": 1,
        }
    )
    df.columns = ["a", "a", "d"]
    ddf = dd.from_pandas(df, npartitions=1)

    assert_eq(
        getattr(df, func)(**kwargs),
        getattr(ddf, func)(**kwargs),
    )


def test_mask_where_callable():
    """https://github.com/dask/dask/issues/10282"""
    pdf = pd.DataFrame({"x": [1, None]})
    ddf = dd.from_pandas(pdf, npartitions=1)

    # dataframe
    assert_eq(pdf.where(lambda d: d.isna(), 3), ddf.where(lambda d: d.isna(), 3))

    #  series
    assert_eq(pdf.x.where(lambda d: d == 1, 2), ddf.x.where(lambda d: d == 1, 2))
Repository: dask/dask
File: dask/core.py (Python)

Code before:
    """
    typ = type(key)
    if typ is tuple:
        return all(iskey(i) for i in cast(tuple, key))
    return typ in {bytes, int, float, str}


def validate_key(key: object) -> None:
    """Validate the format of a dask key.

    See Also
    --------
    iskey
    """
    if iskey(key):
        return
    typ = type(key)

    if typ is tuple:
        _index = None
        try:
            for _index, part in enumerate(cast(tuple, key)):
                validate_key(part)
        except TypeError as e:
            raise TypeError(
                f"Composite key contains unexpected key type at {_index=} ({key=!r})"
            ) from e
    raise TypeError(f"Unexpected key type {typ} ({key=!r})")


@overload
def get_dependencies(
    dsk: Graph,
    key: Key | None = ...,
    task: Key | NoDefault = ...,
    as_list: Literal[False] = ...,
) -> set[Key]:
    ...


@overload
def get_dependencies(
    dsk: Graph,
    key: Key | None,
    task: Key | NoDefault,
    as_list: Literal[True],
) -> list[Key]:
    ...


def get_dependencies(

Review comment (suggestion):
```suggestion
                f"Composite key contains unexpected key type at index={_index} ({key=!r})"
```

That would get ci passing

Code after:
    """
    typ = type(key)
    if typ is tuple:
        return all(iskey(i) for i in cast(tuple, key))
    return typ in {bytes, int, float, str}


def validate_key(key: object) -> None:
    """Validate the format of a dask key.

    See Also
    --------
    iskey
    """
    if iskey(key):
        return
    typ = type(key)

    if typ is tuple:
        index = None
        try:
            for index, part in enumerate(cast(tuple, key)):  # noqa: B007
                validate_key(part)
        except TypeError as e:
            raise TypeError(
                f"Composite key contains unexpected key type at {index=} ({key=!r})"
            ) from e
    raise TypeError(f"Unexpected key type {typ} ({key=!r})")


@overload
def get_dependencies(
    dsk: Graph,
    key: Key | None = ...,
    task: Key | NoDefault = ...,
    as_list: Literal[False] = ...,
) -> set[Key]:
    ...


@overload
def get_dependencies(
    dsk: Graph,
    key: Key | None,
    task: Key | NoDefault,
    as_list: Literal[True],
) -> list[Key]:
    ...


def get_dependencies(
Repository: dask/dask
File: dask/base.py (Python)

Code before:
from operator import getitem
from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar

from tlz import curry, groupby, identity, merge
from tlz.functoolz import Compose

from dask import config, local
from dask._compatibility import EMSCRIPTEN
from dask.core import flatten
from dask.core import get as simple_get
from dask.core import literal, quote
from dask.hashing import hash_buffer_hex
from dask.system import CPU_COUNT
from dask.typing import Key, SchedulerGetCallable
from dask.utils import (
    Dispatch,
    apply,
    ensure_dict,
    is_namedtuple_instance,
    key_split,
    shorten_traceback,
)

DistributedClient = None
get_distributed_client = None
DISTRIBUTED_AVAILABLE = None


def _distributed_available() -> bool:
    # Lazy import in get_scheduler can be expensive
    global DistributedClient, get_distributed_client, DISTRIBUTED_AVAILABLE
    if DISTRIBUTED_AVAILABLE is not None:
        return DISTRIBUTED_AVAILABLE  # type: ignore[unreachable]
    try:
        from distributed import Client as DistributedClient
        from distributed.worker import get_client as get_distributed_client

        DISTRIBUTED_AVAILABLE = True
    except ImportError:
        DISTRIBUTED_AVAILABLE = False
    return DISTRIBUTED_AVAILABLE


__all__ = (
    "DaskMethodsMixin",
    "annotate",
    "get_annotations",
    "is_dask_collection",
    "compute",
    "persist",
    "optimize",

Review comment (suggestion):
The only way to set these variables is currently private, so I suggest to make them private as well.
```suggestion
_DistributedClient = None
_get_distributed_client = None
_DISTRIBUTED_AVAILABLE = None
```

Code after:
from operator import getitem
from typing import TYPE_CHECKING, Any, Literal, Protocol, TypeVar

from tlz import curry, groupby, identity, merge
from tlz.functoolz import Compose

from dask import config, local
from dask._compatibility import EMSCRIPTEN
from dask.core import flatten
from dask.core import get as simple_get
from dask.core import literal, quote
from dask.hashing import hash_buffer_hex
from dask.system import CPU_COUNT
from dask.typing import Key, SchedulerGetCallable
from dask.utils import (
    Dispatch,
    apply,
    ensure_dict,
    is_namedtuple_instance,
    key_split,
    shorten_traceback,
)

_DistributedClient = None
_get_distributed_client = None
_DISTRIBUTED_AVAILABLE = None


def _distributed_available() -> bool:
    # Lazy import in get_scheduler can be expensive
    global _DistributedClient, _get_distributed_client, _DISTRIBUTED_AVAILABLE
    if _DISTRIBUTED_AVAILABLE is not None:
        return _DISTRIBUTED_AVAILABLE  # type: ignore[unreachable]
    try:
        from distributed import Client as _DistributedClient
        from distributed.worker import get_client as _get_distributed_client

        _DISTRIBUTED_AVAILABLE = True
    except ImportError:
        _DISTRIBUTED_AVAILABLE = False
    return _DISTRIBUTED_AVAILABLE


__all__ = (
    "DaskMethodsMixin",
    "annotate",
    "get_annotations",
    "is_dask_collection",
    "compute",
    "persist",
    "optimize",