pivot tables with sqlalchemy

If your database doesn’t support pivots, here is a quick technique to get pivot columns with sqlalchemy

import operator
from sqlalchemy.sql import case, func, select

def pivot_report(report, pivot_on=None, pivot_columns=None, pivot_func=func.sum,
                    non_pivot_columns=None, group_by=None):
    """ produce a pivot on a select

    if we have a report: 

        id, type, count
        1, white, 10
        1, black, 20
        2, white, 12
        2, black, 20

    and we want 

        id, black count, white count
        1, 20, 10
        2, 20, 12

    pass in type as the pivot_on, and [count] as pivot columns  

    """

    # find all possible values of the pivot
    pivot_values = map(
        operator.itemgetter(0),
        select([pivot_on], from_obj=[report]).distinct().execute()
    )

    # build the new pivot columns
    new_columns = [
        pivot_func(case([(pivot_on == value, column)])).label("%s %s" % (value, column))
        for value in pivot_values
        for column in pivot_columns
    ]

    return select(
        non_pivot_columns + new_columns,
        from_obj=[report],
        group_by=group_by
    )

# example code
from sqlalchemy import Table, Column, Integer, String, MetaData
from sqlalchemy import create_engine

metadata = MetaData()
example = Table('example', metadata,
    Column('id', Integer),
    Column('type', String),
    Column('count', Integer),
)

engine = create_engine('sqlite:///:memory:')
metadata.bind = engine
example.create()
example.insert().execute(
    dict(id=1, type='white', count=10),
    dict(id=1, type='black', count=20),
    dict(id=2, type='white', count=30),
    dict(id=2, type='black', count=40),
)

report = example.select()

# now build the pivot
report = pivot_report(
    report,
    pivot_on=report.c.type,
    pivot_columns=[report.c.count],
    non_pivot_columns=[report.c.id],
    group_by=[report.c.id])

for r in report.execute():
    print r.items()

Post a Comment

Your email is never shared. Required fields are marked *

*
*