仅不为字符串列保留在PySpark collect_set中的订单

Question

我正在DataFrame上使用collect_set方法并添加3列。

我的df如下：

id  acc_no  acc_name  cust_id    
1    111      ABC       88    
1    222      XYZ       99

下面是代码段：

from pyspark.sql import Window
import pyspark.sql.functions as F

w = Window.partitionBy('id').orderBy('acc_no')
df1 = df.withColumn(
    'cust_id_new',
    F.collect_set(cust_id).over(w)
).withColumn(
    'acc_no_new',
    F.collect_set(acc_no).over(w)
).withColumn(
    'acc_name_new',
    F.collect_set(acc_name).over(w)
).drop('cust_id').drop('acc_no').drop('acc_name')

在这种情况下，我的输出如下：

id    acc_no     acc_name    cust_id   
1    [111,222]  [XYZ,ABC]    [88,99]

因此，这里的acc_no和cust_id是正确的，但是acc_name的顺序不正确。acc_no 111对应于acc_name ABC，但我们得到XYZ。

有人可以让我知道为什么会这样，怎么解决？

我怀疑此问题仅在字符串列中发生，但是我可能是错的。请帮助...

这类似于下面的线程，但仍然出现错误。

How to maintain sort order in PySpark collect_list and collect multiple lists

Answer 1

我们可以在id列中使用row_number函数，并使用collect_list和sorted_array保留顺序。

from pyspark.sql import Window
import pyspark.sql.functions as F

w = Window.partitionBy('id').orderBy('cust_id')
df1 = df.withColumn(
    'rn',
    row_number(cust_id).over(w)
).groupBy("id").agg(sort_array(collect_list(struct('rn','acc_no','acc_name','cust_id'))).alias('data')) \
.withColumn('grp_acc_no',col('data.acc_no')) \
.withColumn('grp_acc_name',col('data.acc_name')) \
.withColumn('grp_cust_id',col('data.cust_id')) \
.drop('data','acc_no','acc_name','cust_id').show(truncate=False) `````