How to transform json column in databricks with scala

user_8958704 0
Viewed 43

I have two similar columns of json in a dataframe:

Column A's values look like this:

{"categoryCode": "ZZ","yesNoCode": "Y","conditionIndicator1": "ST"}

Column B's values look like this: {"categoryCode": "ZZ","yesNoCode": "Y","conditionIndicator1": "ST", "conditionIndicator2": "RN"}

I have other columns as well, whose values are similar, but have 3 condition indicators, 4 condition indicators or 5 condition indicators. Each column's values always have the same number of condition indicators.

I'd like to do something easy like val resultDf = df.select(array($"col1", $"col2", ...)) but this fails because the array function won't allow you to select columns of different types. So I need to call some function f to do a transformation on each column's json, so that all of them have five condition indicators.

val resultDf = df.select(array(f($"col1"), f($"col2"), ...)).alias("normalizedData")

This would yield:

[
  {"categoryCode": "ZZ", "yesNoCode": "Y", "conditionIndicator1": "ST", "conditionIndicator2": "", "conditionIndicator3": "", "conditionIndicator4": "", "conditionIndicator5": ""},
  {"categoryCode": "KF", "yesNoCode": "Y", "conditionIndicator1": "ST", "conditionIndicator2": "RN", "conditionIndicator3": "", "conditionIndicator4": "", "conditionIndicator5": ""},
  {"categoryCode": "64", "yesNoCode": "N", "conditionIndicator1": "2M", "conditionIndicator2": "7X", "conditionIndicator3": "34", "conditionIndicator4": "22", "conditionIndicator5": "AE"}
]

So far what I am using is a udf, but it seems like this is the Really Hard Way, worse than string manipulation!

val createCondition = udf {(theCondition: Row) => 
  if (theCondition != null) {
    var m = Map(
      "categoryCode" -> theCondition.getAs[String]("categoryCode"),
      "yesNoCode" -> theCondition.getAs[String]("yesNoCode"),
      "conditionIndicator1" -> theCondition.getAs[String]("conditionIndicator1")
    )
    var c2 = theCondition.getAs[String]("conditionIndicator2")
    if (c2 == null) {
      c2 = ""
    }
    m = m + ("conditionIndicator2", c2)  
    var c3 = theCondition.getAs[String]("conditionIndicator3")
    if (c3 == null) {
      c3 = ""
    }
    m = m + ("conditionIndicator3", c3)  
    // et cetera...and way worse if we had a more complex object graph...which I do!
    // then return my map, m.
    m
  } else {
    null
  }
}

val result = df.select(array(createCondition($"col1"), createCondition($"col2")))

Seems like transformation of complex data is why we have databricks, so I'm certain I am doing this wrong. Should be a way have a function maps objects of type x into type y.

apache-spark-sql apache-spark scala json
1 Answers

Let's write a function that generates a schema that matches your data and that can have a variable number of conditions:

import org.apache.spark.sql.types._
def gen_schema(numberOfConditions : Int) = {
    val root = Array(StructField("categoryCode", StringType), StructField("yesNoCode", StringType))
    val conditions = (1 to numberOfConditions)
        .map(i => StructField(s"conditionIndicator$i", StringType))
    StructType(root ++ conditions)
}

// Let's define sample data and try the function:
val df = Seq(("""{"categoryCode": "ZZ","yesNoCode": "Y","conditionIndicator1": "ST"}""",
              """{"categoryCode": "ZZ","yesNoCode": "Y","conditionIndicator1": "ST", "conditionIndicator2": "RN"}"""))
     .toDF("col1", "col2")
df
    .withColumn("col1", from_json('col1, gen_schema(1)))
    .withColumn("col2", from_json('col2, gen_schema(2)))
    .show(false)

+-----------+---------------+
|col1       |col2           |
+-----------+---------------+
|[ZZ, Y, ST]|[ZZ, Y, ST, RN]|
+-----------+---------------+

By doing this, col1 and col2 do not have the same type as you mention it. But if you use gen_schema(5) for all columns, empty columns will be generated if the json entry does not exists. You will then be able to put them all inside an array:

val result = df.select(array(from_json('col1, gen_schema(5)),
                             from_json('col2, gen_schema(5))))
result.show(false)
result.printSchema

+-----------------------------------------------+
|array(jsontostructs(col1), jsontostructs(col2))|
+-----------------------------------------------+
|[[ZZ, Y, ST,,,,], [ZZ, Y, ST, RN,,,]]          |
+-----------------------------------------------+
root
 |-- array(jsontostructs(col1), jsontostructs(col2)): array (nullable = false)
 |    |-- element: struct (containsNull = true)
 |    |    |-- categoryCode: string (nullable = true)
 |    |    |-- yesNoCode: string (nullable = true)
 |    |    |-- conditionIndicator1: string (nullable = true)
 |    |    |-- conditionIndicator2: string (nullable = true)
 |    |    |-- conditionIndicator3: string (nullable = true)
 |    |    |-- conditionIndicator4: string (nullable = true)
 |    |    |-- conditionIndicator5: string (nullable = true)

result

user_88936860
Related