Вы можете просто использовать select
вместе с map
, как показано в следующем примере:
import org.apache.spark.sql.functions._
import spark.implicits._
val df = Seq(
("a", 1.22, 2.333, 3.4444),
("b", 4.55, 5.666, 6.7777)
).toDF("id", "v1", "v2", "v3")
val roundCols = df.columns.filter(_.startsWith("v")) // Or filter with other conditions
val otherCols = df.columns diff roundCols
df.select(otherCols.map(col) ++ roundCols.map(c => round(col(c), 2).as(c)): _*).show
// +---+----+----+----+
// | id| v1| v2| v3|
// +---+----+----+----+
// | a|1.22|2.33|3.44|
// | b|4.55|5.67|6.78|
// +---+----+----+----+
Как сделать это методом:
import org.apache.spark.sql.DataFrame
def doubleToRound(df: DataFrame, roundCols: Array[String]): DataFrame = {
val otherCols = df.columns diff roundCols
df.select(otherCols.map(col) ++ roundCols.map(c => round(col(c), 2).as(c)): _*)
}
В качестве альтернативы используйте foldLeft
и withColumn
следующим образом:
def doubleToRound(df: DataFrame, roundCols: Array[String]): DataFrame =
roundCols.foldLeft(df)((acc, c) => acc.withColumn(c, round(col(c), 2)))