Создать все комбинации на основе подмножества переменных с помощью Polars?Python

Программы на Python
Anonymous
Создать все комбинации на основе подмножества переменных с помощью Polars?

Сообщение Anonymous »

У меня есть DataFrame, который выглядит следующим образом:
import polars as pl

df = pl.DataFrame(
{
"country": ["France", "France", "UK", "UK", "Spain"],
"year": [2020, 2021, 2019, 2020, 2022],
"value": [1, 2, 3, 4, 5],
}
)

df

shape: (5, 3)
┌─────────┬──────┬───────┐
│ country ┆ year ┆ value │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞═════════╪══════╪═══════╡
│ France ┆ 2020 ┆ 1 │
│ France ┆ 2021 ┆ 2 │
│ UK ┆ 2019 ┆ 3 │
│ UK ┆ 2020 ┆ 4 │
│ Spain ┆ 2022 ┆ 5 │
└─────────┴──────┴───────┘

Я хотел бы создать сбалансированную панель, создав все пары «страна-год». В R я мог бы использовать для этого tidyr::complete(), но в Polars я не нашел встроенного способа сделать это. Есть ли что-то подобное? Если нет, то какой самый быстрый способ имитировать это?
Ожидаемый результат:
shape: (12, 3)
┌─────────┬──────┬───────┐
│ country ┆ year ┆ value │
│ --- ┆ --- ┆ --- │
│ str ┆ i64 ┆ i64 │
╞═════════╪══════╪═══════╡
│ France ┆ 2019 ┆ null │
│ France ┆ 2020 ┆ 1 │
│ France ┆ 2021 ┆ 2 │
│ France ┆ 2022 ┆ null │
│ UK ┆ 2019 ┆ 3 │
│ UK ┆ 2020 ┆ 4 │
│ UK ┆ 2021 ┆ null │
│ UK ┆ 2022 ┆ null │
│ Spain ┆ 2019 ┆ null │
│ Spain ┆ 2020 ┆ null │
│ Spain ┆ 2021 ┆ null │
│ Spain ┆ 2022 ┆ 5 │
└─────────┴──────┴───────┘



Изменить: приведенный выше пример довольно прост, поскольку в нем нужно выполнить только 2 переменных. но с тремя переменными стало сложнее, и я не понимаю, как адаптировать Pivot() + unpivot():
import polars as pl

df = pl.DataFrame(
{
"orig": ["France", "France", "UK", "UK", "Spain"],
"dest": ["Japan", "Vietnam", "Japan", "China", "China"],
"year": [2020, 2021, 2019, 2020, 2022],
"value": [1, 2, 3, 4, 5],
}
)
df

shape: (5, 4)
┌────────┬─────────┬──────┬───────┐
│ orig ┆ dest ┆ year ┆ value │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 ┆ i64 │
╞════════╪═════════╪══════╪═══════╡
│ France ┆ Japan ┆ 2020 ┆ 1 │
│ France ┆ Vietnam ┆ 2021 ┆ 2 │
│ UK ┆ Japan ┆ 2019 ┆ 3 │
│ UK ┆ China ┆ 2020 ┆ 4 │
│ Spain ┆ China ┆ 2022 ┆ 5 │
└────────┴─────────┴──────┴───────┘

Хотя оригинал работает, он намного медленнее, чем tidyr::complete() (66 мс для Polars, 1,8 мс для tidyr::complete()) :
import time

tic = time.perf_counter()
(
df
.select("orig")
.unique()
.join(df.select("dest").unique(), how="cross")
.join(df.select("year").unique(), how="cross")
.join(df, how="left", on=["country", "year"])
)
toc = time.perf_counter()
print(f"Lazy eval: {toc - tic:0.4f} seconds")

shape: (36, 4)
┌───────┬─────────┬──────┬───────┐
│ orig ┆ dest ┆ year ┆ value │
│ --- ┆ --- ┆ --- ┆ --- │
│ str ┆ str ┆ i64 ┆ i64 │
╞═══════╪═════════╪══════╪═══════╡
│ Spain ┆ Japan ┆ 2021 ┆ null │
│ Spain ┆ Japan ┆ 2022 ┆ null │
│ Spain ┆ Japan ┆ 2019 ┆ null │
│ Spain ┆ Japan ┆ 2020 ┆ null │
│ … ┆ … ┆ … ┆ … │
│ UK ┆ Vietnam ┆ 2021 ┆ null │
│ UK ┆ Vietnam ┆ 2022 ┆ null │
│ UK ┆ Vietnam ┆ 2019 ┆ null │
│ UK ┆ Vietnam ┆ 2020 ┆ null │
└───────┴─────────┴──────┴───────┘
>>>
>>> toc = time.perf_counter()
>>> print(f"Lazy eval: {toc - tic:0.4f} seconds")
Lazy eval: 0.0669 seconds

В R:
test # A tibble: 1 × 6
#> expression min median `itr/sec` mem_alloc `gc/sec`
#>
#> 1 test 1.61ms 1.81ms 496. 4.6MB 10.1


Подробнее здесь: https://stackoverflow.com/questions/766 ... ith-polars

Вернуться в «Python»