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

Code before:
        "dependencies": dependencies,
        "dependents": dependents,
        "waiting": waiting,
        "waiting_data": waiting_data,
        "cache": cache,
        "ready": ready,
        "running": set(),
        "finished": set(),
        "released": set(),
    }

    return state


"""
Running tasks
-------------

When we execute tasks we both

1.  Perform the actual work of collecting the appropriate data and calling the function
2.  Manage administrative state to coordinate with the scheduler
"""


def execute_task(key, task_info, dumps, loads, get_id, pack_exception):
    """
    Compute task and handle all administration

    See Also
    --------
    _execute_task : actually execute task
    """
    try:
        task, data = loads(task_info)
        result = task(data)
        id = get_id()
        result = dumps((result, id))
        failed = False
    except BaseException as e:  # noqa: B036
        result = pack_exception(e, dumps)
        failed = True
    return key, result, failed


def batch_execute_tasks(it):
    """
    Batch computing of multiple tasks with `execute_task`
    """
    return [execute_task(*a) for a in it]


Review comment (suggestion):
The original code filtered out tasks with empty waiting sets; consider filtering 'waiting' to only include keys with non-empty dependency sets to avoid potential misinterpretation later in the scheduler logic.

Code after:
            for d in dependents[key]:
                if d in waiting:
                    waiting[d].remove(key)
                    if not waiting[d]:
                        del waiting[d]
                        ready_set.add(d)
                else:
                    ready_set.add(d)
        else:
            _wait = task.dependencies - set(cache)
            if not _wait:
                ready_set.add(key)
            else:
                waiting[key] = set(_wait)
            for dep in task.dependencies:
                dependencies[key].add(dep)
                dependents[dep].add(key)
                waiting_data[dep].add(key)
                stack.append(dep)

    ready = sorted(ready_set, key=sortkey, reverse=True)

    state = {
        "dependencies": dict(dependencies),
        "dependents": dict(dependents),
        "waiting": dict(waiting),
        "waiting_data": dict(waiting_data),
        "cache": cache,
        "ready": ready,
        "running": set(),
        "finished": set(),
        "released": set(),
    }

    return state


"""
Running tasks
-------------

When we execute tasks we both

1.  Perform the actual work of collecting the appropriate data and calling the function
2.  Manage administrative state to coordinate with the scheduler
"""


def execute_task(key, task_info, dumps, loads, get_id, pack_exception):
    """
    Compute task and handle all administration
Repository: dask/dask
File: dask/dataframe/dask_expr/io/tests/test_io.py (Python)

Code before:

    assert ddf.to_bag().compute() == list(a.itertuples(False))
    assert ddf.to_bag(True).compute() == list(a.itertuples(True))
    assert ddf.to_bag(format="dict").compute() == [
        {"x": "a", "y": 2},
        {"x": "b", "y": 3},
        {"x": "c", "y": 4},
        {"x": "d", "y": 5},
    ]


def test_from_array_partition_pruning():
    arr = np.random.random(size=(200, 2))
    df = dd.from_array(arr, chunksize=100)
    result = df.partitions[1]
    expected = arr[100:]
    assert_eq(result, pd.DataFrame(expected, index=list(range(100, 200))))


def test_map_partitions_assign_fusedio(tmpdir):
    data = {"country": ["se", "dk", "no"], "value": [0.1, 0.2, 0.3]}
    pd.DataFrame(data).to_parquet(tmpdir, partition_cols=["country"])
    df = dd.read_parquet(tmpdir)
    y_pred = df.map_partitions(len)
    df["y_pred"] = y_pred
    assert assert_eq(
        df["y_pred"], pd.Series([1, 1, 1], name="y_pred"), check_index=False
    )


Review comment (suggestion):
It looks like there is a redundant 'assert' before the assert_eq call. Please remove the extra 'assert' to ensure the test reads correctly as assert_eq(...).
```suggestion
    assert_eq(
```

Code after:

    assert ddf.to_bag().compute() == list(a.itertuples(False))
    assert ddf.to_bag(True).compute() == list(a.itertuples(True))
    assert ddf.to_bag(format="dict").compute() == [
        {"x": "a", "y": 2},
        {"x": "b", "y": 3},
        {"x": "c", "y": 4},
        {"x": "d", "y": 5},
    ]


def test_from_array_partition_pruning():
    arr = np.random.random(size=(200, 2))
    df = dd.from_array(arr, chunksize=100)
    result = df.partitions[1]
    expected = arr[100:]
    assert_eq(result, pd.DataFrame(expected, index=list(range(100, 200))))


def test_map_partitions_assign_fusedio(tmpdir):
    data = {"country": ["se", "dk", "no"], "value": [0.1, 0.2, 0.3]}
    pd.DataFrame(data).to_parquet(tmpdir, partition_cols=["country"])
    df = dd.read_parquet(tmpdir)
    y_pred = df.map_partitions(len)
    df["y_pred"] = y_pred
    assert_eq(df["y_pred"], pd.Series([1, 1, 1], name="y_pred"), check_index=False)
Repository: dask/dask
File: dask/dataframe/dask_expr/tests/test_map_partitions_overlap.py (Python)

Code before:
    res = ddf.map_partitions(lambda df: df.rename_axis("newindex"))
    sol = df.rename_axis("newindex")
    assert_eq(res, sol)


def test_token_given(df, pdf):
    def my_func(df):
        return df + 1

    result = df.map_partitions(my_func, token="token_given")
    assert result._name.split("-")[0] == "token_given"
    assert_eq(result, pdf + 1)

    result = df.map_partitions(my_func)
    assert result._name.split("-")[0] == "my_func"
    assert_eq(result, pdf + 1)


@pytest.mark.parametrize(
    "overlap_setup",
    [
        (2, 0),
        (0, 2),
    ],
)
def test_map_overlap_independent_operations_on_same_input(df, pdf, overlap_setup):
    def count(part):
        return_value = part.copy()
        return_value["count"] = len(return_value)
        return return_value

    before, after = overlap_setup

    df_2 = map_overlap(count, df, before, after)
    df_4 = map_overlap(count, df, before * 2, after * 2)

    pdf_2, pdf_4 = dask.compute(df_2, df_4)

    partition_size = len(pdf) // df.npartitions

    overlap = max(before, after)
    no_overlap_length = (partition_size, partition_size)
    overlap_length = (partition_size + overlap, partition_size + overlap * 2)

    for idx, (division, next_division) in enumerate(
        zip(df.divisions, df.divisions[1:])
    ):
        if before != 0:
            expected_counts = overlap_length if idx > 0 else no_overlap_length
        else:
            expected_counts = (

Review comment (suggestion):
Non-blocker: This test feels like it could be simplified a bit, but I don't have any quick suggestions.

Code after:
    res = ddf.map_partitions(lambda df: df.rename_axis("newindex"))
    sol = df.rename_axis("newindex")
    assert_eq(res, sol)


def test_token_given(df, pdf):
    def my_func(df):
        return df + 1

    result = df.map_partitions(my_func, token="token_given")
    assert result._name.split("-")[0] == "token_given"
    assert_eq(result, pdf + 1)

    result = df.map_partitions(my_func)
    assert result._name.split("-")[0] == "my_func"
    assert_eq(result, pdf + 1)


@pytest.mark.parametrize(
    "overlap_setup",
    [
        (2, 0),
        (0, 2),
    ],
)
def test_map_overlap_independent_operations_on_same_input(df, pdf, overlap_setup):
    def count(part):
        return_value = part.copy()
        return_value["count"] = len(return_value)
        return return_value

    before, after = overlap_setup

    df_2 = map_overlap(count, df, before, after)
    df_4 = map_overlap(count, df, before * 2, after * 2)

    pdf_2, pdf_4 = dask.compute(df_2, df_4)

    partition_size = len(pdf) // df.npartitions
    overlap = max(before, after)

    expected_count_2 = pd.Series(
        partition_size, index=pdf.index, dtype=np.int64, name="count"
    )
    expected_count_4 = pd.Series(
        partition_size, index=pdf.index, dtype=np.int64, name="count"
    )

    x, y = (partition_size, len(pdf)) if before != 0 else (0, -partition_size)

    expected_count_2.iloc[x:y] = partition_size + overlap
Repository: dask/dask
File: dask/dataframe/backends.py (Python)

Code before:
##########


@make_scalar.register(np.dtype)
def _(dtype):
    return _scalar_from_dtype(dtype)


@make_scalar.register(pd.Timestamp)
@make_scalar.register(pd.Timedelta)
@make_scalar.register(pd.Period)
@make_scalar.register(pd.Interval)
def _(x):
    return x


@make_meta_dispatch.register((pd.Series, pd.DataFrame))
def _(x, index=None):
    out = x.iloc[:0].copy(deep=True)

    # https://github.com/pandas-dev/pandas/issues/61930
    # pandas shallow copies arrow-backed extension arrays.
    # Use pyarrow.compute.take to get a new array that doesn't
    # share any memory with the original array.

    for k, v in out.items():
        if isinstance(v.array, pd.arrays.ArrowExtensionArray):
            values = pyarrow.compute.take(
                pyarrow.array(v.array), pyarrow.array([], type="int32")
            )
            out[k] = v._constructor(
                pd.array(values, dtype=v.array.dtype), index=v.index, name=v.name
            )

    # index isn't copied by default in pandas, even if deep=true
    out.index = out.index.copy(deep=True)
    return out


@make_meta_dispatch.register(pd.Index)
def _(x, index=None):
    return x[0:0].copy(deep=True)


meta_object_types: tuple[type, ...] = (pd.Series, pd.DataFrame, pd.Index, pd.MultiIndex)
try:
    import scipy.sparse as sp

    meta_object_types += (sp.spmatrix,)
except ImportError:
    pass

Review comment (suggestion):
I haven't tried out the suggestion in https://github.com/pandas-dev/pandas/issues/61930#issuecomment-3109392476 but just wanted to flag that as it's a lot more straightforward (if it also does the same thing) 

Code after:
##########


@make_scalar.register(np.dtype)
def _(dtype):
    return _scalar_from_dtype(dtype)


@make_scalar.register(pd.Timestamp)
@make_scalar.register(pd.Timedelta)
@make_scalar.register(pd.Period)
@make_scalar.register(pd.Interval)
def _(x):
    return x


@make_meta_dispatch.register((pd.Series, pd.DataFrame))
def _(x, index=None):
    out = x.iloc[:0].copy(deep=True)

    # https://github.com/pandas-dev/pandas/issues/61930
    # pandas shallow copies arrow-backed extension arrays.
    # Use pyarrow.compute.take to get a new array that doesn't
    # share any memory with the original array.

    for k, v in out.items():
        if isinstance(v.array, pd.arrays.ArrowExtensionArray):
            values = pa.chunked_array([v.array]).combine_chunks()
            out[k] = v._constructor(
                pd.array(values, dtype=v.array.dtype), index=v.index, name=v.name
            )

    # index isn't copied by default in pandas, even if deep=true
    out.index = out.index.copy(deep=True)
    return out


@make_meta_dispatch.register(pd.Index)
def _(x, index=None):
    return x[0:0].copy(deep=True)


meta_object_types: tuple[type, ...] = (pd.Series, pd.DataFrame, pd.Index, pd.MultiIndex)
try:
    import scipy.sparse as sp

    meta_object_types += (sp.spmatrix,)
except ImportError:
    pass
Repository: dask/dask
File: dask/bytes/tests/test_http.py (Python)

Code before:
@pytest.mark.network
@pytest.mark.parametrize(
    "engine",
    ["pyarrow"],
)
def test_parquet(engine):
    pytest.importorskip("requests", minversion="2.21.0")
    dd = pytest.importorskip("dask.dataframe")
    pytest.importorskip(engine)
    df = dd.read_parquet(
        [
            "https://github.com/Parquet/parquet-compatibility/raw/"
            "master/parquet-testdata/impala/1.1.1-NONE/"
            "nation.impala.parquet"
        ],
        engine=engine,
    ).compute()
    assert df.n_nationkey.tolist() == list(range(25))
    assert df.columns.tolist() == ["n_nationkey", "n_name", "n_regionkey", "n_comment"]


@pytest.mark.flaky(
    reruns=10, reruns_delay=5, reason="https://github.com/dask/dask/issues/3696"
)
@pytest.mark.network
def test_bag():
    # This test pulls from different hosts
    urls = [
        "https://raw.githubusercontent.com/weierophinney/pastebin/"
        "master/public/js-src/dojox/data/tests/stores/patterns.csv",
        "https://en.wikipedia.org",
    ]
    b = db.read_text(urls)
    assert b.npartitions == 2
    b.compute()


@pytest.mark.network
def test_read_csv():
    dd = pytest.importorskip("dask.dataframe")
    url = (
        "https://raw.githubusercontent.com/weierophinney/pastebin/"
        "master/public/js-src/dojox/data/tests/stores/patterns.csv"
    )
    b = dd.read_csv(url)
    b.compute()


Review comment (suggestion):
All my `requests`-based HTTP requests to `https://en.wikipedia.org` are blocked with a 403 and the message `Please set a user-agent and respect our robot policy https://w.wiki/4wJS. See also T400119.`.

Maybe worth resurrecting the test, but not right now.

Code after:
@pytest.mark.network
@pytest.mark.parametrize(
    "engine",
    ["pyarrow"],
)
def test_parquet(engine):
    pytest.importorskip("requests", minversion="2.21.0")
    dd = pytest.importorskip("dask.dataframe")
    pytest.importorskip(engine)
    df = dd.read_parquet(
        [
            "https://github.com/Parquet/parquet-compatibility/raw/"
            "master/parquet-testdata/impala/1.1.1-NONE/"
            "nation.impala.parquet"
        ],
        engine=engine,
    ).compute()
    assert df.n_nationkey.tolist() == list(range(25))
    assert df.columns.tolist() == ["n_nationkey", "n_name", "n_regionkey", "n_comment"]


@pytest.mark.flaky(
    reruns=10, reruns_delay=5, reason="https://github.com/dask/dask/issues/3696"
)
@pytest.mark.network
@pytest.mark.skip(reason="Flaky test.")
def test_bag():
    # This test pulls from different hosts
    urls = [
        "https://raw.githubusercontent.com/weierophinney/pastebin/"
        "master/public/js-src/dojox/data/tests/stores/patterns.csv",
        "https://en.wikipedia.org",
    ]
    b = db.read_text(urls)
    assert b.npartitions == 2
    b.compute()


@pytest.mark.network
def test_read_csv():
    dd = pytest.importorskip("dask.dataframe")
    url = (
        "https://raw.githubusercontent.com/weierophinney/pastebin/"
        "master/public/js-src/dojox/data/tests/stores/patterns.csv"
    )
    b = dd.read_csv(url)
    b.compute()
Repository: dask/dask
File: docs/source/conf.py (Python)

Code before:
redirect_template = """\
<html>
  <head>
    <meta http-equiv="refresh" content="1; url={new}" />
    <script>
      window.location.href = "{new}"
    </script>
  </head>
</html>
"""

# Rate limiting issue for github: https://github.com/sphinx-doc/sphinx/issues/7388
linkcheck_ignore = [
    r"^https?:\/\/(?:www\.)?github.com\/",
    r"^https?:\/\/localhost(?:[:\/].+)?$",
]

doctest_global_setup = """
import numpy as np
"""


def copy_legacy_redirects(app, docname):
    if app.builder.name == "html":
        for html_src_path, new in redirect_files:
            # add ../ to old nested paths
            new = f"{'../' * html_src_path.count('/')}{new}"
            page = redirect_template.format(new=new)
            target_path = pathlib.Path(app.outdir) / html_src_path
            target_path.parent.mkdir(parents=True, exist_ok=True)
            target_path.write_text(page)


def setup(app):
    app.connect("build-finished", copy_legacy_redirects)


Review comment (suggestion):
`app.outdir` is changing type in Sphinx 9 and is raising a deprecation warning when being used as a string. Explicitly casting to a `Path` resolves this.

Code after:
      window.location.href = "{new}"
    </script>
  </head>
</html>
"""

# Rate limiting issue for github: https://github.com/sphinx-doc/sphinx/issues/7388
linkcheck_ignore = [
    r"^https?:\/\/(?:www\.)?github.com\/",
    r"^https?:\/\/localhost(?:[:\/].+)?$",
]

doctest_global_setup = """
import numpy as np
"""


def copy_legacy_redirects(app, docname):
    if app.builder.name == "html":
        for html_src_path, new in redirect_files:
            # add ../ to old nested paths
            new = f"{'../' * html_src_path.count('/')}{new}"
            page = redirect_template.format(new=new)
            target_path = pathlib.Path(app.outdir) / html_src_path
            target_path.parent.mkdir(parents=True, exist_ok=True)
            target_path.write_text(page)


def setup(app):
    app.connect("build-finished", copy_legacy_redirects)
Repository: dask/dask
File: dask/utils.py (Python)

Code before:
    Works for standard sequences (lists, tuples, etc.), NumPy arrays,
    and sparse-like objects (e.g., SciPy sparse arrays).

    The function checks:
        1. If the object supports len(), returns True if len(obj) == 0.
        2. If the object has a `.nnz` attribute (number of non-zero elements),
           returns True if `.nnz == 0`.
        3. If the object has a `.shape` attribute, returns True if any
           dimension is zero.
        4. Otherwise, returns False (assumes non-empty).

    Parameters
    ----------
    obj : any
        The object to check for emptiness.

    Returns
    -------
    bool
        True if the object is considered empty, False otherwise.
    """
    # Check standard sequences
    try:
        return len(obj) == 0
    except TypeError:
        pass

    # Sparse-like objects
    if hasattr(obj, "nnz"):
        return obj.nnz == 0
    if hasattr(obj, "shape"):
        return 0 in obj.shape

    # Fallback: assume non-empty
    return False


Review comment (bug):
I dislike that these can fail with an exception. We don't know that just because an object has `.nnz` that it's comparable to an int, and we don't know that `obj.shape` is a sequence. I'd feel better if these also catch exceptions (probably broad ones) so that we get through to the fallback.

Code after:
        2. If the object has a `.nnz` attribute (number of non-zero elements),
           returns True if `.nnz == 0`.
        3. If the object has a `.shape` attribute, returns True if any
           dimension is zero.
        4. Otherwise, returns False (assumes non-empty).

    Parameters
    ----------
    obj : any
        The object to check for emptiness.

    Returns
    -------
    bool
        True if the object is considered empty, False otherwise.
    """
    # Check standard sequences
    with contextlib.suppress(Exception):
        return len(obj) == 0

    # Sparse-like objects
    with contextlib.suppress(Exception):
        return obj.nnz == 0

    with contextlib.suppress(Exception):
        return 0 in obj.shape

    # Fallback: assume non-empty
    return False
Repository: dask/dask
File: dask/array/tests/test_array_core.py (Python)

Code before:
        chunks=zarr_chunk_shape,
        shards=zarr_shard_shape,
        dtype=arr.dtype,
    )

    with dask.config.set({"array.chunk-size": 1}):
        with pytest.raises(
            PerformanceWarning,
            match="The input Dask array will be rechunked along axis",
        ):
            arr.to_zarr(z)


def test_zarr_nocompute():
    pytest.importorskip("zarr")
    with tmpdir() as d:
        a = da.zeros((3, 3), chunks=(1, 1))
        out = a.to_zarr(d, compute=False)
        assert isinstance(out, Array)
        dask.compute(out)
        a2 = da.from_zarr(d)
        assert_eq(a, a2)
        assert a2.chunks == a.chunks


def test_zarr_regions():
    zarr = pytest.importorskip("zarr")

    a = da.arange(16).reshape((4, 4)).rechunk(2)
    z = zarr.zeros_like(a, chunks=2)

    a[:2, :2].to_zarr(z, region=(slice(2), slice(2)))
    a2 = da.from_zarr(z)
    expected = [[0, 1, 0, 0], [4, 5, 0, 0], [0, 0, 0, 0], [0, 0, 0, 0]]
    assert_eq(a2, expected)
    assert a2.chunks == a.chunks

    a[:3, 3:4].to_zarr(z, region=(slice(1, 4), slice(2, 3)))
    a2 = da.from_zarr(z)
    expected = [[0, 1, 0, 0], [4, 5, 3, 0], [0, 0, 7, 0], [0, 0, 11, 0]]
    assert_eq(a2, expected)
    assert a2.chunks == a.chunks

    a[3:, 3:].to_zarr(z, region=(slice(2, 3), slice(1, 2)))
    a2 = da.from_zarr(z)
    expected = [[0, 1, 0, 0], [4, 5, 3, 0], [0, 15, 7, 0], [0, 0, 11, 0]]
    assert_eq(a2, expected)
    assert a2.chunks == a.chunks

    with pytest.raises(ValueError):
        with tmpdir() as d:

Review comment (suggestion):
```suggestion
```

Shouldn't this warning be raised when existing dask chunks aren't aligned with zarr shards? IIUC it shouldn't require the `array.chunk-size` setting.

Code after:
    assert_eq(z[sel], arr.compute())


def test_zarr_risky_shards_warns():
    """
    Test that we see a performance warning when dask chooses a chunk size that will cause data loss
    for zarr arrays.
    """
    zarr = pytest.importorskip("zarr", minversion="3.0.0")

    shape = (100,)
    dask_chunks = (10,)
    zarr_chunk_shape = (3,)
    zarr_shard_shape = (6,)

    arr = da.arange(shape[0], chunks=dask_chunks)

    z = zarr.create_array(
        store={},
        shape=shape,
        chunks=zarr_chunk_shape,
        shards=zarr_shard_shape,
        dtype=arr.dtype,
    )

    with dask.config.set({"array.chunk-size": 1}):
        with pytest.raises(
            PerformanceWarning,
            match="The input Dask array will be rechunked along axis",
        ):
            arr.to_zarr(z)


def test_zarr_nocompute():
    pytest.importorskip("zarr")
    with tmpdir() as d:
        a = da.zeros((3, 3), chunks=(1, 1))
        out = a.to_zarr(d, compute=False)
        assert isinstance(out, Array)
        dask.compute(out)
        a2 = da.from_zarr(d)
        assert_eq(a, a2)
        assert a2.chunks == a.chunks


def test_zarr_regions():
    zarr = pytest.importorskip("zarr")

    a = da.arange(16).reshape((4, 4)).rechunk(2)
    z = zarr.zeros_like(a, chunks=2)