Initial edits to Cython layer to support size_type.

rapidsai · Jun 2, 2020 · d78b3a2 · d78b3a2
1 parent 0969c70
commit d78b3a2
Show file tree

Hide file tree

Showing 3 changed files with 6 additions and 16 deletions.
diff --git a/python/cudf/cudf/_lib/cpp/filling.pxd b/python/cudf/cudf/_lib/cpp/filling.pxd
@@ -38,5 +38,5 @@ cdef extern from "cudf/filling.hpp" namespace "cudf" nogil:
 
     cdef unique_ptr[table] repeat(
         const table_view & input,
-        const scalar & count
+        size_type count
     ) except +
diff --git a/python/cudf/cudf/_lib/filling.pyx b/python/cudf/cudf/_lib/filling.pyx
@@ -56,14 +56,8 @@ def fill(Column destination, int begin, int end, Scalar value):
 def repeat(Table inp, object count, bool check_count=False):
     if isinstance(count, Column):
         return _repeat_via_column(inp, count, check_count)
-
-    if isinstance(count, Scalar):
-        return _repeat_via_scalar(inp, count)
-
-    raise TypeError(
-        "Expected `count` to be Column or Scalar but got {}"
-        .format(type(count))
-    )
+    else:
+        return _repeat_via_size_type(inp, count)
 
 
 def _repeat_via_column(Table inp, Column count, bool check_count):
@@ -86,15 +80,14 @@ def _repeat_via_column(Table inp, Column count, bool check_count):
     )
 
 
-def _repeat_via_scalar(Table inp, Scalar count):
+def _repeat_via_size_type(Table inp, size_type count):
     cdef table_view c_inp = inp.view()
-    cdef scalar* c_count = count.c_value.get()
     cdef unique_ptr[table] c_result
 
     with nogil:
         c_result = move(cpp_filling.repeat(
             c_inp,
-            c_count[0]
+            count
         ))
 
     return Table.from_unique_ptr(

diff --git a/python/cudf/cudf/core/frame.py b/python/cudf/cudf/core/frame.py
@@ -10,7 +10,6 @@
 import cudf
 import cudf._lib as libcudf
 from cudf._lib.nvtx import annotate
-from cudf._lib.scalar import as_scalar
 from cudf.core.column import as_column, build_categorical_column
 from cudf.utils.dtypes import (
     is_categorical_dtype,
@@ -1140,9 +1139,7 @@ def repeat(self, repeats, axis=None):
         return self._repeat(repeats)
 
     def _repeat(self, count):
-        if is_scalar(count):
-            count = as_scalar(count)
-        else:
+        if not is_scalar(count):
             count = as_column(count)
 
         result = self.__class__._from_table(