How to add a new Struct column to a DataFrame

Question

I'm currently trying to extract a database from MongoDB and use Spark to ingest into ElasticSearch with geo_points.

The Mongo database has latitude and longitude values, but ElasticSearch requires them to be casted into the geo_point type.

Is there a way in Spark to copy the lat and lon columns to a new column that is an array or struct?

Any help is appreciated!

zero323 · Accepted Answer · 2016-06-16T12:24:11.457

I assume you start with some kind of flat schema like this:

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

First lets create example data:

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)

An easy way is to use an udf and case class:

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

and we get

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

A hard way is to transform your data and apply schema afterwards:

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

and we get an expected output

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

Creating nested schema from scratch can be tedious so if you can I would recommend the first approach. It can be easily extended if you need more sophisticated structure:

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

and we get expected output:

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

Unfortunately you have no control over nullable field so if is important for your project you'll have to specify schema.

Finally you can use struct function introduced in 1.4:

import org.apache.spark.sql.functions.struct

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

Thanks @zero323 for your thorough answer! This helps a bunch. Would you know how I could do this mapping recursively for nested types? This data is uglier than I hoped. — Kim Ngo, Jul 27 '15 at 16:43
Hi @zero323 - Do you know if there is anyway to use your UDF method of creating a struct if there is more than 10 columns in the new struct? UDFs seem to have a limitation on 10 input variables. — Patrick McGloin, Mar 09 '17 at 08:30
@PatrickMcGloin Does http://stackoverflow.com/q/39494620/1560062 help? — zero323, Mar 09 '17 at 11:19
Hi @zero323. Actually your last suggestion of the "struct" function helped me. I should have read to the bottom first! — Patrick McGloin, Mar 09 '17 at 11:58
@zero323 any advice how to solve this in spark ... https://stackoverflow.com/questions/62933135/dataframe-look-up-and-optimization — BdEngineer, Jul 16 '20 at 10:55

score 7 · Answer 2 · edited Oct 23 '17 at 05:16

7

Try this:

import org.apache.spark.sql.functions._

df.registerTempTable("dt")

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

edited Oct 23 '17 at 05:16

Tobias Glaus

3,008
3
18
37

answered Oct 23 '17 at 04:59

user8817325

71
1
1

How to add a new Struct column to a DataFrame

2 Answers2

Linked