[SparkML] VectorAssembler 알아보기

- 목차

 

소개.

SparkML 에서 VectorAssembler 는 feature 들을 하나의 Vector 로 만들어주는 모듈입니다.

Vector 와 Assemble 은 이름 그대로 feature 들을 모아 하나의 Vector 를 만듭니다.

Assemble 이 "모으다", "집합시키다", "조합하다" 뜻을 가지듯 말이죠.

VectorAssembler 를 통해서 DataFrame 의 Column 들은 하나의 Vector 됩니다.

 

아래 예시는 Boolean, Integer, Double 타입의 칼럼 3개를 하나의 Vector 로 만드는 예시입니다.

from pyspark.ml.feature import VectorAssembler
from pyspark.sql import SparkSession
from pyspark.sql.types import StructType, StructField, BooleanType, IntegerType, DoubleType

spark = SparkSession.builder \
    .appName("test") \
    .master("local[*]") \
    .config("spark.driver.bindAddress", "localhost") \
    .getOrCreate()

data = [
    (True, 20, 15.1),
    (False, 21, 1.14),
    (True, 30, 30.3),
    (False, 31, 2.2)
]

schema = StructType([
    StructField("marriage", BooleanType(), True),
    StructField("age", IntegerType(), True),
    StructField("score", DoubleType(), True),
])

df = spark.createDataFrame(data, schema=schema)

assembler = VectorAssembler(inputCols=["marriage", "age", "score"], outputCol="vector")
transform_df = assembler.transform(df)
vector_df = transform_df.select("vector")
vector_df.show()
spark.stop()

+---------------+
|         vector|
+---------------+
|[1.0,20.0,15.1]|
|[0.0,21.0,1.14]|
|[1.0,30.0,30.3]|
| [0.0,31.0,2.2]|
+---------------+

 

위 결과처럼 True, False 인 BooleanType 의 값은 1.0, 0.0 으로 변환됩니다.

그리고 Boolean, Integer, Double 칼럼들이 모여 하나의 Vector 를 구성합니다.