Operator classes in index produces wrong SQL when using column labels

[jleguina]

Describe the bug
Operator classes using labels don't generate the correct SQL code.

Simple case:

Index(
    "ix_chunks_vector",
    Chunk.vector.label("vector"), 
    postgresql_using="hnsw",
    postgresql_with={"m": 16, "ef_construction": 64},
    postgresql_ops={"vector": "halfvec_cosine_ops"},
)

Translates to Alembic as:

    op.create_index(
        "ix_chunks_vector",
        "chunks",
        [sa.text("vector")],
        unique=False,
        postgresql_using="hnsw",
        postgresql_with={"m": 16, "ef_construction": 64},
        postgresql_ops={"vector": "halfvec_cosine_ops"},
    )

Which becomes:

CREATE INDEX ix_chunks_vector 
ON chunks
USING hnsw (vector) 
WITH (m = 16, ef_construction = 64)

The operator class is not inserted correctly into the query: hnsw (vector) should be hnsw (vector halfvec_cosine_ops)

Expected behavior
The correct query generated should've been:

CREATE INDEX ix_chunks_vector 
ON chunks
USING hnsw (vector halfvec_cosine_ops) 
WITH (m = 16, ef_construction = 64)

To Reproduce

from __future__ import annotations
import uuid


from pgvector.sqlalchemy import Vector
from sqlalchemy import UUID, Index
from sqlalchemy.orm import DeclarativeBase, Mapped, mapped_column

from src.settings import settings


class Base(DeclarativeBase):
    __abstract__ = True
    id: Mapped[uuid.UUID] = mapped_column(
        UUID(as_uuid=True),
        primary_key=True,
        default=uuid.uuid4,
        unique=True,
        index=True,
    )


class Chunk(Base):
    __tablename__ = "chunks"

    vector: Mapped[list[float]] = mapped_column(
        Vector(settings.EMBEDDING_DIMS), nullable=False
    )


Index(
    "ix_chunks_vector",
    Chunk.vector.label("vector"),
    postgresql_using="hnsw",
    postgresql_with={"m": 16, "ef_construction": 64},
    postgresql_ops={"vector": "halfvec_cosine_ops"},
)

Migration:

"""test

Revision ID: f4718fa0302f
Revises: None
Create Date: 2025-02-07 12:50:53.195748+00:00

"""

from typing import Sequence, Union

from alembic import op
import sqlalchemy as sa
import pgvector.sqlalchemy
from sqlalchemy.dialects import postgresql

# revision identifiers, used by Alembic.
revision: str = "f4718fa0302f"
down_revision: Union[str, None] = None
branch_labels: Union[str, Sequence[str], None] = None
depends_on: Union[str, Sequence[str], None] = None


def upgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.create_table(
        "chunks",
        sa.Column(
            "vector", pgvector.sqlalchemy.vector.VECTOR(dim=3072), nullable=False
        ),
        sa.Column("id", sa.UUID(), nullable=False),
        sa.PrimaryKeyConstraint("id"),
    )
    op.create_index(op.f("ix_chunks_id"), "chunks", ["id"], unique=True)
    op.create_index(
        "ix_chunks_vector",
        "chunks",
        [sa.text("vector")],
        unique=False,
        postgresql_using="hnsw",
        postgresql_with={"m": 16, "ef_construction": 64},
        postgresql_ops={"vector": "halfvec_cosine_ops"},
    )
    # ### end Alembic commands ###


def downgrade() -> None:
    # ### commands auto generated by Alembic - please adjust! ###
    op.drop_index(
        "ix_chunks_vector",
        table_name="chunks",
        postgresql_using="hnsw",
        postgresql_with={"m": 16, "ef_construction": 64},
        postgresql_ops={"vector": "halfvec_cosine_ops"},
    )
    op.drop_index(op.f("ix_chunks_id"), table_name="chunks")
    op.drop_table("chunks")
    # ### end Alembic commands ###

Error

INFO  [alembic.runtime.migration] Context impl PostgresqlImpl.
INFO  [alembic.runtime.migration] Will assume transactional DDL.
INFO  [alembic.runtime.migration] Running upgrade k631fb1nmb08 -> f4718fa0302f, test
Traceback (most recent call last):
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1964, in _exec_single_context
    self.dialect.do_execute(
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 942, in do_execute
    cursor.execute(statement, parameters)
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/psycopg/cursor.py", line 97, in execute
    raise ex.with_traceback(None)
psycopg.errors.UndefinedObject: data type vector has no default operator class for access method "hnsw"
HINT:  You must specify an operator class for the index or define a default operator class for the data type.

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/Users/javier/Projects/modelml-app/backend/.venv/bin/alembic", line 8, in <module>
    sys.exit(main())
             ^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/config.py", line 636, in main
    CommandLine(prog=prog).main(argv=argv)
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/config.py", line 626, in main
    self.run_cmd(cfg, options)
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/config.py", line 603, in run_cmd
    fn(
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/command.py", line 406, in upgrade
    script.run_env()
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/script/base.py", line 586, in run_env
    util.load_python_file(self.dir, "env.py")
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/util/pyfiles.py", line 95, in load_python_file
    module = load_module_py(module_id, path)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/util/pyfiles.py", line 113, in load_module_py
    spec.loader.exec_module(module)  # type: ignore
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "<frozen importlib._bootstrap_external>", line 995, in exec_module
  File "<frozen importlib._bootstrap>", line 488, in _call_with_frames_removed
  File "/Users/javier/Projects/modelml-app/backend/alembic/env.py", line 82, in <module>
    run_migrations_online()
  File "/Users/javier/Projects/modelml-app/backend/alembic/env.py", line 76, in run_migrations_online
    context.run_migrations()
  File "<string>", line 8, in run_migrations
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/runtime/environment.py", line 946, in run_migrations
    self.get_context().run_migrations(**kw)
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/runtime/migration.py", line 623, in run_migrations
    step.migration_fn(**kw)
  File "/Users/javier/Projects/modelml-app/backend/alembic/versions/2025_02_07_1250_test_f4718fa0302f.py", line 34, in upgrade
    op.create_index(
  File "<string>", line 8, in create_index
  File "<string>", line 3, in create_index
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/operations/ops.py", line 999, in create_index
    return operations.invoke(op)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/operations/base.py", line 442, in invoke
    return fn(self, operation)
           ^^^^^^^^^^^^^^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/operations/toimpl.py", line 114, in create_index
    operations.impl.create_index(idx, **kw)
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/ddl/postgresql.py", line 96, in create_index
    self._exec(CreateIndex(index, **kw))
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/alembic/ddl/impl.py", line 247, in _exec
    return conn.execute(construct, params)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1416, in execute
    return meth(
           ^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/sql/ddl.py", line 180, in _execute_on_connection
    return connection._execute_ddl(
           ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1527, in _execute_ddl
    ret = self._execute_context(
          ^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1843, in _execute_context
    return self._exec_single_context(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1983, in _exec_single_context
    self._handle_dbapi_exception(
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 2352, in _handle_dbapi_exception
    raise sqlalchemy_exception.with_traceback(exc_info[2]) from e
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/base.py", line 1964, in _exec_single_context
    self.dialect.do_execute(
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/sqlalchemy/engine/default.py", line 942, in do_execute
    cursor.execute(statement, parameters)
  File "/Users/javier/Projects/modelml-app/backend/.venv/lib/python3.12/site-packages/psycopg/cursor.py", line 97, in execute
    raise ex.with_traceback(None)
sqlalchemy.exc.ProgrammingError: (psycopg.errors.UndefinedObject) data type vector has no default operator class for access method "hnsw"
HINT:  You must specify an operator class for the index or define a default operator class for the data type.
[SQL: CREATE INDEX ix_chunks_vector ON chunks USING hnsw (vector) WITH (m = 16, ef_construction = 64)]
(Background on this error at: https://sqlalche.me/e/20/f405)

Versions.

• OS: MacOS 15.3
• Python: 3.12.2
• Alembic: 1.14.1
• SQLAlchemy: 2.0.36
• Database: PostgreSQL 16.4 (Debian 16.4-1.pgdg120+2)
• DBAPI: psycopg 3.2.3

Additional context
Using pgvector extension to define Vector operations, but does not make a difference.

Have a nice day!

[zzzeek]

hi -

can't reproduce this, the issue would be in SQLAlchemy where the rendering is but there is no such problem and I can't find any fixed issues where we might have fixed this. I recommend upgrading to SQLAlchemy 2.0.38

demo

from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import text, create_engine

e = create_engine("postgresql://scott:tiger@localhost/test", echo=True)
conn = e.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)

op.create_index(
        "ix_chunks_vector",
        "chunks",
        [text("vector")],
        unique=False,
        postgresql_using="hnsw",
        postgresql_with={"m": 16, "ef_construction": 64},
        postgresql_ops={"vector": "halfvec_cosine_ops"},
    )

the output SQL:

CREATE INDEX ix_chunks_vector ON chunks USING hnsw (vector halfvec_cosine_ops) WITH (m = 16, ef_construction = 64)

[zzzeek]

works in 2.0.36 also

try the above test case

[jleguina]

My bad, in trying to simplify the code to the minimum I removed the part that seems to be the cause of the issue. Doing this reproduces it:

from alembic.migration import MigrationContext
from alembic.operations import Operations
from sqlalchemy import UUID, Column, Integer, PrimaryKeyConstraint, text, create_engine

e = create_engine("postgresql+psycopg://postgres:password123@localhost/new", echo=True)
conn = e.connect()
ctx = MigrationContext.configure(conn)
op = Operations(ctx)

op.create_table(
    "chunks",
    Column("vector", Integer, nullable=False),
    Column("id", UUID(), nullable=False),
    PrimaryKeyConstraint("id"),
)
op.create_index(op.f("ix_chunks_id"), "chunks", ["id"], unique=True)
op.create_index(
    "ix_chunks_vector",
    "chunks",
    [text("CAST(vector AS VARCHAR)")],  # <---- HERE
    unique=False,
    postgresql_using="hnsw",
    postgresql_with={"m": 16, "ef_construction": 64},
    postgresql_ops={"vector": "halfvec_cosine_ops"},
)

Generates

CREATE INDEX ix_chunks_vector ON chunks USING hnsw (CAST(vector AS VARCHAR)) WITH (m = 16, ef_construction = 64)

instead of

CREATE INDEX ix_chunks_vector ON chunks USING hnsw (CAST(vector AS VARCHAR) halfvec_cosine_ops) WITH (m = 16, ef_construction = 64)

---

PS. hnsw is an index for pgvector.Vector so the code above isn't strictly correct, but is good enough to demonstrate the failed SQL. The original index definition was

Index(
    "ix_chunks_vector",
    func.cast(Chunk.vector, HALFVEC(settings.EMBEDDING_DIMS)).label("vector"),
    postgresql_using="hnsw",
    postgresql_with={"m": 16, "ef_construction": 64},
    postgresql_ops={"vector": "halfvec_cosine_ops"},
)

mapping to migration

    op.create_index(
        "ix_chunks_vector",
        "chunks",
        [sa.text("CAST(vector AS HALFVEC(3072))")],
        unique=False,
        postgresql_using="hnsw",
        postgresql_with={"m": 16, "ef_construction": 64},
        postgresql_ops={"vector": "halfvec_cosine_ops"},
    )

and resulting in

 CREATE INDEX ix_chunks_vector ON chunks USING hnsw (CAST(vector AS HALFVEC(3072))) WITH (m = 16, ef_construction = 64)

Thanks!

[zzzeek]

if there's a custom textual SQL expression as the column in this index using text(), that needs to be written out identically in the postgresql_ops, otherwise there's no magic way that the SQLAlchemy compiler would know what "vector" means, since there is no "vector" string, there's "CAST(vector AS HALFVEC(3072))", so put that string verbatim

op.create_index(
    "ix_chunks_vector",
    "chunks",
    [text("CAST(vector AS HALFVEC(3072))")],
    unique=False,
    postgresql_using="hnsw",
    postgresql_with={"m": 16, "ef_construction": 64},
    postgresql_ops={"CAST(vector AS HALFVEC(3072))": "halfvec_cosine_ops"},
)

otherwise to continue using the form at https://docs.sqlalchemy.org/en/20/dialects/postgresql.html#operator-classes, the .label() use has to be maintained. There is no way for autogenerate to render this Python code automatically from your model, so you'd need to lift and move it

from sqlalchemy import column

op.create_index(
    "ix_chunks_vector",
    "chunks",
    [
            func.cast(column("vector"), HALFVEC(settings.EMBEDDING_DIMS)).label("vector") 
    ],
    unique=False,
    postgresql_using="hnsw",
    postgresql_with={"m": 16, "ef_construction": 64},
    postgresql_ops={"vector": "halfvec_cosine_ops"},
)

[CaselIT]

This seems to be quite easy to support. It's complicated a bit by the fact that text().label() is not implemented.

Patch incoming