groupBy 후 group 당 N 개의 record 만 남기기

groupBy 후 group 당 최대 N=300 개만 남긴다고 할 때 

 

방법1

groupByKey + mapGroups 사용

val urlLineDs: Dataset[(String, Seq[String])] = spark.read.text(path)
	.toDF("url", "item")
	.groupByKey(row => row.getAs[String]("url"))
	.mapGroups { case (url, rowIter:  Iterator[Row]) =>
		val itemList: Seq[String] = rowIter.take(300).map(_.getAs[String]("item")).toList
		(url, itemList)
	}

 

 

방법2

window function 사용

val window = Window.partitionBy("url").orderBy(col("item"))
spark.read.text(path)
	.toDF("url", "item")
	.withColumn("_row_number", row_number().over(window))
	.where(col("_row_number") <= 300)
	.groupBy("url")
	.agg(collect_list("item").as("itemList"))