解决cuDF apply函数调用Numba JIT函数返回字符串类型的兼容性问题

import numba, cudf

# 打印当前使用的库版本，有助于问题诊断
print(f"Numba version: {numba.__version__}")    # 示例输出: 0.58.1
print(f"cuDF version: {cudf.__version__}")     # 示例输出: 23.08.00

@numba.njit
def process_string_udf(row):
    """
    一个Numba JIT编译的函数，根据行数据返回字符串。
    """
    input_str = row['str_col']
    scale_val = row['scale']

    if len(input_str) == 0:
        return 'a' + str(scale_val) + 'x'
    elif input_str.startswith('a'):
        return 'b_prefix'
    elif 'example' in input_str:
        return 'c_contains'
    else:
        return 'd_default'

# 创建一个cuDF DataFrame用于演示
data_df = cudf.DataFrame({
    'str_col': ['', 'apple', 'some_example_text', 'banana'],
    'scale': [1, 2, 3, 4]
})

print("原始DataFrame:")
print(data_df)

# 尝试直接应用Numba JIT函数，将导致错误
print("\n尝试直接应用Numba JIT函数 (预期会失败):")
try:
    data_df['processed_col'] = data_df.apply(process_string_udf, axis=1)
    print("应用成功 (不应发生此情况，除非环境或版本有特殊兼容性):")
    print(data_df.head())
except (AttributeError, ValueError) as e:
    print(f"捕获到预期错误：{e}")
    print("错误类型通常为 'AttributeError: 'CPUDispatcher' object has no attribute '__closure__''")
    print("或 'ValueError: user defined function compilation failed.'")

import numba, cudf

@numba.njit
def process_string_udf(row):
    """
    一个Numba JIT编译的函数，根据行数据返回字符串。
    """
    input_str = row['str_col']
    scale_val = row['scale']

    if len(input_str) == 0:
        return 'a' + str(scale_val) + 'x'
    elif input_str.startswith('a'):
        return 'b_prefix'
    elif 'example' in input_str:
        return 'c_contains'
    else:
        return 'd_default'

data_df = cudf.DataFrame({
    'str_col': ['', 'apple', 'some_example_text', 'banana'],
    'scale': [1, 2, 3, 4]
})

# 核心解决方案：使用 .py_func 属性来调用原始的Python函数
data_df['processed_col'] = data_df.apply(process_string_udf.py_func, axis=1)

print("\n成功应用函数后的DataFrame:")
print(data_df.head())