web-dev-qa-db-fra.com

Comment ajouter une nouvelle colonne Struct à un DataFrame

J'essaie actuellement d'extraire une base de données de MongoDB et d'utiliser Spark pour ingérer dans ElasticSearch avec geo_points.

La base de données Mongo a des valeurs de latitude et de longitude, mais ElasticSearch exige qu'elles soient castées dans le geo_point type.

Existe-t-il un moyen dans Spark pour copier les colonnes lat et lon dans une nouvelle colonne qui est un array ou struct?

Toute aide est appréciée!

scalaelasticsearchapache-sparketlapache-spark-sql
19
Kim Ngo

Je suppose que vous commencez avec une sorte de schéma plat comme celui-ci:

root
 |-- lat: double (nullable = false)
 |-- long: double (nullable = false)
 |-- key: string (nullable = false)

Permet d'abord de créer des exemples de données:

import org.Apache.spark.sql.Row
import org.Apache.spark.sql.functions.{col, udf}
import org.Apache.spark.sql.types._

val rdd = sc.parallelize(
    Row(52.23, 21.01, "Warsaw") :: Row(42.30, 9.15, "Corte") :: Nil)

val schema = StructType(
    StructField("lat", DoubleType, false) ::
    StructField("long", DoubleType, false) ::
    StructField("key", StringType, false) ::Nil)

val df = sqlContext.createDataFrame(rdd, schema)

Un moyen simple consiste à utiliser une classe udf et case:

case class Location(lat: Double, long: Double)
val makeLocation = udf((lat: Double, long: Double) => Location(lat, long))

val dfRes = df.
   withColumn("location", makeLocation(col("lat"), col("long"))).
   drop("lat").
   drop("long")

dfRes.printSchema

et nous obtenons

root
 |-- key: string (nullable = false)
 |-- location: struct (nullable = true)
 |    |-- lat: double (nullable = false)
 |    |-- long: double (nullable = false)

Un moyen difficile consiste à transformer vos données et à appliquer le schéma par la suite:

val rddRes = df.
    map{case Row(lat, long, key) => Row(key, Row(lat, long))}

val schemaRes = StructType(
    StructField("key", StringType, false) ::
    StructField("location", StructType(
        StructField("lat", DoubleType, false) ::
        StructField("long", DoubleType, false) :: Nil
    ), true) :: Nil 
)

sqlContext.createDataFrame(rddRes, schemaRes).show

et nous obtenons une sortie attendue

+------+-------------+
|   key|     location|
+------+-------------+
|Warsaw|[52.23,21.01]|
| Corte|  [42.3,9.15]|
+------+-------------+

Créer un schéma imbriqué à partir de zéro peut être fastidieux, donc si vous le pouvez, je recommanderais la première approche. Il peut être facilement étendu si vous avez besoin d'une structure plus sophistiquée:

case class Pin(location: Location)
val makePin = udf((lat: Double, long: Double) => Pin(Location(lat, long))

df.
    withColumn("pin", makePin(col("lat"), col("long"))).
    drop("lat").
    drop("long").
    printSchema

et nous obtenons la sortie attendue:

root
 |-- key: string (nullable = false)
 |-- pin: struct (nullable = true)
 |    |-- location: struct (nullable = true)
 |    |    |-- lat: double (nullable = false)
 |    |    |-- long: double (nullable = false)

Malheureusement, vous n'avez aucun contrôle sur le champ nullable donc si c'est important pour votre projet, vous devrez spécifier le schéma.

Enfin, vous pouvez utiliser la fonction struct introduite dans 1.4:

import org.Apache.spark.sql.functions.struct

df.select($"key", struct($"lat", $"long").alias("location"))
50
zero323

Essaye ça:

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

df.registerTempTable("dt")

dfres = sql("select struct(lat,lon) as colName from dt")
5
user8817325