Processing several million tweets - loop is very slow. Any ideas on how to speed it up?

Question

As the title suggests, I'm processing several million tweets and one of the data points is whether or not any of the words exist in two different lists (each list contains about 500 words). It's understandably pretty slow, but I'll be doing this regularly so I'd like to speed it up. Any thoughts on how I could so?

lista = ['word1', 'word2', ... 'word500']
listb = ['word1', 'word2', ..., 'word500']

def token_list_count(df):

    for i, t in df.iterrows():

        list_a = 0
        list_b = 0

        for tok in t['tokens']:
            if tok in lista: list_a += 1
            elif tok in listb: list_b += 1

        df.loc[i, 'token_count'] = int(len(t['tokens']))
        df.loc[i, 'lista_count'] = int(list_a)
        df.loc[i, 'listb_count'] = int(list_b)

        if i % 25000 == 0: print('25k more processed...')

    return df

Edit:

Input / Before:

Output / After:

Answer 1

In order to avoid iterating over a dataframe you might try the function below. The first line expands the tokens column into a dataframe where every first token will be in the first column, every second one in the second column and so on. Total number of columns will be equal to the largest token_count in df. From there it's easy to calculate lista and listb counts using .isin method and summing that for every row.

The tok dataframe made from full df might take a lot of memory in which case it might be necessary to split the df into several chunks and process them separately.

def token_list_count(df):
    tok = df['tokens'].apply(pd.Series)
    df['token_count'] = tok.notnull().sum(axis=1)
    df['lista_count'] = tok.isin(lista).sum(axis=1)
    df['listb_count'] = tok.isin(listb).sum(axis=1)
    return df

Your original function also has some room for improvement:

• it's not necessary to iterate over full df using iterrows() because you only need one column to do the calculations. So you can use df.tokens.iteritems() instead. This will save you creating a new Series object at every iteration.
• since you're only setting new values one at a time you can use at indexer instead of loc. See Fast scalar value getting and setting in pandas docs.

I don't know how much this will help but probably some =).

def token_list_count(df):
    for i, tok in df.tokens.iteritems():
        list_a = sum((t in lista) for t in tok)
        list_b = sum((t in listb) for t in tok)

        df.at[i, 'token_count'] = int(len(tok))
        df.at[i, 'lista_count'] = int(list_a)
        df.at[i, 'listb_count'] = int(list_b)

        if i % 25000 == 0: print('25k more processed...')
    return df

You could also write this without a for loop but using apply:

df['token_count'] = df.tokens.apply(len)
df['lista_count'] = df.tokens.apply(lambda tok: sum((t in lista) for t in tok))
df['listb_count'] = df.tokens.apply(lambda tok: sum((t in listb) for t in tok))

Answer 2

Using the code below should give you some tiny time advantage. NOTICE that your count of occurrences in listb is not correct if you use this code below as it is.

You don't get much speedup from the code below because creating the sets will take away almost all of the time advantage coming from faster lookups, BUT you will save time if you reuse seta, setb also for another lookup loops if further code.

Maybe you can arrange, that you get your lists directly as sets, so you save the time for creating the sets?

lista = ['word1a', 'word2a', 'word500a']
listb = ['word1b', 'word2a', 'word500b']

seta = set(lista)
setb = set(listb)

def token_list_count(df):

    dfIterrows = df.iterrows() # TRY THIS !!! 
    for i, t in dfIterrows:    # TRY THIS !!! 
        list_a = 0
        list_b = 0
        tTokens = t['tokens']  # TRY THIS !!!
        for tok in tTokens:    # TRY THIS !!!
            if   tok in seta: list_a += 1
            elif tok in setb: list_b += 1
            # why not "if tok in setb: list_b +=1" ???
        df.loc[i, 'token_count'] = int(len(t['tokens']))
        df.loc[i, 'lista_count'] = int(list_a)
        df.loc[i, 'listb_count'] = int(list_b)

        if i % 25000 == 0: print('25k more processed...')

    return df

If the speed really matters to you I suggest you use Python only as a frame and do the lookups using appropriate executables or a a database query in addition to multiprocessing used for lookups in seta and setb.

Another option is to consider what root suggested in a comment: """The biggest issue here is that you're iterating over a DataFrame, which should almost always be avoided. Using a vectorized approach is what will yield the best performance improvements. """