Repository: pola-rs/polars
File: py-polars/tests/unit/datatypes/test_list.py (Python)

Code before:


def test_list_sum_and_dtypes() -> None:
    # ensure the dtypes of sum align with normal sum
    for dt_in, dt_out in [
        (pl.Int8, pl.Int64),
        (pl.Int16, pl.Int64),
        (pl.Int32, pl.Int32),
        (pl.Int64, pl.Int64),
        (pl.Int128, pl.Int128),
        (pl.UInt8, pl.Int64),
        (pl.UInt16, pl.Int64),
        (pl.UInt32, pl.UInt32),
        (pl.UInt64, pl.UInt64),
        (pl.Float32, pl.Float32),
        (pl.Float64, pl.Float64),
    ]:
        df = pl.DataFrame(
            {
                "a": [[1], [1, 2, 3], [1, 2, 3, 4], [1, 2, 3, 4, 5]],
                "b": [[1], [1, 2, None], [1, 2, 3, None], [1, 2, 3, 4, None]],
            },
            schema={
                "a": pl.List(dt_in),
                "b": pl.List(dt_in),
            },
        )

        summed_a = df.select("a").explode("a").sum()
        assert summed_a.dtypes == [dt_out]
        assert summed_a.item() == 32
        assert df.select(pl.col("a").list.sum()).dtypes == [dt_out]

        # include nulls in the list
        summed_b = df.select("b").explode("b").sum()
        assert summed_b.dtypes == [dt_out]
        assert summed_b.item() == 20
        assert df.select(pl.col("b").list.sum()).dtypes == [dt_out]

    assert df.select(pl.col("a").list.sum()).dtypes == [dt_out]

    # include nulls
    assert pl.DataFrame(
        {"a": [[1], [1, 2, 3], [1, 2, 3, 4], [1, 2, 3, 4, 5], None]}
    ).select(pl.col("a").list.sum()).to_dict(as_series=False) == {
        "a": [1, 6, 10, 15, None]
    }

    # Booleans
    assert pl.DataFrame(
        {"a": [[True], [True, True], [True, False, True], [True, True, True, None]]},

Review comment (refactor):
We shouldn't remove the check on the exact values here

Could we refactor this test case to use `assert_series/frame_equal` instead? That will check both the values and dtype correctness


Code after:


def test_list_sum_and_dtypes() -> None:
    # ensure the dtypes of sum align with normal sum
    for dt_in, dt_out in [
        (pl.Int8, pl.Int64),
        (pl.Int16, pl.Int64),
        (pl.Int32, pl.Int32),
        (pl.Int64, pl.Int64),
        (pl.Int128, pl.Int128),
        (pl.UInt8, pl.Int64),
        (pl.UInt16, pl.Int64),
        (pl.UInt32, pl.UInt32),
        (pl.UInt64, pl.UInt64),
        (pl.UInt128, pl.UInt128),
        (pl.Float32, pl.Float32),
        (pl.Float64, pl.Float64),
    ]:
        df = pl.DataFrame(
            {
                "a": [[1], [1, 2, 3], [1, 2, 3, 4], [1, 2, 3, 4, 5]],
                "b": [[None], [1, 2, None], [1, 2, 3, None], [1, 2, 3, 4, None]],
            },
            schema={
                "a": pl.List(dt_in),
                "b": pl.List(dt_in),
            },
        )

        assert_frame_equal(
            df.select("a").explode("a").sum(),
            pl.DataFrame(pl.Series("a", [32], dtype=dt_out)),
            check_dtypes=True,
            check_exact=True,
        )
        assert_series_equal(
            df.get_column("a").list.sum(),
            pl.Series("a", [1, 6, 10, 15], dtype=dt_out),
            check_dtypes=True,
            check_exact=True,
        )

        # include nulls in the list
        assert_frame_equal(
            df.select("b").explode("b").sum(),
            pl.DataFrame(pl.Series("b", [19], dtype=dt_out)),
            check_dtypes=True,
            check_exact=True,
        )
        assert_series_equal(
            df.get_column("b").list.sum(),