[Spark] dataframe내 Vector타입 컬럼 필터링

spark dataframe의 컬럼이 string이나 단순 numeric 타입인 경우 filter 메서드를 사용하여 쉽게 필터링 할 수 있다.

 

df.filter("컬럼명 != '필터값'")

그러나 컬럼이 Vector타입인 경우 동일한 방식으로는 필터링 할 수 없기 때문에 아래의 방법을 적용해야 한다.

 

 

 

---

 

 

1. 예제 데이터 프레임 생성

%spark
// 임포트 라이브러리
import org.apache.spark.ml.feature.VectorAssembler
import org.apache.spark.ml.linalg.Vectors
import org.apache.spark.ml.linalg.Vector
import org.apache.spark.sql.functions.udf


// 데이터프레임 생성
val df = sc.parallelize(Seq(
  (1, 1, 1), (1, 2, 3), (1, 3, 5), (2, 4, 6),
  (2, 5, 2), (2, 6, 1), (3, 7, 5), (3, 8, 16),
  (1, 1, 1))).toDF("c1", "c2", "c3")


// VectorAssembler 활용 벡터컬럼 생성
val dfVec =  new VectorAssembler()
  .setInputCols(Array("c1", "c2", "c3"))
  .setOutputCol("features")
  .transform(df)

 

 

2. udf 및 제거할 벡터 정의, 벡터 필터링

// udf 정의
def vectors_unequal(vec1: Vector) = udf((vec2: Vector) => !vec1.equals(vec2))


// 제거할 벡터 정의
val vecToRemove = Vectors.dense(1,1,1)




// 벡터 필터링 (from dataframe column)
val filtered = dfVec.where(vectors_unequal(vecToRemove)(dfVec.col("features")))  
val filtered2 = dfVec.filter(vectors_unequal(vecToRemove)($"features")) // Also possible

 

 

 

3. 결과

원본 예제 데이터프레임:

dfVec.show

 

 

 

필터링 결과 데이터프레임 :

filtered.show

 

---

참고.

https://stackoverflow.com/questions/44956855/filter-a-features-column-of-vector-type/45056375#45056375

Filter a "features" column of Vector type