如何在pyspark中动态修改嵌套列

Question

我有深度级别为 6 的嵌套架构。我在遍历架构中的每个元素以修改列时遇到问题。我有一个列表，其中包含需要修改（散列/匿名）的列名。

我最初的想法是遍历架构中的每个元素并将列与列表项进行比较，并在匹配后进行修改。但我不知道该怎么做。

列表值：['type','name','work','email']

示例架构：

-- abc: struct (nullable = true)
|    |-- xyz: struct (nullable = true)
|    |    |-- abc123: string (nullable = true)
|    |    |-- services: struct (nullable = true)
|    |    |    |-- service: array (nullable = true)
|    |    |    |    |-- element: struct (contains Null = true)
|    |    |    |    |    |-- type: string (nullable = true)
|    |    |    |    |    |-- subtype: string (nullable = true)
|    |    |    |-- name : string(nullable = true)
|    |    |-- details: struct (nullable =true)
|    |    |    |  -- work: (nullable = true)

注意：如果我将架构展平，它会创建 600 多列。因此，我正在寻找一种动态修改列的解决方案（无硬编码）

编辑： 如果它有帮助，我将在我尝试修改值的地方共享我的代码，但它已损坏

def change_nested_field_value(schema, new_df,fields_to_change, parent=""):
    new_schema = []
    
    if isinstance(schema, StringType):
        return schema

    for field in schema:
        full_field_name = field.name
        short_name = full_field_name.split('.')
        
        if parent:
            full_field_name = parent + "." + full_field_name
        
        #print(full_field_name)
        if short_name[-1] not in fields_to_change:            
            if isinstance(field.dataType, StructType):
                inner_schema = change_nested_field_value(field.dataType,new_df, fields_to_change, full_field_name)
                new_schema.append(StructField(field.name, inner_schema))
            elif isinstance(field.dataType, ArrayType):
                inner_schema = change_nested_field_value(field.dataType.elementType, new_df,fields_to_change, full_field_name)
                new_schema.append(StructField(field.name, ArrayType(inner_schema)))
            else:
                new_schema.append(StructField(field.name, field.dataType))
#         else: 
############ this is where I have access to the nested element. I need to modify the value here
#             print(StructField(field.name, field.dataType))
            

    return StructType(new_schema) 

Answer 1

在使用to_json 将结构列转换为json 并最终使用from_json 将json 字符串转换回结构后，您可以使用here 的答案。 from_json 的架构可以从原始结构中推断出来。

from typing import Dict, Any, List
from pyspark.sql import functions as F
from pyspark.sql.types import StringType
import json

def encrypt(s: str) -> str:
    return f"encrypted_{s}"

def walk_dict(struct: Dict[str, Any], fields_to_encrypt: List[str]):
    keys_copy = set(struct.keys())
    for k in keys_copy:
        if k in fields_to_encrypt and isinstance(struct[k], str):
            struct[k] = encrypt(struct[k])
        else:
            walk_fields(struct[k], fields_to_encrypt)

def walk_fields(field: Any, fields_to_encrypt: List[str]):
    if isinstance(field, dict):
        walk_dict(field, fields_to_encrypt)
    if isinstance(field, list):
        [walk_fields(e, fields_to_encrypt) for e in field]
        
def encrypt_fields(json_string: str) -> str:
    fields_to_encrypt = ["type", "subtype", "work", ]
    as_json = json.loads(json_string)
    walk_fields(as_json, fields_to_encrypt)
    return json.dumps(as_json)

field_encryption_udf = F.udf(encrypt_fields, StringType())

data = [{
    "abc": {
        "xyz": {
            "abc123": "abc123Value",
            "services": {
                "service": [
                    {
                        "type": "type_element_1",
                        "subtype": "subtype_element_1",
                    },
                    {
                        "type": "type_element_2",
                        "subtype": "subtype_element_2",
                    }
                ],
                "name": "nameVal"
            },
            "details": {
                "work": "workVal"
            }
        }
    }
}, ]

df = spark.read.json(spark.sparkContext.parallelize(data))

schema_for_abc_column = StructType.fromJson([x.jsonValue()["type"] for x in df.schema.fields if x.name == "abc"][0])

df.withColumn("value_changes_column", F.from_json(field_encryption_udf(F.to_json("abc")), schema_for_abc_column)).show(truncate=False)

"""
+-----------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|abc                                                                                                              |value_changes_column                                                                                                                                               |
+-----------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|{{abc123Value, {workVal}, {nameVal, [{subtype_element_1, type_element_1}, {subtype_element_2, type_element_2}]}}}|{{abc123Value, {encrypted_workVal}, {nameVal, [{encrypted_subtype_element_1, encrypted_type_element_1}, {encrypted_subtype_element_2, encrypted_type_element_2}]}}}|
+-----------------------------------------------------------------------------------------------------------------+-------------------------------------------------------------------------------------------------------------------------------------------------------------------+
"""