时间戳

首页 > 解决方案 > Pyspark数据框：计算数组或列表中的元素

arrays - Pyspark数据框：计算数组或列表中的元素

问题描述

让我们假设数据框df为：

df.show()

输出：

+------+----------------+
|letter| list_of_numbers|
+------+----------------+
|     A|    [3, 1, 2, 3]|
|     B|    [1, 2, 1, 1]|
+------+----------------+

我想要做的是对countcolumn 中的特定元素进行编号list_of_numbers。像这样的东西：

+------+----------------+----+
|letter| list_of_numbers|ones|
+------+----------------+----+
|     A|    [3, 1, 2, 3]|   1|
|     B|    [1, 2, 1, 1]|   3|
+------+----------------+----+

到目前为止，我已经尝试过创建udf并且效果很好，但我想知道是否可以在不定义任何udf.

标签： arrayslistdataframepysparkcounting

解决方案

您可以分解数组并过滤1. 然后groupBy和count：

from pyspark.sql.functions import col, count, explode

df.select("*", explode("list_of_numbers").alias("exploded"))\
    .where(col("exploded") == 1)\
    .groupBy("letter", "list_of_numbers")\
    .agg(count("exploded").alias("ones"))\
    .show()
#+------+---------------+----+
#|letter|list_of_numbers|ones|
#+------+---------------+----+
#|     A|   [3, 1, 2, 3]|   1|
#|     B|   [1, 2, 1, 1]|   3|
#+------+---------------+----+

为了保留所有行，即使计数为 0，也可以将exploded列转换为指示变量。然后groupBy和sum。

from pyspark.sql.functions import col, count, explode, sum as sum_

df.select("*", explode("list_of_numbers").alias("exploded"))\
    .withColumn("exploded", (col("exploded") == 1).cast("int"))\
    .groupBy("letter", "list_of_numbers")\
    .agg(sum_("exploded").alias("ones"))\
    .show()

请注意，我已导入pyspark.sql.functions.sum以sum_不覆盖内置sum函数。

推荐阅读