it-swarm.com.ru

Добавление столбца строк в список столбцов в Spark Dataframe

У меня есть датафрейм Spark с несколькими столбцами. Я хочу добавить столбец к фрейму данных, который представляет собой сумму определенного количества столбцов. 

Например, мои данные выглядят так:

ID var1 var2 var3 var4 var5
a   5     7    9    12   13
b   6     4    3    20   17
c   4     9    4    6    9
d   1     2    6    8    1

Я хочу добавить столбец суммирования строк для определенных столбцов:

ID var1 var2 var3 var4 var5   sums
a   5     7    9    12   13    46
b   6     4    3    20   17    50
c   4     9    4    6    9     32
d   1     2    6    8    10    27

Я знаю, что можно добавлять столбцы вместе, если вы знаете конкретные столбцы для добавления: 

val newdf = df.withColumn("sumofcolumns", df("var1") + df("var2"))

Но можно ли передать список имен столбцов и добавить их вместе? Основываясь на этом ответе, который в основном то, что я хочу, но он использует Python API вместо scala ( Добавить сумму столбца в качестве нового столбца в PySpark dataframe ) Я думаю, что-то вроде этого будет работать: 

//Select columns to sum
val columnstosum = ("var1", "var2","var3","var4","var5")

// Create new column called sumofcolumns which is sum of all columns listed in columnstosum
val newdf = df.withColumn("sumofcolumns", df.select(columstosum.head, columnstosum.tail: _*).sum)

Это выбрасывает значение ошибки сумма не является членом org.Apache.spark.sql.DataFrame. Есть ли способ суммировать по столбцам?

Заранее спасибо за помощь. 

16
Sarah

Вы должны попробовать следующее:

import org.Apache.spark.sql.functions._

val sc: SparkContext = ...
val sqlContext = new SQLContext(sc)

import sqlContext.implicits._

val input = sc.parallelize(Seq(
  ("a", 5, 7, 9, 12, 13),
  ("b", 6, 4, 3, 20, 17),
  ("c", 4, 9, 4, 6 , 9),
  ("d", 1, 2, 6, 8 , 1)
)).toDF("ID", "var1", "var2", "var3", "var4", "var5")

val columnsToSum = List(col("var1"), col("var2"), col("var3"), col("var4"), col("var5"))

val output = input.withColumn("sums", columnsToSum.reduce(_ + _))

output.show()

Тогда результат:

+---+----+----+----+----+----+----+
| ID|var1|var2|var3|var4|var5|sums|
+---+----+----+----+----+----+----+
|  a|   5|   7|   9|  12|  13|  46|
|  b|   6|   4|   3|  20|  17|  50|
|  c|   4|   9|   4|   6|   9|  32|
|  d|   1|   2|   6|   8|   1|  18|
+---+----+----+----+----+----+----+
29
Paweł Jurczenko

Легко и просто:

import org.Apache.spark.sql.Column
import org.Apache.spark.sql.functions.{lit, col}

def sum_(cols: Column*) = cols.foldLeft(lit(0))(_ + _)

val columnstosum = Seq("var1", "var2", "var3", "var4", "var5").map(col _)
df.select(sum_(columnstosum: _*))

с эквивалентом Python:

from functools import reduce
from operator import add
from pyspark.sql.functions import lit, col

def sum_(*cols):
    return reduce(add, cols, lit(0))

columnstosum = [col(x) for x in ["var1", "var2", "var3", "var4", "var5"]]
select("*", sum_(*columnstosum))

Оба будут по умолчанию NA, если в строке отсутствует значение. Вы можете использовать функцию DataFrameNaFunctions.fill или coalesce, чтобы избежать этого.

8
zero323

Я предполагаю, что у вас есть датафрейм df. Затем вы можете суммировать все столбцы, кроме вашего идентификатора. Это полезно, когда у вас много столбцов, и вы не хотите вручную указывать имена всех столбцов, как все упомянутые выше. Этот пост имеет тот же ответ. 

val sumAll = df.columns.collect{ case x if x != "ID" => col(x) }.reduce(_ + _)
df.withColumn("sum", sumAll)
2
Abu Shoeb

Вот элегантное решение с использованием Python:

NewDF = OldDF.withColumn('sums', sum(OldDF[col] for col in OldDF.columns[1:]))

Надеюсь, это повлияет на что-то похожее в Spark ... кто-нибудь?.

0
Aerianis