Início / coding / Perguntas / 79002206
Accepted
Andi
Asked: 2024-09-19 19:02:26 +0800 CST

Concatenar polares dataframe com colunas do dtype ENUM

  • 772

Considere ter dois pl.DataFrames com esquema idêntico. Uma das colunas tem 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   │
└──────────┴───────┘

Se eu tentar fazer algo simples pl.concat([df1, df2]), recebo o seguinte erro:

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

Você pode contornar esse problema "ampliando" as enumerações assim:

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

Acho que existe uma maneira mais pythonica de fazer isso?

python

1 respostas

  1. Best Answer

    você pode converter enum_col para o tipo de enumeração combinado:

    enum_col = enum_col1 | enum_col2

pl.concat(
    df.with_columns(pl.col.enum_col.cast(enum_col)) for df in [df1, df2]
)

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

    Você também pode criar novos enum_coldinamicamente, por exemplo:

    from functools import reduce

enum_col = reduce(lambda x,y: x | y, [df.schema["enum_col"] for df in [df1, df2]])

    Enum(categories=['type1', 'type2'])
    • 3

relate perguntas