Rdd 到 Dataframe，其中架构以编程方式基于标头答案

【问题标题】：Rdd to Dataframe where schema is programmatically based on headerRdd 到 Dataframe，其中架构以编程方式基于标头
【发布时间】：2020-03-14 11:32:09
【问题描述】：

我有以下 RDD，很多人都喜欢它：

val csv = sc.parallelize(Array(
  "col1, col2, col3",
  "1, cat, dog",
  "2, bird, bee"))

我想将 RDD 转换为数据帧，其中架构是基于 RDD 的第一行动态/以编程方式创建的。

我想将该逻辑应用于多个类似的 RDD，并且不能使用案例类以编程方式指定架构，也不能使用 spark-csv 从一开始就将数据作为数据框加载。

我创建了一个扁平化的数据框，但想知道在创建数据框时如何拆分各个列？

当前代码：

val header= file.first()
val data = file.mapPartitionsWithIndex {
  (idx, iter) => if (idx == 0) iter.drop(1) else iter 
}.toDF(header).show()

当前输出：

+----------------+
|col1, col2, col3|
+----------------+
|     1, cat, dog|
|    2, bird, bee|
+----------------+

【问题讨论】：

如果这是来自 csv 文件，则直接将其作为数据帧读取：stackoverflow.com/questions/29704333/…
我无法使用 spark-csv 将其直接读入数据帧
怎么回事？另请注意，在较新的 spark 版本中，无需使用 spark-csv，因为它可以直接使用而无需额外的包，例如：spark.read.format("csv").option("header", "true").load("csvfile.csv")。
是的，我就是这个意思。有一些限制，我不能这样做

标签： scala dataframe apache-spark apache-spark-sql rdd

【解决方案1】：

在大多数情况下，最好将 csv 文件直接作为数据帧读取，例如：Spark - load CSV file as DataFrame?。

首先您需要将数据拆分为数组，对于标头和 RDD 本身都是如此：

val header = csv.first().split(", ")

val data = full_csv.mapPartitionsWithIndex {
  (idx, iter) => if (idx == 0) iter.drop(1) else iter 
}.map(_.split(", ")).toDF("arr")

请注意，上面的代码会将 RDD 转换为数据帧，但是，它只有一个名为 arr 的列。另一方面，header 将是 Array[String]。

下一步是将具有单个数组列的数据框转换为正确的列数和正确的名称（基于header）：

data.select((0 until header.size).map(i => col("arr")(i).alias(header(i))): _*)

这将产生所需的输出数据帧：

+----+----+----+
|col1|col2|col3|
+----+----+----+
|col1|col2|col3|
|   1| cat| dog|
|   2|bird| bee|
+----+----+----+

【讨论】：

【解决方案2】：

您可以使用此代码：

import org.apache.spark.sql.types.{StringType, StructField, StructType}
import org.apache.spark.sql.{Row, SparkSession}

val csv = spark.sparkContext.parallelize(Array(
    "col1, col2, col3",
    "1, cat, dog",
    "2, bird, bee"))

val header = csv.first.split(",").map(_.trim)

val resultDF = spark.createDataFrame(
    csv.zipWithIndex
    .filter(_._2 > 0)
    .map{case (str, _) => Row.fromSeq(str.split(",").map(_.trim))}
  ,
    StructType(header.map(c => StructField(c, StringType)))
  )

resultDF.show(false)

输出：

+----+----+----+
|col1|col2|col3|
+----+----+----+
|1   |cat |dog |
|2   |bird|bee |
+----+----+----+

【讨论】：

请注意，您需要 zipWithIndex 来过滤掉标头，mapPartitionsWithIndex 将为每个分区分配一个索引，而不是每个记录。