连接两个数据框并存储在新数据框中

Question

我有两个Spark数据帧。
数据帧 A：

Col_A1     Col_A2  
1      ["x", "y", "z"]  
2      ["a", "x", "y"]  
3      ["a", "b", "c"]

数据帧 B：

Col_B1  
"x"  
"a"  
"y"

我想检查数据帧 A 的哪些条目在其Col_A2中具有数据帧 B 的"x"，并将其作为新数据帧本身返回。我想反复对数据框 B 的其余条目执行相同的操作。

输出需要如下所示：

数据帧A_x：

Col_A1     Col_A2  
1      ["x", "y", "z"]  
2      ["a", "x", "y"]

数据帧A_a：

Col_A1     Col_A2  
2      ["a", "x", "y"]  
3      ["a", "b", "c"]

数据帧A_y

Col_A1     Col_A2  
1        ["x", "y", "z"]  
2        ["a", "x", "y"]

我尝试使用 udfs 和地图函数，但并没有真正得到我想要的东西。提前谢谢。

Answer 1

如果你的数据帧 B 很小并且可以收集到一个列表中，而且它的不同值的数量很少，你可以为其每个元素编写一个简单的 UDF [更新：请参阅文章末尾以获得更简单的方法];下面是一个'x'示例：

spark.version
# u'2.2.0'
from pyspark.sql import Row
df_a = spark.createDataFrame([Row(1, ["x", "y", "z"]),
                              Row(2, ["a", "x", "y"]),
                              Row(3, ["a", "b", "c"])],
                              ["col_A1", "col_A2"])
@udf('boolean')
def x_isin(v):
  if 'x' in v:
    return True
  else:
    return False
temp_x = df_a.withColumn('x_isin', x_isin(df_a.col_A2))
temp_x.show()
# +------+---------+------+
# |col_A1|   col_A2|x_isin|
# +------+---------+------+
# |     1|[x, y, z]|  true|
# |     2|[a, x, y]|  true|
# |     3|[a, b, c]| false|
# +------+---------+------+
df_a_x = temp_x.filter(temp_x.x_isin==True).drop('x_isin')
df_a_x.show()
# +------+---------+ 
# |col_A1|   col_A2|
# +------+---------+
# |     1|[x, y, z]|
# |     2|[a, x, y]|
# +------+---------+

更新（在玛丽的评论之后）：

感谢 Marie 指出array_contains函数 - 现在您确实不需要 UDF 来构建temp_x：

import pyspark.sql.functions as func
temp_x = df_a.withColumn('x_isin', func.array_contains(df_a.col_A2, 'x'))
temp_x.show() # same result as shown above