Pandas : Apply Merge operations from a column

Question

I have a data frame, which contains three columns:

import numpy as np
import pandas as pd
np.random.seed(0)
dataframe = pd.DataFrame({'operation': ['data_a', 'data_b', 'avg', 'concat', 'sum', 'take_a', 'concat'], 
             'data_a': list(np.random.uniform(-1,1,[7,2])), 'data_b': list(np.random.uniform(-1,1,[7,2]))})

Column 'operation' represent merge column, so if there is 'data_a' value in Column 'operation', it means take that particular row's data_a value, if there is 'avg' operation, then take the average of 'data_a' and 'data_b' of that particular row so on.

What I am expecting in the output, a new column contains the values as per the operation column's merge functions:

What I have tried:

dataframe['new_column'] = 'dummy_values'

for i in range(len(dataframe)):
    
    if dataframe['operation'].iloc[i]  == 'data_a':
        dataframe['new_column'].iloc[i] = dataframe['data_a'].iloc[i]
    elif dataframe['operation'].iloc[i] == 'data_b':
        dataframe['new_column'].iloc[i] = dataframe['data_b'].iloc[i]
    elif dataframe['operation'].iloc[i] == 'avg':
        dataframe['new_column'].iloc[i] = dataframe[['data_a','data_b']].iloc[i].mean()
    elif dataframe['operation'].iloc[i] == 'sum':
        dataframe['new_column'].iloc[i] = dataframe[['data_a','data_b']].iloc[i].sum()
    elif dataframe['operation'].iloc[i] == 'concat':
        dataframe['new_column'].iloc[i] = np.concatenate([dataframe['data_a'].iloc[i], dataframe['data_b'].iloc[i]], axis=0)

That is obviously not a good way to do this, It's very slow too. I am looking for more pandas version to do this task.

Answer 1

1. Wrap the logic in a function, then .apply that function.

2. Consider then doing another function lookup using a simple dictionary. This keeps your functions very short and simple.

Something like:

def _op_get_a(row):
    return row.data_a

def _op_get_b(row):
    return row.data_b

def _op_avg(row):
    return (row.data_a + row.data_b) / 2.0

def _op_sum(row):
    return row.data_a + row.data_b

def _op_concat(row):
    return np.concatenate([row.data_a, row.data_b], axis=0)

def process_row(row):
    return {
        'data_a': _op_get_a,
        'data_b': _op_get_b,
        'avg': _op_avg,
        'sum': _op_sum,
        'concat': _op_concat,
    }[row.operation](row)

# The "axis=1" is critical to apply this row-wise rather than column-wise
dataframe['new_column'] = dataframe.apply(process_row, axis=1)

I haven't tested this code, but you get the idea. Supporting new row operations is as simple as writing it as a row-wise function and hooking it into the lookup dictionary. The code is written at the smallest possible unit and able to be individually documented without creating a big cluttered mess.