Applying schema at message level instead of dataframe level in Spark Structured Streaming - scala

Question

So, the thing is since my schema could depend on kafka header/key, I want to apply schema at message level rather then dataframe level. How to achieve this? Thanks
The code snippet to apply schema for dataframe level is:
val ParsedDataFrame = kafkaStreamData.selectExpr("CAST(value AS STRING)", "CAST(key AS STRING)")
.select(from_json(col("value"), Schema), col("key"))
.select("value.*","key")
I want something like,
if(key == 'a'){
use Schema1
}
else{
use Schema2
}
P.S: I tried using foreach and map function but none worked, maybe not using them correctly

Answer 1

It is not possible to apply a different schema in the same row as you would end up getting an AnalysisException due to data type mismatch.
To test this you can do the following experiment.
Have the following data in Kafka topic in the forma key:::value:
a:::{"a":"foo","b":"bar"}
b:::{"a":"foo","b":"bar"}
In your streaming query you define:
val schemaA = new StructType().add("a", StringType)
val schemaB = new StructType().add("b", StringType)
val df = spark.readStream
.format("kafka")
.[...]
.selectExpr("CAST(key AS STRING)", "CAST(value AS STRING)")
.withColumn("parsedJson",
when(col("key") === "a", from_json(col("value"), schemaA))
.otherwise(from_json(col("value"), schemaB)))
This will result in an AnalysisException:
org.apache.spark.sql.AnalysisException: cannot resolve 'CASE WHEN (`key` = 'a') THEN jsontostructs(`value`) ELSE jsontostructs(`value`) END' due to data type mismatch: THEN and ELSE expressions should all be same type or coercible to a common type;;
'Project [key#21, value#22, CASE WHEN (key#21 = a) THEN jsontostructs(StructField(a,StringType,true), value#22) ELSE jsontostructs(StructField(b,StringType,true), value#22) END AS parsedJson#27]
As #OneCricketeer mentioned in the comments, you need to seperate the kafka input stream into several Dataframes based on a filter first and then apply different schemas for parsing a column with a json string.