在pyspark中从Spark数据框中仅选择数字/字符串列名

16

我在 Pyspark (2.1.0) 中有一个 Spark DataFrame,我想要获取仅为数字列或字符串列的名称。

例如,这是我的 DF 的模式:

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。文档:http://spark.apache.org/docs/latest/api/python/pyspark.sql.html#pyspark.sql.types.StructField - Daniel de Paula
3个回答

27

dtypes是一个元组列表(columnName,type),你可以使用简单的筛选器。

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

3
短一点的写法:[c for c, t in df.dtypes if t.startswith('string')] - ZygD

18

PySpark提供了丰富的与模式相关的API types。如@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']

1
您可以按照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。

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接