Skip to content

Commit

Permalink
Add tests of implemented StringFunctions
Browse files Browse the repository at this point in the history
Additionally, assert that we raise during translation for an unhandled function.
  • Loading branch information
wence- committed Jun 12, 2024
1 parent c0c2ad3 commit 550debc
Show file tree
Hide file tree
Showing 2 changed files with 44 additions and 1 deletion.
4 changes: 3 additions & 1 deletion python/cudf_polars/cudf_polars/dsl/expr.py
Original file line number Diff line number Diff line change
Expand Up @@ -689,7 +689,9 @@ def do_evaluate(
)
)
else:
raise NotImplementedError(f"StringFunction {self.name}")
raise NotImplementedError(
f"StringFunction {self.name}"
) # pragma: no cover; handled by init raising


class Sort(Expr):
Expand Down
41 changes: 41 additions & 0 deletions python/cudf_polars/tests/expressions/test_stringfunction.py
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
# SPDX-FileCopyrightText: Copyright (c) 2024 NVIDIA CORPORATION & AFFILIATES.
# SPDX-License-Identifier: Apache-2.0
from __future__ import annotations

import pytest

import polars as pl

from cudf_polars import translate_ir
from cudf_polars.testing.asserts import assert_gpu_result_equal


def test_supported_stringfunction_expression():
ldf = pl.LazyFrame(
{
"a": ["a", "b", "cdefg", "h", "Wıth ünιcοde"], # noqa: RUF001
"b": [0, 3, 1, -1, None],
}
)

query = ldf.select(
pl.col("a").str.starts_with("Z"),
pl.col("a").str.ends_with("h").alias("endswith_h"),
pl.col("a").str.to_lowercase().alias("lower"),
pl.col("a").str.to_uppercase().alias("upper"),
)
assert_gpu_result_equal(query)


def test_unsupported_stringfunction():
ldf = pl.LazyFrame(
{
"a": ["a", "b", "cdefg", "h", "Wıth ünιcοde"], # noqa: RUF001
"b": [0, 3, 1, -1, None],
}
)

q = ldf.select(pl.col("a").str.count_matches("e", literal=True))

with pytest.raises(NotImplementedError):
_ = translate_ir(q._ldf.visit())

0 comments on commit 550debc

Please sign in to comment.