python - 如果它们都为空，如何在 PySpark 中一次替换多个列中的值？

Question

给定数据集

富	酒吧
1	2
无效的	无效的
3	4

foo bar如果其中一个或两个都是空值，如何立即替换为其他值，例如 (5, 6)？

富	酒吧
1	2
5	6
3	4

这适用于地理数据集，当 lat/lng 未知并且应该在其他地方获得时。所以 udf 很耗时，我想确保它只为必要的行调用（其中 foo 和 bar 都为空）

以下代码

from pyspark.sql import SparkSession
from pyspark.sql.functions import when, udf

spark = (SparkSession.builder.master("local")
         .appName("SimpleApp")
         .getOrCreate()
         )


def my_udf(): return 0, 0


df = spark.createDataFrame([[1, 2], [None, None], [3, 4]], schema=['foo', 'bar'])
df = df.withColumn("result", when(df['foo'].isNull() | df['bar'].isNull(), udf(my_udf)()))

df.show()

不好尝试

富	酒吧	结果
1	2	无效的
无效的	无效的	[Ljava.lang.Object...
3	4	无效的

所以有必要以某种方式将数组解包到列中。

考虑到这一点，Apache Spark 无法一步完成 ——将 UDF 的结果分配给多个数据帧列

但即使我将返回 struct 并将其解包，如何单独留下不受影响的列？

我尝试过的另一种方法（考虑到我需要进一步处理foo bar）：

def some(baz): return 'some'
def something_else(foo, bar): return 'something'
def my_udf(_):
    foo, bar, baz = _
    return some(baz) if foo is None and bar is None else something_else(foo, bar)


df = spark.createDataFrame([[1, 2, 3], [None, None, 4], [3, 4, 5]], schema=['foo', 'bar', 'baz'])
df = df.withColumn("result", udf(my_udf)(array('foo', 'bar', 'baz')))

df.show()

但我觉得它不是那么优化，即使我们不需要baz大多数行我们仍然将它传递给 udf，我认为它会阻止优化请求。

当然，我可以为不同的列一一应用不同的 udf，但它似乎也不是那么理想。

那么有没有办法一次替换两列中的值？

标签： pythonpyspark

Answer 1

我会坚持你的第一个想法，然后使用coalesce从 udf 的结果中填充空行：

from pyspark.sql import functions as F
from pyspark.sql import types as T

@F.udf(returnType=T.StructType([T.StructField("foo",T.DoubleType(), True), 
                  T.StructField("bar",T.DoubleType(), True)]))
def my_udf(foo, bar): 
    return (0.0, 1.0)


df = spark.createDataFrame([[1, 2], [None, None], [3, 4]], schema=['foo', 'bar'])
df.withColumn("result", F.when(df['foo'].isNull() | df['bar'].isNull(), 
                                                    my_udf("foo", "bar"))) \
    .withColumn("foo", F.coalesce("foo", "result.foo")) \
    .withColumn("bar", F.coalesce("bar", "result.bar")) \
    .show()

输出：

+---+---+----------+                                                            
|foo|bar|    result|
+---+---+----------+
|1.0|2.0|      null|
|0.0|1.0|{0.0, 1.0}|
|3.0|4.0|      null|
+---+---+----------+

使用coalesce不会是性能问题，因为此功能是一个狭窄的转换，因此不会导致洗牌。