Label each row of a dataframe with EventXHoursFromNow or EventWithinXHours

event_label.py

import pyspark.sql.functions as F
from pyspark.sql import Window
from pyspark.sql.column import Column


def overlaps(start_first, end_first, start_second, end_second):
    return (end_first >= start_second) & (end_second >= start_first)


def eventXHrFromNow(hours: int, time_col: str, start_col: str, end_col: str) -> Column:
    xhrs_from_now = F.col(time_col) + F.lit(hours)
    return F.when(
        (xhrs_from_now >= F.col(start_col)) & (xhrs_from_now <= F.col(end_col)),
        F.lit(1),
    ).otherwise(F.lit(0))


def eventWithinXHrs(hours: int, time_col: str, start_col: str, end_col: str) -> Column:
    return F.when(
        overlaps(
            F.col(time_col) + F.lit(1),
            F.col(time_col) + F.lit(1 + hours),
            F.col(start_col),
            F.col(end_col)
        ),
        F.lit(1),
    ).otherwise(F.lit(0))