Question

Suppose I have two tables A and B.

Table A has a multi-level index (a, b) and one column (ts).
b determines univocally ts.

A = pd.DataFrame(
     [('a', 'x', 4), 
      ('a', 'y', 6), 
      ('a', 'z', 5), 
      ('b', 'x', 4), 
      ('b', 'z', 5), 
      ('c', 'y', 6)], 
     columns=['a', 'b', 'ts']).set_index(['a', 'b'])
AA = A.reset_index()

Table B is another one-column (ts) table with non-unique index (a).
The ts’s are sorted “inside” each group, i.e., B.ix[x] is sorted for each x.
Moreover, there is always a value in B.ix[x] that is greater than or equal to
the values in A.

B = pd.DataFrame(
    dict(a=list('aaaaabbcccccc'), 
         ts=[1, 2, 4, 5, 7, 7, 8, 1, 2, 4, 5, 8, 9])).set_index('a')

The semantics in this is that B contains observations of occurrences of an event of type indicated by the index.

I would like to find from B the timestamp of the first occurrence of each event type after the timestamp specified in A for each value of b. In other words, I would like to get a table with the same shape of A, that instead of ts contains the “minimum value occurring after ts” as specified by table B.

So, my goal would be:

C: 
('a', 'x') 4
('a', 'y') 7
('a', 'z') 5
('b', 'x') 7
('b', 'z') 7
('c', 'y') 8

I have some working code, but is terribly slow.

C = AA.apply(lambda row: (
    row[0], 
    row[1], 
    B.ix[row[0]].irow(np.searchsorted(B.ts[row[0]], row[2]))), axis=1).set_index(['a', 'b'])

Profiling shows the culprit is obviously B.ix[row[0]].irow(np.searchsorted(B.ts[row[0]], row[2]))). However, standard solutions using merge/join would take too much RAM in the long run.

Consider that now I have 1000 a‘s, assume constant the average number of b’s per a (probably 100-200), and consider that the number of observations per a is probably in the order of 300. In production I will have 1000 more a‘s.

1,000,000 x 200 x 300 = 60,000,000,000 rows

may be a bit too much to keep in RAM, especially considering that the data I need is perfectly described by a C like the one I discussed above.

How would I improve the performance?

Answer 1

Thanks for providing sample data. I’ve updated this answer with general
suggestions given anticipated array sizes in the 100’s of million.

1. Line profile

   Line profiling the guts of your lambda function shows that most time is spent
   in B.ix[] (which has been refactored here to only be called once).

   In [91]: lprun -f stack.foo1 AA.apply(stack.foo1, B=B, axis=1)
   Timer unit: 1e-06 s
   
   File: stack.py
   Function: foo1 at line 4
   Total time: 0.006651 s
   
   Line #      Hits         Time  Per Hit   % Time  Line Contents
   ==============================================================
        4                                           def foo1(row, B):
        5         6         6158   1026.3     92.6      subset = B.ix[row[0]].ts
        6         6          418     69.7      6.3      idx = np.searchsorted(subset, row[2])
        7         6           56      9.3      0.8      val = subset.irow(idx)
        8         6           19      3.2      0.3      return val
   

2. Consider built-in data types and raw numpy arrays over higher-level constructs.

   Since B behaves like a dict here and the same key is accessed many times, let’s compare df.ix to a normal Python
   dictionary (precomputed elsewhere). A dictionary with 1M keys (unique A values) should only require ~34MB (33% capacity: 3 * 1e6 * 12 bytes).

   In [102]: timeit B.ix['a']
   10000 loops, best of 3: 122 us per loop
   
   In [103]: timeit dct['a']
   10000000 loops, best of 3: 53.2 ns per loop
   

3. Replace function calls with loops

   The last major improvement I can think of would be to replace df.apply() with a for loop to avoid calling any function 200M times (or however large A is).

Hopefully these ideas help.

---

Original, expressive solution, though not memory efficient:

In [5]: CC = AA.merge(B, left_on='a', right_index=True)

In [6]: CC[CC.ts_x <= CC.ts_y].groupby(['a', 'b']).first()
Out[6]: 
     ts_x  ts_y
a b            
a x     4     4
  y     6     7
  z     5     5
b x     4     7
  z     5     7
c y     6     8