如何在 PySpark Dataframe 中的 groupBy 之后计算唯一 ID？

在 PySpark Dataframe 中，使用 groupBy 可以对数据进行分组，然后进行统计、计算等操作。在这个过程中，经常需要对每个组内唯一的 ID 进行计数。

以下是如何在 PySpark Dataframe 中的 groupBy 之后计算唯一 ID 的方法。

示例数据

首先，我们需要准备一个示例数据。假设我们有一个包含用户 ID 和年龄的 PySpark Dataframe，如下所示：

from pyspark.sql.types import StructType, StructField, IntegerType

schema = StructType([
    StructField("user_id", IntegerType(), True),
    StructField("age", IntegerType(), True)
])

data = [
    (1, 20),
    (2, 25),
    (3, 20),
    (4, 30),
    (5, 25),
    (6, 20)
]

df = spark.createDataFrame(data=data, schema=schema)
df.show()

输出：

+-------+---+
|user_id|age|
+-------+---+
|      1| 20|
|      2| 25|
|      3| 20|
|      4| 30|
|      5| 25|
|      6| 20|
+-------+---+

计算每个年龄组的唯一用户数

接下来，我们要计算每个年龄组的唯一用户数。使用 groupBy() 方法对 age 进行分组，然后使用 countDistinct() 方法计算唯一的 user_id 数量。代码如下：

from pyspark.sql.functions import countDistinct

df.groupBy("age").agg(countDistinct("user_id")).show()

输出：

+---+---------------------+
|age|count(DISTINCT user_id)|
+---+---------------------+
| 25|                    2|
| 20|                    3|
| 30|                    1|
+---+---------------------+

计算每个年龄组的所有用户数

如果我们想计算每个年龄组的所有用户数，不需要使用 countDistinct() 方法。只需要使用 count() 方法即可。代码如下：

from pyspark.sql.functions import count

df.groupBy("age").agg(count("user_id")).show()

输出：

+---+--------------+
|age|count(user_id)|
+---+--------------+
| 25|             2|
| 20|             3|
| 30|             1|
+---+--------------+

结论

如此简单，我们就可以使用 PySpark Dataframe 中的 groupBy() 方法和各种聚合函数来计算唯一 ID 或其他统计值了。函数的使用类似 SQL 中的聚合函数，非常方便。