How to write a set of fields to JSON? - json

Question

I am trying to write few fields in my data frame to JSON. My data structure in the data frame is
Key|col1|col2|col3|col4
key|a |b |c |d
Key|a1 |b1 |c1 |d1
Now i am trying to convert just the col1 to col4 fields to JSON and give a name to the Json field
expected output
[Key,{cols:[{col1:a,col2:b,col3:c,col4:d},{col1:a1,col2:b1,col3:c1,col4:d1}]
I wrote a udf for this.
val summary = udf(
(col1:String, col2:String, col3:String, col4:String) => "{\"cols\":[" + " {\"col1\":" + col1 + ",\"col2\":" + col2 + ",\"col3\":" + col3 + ",\"col4\":" + col4 + "}]}"
)
val result = input.withColumn("Summary",summary('col1,'col2,'col3,'col4))
val result1 = result.select('Key,'Summary)
result1.show(10)
This is my result
[Key,{cols:[{col1:a,col2:b,col3:c,col4:d}]}]
[Key,{cols:[{col1:a1,col2:b1,col3:c1,col4:d1}]}]
As you can see, they are not grouped. Is there a way to group these rows using the UDF itself. I am new to scala/Spark and not able to figure out the proper udf.

Answer 1

// Create your dataset
scala> val ds = Seq((1, "hello", 1L), (2, "world", 2L)).toDF("id", "token", "long")
ds: org.apache.spark.sql.DataFrame = [id: int, token: string ... 1 more field]
// select the fields you want to map to json
scala> ds.select('token, 'long).write.json("your-json")
// check the result
➜ spark git:(master) ✗ ls -ltr your-json/
total 16
-rw-r--r-- 1 jacek staff 27 11 kwi 17:18 part-r-00007-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 0 11 kwi 17:18 part-r-00006-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 0 11 kwi 17:18 part-r-00005-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 0 11 kwi 17:18 part-r-00004-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 27 11 kwi 17:18 part-r-00003-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 0 11 kwi 17:18 part-r-00002-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 0 11 kwi 17:18 part-r-00001-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 0 11 kwi 17:18 part-r-00000-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
-rw-r--r-- 1 jacek staff 0 11 kwi 17:18 _SUCCESS
➜ spark git:(master) ✗ cat your-json/part-r-00003-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
{"token":"hello","long":1}
➜ spark git:(master) ✗ cat your-json/part-r-00007-91f81f62-54bb-42ae-bddc-33829a0e3c16.json
{"token":"world","long":2}

Answer 2

UDFs will map one row to one row. If you have multiple rows in your DataFrame that you want to combine into one element, you're going to need to use a function like reduceByKey that aggregates multiple rows.
There may be a DataFrame specific function to do this, but I would do this processing with the RDD functionality, like so:
val colSummary = udf(
(col1:String, col2:String, col3:String, col4:String) => "{\"col1\":" + col1 + ",\"col2\":" + col2 + ",\"col3\":" + col3 + ",\"col4\":" + col4 + "}"
)
val colRDD = input.withColumn("Summary",summary('col1,'col2,'col3,'col4)).rdd.map(x => (x.getString(0),x.getString(5)))
This gives us a RDD[(String,String)], which will allow us to use the PairRDDFunctions like reduceByKey (see the docs). The key of the tuple is the original key, and the value is the json encoding for a single element that we need to aggregate together to make the cols list. We glue them all together into a comma-separated list, and then we add the beginning and end, and then we're done.
val result = colRDD.reduceByKey((x,y) => (x+","+y)).map(x => "["+x._1+",{\"cols\":["+x._2+"]}]")
result.take(10)