Как объединить два кадра данных Polars с несколькими типами условий? (Равенства и сравнения) ⇐ Python

[Anonymous]

Я пытаюсь преобразовать SQL-запрос в полярный код и застрял на одной строке запроса, которая в основном объединяет две таблицы с определенными условиями. Вот пример SQL-запроса, над которым я работал:

Код: Выделить всё

SELECT *
FROM table_1 as tab1
LEFT JOIN table_2 as tab2
ON tab1.article_no = tab2.article_no -- condition 1
AND DIV(tab1.variant_no, 1000) = tab2.variant_no -- condition 2
AND tab1.date_of_day BETWEEN tab2.date_from AND tab2.date_to -- condition 3

Итак, здесь при выполнении LEFT JOIN есть три условия, из которых два являются проверками на равенство, а одно проверяет, находится ли date_of_day между двумя другими столбцами даты из table_2.
В полярах, когда вам нужно соединить два фрейма данных/ленивых фреймов, вы должны использовать либо tab_1.join(tab_2, on=[certain_column]) или tab_1.join(tab_2, left_on=[certain_column], right_on=[certain_column]).
Мы не можем использовать on= и [left_on, right_on] вместе, но в этом случае я не могу поместить все три условия ни внутри on=, ни в лево_включено, right_on, поэтому я не уверен, как мне этого добиться.
Хотя я пробовал выполнять сравнение дат отдельно после объединения с первыми двумя условиями:

Код: Выделить всё

table_1 = pl.read_csv('table_1.csv')
table_2 = pl.read_csv('table_2.csv')

table_joined = (
table_1
.join(
table_2,
how='left',
left_on=[pl.col('article_no', pl.col('variant_no')//1000)],  #Giving first two conditions in join
right_on=[pl.col('article_no'), pl.col('variant_no')],
suffix="_tab2"
)
)

columns_to_check = ['price_1', 'price_2', 'price_3']

final_df = (
table_joined
.with_columns(
[
pl.when(
(pl.col('date_of_day').ge(pl.col('date_from'))) &   # giving the third condition after joining, and converting the other columns to None since it's a left join
(pl.col('date_of_day').le(pl.col('date_to')))
).then(pl.col(col)).otherwise(None).alias(col)
for col in columns_to_check   # using for loop to convert other columns to None (Null)
]
)
)

Хотя этот метод работает нормально, он требует много памяти при работе с миллионами строк, а также это не очень последовательный способ соединения (поскольку иногда в конечном кадре данных бывает больше строк, чем ожидалось). Я пытался искать в Интернете способы задать несколько типов условий, но не нашел таких примеров.
Мне нужно применить все три условия во время самого левого соединения, чтобы избежать такого прохождения и потребления памяти.
Может кто-нибудь мне помочь здесь. Заранее спасибо.
Изменить:
Ниже приведен пример данных и кода:

Код: Выделить всё

import polars as pl
from datetime import date

table_1_data = {
"article_no": [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008],
"variant_no": [5000, 6000, 7000, 8000, 9000, 6000, 4000, 6000],
"date_of_day": [date(2023, 1, 2), date(2023, 1, 12), date(2023, 1, 15), date(2023, 2, 5), date(2023, 3, 15), date(2023, 1, 2), date(2023, 1, 12), date(2023, 1, 15)],
"quantity": [10, 15, 20, 25, 30, 45, 50, 60]
}

table_1 = pl.DataFrame(table_1_data)

table_2_data = {
"article_no": [1001, 1002, 1003, 1004, 1006],
"variant_no": [5, 6, 7, 8, 9],
"date_from": [date(2023, 1, 1), date(2023, 1, 10), date(2023, 1, 1), date(2023, 1, 1), date(2023, 3, 1)],
"date_to": [date(2023, 2, 1), date(2023, 1, 15), date(2023, 1, 10), date(2023, 1, 15), date(2023, 3, 31)],
"price": [100, 110, 120, 130, 140]
}

table_2 = pl.DataFrame(table_2_data)

# Print sample data
print("Table 1:")
print(table_1)
print("\nTable 2:")
print(table_2)

joined_table = (
table_1
.join(table_2,
how='left',
left_on=[pl.col('article_no'), pl.col('variant_no')//1000],
right_on=[pl.col('article_no'), pl.col('variant_no')],
suffix='_tab2'
)
.select('article_no', 'variant_no', 'date_of_day', 'quantity', 'price')
)

print('\nJoined Table:')
print(joined_table)

exptected = {
"article_no": [1001, 1002, 1003, 1004, 1005, 1006, 1007, 1008],
"variant_no": [5000, 6000, 7000, 8000, 9000, 6000, 4000, 6000],
"date_of_day": [date(2023, 1, 2), date(2023, 1, 12), date(2023, 1, 15), date(2023, 2, 5), date(2023, 3, 15), date(2023, 1, 2), date(2023, 1, 12), date(2023, 1, 15)],
"quantity": [10, 15, 20, 25, 30, 45, 50, 60],
"price": [100, 110, None, None, None, None, None, None]
}

exptected = pl.DataFrame(exptected)

print('\nExpected Output:')
print(exptected)

Вывод этого кода:

Код: Выделить всё

Table 1:
shape: (8, 4)
┌────────────┬────────────┬─────────────┬──────────┐
│ article_no ┆ variant_no ┆ date_of_day ┆ quantity │
│ ---        ┆ ---        ┆ ---         ┆ ---      │
│ i64        ┆ i64        ┆ date        ┆ i64      │
╞════════════╪════════════╪═════════════╪══════════╡
│ 1001       ┆ 5000       ┆ 2023-01-02  ┆ 10       │
│ 1002       ┆ 6000       ┆ 2023-01-12  ┆ 15       │
│ 1003       ┆ 7000       ┆ 2023-01-15  ┆ 20       │
│ 1004       ┆ 8000       ┆ 2023-02-05  ┆ 25       │
│ 1005       ┆ 9000       ┆ 2023-03-15  ┆ 30       │
│ 1006       ┆ 6000       ┆ 2023-01-02  ┆ 45       │
│ 1007       ┆ 4000       ┆ 2023-01-12  ┆ 50       │
│ 1008       ┆ 6000       ┆ 2023-01-15  ┆ 60       │
└────────────┴────────────┴─────────────┴──────────┘

Table 2:
shape: (5, 5)
┌────────────┬────────────┬────────────┬────────────┬───────┐
│ article_no ┆ variant_no ┆ date_from  ┆ date_to    ┆ price │
│ ---        ┆ ---        ┆ ---        ┆ ---        ┆ ---   │
│ i64        ┆ i64        ┆ date       ┆ date       ┆ i64   │
╞════════════╪════════════╪════════════╪════════════╪═══════╡
│ 1001       ┆ 5          ┆ 2023-01-01 ┆ 2023-02-01 ┆ 100   │
│ 1002       ┆ 6          ┆ 2023-01-10 ┆ 2023-01-15 ┆ 110   │
│ 1003       ┆ 7          ┆ 2023-01-01 ┆ 2023-01-10 ┆ 120   │
│ 1004       ┆ 8          ┆ 2023-01-01 ┆ 2023-01-15 ┆ 130   │
│ 1006       ┆ 9          ┆ 2023-03-01 ┆ 2023-03-31 ┆ 140   │
└────────────┴────────────┴────────────┴────────────┴───────┘

Joined Table:
shape:  (8, 5)
┌────────────┬────────────┬─────────────┬──────────┬───────┐
│ article_no ┆ variant_no ┆ date_of_day ┆ quantity ┆ price │
│ ---        ┆ ---        ┆ ---         ┆ ---      ┆ ---   │
│ i64        ┆ i64        ┆ date        ┆ i64      ┆ i64   │
╞════════════╪════════════╪═════════════╪══════════╪═══════╡
│ 1001       ┆ 5000       ┆ 2023-01-02  ┆ 10       ┆ 100   │
│ 1002       ┆ 6000       ┆ 2023-01-12  ┆ 15       ┆ 110   │
│ 1003       ┆ 7000       ┆ 2023-01-15  ┆ 20       ┆ 120   │
│ 1004       ┆ 8000       ┆ 2023-02-05  ┆ 25       ┆ 130   │
│ 1005       ┆ 9000       ┆ 2023-03-15  ┆ 30       ┆ null  │
│ 1006       ┆ 6000       ┆ 2023-01-02  ┆ 45       ┆ null  │
│ 1007       ┆ 4000       ┆ 2023-01-12  ┆ 50       ┆ null  │
│ 1008       ┆ 6000       ┆ 2023-01-15  ┆ 60       ┆ null  │
└────────────┴────────────┴─────────────┴──────────┴───────┘

Expected Output:
shape: (8, 5)
┌────────────┬────────────┬─────────────┬──────────┬───────┐
│ article_no ┆ variant_no ┆ date_of_day ┆ quantity ┆ price │
│ ---        ┆ ---        ┆ ---         ┆ ---      ┆ ---   │
│ i64        ┆ i64        ┆ date        ┆ i64      ┆ i64   │
╞════════════╪════════════╪═════════════╪══════════╪═══════╡
│ 1001       ┆ 5000       ┆ 2023-01-02  ┆ 10       ┆ 100   │
│ 1002       ┆ 6000       ┆ 2023-01-12  ┆ 15       ┆ 110   │
│ 1003       ┆ 7000       ┆ 2023-01-15  ┆ 20       ┆ null  │
│ 1004       ┆ 8000       ┆ 2023-02-05  ┆ 25       ┆ null  │
│ 1005       ┆ 9000       ┆ 2023-03-15  ┆ 30       ┆ null  │
│ 1006       ┆ 6000       ┆ 2023-01-02  ┆ 45       ┆ null  │
│ 1007       ┆ 4000       ┆ 2023-01-12  ┆ 50       ┆ null  │
│ 1008       ┆ 6000       ┆ 2023-01-15  ┆ 60       ┆ null  │
└────────────┴────────────┴─────────────┴──────────┴───────┘

здесь в выходных данных строки с номерами 3 и 4 не удовлетворяют третьему условию, а именно date_from

Подробнее здесь: https://stackoverflow.com/questions/788 ... equalities