Scala & Spark: Cast multiple columns at once
casting of all columns with idiomatic approach in scala
def castAllTypedColumnsTo(df: DataFrame, sourceType: DataType, targetType: DataType) = {
df.schema.filter(_.dataType == sourceType).foldLeft(df) {
case (acc, col) => acc.withColumn(col.name, df(col.name).cast(targetType))
}
}
Based on the comments (thanks!) I came up with the following code (no error handling implemented):
def castAllTypedColumnsTo(df: DataFrame,
sourceType: DataType, targetType: DataType) : DataFrame = {
val columnsToBeCasted = df.schema
.filter(s => s.dataType == sourceType)
//if(columnsToBeCasted.length > 0) {
// println(s"Found ${columnsToBeCasted.length} columns " +
// s"(${columnsToBeCasted.map(s => s.name).mkString(",")})" +
// s" - casting to ${targetType.typeName.capitalize}Type")
//}
columnsToBeCasted.foldLeft(df){(foldedDf, col) =>
castColumnTo(foldedDf, col.name, LongType)}
}
Thanks for the inspiring comments. foldLeft
(explained here and here) saves a for
loop to iterate over a var
dataframe.