首页 > 解决方案 > 如何将列表列与 pyspark 数据框列相交?

问题描述

我有一个下面的 pyspark 数据框,我需要创建新列 (new_col),它是 X 和 Y 列中的常见项目,不包括 Z 中的项目。

df

id X             Y                    Z            new_col
1 [12,23,1,24]  [13,412,12,23,24]     [12]         [23,24]
2 [1,2,3]       [2,4,5,6]             []           [2]

标签: dataframepyspark

解决方案


如果您的架构如下:

df.printSchema()
#root
# |-- id: long (nullable = true)
# |-- X: array (nullable = true)
# |    |-- element: long (containsNull = true)
# |-- Y: array (nullable = true)
# |    |-- element: long (containsNull = true)
# |-- Z: array (nullable = true)
# |    |-- element: long (containsNull = true)

和您的 pyspark 版本 2.4+,您可以使用array_intersectarray_except

from pyspark.sql.functions import array_except, array_intersect
df=df.withColumn("new_col", array_except(array_intersect("X", "Y"), "Z"))
df.show()
#+---+---------------+---------------------+----+--------+
#|id |X              |Y                    |Z   |new_col |
#+---+---------------+---------------------+----+--------+
#|1  |[12, 23, 1, 24]|[13, 412, 12, 23, 24]|[12]|[23, 24]|
#|2  |[1, 2, 3]      |[2, 4, 5, 6]         |[]  |[2]     |
#+---+---------------+---------------------+----+--------+

推荐阅读