Pyspark Dataframe 的 GroupBy 替代方案？

Question

我有一个这样的数据集：

timestamp     vars 
2             [1,2]
2             [1,2]
3             [1,2,3]
3             [1,2]

我想要一个这样的数据框。基本上，上述数据框中的每个值都是一个索引，该值的频率是该索引处的值。这种计算是在每个唯一的时间戳上完成的。

timestamp     vars 
2             [0, 2, 2]
3             [0,2,2,1]

现在，我按时间戳分组，并聚合/展平变量（得到类似（1,2,1,2 用于时间戳 2 或 1,2,3,1,2 用于时间戳 3）然后我有一个使用 collections.Counter 来获取 key->value dict 的 udf。然后我将这个 dict 转换为我想要的格式。

groupBy/agg 可以任意大（数组大小可以达到数百万），这似乎是 Window 函数的一个很好的用例，但我不知道如何将它们放在一起。

认为还值得一提的是，我尝试了重新分区、转换为 RDD 并使用 groupByKey。在大型数据集上，两者都非常慢（>24 小时）。

Answer 1

编辑： 正如 cmets 中所讨论的，原始方法的问题可能来自 count，使用了触发不必要的数据扫描的过滤器或聚合函数。下面我们在创建最终数组列之前分解数组并进行聚合（计数）：

from pyspark.sql.functions import collect_list, struct  

df = spark.createDataFrame([(2,[1,2]), (2,[1,2]), (3,[1,2,3]), (3,[1,2])],['timestamp', 'vars'])

df.selectExpr("timestamp", "explode(vars) as var") \
    .groupby('timestamp','var') \
    .count() \
    .groupby("timestamp") \
    .agg(collect_list(struct("var","count")).alias("data")) \
    .selectExpr(
        "timestamp",
        "transform(data, x -> x.var) as indices",
        "transform(data, x -> x.count) as values"
    ).selectExpr(
        "timestamp",
        "transform(sequence(0, array_max(indices)), i -> IFNULL(values[array_position(indices,i)-1],0)) as new_vars"
    ).show(truncate=False)
+---------+------------+
|timestamp|new_vars    |
+---------+------------+
|3        |[0, 2, 2, 1]|
|2        |[0, 2, 2]   |
+---------+------------+

地点：

(1) 我们分解数组并对每个 timestamp + var 执行 count()

(2) groupby timestamp 并创建一个结构数组，其中包含两个字段var 和count

(3) 将structs数组转换成两个数组：indices和values（类似于我们定义的SparseVector）

(4)变换序列sequence(0, array_max(indices))，对于序列中的每一个i，使用array_position在indices数组中找到i的索引，然后同时从values数组中取值位置，见下文：

IFNULL(values[array_position(indices,i)-1],0)

注意函数 array_position 使用从 1 开始的索引，而数组索引是从 0 开始的，因此我们在上面的表达式中有一个 -1。

旧方法：

(1) 使用变换+滤镜/尺寸

from pyspark.sql.functions import flatten, collect_list

df.groupby('timestamp').agg(flatten(collect_list('vars')).alias('data')) \
  .selectExpr(
    "timestamp", 
    "transform(sequence(0, array_max(data)), x -> size(filter(data, y -> y = x))) as vars"
  ).show(truncate=False)
+---------+------------+
|timestamp|vars        |
+---------+------------+
|3        |[0, 2, 2, 1]|
|2        |[0, 2, 2]   |
+---------+------------+

(2)使用aggregate函数：

df.groupby('timestamp').agg(flatten(collect_list('vars')).alias('data')) \
   .selectExpr("timestamp", """ 

     aggregate(   
       data,         
       /* use an array as zero_value, size = array_max(data))+1 and all values are zero */
       array_repeat(0, int(array_max(data))+1),       
       /* increment the ith value of the Array by 1 if i == y */
       (acc, y) -> transform(acc, (x,i) -> IF(i=y, x+1, x))       
     ) as vars   

""").show(truncate=False)