PySpark-переименование нескольких столбцов с помощью withColumnRenamed

Я хочу изменить имена двух столбцов с помощью функции spark withColumnRenamed. Конечно, я могу написать:

data = sqlContext.createDataFrame([(1,2), (3,4)], ['x1', 'x2'])
data = (data
       .withColumnRenamed('x1','x3')
       .withColumnRenamed('x2', 'x4'))

но я хочу сделать это за один шаг (имея список/кортеж новых имен). К сожалению, ни это:

data = data.withColumnRenamed(['x1', 'x2'], ['x3', 'x4'])

ни этот:

data = data.withColumnRenamed(('x1', 'x2'), ('x3', 'x4'))

работает. Можно ли сделать это таким образом?

9
apache-spark pyspark rename

2 ответов

невозможно использовать withColumnRenamed. Вы можете использовать метод, потому что он принимает в качестве аргумента список с именами столбцов:

data.toDF(['x3', 'x4'])

также можно переименовать с помощью simple select:

from pyspark.sql.functions import col

mapping = dict(zip(['x1', 'x2'], ['x3', 'x4']))
data.select([col(c).alias(mapping.get(c, c)) for c in data.columns])
19

Я также не мог найти простое решение pyspark, поэтому просто построил свой собственный, похожий на pandas'df.rename(columns={'old_name_1':'new_name_1', 'old_name_2':'new_name_2'}).

def rename_columns(df, columns):
    if isinstance(columns, dict):
        for old_name, new_name in columns.items():
            df = df.withColumnRenamed(old_name, new_name)
        return df
    else:
        raise ValueError("'columns' should be a dict, like {'old_name_1':'new_name_1', 'old_name_2':'new_name_2'}")

таким образом, ваше решение будет выглядеть как data = rename_columns(data, {'x1': 'x3', 'x2': 'x4'})

Это экономит мне несколько строк кода, надеюсь, что это поможет вам тоже.

4