Using pyspark, how to add a column to a DataFrame as a key-value map of multiple known columns in the same DataFrame excluding nulls?

Question

Given the following example:

d = [{'asset': '2', 'ts': 6,  'B':'123','C':'234'}, 
     {'asset': '1', 'ts': 5, 'C.1':'999', 'B':'888','F':'999'}]
df = spark.createDataFrame(d)
df.show(truncate=False)

+---+----+-----+---+----+----+
|B  |C   |asset|ts |C.1 |F   |
+---+----+-----+---+----+----+
|123|234 |2    |6  |null|null|
|888|null|1    |5  |999 |999 |
+---+----+-----+---+----+----+

I want to create the following output:

+-----+---+--------------------------------+
|asset|ts |signals                         |
+-----+---+--------------------------------+
|2    |6  |[B -> 123, C -> 234]            |
|1    |5  |[B -> 888, C.1 -> 999, F -> 999]|
+-----+---+--------------------------------+

I tried the following:

from itertools import chain
from pyspark.sql.functions import *
all_signals=['B','C','C.1','F']
key_values = create_map(*(chain(*[(lit(name), col("`"+name+"`"))
                                  for name in all_signals])))

new_df = df.withColumn('signals',key_values).drop(*all_signals).show(truncate=False)

+-----+---+--------------------------------------+
|asset|ts |signals                               |
+-----+---+--------------------------------------+
|2    |6  |[B -> 123, C -> 234, C.1 ->, F ->]    |
|1    |5  |[B -> 888, C ->, C.1 -> 999, F -> 999]|
+-----+---+--------------------------------------+

But I don't want the keys with null values. So I tried so many ways to exclude null or None. I tried "if" conditions, when/otherwise but none seem to work. Here is one attempt:

key_values = create_map(*(chain(*[(lit(name), col("`"+name+"`")) 
                                  for name in all_signals 
                                  if col("`"+name+"`").isNotNull()])))
new_df = df.withColumn('signals',key_values).drop(*all_signals).show(truncate=False)


ValueError: Cannot convert column into bool: please use '&' for 'and', '|' for 'or', '~' for 'not' when building DataFrame boolean expressions.

I got it to work using a circular way which I am not happy with:

new_df= df.withColumn("signals", from_json(
                       to_json(struct(["`"+x+"`" for x in all_signals])),"MAP<STRING,STRING>"))
                      
new_df = new_df.drop(*all_signals)
new_df.show(truncate=False)

+-----+---+--------------------------------+
|asset|ts |signals                         |
+-----+---+--------------------------------+
|2    |6  |[B -> 123, C -> 234]            |
|1    |5  |[B -> 888, C.1 -> 999, F -> 999]|
+-----+---+--------------------------------+

But there must be a way to exclude null without going to json and back!

score 2 · Accepted Answer · answered Jul 20 '20 at 03:19

No need for UDF, Use higher order function filter, with arrays_zip and map_from_entries to get your desired output.(spark2.4+)

from pyspark.sql import functions as F

all_singals=['B','C','C.1','F']

df.withColumn("all_one", F.array(*[F.lit(x) for x in all_signals]))\
  .withColumn("all_two", F.array(*["`"+x+"`" for x in all_signals]))\
  .withColumn("signals", F.expr("""map_from_entries(filter(arrays_zip(all_one,all_two),x-> x.all_two is not null))"""))\
  .drop("all_one","all_two").show(truncate=False)

#+---+----+-----+---+----+----+--------------------------------+
#|B  |C   |asset|ts |C.1 |F   |signals                         |
#+---+----+-----+---+----+----+--------------------------------+
#|123|234 |2    |6  |null|null|[B -> 123, C -> 234]            |
#|888|null|1    |5  |999 |999 |[B -> 888, C.1 -> 999, F -> 999]|
#+---+----+-----+---+----+----+--------------------------------+

I could not resist but to smile at your solution. Makes me feel dumb. Impressive! — ziad.rida, Jul 20 '20 at 03:27
For some reason from_json and to_json give better performance, but this does answer my question. — ziad.rida, Aug 07 '20 at 21:41

score 1 · Answer 2 · answered Jul 20 '20 at 03:04

I have an alternative solution. First build the map with null values and then drop the null values.

from pyspark.sql.types import MapType, StringType
from pyspark.sql import functions as F

# Original dataframe
data = [{'asset': '2', 'ts': 6, 'B': '123', 'C': '234'},
        {'asset': '1', 'ts': 5, 'C.1': '999', 'B': '888', 'F': '999'}]
df = spark.createDataFrame(data)
df.show(truncate=False)

# Create a map that includes null values
# Backticks are needed because spark is weird
# https://stackoverflow.com/questions/44367019/column-name-with-dot-spark
names = ['B', 'C', 'C.1', 'F']
key_value_list = []
for name in names:
    key_value_list += [F.lit(name)]
    key_value_list += [df["`{}`".format(name)]]
map_column = F.create_map(*key_value_list)

# UDF that drops null values
remove_null_values_udf = F.udf(
    lambda d: {k: v for k, v in d.items() if v is not None},
    MapType(StringType(), StringType())
)

# Apply both of the above
df = df.withColumn('map', remove_null_values_udf(map_column)).drop(*names)
df.show()
# +-----+---+--------------------+
# |asset| ts|                 map|
# +-----+---+--------------------+
# |    2|  6|[B -> 123, C -> 234]|
# |    1|  5|[B -> 888, F -> 9...|
# +-----+---+--------------------+

This looks great. I will test its performance and let you know. You're amazing. Thanks! — ziad.rida, Jul 20 '20 at 03:17

Using pyspark, how to add a column to a DataFrame as a key-value map of multiple known columns in the same DataFrame excluding nulls?

2 Answers2