happysundar
9/28/2014 - 5:41 PM

Example of a PLPythonU method that returns a record

Example of a PLPythonU method that returns a record

DROP FUNCTION IF EXISTS get_role_to_actor_and_actor_to_role( INOUT BIGINT, OUT JSONB, OUT JSONB );
CREATE OR REPLACE FUNCTION get_role_to_actor_and_actor_to_role(
    INOUT program_id    BIGINT,
    OUT   actor_to_role JSONB,
    OUT   role_to_actor JSONB)
    RETURNS RECORD IMMUTABLE
AS $plpython_function$
import json


def uniq(seq):
    seen = set()
    seen_add = seen.add
    return [x for x in seq if not (x in seen or seen_add(x))]


def get_role_to_actor_and_actor_to_role(program_id):
    plpy.log("getting get_role_to_actor_and_actor_to_role for program_id %s", program_id)
    plan = plpy.prepare("""
            SELECT
                credit_id,
                first_name,
                last_name,
                full_name,
                character_name,
                sequence_number
            FROM program_credits_select
            WHERE program_credits_select.program_id = $1 AND credit_type_id = 11515
            ORDER BY program_credits_select.sequence_number ASC NULLS LAST;""", ["BIGINT"])

    resultSet = plpy.execute(plan, [program_id])

    actor_to_role = dict()
    role_to_actor = dict()

    for row in resultSet:
        # plpy.log(x)

        if 'character_name' in row and row['character_name'] and 'full_name' in row and row['full_name']:
            role_name = row['character_name'].lower().strip('"')
            actor_name = row['full_name'].lower().strip('"')

            if role_name not in role_to_actor:
                role_to_actor[role_name] = []

            role_to_actor[role_name].append(actor_name)

            if actor_name not in actor_to_role:
                actor_to_role[actor_name] = []

            actor_to_role[actor_name].append(role_name)

    for actor in actor_to_role:
        actor_to_role[actor] = uniq(actor_to_role[actor])

    for role in role_to_actor:
        role_to_actor[role] = uniq(role_to_actor[role])

    rv = {'program_id':program_id, 'actor_to_role': json.dumps(actor_to_role), 'role_to_actor': json.dumps(role_to_actor)}
    # plpy.log(rv)
    return rv
    
return get_role_to_actor_and_actor_to_role(program_id)
$plpython_function$ LANGUAGE plpythonu;