pyspark - 加入 OR 條件

Question

如果至少滿足兩個條件之一，我想加入兩個 pyspark 數據幀。

玩具數據：

df1 = spark.createDataFrame([
    (10, 1, 666),
    (20, 2, 777),
    (30, 1, 888),
    (40, 3, 999),
    (50, 1, 111),
    (60, 2, 222),
    (10, 4, 333),
    (50, None, 444),
    (10, 0, 555),
    (50, 0, 666)
    ],
    ['var1', 'var2', 'other_var'] 
)

df2 = spark.createDataFrame([
    (10, 1),
    (20, 2),
    (30, None),
    (30, 0)
    ],
    ['var1_', 'var2_'] 
)

我想保留df1的所有那些行，其中var1存在於df2.var1_的不同值中，或者var2存在於df2.var2_的不同值中（但不是在這種值為 0 的情況下）。

因此，預期的 output 將是

+----+----+---------+-----+-----+
|var1|var2|other_var|var1_|var2_|
+----+----+---------+-----+-----+
|  10|   1|      666|   10|    1|   # join on both var1 and var2
|  20|   2|      777|   20|    2|   # join on both var1 and var2
|  30|   1|      888|   10|    1|   # join on both var1 and var2
|  50|   1|      111|   10|    1|   # join on var2
|  60|   2|      222|   20|    2|   # join on var2
|  10|   4|      333|   10|    1|   # join on var1
|  10|   0|      555|   10|    1|   # join on var1
+----+----+---------+-----+-----+

在其他嘗試中，我嘗試了

cond = [(df1.var1 == (df2.select('var1_').distinct()).var1_) | (df1.var2 == (df2.filter(F.col('var2_') != 0).select('var2_').distinct()).var2_)]
df1\
    .join(df2, how='inner', on=cond)\
    .show()

+----+----+---------+-----+-----+
|var1|var2|other_var|var1_|var2_|
+----+----+---------+-----+-----+
|  10|   1|      666|   10|    1|
|  20|   2|      777|   20|    2|
|  30|   1|      888|   10|    1|
|  50|   1|      111|   10|    1|
|  30|   1|      888|   30| null|
|  30|   1|      888|   30|    0|
|  60|   2|      222|   20|    2|
|  10|   4|      333|   10|    1|
|  10|   0|      555|   10|    1|
|  10|   0|      555|   30|    0|
|  50|   0|      666|   30|    0|
+----+----+---------+-----+-----+

但是我獲得的行數比預期的要多，並且var2 == 0的行也被保留了。

我究竟做錯了什么？

注意：我沒有使用.isin方法，因為我的實際df2大約有 20k 行，並且我在這里讀到過這種具有大量 ID 的方法的性能可能會很差。

Answer 1

試試下面的條件：

cond = (df2.var2_ != 0) & ((df1.var1 == df2.var1_) | (df1.var2 == df2.var2_))
df1\
    .join(df2, how='inner', on=cond)\
    .show()

+----+----+---------+-----+-----+
|var1|var2|other_var|var1_|var2_|
+----+----+---------+-----+-----+
|  10|   1|      666|   10|    1|
|  30|   1|      888|   10|    1|
|  20|   2|      777|   20|    2|
|  50|   1|      111|   10|    1|
|  60|   2|      222|   20|    2|
|  10|   4|      333|   10|    1|
|  10|   0|      555|   10|    1|
+----+----+---------+-----+-----+

條件應僅包括要連接的兩個數據框中的列。 如果要刪除var2_ = 0 ，可以將它們作為連接條件，而不是作為過濾器。

也不需要指定distinct ，因為它不影響相等條件，而且還增加了不必要的步驟。