Write function to add columns to dataframe

Question

Hi I would like to write a function to add two columns to my dataframe below:

ID             NUMBER      OBJECT
1345471886     SIZE-43     GHJ  
1481654311     SIZE-48     IJF   
8620787660     SIZE-67     EFH

This is my code:

def add_columns(res):
    
    h = res.assign(Column_1='Value_1', Column_2='Value_2')
    return h
print(h)

I get NameError: name 'h' is not defined

My expected output would be:


ID             NUMBER      OBJECT  Column_1    Column_2
1345471886     SIZE-43     GHJ     Value_1     Value_2
1481654311     SIZE-48     IJF     Value_1     Value_2
8620787660     SIZE-67     EFH     Value_1     Value_2

Can you please help? Thanks a lot !!

score 0 · Accepted Answer · answered Sep 14 '20 at 15:58

Inspired by: Add column with constant value to pandas dataframe

Code

from io import StringIO

def add_column(df, name, value):
    '''
        Adds a single column to DataFrame
    '''
    df[name] = pd.Series([value for x in range(len(df.index))], index=df.index)
    
    return df

Usage

s = '''ID             NUMBER      OBJECT
1345471886     SIZE-43     GHJ
1481654311     SIZE-48     IJF
8620787660     SIZE-67     EFH '''

# Create DataFrame
df = df =pd.read_csv(StringIO(s), delim_whitespace=True) 

# Use function above to add 2 columns
add_column(df, "Column_1", "Value_1")
add_column(df, "Column_2", "Value_2")

print(df)

Output

          ID   NUMBER OBJECT Column_1 Column_2
0  1345471886  SIZE-43    GHJ  Value_1  Value_2
1  1481654311  SIZE-48    IJF  Value_1  Value_2
2  8620787660  SIZE-67    EFH  Value_1  Value_2

Write function to add columns to dataframe

1 Answers1