Spark SQL вложен с помощью столбца

У меня есть 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)

Я хочу применить UserDefinedFunction в столбце baz, чтобы заменить 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 для этой проблемы.

Может кто-нибудь помочь мне с этим?

Спасибо

ОТВЕТЫ

Ответ 1

это просто, просто используйте точку для выбора вложенных структур, например. $"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"))

EDIT: Замена существующего атрибута в структуре результатом 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")):_*)
)