How to add field within nested JSON when reading from/writing to Kafka via a Spark dataframe

Question

I've a Spark (v.3.0.1) job written in Java that reads Json from Kafka, does some transformation and then writes it back to Kafka. For now, the incoming message structure in Kafka is something like: {"catKey": 1}. The output from the Spark job that's written back to Kafka is something like: {"catKey":1,"catVal":"category-1"}. The code for processing input data from Kafka goes something as follows:

DataFrameReader dfr = putSrcProps(spark.read().format("kafka"));

for (String key : srcProps.stringPropertyNames()) {
  dfr = dfr.option(key, srcProps.getProperty(key));
}

Dataset<Row> df = dfr.option("group.id", getConsumerGroupId())
                     .load()
                     .selectExpr("CAST(value AS STRING) as value")
                     .withColumn("jsonData", from_json(col("value"), schemaHandler.getSchema()))
                     .select("jsonData.*");

// transform df

df.toJSON().write().format("kafka").option("key", "val").save()

I want to change the message structure in Kafka. Now, it should be of the format: {"metadata": <whatever>, "payload": {"catKey": 1}}. While reading, we need to read only the contents of the payload, so the dataframe remains similar. Also, while writing back to Kafka, first I need to wrap the msg in payload, add a metadata. The output will have to be of the format: {"metadata": <whatever>, "payload": {"catKey":1,"catVal":"category-1"}}. I've tried manipulating the contents of the selectExpr and from_json method, but no luck so far. Any pointer on how to achieve the functionality would be very much appreciated.

question from:https://stackoverflow.com/questions/65952271/how-to-add-field-within-nested-json-when-reading-from-writing-to-kafka-via-a-spa

Answer 1

To extract the content of payload in your JSON you can use get_json_object. And to create the new output you can use the built-in functions struct and to_json.

Given a Dataframe:

val df = Seq(("""{"metadata": "whatever", "payload": {"catKey": 1}}""")).toDF("value").as[String]

df.show(false)
+--------------------------------------------------+
|value                                             |
+--------------------------------------------------+
|{"metadata": "whatever", "payload": {"catKey": 1}}|
+--------------------------------------------------+

Then creating the new column called "value"

val df2 = df
  .withColumn("catVal", lit("category-1")) // whatever your logic is to fill this column
  .withColumn("payload",
      struct(
        get_json_object(col("value"), "$.payload.catKey").as("catKey"),
        col("catVal").as("catVal")
      )
  )
  .withColumn("metadata",
        get_json_object(col("value"), "$.metadata"),
  ).select("metadata", "payload")

df2.show(false)
+--------+---------------+
|metadata|payload        |
+--------+---------------+
|whatever|[1, category-1]|
+--------+---------------+

val df3 = df2.select(to_json(struct(col("metadata"), col("payload"))).as("value"))
df3.show(false)
+----------------------------------------------------------------------+
|value                                                                 |
+----------------------------------------------------------------------+
|{"metadata":"whatever","payload":{"catKey":"1","catVal":"category-1"}}|
+----------------------------------------------------------------------+