How to select multiple columns from two different Pandas dataframes in Python

user_15682659 0
Viewed 50

I have the following dataframes D1 and D2.

# importing sparksession from pyspark.sql module
from pyspark.sql import SparkSession
  
# creating sparksession and giving an app name
spark = SparkSession.builder.appName('sparkdf').getOrCreate()
  
# list  of employee data
data = [["1", "sravan", "1","2","3"],
        ["2", "ojaswi", "2","3","3"], 
        ["3", "rohith", "3","4","3"],
        ["4", "sridevi", "4","5","3"], 
        ["5", "bobby", "5","2","3"]]
  
# specify column names
columns = ['ID', 'NAME', 'Company','Company1','Company2']
  
# creating a dataframe from the lists of data
dataframe = spark.createDataFrame(data, columns)
  
dataframe.show()



data1 = [["1", "45000", "1","1","2","1"],
         ["2", "145000", "2","1","2","2"],
         ["6", "45000", "6","1","2","3"],
         ["5", "34000", "7","1","2","3"]]
  
# specify column names
columns = ['ID', 'salary', 'department','Branch','Branch1','Branch2']
  
# creating a dataframe from the lists of data
dataframe1 = spark.createDataFrame(data1, columns)
  
dataframe1.show()

dataframe = dataframe.join(dataframe1,
               dataframe.Company == dataframe1.department,
               "inner").select(dataframe.ID, dataframe.NAME, dataframe.Company,round(dataframe1.salary, 2).alias("NewFlexOne"))


display(dataframe)



+---+-------+-------+--------+--------+
| ID|   NAME|Company|Company1|Company2|
+---+-------+-------+--------+--------+
|  1| sravan|      1|       2|       3|
|  2| ojaswi|      2|       3|       3|
|  3| rohith|      3|       4|       3|
|  4|sridevi|      4|       5|       3|
|  5|  bobby|      5|       2|       3|
+---+-------+-------+--------+--------+

+---+------+----------+------+-------+-------+
| ID|salary|department|Branch|Branch1|Branch2|
+---+------+----------+------+-------+-------+
|  1| 45000|         1|     1|      2|      1|
|  2|145000|         2|     1|      2|      2|
|  6| 45000|         6|     1|      2|      3|
|  5| 34000|         7|     1|      2|      3|
+---+------+----------+------+-------+-------+

I want to join the D1 and D2 and select few columns from D1 and D2 dataframe and rename one column in D2 dataframe in a single line. Below how it is done in Pyspark.

D1 = D1.join(D2, (when(D1.Company > 15, 16).otherwise(D1.Company)) == D2.department, 'inner').select(D1.ID, D1.Name, D1.Company, D1.Company1, D1.Company2, round(D2.Branch, 2).alias("NewBranch"))


D1 = D1.join(D2, (when(D1.FlexTwo > 15, 16).otherwise(D1.company1)) == D2.department, 'inner').select(D1.ID, D1.Name, D1.Company, D1.Company1, D1.Company2,D1.NewBranch, round(D2.Branch1, 2).alias("NewBranch1"))

Below is the code I tried in Python which does not work as expected.

D1 = pd.merge(D1, D2, how='inner', left_on = np.where((D1['company']) > 15, 16, D1['company']), right_on = 'department').loc[:, ["company"  ,"ID","Name","Company1","Company2","Branch"]]

D1=round(D1.rename(columns = {'Branch':'NewBranch'}),2)

D1 = pd.merge(D1, D2, how='inner', left_on = np.where((D1['company1']) > 15, 16, D1['company1']), right_on = 'department').loc[:, ["company"  ,"ID","Name","Company1","Company2","NewBranch"]]

Kindly let me know how to achieve this in Python.

Thank you.

pyspark pandas dataframe join python
1 Answers

As per the syntax of pd.merge(), the left_on and right_on can be column names or pre-computed list of values with length same as dataframe (official documentation). So, a little work around is required. Pre-compute the condition when(D1.Company > 15, 16).otherwise(D1.Company) in a new column and use that column in join:

D1 = pd.DataFrame(data=data, columns=columns)
D2 = pd.DataFrame(data=data1, columns=columns1)

D1['Company_join'] = D1[['Company']].apply(lambda row: "16" if int(row['Company']) > 15 else row['Company'], axis=1)

D1_D2_join = pd.merge(D1, D2, how='inner', left_on="Company_join", right_on='department')

print(D1_D2_join)
>>   ID_x    NAME Company Company1 Company2 Company_join ID_y  salary department   Branch Branch1 Branch2 
>> 0    1  sravan       1        2        3            1    1   45000          1        1       2       1 
>> 1    2  ojaswi       2        3        3            2    2  145000          2        1       2       2

Rest of the things round-off and column renaming are trivial.

user_28473300
Related