PySpark - 按类型选择列

在 PySpark 中，我们可以根据列的类型来选择列，这对于大型数据集非常有用。在本文中，我们将探讨如何在 PySpark 中按类型选择列。

代码示例

以下是一个示例代码片段，演示如何在 PySpark 中按类型选择列：

# 导入必要的库
from pyspark.sql.functions import col

# 创建一个示例 DataFrame
data = [("Bob", 25, True), ("Alice", 30, False), ("Charlie", 35, True)]
df = spark.createDataFrame(data, ["name", "age", "married"])

# 按类型选择列
string_cols = [col for (col, data_type) in df.dtypes if data_type == "string"]
numeric_cols = [col for (col, data_type) in df.dtypes if data_type == "int" or data_type == "double"]
boolean_cols = [col for (col, data_type) in df.dtypes if data_type == "boolean"]

在上面的代码中，我们首先创建了一个示例 DataFrame，包含名字、年龄和婚姻状态三个字段。然后，我们使用 df.dtypes 方法来获取 DataFrame 的列和每列的数据类型。接下来，我们使用列表推导式来选择特定类型的列：字符串、数值和布尔值。

最后，我们可以打印选定的列，例如：

print(string_cols)
# ['name']

print(numeric_cols)
# ['age']

print(boolean_cols)
# ['married']

这样，我们就能够按照列的数据类型选择特定的列了。

总结

在 PySpark 中，我们可以使用列表推导式来选择按列类型分组的列。这对于大型数据集非常有用，因为它允许我们轻松地选择相同类型的列，以方便数据处理和分析。