web-dev-qa-db-fra.com

Spark SQL imbriqué avecColumn

J'ai un DataFrame qui a plusieurs colonnes dont certaines sont des structures. Quelque chose comme ça

root
 |-- foo: struct (nullable = true)
 |    |-- bar: string (nullable = true)
 |    |-- baz: string (nullable = true)
 |-- abc: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- def: struct (nullable = true)
 |    |    |    |-- a: string (nullable = true)
 |    |    |    |-- b: integer (nullable = true)
 |    |    |    |-- c: string (nullable = true)

Je veux appliquer un UserDefinedFunction sur la colonne baz pour remplacer baz par une fonction de baz, mais je ne sais pas comment faire. Voici un exemple de la sortie souhaitée (notez que baz est maintenant un int)

root
 |-- foo: struct (nullable = true)
 |    |-- bar: string (nullable = true)
 |    |-- baz: int (nullable = true)
 |-- abc: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- def: struct (nullable = true)
 |    |    |    |-- a: string (nullable = true)
 |    |    |    |-- b: integer (nullable = true)
 |    |    |    |-- c: string (nullable = true)

On dirait DataFrame.withColumn ne fonctionne que sur les colonnes de niveau supérieur mais pas sur les colonnes imbriquées. J'utilise Scala pour ce problème.

Est-ce que quelqu'un peut m'aider avec ça?

Merci

scalaapache-sparkdataframeudf
16
Jon

c'est simple, il suffit d'utiliser un point pour sélectionner les structures imbriquées, par ex. $"foo.baz":

case class Foo(bar:String,baz:String)
case class Record(foo:Foo)

val df = Seq(
   Record(Foo("Hi","There"))
).toDF()


df.printSchema

root
 |-- foo: struct (nullable = true)
 |    |-- bar: string (nullable = true)
 |    |-- baz: string (nullable = true)


val myUDF = udf((s:String) => {
 // do something with s 
  s.toUpperCase
})


df
.withColumn("udfResult",myUDF($"foo.baz"))
.show

+----------+---------+
|       foo|udfResult|
+----------+---------+
|[Hi,There]|    THERE|
+----------+---------+

Si vous souhaitez ajouter le résultat de votre UDF à la structure existante foo, c'est-à-dire pour obtenir:

root
 |-- foo: struct (nullable = false)
 |    |-- bar: string (nullable = true)
 |    |-- baz: string (nullable = true)
 |    |-- udfResult: string (nullable = true)

il y a deux options:

avec withColumn:

df
.withColumn("udfResult",myUDF($"foo.baz"))
.withColumn("foo",struct($"foo.*",$"udfResult"))
.drop($"udfResult")

avec select:

df
.select(struct($"foo.*",myUDF($"foo.baz").as("udfResult")).as("foo"))

EDIT: Remplacement de l'attribut existant dans la structure par le résultat de l'UDF: malheureusement, cela ne pas fonctionne:

df
.withColumn("foo.baz",myUDF($"foo.baz"))

mais peut être fait comme ceci:

// get all columns except foo.baz
val structCols = df.select($"foo.*")
    .columns
    .filter(_!="baz")
    .map(name => col("foo."+name))

df.withColumn(
    "foo",
    struct((structCols:+myUDF($"foo.baz").as("baz")):_*)
)
19
Raphael Roth

Vous pouvez le faire en utilisant la fonction struct comme Raphael Roth a déjà été démontré dans leur réponse ci-dessus. Il existe un moyen plus simple de le faire en utilisant la bibliothèque Make Structs Easy *. La bibliothèque ajoute une méthode withField à la classe Column vous permettant d'ajouter/remplacer des colonnes dans une colonne StructType, de la même manière que la méthode withColumn de la classe DataFrame vous permet d'ajouter/remplacer les colonnes à l'intérieur d'un DataFrame. Pour votre cas d'utilisation spécifique, vous pouvez faire quelque chose comme ceci:

import org.Apache.spark.sql.functions._
import com.github.fqaiser94.mse.methods._

// generate some fake data
case class Foo(bar: String, baz: String)
case class Record(foo: Foo, arrayOfFoo: Seq[Foo])

val df = Seq(
   Record(Foo("Hello", "World"), Seq(Foo("Blue", "Red"), Foo("Green", "Yellow")))
).toDF

df.printSchema

// root
//  |-- foo: struct (nullable = true)
//  |    |-- bar: string (nullable = true)
//  |    |-- baz: string (nullable = true)
//  |-- arrayOfFoo: array (nullable = true)
//  |    |-- element: struct (containsNull = true)
//  |    |    |-- bar: string (nullable = true)
//  |    |    |-- baz: string (nullable = true)

df.show(false)

// +--------------+------------------------------+
// |foo           |arrayOfFoo                    |
// +--------------+------------------------------+
// |[Hello, World]|[[Blue, Red], [Green, Yellow]]|
// +--------------+------------------------------+

// example user defined function that capitalizes a given string
val myUdf = udf((s: String) => s.toUpperCase)

// capitalize value of foo.baz
df.withColumn("foo", $"foo".withField("baz", myUdf($"foo.baz"))).show(false)

// +--------------+------------------------------+
// |foo           |arrayOfFoo                    |
// +--------------+------------------------------+
// |[Hello, WORLD]|[[Blue, Red], [Green, Yellow]]|
// +--------------+------------------------------+

J'ai remarqué que vous aviez une question de suivi sur le remplacement d'une colonne imbriquée dans une structure imbriquée dans un tableau. Cela peut également être fait en combinant les fonctions fournies par la bibliothèque Make Structs Easy avec les fonctions fournies par la bibliothèque spark-hofs , comme suit:

import za.co.absa.spark.hofs._

// capitalize the value of foo.baz in each element of arrayOfFoo
df.withColumn("arrayOfFoo", transform($"arrayOfFoo", foo => foo.withField("baz", myUdf(foo.getField("baz"))))).show(false)

// +--------------+------------------------------+
// |foo           |arrayOfFoo                    |
// +--------------+------------------------------+
// |[Hello, World]|[[Blue, RED], [Green, YELLOW]]|
// +--------------+------------------------------+

* Divulgation complète: je suis l'auteur de la bibliothèque Make Structs Easy qui est référencée dans cette réponse.

1
fqaiser94