New/override SQLAlchemy operator compiler output

Question

The default SQLAlchemy behavior for compiling in_ expressions is pathological for very large lists, and I want to create a custom, faster, compiler for the operator. It doesn't matter to the application if the solution is a new operator (i.e.: in_list_) or if it overrides the default compiler for in_. However, I haven't been able to find any documentation on how to do this specifically.

The subclassing guidelines for compilation extension don't include anything about operators, suggesting that is not the place to start. The documentation on redefining and creating new operators is focused on changing or creating new operator behavior, but the behavior of the operator is not the problem, just the compiler.

Here is a very non-working example of what I'm trying to accomplish:

from sqlalchemy.types import TypeEngine

class in_list_(TypeEngine.Comparator):
  pass

@compiles(in_list_)
def in_list_impl(element, compiler, **kwargs):
  return "IN ('Now', 'I', 'can', 'inline', 'the', 'list')"

And then in an expression:

select([mytable.c.x, mytable.c.y]).where(mytable.c.x.in_list_(long_list))

Answer 1

Using IN for very large lists is indeed pathological, and you might be better served using a temporary table and IN against a subquery or a join. But the question was "how to override compiler output for specific operator". In case of binary operators such as IN and NOT IN what you need to override is how SQLAlchemy handles compiling BinaryExpressions:

from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.elements import BinaryExpression
from sqlalchemy.sql.operators import in_op, notin_op

def visit_in_op_binary(compiler, binary, operator, **kw):
    return "%s IN %s" % (
        compiler.process(binary.left, **kw),
        compiler.process(binary.right, **{**kw, "literal_binds": True}))

def visit_notin_op_binary(compiler, binary, operator, **kw):
    return "%s NOT IN %s" % (
        compiler.process(binary.left, **kw),
        compiler.process(binary.right, **{**kw, "literal_binds": True}))

@compiles(BinaryExpression)
def compile_binary(binary, compiler, override_operator=None, **kw):
    operator = override_operator or binary.operator

    if operator is in_op:
        return visit_in_op_binary(
            compiler, binary, operator, override_operator=override_operator,
            **kw)

    if operator is notin_op:
        return visit_notin_op_binary(
            compiler, binary, operator, override_operator=override_operator,
            **kw)

    return compiler.visit_binary(binary, override_operator=override_operator, **kw)

Note that simply producing the binary expression containing the grouping and clause list of bind params takes a surprisingly lot of time for very large lists, not to mention compiling all that even if using literal binds, so you might not observe significant performance gains. On the other hand many implementations have limits on how many placeholders / parameters you can use in a statement, and so inlining the binds allows such queries to run at all.

If on the other hand your list does fit within the limits set by your implementation (Postgresql seems to be only limited by available RAM), you may not need any compiler workarounds with a recent enough SQLAlchemy; use expanding bind parameters instead:

In [15]: %%time
    ...: session.query(Foo).\
    ...:     filter(Foo.data.in_(range(250000))).\
    ...:     all()
    ...: 
CPU times: user 5.09 s, sys: 91.9 ms, total: 5.18 s
Wall time: 5.18 s
Out[15]: []

In [16]: %%time
    ...: session.query(Foo).\
    ...:     filter(Foo.data.in_(bindparam('xs', range(250000), expanding=True))).\
    ...:     all()
    ...: 
CPU times: user 310 ms, sys: 8.05 ms, total: 318 ms
Wall time: 317 ms
Out[16]: []

And as mentioned in comments, in version 1.4 an expanding bindparam will support literal execution out of the box:

In [4]: session.query(Foo).\
   ...:     filter(Foo.data.in_(
   ...:         bindparam('xs', range(10), expanding=True, literal_execute=True))).\
   ...:     all()
2019-09-07 20:35:04,560 INFO sqlalchemy.engine.base.Engine BEGIN (implicit)
2019-09-07 20:35:04,561 INFO sqlalchemy.engine.base.Engine SELECT foo.id AS foo_id, foo.data AS foo_data 
FROM foo 
WHERE foo.data IN (0, 1, 2, 3, 4, 5, 6, 7, 8, 9)
2019-09-07 20:35:04,561 INFO sqlalchemy.engine.base.Engine ()
Out[4]: []

Answer 2

One thing you can do is to drop to raw sql and build the query by hand. But, and this is crucial, YOU HAVE TO USE BINDS.

And once you commit to that, you need to manage the entire query's variables, not just the IN list bits. So you're totally on your own. It's unrealistic to do this frequently, unless you have a specialized, highly tested utility function, but it does work. And pretty quickly too: I have to chunk at 999 because Oracle doesn't go above, but neither Postgresql or Oracle have complained much otherwise. And, yes, this is under SQLAlchemy (1.3.8).

Here's some sample code. The bits for the list were all generated dynamically, to target Postgresql. Unfortunately, each RDBMS has its own flavor of placeholder format and bind variable. Refer to PEP249 paramstyles for more details.

what the generated query looks like:

qry = """select recname, objecttype
            from bme_mvprd
            where ignore = false
YOU HAVE TO BUILD THIS 
            and objecttype 
in ( 
%(objecttypes_000__)s
, %(objecttypes_001__)s
, %(objecttypes_002__)s
, %(objecttypes_003__)s
)
...
"""

The criteria originally came from this array: [0, 1, 2, 4]

and what you are passing as a bind parameter to execute look like this:

and what the bind parameters look like, again, Postgres-specific:

(yes, you need to generate this too)

sub = {
    'objecttypes_000__': 0,
    'objecttypes_001__': 1,
    'objecttypes_002__': 2,
    'objecttypes_003__': 4,
}

AND you have to execute(qry, sub) to use binding.

NO CHEATING with execute(qry % sub) which would work in Postgresql but would drop you right back into SQL Injection land.

Oracle uses :1, :2, :3 type placeholders so that would error out but Postgresql works with Python-type placeholders so you need to be super-careful you're not accidentally bypassing parameter binding.

Note: large IN SQL operations can sometimes be replaced with EXISTS tests and those ought to be preferred if possible. My chunk-by-999 example above is because there's no other way other than first populating a temp tables: it's an example, not best practice.

PPS: and what if your list is empty? I asked a question about that => answer = "... in (%(var001)s)...", {'var001':None} but you could just not add the IN (...) at all.

Answer 3

Full disclosure, I don't know how to override the compiler of in_(), and I know I'm at risk of oversimplifying, so please accept this in good faith, but given your example above I wouldn't even try. Rather, just create a helper function:

from sqlalchemy import Table, Column, String, MetaData
from sqlalchemy.sql import text


long_list = ['Now', 'I', 'can', 'inline', 'the', 'list']
tbl = Table("mytable", MetaData(), Column("x", String), Column("y", String))


def col_in_list(col, l):
    # do something safe to generate your in clause here.
    return text("IN ('Now', 'I', 'can', 'inline', 'the', 'list')")


if __name__ == "__main__":
    print(tbl.select().where(col_in_list(tbl.c.x, long_list)))

Renders:

SELECT mytable.x, mytable.y
FROM mytable
WHERE mytable.x IN ('Now','I','can','inline','the','list')