arrays - Pyspark 数据框 : Count elements in array or list

Question

让我们假设数据帧 df作为:

df.show()

输出:

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

我想做的是 count列中特定元素的编号 list_of_numbers .像这样的东西:

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

到目前为止，我已经尝试创建 udf它完美地工作，但我想知道我是否可以在不定义任何 udf 的情况下做到这一点。 .

Answer 1

您可以分解数组并过滤 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功能。