如何不调用DataFrame获取PySpark UDF返回的pyspark.sql.column.Column的dtype

PySpark 装饰生成的 UDF 对象、以及 UDF 调用后生成的 Column 对象本身都内置了返回类型属性，不需要作用于 DataFrame 也不需要额外单独存储参数即可直接读取。

核心原理 #

• 装饰后的 UDF 对象可以直接通过 .returnType 属性获取定义时指定的返回类型
• UDF 作用于列后得到的 pyspark.sql.column.Column 对象，可以通过 .expr.dataType 属性直接获取该表达式的返回类型

完整实现代码 #

import pyspark.sql.types as st

def get_dtype_return_type_abbreviation(dtype) -> str:
    # 支持直接传入UDF对象、Column表达式对象、或者DataType实例
    if hasattr(dtype, 'returnType'):
        dtype = dtype.returnType
    elif hasattr(dtype, 'expr'):
        dtype = dtype.expr.dataType
    
    # 类型映射逻辑，可按需扩展
    if isinstance(dtype, st.ArrayType):
        element_abbr = get_dtype_return_type_abbreviation(dtype.elementType)
        # 适配示例中list_of_strings的输出要求
        if element_abbr == 'string':
            element_abbr = 'strings'
        return f"list_of_{element_abbr}"
    elif isinstance(dtype, st.MapType):
        key_abbr = get_dtype_return_type_abbreviation(dtype.keyType)
        value_abbr = get_dtype_return_type_abbreviation(dtype.valueType)
        return f"map_of_{key_abbr}_to_{value_abbr}"
    elif isinstance(dtype, st.StringType):
        return "string"
    elif isinstance(dtype, st.IntegerType):
        return "int"
    elif isinstance(dtype, st.LongType):
        return "long"
    elif isinstance(dtype, st.FloatType):
        return "float"
    elif isinstance(dtype, st.DoubleType):
        return "double"
    elif isinstance(dtype, st.BooleanType):
        return "bool"
    # 其他未匹配类型返回spark原生类型简名
    else:
        return dtype.simpleString()

调用验证示例 #

import pyspark.sql.functions as sf
from typing import List

# 你的原始UDF定义
@sf.udf(returnType=st.ArrayType(st.StringType()))
def some_function(text: str) -> List[str]:
    return text.split(' ')

# 按你给出的伪代码逻辑执行
input_column_name = 'some_text_column'
expr = some_function(sf.col(input_column_name))
dtype_abbreviation = get_dtype_return_type_abbreviation(expr) 
expr_renamed = expr.alias(f"{input_column_name}_{dtype_abbreviation}")

# 输出验证
print(expr_renamed.alias())
# 输出结果符合预期：some_text_column_list_of_strings

内容的提问来源于stack exchange，提问作者Joop