Индикатор объединения в Polars: Полное руководство

import polars as pl

# Пример данных
df1 = pl.DataFrame({
    'make': ['Toyota', 'Honda', 'Ford'],
    'size': ['medium', 'small', 'large']
})

df2 = pl.DataFrame({
    'make': ['Toyota', 'Honda', 'BMW'],
    'expense': [25000, 18000, 35000]
})

# Выполнение left join
result = df1.join(df2, on='make', how='left')

# Добавление индикаторного столбца вручную
result = result.with_columns(
    pl.when(pl.col('expense').is_null())
    .then('left_only')
    .when(pl.col('make').is_in(df2['make']))
    .then('both')
    .otherwise('right_only')
    .alias('merge_indicator')
)

print(result)

def analyze_join_result(left_df, right_df, join_result, join_type='left'):
    """
    Анализ результата объединения для отображения количества совпавших/несопоставленных строк
    """
    # Подсчет уникальных значений в ключевом столбце для каждого dataframe
    left_unique = left_df.select('make').n_unique()
    right_unique = right_df.select('make').n_unique()
    
    # Подсчет совпавших строк (не-null в столбцах правого dataframe)
    if join_type in ['left', 'inner', 'full']:
        matched_rows = join_result.filter(
            ~pl.col('expense').is_null()
        ).height
    else:
        matched_rows = 0
    
    # Расчет несопоставленных строк на основе типа объединения
    if join_type == 'left':
        unmatched_left = left_unique - matched_rows
        unmatched_right = 0
    elif join_type == 'right':
        unmatched_left = 0
        unmatched_right = right_unique - matched_rows
    elif join_type == 'full':
        unmatched_left = left_unique - matched_rows
        unmatched_right = right_unique - matched_rows
    else:  # inner
        unmatched_left = left_unique - matched_rows
        unmatched_right = right_unique - matched_rows
    
    # Отображение результатов, аналогично выводу STATA
    print(f"Тип объединения: {join_type.upper()}")
    print(f"Строк из левого набора данных: {left_unique}")
    print(f"Строк из правого набора данных: {right_unique}")
    print(f"Совпавших строк: {matched_rows}")
    print(f"Несопоставленных строк из левого: {unmatched_left}")
    print(f"Несопоставленных строк из правого: {unmatched_right}")
    print(f"Всего строк в результате: {join_result.height}")

# Использование
result = df1.join(df2, on='make', how='left')
analyze_join_result(df1, df2, result, 'left')

def smart_join(left_df, right_df, on=None, how='left', **kwargs):
    """
    Расширенная функция объединения с отслеживанием строк и анализом
    """
    # Выполнение объединения
    result = left_df.join(right_df, on=on, how=how, **kwargs)
    
    # Получение имен столбцов из правого dataframe (исключая ключи объединения)
    right_cols = [col for col in right_df.columns if col not in (on if isinstance(on, list) else [on])]
    
    # Создание индикаторного столбца
    if how == 'left':
        indicator = pl.when(pl.col(right_cols[0]).is_null()).then('left_only').otherwise('both')
    elif how == 'right':
        indicator = pl.when(pl.col(right_cols[0]).is_null()).then('right_only').otherwise('both')
    elif how == 'full':
        indicator = pl.when(
            pl.col(right_cols[0]).is_null() & pl.col(on if isinstance(on, list) else [on]).is_not_null()
        ).then('left_only').when(
            pl.col(right_cols[0]).is_not_null() & pl.col(on if isinstance(on, list) else [on]).is_null()
        ).then('right_only').otherwise('both')
    else:  # inner
        indicator = pl.lit('both')
    
    result = result.with_columns(indicator.alias('merge_indicator'))
    
    # Выполнение анализа
    left_count = left_df.height
    right_count = right_df.height
    matched_count = result.filter(pl.col('merge_indicator') == 'both').height
    unmatched_left = result.filter(pl.col('merge_indicator') == 'left_only').height
    unmatched_right = result.filter(pl.col('merge_indicator') == 'right_only').height
    
    # Вывод анализа
    print(f"\n{'='*50}")
    print(f"АНАЛИЗ ОБЪЕДИНЕНИЯ - {how.upper()} JOIN")
    print(f"{'='*50}")
    print(f"Строк в левом наборе данных: {left_count}")
    print(f"Строк в правом наборе данных: {right_count}")
    print(f"Совпавших строк: {matched_count}")
    print(f"Несопоставленных строк из левого: {unmatched_left}")
    print(f"Несопоставленных строк из правого: {unmatched_right}")
    print(f"Всего строк в результате: {result.height}")
    print(f"{'='*50}\n")
    
    return result

# Использование
result = smart_join(df1, df2, on='make', how='left')

# Подход с pandas
import pandas as pd

df1_pd = df1.to_pandas()
df2_pd = df2.to_pandas()

pandas_result = pd.merge(
    df1_pd, df2_pd, 
    on='make', 
    how='left', 
    indicator=True
)

print(pandas_result['_merge'].value_counts())