用 spark 中的 group-by 其他列填充 NaN

Question

数据

 Col1 Col2    result
0  a    x      123.0 
1  a    y     NaN    
2  a    x      453.0 
3  a    y      675.0 
4  b    z      786.0 
5  b    z      332.0 

我想用 675.0 填充 NaN，首先按 col1 分组，然后按 col2 并填充 NaN 值

在熊猫中

df['result'] = df['result'].fillna(df.groupby(['col1', 'col2', ])['result'].bfill())


df['result'] = df['result'].fillna(df.groupby(['col1', 'col2', ])['result'].ffill())

如何在 pyspark 中实现它？

Answer 1

这可以通过pandas udfs 完成。然后就可以直接使用你想要的功能了。

[IN]
from pyspark.sql.functions import pandas_udf, PandasUDFType
import pandas as pd, numpy as np

s = pd.DataFrame({"col1": ["a","a","a","a","b","b"],
              "col2": ["x", "y", "x", "y", "z", "z"],
              "result": [123, np.nan, 453, 675, 786, 332]})

spark_df = spark.createDataFrame(s)

grouped_spark_df = spark_df.groupBy("col1", "col2")

@pandas_udf("col1 string, col2 string, result float", PandasUDFType.GROUPED_MAP)
def fillnaspark(df):
    df['result'] = df['result'].bfill()
    df['result'] = df['result'].ffill()
    return df

grouped_spark_df.apply(fillnaspark).show()

[OUT]
+----+----+------+
|col1|col2|result|
+----+----+------+
|   a|   x| 123.0|
|   a|   x| 453.0|
|   b|   z| 786.0|
|   b|   z| 332.0|
|   a|   y| 675.0|
|   a|   y| 675.0|
+----+----+------+

Answer 2

您可以使用nanvl 将NaN 替换为结果的lead 值（等效于ffill；bfill 将是lag）：

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'id', F.monotonically_increasing_id()
).withColumn(
    'result', 
    F.nanvl(
        'result',
        F.coalesce(
            F.lag('result').over(Window.partitionBy('Col1', 'Col2').orderBy('id')),
            F.lead('result').over(Window.partitionBy('Col1', 'Col2').orderBy('id'))
        )
    )
).orderBy('id')

df2.show()
+---+----+----+------+
| id|Col1|Col2|result|
+---+----+----+------+
|  0|   a|   x| 123.0|
|  1|   a|   y| 675.0|
|  2|   a|   x| 453.0|
|  3|   a|   y| 675.0|
|  4|   b|   x| 786.0|
|  5|   b|   y| 332.0|
+---+----+----+------+