PySpark DataFrame拆分pred_label字段为score和label列的方法

问题背景 #

现有PySpark DataFrame包含update_ts、text、pred_label三列，其中pred_label为特殊格式的字符串（示例：[{score=0.998854398727417, label=POSITIVE}]），需要将其拆分为label和score独立列，最终保留update_ts、text、label、score四列。

已实现纯Python版本处理逻辑，但在PySpark（Databricks）环境中需要适配。补充说明：pred_label字段为StringDataType，由返回类型为StringType的UDF生成，若修改该UDF的返回类型可简化后续处理。

方案一：修改原UDF返回类型（推荐） #

直接让生成pred_label的UDF返回结构化数据（结构体或结构体数组），避免字符串解析的麻烦，是最简洁的方案。

情况1：UDF返回单个结构体 #

如果原预测逻辑只返回一组score和label，可让UDF返回字典，指定返回类型为StructType：

from pyspark.sql.types import StructType, StructField, StringType, FloatType
from pyspark.sql import functions as F

# 定义预测结果的结构体 schema
pred_schema = StructType([
    StructField("score", FloatType(), nullable=False),
    StructField("label", StringType(), nullable=False)
])

# 修改原UDF的逻辑，返回字典形式的结构化数据
def my_fn(description):
    # 此处替换为你的实际预测逻辑
    score = 0.998854398727417
    label = "POSITIVE"
    return {"score": score, "label": label}

# 重新定义UDF，指定返回类型为pred_schema
my_udf = F.udf(my_fn, pred_schema)

# 生成包含结构化pred_label的DataFrame
df = df.withColumn('pred_label', my_udf(F.col("description")))

# 直接提取结构体中的字段，生成最终DataFrame
final_df = df.select(
    "update_ts",
    "text",
    F.col("pred_label.label").alias("label"),
    F.col("pred_label.score").alias("score")
)

情况2：UDF返回结构体数组 #

如果原逻辑返回的是数组形式（如示例中的[{...}]），则指定返回类型为ArrayType(pred_schema)，再取数组第一个元素：

from pyspark.sql.types import StructType, StructField, StringType, FloatType, ArrayType
from pyspark.sql import functions as F

pred_schema = StructType([
    StructField("score", FloatType(), nullable=False),
    StructField("label", StringType(), nullable=False)
])

def my_fn(description):
    # 原逻辑返回数组形式的预测结果
    return [{"score": 0.998854398727417, "label": "POSITIVE"}]

my_udf = F.udf(my_fn, ArrayType(pred_schema))

df = df.withColumn('pred_label', my_udf(F.col("description")))

# 提取数组第一个元素的字段
final_df = df.select(
    "update_ts",
    "text",
    F.col("pred_label")[0]["label"].alias("label"),
    F.col("pred_label")[0]["score"].alias("score")
)

方案二：处理现有String类型的pred_label #

若无法修改原UDF，可通过字符串解析处理现有格式的pred_label字段。

方法1：内置正则函数解析 #

利用PySpark的正则提取函数直接匹配score和label的值：

from pyspark.sql import functions as F
from pyspark.sql.types import FloatType

# 正则提取label（匹配label=后的大写字母，适配示例格式）
extract_label = F.regexp_extract(F.col("pred_label"), r'label=([A-Z]+)[,\}]', 1).alias("label")
# 正则提取score并转换为浮点型
extract_score = F.regexp_extract(F.col("pred_label"), r'score=([\d\.]+)[,\}]', 1).cast(FloatType()).alias("score")

# 生成最终DataFrame
final_df = df.select(
    "update_ts",
    "text",
    extract_label,
    extract_score
).drop("pred_label")

方法2：自定义UDF解析字符串 #

如果格式更复杂，正则难以覆盖，可编写UDF解析字符串：

from pyspark.sql.types import StructType, StructField, StringType, FloatType
from pyspark.sql import functions as F

# 定义UDF的返回结构
result_schema = StructType([
    StructField("label", StringType(), nullable=True),
    StructField("score", FloatType(), nullable=True)
])

def parse_pred_label(pred_str):
    # 处理空或无效字符串
    if not pred_str or pred_str.strip() == "[]":
        return (None, None)
    # 去掉前后的[{和}]，拆分键值对
    content = pred_str.strip()[1:-1].strip()
    pairs = content.split(", ")
    pred_dict = {}
    for pair in pairs:
        key, value = pair.split("=")
        pred_dict[key] = value
    # 转换数据类型并返回
    label = pred_dict.get("label")
    score = float(pred_dict.get("score")) if pred_dict.get("score") else None
    return (label, score)

# 注册UDF
parse_udf = F.udf(parse_pred_label, result_schema)

# 应用UDF并拆分字段
final_df = df.withColumn("parsed", parse_udf(F.col("pred_label")))\
    .select(
        "update_ts",
        "text",
        F.col("parsed.label").alias("label"),
        F.col("parsed.score").alias("score")
    ).drop("pred_label", "parsed")

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