如何将所有列都是字符串的 DataFrame 转换为具有特定模式的 DataFrame

Question

想象以下输入：

val data = Seq (("1::Alice"), ("2::Bob"))
val dfInput = data.toDF("input")
val dfTwoColTypeString = dfInput.map(row => row.getString(0).split("::")).map{ case Array(id, name) => (id, name) }.toDF("id", "name")

现在我有一个包含所需列的 DataFrame：

scala> dfTwoColTypeString.show
+---+-----+
| id| name|
+---+-----+
|  1|Alice|
|  2|  Bob|
+---+-----+

我当然想要int类型的列id，但它是String类型的：

scala> dfTwoColTypeString.printSchema
root
 |-- id: string (nullable = true)
 |-- name: string (nullable = true)

因此我定义了这个架构：

val mySchema = StructType(Array(
    StructField("id", IntegerType, true),
    StructField("name", StringType, true)
    ))

将 DataFrame dfTwoColTypeString 转换或转换为给定目标架构的最佳方法是什么。

奖励：如果给定的输入无法转换或转换为目标模式，我希望得到一个空行，其中包含一个包含错误输入数据的额外列“bad_record”。也就是说，我想完成与 PERMISSIVE 模式下的 CSV 解析器相同的操作。

非常感谢任何帮助。

Answer 1

如果读取数据时需要转换，可以使用这样的代码：

val resultDF = mySchema.fields.foldLeft(dfTwoColTypeString)((df, c) => df.withColumn(c.name, col(c.name).cast(c.dataType)))
resultDF.printSchema()

输出：

root
 |-- id: integer (nullable = true)
 |-- name: string (nullable = true)

为了检查值匹配类型，可以使用这样的代码：

  val dfTwoColTypeString = dfInput.map(
  row =>
    row.getString(0).split("::"))
  .map {
        case Array(id, name) =>
          if (ConvertUtils.canBeCasted((id, name), mySchema))
            (id, name, null)
          else (null, null, id + "::" + name)}
  .toDF("id", "name", "malformed")

可以在自定义类（此处为 ConvertUtils）中创建两个新的静态函数：

def canBeCasted(values: Product, mySchema: StructType): Boolean = {
    mySchema.fields.zipWithIndex.forall(v => canBeCasted(values.productElement(v._2).asInstanceOf[String], v._1.dataType))
  }

import scala.util.control.Exception.allCatch

def canBeCasted(value: String, dtype: DataType): Boolean = dtype match {
    case StringType => true
    case IntegerType => (allCatch opt value.toInt).isDefined
    // TODO add other types here
    case _ => false
  }

输出错误的“cc::Bob”值：

+----+-----+---------+
|id  |name |malformed|
+----+-----+---------+
|1   |Alice|null     |
|null|null |cc::Bob  |
+----+-----+---------+

Answer 2

如果需要读取 CSV，并且架构已知，则可以在读取期间分配：

spark.read.schema(mySchema).csv("filename.csv")

Answer 3

val cols = Array(col("id").cast(IntegerType),col("name"))
dfTwoColTypeString.select(cols:_*).printSchema

根 |-- id: 整数（可为空=真） |-- 名称：字符串（可为空=真）

//另一种方法

import org.apache.spark.sql.types.{StringType,IntegerType,StructType,StructField}
val mySchema = StructType(Array(StructField("id", IntegerType, true),StructField("name", StringType, true)))
val df = spark.createDataFrame(dfTwoColTypeString.rdd,mySchema)
df.printSchema

根 |-- id: 整数（可为空=真） |-- 名称：字符串（可为空=真）

Answer 4

考虑到dfTwoColTypeString 是一个数据框，您还可以将其架构类型转换如下。

dfTwoColTypeString.withColumn("id", col("id").cast("Int"))