web-dev-qa-db-fra.com

Comment ajouter un élément à une colonne de tableau d'un Spark Dataframe?

Supposons que j'ai le DataFrame suivant:

scala> val df1 = Seq("a", "b").toDF("id").withColumn("nums", array(lit(1)))
df1: org.Apache.spark.sql.DataFrame = [id: string, nums: array<int>]

scala> df1.show()
+---+----+
| id|nums|
+---+----+
|  a| [1]|
|  b| [1]|
+---+----+

Et je veux ajouter des éléments au tableau dans la colonne nums, afin d'obtenir quelque chose comme ceci:

+---+-------+
| id|nums   |
+---+-------+
|  a| [1,5] |
|  b| [1,5] |
+---+-------+

Existe-t-il un moyen de le faire en utilisant la méthode .withColumn() du DataFrame? Par exemple.

val df2 = df1.withColumn("nums", append(col("nums"), lit(5)))

J'ai parcouru la documentation de l'API pour Spark, mais je ne trouve rien qui me permettrait de le faire. Je pourrais probablement utiliser split et concat_ws Pour pirater quelque chose ensemble, mais je préférerais une solution plus élégante si possible. Merci.

scalaapache-spark
8
Shafique Jamal

Vous pouvez le faire en utilisant une fonction udf comme

def addValue = udf((array: Seq[Int])=> array ++ Array(5))

df1.withColumn("nums", addValue(col("nums")))
  .show(false)

et vous devriez

+---+------+
|id |nums  |
+---+------+
|a  |[1, 5]|
|b  |[1, 5]|
+---+------+

Mise à jour La méthode alternative consiste à suivre la méthode du jeu de données et à utiliser

df1.map(row => add(row.getAs[String]("id"), row.getAs[Seq[Int]]("nums")++Seq(5)))
  .show(false)

où add est une classe de cas

case class add(id: String, nums: Seq[Int])

J'espère que la réponse est utile

4
Ramesh Maharjan
import org.Apache.spark.sql.functions.{lit, array, array_union}

val df1 = Seq("a", "b").toDF("id").withColumn("nums", array(lit(1)))
val df2 = df1.withColumn("nums", array_union($"nums", lit(Array(5))))
df2.show

+---+------+
| id|  nums|
+---+------+
|  a|[1, 5]|
|  b|[1, 5]|
+---+------+

La array_union() a été ajoutée depuis spark 2.4.0 release le 11/2/2018, 7 mois après avoir posé la question, :) voir https://spark.Apache.org/news/index.html

0
Dorren Chen