Aggregate by several keys

Question

I have this rather verbose piece of code. How can I get it shorter and simpler?

This function aggregate results of test by population and by different keys:

def aggregate(f):
  """ Aggregate population by several keys """
  byMerchant  = col.defaultdict(lambda: col.defaultdict(lambda: [0,0,0,0]))
  byPlat      = col.defaultdict(lambda: col.defaultdict(lambda: [0,0,0,0]))
  byDay       = col.defaultdict(lambda: col.defaultdict(lambda: [0,0,0,0]))
  byAffiliate = col.defaultdict(lambda: col.defaultdict(lambda: [0,0,0,0]))
  byCountry   = col.defaultdict(lambda: col.defaultdict(lambda: [0,0,0,0]))
  byType      = col.defaultdict(lambda: col.defaultdict(lambda: [0,0,0,0]))
  byPop       = col.defaultdict(lambda: [0,0,0,0])
  for line in f:
    fields = line.strip().split("^A")
    if fields[2] == '\\N':  continue

    day          = fields[0]
    popid        = int(fields[1])
    affiliate_id = int(fields[2])
    merchant_id  = int(fields[3])
    clicks       = int(fields[4])
    displays     = int(fields[5])
    sales        = float(fields[6])
    gross        = float(fields[7])

    merchant_plat = partners_info[merchant_id][1]
    merchant_name = partners_info[merchant_id][0]

    publisher_name    = publisher_info[affiliate_id][0]
    publisher_country = publisher_info[affiliate_id][1]
    publisher_type    = publisher_info[affiliate_id][2]

    data = [clicks,displays,sales,gross]
    for k in xrange(len(data)):
      byMerchant[merchant_name][popid][k]    += data[k]
      byPlat[merchant_plat][popid][k]        += data[k]
      byAffiliate[publisher_name][popid][k]  += data[k]
      byCountry[publisher_country][popid][k] += data[k]
      byDay[day][popid][k]                   += data[k]
      byType[publisher_type][popid][k]       += data[k]
      byPop[popid][k]                        += data[k]

return byPop, byMerchant, byPlat, byDay, byAffiliate, byCountry, byType

Answer 1

If I try to split on the empty string I get an error:

>>> 'a b c'.split('')
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
ValueError: empty separator

Perhaps you meant .split()?

Anyway, here are some ideas for simplifying your code.

1. Use from collections import defaultdict to avoid the need to refer to col (I presume you have import collections as col).

2. Put a dummy key into the byPop aggregation so that its structure matches the others.

3. Put the aggregates into a dictionary instead of having a bunch of variables. That way you can manipulate them systematically in loops.

4. Use enumerate in your last loop to get the index as well as the value from the data list.

That leads to this:

from collections import defaultdict

def aggregate(f):
    """
    Aggregate population by several keys.
    """
    aggregates = 'pop merchant plat day affiliate country type'.split()
    result = dict()
    for a in aggregates:
        result[a] = defaultdict(lambda: defaultdict(lambda: [0,0,0,0]))

    for line in f:
        fields = line.strip().split()
        if fields[2] == '\\N': continue

        # Decode fields from this record.
        day          = fields[0]
        popid        = int(fields[1])
        affiliate_id = int(fields[2])
        merchant_id  = int(fields[3])
        clicks       = int(fields[4])
        displays     = int(fields[5])
        sales        = float(fields[6])
        gross        = float(fields[7])

        # Keys under which to aggregate the data from this record.
        keys = dict(merchant = partners_info[merchant_id][0],
                    plat = partners_info[merchant_id][1],
                    day = day,
                    affiliate = publisher_info[affiliate_id][0],
                    country = publisher_info[affiliate_id][1],
                    type = publisher_info[affiliate_id][2],
                    pop = 0)       # dummy key

        # Update all aggregates.
        for i, value in enumerate([clicks, displays, sales, gross]):
            for a in aggregates:
                result[a][keys[a]][popid][i] += value

    return (result['pop'][0],) + tuple(result[a] for a in aggregates[1:])

It's not really a vast improvement.