python - pySpark 替换行子集上的空值

Question

我有一个 pySpark 数据框，其中有我想要替换的空值 - 但是要替换的值对于不同的组是不同的。

我的数据看起来像这样(抱歉，我没有办法将它作为文本传递):

对于 A 组，我想用 -999 替换空值；而对于 B 组，我想用 0 替换空值。

目前，我将数据分成几个部分，然后执行 df = df.fillna(-999) 。

有没有更有效的方法呢？在伪代码中，我在考虑 df = df.where(col('group') == A).fillna(lit(-999)).where(col('group') = = B).fillna(lit(0)) 但当然，这是行不通的。

Answer 1

您可以使用 when :

from pyspark.sql import functions as F

# Loop over all the columns you want to fill
for col in ('Col1', 'Col2', 'Col3'):
    # compute here conditions to fill using a value or another
    fill_a = F.col(col).isNull() & (F.col('Group') == 'A')
    fill_b = F.col(col).isNull() & (F.col('Group') == 'B')

    # Fill the column based on the different conditions 
    # using nested `when` - `otherwise`.
    #
    # Do not forget to add the last `otherwise` with the original 
    # values if none of the previous conditions have been met
    filled_col = (
        F.when(fill_a, -999)
        .otherwise(
            F.when(fill_b, 0)
            .otherwise(F.col(col))
        )
    )

    # 'overwrite' the original column with the filled column
    df = df.withColumn(col, filled_col)