How to speed up splitting one column to create a new one?-CodePudding

There is a DataFrame containing last_element and root.

import pandas as pd

df_test = pd.DataFrame({"last_element":["2000A01.2 B003", "8000N02 B001"]*100000, "root":[None, "8000N02"]*100000})

df_test.head()

 ---- ---------------- --------- 
|    | last_element   | root    |
|---- ---------------- ---------|
|  0 | 2000A01.2 B003 |         |
|  1 | 8000N02 B001   | 8000N02 |
|  2 | 2000A01.2 B003 |         |
|  3 | 8000N02 B001   | 8000N02 |
|  4 | 2000A01.2 B003 |         |
 ---- ---------------- ---------

If root is empty, we can take it from last_element. The value before space is always the root.

I already have a solution, but want to ask if someone has a faster one as on real-world data it takes more than a minute to calculate.

Solution 1 - Function

%%timeit

df_test = pd.DataFrame({"last_element":["2000A01.2 B003", "8000N02 B001"]*100000, "root":[None, "8000N02"]*100000})

def fill_root(row):
    if pd.isna(row.root) & pd.notna(row.last_element):
        return row.last_element.split(' ')[0]
    else:
        return row.root

df_test.assign(new_root = lambda x: x.apply(fill_root, axis=1))

Time:

5.14 s ± 41.4 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

Solution 2 - Apply

%%timeit

df_test['new_root'] = df_test.apply(lambda row: row.last_element.split(' ')[0] if pd.isna(row.root) else row.root, axis=1)

Time:

3.67 s ± 21.3 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)

CodePudding user response：

You could try with np.where

df['new'] = np.where(df['root'].isna(),df['last_element'].str.split().str[0],df['root'])

df['root'].mask(df['root'].isna(),df['last_element'].str.split().str[0],inplace=True)

Timing

#%timeit df_test['new_root'] = df_test.apply(lambda row: row.last_element.split(' ')[0] if pd.isna(row.root) else row.root, axis=1)
2.97 s ± 101 ms per loop (mean ± std. dev. of 7 runs, 1 loop each)
#%timeit df['new_root'] = df['root'].mask(df['root'].isna(),df['last_element'].str.split().str[0])
170 ms ± 348 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)
#%timeit df['new_root'] =  np.where(df['root'].isna(),df['last_element'].str.split().str[0],df['root'])
172 ms ± 2.75 ms per loop (mean ± std. dev. of 7 runs, 10 loops each)

CodePudding user response：

You can use mask and df.loc

mask = df_test['root'].isnull()
df_test.loc[mask,'root'] = df_test.loc[mask, 'last_element'].str.split(' ').str[0]


# timeit
def function(df_test):
    mask = df_test['root'].isnull()
    df_test.loc[mask,'root'] = df_test.loc[mask, 'last_element'].str.split(' ').str[0]

%timeit function(df_test)
11.6 ms ± 494 µs per loop (mean ± std. dev. of 7 runs, 100 loops each)