Pandas - 3 iterations per second on simple conditional

Question

The function removes rows from a pandas df if that row doesn't have the value of important_1 inside of important_2. For example if important_1 is "blue" and important_2 is "M" then that row would be removed, but if important_2 were "redbluegreen" then the row would be kept.

On my ~125mb files this code runs really slow. Only about 3 iterations a second. In this example it runs faster maybe because it is a much smaller csv. How can I speed this up?

import pandas as pd
import numpy as np
from io import StringIO


df = """
valid,important_1,important_2,unrelated_1,unrelated_2
2015-07-10 01:47:00,blue,,blabla,foobar56
2015-07-10 01:51:00,blue,M,blabla,foobar32
2015-07-10 02:37:00,blue,M,blab004la,foobar
2015-07-10 02:51:00,blue,M,blabla,foobar343
2015-07-10 03:19:00,blue,blue green,blabla,foobar
2015-07-10 03:51:00,blue,,blabla1,foobar6543
2015-07-10 04:11:00,blue,green red,blabla,foobar
2015-07-10 04:51:00,blue,red,blabla,foobar2466
2015-07-10 05:51:00,blue,blue,blabla,foobar
2015-07-10 06:27:00,blue,,blabla4,foobar
2015-07-10 06:51:00,blue,,blab605la3,foobar6543
2015-07-10 07:27:00,blue,M,blabla,foobar
2015-07-10 07:51:00,blue,M,blab445la2,foobar2334
2015-07-10 08:51:00,blue,blue green red,blabla,foobar7666
2015-07-10 09:51:00,blue,blue green,blabla,foobar
"""
df = StringIO(df)
df = pd.read_csv(df)


def remove_rows_that_dont_have_important_1_in_important_2(df):
    """
    Removes all rows of the dataframe that do not have important_1 inside of important_2.
    Note: All rows have the same value for important_1.
    Note: Every row can have a different value for important_2.
    Note: There are unrelated columns in the df that cannot be dropped from the df. But rows still can.
    """
    important_1 = df.at[0, 'important_1']  # get the icao fro the df from the first ob
    important_2_column_index = df.columns.get_loc('important_2')  # column order won't always be the same, so find the right column first
    iter_index = len(df)  # the start point of the loop

    while iter_index > 0:  # start from the end of the df and work forwards
        print(iter_index)  # DEBUG printing
        iter_index -= 1  # every loop move forward one item. This is at the top of the loop because that way every loop we ensure to move down one (or up a row depending on how you look at it).

        value_of_important_2_for_this_row = df.iat[iter_index, important_2_column_index]  # the value of important_2 for the row we are currently on

        # check for np.NaN first because it is cheaper then checking str in str
        if value_of_important_2_for_this_row is np.NaN:  # if the value of important_2 is np.NaN then there is now way important_1 can be in it
            pass

        # str in str check
        elif important_1 not in value_of_important_2_for_this_row:  # safe to assume the important_2 is off type str
            pass

        else:  # safe to assume that important_1 is indeed in important_2
            continue  # skip deletion because important_1 is in important_2 Yay!

        df.drop([iter_index], inplace=True)  # delete the row

    return df  # reset the index before returning the df


df = remove_rows_that_dont_have_important_1_in_important_2(df=df)
print(df)

OUTPUT

                  valid important_1     important_2 unrelated_1 unrelated_2
4   2015-07-10 03:19:00        blue      blue green      blabla      foobar
8   2015-07-10 05:51:00        blue            blue      blabla      foobar
13  2015-07-10 08:51:00        blue  blue green red      blabla  foobar7666
14  2015-07-10 09:51:00        blue      blue green      blabla      foobar

Answer 1

If you use the .str attribute of the column, you get most of the standard Python string functions. In particular, with Python strings you can ask if a string contains another string with the __contains__() method (i.e. the in operator):

>>> "asdf" in "asdfqwerty"
True
>>> "asdfqwerty".__contains__("asdf")  # equivalently
True

Pandas exposes this as the .contains() method on the .str attribute, as discussed in the working with text data section of the docs. Annoyingly, you cannot operate on two columns---but your code snippet specifically states:

Note: All rows have the same value for important_1.

So, you actually just need to operate on important_2 and check if the single string in important_1 is contained in each row, and that you can do with the string methods. This one-liner would do what you want:

reduced_df = df[df["important_2"].str.contains(df["important_1"][0]) == True]

Most of the work is done via

df["important_2"].str.contains(df["important_1"][0])

which is checking if the strings in important_2 have the string which is in the first row of important_1. Since your column has NaN in it¹, you will get NaN values on the comparison, so you have to specifically check if the value is equal to True (or otherwise cast to boolean) to get a boolean array you can index with. Then, using that boolean result you can index your dataframe to drop the irrelevant rows. So to fully explain the line in one sentence, it's "select the rows of my dataframe where the column important_2 contains the string in the first row of important_1".

---

¹ this is bad practice FWIW, as you're mixing datatypes--NaN is a floating point value and you have it in a column of strings. You can use empty strings for null-values sometimes with strings, but better practice altogether is to have another column which tells you whether or not a value is present; that way an empty string could still be a valid input.

Answer 2

I don't really know the libraries that you're using, but maybe the overhead of dropping rows from the CSV one by one is significant? You could try batching the drop (it looks like that drop function takes a list of indices but you're only passing it one at a time) and see if that speeds things up:

from pandas import DataFrame
from typing import List

def remove_rows_that_dont_have_important_1_in_important_2(df: DataFrame) -> DataFrame:
    """
    Removes all rows of the dataframe that do not have important_1 inside of important_2.
    Note: All rows have the same value for important_1.
    Note: Every row can have a different value for important_2.
    Note: There are unrelated columns in the df that cannot be dropped from the df. But rows still can.
    """
    important_1 = df.at[0, 'important_1']  # get the icao fro the df from the first ob
    important_2_col = df.columns.get_loc('important_2')  # column order won't always be the same, so find the right column first

    indices_to_drop: List[int] = []

    for i in range(len(df), 0, -1):
        print(i)  # DEBUG printing

        important_2  = df.iat[i, important_2_col]
        if important_2 == np.NaN or important_1 not in important_2:
            indices_to_drop.append(i)

    df.drop(indices_to_drop, inplace=True)
    return df