How do I integrate or use SQLModel with Pandas?

Came here for the sqmodel_to_df() function, thank you. I found that the created dataframe does not retain the column order, so here is a slight improvement which does using the first object’s schema attribute.

def sqlmodel_to_df(objects: List[SQLModel], set_index: bool = True) -> pd.DataFrame:
    """Converts SQLModel objects into a Pandas DataFrame.
    Usage
    ----------
    df = sqlmodel_to_df(list_of_sqlmodels)
    Parameters
    ----------
    :param objects: List[SQLModel]: List of SQLModel objects to be converted.
    :param set_index: bool: Sets the first column, usually the primary key, to dataframe index."""

    records = [obj.dict() for obj in objects]
    columns = list(objects[0].schema()["properties"].keys())
    df = pd.DataFrame.from_records(records, columns=columns)
    return df.set_index(columns[0]) if set_index else df

If you want to convert a DataFrame back to a list of SQLModel objects you could do this.

from typing import List, Optional

import pandas as pd
from sqlmodel import Field, SQLModel


class User(SQLModel, table=True):
    id: Optional[int] = Field(default=None, primary_key=True)
    email: str
    password: str
    name: str
    

# Create a few users for example
users = [
    User(email="sam@gmail.com", password="1234", name="sam"),
    User(email="jake@gmail.com", password="abcd", name="jake"),
    User(email="olivia@gmail.com", password="password", name="olivia"),
]


def sqmodel_to_df(objs: List[SQLModel]) -> pd.DataFrame:
    """Convert a SQLModel objects into a pandas DataFrame."""
    records = [i.dict() for i in objs]
    df = pd.DataFrame.from_records(records)
    return df

df = sqmodel_to_df(users)
print(df)

#      id             email  password    name
# 0  None     sam@gmail.com      1234     sam
# 1  None    jake@gmail.com      abcd    jake
# 2  None  olivia@gmail.com  password  Olivia

def df_to_sqlmodel(df: pd.DataFrame) -> List[SQLModel]:
    """Convert a pandas DataFrame into a a list of SQLModel objects."""    
    objs = [User(**row) for row in df.to_dict('records')]
    return objs

sql_model_objs = df_to_sqlmodel(df)
for i in sql_model_objs:
    print(i)
    
# id=None email='sam@gmail.com' password='1234' name='sam'
# id=None email='jake@gmail.com' password='abcd' name='jake'
# id=None email='olivia@gmail.com' password='password' name='olivia'

Not sure how to do this without a loop of some kind. Maybe you could come up with something that uses pandas.DataFrame.apply() but I am not sure if that would be faster than the list comprehension in my example.