Pandas: how to pass multiple column values into apply function?

Question:

I have an task where I have to find a weight value in dataframe, it can be found in several columns (but not every column though)

Simplified dataframe looks like this:

          str1         str2
0  some string         2kg.
1      got 5kg  some string
2  some string  some string

Simplified example of how it works right now:

import re
import pandas as pd
from pandas import Series

regex = re.compile(r'(d+)kg')

def find_weight(row: Series) -> int | None:
    # first check in 'str1'
    result = re.search(pattern=regex, string=row['str1'])
    if result:
        return int(result.group(1))
    # if not found, check 'str2'
    result = re.search(pattern=regex, string=row['str2'])
    if result:
        return int(result.group(1))
    # if nothing found in both strings - return None
    return None


df = pd.DataFrame([['some string', '2kg.'], ['got 5kg', 'some string'], ['some string', 'some string']], columns=['str1', 'str2'])
df['weight'] = df.apply(find_weight, axis=1)
print(df)

However, in a real-life case, I have 40+ columns and weight value can be found only in 2 of them. Passing an entire row into a function every time feels wasteful. How can I pass only arguments from several columns instead of the entire row?

How to pass values only of colmns ‘str1’ and ‘str2’ into function and then apply it to every row?

Answers:

SOLUTION IS:

df['weight'] = df[['str1', 'str2']].apply(find_weight, axis=1)

Thanks to @HenryEcker!

Categories: questions Tags: , ,
Answers are sorted by their score. The answer accepted by the question owner as the best is marked with
at the top-right corner.