通过在pyspark代码中使用sql来交叉两个数据帧之间的操作

时间:2018-02-13 16:58:28

标签: python pyspark spark-dataframe pyspark-sql intersect

我有两个数据框

第一个数据帧:

+---+---+----------------+
|  p|  o| collect_list(s)|
+---+---+----------------+
|  T| V2|[c1, c5]        |
|  T| V1|[c2, c3, c4, c6]|
+---+---+----------------+

第二个数据帧:

+---+---+--------------------+
|  p|  o|     collect_list(s)|
+---+---+--------------------+
|  A| V3|[c1, c2, c3, c4]    |
|  B| V3|[c1, c2, c3, c5, c6]|
+---+---+--------------------+

我们如何根据collect_list列在上述数据框之间进行交叉操作? 如果交叉操作的长度大于最小支持2,则结果应该是在项之间连接的另一个数据帧 如下:

+----------------------------+
|  2-Itemset    |TID         |
+----------------------------+
|[(T,V2),(B,V3)]|[c1, c5]    |
|[(T,V1),(A,V3)]|[c2, c3,c4] |
|[(T,V1),(B,V3)]|[c2,c3,c6]  |
+----------------------------+ 

1 个答案:

答案 0 :(得分:1)

我认为你需要做一个crossJoin(),有些人会使用一些udf来完成这个任务。请考虑以下示例:

创建示例数据

import pyspark.sql.functions as f

data1 = [
    ('T', 'V2', ['c1', 'c5']),
    ('T', 'V1', ['c2', 'c3', 'c4', 'c6'])
]
df1 = sqlCtx.createDataFrame(data1, ["p", "o", "c"])

data2 = [
    ('A', 'V3', ['c1', 'c2', 'c3', 'c4'] ),
    ('B', 'V3', ['c1', 'c2', 'c3', 'c5', 'c6'])
]
df2 = sqlCtx.createDataFrame(data2, ["p", "o", "c"])

定义一些UDF

intersection_udf = f.udf(lambda u, v: list(set(u) & set(v)), ArrayType(StringType()))
intersection_length_udf = f.udf(lambda u, v: len(set(u) & set(v)), IntegerType())

笛卡尔积

df1.alias("l")\
    .crossJoin(df2.alias("r"))\
    .select(
        f.col('l.p').alias('lp'),
        f.col('l.o').alias('lo'),
        f.col('r.p').alias('rp'),
        f.col('r.o').alias('ro'),
        intersection_udf(f.col('l.c'), f.col('r.c')).alias('TID'),
        intersection_length_udf(f.col('l.c'), f.col('r.c')).alias('len')
    )\
    .where(f.col('len') > 1)\
    .select(
        f.struct(f.struct('lp', 'lo'), f.struct('rp', 'ro')).alias('2-Itemset'), 
        'TID'
    )\
    .show()

<强>输出

#+---------------+------------+
#|      2-Itemset|         TID|
#+---------------+------------+
#|[[T,V2],[B,V3]]|    [c1, c5]|
#|[[T,V1],[A,V3]]|[c3, c2, c4]|
#|[[T,V1],[B,V3]]|[c3, c2, c6]|
#+---------------+------------+