pythonpython-polars

Expanding polars dataframe with cartesian product of two columns


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   │
└──────┴─────┴─────┘

Solution

  • 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    │
    └─────┴─────┴──────┘