scalars/mappings performance improvement. - #13415

"""Micro-benchmark: Core .scalars().all() / .mappings().all() vs
plain .all(), 10k rows, in-memory SQLite, min-of-rounds.

Run identically on `main` and on the branch (rebuild the extensions
at -O3 after each checkout!). Compare the `min ms` column only —
the machine is noisy.
"""

import gc
import statistics
import time

from sqlalchemy import Column
from sqlalchemy import create_engine
from sqlalchemy import Integer
from sqlalchemy import MetaData
from sqlalchemy import select
from sqlalchemy import String
from sqlalchemy import Table
from sqlalchemy.types import TypeDecorator


class UpperString(TypeDecorator):
    impl = String(64)
    cache_ok = True

    def process_result_value(self, value, dialect):
        return value.upper() if value is not None else None


N_ROWS = 10_000

engine = create_engine("sqlite://")
metadata = MetaData()
t = Table(
    "bench",
    metadata,
    Column("id", Integer, primary_key=True, autoincrement=False),
    Column("a", Integer),
    Column("b", Integer),
    Column("name", String(64)),
    Column("deco", UpperString()),
)
metadata.create_all(engine)
with engine.begin() as conn:
    conn.execute(
        t.insert(),
        [
            {
                "id": i,
                "a": i % 7,
                "b": i * 2,
                "name": "name-%d" % i,
                "deco": "deco-%d" % i,
            }
            for i in range(N_ROWS)
        ],
    )

CASES = {}


def case(fn):
    CASES[fn.__name__] = fn
    return fn


@case
def scalars_int(conn):
    return conn.execute(select(t.c.a)).scalars().all()


@case
def scalars_processor(conn):
    return conn.execute(select(t.c.deco)).scalars().all()


@case
def scalars_index_2(conn):
    return conn.execute(select(t.c.id, t.c.a, t.c.name)).scalars(2).all()


@case
def scalars_unique(conn):
    return conn.execute(select(t.c.a)).scalars().unique().all()


@case
def mappings_5col(conn):
    return conn.execute(select(t)).mappings().all()


@case
def rows_5col_baseline(conn):
    return conn.execute(select(t)).all()


def run(repeats=30, warmup=3):
    import sqlalchemy
    from sqlalchemy.engine import _result_cy

    print(
        "sqlalchemy %s  compiled=%s"
        % (sqlalchemy.__version__, _result_cy._is_compiled())
    )
    assert _result_cy._is_compiled(), (
        "extensions not compiled - rebuild before benchmarking"
    )
    results = {}
    with engine.connect() as conn:
        for name, fn in CASES.items():
            for _ in range(warmup):
                fn(conn)
            timings = []
            for _ in range(repeats):
                gc.collect()
                t0 = time.perf_counter()
                fn(conn)
                timings.append(time.perf_counter() - t0)
            results[name] = timings
    print("%-24s%10s%12s" % ("case", "min ms", "median ms"))
    for name, timings in results.items():
        print(
            "%-24s%10.3f%12.3f"
            % (
                name,
                min(timings) * 1000,
                statistics.median(timings) * 1000,
            )
        )


if __name__ == "__main__":
    run()