[Pyspark] Vector Column에서 값 가져오기

업무 관련 작업을 하다가 Vector형식의 Column에서 값을 가져오고싶은데 뜻대로 작동이 되지않았다.

찾아보니까 두가지 방법이 있는것같은데...

제일 간단한건 일단

https://spark.apache.org/docs/3.1.1/api/python/reference/api/pyspark.ml.functions.vector_to_array.html

pyspark.ml.functions.vector_to_array — PySpark 3.1.1 documentation

pyspark에 내장되어있는 vector_to_array함수를 사용하는것인데, 작업환경에는 동작하지않아서 다른 방법을 시도했다.

 

대다수는 UDF함수를 만들어서 작업하는걸로 나오는데, UDF함수에서 단순하게 몇번째 item을 가져오라고 작업할경우 정상적으로 동작하지않았다.

 

에러문을 읽어보니 가져오는 vector가 numpy형식으로 이뤄져있다고 쓰여있어서 tolist함수를 사용했더니 정상적으로 vector to array가 이뤄졌다.

# 동작 안함 x
def getValue(x):
    return x[0]
getVector = F.udf(lambda x: getValue(x), FloatType())

# 정상적으로 동작
def getValue(x):
    return x.tolist()
getVector = F.udf(lambda x: getValue(x), ArrayType(FloatType()))