Python pyspark : withColumn (spark dataframe에 새로운 컬럼 추가하기)

 

 

 

 

spark dataframe의 어떤 컬럼의 모든 값에 1을 더한 값을 새로운 컬럼으로 추가하고 싶은 상황에선 어떻게 해야할까요?

withColumn method를 사용하면 됩니다.

 

from pyspark.sql import SparkSession
from pyspark.sql.functions import col
import pandas as pd

spark = SparkSession.builder.getOrCreate()

df_test = pd.DataFrame({
    'a': [1, 2, 3],
    'b': [10.0, 3.5, 7.315],
    'c': ['apple', 'banana', 'tomato']
})

df_spark = spark.createDataFrame(df_test)

df_spark_new = df_spark.withColumn( # 1
    'd', df_spark.a + 1
)

df_spark_new = df_spark_new.withColumn( # 2
    'd', col('a') + 2
)

def my_func(x):
    return x + 3

df_spark_new = df_spark_new.withColumn( # 3
    'd', my_func(df_spark_new.a)
)

df_spark_new = df_spark_new.withColumn( # 3
    'd', my_func(col('a'))
)

df_spark_new.show()


-- Result
+---+------+-------+---+---+---+---+
|  a|     b|      c|  d|  e|  f|  g|
+---+------+-------+---+---+---+---+
|  1|  10.0|  apple|  2|  3|  4|  4|
|  2|   3.5| banana|  3|  4|  5|  5|
|  3| 7.315| tomato|  4|  5|  6|  6|
+---+------+-------+---+---+---+---+

1. withColumn(<new_column_name>, <original_column>)와 같은 구문을 사용할 수 있습니다.

위 예시 코드를 해석해보면

df_spark dataframe에 'd'라는 컬럼을 추가하고

이 컬럼에는 df_spark dataframe에 있는 'a' 컬럼의 모든 값에 1을 더한 값을 넣으라는 뜻이죠.

 

2. 마찬가지로 col 키워드도 사용 가능합니다.

 

3. 함수도 적용 가능합니다.