时间戳

首页 > 解决方案 > 如何获取 Dataframe 上的所有数组名称

scala - 如何获取 Dataframe 上的所有数组名称

问题描述

如何获取 Dataframe 上的所有数组名称？

问题是我试图爆炸所有阵列。

import org.apache.spark.sql.{Column, DataFrame, SparkSession}
import org.apache.spark.sql.functions.col
import org.apache.spark.sql.types.{ArrayType, StructField, StructType}

val providersDF=SIWINSDF.select(explode(col("**providers**")).as("collection")).select(col("collection.*"))
 def flattenSchema(schema: StructType, prefix: String = null) : Array[Column] = {
    schema.fields.flatMap(f => {
      val colName = if (prefix == null) f.name else (prefix + "." + f.name)

      f.dataType match {
        case st: StructType => flattenSchema(st, colName)
        case _ => Array(col(colName).alias(colName))
      }
    })
  }

  val newDF=providersDF.select(flattenSchema(providersDF.schema):_*)
  newDF.toDF(newDF.columns.map(_.replace(".", "_")): _*).printSchema

要获取我想做的所有数组名称：

获取我正在做的名称 df.schema.filter(st => st.dataType.isInstanceOf[ArrayType]).flatMap(.dataType.asInstanceOf[StructType].fields).map(.name)

任何帮助表示赞赏。

标签： scalaapache-sparkapache-spark-sql

解决方案

ArrayType这是一个从 DataFrame中提取所有嵌套列的递归方法：

import org.apache.spark.sql.types._

def extractArrayCols(schema: StructType, prefix: String): Seq[String] =
  schema.fields.flatMap {
    case StructField(name, struct: StructType, _, _) => extractArrayCols(struct, prefix + name + ".")
    case StructField(name, ArrayType(_, _), _, _) => Seq(s"$prefix$name")
    case _ => Seq.empty[String]
  }

测试方法：

import org.apache.spark.sql.functions._

case class W(u: Int, v: Seq[String])

val df = Seq(
  (10, Seq(1, 2), W(1, Seq("a", "b"))),
  (20, Seq(3), W(2, Seq("c", "d")))
).toDF("c1", "c2", "c3")

val arrayCols = extractArrayCols(df.schema, "")
// arrayCols: Seq[String] = ArraySeq(c2, c3.v)

推荐阅读