时间戳

首页 > 解决方案 > 用 spark 中的 group-by 其他列填充 NaN

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

问题描述

数据

 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 中实现它？

标签： pythonapache-sparkpysparkapache-spark-sql

解决方案

这可以通过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|
+----+----+------+

推荐阅读