chat1

import polars as pl
from datetime import timedelta

def propagate_indicator(df):
    # Ensure the dataframe is sorted by timestamp
    df = df.sort('timestamp')

    # Create a helper column to track where indicator is 1
    df = df.with_column(
        (pl.col('indicator') == 1).alias('start_indicator')
    )

    # Calculate the timestamp up to which the value 1 should be propagated
    df = df.with_column(
        pl.when(pl.col('start_indicator'))
        .then(pl.col('timestamp') + timedelta(seconds=30))
        .otherwise(None)
        .alias('propagate_until')
    )

    # Forward fill the 'propagate_until' to get the maximum propagate time
    df = df.with_column(
        pl.col('propagate_until').fill_null('forward').alias('propagation_time')
    )

    # Use a window function to propagate the indicator
    df = df.with_columns(
        pl.when(pl.col('timestamp') <= pl.col('propagation_time'))
        .then(1)
        .otherwise(pl.col('indicator'))
        .alias('indicator_propagated')
    )

    # Drop helper columns
    df = df.drop(['start_indicator', 'propagate_until', 'propagation_time'])

    return df

# Example DataFrame
data = {
    'timestamp': [pl.datetime("2023-01-01 00:00:00") + timedelta(seconds=i) for i in range(120)],
    'indicator': [1 if i in [0, 50, 70, 100] else None for i in range(120)]
}
df = pl.DataFrame(data)

# Apply the function
result_df = propagate_indicator(df)
print(result_df)