Composing Polymorphic Unions With Lambda Statements

[sonthonaxrk]

Background

A few months ago I was asking about poor performance generating a complex query. And I still haven't been able to reproduce it outside my actual codebase. It involves a lot of different entities and recursive ctes. I think the SQLAlchemy 'bug' lies in how in just how much the compiler traverses.

From cProfile:

   ncalls  tottime  percall  cumtime  percall filename:lineno(function)
   416141    0.027    0.000    0.027    0.000 annotation.py:216(__hash__)
   301121    0.036    0.000    0.036    0.000 {built-in method builtins.isinstance}
   199574    0.008    0.000    0.008    0.000 {built-in method builtins.id}
 185614/1    0.354    0.000    0.501    0.501 traversals.py:159(_gen_cache_key)

What results in a massive allocation to _gen_cache_key. It's absolutely huge. One call, then nearly 200'000 recursive calls to the CacheKey named tuple. It's 0.5s of python.

The AST is quite large, but I don't have close 200'000 elements in my query and I'm a little confused as to why this got so big. But anyway, the only quick solution to this was a lambda statement.

Lambda statements are great but I've struggled to make the best use of them. For instance, I cannot figure out how to make a polymorphic_union from multiple lambda statements. To explain further, I can make a polymorphic union from a lambda statement, but each filter within that polymorphic union will create a separate lambda statement.

I agree this is an unusual problem to have, but filtering each subquery of a union statement is a lot faster.

Code attempt.

This has some idiosyncrasies, but it's meant to model what I'm running in production.

import pytest
from sqlalchemy import sql, util
from sqlalchemy.sql import util as sql_util
from sqlalchemy.sql import coercions, roles
from sqlalchemy import *
from sqlalchemy.orm import *
from sqlalchemy.ext.declarative import (
    declarative_base,
    declared_attr,
    AbstractConcreteBase,
)
from sqlalchemy import lambda_stmt
import copy
from sqlalchemy.sql import visitors
from sqlalchemy.sql.elements import ColumnClause, True_, ClauseElement, ClauseList
from sqlalchemy.sql.elements import Null


engine = create_engine("sqlite://", echo=True)
Base = declarative_base()
Base.metadata = MetaData()


class CommonMetadata(Base):
    __tablename__ = "metadata"

    @declared_attr
    def id(cls):
        return Column(
            Integer,
            primary_key=True,
        )

    data = Column(Text)


class Product(AbstractConcreteBase, Base):
    @declared_attr
    def id(cls):
        return Column(
            Integer,
            ForeignKey(
                "metadata.id",
                deferrable=True,
                initially="DEFERRED",
            ),
            primary_key=True,
        )

    def __repr__(self):
        return f"<{self.__class__.__name__} id={self.id}>"

    @classmethod
    def get_mappers(cls):
        from sqlalchemy.orm.base import _mapper_or_none

        mappers = []
        stack = list(cls.__subclasses__())
        while stack:
            klass = stack.pop()
            stack.extend(klass.__subclasses__())
            mn = _mapper_or_none(klass)
            if mn is not None:
                mappers.append(mn)

        return mappers


class Electronics(Product):
    __tablename__ = "electronics"

    clock_speed = Column(Numeric)

    __mapper_args__ = {"polymorphic_identity": "electronics", "concrete": True}


class Clothing(Product):
    __tablename__ = "clothing"

    __mapper_args__ = {"polymorphic_identity": "clothing", "concrete": True}

    size = Column(Integer)


configure_mappers()
Base.metadata.create_all(engine)


e = create_engine("sqlite://", echo=True)
Base.metadata.create_all(e)
configure_mappers()
session = Session(e)


# Create some example data
cm1 = CommonMetadata(id=1, data="cm1")
cm2 = CommonMetadata(id=2, data="cm2")
electronic = Electronics(id=1, clock_speed=1)
clothing = Clothing(id=2, size=1)


session.add_all([cm1, cm2, electronic, clothing])
session.commit()


# Query generation
cache = {}


class PolyMorphicUnionGeneratorSimple:  # simple example
    @classmethod
    def polymorphic_union(
        cls,
        table_map,
        filter_stmt,
        typecolname="type",
        aliasname="pjoin",
    ):

        colnames = util.OrderedSet()
        colnamemaps = {}
        types = {}
        for key in table_map:
            table = table_map[key]
            table = coercions.expect(
                roles.StrictFromClauseRole, table, allow_select=True
            )
            table_map[key] = table

            m = {}
            for c in table.c:
                colnames.add(c.key)
                m[c.key] = c
                types[c.key] = c.type
            colnamemaps[table] = m

        def col(name, table):
            try:
                return colnamemaps[table][name]
            except KeyError:
                return sql.cast(sql.null(), types[name]).label(name)

        result = []
        for type_, table in table_map.items():
            query = sql.select(
                *(
                    [col(name, table) for name in colnames]
                    + [
                        sql.literal_column(sql_util._quote_ddl_expr(type_)).label(
                            typecolname
                        )
                    ]
                )
            )

            # So this is the problem I have with the lambda statement:
            # I can't get this alias AoT because it's chained from this
            # arbitary select
            base_product = aliased(Product, alias=query, adapt_on_names=True)
            query = query.where(*filter_stmt(base_product))

            result.append(query)

        return sql.union_all(*result).alias(aliasname)

    @classmethod
    def create_query(cls, filter_stmt):
        table_map = {
            mp.polymorphic_identity: mp.local_table.alias(
                f"product_{mp.local_table.name}"
            )
            for mp in Product.get_mappers()
        }

        union_stmt = cls.polymorphic_union(table_map, filter_stmt)
        ProductAlias = aliased(Product, alias=union_stmt, adapt_on_names=True)
        return select(ProductAlias)

    @classmethod
    def lambda_stmt_query(cls, clock_speed):
        return lambda_stmt(
            lambda: cls.create_query(
                lambda product: [product.clock_speed == clock_speed]
            ),
            lambda_cache=cache,
            track_on=[clock_speed],
        )


class PolyMorphicUnionGeneratorExpression:
    @staticmethod
    def remap_exp(filter_by_exp, base_product):
        def visit(element):
            if isinstance(element, ColumnClause):
                new = copy.copy(element)
                local_table = base_product._aliased_insp.local_table

                # If this column doesn't exist just don't filter anything
                if not element.name in local_table.columns:
                    return Null()
                else:
                    new.table = local_table
                return new

        new = copy.copy(filter_by_exp)
        return visitors.replacement_traverse(new, {}, visit)

    @classmethod
    def polymorphic_union(
        cls,
        table_map,
        filter_stmt,
        typecolname="type",
        aliasname="pjoin",
    ):

        colnames = util.OrderedSet()
        colnamemaps = {}
        types = {}
        for key in table_map:
            table = table_map[key]
            table = coercions.expect(
                roles.StrictFromClauseRole, table, allow_select=True
            )
            table_map[key] = table

            m = {}
            for c in table.c:
                colnames.add(c.key)
                m[c.key] = c
                types[c.key] = c.type
            colnamemaps[table] = m

        def col(name, table):
            try:
                return colnamemaps[table][name]
            except KeyError:
                return sql.cast(sql.null(), types[name]).label(name)

        result = []
        for type_, table in table_map.items():
            query = sql.select(
                *(
                    [col(name, table) for name in colnames]
                    + [
                        sql.literal_column(sql_util._quote_ddl_expr(type_)).label(
                            typecolname
                        )
                    ]
                )
            )

            # So this is the problem I have with the lambda statement:
            # I can't get this alias AoT because it's chained from this
            # arbitary select
            base_product = aliased(Product, alias=table, adapt_on_names=True)

            """
            So here if I do

            >>> str(filter_stmt)
            'pjoin.clock_speed = :bind'

            Which is correct.  But because I need to filter each interior subquery
            within a CTE I have to remap it.  Because expressions passed into filter_stmt
            can have an arbitrary complexity we have to recursively iterate though it to remap it
            """
            exp = cls.remap_exp(filter_stmt, base_product)
            query = query.where(exp)
            result.append(query)

        return sql.union_all(*result).alias(aliasname)

    @classmethod
    def create_query(cls, filter_stmt):
        table_map = {
            mp.polymorphic_identity: mp.local_table.alias(
                f"product_{mp.local_table.name}"
            )
            for mp in Product.get_mappers()
        }

        union_stmt = cls.polymorphic_union(table_map, filter_stmt)
        ProductAlias = aliased(Product, alias=union_stmt, adapt_on_names=True)
        return select(ProductAlias)

    @classmethod
    def lambda_stmt_query(cls, exp):
        return lambda_stmt(
            lambda: cls.create_query(exp), lambda_cache=cache, track_on=[exp]
        )



## Test results


Lambda statements are really amazing in how they handle complex cases.

assert len(cache) == 0
PolyMorphicUnionGeneratorExpression.lambda_stmt_query(Product.clock_speed == 1)
PolyMorphicUnionGeneratorExpression.lambda_stmt_query(Product.clock_speed == 2)
assert len(cache) == 1

However the problem with the above, is that the cache should ideally have two components.  The CTE part, and the filter part.

I end up in situations where small changes to the query trigger an expensive regeneration.


```python
stmt4 = PolyMorphicUnionGeneratorExpression.lambda_stmt_query(Product.size == 1)
assert len(cache) == 2

I fully appreciate this is a complex issue. It's also not a blocking issue for me anymore since I can just trigger some queries on app startup to warm the lambda cache up. However if someone gives a query filter combination I haven't seen before I have to recalculate everything - which in my case seems to take 2s.

[zzzeek]

Background

A few months ago I was asking about poor performance generating a complex query. And I still haven't been able to reproduce it outside my actual codebase. It involves a lot of different entities and recursive ctes. I think the SQLAlchemy 'bug' lies in how in just how much the compiler traverses.

well the cache key i thought. the compiler should not be traversing any more beyond the SQL strings it outputs so should be proportional to the size of your query.

What results in a massive allocation to _gen_cache_key. It's absolutely huge. One call, then nearly 200'000 recursive calls to the CacheKey named tuple. It's 0.5s of python.

The AST is quite large, but I don't have close 200'000 elements in my query and I'm a little confused as to why this got so big.

sure, I continue to be completely alarmed by your problem as the cache key gen has many guards in place to prevent excess recursion. it remains absolutely critical that I get a test case that can reproduce this and if not, if you can work with your company to get me temporary access into a pdb session with your codebase so that I can see what the problem is. I am offering support for free here that should be seen as extremely high value to your organization.

But anyway, the only quick solution to this was a lambda statement.

im not sure it was "quick" bc we had a lot of problems with lambda also as lambda statements aren't really that mature or ready for intensive use in all cases.

Lambda statements are great but I've struggled to make the best use of them. For instance, I cannot figure out how to make a polymorphic_union from multiple lambda statements. To explain further, I can make a polymorphic union from a lambda statement, but each filter within that polymorphic union will create a separate lambda statement.

I can look at this code in the coming days but if there are more inherent limitations in lambdas, it's not priority on this end to make more unusual cases work with them, the priority is to fix the caching bug so I here continue my appeal to you to please get me a reproducing test case and/or the access I need in order to understand the problem. thanks!

[zzzeek]

just a note, I likely mentioned this earlier, if you can just provide me with a cProfile output file (the binary file), I can likely learn a lot about the cache key w/ 200K entries just from that alone. Just need some information, thanks.

[zzzeek]

is the problem here simply that there are four cache entries and not two? There's lots of weird little reasons that can happen and unless it's putting in hundreds of items for a particular statement that has only one or two varieties, this is likely safe to ignore.

[zzzeek]

also i tried looking up where we had details on your nested CTE case w/ cache key and I think that was only in chat ? I really really would like to get a cprofile file, just that alone and I can at least get a sense of what it's traversing.

[sonthonaxrk]

Hey @zzzeek I'm really sorry for just getting back to you know. I have authorisation to actually send you the code.

I'll send you a cprofile now.

[sonthonaxrk]

Sent you the full Cprofile on gitter.

But you can see how long traversals takes. I'd love to share with you the full code, but need to work out a way to do it securely.