handle None

rapidsai · Aug 2, 2021 · a7483c0 · a7483c0
1 parent 63d8a06
commit a7483c0
Show file tree

Hide file tree

Showing 2 changed files with 18 additions and 2 deletions.
diff --git 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
@@ -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)