如何将列表列与 pyspark 数据框列相交？答案

【问题标题】：how to do intersection of list columns with pyspark dataframe columns?如何将列表列与 pyspark 数据框列相交？
【发布时间】：2019-07-23 10:09:35
【问题描述】：

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

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]

【问题讨论】：

标签： dataframe pyspark

【解决方案1】：

如果您的架构如下：

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_intersect 和 array_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]     |
#+---+---------------+---------------------+----+--------+

【讨论】：

感谢 pault，如果 X 或 Y 中的一个也为空，那么有什么方法可以包含 X 或 Y，然后减去 Z？
使用size 检查数组是否为空，同时使用when 执行pyspark equivalent of if-then-else。

【解决方案2】：

你可以使用 withcolumn + udf

import pyspark.sql.functions as F
from pyspark.sql.types import IntegerType, ArrayType

def intersection_function(list1, list2):
    intersection_list = [value for value in list1 if value in list2]
    return intersection_list

udf_intersection = F.udf(intersection_function, ArrayType(IntegerType())

newdf = df.withColumn("new_col", udf_intersection(df["ListColumn1"], df["ListColumn2"]))

【讨论】：

是的，看@pault的回答