update value in specific row by checking condition for another row, pyspark

user_14800042 0
Viewed 844

So I tried a bunch of stuff but nothing seems to work. See the dataframe.

df

index Bool New_Bool
1 True True
2 True True
3 True True
4 False True

I want to update a column (New_bool). It should contain False if column Bool is false already AND it also should contain false in row with index=1 if Bool in row with index=4 is false as well. All other rows should be kept as are. I cannot manage to do this...see expected output down below.

index Bool New_Bool
1 True False
2 True True
3 False False
4 False False

I tried different combinations of .when statements but that did not work. Best solution so far:

df = df.withColumn('New_Bool', F.when((F.col('index')==4) &
                               ((F.col('Bool')==False), False))

But this produces:

index Bool New_Bool
1 True None
2 True None
3 True None
4 False False

If I were able to now fill all Nones with the value from row 4 in New_Bool guess that would help. I'd then set up another final column with:

df = df.withColumn('Final_Bool', F.when((F.col('index')==1) &
                               ((F.col('New_Bool')==False), False).otherwise(F.col('Bool'))

Not sure about anything anymore though because I tried so many combinations....by the way I will have many groups of those dataframes that's why I need a solution that works with groupby or a window function.

Happy for every advice :)

pyspark apache-spark-sql apache-spark
2 Answers

You can use lead to check the row with index = 4 when index = 1. For other indices, just keep the bool column.

from pyspark.sql import functions as F, Window

df2 = df.withColumn(
    'New_Bool',
    F.when(
        (F.col('index') == 1) & 
        (F.lead('Bool', 3).over(Window.orderBy('index')) == False),
        False
    ).otherwise(F.col('Bool'))
)

df2.show()
+-----+-----+--------+
|index| Bool|New_Bool|
+-----+-----+--------+
|    1| true|   false|
|    2| true|    true|
|    3|false|   false|
|    4|false|   false|
+-----+-----+--------+

Another way is to do a cross join to the bool of the row with index = 4:

df2 = df.crossJoin(
    df.filter('index = 4')
      .select(F.col('Bool').alias('id_4'))
).withColumn(
    'New_Bool',
    F.when(
        (F.col('index') == 1) & (F.col('id_4') == False), 
        False
    ).otherwise(F.col('Bool'))
).drop('id_4')

If 4 is always the last index, then you can use last:

df2 = df.withColumn(
    'New_Bool',
    F.when(
        (F.col('index') == 1) & 
        (F.last('Bool').over(Window.orderBy('index').rowsBetween(Window.unboundedPreceding, Window.unboundedFollowing)) == False),
        False
    ).otherwise(F.col('Bool'))
)
user_141657300

You can also select the last index's Bool value into a variable and use it in when expression like this:

from pyspark.sql import functions as F

# or actually, if the index is always = 4 you can just filter without ordering
last_bool = df.orderBy(F.desc("index")).limit(1).select("Bool").first().Bool

df2 = df.withColumn(
    'New_Bool',
    F.when(
        F.col('index') == 1, 
        F.lit(last_bool) & F.col("New_bool") # keep New_bool or update to false if last is false  
    ).otherwise(
        F.col("New_bool") & F.col("Bool") # keep New_bool or update to false if Bool is false 
    )
)

df2.show()
# +-----+-----+--------+
# |index| Bool|New_Bool|
# +-----+-----+--------+
# |    1| true|   false|
# |    2| true|    true|
# |    3|false|   false|
# |    4|false|   false|
# +-----+-----+--------+
user_13865510
Related