Dataframe 에서 그룹 별 첫 번째 행 추출하는 방법

데이터 프레임에서 그룹별 첫 번째 행 추출하는 방법을 설명합니다.

SparkSQL의 window 함수를 이용해 DataFrame 에서 각 그룹별 최상위 행을 추출하는 방법을 설명합니다.

테스트 셋 준비

val testData = Seq(
    ("Nike Dunk Low Retro Black", "Nike", "179,000"),
    ("Nike x Off-White Air Force 1 Mid SP White", "Nike", "234,000"),
    ("Nike Air Force 1 '07 Low White", "Nike", "127,000"),
    ("Nike Dunk Low Retro Valerian Blue", "Nike", "166,000"),
    ("Adidas Yeezy Foam RNNR Onyx", "Adidas", "190,000"),
    ("Adidas Yeezy Foam RNNR Stone Sage", "Adidas", "197,000"),
    ("Adidas Yeezy Boost 350 V2 Onyx", "Adidas", "322,000"),
    ("Adidas Yeezy Slide Pure - Re-Release Ver.", "Adidas", "161,000"),
    ("Puma x Maison Kitsune Roma White", "Puma", "200,000"),
    ("Puma x Rick and Morty MB.01 Jasmine Green Energy Rose", "Puma", "319,000"),
    ("Puma x Ader Error Vaderon White", "Puma", "277,000")
)

import spark.implicits._
val df = testData.toDF("model", "brand", "price")
df.show()

// +--------------------+------+-------+
// |               model| brand|  price|
// +--------------------+------+-------+
// |Nike Dunk Low Ret...|  Nike|179,000|
// |Nike x Off-White ...|  Nike|234,000|
// |Nike Air Force 1 ...|  Nike|127,000|
// |Nike Dunk Low Ret...|  Nike|166,000|
// |Adidas Yeezy Foam...|Adidas|190,000|
// |Adidas Yeezy Foam...|Adidas|197,000|
// |Adidas Yeezy Boos...|Adidas|322,000|
// |Adidas Yeezy Slid...|Adidas|161,000|
// |Puma x Maison Kit...|  Puma|200,000|
// |Puma x Rick and M...|  Puma|319,000|
// |Puma x Ader Error...|  Puma|277,000|
// +--------------------+------+-------+

위 테스트 셋은 스포츠웨어 브랜드의 모델별 가격을 정리한 데이터 셋으로 각 브랜드 별 최고 모델에 대해 추출해 봅니다.

‘brand’ 그룹 별 ‘price’ 가 가장 높은 행 추출

import org.apache.spark.sql.expressions.Window

val window = Window.partitionBy("brand").orderBy(col("price").desc)
val extractedDF = df.withColumn("row", row_number.over(window)).
where(col("row").equalTo(1)).drop("row")
extractedDF.show()

// +--------------------+------+-------+
// |               model| brand|  price|
// +--------------------+------+-------+
// |Nike x Off-White ...|  Nike|234,000|
// |Puma x Rick and M...|  Puma|319,000|
// |Adidas Yeezy Boos...|Adidas|322,000|
// +--------------------+------+-------+

먼저 ‘brand’ 그룹을 지정하고 ‘price’ 컬럼에 대해 오름차순으로 정리하는 windowSpec을 생성합니다. DataFrame에서 window 함수인 row_number를 이용하여 새로운 컬럼 ‘row’를 추가합니다. where 조건을 통해 ‘row’ == 1 조건의 행만 추출한 후, “row” 컬럼을 삭제합니다.