scala - 如何在字符串类型的嵌套结构中获取spark中的所有列名称

Question

所以我有这个模式，我想获取所有包含字符串的列名

  val schema = StructType(Array(
    StructField("id", IntegerType),
    StructField("nested", StructType(Array(
      StructField("value1", StringType),
      StructField("value2", StringType)
    ))),
    StructField("name", StringType)
    ))

我想得到

  Seq("nested.value1","nested.value2","name")

这只是一个示例，但它应该可以工作并且可以嵌套多个级别

Answer 1

def extractNames(schema: StructType): Seq[String] = {
  schema.fields.flatMap {
    field =>
      field.dataType match {
        case structType: StructType =>
          extractNames(structType).map(field.name + "." + _)
        case _: StringType =>
          field.name :: Nil
        case _ =>
          Nil
      }
  }
}

测试用例:

val schema = StructType(Array(
  StructField("id", IntegerType),
  StructField("nested", StructType(Array(
    StructField("value1", StringType),
    StructField("value2", StringType),
    StructField("struct", StructType(Array(
      StructField("value1", IntegerType),
      StructField("value2", StringType)
    )))
  ))),
  StructField("name", StringType)
))

val names = extractNames(schema)
println(names.mkString(", "))

输出:

nested.value1, nested.value2, nested.struct.value2, name