使用Polars基于变量子集创建所有组合。

huangapple
go评论102阅读模式
英文:

Create all combinations based on a subset of variables with Polars?

问题

以下是您提供的代码部分的翻译：

  1. import polars as pl
  3. df = pl.DataFrame(
  4.     {
  5.         "country": ["France", "France", "UK", "UK", "Spain"],
  6.         "year": [2020, 2021, 2019, 2020, 2022],
  7.         "value": [1, 2, 3, 4, 5],
  8.     }
  9. )
  11. df
  13. shape: (5, 3)
  14. ┌─────────┬──────┬───────┐
  15. │ country ┆ year ┆ value │
  16. │ ---     ┆ ---  ┆ ---   │
  17. │ str     ┆ i64  ┆ i64   │
  18. ╞═════════╪══════╪═══════╡
  19. │ France  ┆ 2020 ┆ 1     │
  20. │ France  ┆ 2021 ┆ 2     │
  21. │ UK      ┆ 2019 ┆ 3     │
  22. │ UK      ┆ 2020 ┆ 4     │
  23. │ Spain   ┆ 2022 ┆ 5     │
  24. └─────────┴──────┴───────┘
  1. import time
  3. tic = time.perf_counter()
  4. (
  5.     df
  6.     .select("country")
  7.     .unique()
  8.     .join(df.select("year").unique(), how="cross")
  9.     .join(df, how="left", on=["country", "year"])
  10. )
  11. toc = time.perf_counter()
  12. print(f"Lazy eval: {toc - tic:0.4f} seconds")
  14. shape: (36, 4)
  15. ┌───────┬─────────┬──────┬───────┐
  16. │ country │ year │ value │
  17. │ ---     │ ---  │ ---   │
  18. │ str     │ i64  │ i64   │
  19. ╞═════════╪══════╪═══════╡
  20. │ Spain   │ 2021 │ null  │
  21. │ Spain   │ 2022 │ 5     │
  22. │ Spain   │ 2019 │ null  │
  23. │ Spain   │ 2020 │ null  │
  24. │ …       │ …    │ …     │
  25. │ UK      │ 2021 │ null  │
  26. │ UK      │ 2022 │ null  │
  27. │ UK      │ 2019 │ 3     │
  28. │ UK      │ 2020 │ 4     │
  29. └───────┴─────────┴──────┴───────┘

希望这有助于您理解代码。如果您有任何其他问题，请随时提出。

英文:

I have a DataFrame that looks like this:

  1. import polars as pl
  3. df = pl.DataFrame(
  4.     {
  5.         "country": ["France", "France", "UK", "UK", "Spain"],
  6.         "year": [2020, 2021, 2019, 2020, 2022],
  7.         "value": [1, 2, 3, 4, 5],
  8.     }
  9. )
  11. df
  13. shape: (5, 3)
  14. ┌─────────┬──────┬───────┐
  15. │ country ┆ year ┆ value │
  16. │ ---     ┆ ---  ┆ ---   │
  17. │ str     ┆ i64  ┆ i64   │
  18. ╞═════════╪══════╪═══════╡
  19. │ France  ┆ 2020 ┆ 1     │
  20. │ France  ┆ 2021 ┆ 2     │
  21. │ UK      ┆ 2019 ┆ 3     │
  22. │ UK      ┆ 2020 ┆ 4     │
  23. │ Spain   ┆ 2022 ┆ 5     │
  24. └─────────┴──────┴───────┘

I'd like to make a balanced panel by creating all country-year pairs. In R, I could use tidyr::complete() for this, but I didn't find a built-in way to do this in Polars. Is there something like this? If not, what would be the fastest way to mimick it?

Expected output:

  1. shape: (12, 3)
  2. ┌─────────┬──────┬───────┐
  3. │ country ┆ year ┆ value │
  4. │ ---     ┆ ---  ┆ ---   │
  5. │ str     ┆ i64  ┆ i64   │
  6. ╞═════════╪══════╪═══════╡
  7. │ France  ┆ 2019 ┆ null  │
  8. │ France  ┆ 2020 ┆ 1     │
  9. │ France  ┆ 2021 ┆ 2     │
  10. │ France  ┆ 2022 ┆ null  │
  11. │ UK      ┆ 2019 ┆ 3     │
  12. │ UK      ┆ 2020 ┆ 4     │
  13. │ UK      ┆ 2021 ┆ null  │
  14. │ UK      ┆ 2022 ┆ null  │
  15. │ Spain   ┆ 2019 ┆ null  │
  16. │ Spain   ┆ 2020 ┆ null  │
  17. │ Spain   ┆ 2021 ┆ null  │
  18. │ Spain   ┆ 2022 ┆ 5     │
  19. └─────────┴──────┴───────┘

Edit: the example above is quite simple because it only has 2 vars to complete but it started being trickier with 3 vars and I don't see how to adapt the pivot() + melt():

  1. import polars as pl
  3. df = pl.DataFrame(
  4.     {
  5.         "orig": ["France", "France", "UK", "UK", "Spain"],
  6.         "dest": ["Japan", "Vietnam", "Japan", "China", "China"],
  7.         "year": [2020, 2021, 2019, 2020, 2022],
  8.         "value": [1, 2, 3, 4, 5],
  9.     }
  10. )
  11. df
  13. shape: (5, 4)
  14. ┌────────┬─────────┬──────┬───────┐
  15. │ orig   ┆ dest    ┆ year ┆ value │
  16. │ ---    ┆ ---     ┆ ---  ┆ ---   │
  17. │ str    ┆ str     ┆ i64  ┆ i64   │
  18. ╞════════╪═════════╪══════╪═══════╡
  19. │ France ┆ Japan   ┆ 2020 ┆ 1     │
  20. │ France ┆ Vietnam ┆ 2021 ┆ 2     │
  21. │ UK     ┆ Japan   ┆ 2019 ┆ 3     │
  22. │ UK     ┆ China   ┆ 2020 ┆ 4     │
  23. │ Spain  ┆ China   ┆ 2022 ┆ 5     │
  24. └────────┴─────────┴──────┴───────┘

While the original works, it is much slower than tidyr::complete() (66ms for Polars, 1.8ms for tidyr::complete()):

  1. import time
  3. tic = time.perf_counter()
  4. (
  5.     df
  6.     .select("orig")
  7.     .unique()
  8.     .join(df.select("dest").unique(), how="cross")
  9.     .join(df.select("year").unique(), how="cross")
  10.     .join(df, how="left", on=["country", "year"])
  11. )
  12. toc = time.perf_counter()
  13. print(f"Lazy eval: {toc - tic:0.4f} seconds")
  15. shape: (36, 4)
  16. ┌───────┬─────────┬──────┬───────┐
  17. │ orig  ┆ dest    ┆ year ┆ value │
  18. │ ---   ┆ ---     ┆ ---  ┆ ---   │
  19. │ str   ┆ str     ┆ i64  ┆ i64   │
  20. ╞═══════╪═════════╪══════╪═══════╡
  21. │ Spain ┆ Japan   ┆ 2021 ┆ null  │
  22. │ Spain ┆ Japan   ┆ 2022 ┆ null  │
  23. │ Spain ┆ Japan   ┆ 2019 ┆ null  │
  24. │ Spain ┆ Japan   ┆ 2020 ┆ null  │
  25. │ …     ┆ …       ┆ …    ┆ …     │
  26. │ UK    ┆ Vietnam ┆ 2021 ┆ null  │
  27. │ UK    ┆ Vietnam ┆ 2022 ┆ null  │
  28. │ UK    ┆ Vietnam ┆ 2019 ┆ null  │
  29. │ UK    ┆ Vietnam ┆ 2020 ┆ null  │
  30. └───────┴─────────┴──────┴───────┘
  31. >>>
  32. >>> toc = time.perf_counter()
  33. >>> print(f"Lazy eval: {toc - tic:0.4f} seconds")
  34. Lazy eval: 0.0669 seconds

In R:

  1. test <- data.frame(
  2.   orig = c("France", "France", "UK", "UK", "Spain"),
  3.   dest = c("Japan", "Vietnam", "Japan", "China", "China"),
  4.   year = c(2020, 2021, 2019, 2020, 2022),
  5.   value = c(1, 2, 3, 4, 5)
  6. )
  8. bench::mark(
  9.   test = tidyr::complete(test, orig, dest, year),
  10.   iterations = 100
  11. )
  12. #> # A tibble: 1 × 6
  13. #>   expression      min   median `itr/sec` mem_alloc `gc/sec`
  14. #>   <bch:expr> <bch:tm> <bch:tm>     <dbl> <bch:byt>    <dbl>
  15. #> 1 test         1.61ms   1.81ms      496.     4.6MB     10.1

答案1

得分: 2

也许有一个更简单的方法，但组合是唯一值的交叉连接。

  1. df.select('country').unique().join(
  2.    df.select('year').unique(),
  3.    how='cross'
  4. )
  1. shape: (12, 2)
  2. ┌─────────┬──────┐
  3. │ country ┆ year │
  4. │ ---     ┆ ---  │
  5. │ str     ┆ i64  │
  6. ╞═════════╪══════╡
  7. │ UK      ┆ 2021 │
  8. │ UK      ┆ 2022 │
  9. │ UK      ┆ 2019 │
  10. │ UK      ┆ 2020 │
  11. │ Spain   ┆ 2021 │
  12. │ Spain   ┆ 2022 │
  13. │ Spain   ┆ 2019 │
  14. │ Spain   ┆ 2020 │
  15. │ France  ┆ 2021 │
  16. │ France  ┆ 2022 │
  17. │ France  ┆ 2019 │
  18. │ France  ┆ 2020 │
  19. └─────────┴──────┘

然后，您可以与原始数据进行左连接：

  1. df.select('country').unique().join(
  2.    df.select('year').unique(),
  3.    how='cross'
  4. ).join(df, how='left', on=['country', 'year'])
  1. shape: (12, 3)
  2. ┌─────────┬──────┬───────┐
  3. │ country ┆ year ┆ value │
  4. │ ---     ┆ ---  ┆ ---   │
  5. │ str     ┆ i64  ┆ i64   │
  6. ╞═════════╪══════╪═══════╡
  7. │ UK      ┆ 2021 ┆ null  │
  8. │ UK      ┆ 2019 ┆ 3     │
  9. │ UK      ┆ 2022 ┆ null  │
  10. │ UK      ┆ 2020 ┆ 4     │
  11. │ France  ┆ 2021 ┆ 2     │
  12. │ France  ┆ 2019 ┆ null  │
  13. │ France  ┆ 2022 ┆ null  │
  14. │ France  ┆ 2020 ┆ 1     │
  15. │ Spain   ┆ 2021 ┆ null  │
  16. │ Spain   ┆ 2019 ┆ null  │
  17. │ Spain   ┆ 2022 ┆ 5     │
  18. │ Spain   ┆ 2020 ┆ null  │
  19. └─────────┴──────┴───────┘
英文:

Perhaps there is a simpler way but the combinations are a cross join of the unique values.

  1. df.select('country').unique().join(
  2. df.select('year').unique(),
  3. how = 'cross'
  4. )
  1. shape: (12, 2)
  2. ┌─────────┬──────┐
  3. │ country ┆ year │
  4. │ ---     ┆ ---  │
  5. │ str     ┆ i64  │
  6. ╞═════════╪══════╡
  7. │ UK      ┆ 2021 │
  8. │ UK      ┆ 2022 │
  9. │ UK      ┆ 2019 │
  10. │ UK      ┆ 2020 │
  11. │ Spain   ┆ 2021 │
  12. │ Spain   ┆ 2022 │
  13. │ Spain   ┆ 2019 │
  14. │ Spain   ┆ 2020 │
  15. │ France  ┆ 2021 │
  16. │ France  ┆ 2022 │
  17. │ France  ┆ 2019 │
  18. │ France  ┆ 2020 │
  19. └─────────┴──────┘

Which you can left join with the original:

  1. df.select('country').unique().join(
  2. df.select('year').unique(),
  3. how = 'cross'
  4. ).join(df, how='left', on=['country', 'year'])
  1. shape: (12, 3)
  2. ┌─────────┬──────┬───────┐
  3. │ country ┆ year ┆ value │
  4. │ ---     ┆ ---  ┆ ---   │
  5. │ str     ┆ i64  ┆ i64   │
  6. ╞═════════╪══════╪═══════╡
  7. │ UK      ┆ 2021 ┆ null  │
  8. │ UK      ┆ 2019 ┆ 3     │
  9. │ UK      ┆ 2022 ┆ null  │
  10. │ UK      ┆ 2020 ┆ 4     │
  11. │ France  ┆ 2021 ┆ 2     │
  12. │ France  ┆ 2019 ┆ null  │
  13. │ France  ┆ 2022 ┆ null  │
  14. │ France  ┆ 2020 ┆ 1     │
  15. │ Spain   ┆ 2021 ┆ null  │
  16. │ Spain   ┆ 2019 ┆ null  │
  17. │ Spain   ┆ 2022 ┆ 5     │
  18. │ Spain   ┆ 2020 ┆ null  │
  19. └─────────┴──────┴───────┘

答案2

得分: 2

以下是已翻译的代码部分：

  1. (
  2.     df.select(pl.col(["orig", "dest", "year"]).unique().sort().implode())
  3.     .explode("orig")
  4.     .explode("dest")
  5.     .explode("year")
  6.     .join(df, how="left", on=["orig", "dest", "year"])
  7. )
英文:

As pointed out by @jqurious in the comments of their answer, it is faster to use .implode() and .explode() (it isn't faster with the small example I gave but I can clearly see the difference with larger data):

  1. (
  2.     df.select(pl.col(["orig", "dest", "year"]).unique().sort().implode())
  3.     .explode("orig")
  4.     .explode("dest")
  5.     .explode("year")
  6.     .join(df, how="left", on=["orig", "dest", "year"])
  7. )

huangapple
  • 本文由 发表于 2023年7月4日 20:45:01
  • 转载请务必保留本文链接：https://go.coder-hub.com/76612786.html
  • python
  • python-polars
匿名

发表评论

匿名网友

:?: :razz: :sad: :evil: :!: :smile: :oops: :grin: :eek: :shock: :???: :cool: :lol: :mad: :twisted: :roll: :wink: :idea: :arrow: :neutral: :cry: :mrgreen:

确定