如何将列添加到 Spark 中的爆炸结构？

Question

假设我有以下数据：

{"id":1, "payload":[{"foo":1, "lol":2},{"foo":2, "lol":2}]}

我想分解有效载荷并为其添加一列，如下所示：

df = df.select('id', F.explode('payload').alias('data'))
df = df.withColumn('data.bar', F.col('data.foo') * 2)

但是这会导致数据框包含三列：

• id
• data
• data.bar

我希望 data.bar 成为 data 结构的一部分...

如何在展开的结构中添加列，而不是添加顶级列？

Answer 1

df = df.withColumn('data', f.struct(
    df['data']['foo'].alias('foo'),
   (df['data']['foo'] * 2).alias('bar')
))

这将导致：

root
 |-- id: long (nullable = true)
 |-- data: struct (nullable = false)
 |    |-- col1: long (nullable = true)
 |    |-- bar: long (nullable = true)

更新：

def func(x):
    tmp = x.asDict()
    tmp['foo'] = tmp.get('foo', 0) * 100
    res = zip(*tmp.items())
    return Row(*res[0])(*res[1])

df = df.withColumn('data', f.UserDefinedFunction(func, StructType(
    [StructField('foo', StringType()), StructField('lol', StringType())]))(df['data']))

附：

Spark 几乎不支持 inplace 操作。

所以每次你想inplace，你实际上需要replace。