Как передать столбец df в качестве параметра в функцию?

Question

Я написал ниже функцию

object AgeClassification {
  def AgeCategory(age:Int) : String = {
    if(age<=30)
      return "Young" 
    else if(age>=65)
      return "Older" 
    else
      return "Mid-age"
  }
}

и я пытаюсь передать столбец данных в качестве параметра

val df_new = df
  .withColumn("Age_Category", AgeClassification.AgeCategory(df("age")))

но получаю ошибку

: 33: ошибка: несоответствие типов;
найдено: org.apache.spark.sql.Column
требуется: Int
val df_new = df.withColumn ("Age_Category", AgeClassification.AgeCategory (df ("age")))

Как передать столбец в качестве параметра?

val df_new = df
  .withColumn("Age_Category",AgeClassification.AgeCategory(df.age.cast(IntegerType)))   

: 33: ошибка: значение age не является членом org.apache.spark.sql.DataFrame
val df_new = df.withColumn ("Age_Category", AgeClassification.AgeCategory (df.age.cast (IntegerType))))

val df_new = df
   .withColumn("Age_Category", AgeClassification.AgeCategory(df("age").cast(Int)))

: 33: ошибка: приведено значение перегруженного метода с альтернативами:
(to: String) org.apache.spark.sql.Column
(to: org.apache.spark.sql.types.DataType) org.apache.spark.sql.Column
не может применяться к (Int.type)
val df_new = df.withColumn ("Age_Category", AgeClassification.AgeCategory (df ("age"). cast (Int)))

Answer 1

Вы не можете использовать функции scala напрямую при манипуляции с фреймом данных с помощью API SparkSQL.Вы можете использовать только функции «столбца», определенные в классе Column или в классе functions.Они в основном преобразуют столбцы в столбцы.Фактические вычисления обрабатываются в Spark.

Чтобы проиллюстрировать это, вы можете попробовать это в REPL:

scala> df("COL1").cast("int")
res6: org.apache.spark.sql.Column = CAST(COL1 AS INT)

Тип Column, а не int, и именно поэтомуscala отказывается применять вашу функцию (которая ожидает целое число) к такому объекту.

Чтобы использовать пользовательскую функцию, вам нужно обернуть ее в UDF следующим образом:

val ageUDF = udf((age : Int) => AgeClassification.AgeCategory(age))
// or shorter
val ageUDF = udf(AgeClassification.AgeCategory _)

// The you may use it this way:
df.withColumn("classif", ageUDF(df("age")))

Примечаниетакже, что df.age работает в pyspark, но это не является действительным ни scala.Для краткого доступа к столбцам по имени вы можете импортировать spark.implicits._ и писать $"age" или даже короче 'age.

Answer 2

import org.apache.spark.sql.Column

def AgeCategory(age:Column) : String