def test_argsort_window_functions() -> None:
df = pl.DataFrame({"Id": [1, 1, 2, 2, 3, 3], "Age": [1, 2, 3, 4, 5, 6]})
out = df.select([
pl.col("Age").arg_sort().over("Id").alias("arg_sort"),
pl.argsort_by("Age").over("Id").alias("argsort_by"),
])
assert (out["arg_sort"].to_list() == out["argsort_by"].to_list() ==
[0, 1, 0, 1, 0, 1])
def test_argsort_by(df):
a = df[pl.argsort_by(["int_nulls", "floats"], reverse=[False,
True])]["int_nulls"]
assert a == [1, 0, 3]
def test_argsort_by():
df = get_complete_df()
a = df[pl.argsort_by(["int_nulls", "floats"], reverse=[False, True])]["int_nulls"]
assert a == [1, 0, 3]