Создайте новый столбец в DataFrame Spark, используя переменную с постоянным значением

Question

Я пытаюсь определить новый столбец в кадре данных Spark, используя константу, определенную как var.Я использую Zeppelin - в начальной ячейке он начинается с

%spark
import org.apache.spark.sql.functions._
var year : Int = 2016
spark.read.parquet("<path/to/file>")

Файл содержит столбец с именем birth_year;Я хочу создать новый столбец с именем age, определенный как $year - birth_year, где birth_year - это строковый столбец.Я не совсем понимаю, как это сделать, когда входной аргумент UDF является параметром.Я провел пару часов поиска и создал UDF, но получил сообщение об ошибке, основная часть которого

<console>:71: error: type mismatch;
 found   : Int
 required: org.apache.spark.sql.Column
       spark.read.parquet("path/to/file").withColumn("birth_year", $"birth_year" cast "Int").withColumn("age", createAge(year, col("birth_year"))).createOrReplaceTempView("tmp")

, и каретку прямо под 'year'.

Я подозреваю, что $year не отображается в переменную той же длины, что и birth_year;Я видел функцию lit(), которая работает со строками - работает ли она также с целочисленными значениями, или для этой цели есть другая функция?

Я попробовал следующее:

%spark
import org.apache.spark.sql.functions._
var year : Int = 2016
def createAge = udf((yr : Int, dob : Int) => {yr - dob})
spark.read.parquet("<path/to/file>").withColumn("birth_year", $"birth_year" cast "Int").withColumn("age", createAge($"year", col("birth_year"))).createOrReplaceTempView("tmp")

Любые предложения приветствуются - заранее спасибо за любую помощь.

Answer 1

Вы не можете использовать year напрямую как вход для UDF, поскольку он ожидает, что столбцы будут работать.Для создания столбца с постоянным значением используйте lit().Вы можете вызвать UDF следующим образом:

df.withColumn("age", createAge(lit(year), $"birth_year".cast("int")))

---

Однако, когда это возможно, всегда предпочтительнее использовать встроенные функции в Spark, когда это возможно.В этом случае вам не нужно UDF.Просто сделайте:

df.withColumn("age", lit(year) - $"birth_year".cast("int"))

Это должно быть намного быстрее.