按类型触发选择列

时间:2019-02-20 07:53:37

标签: scala apache-spark generics

我想要一个功能来根据其数据类型动态选择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

如何在没有警告的情况下使其工作和编译?

1 个答案:

答案 0 :(得分:5)

想法是仅过滤具有所需类型的列,然后进行选择。

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)