如何使用字符串（或其他某种元数据）中的逻辑将新列添加到（PySpark）Dataframe？

Question

I am trying to add a new column to a dataframe dynamically using logic stored elsewhere.

I want to be able to loop over a new column name and the new column logic contained in an array or lsit and use these values as the parameters in the withColumn function.

Using an example dataframe from the titanic dataset I have been attempting to use the exec() function to execute a string creating a new dataframe with a column using logic defined in a string.

#create the spark titanic dataframe
import pandas as pd

data1 = {'PassengerId': {0: 1, 1: 2, 2: 3, 3: 4, 4: 5},
         'Name': {0: 'Owen', 1: 'Florence', 2: 'Laina', 3: 'Lily', 4: 'William'},
         'Sex': {0: 'male', 1: 'female', 2: 'female', 3: 'female', 4: 'male'},
         'Survived': {0: 0, 1: 1, 2: 1, 3: 1, 4: 0}}

df1 = spark.createDataFrame(pd.DataFrame(data1, columns=data1.keys()))
df1.show()

Below is a function that takes the old dataframe name, new column name and the logic used to calculate the new column. The function creates the string: df3=df1.withColumn('diff_PassengerId',df1.PassengerId)

The function then executes the string.

def testfunc(dfname,colname,col_logic):
  print("datafram:",dfname,"colname:",colname,"collogic",col_logic)
  string="df3="+dfname+".withColumn("+"'diff_PassengerId'"+","+col_logic+")"
  print(string)
  return exec(string)

testfunc('df1','diff_PassengerId','df1.PassengerId+1')

df3.show()

I expected a new dataframe df3 to be created with a new column "diff_PassangerId".

However instead I get the error on execution:

NameError: name 'df3' is not defined
---------------------------------------------------------------------------
NameError                                 Traceback (most recent call last)
<command-3662686508692761> in <module>()
      9 
     10 #df3=df1.withColumn('diff_PassengerId',df1.PassengerId)
---> 11 df3.show()

NameError: name 'df3' is not defined

I when I use the show() function inside the string ie when string="df3="+dfname+".withColumn("+"'diff_PassengerId'"+","+col_logic+").show()")

it will print the dataframe. Therfore the string is being executed. However the df3 dataframe is not being created outside of the exec function.

Any help will be highly appreciated. Many thanks.

Answer 1

One reason might be because you do not pass any dataframe to your function, and then want to execute df3 = dfname.withColumn(...) .

One option is to remove function and do the following

dfname,colname,col_logic = 'df1','diff_PassengerId','df1.PassengerId+1'
string="df3="+dfname+".withColumn("+"'diff_PassengerId'"+","+col_logic+")"
exec(string)
df3.show()

Or move execution outside the function:

def testfunc(dfname,colname,col_logic):
    print("datafram:",dfname,"colname:",colname,"collogic",col_logic)
    string="df3="+dfname+".withColumn("+"'diff_PassengerId'"+","+col_logic+")"
    print(string)
#     exec(string)
    return string
exec(testfunc('df1','diff_PassengerId','df1.PassengerId+1'))
df3.show()

In both cases you get the following output:

+-----------+--------+------+--------+----------------+
|PassengerId|    Name|   Sex|Survived|diff_PassengerId|
+-----------+--------+------+--------+----------------+
|          1|    Owen|  male|       0|               2|
|          2|Florence|female|       1|               3|
|          3|   Laina|female|       1|               4|
|          4|    Lily|female|       1|               5|
|          5| William|  male|       0|               6|
+-----------+--------+------+--------+----------------+

Or pass a dataframe to your function as input.