Spark Dataframe 选择列使用案例

Question

我想实现以下内容
例如我有 Emp 文件（2 个文件） 我只想选择 2 列，例如 Empid 和 EmpName 如果文件没有 EmpName 它应该选择 Empid 数据框的一列

1) Emp1.csv（文件）

Empid   EmpName Dept
1       ABC     IS
2       XYZ     COE

2) Emp.csv（文件）

 Empid  EmpName
 1      ABC
 2      XYZ

代码尝试到现在

scala>  val SourceData = spark.read.format("com.databricks.spark.csv").option("inferSchema", "true").option("delimiter", ",").option("header", "true").load("/root/Empfiles/")
SourceData: org.apache.spark.sql.DataFrame = [Empid: string, EmpName: string ... 1 more field]

scala> SourceData.printSchema
root
|-- Empid: string (nullable = true)
|-- EmpName: string (nullable = true)
|-- Dept: string (nullable = true)

如果指定文件的所有列名，则此代码有效

 scala> var FormatedColumn = SourceData.select(
 |             SourceData.columns.map {
| case "Empid"                     => SourceData("Empid").cast(IntegerType).as("empid")
 | case "EmpName"                     => SourceData("EmpName").cast(StringType).as("empname")
 | case "Dept"                     => SourceData("Dept").cast(StringType).as("dept")
 | }: _*
 | )
 FormatedColumn: org.apache.spark.sql.DataFrame = [empid: int, empname: string ... 1 more field]

但我只想要特定的 2 列它会失败（如果列可用，它会显示选择并更改数据类型和列名）

 scala> var FormatedColumn = SourceData.select(
 | SourceData.columns.map {
 | case "Empid"                     => SourceData("Empid").cast(IntegerType).as("empid")
 | case "EmpName"                     => SourceData("EmpName").cast(StringType).as("empname")
 | }: _*
 | )
 scala.MatchError: Dept (of class java.lang.String)
 at $anonfun$1.apply(<console>:32)
 at $anonfun$1.apply(<console>:32)
 at scala.collection.TraversableLike$$anonfun$map$1.apply(TraversableLike.scala:234)
 at scala.collection.TraversableLike$$anonfun$map$1.apply(TraversableLike.scala:234)
 at scala.collection.IndexedSeqOptimized$class.foreach(IndexedSeqOptimized.scala:33)
 at scala.collection.mutable.ArrayOps$ofRef.foreach(ArrayOps.scala:186)
 at scala.collection.TraversableLike$class.map(TraversableLike.scala:234)
 at scala.collection.mutable.ArrayOps$ofRef.map(ArrayOps.scala:186)
  ... 53 elided

Answer 1

所有其他列也需要匹配：

var formattedColumn = sourceData.select(
  sourceData.columns.map {
      case "Empid"   => sourceData("Empid").cast(IntegerType).as("empid")
      case "EmpName" => sourceData("EmpName").cast(StringType).as("empname")
      case other: String => sourceData(other)
  }: _*
)

更新 1。如果只想选择“Empid”和“EmpName”两列，则无需使用匹配器：

val formattedColumn = sourceData.select(
  sourceData("Empid").cast(IntegerType).as("empid"),
  sourceData("EmpName").cast(StringType).as("empname")
)

更新 2。如果您想根据它们的存在来选择列，我可以建议以下内容：

val colEmpId = "Empid"
val colEmpName = "EmpName"
// list of possible expected column names
val selectableColums = Seq(colEmpId, colEmpName)
// take only the ones that are in the list
val foundColumns = sourceData.columns.filter(column => selectableColums.contains(column))
// create the target dataframe
val formattedColumn = sourceData.select(
  foundColumns.map(column =>
    column match {
      case colEmpId   => sourceData(colEmpId).cast(IntegerType).as("empid")
      case colEmpName => sourceData(colEmpName).cast(StringType).as("empname")
      case _ => throw new IllegalArgumentException("Unexpected column: " + column)
    }
  ): _*
)

附言请为vals 和vars 使用传统的驼峰命名法。

Answer 2

如果你用这个查询替换你的语句，它应该可以工作。 它会过滤掉所有不属于 match 子句的列。这样可以避免您看到的 MatchError。

df.select($"Empid", $"EmpName").select(df.columns.map {
    case "Empid" => df("Empid").cast(IntegerType).as("empid")
    case "EmpName" => df("EmpName").cast(StringType).as("empname")
}: _*)

Answer 3

我不知道为什么这么复杂..

为什么不这样做？

df
  .withColumn("empid", $"EmpId".cast(IntegerType))
  .withColumn("empname", $"EmpName".cast(StringType))