Repository: dask/dask
File: dask/dataframe/dask_expr/tests/test_shuffle.py (Python)

Code before:

def test_filter_sort(df):
    a = df.sort_values("x")
    a = a[a.y > 40]

    b = df[df.y > 40]
    b = b.sort_values("x")

    assert a.optimize()._name == b.optimize()._name


def test_sort_values_add():
    pdf = pd.DataFrame({"x": [1, 2, 3, 0, 1, 2, 4, 5], "y": 1})
    df = from_pandas(pdf, npartitions=2, sort=False)
    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
        df = df.sort_values("x")
        df["z"] = df.x + df.y
        pdf = pdf.sort_values("x")
        pdf["z"] = pdf.x + pdf.y
        assert_eq(df, pdf, sort_results=False)


@pytest.mark.parametrize("null_value", [None, pd.NaT, pd.NA])
def test_index_nulls(null_value):
    "Setting the index with some non-numeric null raises error"
    df = pd.DataFrame(
        {"numeric": [1, 2, 3, 4], "non_numeric": ["foo", "bar", "foo", "bar"]}
    )
    ddf = from_pandas(df, npartitions=2)
    with pytest.raises(NotImplementedError, match="presence of nulls"):
        ddf.set_index(
            ddf["non_numeric"].map(
                {"foo": "foo", "bar": null_value}, meta=ddf["non_numeric"]._meta
            )
        ).compute()


@pytest.mark.parametrize("freq", ["16h", "-16h"])
def test_set_index_with_dask_dt_index(freq):
    values = {
        "x": [1, 2, 3, 4] * 3,
        "y": [10, 20, 30] * 4,
        "name": ["Alice", "Bob"] * 6,
    }
    date_index = pd.date_range(
        start="2022-02-22", freq=freq, periods=12
    ) - pd.Timedelta(seconds=30)
    df = pd.DataFrame(values, index=date_index)
    ddf = from_pandas(df, npartitions=3, sort=False)
    # specify a different date index entirely
    day_index = ddf.index.dt.floor("D")

Review comment (suggestion):
This change needs some careful thought.

In https://github.com/dask/dask/pull/12025/files#diff-864583c68ab3af8a180f3fff83a028c964a5ffda839bfb71acfa0f057952e1e3R280, I'm proposing that we use an actdual `object()` for `meta_nonempty` with object dtype. Historically, `object` has typically meant "probably strings, but possibly whatever".

Somewhere way down in `dask_expr` when figuring out the number of partitions or divisions, this was erroring when we tried to do a comparison operation between two objects, which raised an error.

The solution is to specify the `meta` manually, which is technically correct (`Series.map(dict)` can return arbitrary things), but a bit of a pain. With `meta` specified we don't get the user warning (not really relevant to this test) and we raise a `NotImplementedError` in `set_index` like the test expects.

So the things to figure out:

1. Are we OK with using `object` for the `meta_nonempty` of object-dtype arrays?
2. Are we OK with requiring the user to pass `meta` here in `.map`? (previously, I think it inferred `.map` was `object -> object`, which in essence meant `string -> string`, which worked fine).

Code after:
    b = df[df.y > 40]
    b = b.sort_values("x")

    assert a.optimize()._name == b.optimize()._name


def test_sort_values_add():
    pdf = pd.DataFrame({"x": [1, 2, 3, 0, 1, 2, 4, 5], "y": 1})
    df = from_pandas(pdf, npartitions=2, sort=False)
    with dask.config.set({"dataframe.shuffle.method": "tasks"}):
        df = df.sort_values("x")
        df["z"] = df.x + df.y
        pdf = pdf.sort_values("x")
        pdf["z"] = pdf.x + pdf.y
        assert_eq(df, pdf, sort_results=False)


@pytest.mark.parametrize("null_value", [None, pd.NaT, pd.NA])
def test_index_nulls(null_value):
    "Setting the index with some non-numeric null raises error"
    df = pd.DataFrame(
        {"numeric": [1, 2, 3, 4], "non_numeric": ["foo", "bar", "foo", "bar"]}
    )
    ddf = from_pandas(df, npartitions=2)
    with pytest.raises(NotImplementedError, match="presence of nulls"):
        ddf.set_index(
            ddf["non_numeric"].map(
                {"foo": "foo", "bar": null_value}, meta=ddf["non_numeric"]._meta
            )
        ).compute()


@pytest.mark.parametrize("freq", ["16h", "-16h"])
def test_set_index_with_dask_dt_index(freq):
    values = {
        "x": [1, 2, 3, 4] * 3,
        "y": [10, 20, 30] * 4,
        "name": ["Alice", "Bob"] * 6,
    }
    date_index = pd.date_range(
        start="2022-02-22", freq=freq, periods=12
    ) - pd.Timedelta(seconds=30)
    df = pd.DataFrame(values, index=date_index)
    ddf = from_pandas(df, npartitions=3, sort=False)
    # specify a different date index entirely
    day_index = ddf.index.dt.floor("D")
    result = ddf.set_index(day_index)
    assert_eq(result, pd.DataFrame(values, index=date_index.floor("D")))


def test_set_index_sort_values_shuffle_options(df, pdf):