Pandas filtering for multiple substrings in series

前端 未结 3 1792
说谎
说谎 2020-11-22 04:08

I need to filter rows in a pandas dataframe so that a specific string column contains at least one of a list of provided substrings. The substrings may have unu

3条回答
  •  栀梦
    栀梦 (楼主)
    2020-11-22 04:32

    You could try using the Aho-Corasick algorithm. In the average case, it is O(n+m+p) where n is length of the search strings and m is the length of the searched text and p is the number of output matches.

    The Aho-Corasick algorithm is often used to find multiple patterns (needles) in an input text (the haystack).

    pyahocorasick is a Python wrapper around a C implementation of the algorithm.


    Let's compare how fast it is versus some alternatives. Below is a benchmark showing using_aho_corasick to be over 30x faster than the original method (shown in the question) on a 50K-row DataFrame test case:

    |                    |     speed factor | ms per loop |
    |                    | compared to orig |             |
    |--------------------+------------------+-------------|
    | using_aho_corasick |            30.7x |         140 |
    | using_regex        |             2.7x |        1580 |
    | orig               |             1.0x |        4300 |
    

    In [89]: %timeit using_ahocorasick(col, lst)
    10 loops, best of 3: 140 ms per loop
    
    In [88]: %timeit using_regex(col, lst)
    1 loop, best of 3: 1.58 s per loop
    
    In [91]: %timeit orig(col, lst)
    1 loop, best of 3: 4.3 s per loop
    

    Here the setup used for the benchmark. It also verifies that the output matches the result returned by orig:

    import numpy as np
    import random
    import pandas as pd
    import ahocorasick
    import re
    
    random.seed(321)
    
    def orig(col, lst):
        mask = np.logical_or.reduce([col.str.contains(i, regex=False, case=False) 
                                     for i in lst])
        return mask
    
    def using_regex(col, lst):
        """https://stackoverflow.com/a/48590850/190597 (Alex Riley)"""
        esc_lst = [re.escape(s) for s in lst]
        pattern = '|'.join(esc_lst)
        mask = col.str.contains(pattern, case=False)
        return mask
    
    def using_ahocorasick(col, lst):
        A = ahocorasick.Automaton(ahocorasick.STORE_INTS)
        for word in lst:
            A.add_word(word.lower())
        A.make_automaton() 
        col = col.str.lower()
        mask = col.apply(lambda x: bool(list(A.iter(x))))
        return mask
    
    N = 50000
    # 100 substrings of 5 characters
    lst = [''.join([chr(random.randint(0, 256)) for _ in range(5)]) for _ in range(100)]
    
    # N strings of 20 characters
    strings = [''.join([chr(random.randint(0, 256)) for _ in range(20)]) for _ in range(N)]
    # make about 10% of the strings match a string from lst; this helps check that our method works
    strings = [_ if random.randint(0, 99) < 10 else _+random.choice(lst) for _ in strings]
    
    col = pd.Series(strings)
    
    expected = orig(col, lst)
    for name, result in [('using_regex', using_regex(col, lst)),
                         ('using_ahocorasick', using_ahocorasick(col, lst))]:
        status = 'pass' if np.allclose(expected, result) else 'fail'
        print('{}: {}'.format(name, status))
    

提交回复
热议问题