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

Code before:

def test_isnull():
    pdf = pd.DataFrame(
        {
            "A": range(10),
            "B": [None] * 10,
            "C": [1] * 4 + [None] * 4 + [2] * 2,
        }
    )
    df = from_pandas(pdf, npartitions=2)
    assert_eq(df.notnull(), pdf.notnull())
    assert_eq(df.isnull(), pdf.isnull())
    for c in pdf.columns:
        assert_eq(df[c].isnull(), pdf[c].isnull())


def test_scalar_to_series():
    sc = from_pandas(pd.Series([1])).sum()
    ss1 = sc.to_series()
    ss2 = sc.to_series("xxx")
    assert_eq(ss1, pd.Series([1]))
    assert_eq(ss2, pd.Series([1], index=["xxx"]))


@pytest.mark.parametrize("data_freq, divs1", [("B", False), ("D", True), ("h", True)])
def test_shift_with_freq_datetime(pdf, data_freq, divs1):
    pdf.index = pd.date_range(start="2020-01-01", periods=len(pdf), freq=data_freq)
    df = from_pandas(pdf, npartitions=4)
    for freq, divs2 in [("s", True), ("W", False), (pd.Timedelta(10, unit="h"), True)]:
        for d, p in [(df, pdf), (df.x, pdf.x)]:
            res = d.shift(2, freq=freq)
            assert_eq(res, p.shift(2, freq=freq))
            assert res.known_divisions == divs2

    res = df.index.shift(2)
    assert_eq(res, df.index.shift(2))
    assert res.known_divisions == divs1


def test_isin_head(df, pdf):
    df["foo"] = df.x.isin([1, 2])
    pdf["foo"] = pdf.x.isin([1, 2])
    assert_eq(df.head(), pdf.head())


def test_head_npartitions_divisions(pdf):
    df = from_pandas(pdf, npartitions=1)
    assert_eq(df.head(3, compute=False), pdf.head(3))


@pytest.mark.parametrize("data_freq,divs", [("D", True), ("h", True)])

Review comment (bug):
These fixture changes should be equivalent to the nested for loops.

We end up doing a bit more work (running the lines outside the for loop multiple times), but it made debugging easier .

Code after:

def test_isnull():
    pdf = pd.DataFrame(
        {
            "A": range(10),
            "B": [None] * 10,
            "C": [1] * 4 + [None] * 4 + [2] * 2,
        }
    )
    df = from_pandas(pdf, npartitions=2)
    assert_eq(df.notnull(), pdf.notnull())
    assert_eq(df.isnull(), pdf.isnull())
    for c in pdf.columns:
        assert_eq(df[c].isnull(), pdf[c].isnull())


def test_scalar_to_series():
    sc = from_pandas(pd.Series([1])).sum()
    ss1 = sc.to_series()
    ss2 = sc.to_series("xxx")
    assert_eq(ss1, pd.Series([1]))
    assert_eq(ss2, pd.Series([1], index=["xxx"]))


@pytest.mark.parametrize("data_freq, divs1", [("B", False), ("D", True), ("h", True)])
@pytest.mark.parametrize(
    "freq_divs2", [("s", True), ("W", False), (pd.Timedelta(10, unit="h"), True)]
)
@pytest.mark.parametrize("as_series", [True, False])
def test_shift_with_freq_datetime(pdf, data_freq, divs1, freq_divs2, as_series):
    pdf.index = pd.date_range(start="2020-01-01", periods=len(pdf), freq=data_freq)
    df = from_pandas(pdf, npartitions=4)

    freq, divs2 = freq_divs2
    if as_series:
        pd_obj = pdf["x"]
        dask_obj = df["x"]
    else:
        pd_obj = pdf
        dask_obj = df

    res = dask_obj.shift(2, freq=freq)
    assert_eq(res, pd_obj.shift(2, freq=freq))
    assert res.known_divisions == divs2

    res = df.index.shift(2)
    assert_eq(res, df.index.shift(2))
    assert res.known_divisions == divs1


def test_isin_head(df, pdf):