Dropping nested struct elements from DataFrame when comparing to another schema - scala

Question

On Spark version 2.4.5:
I have DataFrames that can have arbitrary amounts of nesting:
[scalatest] root
[scalatest] |-- nestedStruct: struct (nullable = true)
[scalatest] | |-- a: string (nullable = true)
[scalatest] | |-- b: integer (nullable = true)
and I have another StructType schema that will always be a subset of the DataFrame's schema, e.g.
[scalatest] root
[scalatest] |-- nestedStruct: struct (nullable = true)
[scalatest] | |-- a: string (nullable = true)
how would I go about dropping column b or any level of nested columns in the initial DataFrame when comparing it to the subset schema?
I've done something like this so far:
def dropColumnsNotInSubsetSchema(df: DataFrame, subsetSchema: StructType): DataFrame = {
val missingColumns = new ArrayBuffer[String]()
findMissingColumns(df.schema, subsetSchema, missingColumns)
df.drop(missingColumns.toArray: _*)
}
private def findMissingColumns(nestedSchema: StructType,
subsetSchema: StructType,
missingColumns: ArrayBuffer[String]): Unit = {
nestedSchema.fields.foreach(field => {
field.dataType match {
case structType: StructType =>
findMissingColumns(structType, subsetSchema, missingColumns)
case _ =>
if (!subsetSchema.fields.contains(field)) {
missingColumns += field.name
}
}
})
}
but it seems subsetSchema also has to be recursively gone through since nested columns won't exist in the outer most layer of the subset schema. And even if I accomplish that and identify that b is not in the subset schema, the column name identified would just be b and doing df.drop("b") wouldn't work..