根据 PySpark DataFrame 中的条件计算行数

在 PySpark 中，我们通常需要根据特定的条件来选择 DataFrame 中的行（即筛选），并在此基础上计算满足条件的行数。本文将介绍如何使用 PySpark 来实现这一需求。

示例

假设我们有以下 PySpark DataFrame：

from pyspark.sql import SparkSession
from pyspark.sql.functions import *

spark = SparkSession.builder.appName("Example").getOrCreate()

data = [("A", 10), ("A", 20), ("B", 15), ("B", 25), ("C", 5)]
df = spark.createDataFrame(data, ["group", "value"])

df.show()

输出：

+-----+-----+
|group|value|
+-----+-----+
|    A|   10|
|    A|   20|
|    B|   15|
|    B|   25|
|    C|    5|
+-----+-----+

现在，我们想要计算所有 group 为 A 的行的数量。实现方法如下：

count = df.filter(col("group") == "A").count()
print(count)

输出：

2

解释

代码中，我们使用了 filter 方法来指定条件。该方法接受一个条件表达式，返回一个新的 DataFrame，其中只包含符合条件的行。在本例中，我们使用了 col 方法来获取列 group，并使用 Python 的“等于”运算符 == 来判断其是否为 A。筛选后的 DataFrame 中只包含 group 为 A 的行，然后使用 count 方法计算其行数。

总结

本文介绍了如何在 PySpark 中根据特定条件计算 DataFrame 中的行数。具体而言，我们使用了 filter 方法来筛选行，并使用 count 方法计算行数。该方法非常灵活，可以对多个条件进行复杂的筛选。