如何避免火花循環？

Question

我在 PySpark 中有一個如下用例：

df.show()
emp_no dept_no emp_name emp_address location
 1.     10.     ABC.     AAA.         X
 2.     20.     DEF.     CCC.         Y
 3.     20.     GHI.     DDD.         Z
 4.     10.     JKL.     EEE.         Y

而且，我有以下 2 個列表：

dept_list = [10, 20]
location_list = ['Y', 'Z']

現在我正在迭代數據框列表並按如下方式進行連接：

lst = []
for a, b in zip(dept_list, location_list):
  df1 = df.where(col('dept_no' == a))
  df2 = df.where(col('location' == b))
  join_conditions = [df1.dept_no==df2.dept_no, df1.emp_address==df2.emp_address]
  result_df = df1.join(df2, join_conditions, how="inner").select(df1.emp_id, df1.emp_name)
  lst.append(result_df)

最后做所有的聯合如下：

from functools import reduce
from pyspark.sql import DataFrame
final_df = reduce(DataFrame.union, lst)

現在最終結果：

final_df.show()
emp_no emp_name
  4.   JKL
  3.   GHI

我怎樣才能避免這個 FOR 循環？

EDIT1：如果我們對同一列[即位置]有相似的列表，那么如何使用？ 例如：

location_list_1 = ['X', 'Y']
location_list_2 = ['Z', 'Z'] # value can be repeated here, but len(location_list_1)=len(location_list_2)

lst = []
for a, b in zip(location_list_1, location_list_2):
  df1 = df.where(col('location' == a))
  df2 = df.where(col('location' == b))
  join_conditions = [df1.dept_no==df2.dept_no, df1.emp_address==df2.emp_address]
  result_df = df1.join(df2, join_conditions, how="fullouter").select(nvl(df1.emp_id, df2.emp_id), nvl(df1.emp_name, df2.emp_name))
  lst.append(result_df)

而 output 應該是：

emp_no emp_name
  1.   ABC  # from 1st iteration in FOR loop [for location 'X' & 'Z']
  3.   GHI  # from 1st iteration in FOR loop [for location 'X' & 'Z']
  2.   DEF  # from 2nd iteration in FOR loop [for location 'Y' & 'Z']
  4.   JKL  # from 2nd iteration in FOR loop [for location 'Y' & 'Z']
  3.   GHI  # from 2nd iteration in FOR loop [for location 'Y' & 'Z']

# Here, 3-GHI should come twice.

這里同樣，如何避免FOR循環？

Answer 1

如果emp_address不是唯一的，則對其執行自連接（ df.join(df, 'emp_address') ），然后使用以下條件進行過濾： (dept_no, location) in zip(dept_list, location_list) 。

最簡單的方法是創建一個微型 UDF：

def check(x, y):
  return (x,y) in zip(list1, list2)

f = udf(check, StringType(), IntegerType())
df.filter(f(col(x), col(y))

或者，您可以將兩列連接為字符串（使用一些分隔符） - 然后您可以使用isin ：

my_list = [f'{x}:::{y}' for x, y in zip(list1, list2)]
df.filter(concatenate(col(x), lit(':::'), col(y)).isin(my_list))

您可能需要將一些參數（如r_suffix傳遞給連接，因為連接兩側的列名相同。

如果您的連接不是內部連接，則需要在過濾之前先連接。 否則，請先過濾，這樣您就可以加入更少的行。

Answer 2

我認為這應該有效：

from functools import reduce
from operator import __or__
from pyspark.sql import functions as f

final_df = df.alias("a").join(
  df.alias("b"),
  on=reduce(
    __or__,
    [
      (f.col("a.dept_no") == f.lit(a)) & (f.col("b.location") == f.lit(b))
      for a, b in zip(dept_list, location_list)
    ]
  ) & (
    f.col("a.dept_no") == f.col("b.dept_no")
  ) & (
    f.col("a.emp_address") == f.col("b.emp_address")
  )
).select("a.emp_id", "b.emp_name")