Spark / Scala повторил создание DataFrames, используя одну и ту же функцию на разных подмножествах данных.

Question

Мой текущий код многократно создает новые фреймы данных (df_1, df_2, df_3) с использованием одной и той же функции, но применяется к разным подмножествам исходного df фрейма данных (например, где ("category == 1 ')).

Я хотел бы создать функцию, которая может автоматизировать создание этих DataFrames.

В следующем примере My DataFrame df имеет три столбца: «категория», «id» и «количество». Предположим, яЕсть 10 категорий. Я хочу суммировать значение столбца «категория», а также подсчитать количество вхождений «категории» на основе различных категорий:

val df_1 = df.where("category == 1")
.groupBy("id")
.agg(sum(when(col("amount") > 0,(col("amount")))).alias("total_incoming_cat_1"),
count(when(col("amount") < 0, (col("amount")))).alias("total_outgoing_cat_1"))

val df_2 = df.where("category == 2")
.groupBy("id")
.agg(sum(when(col("amount") > 0,(col("amount")))).alias("total_incoming_cat_2"),
count(when(col("amount") < 0, (col("amount")))).alias("total_outgoing_cat_2"))

val df_3 = df.where("category == 3")
.groupBy("id")
.agg(sum(when(col("amount") > 0, (col("amount")))).alias("total_incoming_cat_3"),
count(when(col("amount") < 0, (col("amount")))).alias("total_outgoing_cat_3"))

Я хотел бы что-то вроде этого:

def new_dfs(L:List, df:DataFrame): DataFrame={
  for l in L{
    val df_+l df.filter($amount == l)
    .groupBy("id")
    .agg(sum(when(col("amount") > 0, (col("amount")))).alias("total_incoming_cat_"+l),
    count(when(col("amount") < 0, (col("amount")))).alias("total_outgoing_cat_"+l))
  }
}

Answer 1

лучше группировать по категориям и идентификатору

df
.groupBy("category","id")
.agg(sum(when(col("amount") > 0,(col("amount")))).alias("total_incoming_cat"),
count(when(col("amount") < 0, (col("amount")))).alias("total_outgoing_cat"))