从 pyspark 中的 Spark DF 中仅选择数字/字符串列名称答案

【问题标题】：Selecting only numeric/string columns names from a Spark DF in pyspark从 pyspark 中的 Spark DF 中仅选择数字/字符串列名称
【发布时间】：2022-01-03 06:00:43
【问题描述】：

我在 Pyspark (2.1.0) 中有一个 Spark DataFrame，我希望只获取数字列或字符串列的名称。

例如，这是我的 DF 的 Schema：

root
 |-- Gender: string (nullable = true)
 |-- SeniorCitizen: string (nullable = true)
 |-- MonthlyCharges: double (nullable = true)
 |-- TotalCharges: double (nullable = true)
 |-- Churn: string (nullable = true)

这是我需要的：

num_cols = [MonthlyCharges, TotalCharges]
str_cols = [Gender, SeniorCitizen, Churn]

我该怎么做？

【问题讨论】：

您可能需要使用df.schema 并根据field.dataType 进行过滤，然后得到field.name。文档：spark.apache.org/docs/latest/api/python/…

标签： python apache-spark pyspark apache-spark-sql

【解决方案1】：

dtypes 是可以使用简单过滤器的元组列表（columnNane,type）

 columnList = [item[0] for item in df.dtypes if item[1].startswith('string')]

【讨论】：

【解决方案2】：

PySpark 提供了与架构 types 相关的丰富 API。正如@DanieldePaula 提到的，您可以通过df.schema.fields 访问字段的元数据。

这是一种基于静态类型检查的不同方法：

from pyspark.sql.types import StringType, DoubleType

df = spark.createDataFrame([
  [1, 2.3, "t1"],
  [2, 5.3, "t2"],
  [3, 2.1, "t3"],
  [4, 1.5, "t4"]
], ["cola", "colb", "colc"])

# get string
str_cols = [f.name for f in df.schema.fields if isinstance(f.dataType, StringType)]
# ['colc']

# or double
dbl_cols = [f.name for f in df.schema.fields if isinstance(f.dataType, DoubleType)]
# ['colb']

【讨论】：

【解决方案3】：

您可以按照 zlidme 的建议仅获取字符串（分类列）。要扩展给出的答案，请查看下面的示例。它将为您提供名为 ContinuousCols 的列表中的所有数字（连续）列、名为 categoricalCols 的列表中的所有分类列以及名为 allCols 的列表中的所有列。

data = {'mylongint': [0, 1, 2],
        'shoes': ['blue', 'green', 'yellow'],
        'hous': ['furnitur', 'roof', 'foundation'],
        'C': [1, 0, 0]}

play_df = pd.DataFrame(data)
play_ddf = spark.createDataFrame(play_df)

#store all column names in a list
allCols = [item[0] for item in play_ddf]

#store all column names that are categorical in a list
categoricalCols = [item[0] for item in play_ddf.dtypes if item[1].startswith('string')]

#store all column names that are continous in a list
continuousCols =[item[0] for item in play_ddf.dtypes if item[1].startswith('bigint')]

print(len(allCols), ' - ', len(continuousCols), ' - ', len(categoricalCols))

这将给出结果：4 - 2 - 2

【讨论】：