PySpark 基於多個參數化條件的join

Question

所以我有兩個 pyspark 數據幀。 我們稱它們為 A 和 B。我想根據多個條件執行左連接。 假設要加入的列名如下：

cond= [A.columnA1==B.columnB1, A.columnA2==B.columnB2]
df=A.join(B,cond,'left')

現在，如果我事先不知道列名，並且想對其進行參數化怎么辦？ 想象一下，允許用戶傳遞兩個包含要加入的列名的列表（每個列表可能超過 2 列，我們不知道）

想象一下，我們有以下要加入的列列表，它接受用戶的輸入：

columnlistA=[]
columnlistB=[]

用戶將為這兩個列表傳遞任意數量的列輸入，但對於這兩個列表，它始終是相同的數字，因此 columnlistA 的第一個元素對應於 columnlistB 的第一個元素，而連接時對應的元素以此類推. 那么，如何編寫連接，以便可以利用這些列列表參數包含在這些數據幀的連接條件中？

Answer 1

您可以通過為數據框使用別名來做到這一點。 像這樣，當您將它們的列名稱為簡單字符串時，您可以訪問它們。

如果我將 dataframe 別名為myDataFrame ，我可以在這樣的字符串中引用它的列：

import pyspark.sql.functions as F
df = spark.createDataFrame(.....)
aliased_df = df.alias("myDataFrame")
F.col("myDataFrame.columnName")  # this is the same as df.columnName

因此，您可以使用它來構造一個動態指定列的列表：

A.alias("dfA").join(
  B.alias("dfB"),
  [F.col("dfA."+col_a) == F.col("dfB."+col_b) for col_a, col_b in zip(columnlistA, columnlistB)],
  'left'
)

Answer 2

以下解決方案基於兩個列表，它將從中生成連接條件。 它假定列之間的相等運算符始終為== 。 您可以通過指定op參數來控制條件之間的二元運算符（僅允許[or, and] ）。

from pyspark.sql.functions import col
from functools import reduce
from pyspark.sql import Column
from pyspark.sql.column import _bin_op

def generate_conditions(left_cols: list, right_cols: list, op: str = "or") -> Column:
  if(not left_cols or not right_cols):
    raise Exception("The lists should not be emtpy.")
  
  if(len(left_cols) != len(right_cols)):
    raise Exception("The lists should have same length.")
    
  if(op not in ["and", "or"]):
    raise Exception("Only [and, or] binary operators are allowed.")
    
  condition_list = reduce(lambda x,y: _bin_op(op)(x, y), [(col(l) == col(r)) for l, r in zip(left_cols, right_cols)])
  
  return condition_list

l = ["a1", "a2", "a3"]
r = ["b1", "b2", "b3"]

join_conditions = generate_conditions(l, r, "or")

print(join_conditions)
# Column<'(((a1 = b1) OR (a2 = b2)) OR (a3 = b3))'>

現在你可以在你的連接中使用它作為A.join(B, join_conditions, 'left')