[Python] How to Exclude Rows Based on a Dynamic Condition in a PySpark Window Function?

I am working with PySpark and need to create a window function that calculates the median of the previous 5 values in a column. However, I want to exclude rows where a specific column feature is True. The challenge is that the value of the feature is being calculated on the go, similar to the median, so I can't filter it beforehand.

Specifically, if the median of the previous 10 values exceeds a certain threshold (let's say 20), the feature becomes True, and I want to exclude this row from future calculations of the median.

The window function does not support a groupBy statement and i can't find a way to dynamically calculate it over window without creating a new dataframe.

from pyspark.sql import SparkSession
from pyspark.sql.functions import col, expr
from pyspark.sql.window import Window
from pyspark.sql.functions import percentile_approx

spark = SparkSession.builder.appName("example").getOrCreate()

data = \[
(1, 10),
(2, 20),
(3, 30),
(4, 40),
(5, 50),
(6, 60),
(7, 70),
(8, 80),
(9, 90),
(10, 100),
(11, 110),
(12, 20)
\]

columns = \["id", "value"\]
df = spark.createDataFrame(data, columns)



window_spec = Window.orderBy("id").rowsBetween(-5, 0)

df = df.withColumn("median_value",
expr("percentile_approx(value, 0.5)").over(window_spec))
df = df.withColumn("feature", median_value > 35)


For example when calculating the median for row 7 the median is 40 which makes it true so on row 8 the median should be calculated from rows 1-6 and 8 skipping 7 because it was true.

Continue reading...