[Python] polars.read_database() Fails with Enum Type when Using schema_overrides

I have a database table with an Enum column. My application does transforms with polars, so I'm reading with the pl.read_database() method. However reading the table when schema_overrides is supplied fails.

Minimum Reproducible Example

Imports


# Standard Imports
import enum
import polars as pl

# Non-Sandard Imports
from sqlalchemy import Enum, Column
from sqlmodel import SQLModel, create_engine
from sqlmodel import insert, select, Field

Define a Python Native Enum


class CompassEnum(enum.Enum):
NORTH = 'N'
EAST = 'E'
SOUTH = 'S'
WEST = 'W'

Define a SQModel Table + Deploy to DB


class Compass(SQLModel, table = True):
id: int | None = Field(default = None, primary_key=True)
compass_direction: CompassEnum = Field(
sa_column = Column(
Enum(CompassEnum),
default = None,
nullable = False,
index = False
)
)

ENGINE = create_engine('my_connection', echo = True)
SQLModel.metadata.create_all(ENGINE)

Insert a Record into the Database


statement = insert(Compass)\
.values(compass_direction = 'NORTH')

conn = ENGINE.connect()
conn.execute(statement)
conn.commit()
conn.close()

Read Back from the Database


query = select(Compass.id, Compass.compass_direction)
schema = {'id': pl.String,'compass_direction': pl.Enum(CompassEnum)}

# This works
df_one = df_two = pl.read_database(query, ENGINE, infer_schema_length=1000)
print(df_one)

# This doesn't work
df_two = pl.read_database(query, ENGINE, schema_overrides = schema)
print(df_two)

The Error


"""ComputeError: could not append value: CompassEnum.NORTH of type: object to the builder; make sure that all rows have the same schema or consider increasing infer_schema_length"""

Note how when we allow schema inference with df_one, the method succeeds. However - I need to supply schema because the table can be empty in production. The database is CockroachDB which is an implementation of Postgresql.

My Essentials

• polars==1.25.2
• Python 3.11.4

Continue reading...