CODEWITHSUNDEEP
apache-sparkpysparkapache-spark-sql
iamyojimbo
iamyojimbo

Reputation: 4669

PySpark Compare Empty Map Literal

I want to drop rows in a PySpark DataFrame where a certain column contains an empty map. How do I do this? I can't seem to declare a typed empty MapType against which to compare my column. I have seen that in Scala, you can use typedLit, but there seems to be no such equivalent in PySpark. I have also tried using lit(...) and casting to a struct<string,int> but I have found no acceptable argument for lit() (tried using None which returns null and {} which is an error).

I'm sure this is trivial but I haven't seen any docs on this!

Upvotes: 2

Views: 3116

Answers (1)

abiratsis
abiratsis

Reputation: 7316

Here is a solution using pyspark size build-in function:

from pyspark.sql.functions import col, size

df = spark.createDataFrame(
  [(1, {1:'A'} ),
  (2, {2:'B'} ), 
  (3, {3:'C'} ),
  (4, {}),
  (5, None)]
).toDF("id", "map")

df.printSchema()
# root
#  |-- id: long (nullable = true)
#  |-- map: map (nullable = true)
#  |    |-- key: long
#  |    |-- value: string (valueContainsNull = true)

df.withColumn("is_empty", size(col("map")) <= 0).show()

# +---+--------+--------+
# | id|     map|is_empty|
# +---+--------+--------+
# |  1|[1 -> A]|   false|
# |  2|[2 -> B]|   false|
# |  3|[3 -> C]|   false|
# |  4|      []|    true|
# |  5|    null|    true|
# +---+--------+--------+

Note that the condition is size <= 0 since in the case of null the function returns -1 (if the spark.sql.legacy.sizeOfNull setting is true otherwise it will return null). Here you can find more details.

Generic solution: comparing Map column and literal Map

For a more generic solution we can use the build-in function size in combination with a UDF which append the string key + value of each item into a sorted list (thank you @jxc for pointing out the problem with the previous version). The hypothesis here will be that two maps are equal when:

  1. they have the same size
  2. the string representation of key + value is identical between the items of the maps

The literal map is created from an arbitrary python dictionary combining keys and values via map_from_arrays:

from pyspark.sql.functions import udf, lit, size, when, map_from_arrays, array

df = spark.createDataFrame([
   [1, {}],
   [2, {1:'A', 2:'B', 3:'C'}],
   [3, {1:'A', 2:'B'}]
  ]).toDF("key", "map")

dict = { 1:'A' , 2:'B' }

map_keys_ = array([lit(k) for k in dict.keys()])
map_values_ = array([lit(v) for v in dict.values()])
tmp_map = map_from_arrays(map_keys_, map_values_) 

to_strlist_udf = udf(lambda d: sorted([str(k) + str(d[k]) for k in d.keys()]))

def map_equals(m1, m2):
  return when(
            (size(m1) == size(m2)) & 
            (to_strlist_udf(m1) == to_strlist_udf(m2)), True
          ).otherwise(False)

df = df.withColumn("equals", map_equals(df["map"], tmp_map))

df.show(10, False)

# +---+------------------------+------+
# |key|map                     |equals|
# +---+------------------------+------+
# |1  |[]                      |false |
# |2  |[1 -> A, 2 -> B, 3 -> C]|false |
# |3  |[1 -> A, 2 -> B]        |true  |
# +---+------------------------+------+

Note: As you can see the pyspark == operator works pretty well for array comparison as well.

Upvotes: 4

Related Questions