Pyspark:获取Column Name的子字符串

Question

我是PySpark的新手，想要更改我的列名，因为大多数列名都有一个烦人的前缀。我的列名是这样的:

e1013_var1
e1014_var2
e1015_var3
Data_date_stamp

如果存在，我想从列名中删除EXXX_，如何做到这一点?因为我还希望所有内容都大写，所以我的代码到目前为止看起来像这样

for col in df.columns:
df= df.withColumnRenamed(col, col.upper())  

帮助是感激的，谢谢!

Answer 1

避免for循环的一个选项是使用toDF重命名Spark数据框的所有列

import re
df_new = df.toDF(*[re.sub('ed+_', '', c).upper() for c in df.columns])
print(df_new.columns)
# ['VAR1', 'VAR2', 'VAR3', 'DATA_DATE_STAMP']

Answer 2

假设恼人的前缀的格式是e[digits]_，我们可以使用re.sub:

import re
for col in df.columns:
new_col = re.sub(r"ed+_", "", col).upper()
df = df.withColumnRenamed(col, new_col)