Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

test(python): Refactor arg_min/max test case #14439

Merged
merged 1 commit into from
Feb 12, 2024
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
168 changes: 52 additions & 116 deletions py-polars/tests/unit/series/test_series.py
Original file line number Diff line number Diff line change
Expand Up @@ -1578,76 +1578,62 @@ def test_arg_sort() -> None:
assert_series_equal(s.arg_sort(descending=True), expected_descending)


def test_arg_min_and_arg_max() -> None:
# numerical no null.
s = pl.Series([5, 3, 4, 1, 2])
assert s.arg_min() == 3
assert s.arg_max() == 0

# numerical has null.
s = pl.Series([None, 5, 1])
assert s.arg_min() == 2
assert s.arg_max() == 1

# numerical all null.
s = pl.Series([None, None], dtype=Int32)
assert s.arg_min() is None
assert s.arg_max() is None

# boolean no null.
s = pl.Series([True, False])
assert s.arg_min() == 1
assert s.arg_max() == 0
s = pl.Series([True, True])
assert s.arg_min() == 0
assert s.arg_max() == 0
s = pl.Series([False, False])
assert s.arg_min() == 0
assert s.arg_max() == 0

# boolean has null.
s = pl.Series([None, True, False, True])
assert s.arg_min() == 2
assert s.arg_max() == 1
s = pl.Series([None, True, True])
assert s.arg_min() == 1
assert s.arg_max() == 1
s = pl.Series([None, False, False])
assert s.arg_min() == 1
assert s.arg_max() == 1

# boolean all null.
s = pl.Series([None, None], dtype=pl.Boolean)
assert s.arg_min() is None
assert s.arg_max() is None

# str no null
s = pl.Series(["a", "c", "b"])
assert s.arg_min() == 0
assert s.arg_max() == 1
@pytest.mark.parametrize(
("series", "argmin", "argmax"),
[
# Numeric
(pl.Series([5, 3, 4, 1, 2]), 3, 0),
(pl.Series([None, 5, 1]), 2, 1),
# Boolean
(pl.Series([True, False]), 1, 0),
(pl.Series([True, True]), 0, 0),
(pl.Series([False, False]), 0, 0),
(pl.Series([None, True, False, True]), 2, 1),
(pl.Series([None, True, True]), 1, 1),
(pl.Series([None, False, False]), 1, 1),
# String
(pl.Series(["a", "c", "b"]), 0, 1),
(pl.Series([None, "a", None, "b"]), 1, 3),
# Categorical
(pl.Series(["c", "b", "a"], dtype=pl.Categorical), 0, 2),
(pl.Series([None, "c", "b", None, "a"], dtype=pl.Categorical), 1, 4),
(pl.Series(["c", "b", "a"], dtype=pl.Categorical(ordering="lexical")), 2, 0),
(
pl.Series(
[None, "c", "b", None, "a"], dtype=pl.Categorical(ordering="lexical")
),
4,
1,
),
],
)
def test_arg_min_arg_max(series: pl.Series, argmin: int, argmax: int) -> None:
assert series.arg_min() == argmin
assert series.arg_max() == argmax

# str has null
s = pl.Series([None, "a", None, "b"])
assert s.arg_min() == 1
assert s.arg_max() == 3

# str all null
s = pl.Series([None, None], dtype=pl.String)
assert s.arg_min() is None
assert s.arg_max() is None
@pytest.mark.parametrize(
("series"),
[
# All nulls
pl.Series([None, None], dtype=pl.Int32),
pl.Series([None, None], dtype=pl.Boolean),
pl.Series([None, None], dtype=pl.String),
pl.Series([None, None], dtype=pl.Categorical),
pl.Series([None, None], dtype=pl.Categorical(ordering="lexical")),
# Empty Series
pl.Series([], dtype=pl.Int32),
pl.Series([], dtype=pl.Boolean),
pl.Series([], dtype=pl.String),
pl.Series([], dtype=pl.Categorical),
],
)
def test_arg_min_arg_max_all_nulls_or_empty(series: pl.Series) -> None:
assert series.arg_min() is None
assert series.arg_max() is None

# test ascending and descending series
s = pl.Series([None, 1, 2, 3, 4, 5])
s.sort(in_place=True) # set ascending sorted flag
assert s.flags == {"SORTED_ASC": True, "SORTED_DESC": False}
assert s.arg_min() == 1
assert s.arg_max() == 5
s = pl.Series([None, 5, 4, 3, 2, 1])
s.sort(descending=True, in_place=True) # set descing sorted flag
assert s.flags == {"SORTED_ASC": False, "SORTED_DESC": True}
assert s.arg_min() == 5
assert s.arg_max() == 1

def test_arg_min_and_arg_max_sorted() -> None:
# test ascending and descending numerical series
s = pl.Series([None, 1, 2, 3, 4, 5])
s.sort(in_place=True) # set ascending sorted flag
Expand All @@ -1672,56 +1658,6 @@ def test_arg_min_and_arg_max() -> None:
assert s.arg_min() == 5
assert s.arg_max() == 1

# test numerical empty series
s = pl.Series([], dtype=pl.Int32)
assert s.arg_min() is None
assert s.arg_max() is None

# test boolean empty series
s = pl.Series([], dtype=pl.Boolean)
assert s.arg_min() is None
assert s.arg_max() is None

# test str empty series
s = pl.Series([], dtype=pl.String)
assert s.arg_min() is None
assert s.arg_max() is None

# categorical empty series
s = pl.Series([], dtype=pl.Categorical)
assert s.arg_min() is None
assert s.arg_max() is None

# categorical with physical ordering no null
s = pl.Series(["c", "b", "a"], dtype=pl.Categorical)
assert s.arg_min() == 0
assert s.arg_max() == 2

# categorical with physical ordering has null
s = pl.Series([None, "c", "b", None, "a"], dtype=pl.Categorical)
assert s.arg_min() == 1
assert s.arg_max() == 4

# categorical with physical all null
s = pl.Series([None, None], dtype=pl.Categorical)
assert s.arg_min() is None
assert s.arg_max() is None

# categorical with lexical ordering no null
s = pl.Series(["c", "b", "a"], dtype=pl.Categorical(ordering="lexical"))
assert s.arg_min() == 2
assert s.arg_max() == 0

# categorical with lexical ordering has null
s = pl.Series([None, "c", "b", None, "a"], dtype=pl.Categorical(ordering="lexical"))
assert s.arg_min() == 4
assert s.arg_max() == 1

# categorical with lexical ordering all null
s = pl.Series([None, None], dtype=pl.Categorical(ordering="lexical"))
assert s.arg_min() is None
assert s.arg_max() is None


def test_is_null_is_not_null() -> None:
s = pl.Series("a", [1.0, 2.0, 3.0, None])
Expand Down