scala - spark 按类型选择列

Question

我想要一个函数来按数据类型动态选择 spark Dataframe 列。

到目前为止，我已经创建了:

object StructTypeHelpers {
  def selectColumnsByType[T <: DataType](schem: StructType):Seq[String] = {
    schem.filter(_.dataType.isInstanceOf[T]).map(_.name)
  }

}

这样一个 StructTypeHelpers. selectColumnsByType[StringType](df.schema)应该管用。但是，编译器警告我:

abstract type T is unchecked since it is eliminated by erasure

尝试使用时:

import scala.reflect.ClassTag
def selectColumnsByType[T <: DataType: ClassTag](schem: StructType):Seq[String]

它失败了

No ClassTag available for T

我怎样才能让它在没有警告的情况下工作和编译？

Answer 1

这个想法是只过滤具有您想要的类型的列，然后选择。

val df  =  Seq(
  (1, 2, "hello")
).toDF("id", "count", "name")

import org.apache.spark.sql.functions.col
def selectByType(colType: DataType, df: DataFrame) = {

  val cols = df.schema.toList
    .filter(x => x.dataType == colType)
    .map(c => col(c.name))
  df.select(cols:_*)

}
val res = selectByType(IntegerType, df)