# функция для расчёта удержания

 

def get_retention(

    profiles,

    sessions,

    observation_date,

    horizon_days,

    dimensions=[],

    ignore_horizon=False,

):

 

    # добавляем столбец payer в передаваемый dimensions список

    dimensions = ['payer'] + dimensions

 

    # исключаем пользователей, не «доживших» до горизонта анализа

    last_suitable_acquisition_date = observation_date

    if not ignore_horizon:

        last_suitable_acquisition_date = observation_date - timedelta(

            days=horizon_days - 1

        )

    result_raw = profiles.query('dt <= @last_suitable_acquisition_date')

 

    # собираем «сырые» данные для расчёта удержания

    result_raw = result_raw.merge(

        sessions[['user_id', 'session_start']], on='user_id', how='left'

    )

    result_raw['lifetime'] = (

        result_raw['session_start'] - result_raw['first_ts']

    ).dt.days

 

    # функция для группировки таблицы по желаемым признакам

    def group_by_dimensions(df, dims, horizon_days):

        result = df.pivot_table(

            index=dims, columns='lifetime', values='user_id', aggfunc='nunique'

        )

        cohort_sizes = (

            df.groupby(dims)

            .agg({'user_id': 'nunique'})

            .rename(columns={'user_id': 'cohort_size'})

        )

        result = cohort_sizes.merge(result, on=dims, how='left').fillna(0)

        result = result.div(result['cohort_size'], axis=0)

        result = result[['cohort_size'] + list(range(horizon_days))]

        result['cohort_size'] = cohort_sizes

        return result

 

    # получаем таблицу удержания

    result_grouped = group_by_dimensions(result_raw, dimensions, horizon_days)

 

    # получаем таблицу динамики удержания

    result_in_time = group_by_dimensions(

        result_raw, dimensions + ['dt'], horizon_days

    )

 

    # возвращаем обе таблицы и сырые данные

    return result_raw, result_grouped, result_in_time