Используйте поляры .when() вместо объединения

У меня есть 3 кадра данных Polars, один из которых содержит 2 IDS, а другие содержат идентификатор и значение. Я хотел бы присоединиться к трем фреймам данных, если идентификатор основной таблицы существует в одной из других таблиц, и перенести значения из нужного столбца.

Мой текущий подход заключается в том, чтобы просто переименовать каждый идентификатор таблицы, а затем сделать .join(how = 'left'), однако я думаю, что переименование и дублирование таблиц — это неправильный способ решения этой проблемы. (Из-за дополнительного кода и потраченной впустую оперативной памяти)

Первый содержит 2 столбца идентификаторов:

data = {
    "ID1" : [1,2,3],
    "ID2" : [1,4,5]
}
df = pl.DataFrame(data)

Второй и третий — это фреймы данных, которые содержат идентификатор и значение:

T1 = {
    "ID" : [9,2,5],
    "Values" : ["A","B","c"],
    "Values II" : ["foo","boo","baz"]
}
T1 = pl.DataFrame(T1)

T2 = {
    "ID" : [1,4,10],
    "Values" : ["X","J","c"]
}
T2 = pl.DataFrame(T2)

Я могу проверить, существует ли ID в других таблицах, подобных этой.

(
    df
    .with_columns(
        ID1_is_on_T1 = pl.col("ID1").is_in(T1.select(pl.col("ID"))),
        ID2_is_on_T1 = pl.col("ID2").is_in(T1.select(pl.col("ID"))),
        ID1_is_on_T2 = pl.col("ID1").is_in(T2.select(pl.col("ID"))),
        ID2_is_on_T2 = pl.col("ID2").is_in(T2.select(pl.col("ID"))),
    )
)

И я хочу сделать что-то вроде этого:

(
    df
    .with_columns(
        pl
        .when(
            pl.col("ID1").is_in(T1.select(pl.col("ID")))
        )
        .then(
            T1.select(pl.col("Values"))
        )
        .otherwise(0)
    )
)

ValueError: can only call .item() if the dataframe is of shape (1, 1), or if explicit row/col values are provided; frame has shape (3, 1)

Текущий .join() подход:

T1_1 = (
    T1
    .rename(
        {"ID": "ID1"}
    )
)

T1_2 = (
    T1
    .rename(
        {"ID": "ID2"}
    )
)

Join_1 = df.join(T1_1,on = "ID1", how = "left").rename({"Values" : "ID1_Values", "Values II" : "ID1_Values II"})
Join_2 = Join_1.join(T1_2, on = "ID2", how = "left").rename({"Values" : "ID2_Values", "Values II" : "ID2_Values II"})

При таком подходе рассматривается только первая таблица, мне нужно будет сделать то же самое и для T2.

🤔 А знаете ли вы, что...
Python позволяет создавать сценарии для автоматизации задач и обработки данных.


1
50
2

Ответы:

Решено

Два соединения здесь будут наиболее эффективными. Вы можете избежать чрезмерного переименования, указав left_on и right_on отдельно, а также используя suffix для уточнения повторяющихся имен.

>>> (df.join(T1, left_on = "ID1", right_on = "ID", how = "left")
...    .join(T2, left_on = "ID2", right_on = "ID", how = "left", suffix = "_T2"))
shape: (3, 5)
┌─────┬─────┬────────┬───────────┬───────────┐
│ ID1 ┆ ID2 ┆ Values ┆ Values II ┆ Values_T2 │
│ --- ┆ --- ┆ ---    ┆ ---       ┆ ---       │
│ i64 ┆ i64 ┆ str    ┆ str       ┆ str       │
╞═════╪═════╪════════╪═══════════╪═══════════╡
│ 1   ┆ 1   ┆ null   ┆ null      ┆ X         │
│ 2   ┆ 4   ┆ B      ┆ boo       ┆ J         │
│ 3   ┆ 5   ┆ null   ┆ null      ┆ null      │
└─────┴─────┴────────┴───────────┴───────────┘

Если вместо хранения T1 и T2 как отдельных переменных вы поместите их в словарь, вы сможете создать функцию, которая сделает все это за вас.

Настройте переменные следующим образом:

df = pl.DataFrame({
    "ID1" : [1,2,3],
    "ID2" : [1,4,5]
})

T = {}
T["T1"]=pl.DataFrame({
    "ID" : [9,2,5],
    "Values" : ["A","B","c"],
    "Values II" : ["foo","boo","baz"]
}
)
T["T2"]=pl.DataFrame({
    "ID" : [1,4,10],
    "Values" : ["X","J","c"]
})

И тогда функция

def big_join(df:pl.DataFrame, T:dict[str, pl.DataFrame])->pl.DataFrame:
    joined=df
    for id_col in df.columns:
        for tbl_name, subdf in T.items():
            joined=joined.join(
                subdf.select("ID", pl.exclude("ID").name.prefix(f"{id_col}_{tbl_name}_")),
                left_on=id_col, right_on='ID', how='left'
                )
    return joined

Если вы хотите присоединиться к ним всем, просто сделайте

big_join(df, T)
shape: (3, 8)
┌─────┬─────┬──────────────┬──────────────┬──────────────┬─────────────┬─────────────┬─────────────┐
│ ID1 ┆ ID2 ┆ ID1_T1_Value ┆ ID1_T1_Value ┆ ID1_T2_Value ┆ ID2_T1_Valu ┆ ID2_T1_Valu ┆ ID2_T2_Valu │
│ --- ┆ --- ┆ s            ┆ s II         ┆ s            ┆ es          ┆ es II       ┆ es          │
│ i64 ┆ i64 ┆ ---          ┆ ---          ┆ ---          ┆ ---         ┆ ---         ┆ ---         │
│     ┆     ┆ str          ┆ str          ┆ str          ┆ str         ┆ str         ┆ str         │
╞═════╪═════╪══════════════╪══════════════╪══════════════╪═════════════╪═════════════╪═════════════╡
│ 1   ┆ 1   ┆ null         ┆ null         ┆ X            ┆ null        ┆ null        ┆ X           │
│ 2   ┆ 4   ┆ B            ┆ boo          ┆ null         ┆ null        ┆ null        ┆ J           │
│ 3   ┆ 5   ┆ null         ┆ null         ┆ null         ┆ c           ┆ baz         ┆ null        │
└─────┴─────┴──────────────┴──────────────┴──────────────┴─────────────┴─────────────┴─────────────┘

Это можно расширить на любое количество таблиц T, если они есть в словаре. Даже если вы хотите сохранить T1 и T2 как отдельные переменные, вместо этого вы можете использовать big_join(df, {"T1":T1, "T2":T2}).

Polars не копирует без необходимости, поэтому вам не нужно беспокоиться о потраченной впустую оперативной памяти. Например, если вы это сделаете df1=df2=df3=df, он не будет копировать df 3 раза, в каждой из «копий» будет просто указатель, указывающий на df. То же самое относится и к столбцам: если вы оставили соединение df с T1, то df не копируется в новую таблицу соединения, а просто указывает на столбцы. Я не очень хорошо разбираюсь в том, когда именно создаются/необходимы копии, но их гораздо меньше, чем в других библиотеках DataFrame на основе медведей.