From 4207be0e1dc2c4754ae20e17cba84f91a38f5bc5 Mon Sep 17 00:00:00 2001 From: Chris Jarrett Date: Thu, 27 Oct 2022 12:07:56 -0700 Subject: [PATCH] Add unit tests --- dask_sql/physical/rex/core/call.py | 10 ++++------ tests/unit/test_call.py | 3 +++ 2 files changed, 7 insertions(+), 6 deletions(-) diff --git a/dask_sql/physical/rex/core/call.py b/dask_sql/physical/rex/core/call.py index 21f5d2ad6..a66b178dc 100644 --- a/dask_sql/physical/rex/core/call.py +++ b/dask_sql/physical/rex/core/call.py @@ -542,12 +542,10 @@ def __init__(self): super().__init__(self.replace) def replace(self, s, pat, repl): - if isinstance(s, str): - return s.replace(pat, repl) - elif isinstance(s, dd.Series): - return s.str.replace(pat, repl) - else: - raise TypeError("The string expression must be a string or a column name") + if is_frame(s): + s = s.str + + return s.replace(pat, repl) class OverlayOperation(Operation): diff --git a/tests/unit/test_call.py b/tests/unit/test_call.py index 0075c5cb5..05b116af8 100644 --- a/tests/unit/test_call.py +++ b/tests/unit/test_call.py @@ -182,6 +182,9 @@ def test_string_operations(): assert ops_mapping["substring"](a, 2) == " normal string" assert ops_mapping["substring"](a, 2, 2) == " n" assert ops_mapping["initcap"](a) == "A Normal String" + assert ops_mapping["replace"](a, "nor", "") == "a mal string" + assert ops_mapping["replace"](a, "normal", "new") == "a new string" + assert ops_mapping["replace"]("hello", "", "w") == "whwewlwlwow" def test_dates():