Conditional filter threshold on pandas dataframe column based on another column value

Question

Let's say I have a dataframe with two columns, and I would like to filter the values of the second column based on different thresholds that are determined by the values of the first column. Such thresholds are defined in a dictionary, whose keys are the first column values, and the dict values are the thresholds. There will be also a default value to match columns that do not have any of the specified values.

So for example:

thresholds_dict = {"A": 5, "B": 2, "C": 4, "default": 0}

sample_dataframe = 
| Column1 | Column2 |
|   A     | 3       |
|   A     | 6       |
|   B     | 4       |
|   B     | 1       |
|   C     | 2       |
|   D     | 0       |

//Get threshold from dict based on value of Column1 on ...
result_dataframe = sample_dataframe[sample_dataframe[Column2] >= ...] 

result_dataframe =
| Column1 | Column2 |
|   A     | 6       |
|   B     | 4       |
|   D     | 0       |

What would be the best way to achieve this? (Not sure what to write in... part).

Answer 1

PySpark version.

Your dataframe:

from pyspark.sql import functions as F

sample_dataframe = spark.createDataFrame(
    [("A", 3),
     ("A", 6),
     ("B", 4),
     ("B", 1),
     ("C", 2),
     ("D", 0)],
    ["Column1", "Column2"]
)
thresholds_dict = {"A": 5, "B": 2, "C": 4, "default": 0}

Script:

comparison = F.when(F.lit(False), None)
for k, v in thresholds_dict.items():
    comparison = comparison.when(F.col("Column1") == k, v)
comparison = comparison.otherwise(thresholds_dict["default"])

result_dataframe = sample_dataframe.filter(F.col("Column2") >= comparison)

result_dataframe.show()
# +-------+-------+
# |Column1|Column2|
# +-------+-------+
# |      A|      6|
# |      B|      4|
# |      D|      0|
# +-------+-------+