您可以从数据帧的架构中检索所有具有 datatype == DecimalType()的列名称,请参见下面的示例(在 Spark 2.4.0 上测试):
更新:只需使用 df.dtypes 就足以检索信息。
from pyspark.sql.functions import col
df = spark.createDataFrame([ (1, 12.3, 1.5, 'test', 13.23) ], ['i1', 'd2', 'f3', 's4', 'd5'])
df = df.withColumn('d2', col('d2').astype('decimal(10,1)')) \
.withColumn('d5', col('d5').astype('decimal(10,2)'))
#DataFrame[i1: bigint, d2: decimal(10,1), f3: double, s4: string, d5: decimal(10,2)]
decimal_cols = [ f[0] for f in df.dtypes if f[1].startswith('decimal') ]
print(decimal_cols)
['d2', 'd5']
只是一个后续:上述方法不适用于array,struct和嵌套数据结构。如果struct中的字段名称不包含空格、点等字符,则可以直接使用 df.dtypes 中的类型。
import re
from pyspark.sql.functions import array, struct, col
decimal_to_double = lambda x: re.sub(r'decimal\(\d+,\d+\)', 'double', x)
df1 = df.withColumn('a6', array('d2','d5')).withColumn('s7', struct('i1','d2'))
# DataFrame[i1: bigint, d2: decimal(10,1), l3: double, s4: string, d5: decimal(10,2), a6: array<decimal(11,2)>, s7: struct<i1:bigint,d2:decimal(10,1)>]
df1.select(*[ col(d[0]).astype(decimal_to_double(d[1])) if 'decimal' in d[1] else col(d[0]) for d in df1.dtypes ])
# DataFrame[i1: bigint, d2: double, l3: double, s4: string, d5: double, a6: array<double>, s7: struct<i1:bigint,d2:double>]
但是,如果任何字段名称StructType()
包含空格、点等,则上述方法可能不起作用。在这种情况下,我建议您检查:df.schema.jsonValue()['fields']
检索和操作 JSON 数据以进行 dtype 转换。