Spark schema nullability modifier

DataFrameHelpers.scala

object DataFrameHelpers {
  import org.apache.spark.sql.types._

  implicit class NullableDataFrame(df: DataFrame) {
    def setNullableStateOfColumns( nullable: Boolean, containsNull: Boolean, cn: String*) : DataFrame = {
      // get schema
      val names = cn.toSet
      val schema = df.schema
      // modify [[StructField] with name `cn`
      val newSchema = schemaWithNullableStates(schema, nullable, containsNull, names)
      // apply new schema
      df.sqlContext.createDataFrame( df.rdd, newSchema )
    }
    
    private def schemaWithNullableStates(schema: StructType, nullable: Boolean, containsNull: Boolean, names: Set[String]): StructType = 
      StructType(schema.map {
        case StructField( name, dataType: ArrayType, _, metadata) if names.exists(_.startsWith(s"$name.")) && dataType.elementType.isInstanceOf[StructType] => 
          val innerColumns = names collect {
            case n if n.startsWith(s"$name.") => n.drop(name.length + 1)
          }
      
          val newArrayType = ArrayType(
            schemaWithNullableStates(dataType.elementType.asInstanceOf[StructType], nullable, containsNull, innerColumns), 
            dataType.containsNull
          )
        
          StructField( name, newArrayType, nullable = nullable, metadata)

        case StructField(name, dataType: StructType, nullability, metadata) if names.exists(_.startsWith(s"$name.")) =>
          val innerColumns = names collect {
            case n if n.startsWith(s"$name.") => n.drop(name.length + 1)
          }

          StructField( name, schemaWithNullableStates(dataType, nullable, containsNull, innerColumns), nullability, metadata)          

        case StructField( name, dataType: ArrayType, _, metadata) if names.contains(name) => 
          StructField( name, dataType.copy(containsNull = containsNull), nullable = nullable, metadata)
        
        case StructField( name, dataType, _, metadata) if names.contains(name) => 
          StructField( name, dataType, nullable = nullable, metadata)
        
        case y: StructField => y
      })
  }
}