web-dev-qa-db-fra.com

Spark UDF pour StructType / Row

J'ai une colonne "StructType" dans spark Dataframe qui a un tableau et une chaîne comme sous-champs. J'aimerais modifier le tableau et renvoyer la nouvelle colonne du même type. Peut Je le traite avec UDF ou quelles sont les alternatives?

import org.Apache.spark.sql.types._
import org.Apache.spark.sql.Row
val sub_schema = StructType(StructField("col1",ArrayType(IntegerType,false),true) :: StructField("col2",StringType,true)::Nil)
val schema = StructType(StructField("subtable", sub_schema,true) :: Nil)
val data = Seq(Row(Row(Array(1,2),"eb")),  Row(Row(Array(3,2,1), "dsf")) )
val rd = sc.parallelize(data)
val df = spark.createDataFrame(rd, schema)
df.printSchema

root
 |-- subtable: struct (nullable = true)
 |    |-- col1: array (nullable = true)
 |    |    |-- element: integer (containsNull = false)
 |    |-- col2: string (nullable = true)

Il semble que j'ai besoin d'un UDF de type Row, quelque chose comme

val u =  udf((x:Row) => x)
       >> Schema for type org.Apache.spark.sql.Row is not supported

Cela a du sens, car Spark ne connaît pas le schéma du type de retour. Malheureusement, udf.register échoue également:

spark.udf.register("foo", (x:Row)=> Row, sub_schema)
     <console>:30: error: overloaded method value register with alternatives: ...
scalaapache-sparkudf
16
Danil Kirsanov

s'avère que vous pouvez passer le schéma de résultat en tant que deuxième paramètre UDF:

val u =  udf((x:Row) => x, sub_schema)
15
Danil Kirsanov

Oui, vous pouvez le faire avec UDF. Pour plus de simplicité, j'ai pris votre exemple avec les classes de cas et j'ai changé le tableau en ajoutant 2 à chaque valeur:

case class Root(subtable: Subtable)
case class Subtable(col1: Seq[Int], col2: String)

val df = spark.createDataFrame(Seq(
  Root(Subtable(Seq(1, 2, 3), "toto")),
  Root(Subtable(Seq(10, 20, 30), "tata"))
))

val myUdf = udf((subtable: Row) =>
  Subtable(subtable.getSeq[Int](0).map(_ + 2), subtable.getString(1))
)
val result = df.withColumn("subtable_new", myUdf(df("subtable")))
result.printSchema()
result.show(false)

imprimera:

root
 |-- subtable: struct (nullable = true)
 |    |-- col1: array (nullable = true)
 |    |    |-- element: integer (containsNull = false)
 |    |-- col2: string (nullable = true)
 |-- subtable_new: struct (nullable = true)
 |    |-- col1: array (nullable = true)
 |    |    |-- element: integer (containsNull = false)
 |    |-- col2: string (nullable = true)

+-------------------------------+-------------------------------+
|subtable                       |subtable_new                   |
+-------------------------------+-------------------------------+
|[WrappedArray(1, 2, 3),toto]   |[WrappedArray(3, 4, 5),toto]   |
|[WrappedArray(10, 20, 30),tata]|[WrappedArray(12, 22, 32),tata]|
+-------------------------------+-------------------------------+
5
L. CWI

Tu es sur la bonne piste. Dans ce scénario, UDF vous facilitera la vie. Comme vous l'avez déjà rencontré, UDF ne peut pas renvoyer de types dont spark ne sait pas. Donc, fondamentalement, vous aurez besoin de renvoyer quelque chose qui spark peut facilement sérialiser. Il peut-être un case class ou vous pouvez retourner un Tuple comme (Seq[Int], String). Voici donc une version modifiée de votre code:

def main(args: Array[String]): Unit = {
  import org.Apache.spark.sql.Row
  import org.Apache.spark.sql.functions._
  import org.Apache.spark.sql.types._
  val sub_schema = StructType(StructField("col1", ArrayType(IntegerType, false), true) :: StructField("col2", StringType, true) :: Nil)
  val schema = StructType(StructField("subtable", sub_schema, true) :: Nil)
  val data = Seq(Row(Row(Array(1, 2), "eb")), Row(Row(Array(3, 2, 1), "dsf")))
  val rd = spark.sparkContext.parallelize(data)
  val df = spark.createDataFrame(rd, schema)

  df.printSchema()
  df.show(false)

  val mapArray = (subRows: Row) => {
    // I prefer reading values from row by specifying column names, you may use index also
    val col1 = subRows.getAs[Seq[Int]]("col1")
    val mappedCol1 = col1.map(x => x * x) // Use map based on your requirements
    (mappedCol1, subRows.getAs[String]("col2")) // now mapping is done for col2
  }
  val mapUdf = udf(mapArray)

  val newDf = df.withColumn("col1_mapped", mapUdf(df("subtable")))
  newDf.show(false)
  newDf.printSchema()
}

Veuillez consulter ces liens, ceux-ci peuvent vous donner plus d'informations.

  1. Réponse la plus complète sur l'utilisation d'un schéma complexe: https://stackoverflow.com/a/33850490/4046067
  2. Types de données pris en charge par Spark: https://spark.Apache.org/docs/latest/sql-programming-guide.html#data-types
4
Tawkir