Добавление нового столбца в Data Frame из других столбцов (Spark)

Я использую Spark 1.3.0 и Python. У меня есть dataframe, и я хочу добавить дополнительный столбец, который получен из других столбцов. Как это,

>>old_df.columns
[col_1, col_2, ..., col_m]

>>new_df.columns
[col_1, col_2, ..., col_m, col_n]

где

col_n = col_3 - col_4

Как это сделать в PySpark?

Ответ 1

Одним из способов достижения этого является использование метода withColumn:

old_df = sqlContext.createDataFrame(sc.parallelize(
    [(0, 1), (1, 3), (2, 5)]), ('col_1', 'col_2'))

new_df = old_df.withColumn('col_n', old_df.col_1 - old_df.col_2)

В качестве альтернативы вы можете использовать SQL в зарегистрированной таблице:

old_df.registerTempTable('old_df')
new_df = sqlContext.sql('SELECT *, col_1 - col_2 AS col_n FROM old_df')

Ответ 2

Кроме того, мы можем использовать udf

from pyspark.sql.functions import udf,col
from pyspark.sql.types import IntegerType
from pyspark import SparkContext
from pyspark.sql import SQLContext

sc = SparkContext()
sqlContext = SQLContext(sc)
old_df = sqlContext.createDataFrame(sc.parallelize(
    [(0, 1), (1, 3), (2, 5)]), ('col_1', 'col_2'))
function = udf(lambda col1, col2 : col1-col2, IntegerType())
new_df = old_df.withColumn('col_n',function(col('col_1'), col('col_2')))
new_df.show()

Ответ 3

Это сработало для меня в кирпичах данных с использованием spark.sql

df_converted = spark.sql('select total_bill, tip, sex, case when sex == "Female" then "0" else "1" end as sex_encoded from tips')