Add read-only functions on string dtypes to DataFrame.apply and Series.apply

python/cudf/cudf/core/indexed_frame.py

@@ -49,7 +49,13 @@
 from cudf.core.index import Index, RangeIndex, _index_from_columns
 from cudf.core.missing import NA
 from cudf.core.multiindex import MultiIndex
-from cudf.core.udf.utils import _compile_or_get, _supported_cols_from_frame
+from cudf.core.udf.utils import (
+    _compile_or_get,
+    _launch_arg_from_col,
+    _post_process_output_col,
+    _return_col_from_dtype,
+    _supported_cols_from_frame,
+)
 from cudf.utils import docutils
 from cudf.utils.utils import _cudf_nvtx_annotate
 
@@ -1797,20 +1803,15 @@ def _apply(self, func, kernel_getter, *args, **kwargs):
             ) from e
 
         # Mask and data column preallocated
-        ans_col = cp.empty(len(self), dtype=retty)
+        ans_col = _return_col_from_dtype(retty, len(self))
         ans_mask = cudf.core.column.column_empty(len(self), dtype="bool")
         launch_args = [(ans_col, ans_mask), len(self)]
         offsets = []
 
         # if _compile_or_get succeeds, it is safe to create a kernel that only
         # consumes the columns that are of supported dtype
         for col in _supported_cols_from_frame(self).values():
-            data = col.data
-            mask = col.mask
-            if mask is None:
-                launch_args.append(data)
-            else:
-                launch_args.append((data, mask))
+            launch_args.append(_launch_arg_from_col(col))
             offsets.append(col.offset)
         launch_args += offsets
         launch_args += list(args)
@@ -1820,7 +1821,8 @@ def _apply(self, func, kernel_getter, *args, **kwargs):
         except Exception as e:
             raise RuntimeError("UDF kernel execution failed.") from e
 
-        col = cudf.core.column.as_column(ans_col)
+        col = _post_process_output_col(ans_col, retty)
+
         col.set_base_mask(libcudf.transform.bools_to_mask(ans_mask))
         result = cudf.Series._from_data({None: col}, self._index)
 

python/cudf/cudf/core/udf/__init__.py

@@ -1 +1,2 @@
-from . import typing, lowering
+# Copyright (c) 2022, NVIDIA CORPORATION.
+from . import masked_typing, masked_lowering, strings_typing, strings_lowering

python/cudf/cudf/core/udf/lowering.py → python/cudf/cudf/core/udf/masked_lowering.py

@@ -11,14 +11,16 @@
 )
 from numba.extending import lower_builtin, types
 
+from strings_udf._typing import string_view
+
 from cudf.core.udf import api
 from cudf.core.udf._ops import (
     arith_ops,
     bitwise_ops,
     comparison_ops,
     unary_ops,
 )
-from cudf.core.udf.typing import MaskedType, NAType
+from cudf.core.udf.masked_typing import MaskedType, NAType
 
 
 @cuda_lowering_registry.lower_constant(NAType)
@@ -62,7 +64,6 @@ def masked_scalar_op_impl(context, builder, sig, args):
         result = cgutils.create_struct_proxy(masked_return_type)(
             context, builder
         )
-
         # compute output validity
         valid = builder.and_(m1.valid, m2.valid)
         result.valid = valid
@@ -280,6 +281,7 @@ def pack_return_masked_impl(context, builder, sig, args):
 @cuda_lower(api.pack_return, types.Number)
 @cuda_lower(api.pack_return, types.NPDatetime)
 @cuda_lower(api.pack_return, types.NPTimedelta)
+@cuda_lower(api.pack_return, string_view)
 def pack_return_scalar_impl(context, builder, sig, args):
     outdata = cgutils.create_struct_proxy(sig.return_type)(context, builder)
     outdata.value = args[0]
@@ -351,6 +353,7 @@ def cast_masked_to_masked(context, builder, fromty, toty, val):
 @lower_builtin(api.Masked, types.Number, types.boolean)
 @lower_builtin(api.Masked, types.NPDatetime, types.boolean)
 @lower_builtin(api.Masked, types.NPTimedelta, types.boolean)
+@lower_builtin(api.Masked, string_view, types.boolean)
 def masked_constructor(context, builder, sig, args):
     ty = sig.return_type
     value, valid = args

python/cudf/cudf/core/udf/typing.py → python/cudf/cudf/core/udf/masked_typing.py

@@ -18,6 +18,8 @@
 from numba.core.typing.typeof import typeof
 from numba.cuda.cudadecl import registry as cuda_decl_registry
 
+from strings_udf._typing import StringView, string_view
+
 from cudf.core.missing import NA
 from cudf.core.udf import api
 from cudf.core.udf._ops import (
@@ -32,9 +34,11 @@
     types.Boolean,
     types.NPDatetime,
     types.NPTimedelta,
+    types.PyObject,
 )
 
 
+# Masked scalars of all types
 class MaskedType(types.Type):
     """
     A Numba type consisting of a value of some primitive type
@@ -44,7 +48,12 @@ class MaskedType(types.Type):
     def __init__(self, value):
         # MaskedType in Numba shall be parameterized
         # with a value type
-        if isinstance(value, SUPPORTED_NUMBA_TYPES):
+
+        # TODO - replace object with stringview immediately
+        if isinstance(value, (types.PyObject, StringView)):
+            self.value_type = string_view
+
+        elif isinstance(value, SUPPORTED_NUMBA_TYPES):
             self.value_type = value
         else:
             # Unsupported Dtype. Numba tends to print out the type info
@@ -54,9 +63,9 @@ def __init__(self, value):
             self.value_type = types.Poison(
                 "\n\n\n Unsupported MaskedType. This is usually caused by "
                 "attempting to use a column of unsupported dtype in a UDF. "
-                f"Supported dtypes are {SUPPORTED_NUMBA_TYPES}"
+                f"Supported dtypes are {SUPPORTED_NUMBA_TYPES}\n\n\n"
             )
-        super().__init__(name=f"Masked{self.value_type}")
+        super().__init__(name=f"Masked({self.value_type})")
 
     def __hash__(self):
         """
@@ -145,6 +154,7 @@ class MaskedConstructor(ConcreteTemplate):
             | datetime_cases
             | timedelta_cases
             | {types.boolean}
+            | {string_view}
         )
     ]
 

python/cudf/cudf/core/udf/row_function.py

@@ -6,14 +6,16 @@
 from numba.np import numpy_support
 from numba.types import Record
 
+from strings_udf._typing import DString
+
 from cudf.core.udf.api import Masked, pack_return
+from cudf.core.udf.masked_typing import MaskedType
 from cudf.core.udf.templates import (
     masked_input_initializer_template,
     row_initializer_template,
     row_kernel_template,
     unmasked_input_initializer_template,
 )
-from cudf.core.udf.typing import MaskedType
 from cudf.core.udf.utils import (
     _all_dtypes_from_frame,
     _construct_signature,
@@ -24,6 +26,8 @@
     _supported_dtypes_from_frame,
 )
 
+dstring = DString()
+
 
 def _get_frame_row_type(dtype):
     """
@@ -45,7 +49,13 @@ def _get_frame_row_type(dtype):
     fields = []
     offset = 0
 
-    sizes = [val[0].itemsize for val in dtype.fields.values()]
+    sizes = []
+    for field in dtype.fields.values():
+        if field[0] == np.dtype("object"):
+            sizes.append(dstring.size_bytes)
+        else:
+            sizes.append(field[0].itemsize)
+
     for i, (name, info) in enumerate(dtype.fields.items()):
         # *info* consists of the element dtype, its offset from the beginning
         # of the record, and an optional "title" containing metadata.
@@ -62,7 +72,11 @@ def _get_frame_row_type(dtype):
         fields.append((name, infos))
 
         # increment offset by itemsize plus one byte for validity
-        offset += elemdtype.itemsize + 1
+        if elemdtype == np.dtype("object"):
+            itemsize = dstring.size_bytes
+        else:
+            itemsize = elemdtype.itemsize
+        offset += itemsize + 1
 
         # Align the next member of the struct to be a multiple of the
         # memory access size, per PTX ISA 7.4/5.4.5
@@ -127,10 +141,8 @@ def _get_row_kernel(frame, func, args):
         np.dtype(list(_all_dtypes_from_frame(frame).items()))
     )
     scalar_return_type = _get_udf_return_type(row_type, func, args)
-
     # this is the signature for the final full kernel compilation
     sig = _construct_signature(frame, scalar_return_type, args)
-
     # this row type is used within the kernel to pack up the column and
     # mask data into the dict like data structure the user udf expects
     np_field_types = np.dtype(

python/cudf/cudf/core/udf/scalar_function.py

@@ -4,12 +4,12 @@
 from numba.np import numpy_support
 
 from cudf.core.udf.api import Masked, pack_return
+from cudf.core.udf.masked_typing import MaskedType
 from cudf.core.udf.templates import (
     masked_input_initializer_template,
     scalar_kernel_template,
     unmasked_input_initializer_template,
 )
-from cudf.core.udf.typing import MaskedType
 from cudf.core.udf.utils import (
     _construct_signature,
     _get_kernel,