Spark SQL 嵌套 withColumn

Question

我有一个 DataFrame，它有多个列，其中一些是结构。像这样的

root
 |-- foo: struct (nullable = true)
 |    |-- bar: string (nullable = true)
 |    |-- baz: string (nullable = true)
 |-- abc: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- def: struct (nullable = true)
 |    |    |    |-- a: string (nullable = true)
 |    |    |    |-- b: integer (nullable = true)
 |    |    |    |-- c: string (nullable = true)

我想在baz 列上应用UserDefinedFunction 以将baz 替换为baz 的函数，但我不知道该怎么做。这是所需输出的示例（请注意，baz 现在是 int）

root
 |-- foo: struct (nullable = true)
 |    |-- bar: string (nullable = true)
 |    |-- baz: int (nullable = true)
 |-- abc: array (nullable = true)
 |    |-- element: struct (containsNull = true)
 |    |    |-- def: struct (nullable = true)
 |    |    |    |-- a: string (nullable = true)
 |    |    |    |-- b: integer (nullable = true)
 |    |    |    |-- c: string (nullable = true)

看起来DataFrame.withColumn 仅适用于顶级列，但不适用于嵌套列。我正在使用 Scala 来解决这个问题。

有人可以帮我解决这个问题吗？

谢谢

Answer 1

您可以使用 struct 函数来执行此操作，因为 Raphael Roth 已经在上面的答案中进行了演示。尽管使用Make Structs Easy* 库，但有一种更简单的方法可以做到这一点。该库向 Column 类添加了一个 withField 方法，允许您在 StructType 列中添加/替换列，这与 DataFrame 类中的 withColumn 方法允许您在 DataFrame 中添加/替换列的方式大致相同。对于您的特定用例，您可以执行以下操作：

import org.apache.spark.sql.functions._
import com.github.fqaiser94.mse.methods._

// generate some fake data
case class Foo(bar: String, baz: String)
case class Record(foo: Foo, arrayOfFoo: Seq[Foo])

val df = Seq(
   Record(Foo("Hello", "World"), Seq(Foo("Blue", "Red"), Foo("Green", "Yellow")))
).toDF

df.printSchema

// root
//  |-- foo: struct (nullable = true)
//  |    |-- bar: string (nullable = true)
//  |    |-- baz: string (nullable = true)
//  |-- arrayOfFoo: array (nullable = true)
//  |    |-- element: struct (containsNull = true)
//  |    |    |-- bar: string (nullable = true)
//  |    |    |-- baz: string (nullable = true)

df.show(false)

// +--------------+------------------------------+
// |foo           |arrayOfFoo                    |
// +--------------+------------------------------+
// |[Hello, World]|[[Blue, Red], [Green, Yellow]]|
// +--------------+------------------------------+

// example user defined function that capitalizes a given string
val myUdf = udf((s: String) => s.toUpperCase)

// capitalize value of foo.baz
df.withColumn("foo", $"foo".withField("baz", myUdf($"foo.baz"))).show(false)

// +--------------+------------------------------+
// |foo           |arrayOfFoo                    |
// +--------------+------------------------------+
// |[Hello, WORLD]|[[Blue, Red], [Green, Yellow]]|
// +--------------+------------------------------+

我注意到您有一个关于替换嵌套在数组中的结构中的列的后续问题。 这也可以通过将Make Structs Easy库提供的函数与spark-hofs库提供的函数结合起来实现，如下：

import za.co.absa.spark.hofs._

// capitalize the value of foo.baz in each element of arrayOfFoo
df.withColumn("arrayOfFoo", transform($"arrayOfFoo", foo => foo.withField("baz", myUdf(foo.getField("baz"))))).show(false)

// +--------------+------------------------------+
// |foo           |arrayOfFoo                    |
// +--------------+------------------------------+
// |[Hello, World]|[[Blue, RED], [Green, YELLOW]]|
// +--------------+------------------------------+

*完全披露：我是此答案中引用的Make Structs Easy 库的作者。

Answer 2

这很简单，只需使用一个点来选择嵌套结构，例如$"foo.baz"：

case class Foo(bar:String,baz:String)
case class Record(foo:Foo)

val df = Seq(
   Record(Foo("Hi","There"))
).toDF()


df.printSchema

root
 |-- foo: struct (nullable = true)
 |    |-- bar: string (nullable = true)
 |    |-- baz: string (nullable = true)


val myUDF = udf((s:String) => {
 // do something with s 
  s.toUpperCase
})


df
.withColumn("udfResult",myUDF($"foo.baz"))
.show

+----------+---------+
|       foo|udfResult|
+----------+---------+
|[Hi,There]|    THERE|
+----------+---------+

如果您想将 UDF 的结果添加到现有结构 foo，即获取：

root
 |-- foo: struct (nullable = false)
 |    |-- bar: string (nullable = true)
 |    |-- baz: string (nullable = true)
 |    |-- udfResult: string (nullable = true)

有两种选择：

withColumn:

df
.withColumn("udfResult",myUDF($"foo.baz"))
.withColumn("foo",struct($"foo.*",$"udfResult"))
.drop($"udfResult")

select:

df
.select(struct($"foo.*",myUDF($"foo.baz").as("udfResult")).as("foo"))

编辑： 用 UDF 的结果替换结构中的现有属性： 不幸的是，这不起作用：

df
.withColumn("foo.baz",myUDF($"foo.baz")) 

但可以这样做：

// get all columns except foo.baz
val structCols = df.select($"foo.*")
    .columns
    .filter(_!="baz")
    .map(name => col("foo."+name))

df.withColumn(
    "foo",
    struct((structCols:+myUDF($"foo.baz").as("baz")):_*)
)