如何将含混合数据类型列的PySpark DataFrame导入Cosmos DB

问题根源 #

Spark DataFrame是强类型数据结构，同一列只能存储单一数据类型。你之前的when/otherwise转换代码中，虽然逻辑上想返回整数或字符串，但Spark会自动进行类型提升，将整数转为字符串，最终整个列仍为字符串类型，因此导入Cosmos后所有值都变成带引号的字符串。

解决方案 #

1. 将rating列转为AnyType支持混合类型 #

通过自定义UDF将rating列的每个元素转换为对应类型，让列的类型为AnyType（允许存储任意类型的对象）：

from pyspark.sql import functions as F
from pyspark.sql.types import AnyType
from pyspark.sql.functions import udf

def parse_rating(val):
    try:
        # 尝试转为整数
        return int(val)
    except (ValueError, TypeError):
        # 转换失败则保留原字符串
        return val

# 注册UDF，指定返回类型为AnyType
parse_rating_udf = udf(parse_rating, AnyType())

# 转换rating列
df = df.withColumn("rating", parse_rating_udf(F.col("rating")))

2. 修正Cosmos写入配置 #

你的配置存在键名错误，Upsert应改为标准的spark.cosmos.write.upsertEnabled，调整后的完整写入配置如下：

write_config = {
    "spark.cosmos.accountEndpoint": "你的Cosmos账户端点",
    "spark.cosmos.accountKey": "你的Cosmos账户密钥",
    "spark.cosmos.database": "目标数据库名",
    "spark.cosmos.container": "目标容器名",
    "spark.cosmos.write.strategy": "ItemOverwrite",
    "spark.cosmos.serialization.inclusionMode": "NonNull",
    "spark.cosmos.write.bulk.enabled": "true",
    "spark.cosmos.write.upsertEnabled": "true",
    "mode": "Append"
}

# 写入Cosmos
df.write.format("cosmos.oltp").options(**write_config).save()

验证结果 #

处理后，DataFrame的rating列中，id1的5会以整数类型存储，id2的bad会以字符串类型存储，导入Cosmos后将保留各自的原始类型。

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