Spark Scala：在同一行的数据框中创建和使用列

Question

When I need a new column in a dataframe to then use in a different computation, my code looks similar to:

var df: DataFrame = ...
df = df.withColumn("new_col", df.col("a") / 2)
println(df.withColumn("res", df.col("b") + df.col("new_col")).head())

How to combine into a single line (and avoid using var )?

The problem is df.col() as I cannot simply do the following because new_col does not exist in df yet:

df.withColumn("new_col", df.col("a"))
  .withColumn("res", df.col("b") + df.col("new_col"))
  .head()

Is there some API I am missing?

Answer 1

You can use $ to make a column instead of df.col ; The former will infer the column from the new data frame instead of df :

df.withColumn("new_col", $"a")
  .withColumn("res", $"b" + $"new_col")
  .head()

Or:

import org.apache.spark.sql.functions.col
df.withColumn("new_col", col("a"))
  .withColumn("res", col("b") + col("new_col"))
  .head()