在 PySpark 数据框中选择列

在 PySpark 中，我们可以通过 select 函数选择我们需要的列。这个函数可以接收一个或多个列名或表达式作为参数，返回一个包含指定列的新数据框。

基本语法

new_df = df.select(columns)

其中，df 是原数据框，new_df 是新的数据框，columns 是我们需要选择的列，可以是列名或表达式，可以是单个列或多个列。

示例代码

假设我们有下面这个包含三列的数据框：

| name | age | gender | | ---- | --- | ------ | | John | 25 | Male | | Mary | 30 | Female | | Jack | 35 | Male |

我们可以通过以下示例代码选择其中的一列或多列：

# 导入 PySpark 模块
from pyspark.sql import SparkSession

# 创建 SparkSession 对象
spark = SparkSession.builder.appName("Select Columns").getOrCreate()

# 创建数据框 df
data = [("John", 25, "Male"), ("Mary", 30, "Female"), ("Jack", 35, "Male")]
df = spark.createDataFrame(data, ["name", "age", "gender"])

# 选择单个列
name_df = df.select("name")
name_df.show()

# 选择多个列
name_age_df = df.select("name", "age")
name_age_df.show()

# 使用表达式选择列
condition = df["age"] > 30
age_df = df.select("age", condition)
age_df.show()

以上代码会分别输出新的数据框，内容分别为：

选取 name 列的数据框：

+----+
|name|
+----+
|John|
|Mary|
|Jack|
+----+

选取 name 和 age 两列的数据框：

+----+---+
|name|age|
+----+---+
|John| 25|
|Mary| 30|
|Jack| 35|
+----+---+

选取 age 列和一个表达式的结果，该表达式计算 age 大于 30 的行：

+---+------+
|age|(age > 30)|
+---+--------+
| 25|  false|
| 30|  false|
| 35|   true|
+---+--------+

注意事项

在 PySpark 中，对于数据框的列，我们可以使用两种方式来引用它们：

• df.columnName：使用属性
• df['columnName']：访问列对象

因此，在选择列的时候，对于列名，这两种方式是等价的，可以根据个人习惯自由选择。