PySpark 将 Dataframe 中的数组中的元素映射到另一个 Dataframe答案

【问题标题】：PySpark Mapping Elements in Array within a Dataframe to another DataframePySpark 将 Dataframe 中的数组中的元素映射到另一个 Dataframe
【发布时间】：2022-01-03 03:45:38
【问题描述】：

我有两个数据框。第一个数据帧有一个数组作为column2 的值，我想将它与第二个数据帧连接起来，以便将数值映射到它们的字符串值。元素的顺序应该保持不变，因为它们按索引对应于column3 中的数组元素。

df_one:

 column1|  column2|        column3
----------------------------------
"thing1"|[1,2,3..]|[0.1,0.2,0.3..]
"thing2"|[1,2,3..]|[0.1,0.2,0.3..]
"thing3"|[1,2,3..]|[0.1,0.2,0.3..]
...

df_two:

columnA|columnB
---------------
      1|"item1"
      2|"item2"
      3|"item3"
...

有没有办法加入这些数据框并像这样选择列：

column1 |                  newColumn|        column3
----------------------------------------------------
"thing1"|["item1","item2","item3"..]|[0.1,0.2,0.3..]
"thing2"|["item1","item2","item3"..]|[0.1,0.2,0.3..]
"thing3"|["item1","item2","item3"..]|[0.1,0.2,0.3..]
...

【问题讨论】：

你可以在 'column2' 上'explode' df_one，然后加入 df_two 并在结果数据集的 'column1' 上使用 'groupBy' 和 'collect_list' 将其作为数组取回。跨度>
'explode' 在性能方面会是一个不错的选择吗？ df_one 中至少有 100 万行，两列中的每个数组都有大约 8k 个元素。
'explode' 不会增加数据量，可能会影响性能的是数据混洗。看到这篇文章：stackoverflow.com/questions/52777421/…

标签： python dataframe apache-spark pyspark

【解决方案1】：

如 cmets 中所述，explode 上 column2 然后 join 上 columnA 是不错的选择。但是，当您对数据进行分组时，我不确定是否会始终保留顺序。

为了确保避免在 python 中使用昂贵的 UDF，您可以使用 posexplode 来跟踪每个元素的位置，然后使用有序窗口函数来构建列表：

df_one = spark.createDataFrame([("thing1", [1, 2, 3], "X"), ("thing2", [1, 2, 3], "Y"), ("thing3", [1, 2, 3], "Z")],
                               ["column1", "column2", "column3"])
df_two = spark.createDataFrame([(1, "item1"), (2, "item2"), (3, "item3")],
                               ["columnA", "columnB"])

w = Window.partitionBy("column1").orderBy("pos")

df_one\
    .select("*", f.posexplode("column2").alias("pos", "columnA"))\
    .join(df_two, ['columnA'])\
    .withColumn("newColumn", f.collect_list("columnB").over(w))\
    .where(f.col("pos")+1 == f.size(f.col("column2")))\
    .select("column1", "newColumn", "column3")\
    .show(truncate=False)

+-------+---------------------+-------+
|column1|newColumn            |column3|
+-------+---------------------+-------+
|thing1 |[item1, item2, item3]|X      |
|thing2 |[item1, item2, item3]|Y      |
|thing3 |[item1, item2, item3]|Z      |
+-------+---------------------+-------+

【讨论】：