pyspark withColumn, как изменить имя столбца

Question

есть ли способ создать / заполнить столбцы с pyspark 2.1.0, где имя столбца является значением другого столбца?Я попробовал следующее

def createNewColumnsFromValues(dataFrame, colName, targetColName):
  """
  Set value of column colName to targetColName's value
  """
  cols = dataFrame.columns
  #df = dataFrame.withColumn(f.col(colName), f.col(targetColName))
  df = dataFrame.withColumn('x', f.col(targetColName))
  return df

Строка с комментариями не работает, при вызове метода я получаю ошибку

TypeError: 'Column' object is not callable

, тогда как с фиксированным именем (в виде строки) проблем нет,Любая идея о том, как сделать так, чтобы имя столбца происходило от другого, а не только от значения?Я также попытался использовать определение функции UDF в качестве обходного пути с тем же безуспешным результатом.

Спасибо за помощь!

Редактировать:

from pyspark.sql import functions as f

Answer 1

Я нашел решение, которое хорошо масштабируется для нескольких (или не многих) различных значений, для которых мне нужны столбцы.Это обязательно так, иначе число столбцов взорвется.

def createNewColumnsFromValues(dataFrame, colName, targetCol):
  distinctValues = dataFrame.select(colName).distinct().collect()
  for value in distinctValues:
    dataFrame = dataFrame.withColumn(str(value[0]), f.when(f.col(colName) == value[0], f.col(targetCol)).otherwise(f.lit(None)))

return dataFrame

Answer 2

Возможно, вы захотите попробовать следующий код:

test_df = spark.createDataFrame([
    (1,"2",5,1),(3,"4",7,8),              
    ], ("col1","col2","col3","col4"))

def createNewColumnsFromValues(dataFrame, sourceCol, colName, targetCol):
    """
    Set value column colName to targetCol
    """
    for value in sourceCol:
        dataFrame = dataFrame.withColumn(str(value[0]), when(col(colName)==value[0], targetCol).otherwise(None))

    return dataFrame

createNewColumnsFromValues(test_df, test_df.select("col4").collect(), "col4", test_df["col3"]).show()   

Хитрость здесь заключается в том, чтобы сделать select("COLUMNNAME").collect(), чтобы получить список значений в столбце.Затем colName содержит этот список, который представляет собой список строк, где каждая строка имеет один элемент.Таким образом, вы можете напрямую выполнять итерацию по списку и обращаться к элементу в позиции 0. В этом случае необходимо преобразование в строку, чтобы убедиться, что имя столбца нового столбца является строкой.Целевой столбец используется для значений для каждого из отдельных столбцов.Таким образом, результат будет выглядеть так:

+----+----+----+----+----+----+
|col1|col2|col3|col4|   1|   8|
+----+----+----+----+----+----+
|   1|   2|   5|   1|   5|null|
|   3|   4|   7|   8|null|   7|
+----+----+----+----+----+----+