Создать новый столбец на основе данных в существующем столбце

У меня есть трехсторонняя иерархия: свойство -> prov -> co. Каждое свойство имеет сегмент, то есть отель / дом. Я написал запрос, чтобы получить количество каждого ниже:

properties = spark.sql("""
    SELECT
        COUNT(ps.property_id) as property_count,
        ps.prov_id,
        c.id as co_id,
        ps.segment
    FROM
        schema.t1 ps
    INNER JOIN
        schema.t2 c
        ON c.id = p.co_id
    GROUP BY
        2,3,4
""")
properties = properties.toPandas()

Это дает мне общее количество свойств на сегмент, на пров, на кооперацию. Из приведенных выше properties df я хочу создать новый df, который выглядит следующим образом:

- prov_id,
- prov_segment,
- co_id,
- co_segment

prov_segment должен быть 'Home', если> 50% свойств в этом pro_id попадают в сегмент Home , в противном случае это должен быть Core . Аналогично, co_segment должен быть Home если> 50% prov_id попадают в Home prov_segment, в противном случае он должен быть основным.

Я знаю, я могу получить общее количество свойств, сгруппировав данные:

prop_total_count = properties.groupby('prov_name')['property_count'].sum()

Тем не менее, я не уверен, как использовать это для создания нового кадра данных.

Пример данных:

properties.show(6) :

| property_count | prov_id | co_id | segment |
|----------------|---------|-------|---------|
| 10             | 1       | ABC   | Core    |
| 200            | 1       | ABC   | Home    |
| 300            | 9       | ABC   | Core    |
| 10             | 9       | ABC   | Home    |
| 100            | 131     | MNM   | Home    |
| 200            | 199     | KJK   | Home    |

Исходя из вышеизложенного, я бы хотел следующий вывод:

| prov_id | prov_segment | co_id | co_segment |
|---------|--------------|-------|------------|
| 1       | Home         | ABC   | Core       |
| 9       | Core         | ABC   | Core       |
| 131     | Home         | MNM   | Home       |
| 199     | Home         | KJK   | Home       |

prov_id 1 получает сегмент Home, поскольку он имеет 200 объектов недвижимости по сравнению с 10 основными объектами. prov_id 9 получает основной сегмент, поскольку он имеет 300 основных свойств для 10 свойств Home.

co_id ABC получает основной сегмент из-за портфеля, имеющего в общей сложности 310 основных свойств по сравнению с 210 домашними объектами.

prov_id 131 и 199 находятся только в одном сегменте, так что этот сегмент остается.

Всего 1 ответ


Хорошо, возможно, можно решить эту проблему «более коротким» способом, но это должно сработать. Он основан на создании двух других DataFrames с сегментами на группу ( co_id или prov_id ) и последующем объединении DataFrames.

Объединение серии, подобной co_id['co_segment'] в DataFrame невозможно в старых версиях pandas поэтому я добавил .to_frame() для совместимости. С версией pandas >= 0.25.1 эта операция разрешена, и вызов этой функции является излишним.

Примечание : этот код предполагает, что единственными сегментами являются Home , Core и Managed .

import pandas as pd

properties = pd.DataFrame(data={'property_count': [10, 200, 300, 10, 100, 200], 
                                'prov_id': [1, 1, 9, 9, 131, 199], 
                                'co_id': ['ABC', 'ABC', 'ABC', 'ABC', 'MNM', 'KJK'], 
                                'segment': ['Core', 'Home', 'Core', 'Home', 'Home', 'Home']})


def get_segment(row):
    if row['home_perc'] > 0.5:
        return 'Home'
    elif row['core_perc'] > 0.5:
        return 'Core'
    else:
        return 'Managed'


def get_grouped_dataframe(properties_df, grouping_col):
    id = pd.DataFrame()
    id['total'] = properties.groupby(grouping_col)['property_count'].sum()
    id['home'] = properties[properties.segment == 'Home'].groupby(grouping_col)['property_count'].sum()
    id['core'] = properties[properties.segment == 'Core'].groupby(grouping_col)['property_count'].sum()
    id['managed'] = properties[properties.segment == 'Managed'].groupby(grouping_col)['property_count'].sum()
    id['home_perc'] = id['home'] / id['total']
    id['home_perc'] = id['home_perc'].fillna(0)
    id['core_perc'] = id['core'] / id['total']
    id['core_perc'] = id['core_perc'].fillna(0)
    id['managed_perc'] = id['core'] / id['total']
    id['managed_perc'] = id['core_perc'].fillna(0)
    id['segment'] = id.apply(get_segment, axis=1)

    return id


prov_id = get_grouped_dataframe(properties, 'prov_id')
prov_id.rename(columns={'segment': 'prov_segment'}, inplace=True)

#          total  home   core  home_perc  core_perc prov_segment
# prov_id                                                  
# 1          210   200   10.0   0.952381   0.047619         Home
# 9          310    10  300.0   0.032258   0.967742         Core
# 131        100   100    NaN   1.000000   0.000000         Home
# 199        200   200    NaN   1.000000   0.000000         Home

co_id = get_grouped_dataframe(properties, 'co_id')
co_id.rename(columns={'segment': 'co_segment'}, inplace=True)

#        total  home   core  home_perc  core_perc co_segment
# co_id                                                  
# ABC      520   210  310.0   0.403846   0.596154       Core
# KJK      200   200    NaN   1.000000   0.000000       Home
# MNM      100   100    NaN   1.000000   0.000000       Home

property_segments = properties.drop(columns=['property_count', 'segment']).drop_duplicates()

property_segments = pd.merge(property_segments, prov_id['prov_segment'].to_frame(), on='prov_id')
property_segments = pd.merge(property_segments, co_id['co_segment'].to_frame(), on='co_id')

#    prov_id co_id co_segment prov_segment
# 0        1   ABC       Core         Home
# 1        9   ABC       Core         Core
# 2      131   MNM       Home         Home
# 3      199   KJK       Home         Home

РЕДАКТИРОВАТЬ : поставить повторяющийся код в функцию, добавил Managed сегмент в соответствии с комментарием. Добавьте дополнительный to_frame() для совместимости.