Выполнять арифметические операции над несколькими столбцами в кадре данных Spark

Question

У меня есть вход spark-dataframe с именем df как

+---------------+---+---+---+-----------+
|Main_CustomerID| P1| P2| P3|Total_Count|
+---------------+---+---+---+-----------+
|         725153|  1|  0|  2|          3|
|         873008|  0|  0|  3|          3|
|         625109|  1|  1|  0|          2|
+---------------+---+---+---+-----------+

Здесь Total_Count - это сумма P1,P2,P3, а P1,P2,P3 - это product names. Мне нужно найти frequency каждого продукта, разделив значения продуктов на Total_Count. Мне нужно создать новый spark-dataframe с именем frequencyTable следующим образом,

+---------------+------------------+---+------------------+-----------+
|Main_CustomerID|                P1| P2|                P3|Total_Count|
+---------------+------------------+---+------------------+-----------+
|         725153|0.3333333333333333|0.0|0.6666666666666666|          3|
|         873008|               0.0|0.0|               1.0|          3|
|         625109|               0.5|0.5|               0.0|          2|
+---------------+------------------+---+------------------+-----------+

Я сделал это, используя Scala as,

val df_columns = df.columns.toSeq
var frequencyTable = df
for (index <- df_columns) {
  if (index != "Main_CustomerID" && index != "Total_Count") {
  frequencyTable = frequencyTable.withColumn(index, df.col(index) / df.col("Total_Count"))
}
}

Но я не предпочитаю это для цикла, потому что мой df имеет больший размер. Что такое оптимизированное решение?

Answer 1

Если у вас есть датафрейм как

val df = Seq(
  ("725153", 1, 0, 2, 3),
  ("873008", 0, 0, 3, 3),
  ("625109", 1, 1, 0, 2)
).toDF("Main_CustomerID", "P1", "P2", "P3", "Total_Count")

+---------------+---+---+---+-----------+
|Main_CustomerID|P1 |P2 |P3 |Total_Count|
+---------------+---+---+---+-----------+
|725153         |1  |0  |2  |3          |
|873008         |0  |0  |3  |3          |
|625109         |1  |1  |0  |2          |
+---------------+---+---+---+-----------+

Вы можете просто использовать foldLeft в столбцах, кроме Main_CustomerID, Total_Count, т.е. в P1 P2 и P3

val df_columns = df.columns.toSet - "Main_CustomerID" - "Total_Count" toList

df_columns.foldLeft(df){(tempdf, colName) => tempdf.withColumn(colName, df.col(colName) / df.col("Total_Count"))}.show(false)

что должно дать вам

+---------------+------------------+---+------------------+-----------+
|Main_CustomerID|P1                |P2 |P3                |Total_Count|
+---------------+------------------+---+------------------+-----------+
|725153         |0.3333333333333333|0.0|0.6666666666666666|3          |
|873008         |0.0               |0.0|1.0               |3          |
|625109         |0.5               |0.5|0.0               |2          |
+---------------+------------------+---+------------------+-----------+

Надеюсь, ответ полезен