将极坐标数据框与 ENUM 类型的列连接起来

假设有两个pl.DataFrame具有相同架构的 。其中一列具有dtype=pl.Enum。

import polars as pl

enum_col1 = pl.Enum(["type1"])
enum_col2 = pl.Enum(["type2"])
df1 = pl.DataFrame(
    {"enum_col": "type1", "value": 10},
    schema={"enum_col": enum_col1, "value": pl.Int64},
)
df2 = pl.DataFrame(
    {"enum_col": "type2", "value": 200},
    schema={"enum_col": enum_col2, "value": pl.Int64},
)

print(df1)
print(df2)

shape: (1, 2)
┌──────────┬───────┐
│ enum_col ┆ value │
│ ---      ┆ ---   │
│ enum     ┆ i64   │
╞══════════╪═══════╡
│ type1    ┆ 10    │
└──────────┴───────┘
shape: (1, 2)
┌──────────┬───────┐
│ enum_col ┆ value │
│ ---      ┆ ---   │
│ enum     ┆ i64   │
╞══════════╪═══════╡
│ type2    ┆ 200   │
└──────────┴───────┘

如果我尝试做一个简单的事情pl.concat([df1, df2])，我会收到以下错误：

polars.exceptions.SchemaError: type Enum(Some(local), Physical) is incompatible with expected type Enum(Some(local), Physical)

您可以通过“扩大”枚举来解决此问题，如下所示：

pl.concat(
    [
        df1.with_columns(pl.col("enum_col").cast(pl.Enum(["type1", "type2"]))),
        df2.with_columns(pl.col("enum_col").cast(pl.Enum(["type1", "type2"]))),
    ]
)

shape: (2, 2)
┌──────────┬───────┐
│ enum_col ┆ value │
│ ---      ┆ ---   │
│ enum     ┆ i64   │
╞══════════╪═══════╡
│ type1    ┆ 10    │
│ type2    ┆ 200   │
└──────────┴───────┘

我想，有更符合 Python 风格的方法来做到这一点？