Question

我有一个Spark数据框，看起来像这样：

order by

我想通过按ID，国家/地区分组并将“操作”列的唯一值收集到一个数组中来减少此数据帧，但此数组应按日期列进行排序。

例如

select rd.*
from #regression_data rd
order by x;

为了更简洁地解释这一点，我有一些SQL（presto）代码完全可以实现我想要的...我只是在PySpark或SparkSQL中努力做到这一点：

id  country  date        action
 1    A   2019-01-01   suppress
 1    A   2019-01-02   suppress
 2    A   2019-01-03   bid-up
 2    A   2019-01-04   bid-down
 3    C   2019-01-01   no-action
 3    C   2019-01-02   bid-up
 4    D   2019-01-01   suppress

现在这是我在PySpark中的尝试：

id  country action_arr
 1    A      [suppress]
 2    A      [bid-up, bid-down]
 3    C      [no-action, bid-up]
 4    D      [suppress]

然后我想按组找出每组动作的出现次数：

SELECT id, country, array_distinct(array_agg(action ORDER BY date ASC)) AS actions
FROM table
GROUP BY id, country

代码会运行，但是在输出中，sorted_list列与没有任何数组聚合的操作基本相同。有人可以帮忙吗？

编辑：我几乎设法得到了我想要的东西。但是结果并不完全符合之前的结果。谁能解释为什么？解决方案如下：

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

w = Window.partitionBy('action').orderBy('date')

sorted_list_df = df.withColumn('sorted_list', F.collect_set('action').over(w))

Answer 1

IMO，您的窗口定义错误。您应该按要分组的列进行分区，然后为每个组收集一组唯一的值。

IIUC，您只需要执行以下操作：

w = Window.partitionBy(['id', 'country']).orderBy('date')

sorted_list_df = df.withColumn('sorted_list', F.collect_set('action').over(w))

df_new = sorted_list_df.select('id', 'country', 'sorted_list').withColumn("count_of_elems", F.size("sorted_list"))

撤回：

如果使用窗口，则每一行都会有一个新的集合，并且行数将与旧的df相同。本身就不会有汇总，因为我认为那不是您想要的。

下一行将每个组的值汇总为一组。我希望它能为您提供所需的一切：

df_new = sorted_list_df.groupby('id', 'country').agg(F.max('sorted_list').alias('sorted_list')).withColumn("count_of_elems", F.size("sorted_list"))

PySpark：基于另一列的顺序在数据框列上的collect_set

1 个答案: