如何在PySpark/Snowpark DataFrame中访问JSON格式列字段

场景说明 #

现有如下Snowpark/PySpark DataFrame，已按table和seq列分区并将数据收集为列表lst：

df = session.create_dataframe(
    [('a', 'xx', {
  "a": {
    "flag": "Y",
    "num": -12,
    "CD": [
      "403",
      "402",
      "405",
      "404"
    ],
    "RANGE_NUM": 999,
    "COL_NM": "COL1"
  }
}),
     ('a', 'xx', {
  "b": {
    "flag": "N",
    "num": -12,
    "CD": [
      "401"
    ],
    "RANGE_NUM": 999,
    "COL_NM": "COL2"
  }
}),
     ('b', 'z', {"c": {
    "flag": "N",
    "num": -12,
    "CD": [
      "411"
    ],
    "RANGE_NUM": 999,
    "COL_NM": "COL3"
  }})],
    ['table', 'seq', 'col']
)

# 分区并收集为列表
w = W.partitionBy(['table','seq']).orderBy(['table','seq'])
lst = df.withColumn('row', F.row_number().over(w)).collect()

需要遍历lst中的每一行，提取COL列里的CD、num、flag值。

解决方案 #

观察COL列的结构：每个元素是仅包含一个顶层键的字典，顶层键对应的值是包含目标字段的子字典。可以通过以下方式提取：

for row in lst:
    tbl = row["TABLE"]
    seq = row["SEQ"]
    # 获取COL列的字典数据
    col_data = row["COL"]
    # 取出字典中唯一的子字典（无需关心顶层键名）
    inner_data = next(iter(col_data.values()))
    # 提取目标字段
    cd = inner_data["CD"]
    num = inner_data["num"]
    flag = inner_data["flag"]
    # 后续可根据需求使用这些变量
    print(f"tbl: {tbl}, seq: {seq}, cd: {cd}, num: {num}, flag: {flag}")

代码说明 #

• row["COL"]直接获取当前行的嵌套字典数据
• next(iter(col_data.values()))：由于每个COL字典只有一个键值对，通过迭代字典的值并取第一个元素，无需硬编码顶层键名（如"a"、"b"、"c"），适配所有行的结构
• 从子字典inner_data中直接通过键名提取CD、num、flag字段

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