Skip to content

Commit

Permalink
Add tests of expression-based sort and sort-by
Browse files Browse the repository at this point in the history
We only need stable vs unstable variants for the sort-by case, since
when sorting a single column by itself there is no distinction between
stable and unstable.
  • Loading branch information
wence- committed Jun 12, 2024
1 parent c0c2ad3 commit 53bfc26
Showing 1 changed file with 53 additions and 0 deletions.
53 changes: 53 additions & 0 deletions python/cudf_polars/tests/expressions/test_sort.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,53 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import itertools

import pytest

import polars as pl

from cudf_polars.testing.asserts import assert_gpu_result_equal


@pytest.mark.parametrize("descending", [False, True])
@pytest.mark.parametrize("nulls_last", [False, True])
def test_sort_expression(descending, nulls_last):
ldf = pl.LazyFrame(
{
"a": [5, -1, 3, 4, None, 8, 6, 7, None],
}
)

query = ldf.select(pl.col("a").sort(descending=descending, nulls_last=nulls_last))
assert_gpu_result_equal(query)


@pytest.mark.parametrize(
"descending", itertools.combinations_with_replacement([False, True], 3)
)
@pytest.mark.parametrize(
"nulls_last", itertools.combinations_with_replacement([False, True], 3)
)
@pytest.mark.parametrize("maintain_order", [False, True], ids=["unstable", "stable"])
def test_sort_by_expression(descending, nulls_last, maintain_order):
ldf = pl.LazyFrame(
{
"a": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10],
"b": [1, 2, 2, 3, 9, 5, -1, 2, -2, 16],
"c": ["a", "A", "b", "b", "c", "d", "A", "Z", "ä", "̈Ä"],
}
)

query = ldf.select(
pl.col("a").sort_by(
pl.col("b"),
pl.col("c"),
pl.col("b") + pl.col("a"),
descending=descending,
nulls_last=nulls_last,
maintain_order=maintain_order,
)
)
assert_gpu_result_equal(query, check_row_order=maintain_order)

0 comments on commit 53bfc26

Please sign in to comment.