如何在pyspark中聚合和重组dataframe数据（列方式）

Question

我正在尝试根据特定条件聚合pyspark数据框中的数据。我正在尝试根据switchOUT数量将acct对齐到switchIN数量。这样，从中转出资金的帐户就变成了帐户，而其他帐户变成了to_accounts。

我要从数据框中获取数据

+--------+------+-----------+----------+----------+-----------+ 
| person | acct | close_amt | open_amt | switchIN | switchOUT | 
+--------+------+-----------+----------+----------+-----------+ 
| A      | 1    |       125 | 50       | 75       | 0         | 
+--------+------+-----------+----------+----------+-----------+ 
| A      | 2    |       100 | 75       | 25       | 0         | 
+--------+------+-----------+----------+----------+-----------+ 
| A      | 3    |       200 | 300      | 0        | 100       | 
+--------+------+-----------+----------+----------+-----------+ 

到此表

+--------+--------+-----------+----------+----------+
| person | from_acct| to_acct | switchIN | switchOUT| 
+--------+----------+--------+----------+-----------+ 
| A      | 3        |      1 | 75       | 100       |
+--------+----------+--------+----------+-----------+
| A      | 3        |      2 | 25       | 100       | 
+--------+----------+--------+----------+-----------+ 

还有如何做到这一点，使其适用于N行（不只是3个帐户）

到目前为止，我已经使用了这段代码

# define udf
def sorter(l):
  res = sorted(l, key=operator.itemgetter(1))
  return [item[0] for item in res]

def list_to_string(l):
  res = 'from_fund_' +str(l[0]) + '_to_fund_'+str(l[1])
  return res

def listfirstAcc(l):
    res = str(l[0])
    return res

def listSecAcc(l):
    res = str(l[1])
    return res


sort_udf = F.udf(sorter)
list_str = F.udf(list_to_string)
extractFirstFund = F.udf(listfirstAcc)
extractSecondFund = F.udf(listSecAcc)


# Add additional columns
df= df.withColumn("move", sort_udf("list_col").alias("sorted_list"))
df= df.withColumn("move_string", list_str("move"))
df= df.withColumn("From_Acct",extractFirstFund("move"))
df= df.withColumn("To_Acct",extractSecondFund("move"))

我得到的当前结果：

+--------+--------+-----------+----------+----------+
| person | from_acct| to_acct | switchIN | switchOUT| 
+--------+----------+--------+----------+-----------+ 
| A      | 3        |    1,2 | 75       | 100       |
+--------+----------+--------+----------+-----------+