From a7483c046980f805336ab88fdbd15d43a2695dd2 Mon Sep 17 00:00:00 2001 From: Sarah Yurick <53962159+sarahyurick@users.noreply.github.com> Date: Mon, 2 Aug 2021 14:31:05 -0400 Subject: [PATCH] handle None --- python/cudf/cudf/core/column/string.py | 8 +++++++- python/cudf/cudf/tests/test_string.py | 12 +++++++++++- 2 files changed, 18 insertions(+), 2 deletions(-) diff --git a/python/cudf/cudf/core/column/string.py b/python/cudf/cudf/core/column/string.py index 92c32131ef9..27827a4c6b4 100644 --- a/python/cudf/cudf/core/column/string.py +++ b/python/cudf/cudf/core/column/string.py @@ -791,12 +791,18 @@ def repeat(self, repeats: Union[int, Sequence],) -> SeriesOrIndex: 2 ccc dtype: object """ + none_flag = False + if repeats is None: + none_flag = True + repeats = [None] * len(self._column) + if can_convert_to_column(repeats): - return self._return_or_inplace( + result = self._return_or_inplace( libstrings.repeat_sequence( self._column, column.as_column(repeats, dtype="int"), ), ) + return result.astype("float64") if none_flag else result return self._return_or_inplace( libstrings.repeat_scalar(self._column, repeats) diff --git a/python/cudf/cudf/tests/test_string.py b/python/cudf/cudf/tests/test_string.py index 59142368451..66df338df25 100644 --- a/python/cudf/cudf/tests/test_string.py +++ b/python/cudf/cudf/tests/test_string.py @@ -857,7 +857,17 @@ def test_string_contains(ps_gs, pat, regex, flags, flags_raise, na, na_raise): ) @pytest.mark.parametrize( "repeats", - [2, 0, -3, [5, 4, 3, 2, 6], [0, 0, 0, 0, 0], [-1, -2, -3, -4, -5]], + [ + 2, + 0, + -3, + None, + [5, 4, 3, 2, 6], + [5, None, 3, 2, 6], + [0, 0, 0, 0, 0], + [-1, -2, -3, -4, -5], + [None, None, None, None, None], + ], ) def test_string_repeat(data, repeats): ps = pd.Series(data)