Allow casting from UDFString back to StringView to call methods in strings_udf

python/cudf/cudf/core/udf/strings_lowering.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 import operator
 
@@ -36,6 +36,7 @@
 
 
 @cuda_lower(len, MaskedType(string_view))
+@cuda_lower(len, MaskedType(udf_string))
 def masked_len_impl(context, builder, sig, args):
     ret = cgutils.create_struct_proxy(sig.return_type)(context, builder)
     masked_sv_ty = sig.args[0]

python/cudf/cudf/core/udf/strings_typing.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 import operator
 
@@ -9,6 +9,7 @@
 
 from strings_udf._typing import (
     StringView,
+    UDFString,
     bool_binary_funcs,
     id_unary_funcs,
     int_binary_funcs,
@@ -25,11 +26,13 @@
 
 masked_typing.MASKED_INIT_MAP[types.pyobject] = string_view
 masked_typing.MASKED_INIT_MAP[string_view] = string_view
+masked_typing.MASKED_INIT_MAP[udf_string] = udf_string
 
 
 def _is_valid_string_arg(ty):
     return (
-        isinstance(ty, MaskedType) and isinstance(ty.value_type, StringView)
+        isinstance(ty, MaskedType)
+        and isinstance(ty.value_type, (StringView, UDFString))
     ) or isinstance(ty, types.StringLiteral)
 
 
@@ -53,7 +56,7 @@ class MaskedStringFunction(AbstractTemplate):
 @register_string_function(len)
 def len_typing(self, args, kws):
     if isinstance(args[0], MaskedType) and isinstance(
-        args[0].value_type, StringView
+        args[0].value_type, (StringView, UDFString)
     ):
         return nb_signature(MaskedType(size_type), args[0])
     elif isinstance(args[0], types.StringLiteral) and len(args) == 1:
@@ -223,4 +226,13 @@ def resolve_valid(self, mod):
         create_masked_unary_attr(f"MaskedType.{func}", udf_string),
     )
 
+
+class MaskedUDFStringAttrs(MaskedStringViewAttrs):
+    key = MaskedType(udf_string)
+
+    def resolve_value(self, mod):
+        return udf_string
+
+
 cuda_decl_registry.register_attr(MaskedStringViewAttrs)
+cuda_decl_registry.register_attr(MaskedUDFStringAttrs)

python/cudf/cudf/tests/test_udf_masked_ops.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2021-2022, NVIDIA CORPORATION.
+# Copyright (c) 2021-2023, NVIDIA CORPORATION.
 import math
 import operator
 
@@ -15,6 +15,7 @@
     comparison_ops,
     unary_ops,
 )
+from cudf.core.udf.api import Masked
 from cudf.core.udf.utils import precompiled
 from cudf.testing._utils import (
     _decimal_series,
@@ -72,6 +73,38 @@ def run_masked_udf_test(func, data, args=(), **kwargs):
     assert_eq(expect, obtain, **kwargs)
 
 
+def run_masked_string_udf_test(func, data, args=(), **kwargs):
+    from strings_udf._typing import sv_to_udf_str
+
+    gdf = data
+    pdf = data.to_pandas(nullable=True)
+
+    def row_wrapper(row):
+        st = row["str_col"]
+        return func(st)
+
+    expect = pdf.apply(row_wrapper, args=args, axis=1)
+
+    func = cuda.jit(device=True)(func)
+    obtain = gdf.apply(row_wrapper, args=args, axis=1)
+    assert_eq(expect, obtain, **kwargs)
+
+    # strings that come directly from input columns are backed by
+    # MaskedType(string_view) types. But new strings that are returned
+    # from functions or operators are backed by MaskedType(udf_string)
+    # types. We need to make sure all of our methods work on both kind
+    # of MaskedType. This function promotes the former to the latter
+    # prior to running the input function
+    def udf_string_wrapper(row):
+        masked_udf_str = Masked(
+            sv_to_udf_str(row["str_col"].value), row["str_col"].valid
+        )
+        return func(masked_udf_str)
+
+    obtain = gdf.apply(udf_string_wrapper, args=args, axis=1)
+    assert_eq(expect, obtain, **kwargs)
+
+
 def run_masked_udf_series(func, data, args=(), **kwargs):
     gsr = data
     psr = data.to_pandas(nullable=True)

python/strings_udf/cpp/src/strings/udf/shim.cu

@@ -1,5 +1,5 @@
 /*
- * Copyright (c) 2022, NVIDIA CORPORATION.
+ * Copyright (c) 2022-2023, NVIDIA CORPORATION.
  *
  * Licensed under the Apache License, Version 2.0 (the "License");
  * you may not use this file except in compliance with the License.
@@ -231,6 +231,17 @@ extern "C" __device__ int udf_string_from_string_view(int* nb_retbal,
   return 0;
 }
 
+extern "C" __device__ int string_view_from_udf_string(int* nb_retval,
+                                                      void const* udf_str,
+                                                      void* str)
+{
+  auto udf_str_ptr = reinterpret_cast<udf_string const*>(udf_str);
+  auto sv_ptr      = new (str) cudf::string_view;
+  *sv_ptr          = cudf::string_view(*udf_str_ptr);
+
+  return 0;
+}
+
 extern "C" __device__ int strip(int* nb_retval,
                                 void* udf_str,
                                 void* const* to_strip,

python/strings_udf/strings_udf/_typing.py

@@ -1,4 +1,4 @@
-# Copyright (c) 2022, NVIDIA CORPORATION.
+# Copyright (c) 2022-2023, NVIDIA CORPORATION.
 
 import operator
 
@@ -39,7 +39,6 @@ def return_type(self):
 
 
 class StringView(types.Type):
-
     np_dtype = np.dtype("object")
 
     def __init__(self):
@@ -120,6 +119,18 @@ def prepare_args(self, ty, val, **kwargs):
 str_view_arg_handler = StrViewArgHandler()
 
 
+# for use in testing only
+def sv_to_udf_str(sv):
+    pass
+
+
+@cuda_decl_registry.register_global(sv_to_udf_str)
+class StringViewToUDFStringDecl(AbstractTemplate):
+    def generic(self, args, kws):
+        if isinstance(args[0], StringView) and len(args) == 1:
+            return nb_signature(udf_string, string_view)
+
+
 # String functions
 @cuda_decl_registry.register_global(len)
 class StringLength(AbstractTemplate):
@@ -242,6 +253,7 @@ def resolve_replace(self, mod):
         create_binary_attr(func, types.boolean),
     )
 
+
 for func in string_return_attrs:
     setattr(
         StringViewAttrs,
@@ -252,9 +264,12 @@ def resolve_replace(self, mod):
 
 for func in int_binary_funcs:
     setattr(
-        StringViewAttrs, f"resolve_{func}", create_binary_attr(func, size_type)
+        StringViewAttrs,
+        f"resolve_{func}",
+        create_binary_attr(func, size_type),
     )
 
+
 for func in id_unary_funcs:
     setattr(
         StringViewAttrs,
@@ -269,7 +284,14 @@ def resolve_replace(self, mod):
         create_identifier_attr(func, udf_string),
     )
 
+
+@cuda_decl_registry.register_attr
+class UDFStringAttrs(StringViewAttrs):
+    key = udf_string
+
+
 cuda_decl_registry.register_attr(StringViewAttrs)
+cuda_decl_registry.register_attr(UDFStringAttrs)
 
 register_stringview_binaryop(operator.eq, types.boolean)
 register_stringview_binaryop(operator.ne, types.boolean)