Splitting an array containing nested arrays using Scala in Azure Databricks - scala

Question

I'm currently working on a project where I'm having to extract some horribly nested data out of a json document (output from Log Analytics REST API call), document structure example below (I have a lot more columns):
{
"tables": [
{
"name": "PrimaryResult",
"columns": [
{
"name": "Category",
"type": "string"
},
{
"name": "count_",
"type": "long"
}
],
"rows": [
[
"Administrative",
20839
],
[
"Recommendation",
122
],
[
"Alert",
64
],
[
"ServiceHealth",
11
]
]
}
] }
I have managed to extract this json document into a data frame but I am stumped as to where to go from here.
The goal I am trying to achieve is an output like the below:
[{
"Category": "Administrative",
"count_": 20839
},
{
"Category": "Recommendation",
"count_": 122
},
{
"Category": "Alert",
"count_": 64
},
{
"Category": "ServiceHealth",
"count_": 11
}]
Ideally, I would like to use my columns array as the headers for each record. Then I want to split out each record array from the parent rows array in to its own record.
So far, I have tried flattening my raw imported data frame but this won't work as the rows data is an array of arrays.
How would I go about solving this conundrum?

Answer 1

It's a bit messy to deal with this, but here's a way to do it:
val df = spark.read.option("multiline",true).json("filepath")
val result = df.select(explode($"tables").as("tables"))
.select($"tables.columns".as("col"), explode($"tables.rows").as("row"))
.selectExpr("inline(arrays_zip(col, row))")
.groupBy()
.pivot($"col.name")
.agg(collect_list($"row"))
.selectExpr("inline(arrays_zip(Category, count_))")
result.show
+--------------+------+
| Category|count_|
+--------------+------+
|Administrative| 20839|
|Recommendation| 122|
| Alert| 64|
| ServiceHealth| 11|
+--------------+------+
To get the JSON output, you can do
val result_json = result.agg(to_json(collect_list(struct("Category", "count_"))).as("json"))
result_json.show(false)
+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|json |
+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
|[{"Category":"Administrative","count_":"20839"},{"Category":"Recommendation","count_":"122"},{"Category":"Alert","count_":"64"},{"Category":"ServiceHealth","count_":"11"}]|
+---------------------------------------------------------------------------------------------------------------------------------------------------------------------------+
Or you can save as JSON, e.g. with result.save.json("output").

Answer 2

Another way using transform function :
import org.apache.spark.sql.functions._
val df = spark.read.option("multiline",true).json(inPath)
val df1 = df.withColumn("tables", explode($"tables"))
.select($"tables.rows".as("rows"))
.select(expr("inline(transform(rows, x -> struct(x[0] as Category, x[1] as _count)))"))
df1.show
//+--------------+------+
//| Category|_count|
//+--------------+------+
//|Administrative| 20839|
//|Recommendation| 122|
//| Alert| 64|
//| ServiceHealth| 11|
//+--------------+------+
Then saving into json file:
df1.save.json(outPath)