從單個 pyspark 數據幀返回多列

Question

我正在嘗試解析單列 pyspark 數據框並獲取具有多列的數據框。我的數據框如下：

   a  b               dic
0  1  2  {'d': 1, 'e': 2}
1  3  4  {'d': 7, 'e': 0}
2  5  6  {'d': 5, 'e': 4}

我想解析 dic 列並按如下方式獲取數據框。 如果可能，我期待使用 Pandas UDF。 我的預期輸出如下：

   a  b  c  d
0  1  2  1  2
1  3  4  7  0
2  5  6  5  4

這是我嘗試解決的方法：

schema = StructType([
    StructField("c", IntegerType()),
    StructField("d", IntegerType())])

@pandas_udf(schema,PandasUDFType.GROUPED_MAP)
def do_someting(dic_col):
    return (pd.DataFrame(dic_col))

df.apply(add_json).show(10)

但這給出了錯誤 'DataFrame' 對象沒有屬性 'apply'

Answer 1

您可以先將簡單引號替換為雙引號將其轉換為 JSON 字符串，然后使用from_json將其轉換為 struct 或 map 列。

如果您知道 dict 的架構，您可以這樣做：

data = [
    (1,   2,  "{'c': 1, 'd': 2}"),
    (3,   4,  "{'c': 7, 'd': 0}"),
    (5,   6,  "{'c': 5, 'd': 4}")
]

df = spark.createDataFrame(data, ["a", "b", "dic"])

schema = StructType([
    StructField("c", StringType(), True),
    StructField("d", StringType(), True)
])

df = df.withColumn("dic", from_json(regexp_replace(col("dic"), "'", "\""), schema))

df.select("a", "b", "dic.*").show(truncate=False)

#+---+---+---+---+
#|a  |b  |c  |d  |
#+---+---+---+---+
#|1  |2  |1  |2  |
#|3  |4  |7  |0  |
#|5  |6  |5  |4  |
#+---+---+---+---+

如果您不知道所有鍵，則可以將其轉換為映射而不是結構，然后將其分解並旋轉以獲取鍵作為列：

df = df.withColumn("dic", from_json(regexp_replace(col("dic"), "'", "\""), MapType(StringType(), StringType())))\
       .select("a", "b", explode("dic"))\
       .groupBy("a", "b")\
       .pivot("key")\
       .agg(first("value"))

Answer 2

嘗試：

#to convert pyspark df into pandas:
df=df.toPandas()

df["d"]=df["dic"].str.get("d")
df["e"]=df["dic"].str.get("e")
df=df.drop(columns=["dic"])

返回：

   a  b  d  e
0  1  2  1  2
1  3  4  7  0
2  5  6  5  4