How does composite index work in SQLModel? Query part

[tepelbaum]

First Check

• I added a very descriptive title here.
• I used the GitHub search to find a similar question and didn't find it.
• I searched the SQLModel documentation, with the integrated search.
• I already searched in Google "How to X in SQLModel" and didn't find any information.
• I already read and followed all the tutorial in the docs and didn't find an answer.
• I already checked if it is not related to SQLModel but to Pydantic.
• I already checked if it is not related to SQLModel but to SQLAlchemy.

Commit to Help

• I commit to help with one of those options 👆

Example Code

"""
The class definition
"""
from sqlmodel import Field
from sqlmodel import Index
from sqlmodel import SQLModel



class Foo(SQLModel, table=True):  # type: ignore
    """
    Exposure db version
    """
    __tablename__ = 'foo'
    id: int | None = Field(default=None, primary_key=True)
    bar_1: str
    bar_2: str
    bar_3: str

    __table_args__ = (
        Index(
            'foo_filter_id',
            'bar_1',
            'bar_2',
            'bar_3'
        ),
    )

"""
Notebook imports
"""
from ecodev_core.db_connection import DB_URL, DB
from sqlmodel import create_engine
engine= create_engine(DB_URL,echo=True)
import psycopg2 as psy
from psycopg2 import extras as psy_extras
from app.db_model import Foo
from sqlmodel import SQLModel, select, Session
conf = {
        'host': DB.db_host,
        'database': DB.db_name,
        'user': DB.db_username,
        'password': DB.db_password,
        'port': DB.db_port,
    }
def db_cursor():
    """
        Get database connection adn cursor based on config file
    """
    connection = psy.connect(**conf)
    return connection, connection.cursor(cursor_factory=psy_extras.RealDictCursor)


"""
psycopg2 direct call, working (meaning using the index)
"""

%%time
db_con, db_cur = db_cursor()
db_cur.execute(f"explain SELECT * from foo where bar_1 = 'a' and bar_2='g' and bar_3 = 's'")
db_cur.fetchall()


"""
SQLModel not working, meaning not using the index
"""
%%time
with Session(engine) as session:
    toto = session.exec(select(Foo.bar_1,Foo.bar_2,Foo.bar_3).where(
        Foo.bar_1 == 'a',
        Foo.bar_2 ==  'g',
        Foo.bar_3 == 's')).all()

Description

Hi.

I successfully created a composite index (see code snippet), but can't manage to forge a query that actually uses it.

Some screenshots

index indeed generated via sqlmodel

composite index working in plain sql

composite index working in psycopg2

composite index not working in SQLModel

Operating System

Linux

Operating System Details

No response

SQLModel Version

0.0.22

Python Version

3.11

Additional Context

No response

[YuriiMotov]

Why do you think it doesn't work?
You didn't use explain statement in case of SQLModel query. If you try you will see it works:

from sqlmodel import Field, Index, Session, SQLModel, col, create_engine, select, text

class Foo(SQLModel, table=True):  # type: ignore
    """
    Exposure db version
    """
    __tablename__ = "foo"
    id: int | None = Field(default=None, primary_key=True)
    bar_1: str
    bar_2: str
    bar_3: str

    __table_args__ = (Index("foo_filter_id", "bar_1", "bar_2", "bar_3"),)


engine = create_engine(
    "postgresql://user:mysecretpassword@localhost:5432/some_db", echo=False
)
SQLModel.metadata.create_all(engine)


with Session(engine) as session:
    st = (
        select(Foo.bar_1, Foo.bar_2, Foo.bar_3)
        .where(col(Foo.bar_1) == "a")
        .where(Foo.bar_2 == "g")
        .where(Foo.bar_3 == "s")
    )
    st_compiled = st.compile(engine, compile_kwargs={"literal_binds": True})
    res = session.exec(text(f"EXPLAIN {st_compiled}")).all()
    print(res)

Output:

[
    ('Index Only Scan using foo_filter_id on foo  (cost=0.15..8.17 rows=1 width=96)',),
    ("  Index Cond: ((bar_1 = 'a'::text) AND (bar_2 = 'g'::text) AND (bar_3 = 's'::text))",)
]