Add Columns in PySpark and Add Columns containing NULLS without casting all NULLS as 0

user_19377450 0
Viewed 35

I am trying to create a new column by adding two existing columns in my dataframe, while preserving NULLS values if both existing columns contain null.

DataFrame 1:

╔═══╦═══╗
║ cola ║ colb ║
╠═══╬═══╣
║ 100 ║ 100 ║
║ null ║ 300 ║
║ 200 ║ null ║
║ 400 ║ 200 ║
║ null ║ null ║
╚═══╩═══╝

Expected Output:
╔═══╦═══╦═══╗
║ cola ║ colb ║ colc ║
╠═══╬═══╬═══╣
║ 100 ║ 100 ║ 200 ║
║ null ║ 300 ║ 300 ║
║ 200 ║ null ║ 200 ║
║ 400 ║ 200 ║ 600 ║
║ null ║ null ║ null ║
╚═══╩═══╩═══╝

I've seen a lot of similar questions to this where nulls are always transformed to 0 before the add.

pyspark python
2 Answers

You can leverage when conditionining.
See if this helps:

>>>from pyspark.sql import functions as f
>>>from pyspark.sql.functions import when

>>> data=[{'colA':10,'colB':None},{'colA':10,'colB':20},{'colA':None,'colB':None}]
>>> dataframe=spark.createDataFrame(data)
>>> dataframe.show()
+----+----+
|colA|colB|
+----+----+
|  10|null|
|  10|  20|
|null|null|
+----+----+
>>> dataframe=dataframe.withColumn("new_count",when(f.col('colA').isNull() & f.col('colB').isNull(), None).when(f.col('colA').isNull(),f.col('colB')).when(f.col('colB').isNull(),f.col('colA')).otherwise(f.col('colA')+f.col('colB')))
>>> dataframe.show()
+----+----+---------+
|colA|colB|new_count|
+----+----+---------+
|  10|null|       10|
|  10|  20|       30|
|null|null|     null|
+----+----+---------+
user_29863440

Just thought of a possible solution that should take less code:

df.withColumn('sum',fun.when(fun.col('col1').isNotNull()&fun.col('col2').isNotNull(),fun.col('col1')+fun.col('col2')).otherwise(fun.greatest(fun.col('col1'),fun.col('col2'))))
user_193774500
Related