The code below shows a solution I have found in order to expand a dataframe to include the cartesian product of columns A
and B
, filling in the other columns with null
values. I'm wondering if there is a better and more efficient way of solving this?
>>> df = pl.DataFrame({'A': [0, 1, 1],
... 'B': [1, 1, 2],
... 'C': [6, 7, 8]})
>>> df
shape: (3, 3)
┌─────┬─────┬─────┐
│ A ┆ B ┆ C │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪═════╡
│ 0 ┆ 1 ┆ 6 │
│ 1 ┆ 1 ┆ 7 │
│ 1 ┆ 2 ┆ 8 │
└─────┴─────┴─────┘
>>> df.join(df.select('A').unique().join(df.select('B').unique(), how='cross'), on=['A','B'], how='right')
shape: (4, 3)
┌──────┬─────┬─────┐
│ C ┆ A ┆ B │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞══════╪═════╪═════╡
│ 6 ┆ 0 ┆ 1 │
│ null ┆ 0 ┆ 2 │
│ 7 ┆ 1 ┆ 1 │
│ 8 ┆ 1 ┆ 2 │
└──────┴─────┴─────┘
This is a requested feature (this is available in R or pandas's janitor as complete
).
An alternative approach mentioned in the feature request would be:
(df.select(pl.col(['A', 'B']).unique().sort().implode())
.explode('A')
.explode('B')
.join(df, how='left', on=['A', 'B'])
)
Which makes it easy to generalize to a greater number of columns.
Output:
┌─────┬─────┬──────┐
│ A ┆ B ┆ C │
│ --- ┆ --- ┆ --- │
│ i64 ┆ i64 ┆ i64 │
╞═════╪═════╪══════╡
│ 0 ┆ 1 ┆ 6 │
│ 0 ┆ 2 ┆ null │
│ 1 ┆ 1 ┆ 7 │
│ 1 ┆ 2 ┆ 8 │
└─────┴─────┴──────┘