過濾嵌套的 JSON 結構並獲取字段名稱作為 Pyspark 中的值

Question

我有以下想要在 PySpark 中解析的復雜數據：

records = '[{"segmentMembership":{"ups":{"FF6KCPTR6AQ0836R":{"lastQualificationTime":"2021-01-16 22:05:11.074357","status":"exited"},"QMS3YRT06JDEUM8O":{"lastQualificationTime":"2021-01-16 22:05:11.074357","status":"realized"},"8XH45RT87N6ZV4KQ":{"lastQualificationTime":"2021-01-16 22:05:11.074357","status":"exited"}}},"_aepgdcdevenablement2":{"emailId":{"address":"stuff@someemail.com"},"person":{"name":{"firstName":"Name2"}},"identities":{"customerid":"PH25PEUWOTA7QF93"}}},{"segmentMembership":{"ups":{"FF6KCPTR6AQ0836R":{"lastQualificationTime":"2021-01-16 22:05:11.074457","status":"realized"},"D45TOO8ZUH0B7GY7":{"lastQualificationTime":"2021-01-16 22:05:11.074457","status":"realized"},"QMS3YRT06JDEUM8O":{"lastQualificationTime":"2021-01-16 22:05:11.074457","status":"existing"}}},"_aepgdcdevenablement2":{"emailId":{"address":"stuff4@someemail.com"},"person":{"name":{"firstName":"TestName"}},"identities":{"customerid":"9LAIHVG91GCREE3Z"}}}]'
df = spark.read.json(sc.parallelize([records]))
df.show()
df.printSchema()

我遇到的問題是segmentMembership object。 JSON object 看起來像這樣：

"segmentMembership": {
      "ups": {
        "FF6KCPTR6AQ0836R": {
          "lastQualificationTime": "2021-01-16 22:05:11.074357",
          "status": "exited"
        },
        "QMS3YRT06JDEUM8O": {
          "lastQualificationTime": "2021-01-16 22:05:11.074357",
          "status": "realized"
        },
        "8XH45RT87N6ZV4KQ": {
          "lastQualificationTime": "2021-01-16 22:05:11.074357",
          "status": "exited"
        }
      }
    }

令人討厭的是，鍵值("FF6KCPTR6AQ0836R", "QMS3YRT06JDEUM8O", "8XH45RT87N6ZV4KQ")最終被定義為 pyspark 中的列。

最后，如果段的狀態是“退出”，我希望得到如下結果。

+--------------------+----------------+---------+------------------+
|address             |customerid      |firstName|segment_id        |
+--------------------+----------------+---------+------------------+
|stuff@someemail.com |PH25PEUWOTA7QF93|Name2    |[8XH45RT87N6ZV4KQ]|
|stuff4@someemail.com|9LAIHVG91GCREE3Z|TestName |[8XH45RT87N6ZV4KQ]|
+--------------------+----------------+---------+------------------+

將數據加載到數據框（上圖）后，我嘗試了以下操作：

dfx = df.select("_aepgdcdevenablement2.emailId.address", "_aepgdcdevenablement2.identities.customerid", "_aepgdcdevenablement2.person.name.firstName", "segmentMembership.ups")
dfx.show(truncate=False)

seg_list = array(*[lit(k) for k in ["8XH45RT87N6ZV4KQ", "QMS3YRT06JDEUM8O"]])
print(seg_list)

# if v["status"] in ['existing', 'realized']

def confusing_compare(ups, seg_list):
    seg_id_filtered_d = dict((k, ups[k]) for k in seg_list if k in ups)

    # This is the line I am having a problem with.
    # seg_id_status_filtered_d = {key for key, value in seg_id_filtered_d.items() if v["status"] in ['existing', 'realized']}
       
    return list(seg_id_filtered_d)

final_conf_dx_pred = udf(confusing_compare, ArrayType(StringType()))
result_df = dfx.withColumn("segment_id", final_conf_dx_pred(dfx.ups, seg_list)).select("address", "customerid", "firstName", "segment_id")

result_df.show(truncate=False)

我無法檢查 dic 值字段中的狀態字段。

Answer 1

您實際上可以在不使用 UDF 的情況下做到這一點。 在這里，我使用架構中存在的所有段名稱並過濾掉那些status = 'exited' 。 您可以根據您想要的細分和狀態對其進行調整。

首先，使用模式字段，獲取所有段名稱的列表，如下所示：

segment_names = df.select("segmentMembership.ups.*").schema.fieldNames()

然后，通過遍歷上面創建的列表並when function 時使用，您可以創建一個列，該列可以將segment_name作為值或 null，具體取決於status ：

active_segments = [
        when(col(f"segmentMembership.ups.{c}.status") != lit("exited"), lit(c)) 
        for c in segment_names
]

最后，添加數組類型的新列segments並使用filter function 從數組中刪除 null 元素（對應於狀態'exited' ）：

dfx = df.withColumn("segments", array(*active_segments)) \
        .withColumn("segments", expr("filter(segments, x -> x is not null)")) \
        .select(
        col("_aepgdcdevenablement2.emailId.address"),
        col("_aepgdcdevenablement2.identities.customerid"),
        col("_aepgdcdevenablement2.person.name.firstName"),
        col("segments").alias("segment_id")
    )

dfx.show(truncate=False)

#+--------------------+----------------+---------+------------------------------------------------------+
#|address             |customerid      |firstName|segment_id                                            |
#+--------------------+----------------+---------+------------------------------------------------------+
#|stuff@someemail.com |PH25PEUWOTA7QF93|Name2    |[QMS3YRT06JDEUM8O]                                    |
#|stuff4@someemail.com|9LAIHVG91GCREE3Z|TestName |[D45TOO8ZUH0B7GY7, FF6KCPTR6AQ0836R, QMS3YRT06JDEUM8O]|
#+--------------------+----------------+---------+------------------------------------------------------+