Skip to content

PGArrayIndex

Jump to bottom
Mike Bayer edited this page Feb 20, 2014 · 2 revisions

PGArrayIndex

Note: This feature is included as of SQLAlchemy 0.8. See http://docs.sqlalchemy.org/en/latest/dialects/postgresql.html#sqlalchemy.dialects.postgresql.ARRAY.

When using the ARRAY type, SQLAlchemy doesn't yet define a meaningful __getitem__() on column expressions such that you can render expressions like foo[1](1) or bar[10](10)[20](20). In the meantime, the recipe below, based on @compiles (as usual), will allow this behavior reasonably well.

#!python
from sqlalchemy import *
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql.expression import ColumnElement
from sqlalchemy.dialects.postgresql import ARRAY, dialect

class as_indexed(ColumnElement):
    def __init__(self, expr, index=None):
        
        if not isinstance(expr, as_indexed):
            assert isinstance(expr.type, ARRAY)
            self.type = expr.type.item_type
        else:
            self.type = expr.type
        self.expr = expr
        self.index = index
    
    def __getitem__(self, index):
        if self.index is not None:
            return as_indexed(self, index)
        else:
            return as_indexed(self.expr, index)
    
@compiles(as_indexed, 'postgresql')
def compile_indexed(element, compiler, **kw):
    ret = compiler.process(element.expr, **kw)
    if element.index is not None:
       ret = "%s[%d]" % (ret, element.index)
    return ret
    
if __name__ == '__main__':
    # demonstration
    metadata = MetaData()
    
    test = Table("test",metadata,
        Column("mykey",Integer,nullable=False),
        Column("mydata",ARRAY(Float))
    )

    expr1 = as_indexed(test.c.mydata)
    assert isinstance(expr1.type, Float)
    expr2 = expr1[1]
    assert isinstance(expr2.type, Float)
    expr3 = expr2[1]
    assert isinstance(expr3.type, Float)
    
    pg_dialect = dialect()
    print select([
                test.c.mykey, 
                func.sum(as_indexed(test.c.mydata)[1][1])
            ]).group_by(test.c.mykey).\
            compile(dialect=pg_dialect)

    print select([
                test.c.mykey, 
                as_indexed(test.c.mydata)[1] + 35.5
            ]).\
            compile(dialect=pg_dialect)
Clone this wiki locally