Spark Scala: создать и использовать столбец в кадре данных в одной строке

Question

Когда мне нужен новый столбец в кадре данных для последующего использования в другом вычислении, мой код выглядит примерно так:

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

Как объединить в одну строку (и избегать использования var)?

Проблема в df.col(), поскольку я не могу просто сделать следующее, потому что new_col еще не существует в df:

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

Есть ли какой-то API, которого мне не хватает?

Answer 1

Вы можете использовать $ для создания столбца вместо df.col;Первый из них выведет столбец из нового фрейма данных вместо df:

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

или:

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