选择列 PySpark

在 PySpark 中，选择列很常见，例如从一个 Spark DataFrame 中选择一个或多个列。这里将介绍如何使用 PySpark 选择列，以及一些常用的列选择操作。

我们将基于以下示例数据集展示如何在 PySpark 中选择列：

+----+-------+---------+-------+
| id | sales | country | month |
+----+-------+---------+-------+
|  1 |   100 |      US |   Jan |
|  2 |   200 |      UK |   Jan |
|  3 |   300 |      IN |   Feb |
|  4 |   400 |      CN |   Feb |
|  5 |   500 |      US |   Mar |
+----+-------+---------+-------+

选择单个列

要选取数据集中的单个列，可以使用 select 方法：

from pyspark.sql import SparkSession

spark = SparkSession.builder.appName("select-col").getOrCreate()

data = spark.read.csv("path/to/data.csv", header=True, inferSchema=True)

sales = data.select("sales")
sales.show()

输出：

+-------+
|  sales|
+-------+
|    100|
|    200|
|    300|
|    400|
|    500|
+-------+

在 select 方法中，我们传入要选择的列的名称，PySpark 将会返回一个新的 DataFrame，只包含我们所选列的数据。

选择多个列

如果需要选择多个列，可以将列名作为参数传入 select 方法：

sales_country = data.select("sales", "country")
sales_country.show()

输出：

+-------+-------+
|  sales|country|
+-------+-------+
|    100|     US|
|    200|     UK|
|    300|     IN|
|    400|     CN|
|    500|     US|
+-------+-------+

在这个示例中，我们选取数据集中的 sales 和 country 两列作为展示。同样，PySpark 将会返回一个新的 DataFrame，只包含我们所选列的数据。

选择除指定列之外的所有列

如果需要选择除指定列之外的所有列，可以使用 drop 方法：

without_sales = data.drop("sales")
without_sales.show()

输出：

+----+---------+-------+
| id | country | month |
+----+---------+-------+
|  1 |      US |   Jan |
|  2 |      UK |   Jan |
|  3 |      IN |   Feb |
|  4 |      CN |   Feb |
|  5 |      US |   Mar |
+----+---------+-------+

在这个示例中，我们从数据集中剔除了 sales 列，返回了一个新的 DataFrame。drop 方法接收要删除的列名作为参数，可以同时删除多列。

选择所有列并添加新列

如果需要选择所有列并添加新列，可以使用 withColumn 方法：

new_df = data.withColumn("year", lit("2020"))
new_df.show()

输出：

+----+-------+---------+-------+----+
| id | sales | country | month |year|
+----+-------+---------+-------+----+
|  1 |   100 |      US |   Jan |2020|
|  2 |   200 |      UK |   Jan |2020|
|  3 |   300 |      IN |   Feb |2020|
|  4 |   400 |      CN |   Feb |2020|
|  5 |   500 |      US |   Mar |2020|
+----+-------+---------+-------+----+

在这个示例中，我们使用 withColumn 方法添加了一个新列 year，并将其值设为固定值 2020。withColumn 方法接收两个参数：要添加的列名和要添加的列的值。这个方法返回一个包含新列的新 DataFrame。

总结

这篇文章介绍了 PySpark 中如何选择列。我们讨论了如何选择单个和多个列、删除指定列以及如何在 DataFrame 中添加新列。这些功能是使用 PySpark 进行数据分析和处理的基本操作。使用这些方法，您可以从数据集中提取出所需的信息，并通过添加新的派生特征对数据集进行丰富。