从 PySpark 数据框中获取特定行

在使用 PySpark 进行数据分析时，经常需要从数据框中获取特定的行。本文将介绍如何使用 PySpark 对数据框进行行的筛选和选择操作。

1. 筛选行

常用的筛选方法有两种：基于条件筛选和基于索引筛选。

1.1 基于条件筛选

基于条件筛选可以使用 where 或 filter 方法来实现。这两个方法的作用是相同的。

例如，假设我们有以下数据框：

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

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

data = [("Alice", "female", 25),
        ("Bob", "male", 30),
        ("Charlie", "male", 20),
        ("David", "male", 18)]

df = spark.createDataFrame(data, ["name", "gender", "age"])

我们可以用 where 或 filter 方法来筛选出符合条件的行：

# 过滤出性别为 female 的行
df.where(col("gender") == "female").show()

# 过滤出年龄大于 25 的行
df.filter(col("age") > 25).show()

输出结果：

+-----+------+---+
| name|gender|age|
+-----+------+---+
|Alice|female| 25|
+-----+------+---+

+----+------+---+
|name|gender|age|
+----+------+---+
| Bob|  male| 30|
+----+------+---+

可以看到，where 和 filter 方法的参数都可以使用 PySpark 内置的 col 函数来指定列名，使用比较运算符来指定比较条件。

1.2 基于索引筛选

基于索引筛选可以使用 take 或 limit 方法来实现。这两个方法的作用是相同的。

例如，假设我们有以下数据框：

data = [(1, "Alice", 25),
        (2, "Bob", 30),
        (3, "Charlie", 20),
        (4, "David", 18)]

df = spark.createDataFrame(data, ["id", "name", "age"])

我们可以用 take 或 limit 方法来筛选出前几行：

# 获取前两行
df.take(2)

# 获取前两行，并转换成 Pandas 数据框
df.limit(2).toPandas()

输出结果：

[(1, 'Alice', 25), (2, 'Bob', 30)]

   id   name  age
0   1  Alice   25
1   2    Bob   30

2. 选择行

选择行可以使用 select 或 selectExpr 方法来实现。这两个方法的作用是相同的。

例如，假设我们有以下数据框：

data = [("Alice", "female", 25),
        ("Bob", "male", 30),
        ("Charlie", "male", 20),
        ("David", "male", 18)]

df = spark.createDataFrame(data, ["name", "gender", "age"])

我们可以用 select 或 selectExpr 方法来选择特定的列：

# 选择名字和年龄两列
df.select("name", "age").show()

# 选择名字和两倍年龄两列，并重命名为 name 和 age_2
df.selectExpr("name", "age * 2 as age_2").show()

输出结果：

+-----+---+
| name|age|
+-----+---+
|Alice| 25|
|  Bob| 30|
|Charlie| 20|
|David| 18|
+-----+---+

+-----+-----+
| name|age_2|
+-----+-----+
|Alice|   50|
|  Bob|   60|
|Charlie|   40|
|David|   36|
+-----+-----+

可以看到，select 方法的参数是列名，多个列名之间用逗号隔开。而 selectExpr 方法的参数是字符串表达式，可以使用 PySpark 内置的函数或算术运算符。重命名列名可以使用 as 关键字。

总结

本文介绍了 PySpark 中如何从数据框中获取特定的行，包括筛选和选择操作。其中，筛选可以基于条件或索引，选择可以选择列名或使用字符串表达式。这些操作在数据分析中非常常用，掌握它们可以提高工作效率。