在pyspark中匹配数组

k0pti3hp  于 2021-07-14  发布在  Spark
关注(0)|答案(2)|浏览(405)

我正在尝试使用pyspark作为aws粘合作业的一部分来操作两个Dataframe。
df1型:

item    tag
1       AB 
2       CD
3       EF
4       QQ

df2型:

key1    key2    tags
A1      B1      [AB]
A1      B2      [AB, CD, EF]
A2      B1      [CD, EF]
A2      B3      [AB, EF, ZZ]

我想通过以下方式将df2中的数组与df1中的标记匹配起来:

item  key1   key2   tag
1     A1     B1     AB
1     A1     B2     AB
2     A1     B2     CD
2     A2     B1     CD
3     A1     B2     EF
3     A2     B1     EF
3     A2     B3     EF

因此,df1中的标记用于根据df2中的标记条目展开行。例如,项1的标记“ab”出现在前两行的df2中的tags数组中。
还要注意4是如何被忽略的,因为标记qq在df2的任何数组中都不存在。
我知道这将是一个内部连接,但我不知道如何匹配df1.tag和df2.tags来拉入key1和key2。任何协助都将不胜感激。

nafvub8i

nafvub8i1#

可以使用 array_contains 条件:

import pyspark.sql.functions as F

result = (df1.join(df2, F.array_contains(df2.tags, df1.tag))
             .select('item', 'key1', 'key2', 'tag')
             .orderBy('item', 'key1', 'key2')
         )

result.show()
+----+----+----+---+
|item|key1|key2|tag|
+----+----+----+---+
|   1|  A1|  B1| AB|
|   1|  A1|  B2| AB|
|   1|  A2|  B3| AB|
|   2|  A1|  B2| CD|
|   2|  A2|  B1| CD|
|   3|  A1|  B2| EF|
|   3|  A2|  B1| EF|
|   3|  A2|  B3| EF|
+----+----+----+---+
rbpvctlc

rbpvctlc2#

import pyspark.sql.functions as F

df = df1.join(
    df2.select('key1', 'key2', F.explode('tags').alias('tag')),
    'tag',
    'inner'
)
df.show()

# +---+----+----+----+

# |tag|item|key1|key2|

# +---+----+----+----+

# | EF|   3|  A1|  B2|

# | EF|   3|  A2|  B1|

# | EF|   3|  A2|  B3|

# | AB|   1|  A1|  B1|

# | AB|   1|  A1|  B2|

# | AB|   1|  A2|  B3|

# | CD|   2|  A1|  B2|

# | CD|   2|  A2|  B1|

# +---+----+----+----+

相关问题