# функция для расчёта удержания 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