Repository: pola-rs/polars
File: py-polars/tests/unit/operations/test_gather.py (Python)

Code before:
    assert_frame_equal(df, expected, check_row_order=maintain_order)


def test_gather_invalid_indices_groupby_24182() -> None:
    df = pl.DataFrame({"x": [1, 2]})
    with pytest.raises(pl.exceptions.InvalidOperationError):
        df.group_by(True).agg(pl.col("x").gather(pl.lit("y")))


@pytest.mark.parametrize("maintain_order", [False, True])
def test_gather_group_by_lit(maintain_order: bool) -> None:
    assert_frame_equal(
        pl.DataFrame(
            {
                "a": [1, 2, 3],
            }
        )
        .group_by("a", maintain_order=maintain_order)
        .agg(pl.lit([1]).gather([0, 0, 0])),
        pl.DataFrame({"a": [1, 2, 3], "literal": [[[1], [1], [1]]] * 3}),
        check_row_order=maintain_order,
    )


def test_get_null_on_oob_true() -> None:
    # OOB get with null_on_oob=True should yield null, not raise.
    df = pl.DataFrame({"value": [1, 2, 10]})

    out = df.select(v=pl.col("value").get(5, null_on_oob=True))

    # index 5 is out of range -> null
    assert out["v"].to_list() == [None]


def test_get_null_on_oob_false_raises_out_of_bounds() -> None:
    # OOB get with null_on_oob=False should raise OutOfBoundsError.
    df = pl.DataFrame({"value": [10, 11]})

    with pytest.raises(OutOfBoundsError, match="gather indices are out of bounds"):
        df.select(pl.col("value").get(5, null_on_oob=False))


def test_get_negative_index() -> None:
    # Negative indices follow Python semantics: -1 = last
    df = pl.DataFrame({"value": [1, 98, 2, 3, 99]})

    out = df.select(last=pl.col("value").get(-1, null_on_oob=True))

    assert out["last"].to_list() == [99]



Review comment (suggestion):
In case a test refers to a specific issue we add the issue number to the test name instead of a comment. So this would be `test_get_window_with_filtered_empty_groups_23029`.

Code after:
    assert_frame_equal(df, expected, check_row_order=maintain_order)


def test_gather_invalid_indices_groupby_24182() -> None:
    df = pl.DataFrame({"x": [1, 2]})
    with pytest.raises(pl.exceptions.InvalidOperationError):
        df.group_by(True).agg(pl.col("x").gather(pl.lit("y")))


@pytest.mark.parametrize("maintain_order", [False, True])
def test_gather_group_by_lit(maintain_order: bool) -> None:
    assert_frame_equal(
        pl.DataFrame(
            {
                "a": [1, 2, 3],
            }
        )
        .group_by("a", maintain_order=maintain_order)
        .agg(pl.lit([1]).gather([0, 0, 0])),
        pl.DataFrame({"a": [1, 2, 3], "literal": [[[1], [1], [1]]] * 3}),
        check_row_order=maintain_order,
    )


def test_get_window_with_filtered_empty_groups_23029() -> None:
    # https://github.com/pola-rs/polars/issues/23029
    df = pl.DataFrame(
        {
            "group": [1, 1, 2, 2, 3, 3],
            "value": [10, 20, 30, 40, 50, 60],
            "filter_condition": [False, True, False, False, True, True],
        }
    )

    result = df.with_columns(
        get_first=(
            pl.col("value")
            .filter(pl.col("filter_condition"))
            .get(0, null_on_oob=True)
            .over("group")
        ),
        first_value=(
            pl.col("value").filter(pl.col("filter_condition")).first().over("group")
        ),
    )

    assert_series_equal(
        result["get_first"],
        result["first_value"],
        check_names=False,
    )