scala - 如何向 DataFrame 类型添加新方法？

Question

假设我有一个在 Spark 数据帧上运行的 Scala 函数:

class MyClass {
      def makeColumnNull(df: DataFrame, columnToMakeNull: String): DataFrame = {
        val colType = df.select(columnToMakeNull).schema.head.dataType
        df.withColumn(columnToMakeNull, lit(null).cast(colType))
      }
}

我这样调用它:

val df = spark.range(0,10).toDF()
val df2 = MyClass.makeColumnNull(df, "id") 

虽然工作正常，但它的工作方式不如 Spark 的 API 流畅。我想要的是以一种使我能够做到这一点的方式重写我的函数:

val df2 = df.makeColumnNull("id")

有人能帮忙吗？

Answer 1

隐式类是可行的方法，我已经用它们扩展了几个 spark 类。所以你需要这个:

package com.mycompany.utils.spark

import org.apache.spark.sql.DataFrame
import org.apache.spark.sql.functions.lit

object DataFrameExtensions {

  implicit class DataFrameWrapper(df: DataFrame) {

    def makeColumnNull(columnToMakeNull: String): DataFrame = {
      val colType = df.select(columnToMakeNull).schema.head.dataType
      df.withColumn(columnToMakeNull, lit(null).cast(colType))
    }

  }
}

然后你必须导入com.mycompany.utils.spark.DataFrameExtensions._你将能够调用 makeColumnNull()针对任何DataFrame对象